"""
Copyright (c) 2025 by FlashInfer team.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import functools
from types import SimpleNamespace
from typing import Callable, Dict, List, Optional, Tuple, Union

import torch

from ..api_logging import flashinfer_api
from ..trace.templates.quantize import (
    fp4_quantize_trace,
    mxfp4_quantize_trace,
    nvfp4_kv_dequantize_paged_trace,
    nvfp4_kv_quantize_trace,
    nvfp4_quantize_trace,
    silu_and_mul_nvfp4_quantize_trace,
)
from ..jit import JitSpec
from ..jit import env as jit_env
from ..jit import (
    gen_jit_spec,
    sm121a_nvcc_flags,
    sm120a_nvcc_flags,
    sm120f_nvcc_flags,
    sm110a_nvcc_flags,
    sm107a_nvcc_flags,
    sm103a_nvcc_flags,
    sm100a_nvcc_flags,
    sm100f_nvcc_flags,
    sm90a_nvcc_flags,
)
from ..jit.cpp_ext import is_cuda_version_at_least
from ..utils import (
    backend_requirement,
    device_support_pdl,
    get_compute_capability,
    get_shuffle_matrix_a_row_indices,
    get_shuffle_matrix_sf_a_row_indices,
    register_custom_op,
    register_fake_op,
    supported_compute_capability,
    round_up,
    _check_kv_layout,
    _unpack_paged_kv_cache,
)
from ..tllm_enums import SfLayout


NVFP4_QUANT_ENV_VARS = (
    "FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH",
    "FLASHINFER_NVFP4_4OVER6",
    "FLASHINFER_NVFP4_4OVER6_ERR_MODE",
    "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH",
    "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256",
)


def _compute_swizzled_layout_sf_size(total_row, total_column, row_size=128):
    padded_row = round_up(total_row, row_size)
    padded_column = round_up(total_column, 4)
    return padded_row * padded_column


def _pad_scale_factors(
    unswizzled_sf: torch.Tensor, m: int, n: int, sf_vec_size: int = 16
) -> torch.Tensor:
    """Pad scale factors tensor to meet alignment requirements.

    Args:
        unswizzled_sf (torch.Tensor): Input scale factors tensor with dtype uint8.
        m (int): M dimension.
        n (int): N dimension.
        sf_vec_size (int, optional): Scale factor vector size. Defaults to 16.

    Returns:
        torch.Tensor: Padded scale factors tensor.
    """
    factor = sf_vec_size * 4
    padded_row = round_up(m, 128)
    padded_col = round_up(n, factor)

    # Pad the input tensor to [padded_row, padded_col // scaling_vector_size]
    pad_rows = padded_row - m
    pad_cols = (padded_col - n) // sf_vec_size
    if pad_rows == 0 and pad_cols == 0:
        return unswizzled_sf
    else:
        return torch.nn.functional.pad(
            unswizzled_sf, (0, pad_cols, 0, pad_rows), mode="constant", value=0
        ).contiguous()


# E2M1 lookup table: 16 possible 4-bit values (index = 4-bit code, value = float)
# Format: bit3=sign, bits2-0=magnitude (exponent+mantissa)
_E2M1_VALUES = torch.tensor(
    [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6],
    dtype=torch.float32,
)
_e2m1_values_cache: dict = {}


def _get_e2m1_values(device: torch.device) -> torch.Tensor:
    if device not in _e2m1_values_cache:
        _e2m1_values_cache[device] = _E2M1_VALUES.to(device)
    return _e2m1_values_cache[device]


def _e2m1_and_ufp8sf_scale_to_float_cpu(
    e2m1_tensor: torch.Tensor,
    ufp8_scale_tensor: torch.Tensor,
    global_scale_tensor: Optional[torch.Tensor],
    sf_vec_size: int,
    ufp8_type: int,
    is_sf_swizzled_layout: bool,
) -> torch.Tensor:
    """Pure-PyTorch CPU-compatible dequantization fallback for arch < SM90.

    Only supports is_sf_swizzled_layout=False (linear SF layout).
    """
    if is_sf_swizzled_layout:
        raise NotImplementedError(
            "CPU fallback for e2m1_and_ufp8sf_scale_to_float does not support "
            "swizzled SF layout. Use a GPU with SM90+ for swizzled layout support."
        )

    device = e2m1_tensor.device
    m, k_half = e2m1_tensor.shape
    k = k_half * 2

    # Unpack two E2M1 nibbles per byte: low nibble = even indices, high nibble = odd
    fp4_vals = torch.empty(m, k, dtype=torch.uint8, device=device)
    fp4_vals[:, 0::2] = e2m1_tensor & 0x0F
    fp4_vals[:, 1::2] = (e2m1_tensor >> 4) & 0x0F

    # Map 4-bit codes to float via LUT
    float_vals = _get_e2m1_values(device)[fp4_vals.long()]  # [M, K]

    # Decode UFP8 scale factors
    if ufp8_type == 1:
        # E4M3: interpret raw bytes as float8_e4m3fn
        sf_float = ufp8_scale_tensor.view(torch.float8_e4m3fn).float()
    else:
        # UE8M0: 2^(byte - 127)
        sf_float = torch.pow(2.0, ufp8_scale_tensor.float() - 127.0)

    # Broadcast each SF over its sf_vec_size consecutive FP4 elements
    sf_expanded = sf_float.repeat_interleave(sf_vec_size, dim=-1)  # [M, K]

    # Apply global scale
    gs = global_scale_tensor.float().item() if global_scale_tensor is not None else 1.0

    return (float_vals * sf_expanded * gs).float()


def gen_fp4_quantization_sm100_module() -> JitSpec:
    return gen_fp4_quantization_module(sm100a_nvcc_flags, "100")


def gen_fp4_quantization_sm103_module() -> JitSpec:
    return gen_fp4_quantization_module(sm103a_nvcc_flags, "103")


def gen_fp4_quantization_sm107_module() -> JitSpec:
    from flashinfer.compilation_context import cutlass_supports_sm107

    nvcc_flags = sm107a_nvcc_flags if cutlass_supports_sm107() else sm100f_nvcc_flags
    return gen_fp4_quantization_module(nvcc_flags, "107")


def gen_fp4_quantization_sm90_module() -> JitSpec:
    return gen_fp4_quantization_module(sm90a_nvcc_flags, "90")


def gen_fp4_quantization_sm110_module() -> JitSpec:
    return gen_fp4_quantization_module(sm110a_nvcc_flags, "110")


def gen_fp4_quantization_sm120_module() -> JitSpec:
    return gen_fp4_quantization_module(sm120a_nvcc_flags, "120")


def gen_fp4_quantization_sm120f_module() -> JitSpec:
    return gen_fp4_quantization_module(sm120f_nvcc_flags, "120f")


def gen_fp4_quantization_sm121_module() -> JitSpec:
    return gen_fp4_quantization_module(sm121a_nvcc_flags, "121")


def gen_fp4_quantization_module(nvcc_flags: List[str], device_arch: str) -> JitSpec:
    return gen_jit_spec(
        f"fp4_quantization_{device_arch}",
        [
            jit_env.FLASHINFER_CSRC_DIR
            / "nv_internal/tensorrt_llm/thop/fp4Quantize.cpp",
            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/tensorrt_llm/thop/fp4Op.cpp",
            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/kernels/quantization.cu",
            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/envUtils.cpp",
            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/logger.cpp",
            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/stringUtils.cpp",
            jit_env.FLASHINFER_CSRC_DIR / "nv_internal/cpp/common/tllmException.cpp",
        ],
        extra_cuda_cflags=nvcc_flags
        + [
            "-DENABLE_BF16",
            "-DENABLE_FP8",
            "-DENABLE_FP4" if is_cuda_version_at_least("12.8") else "",
        ],
        extra_cflags=[
            "-DENABLE_BF16",
            "-DENABLE_FP8",
            "-DENABLE_FP4" if is_cuda_version_at_least("12.8") else "",
        ],
        extra_include_paths=[
            jit_env.FLASHINFER_CSRC_DIR / "nv_internal",
            jit_env.FLASHINFER_CSRC_DIR / "nv_internal" / "include",
        ],
    )


@functools.cache
def get_fp4_quantization_module(backend: str = "100"):
    backend_modules: Dict[str, Callable[..., JitSpec]] = {
        "121": gen_fp4_quantization_sm121_module,
        "120f": gen_fp4_quantization_sm120f_module,
        "120": gen_fp4_quantization_sm120_module,
        "110": gen_fp4_quantization_sm110_module,
        "107": gen_fp4_quantization_sm107_module,
        "103": gen_fp4_quantization_sm103_module,
        "100": gen_fp4_quantization_sm100_module,
        "90": gen_fp4_quantization_sm90_module,
    }

    # Prefer 'f' (family / feature-set) variant for SM12x when CUDA >= 12.9,
    # as it enables native FP4 conversion instructions (cvt.rn.satfinite.e2m1x2.f32).
    # sm_120f covers the entire SM12x family (both SM120 and SM121).
    # See: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/
    if backend in ("120", "121"):
        from ..utils import version_at_least

        if version_at_least(torch.version.cuda, "12.9"):
            backend = "120f"

    if backend not in backend_modules:
        raise ValueError(f"Invalid backend: {backend}")

    module = backend_modules[backend]().build_and_load()

    @register_custom_op(
        "flashinfer::fp4_quantize_sm100",
        mutates_args=(""),
    )
    def fp4_quantize_sm100(
        input: torch.Tensor,
        global_scale: Optional[torch.Tensor] = None,
        sf_vec_size: int = 16,
        sf_use_ue8m0: bool = False,
        is_sf_swizzled_layout: bool = True,
        is_sf_8x4_layout: bool = False,
        is_global_scale_inversed: bool = False,
        enable_pdl: Optional[bool] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Quantize input tensor to FP4 format.

        Args:
            input (torch.Tensor): Input tensor of shape [M, K] with dtype fp16/bf16/fp8_quantized.
            global_scale (torch.Tensor, optional): Global scale factor of shape [1] and dtype float32.
            sf_vec_size (int, optional): Scale factor vector size. Defaults to 16.
            sf_use_ue8m0 (bool, optional): Whether to use UE8M0 format for scale factors. Defaults to False.
            is_sf_swizzled_layout (bool, optional): Whether to use swizzled layout for scale factors. Defaults to True.
            is_sf_8x4_layout (bool, optional): Whether to use 8x4 layout or 128x4 layout for scale factors. Defaults to False.
            enable_pdl (Optional[bool], optional): Whether to enable PDL (Programmatic Dependent Launch).
                If None, automatically detects based on device capability. Defaults to None.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
                - Quantized tensor of shape [M, K/2] with dtype FLOAT4_E2M1X2
                - Scale factors tensor with shape determined by layout and sf_vec_size
        """
        if enable_pdl is None:
            enable_pdl = device_support_pdl(input.device)
        out_val = torch.empty(
            (*input.shape[:-1], input.shape[-1] // 2),
            dtype=torch.uint8,
            device=input.device,
        )
        m = input.numel() // input.shape[-1]
        k = input.shape[-1]
        if is_sf_swizzled_layout:
            out_sf_size = _compute_swizzled_layout_sf_size(
                m, k // sf_vec_size, 8 if is_sf_8x4_layout else 128
            )
            out_sf_size_padded = out_sf_size
        else:
            out_sf_size = m * k // sf_vec_size
            out_sf_size_padded = round_up(m, 16) * k // sf_vec_size
        out_sf = torch.empty(
            (out_sf_size_padded,), dtype=torch.uint8, device=input.device
        )
        module.fp4_quantize(
            input,
            global_scale,
            out_val,
            out_sf,
            sf_vec_size,
            sf_use_ue8m0,
            is_sf_swizzled_layout,
            is_sf_8x4_layout,
            is_global_scale_inversed,
            enable_pdl,
        )
        return out_val, out_sf[:out_sf_size]

    @register_fake_op("flashinfer::fp4_quantize_sm100")
    def _fake_fp4_quantize_sm100(
        input: torch.Tensor,
        global_scale: Optional[torch.Tensor] = None,
        sf_vec_size: int = 16,
        sf_use_ue8m0: bool = False,
        is_sf_swizzled_layout: bool = True,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        m, k = input.shape
        return (
            input.new_empty([m, k // 2], dtype=torch.int64),  # FLOAT4_E2M1X2
            input.new_empty([m * k // sf_vec_size], dtype=torch.int32),  # Scale factors
        )

    @register_custom_op(
        "flashinfer::mxfp4_dequantize_host",
        mutates_args=(""),
    )
    def mxfp4_dequantize_host(
        weight: torch.Tensor,
        scale: torch.Tensor,
        group_size: int = 32,
    ) -> torch.Tensor:
        out = torch.empty(
            (weight.shape[0], weight.shape[1] * 2),
            dtype=torch.float32,
            device=weight.device,
        )
        module.mxfp4_dequantize_host(
            weight,
            scale,
            out,
            group_size,
        )
        return out

    @register_fake_op("flashinfer::mxfp4_dequantize_host")
    def _fake_mxfp4_dequantize_host(
        weight: torch.Tensor,
        scale: torch.Tensor,
        group_size: int = 32,
    ) -> torch.Tensor:
        return weight.new_empty(
            [weight.shape[0], weight.shape[1] * 2], dtype=torch.float32
        )

    @register_custom_op(
        "flashinfer::block_scale_interleave_sm100",
        mutates_args=("",),
    )
    def block_scale_interleave_sm100(
        unswizzled_sf: torch.Tensor,
    ) -> torch.Tensor:
        """Swizzle block scale tensor for FP4 format.

        Args:
            unswizzled_sf (torch.Tensor): unswizzled block scale tensor with dtype uint8 or bfloat16.

        Returns:
            torch.Tensor: output tensor for swizzled block scale with dtype uint8 or bfloat16.
        """
        num_experts = unswizzled_sf.shape[0] if unswizzled_sf.dim() == 3 else 1
        expert_out_size = _compute_swizzled_layout_sf_size(
            unswizzled_sf.shape[-2], unswizzled_sf.shape[-1], 128
        )
        out = torch.empty(
            (num_experts * expert_out_size,),
            dtype=unswizzled_sf.dtype,
            device=unswizzled_sf.device,
        )
        module.block_scale_interleave_sm100(unswizzled_sf, out)
        return out

    @register_fake_op("flashinfer::block_scale_interleave_sm100")
    def _fake_block_scale_interleave_sm100(
        unswizzled_sf: torch.Tensor,
    ) -> torch.Tensor:
        return unswizzled_sf.new_empty(
            [unswizzled_sf.shape[0] * unswizzled_sf.shape[1] // 16], dtype=torch.uint8
        )

    @register_custom_op(
        "flashinfer::fp4_batched_quantize_sm100",
        mutates_args=("",),
    )
    def fp4_batched_quantize_sm100(
        input: torch.Tensor,
        global_scale: Optional[torch.Tensor] = None,
        sf_vec_size: int = 16,
        sf_use_ue8m0: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Quantize a batched tensor to FP4 (E2M1x2) with per-block scale factors.

        This function converts a float/bfloat16 (or FP8-quantized) input tensor into a
        packed FP4 tensor using the E2M1 format (two 4-bit values per byte), along with
        per-block scale factors. Scale factors are encoded as UE4M3 by default, or UE8M0
        when requested, and an optional global scale can be applied.

        Args:
            input (torch.Tensor): Input tensor of shape [B, M, K] with dtype torch.float16,
                torch.bfloat16, or an FP8-quantized dtype supported by the kernel.
            global_scale (torch.Tensor, optional): Global scale factor of shape [1] and
                dtype float32.
            sf_vec_size (int, optional): Scale-factor vector size and alignment unit along K.
                Supported/expected values:
                - 16 (NVFP4 path; supported)
                - 32 (MXFP4 path; not supported yet)
                Defaults to 16.
            sf_use_ue8m0 (bool, optional): Scale-factor encoding type.
                False → UE4M3 (default), True → UE8M0.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]:
                - self_fp4 (torch.Tensor): Packed FP4 tensor in E2M1x2 format of shape
                [B, M, K // 2] with dtype torch.uint8 (two FP4 lanes per byte).
                - self_block_scale_factors (torch.Tensor): Block scale factors with dtype
                uint8 (UE4M3 or UE8M0), laid out as a flat buffer of shape
                [B, ceil(M / 128) * 128 * ceil(K / sf_vec_size / 4) * 4].

        Notes:
            - K must be even (because outputs pack two FP4 values per byte).
            - For best performance, K should be a multiple of sf_vec_size; the scale-factor
            buffer is aligned to sf_vec_size along K, pads M to multiples of 128, and
            rounds (K / sf_vec_size) up to a multiple of 4 for storage.
            - The batch dimension B is preserved for both outputs.
        """
        b, m, k = input.shape
        out_val = torch.empty(
            (b, m, k // 2),
            dtype=torch.uint8,
            device=input.device,
        )
        out_sf = torch.empty(
            (b, _compute_swizzled_layout_sf_size(m, k // sf_vec_size, 128)),
            dtype=torch.uint8,
            device=input.device,
        )
        module.fp4_batched_quantize(
            input,
            global_scale,
            out_val,
            out_sf,
            sf_vec_size,
            sf_use_ue8m0,
        )
        return out_val, out_sf

    @register_fake_op("flashinfer::fp4_batched_quantize_sm100")
    def _fake_fp4_batched_quantize_sm100(
        input: torch.Tensor,
        global_scale: Optional[torch.Tensor] = None,
        sf_vec_size: int = 16,
        sf_use_ue8m0: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        b, m, k = input.shape
        return (
            input.new_empty([b, m, k // 2], dtype=torch.uint8),  # FLOAT4_E2M1X2
            input.new_empty(
                [b, _compute_swizzled_layout_sf_size(m, k // sf_vec_size, 128)],
                dtype=torch.uint8,
            ),  # swizzled SF buffer
        )

    @register_custom_op(
        "flashinfer::nvfp4_quant_and_per_token_scale_sm100",
        mutates_args=(""),
    )
    def nvfp4_quant_and_per_token_scale_sm100(
        input: torch.Tensor,
        scale_inv: float,
        expanded_idx_to_permuted_idx: Optional[torch.Tensor] = None,
        sf_layout: int = SfLayout.layout_linear.value,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        m, k = input.shape
        output = input.new_empty((m, k // 2), dtype=torch.uint8)
        out_scale_cols = (
            round_up(k // 16, 4)
            if sf_layout != SfLayout.layout_linear.value
            else (k // 16)
        )
        out_scale_rows = (
            round_up(
                m,
                128 if sf_layout == SfLayout.layout_128x4.value else 8,
            )
            if sf_layout != SfLayout.layout_linear.value
            else m
        )
        output_scale = input.new_zeros(
            (out_scale_rows, out_scale_cols), dtype=torch.uint8
        )
        output_per_token_scale = input.new_empty((m,), dtype=torch.float32)
        module.nvfp4_quant_and_per_token_scale(
            input,
            scale_inv,
            output,
            output_scale,
            output_per_token_scale,
            expanded_idx_to_permuted_idx,
            sf_layout,
        )
        return output, output_scale, output_per_token_scale

    @register_fake_op("flashinfer::nvfp4_quant_and_per_token_scale_sm100")
    def _fake_nvfp4_quant_and_per_token_scale_sm100(
        input: torch.Tensor,
        scale_inv: float,
        expanded_idx_to_permuted_idx: Optional[torch.Tensor] = None,
        sf_layout: int = SfLayout.layout_linear.value,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        m, k = input.shape
        out_scale_cols = (
            round_up(k // 16, 4)
            if sf_layout != SfLayout.layout_linear.value
            else (k // 16)
        )
        out_scale_rows = (
            round_up(
                m,
                128 if sf_layout == SfLayout.layout_128x4.value else 8,
            )
            if sf_layout != SfLayout.layout_linear.value
            else m
        )
        return (
            input.new_empty((m, k // 2), dtype=torch.uint8),
            input.new_empty((out_scale_rows, out_scale_cols), dtype=torch.uint8),
            input.new_empty((m,), dtype=torch.float32),
        )

    @register_custom_op(
        "flashinfer::silu_and_mul_scaled_nvfp4_experts_quantize_sm100",
        mutates_args=("",),
    )
    def silu_and_mul_scaled_nvfp4_experts_quantize_sm100(
        input: torch.Tensor,
        mask: torch.Tensor,
        global_scale: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Quantize a silu and matmul with masked batched tensor to FP4 (E2M1x2) with per-block scale factors.

        This function first does silu and matmul to a float/bfloat16 input tensor then convect the result
        into a packed FP4 tensor using the E2M1 format (two 4-bit values per byte), along with
        per-block scale factors. Scale factors are encoded as UE4M3 by default, or UE8M0
        when requested, and an optional global scale can be applied.

        Args:
            input (torch.Tensor): Input tensor of shape [B, M, K] with dtype torch.float16,
                torch.bfloat16, or an FP8-quantized dtype supported by the kernel.
            mask (torch.Tensor): mask tensor of shape [B] with dtype torch.int32.
            global_scale (torch.Tensor, optional): Global scale factor of shape [1] and
                dtype float32.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]:
                - self_fp4 (torch.Tensor): Packed FP4 tensor in E2M1x2 format of shape
                [B, M, K // 2] with dtype torch.uint8 (two FP4 lanes per byte).
                - self_block_scale_factors (torch.Tensor): Block scale factors with dtype
                uint8 (UE4M3 or UE8M0), laid out as a flat buffer of shape
                [B, ceil(M / 128) * 128 * ceil(K / sf_vec_size / 4) * 4].

        Notes:
            - K must be even (because outputs pack two FP4 values per byte).
            - For best performance, K should be a multiple of sf_vec_size; the scale-factor
            buffer is aligned to sf_vec_size along K, pads M to multiples of 128, and
            rounds (K / sf_vec_size) up to a multiple of 4 for storage.
            - The batch dimension B is preserved for both outputs.
        """
        device = input.device
        l, m, k_by_2 = input.shape
        k = k_by_2 // 2
        sf_vec_size = 16
        assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}."

        scale_k = k // sf_vec_size
        padded_k = round_up(scale_k, 4)
        padded_k_int32 = padded_k // 4
        padded_m = round_up(m, 128)
        output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8)
        output_scales = torch.empty(
            l, padded_m, padded_k_int32, device=device, dtype=torch.int32
        )

        module.silu_and_mul_scaled_nvfp4_experts_quantize(
            output.view(l * m, k // 2),
            output_scales.view(l * padded_m, padded_k_int32),
            input.view(l * m, k_by_2),
            global_scale,
            mask,
            True,
        )
        output = output.permute(1, 2, 0)
        output_scales = output_scales.view(torch.float8_e4m3fn).view(
            l, padded_m // 128, padded_k // 4, 32, 4, 4
        )
        output_scales = output_scales.permute(3, 4, 1, 5, 2, 0)
        return output, output_scales

    @register_fake_op("flashinfer::silu_and_mul_scaled_nvfp4_experts_quantize_sm100")
    def _fake_silu_and_mul_scaled_nvfp4_experts_quantize_sm100(
        input: torch.Tensor,
        mask: torch.Tensor,
        global_scale: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        device = input.device
        l, m, k_by_2 = input.shape
        k = k_by_2 // 2
        sf_vec_size = 16
        assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}."

        scale_k = k // sf_vec_size
        padded_k = round_up(scale_k, 4)
        padded_k_int32 = padded_k // 4
        padded_m = round_up(m, 128)
        output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8)
        output_scales = torch.empty(
            l, padded_m, padded_k_int32, device=device, dtype=torch.int32
        )

        output_scales = output_scales.view(torch.float8_e4m3fn).view(
            l, padded_m // 128, padded_k // 4, 32, 4, 4
        )
        output_scales = output_scales.permute(3, 4, 1, 5, 2, 0)
        return (output, output_scales)

    @register_custom_op(
        "flashinfer::scaled_fp4_grouped_quant_sm100",
        mutates_args=("",),
    )
    def scaled_fp4_grouped_quant_sm100(
        input_tensor: torch.Tensor,
        input_global_scale: torch.Tensor,
        mask: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Quantize input tensor to FP4 and return quantized tensor and scale, for
        grouped gemm inputs (e.g., grouped_gemm_nt_masked for flashinfer).
        Args:
            input: The input tensor to be quantized to FP4, with shape (l, m, k)
                l is number of groups, m is number of tokens per group, k is number of features.
            input_global_scale: A scalar scaling factor for the entire tensor, with
                shape (l,).
        Outputs:
            output: The quantized tensor in FP4, with shape (m, k // 2, l) but the physical
                layout is (l, m, k // 2). `// 2` is because two fp4 values are packed into
                an uint8.
            output_scales: The blockscale tensor in FP8-E4M3, with shape (32, 4, rm, 4, rk, l)
                but the physical layout is (l, rm, rk, 32, 4, 4).
        Note:
            For the shape of output_scales, `32 * 4 * rm` is a padded m to nearest multiple of 128.
            `4 * rk` is a padded `k // 16` to nearest multiple of 4. These layout constants are
            required by the NVIDIA Blackwell MMA operations.
        """
        device = input_tensor.device
        l, m, k = input_tensor.shape
        sf_vec_size = 16
        assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}."

        scale_k = k // sf_vec_size
        padded_k = round_up(scale_k, 4)
        padded_k_int32 = padded_k // 4
        padded_m = round_up(m, 128)
        output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8)
        output_scales = torch.empty(
            l, padded_m, padded_k_int32, device=device, dtype=torch.int32
        )

        module.silu_and_mul_scaled_nvfp4_experts_quantize(
            output.view(l * m, k // 2),
            output_scales.view(l * padded_m, padded_k_int32),
            input_tensor.view(l * m, k),
            input_global_scale,
            mask,
            False,
        )
        # The physical layout of the output is (l, m, k // 2), but we want to return a
        # logical layout (m, k // 2, l) required by the flashinfer masked group gemm.
        output = output.permute(1, 2, 0)
        # The physical layout of the output scales is already swizzled as (l, rm, rk, 32, 4, 4), a
        # requirement for the flashinfer masked group gemm, where rm=m/128 and rk=k/4. The logic
        # layout is (32, 4, rm, 4, rk, l).
        output_scales = output_scales.view(torch.float8_e4m3fn).view(
            l, padded_m // 128, padded_k // 4, 32, 4, 4
        )
        output_scales = output_scales.permute(3, 4, 1, 5, 2, 0)
        return output, output_scales

    @register_fake_op("flashinfer::scaled_fp4_grouped_quant_sm100")
    def _fake_scaled_fp4_grouped_quant_sm100(
        input_tensor: torch.Tensor,
        input_global_scale: torch.Tensor,
        mask: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        device = input_tensor.device
        l, m, k = input_tensor.shape
        sf_vec_size = 16
        assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}."

        scale_k = k // sf_vec_size
        padded_k = round_up(scale_k, 4)
        padded_k_int32 = padded_k // 4
        padded_m = round_up(m, 128)
        output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8)
        output_scales = torch.empty(
            l, padded_m, padded_k_int32, device=device, dtype=torch.int32
        )

        output = output.permute(1, 2, 0)
        output_scales = output_scales.view(torch.float8_e4m3fn).view(
            l, padded_m // 128, padded_k // 4, 32, 4, 4
        )
        output_scales = output_scales.permute(3, 4, 1, 5, 2, 0)
        return output, output_scales

    @register_custom_op(
        "flashinfer::e2m1_and_ufp8sf_scale_to_float_sm100",
        mutates_args=(""),
    )
    def e2m1_and_ufp8sf_scale_to_float_sm100(
        e2m1_tensor: torch.Tensor,
        ufp8_scale_tensor: torch.Tensor,
        global_scale_tensor: Optional[torch.Tensor] = None,
        sf_vec_size: int = 16,
        ufp8_type: int = 1,
        is_sf_swizzled_layout: bool = True,
        is_sf_8x4_layout: bool = False,
    ) -> torch.Tensor:
        """Convert E2M1 format tensor and UFP8 scale factors to float tensor.

        This function performs dequantization by converting a packed FP4 tensor in E2M1 format
        back to float values using the associated UFP8 scale factors and global scale.

        Args:
            e2m1_tensor (torch.Tensor): Packed FP4 tensor in E2M1 format of shape [M, K/2] with dtype uint8.
            ufp8_scale_tensor (torch.Tensor): Scale factors tensor in UFP8 format with dtype uint8.
            global_scale_tensor (torch.Tensor, optional): Global scale factor of shape [1] and dtype float32.
            sf_vec_size (int, optional): Scale factor vector size. Defaults to 16.
            ufp8_type (int, optional): UFP8 scale factor type (0 for UE8M0, 1 for E4M3). Defaults to 1.
            is_sf_swizzled_layout (bool, optional): Whether scale factors use swizzled layout. Defaults to True.
            is_sf_8x4_layout (bool, optional): Whether swizzled scale factors use the 8x4 layout. Defaults to False.

        Returns:
            torch.Tensor: Dequantized float tensor of shape [M, K] with dtype float32.
        """
        out = torch.zeros(
            (e2m1_tensor.shape[0], e2m1_tensor.shape[1] * 2),
            dtype=torch.float32,
            device="cpu",
        )
        module.e2m1_and_ufp8sf_scale_to_float_sm100(
            e2m1_tensor.cpu(),
            ufp8_scale_tensor.cpu().reshape(-1),
            global_scale_tensor.cpu(),
            out,
            sf_vec_size,
            ufp8_type,
            is_sf_swizzled_layout,
            is_sf_8x4_layout,
        )
        return out

    @register_fake_op("flashinfer::e2m1_and_ufp8sf_scale_to_float_sm100")
    def _fake_e2m1_and_ufp8sf_scale_to_float_sm100(
        e2m1_tensor: torch.Tensor,
        ufp8_scale_tensor: torch.Tensor,
        global_scale_tensor: Optional[torch.Tensor] = None,
        sf_vec_size: int = 16,
        ufp8_type: int = 1,
        is_sf_swizzled_layout: bool = True,
        is_sf_8x4_layout: bool = False,
    ) -> torch.Tensor:
        return e2m1_tensor.new_empty(
            [e2m1_tensor.shape[0], e2m1_tensor.shape[1] * 2], dtype=torch.float32
        )

    # Register the module
    return SimpleNamespace(
        fp4_quantize_sm100=fp4_quantize_sm100,
        block_scale_interleave_sm100=block_scale_interleave_sm100,
        e2m1_and_ufp8sf_scale_to_float_sm100=e2m1_and_ufp8sf_scale_to_float_sm100,
        mxfp4_dequantize_host=mxfp4_dequantize_host,
        fp4_batched_quantize_sm100=fp4_batched_quantize_sm100,
        nvfp4_quant_and_per_token_scale_sm100=nvfp4_quant_and_per_token_scale_sm100,
        silu_and_mul_scaled_nvfp4_experts_quantize_sm100=silu_and_mul_scaled_nvfp4_experts_quantize_sm100,
        scaled_fp4_grouped_quant_sm100=scaled_fp4_grouped_quant_sm100,
    )


@flashinfer_api(trace=silu_and_mul_nvfp4_quantize_trace)
def silu_and_mul_nvfp4_quantize(
    input: torch.Tensor,
    global_scale: torch.Tensor,
    sf_vec_size: int = 16,
    is_sf_swizzled_layout: bool = True,
    is_sf_8x4_layout: bool = False,
    enable_pdl: Optional[bool] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    r"""Apply SwiGLU and NVFP4 quantization in one CuTe-DSL kernel.

    Computes ``silu(input[..., :K]) * input[..., K:]`` before quantization.
    Requires CuTe-DSL and an SM100+ GPU.

    Parameters
    ----------
    input : torch.Tensor
        Contiguous fp16/bf16 tensor of shape ``[..., 2K]``. Leading dimensions
        are flattened into M.
    global_scale : torch.Tensor
        Float32 scale of shape ``[1]``, typically
        ``(448 * 6) / silu_and_mul(input).abs().max()``.
    sf_vec_size : int
        Scale vector size. Only 16 is supported.
    is_sf_swizzled_layout : bool
        Use a swizzled scale layout. Defaults to True.
    is_sf_8x4_layout : bool
        Use the 8x4 rather than 128x4 swizzled layout.
    enable_pdl : bool, optional
        Enable Programmatic Dependent Launch. Auto-detected when None.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        Packed FP4 values of shape ``[M, K/2]`` and their scale factors.
    """
    if sf_vec_size != 16:
        raise NotImplementedError(
            "sf_vec_size can only be 16 for silu_and_mul_nvfp4_quantize"
        )
    assert input.is_contiguous(), "input must be row-major contiguous"
    assert input.shape[-1] % 2 == 0, (
        "the last dimension must be even (gate and up concatenated)"
    )
    k = input.shape[-1] // 2
    assert k % sf_vec_size == 0

    # Reject pre-Blackwell GPUs before CuTe-DSL compilation.
    major, minor = get_compute_capability(input.device)
    if major < 10:
        raise RuntimeError(
            "silu_and_mul_nvfp4_quantize requires a Blackwell GPU (SM100+, compute "
            f"capability >= 10.0); got SM{major}{minor}."
        )

    from ..cute_dsl import is_cute_dsl_available

    if not is_cute_dsl_available():
        raise RuntimeError(
            "silu_and_mul_nvfp4_quantize requires the CuTe-DSL backend, which is not "
            "available. Please install the required dependencies."
        )
    # Qualify NVFP4 layout constants to distinguish them from MXFP4 constants.
    from .kernels import nvfp4_quantize as _nvfp4_kernels

    if not is_sf_swizzled_layout:
        sf_layout = _nvfp4_kernels.SF_LAYOUT_LINEAR
    elif is_sf_8x4_layout:
        sf_layout = _nvfp4_kernels.SF_LAYOUT_8x4
    else:
        sf_layout = _nvfp4_kernels.SF_LAYOUT_128x4

    global_scale = global_scale.to(device=input.device, dtype=torch.float32)
    # The helper returns scales in the final layout-specific 2D shape.
    return _nvfp4_kernels.silu_and_mul_nvfp4_quantize_cute_dsl(
        input,
        global_scale,
        sf_layout=sf_layout,
        enable_pdl=enable_pdl,
    )


@flashinfer_api(trace=fp4_quantize_trace)
def fp4_quantize(
    input: torch.Tensor,
    global_scale: Optional[torch.Tensor] = None,
    sf_vec_size: int = 16,
    sf_use_ue8m0: bool = False,
    is_sf_swizzled_layout: bool = True,
    is_sf_8x4_layout: bool = False,
    is_global_scale_inversed: bool = False,
    enable_pdl: Optional[bool] = None,
    backend: str = "cuda",
) -> Tuple[torch.Tensor, torch.Tensor]:
    r"""Quantize input tensor to FP4 format.

    Implements FP4 quantization that converts input tensors to a
    compressed FP4 format with associated scale factors.  Supports
    various input data types and scale-factor layouts (covering both
    NVFP4 and MXFP4 quantization recipes).

    Parameters
    ----------
    input : torch.Tensor
        Input tensor of shape ``[M, K]`` with dtype fp16/bf16/fp8_quantized.
    global_scale : torch.Tensor, optional
        Global scale factor of shape ``[1]`` and dtype ``float32``.
    sf_vec_size : int
        Scale factor vector size.  Defaults to ``16``.
    sf_use_ue8m0 : bool
        Whether to use UE8M0 format for scale factors.  Defaults to
        ``False``.
    is_sf_swizzled_layout : bool
        Whether to use the swizzled layout for scale factors.  Defaults to
        ``True``.
    is_sf_8x4_layout : bool
        Use the 8x4 swizzled layout instead of 128x4.  Defaults to
        ``False``.
    is_global_scale_inversed : bool
        When ``True``, ``global_scale`` is interpreted as the inverse scale.
        Defaults to ``False``.
    enable_pdl : bool, optional
        Whether to enable Programmatic Dependent Launch.  Auto-detected
        from device capability when ``None``.
    backend : str
        Backend to use for quantization:

        - ``"cuda"``: stable CUDA kernel (default).
        - ``"cute-dsl"``: CuTe-DSL kernel (SM100+, **experimental**).
          Supported combinations:

          * ``sf_vec_size=16, sf_use_ue8m0=False``: all layouts,
            fp16/bf16/fp8 (NVFP4).
          * ``sf_vec_size=32, sf_use_ue8m0=True``: all layouts, fp16/bf16
            (MXFP4).

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        ``(x_q, sf)`` where ``x_q`` has shape ``[M, K/2]`` with dtype
        ``FLOAT4_E2M1X2`` and ``sf`` is the scale-factor tensor whose
        shape depends on the layout and ``sf_vec_size``.

    Raises
    ------
    NotImplementedError
        If the requested feature is not implemented (e.g. BFloat16 input
        when BFloat16 is not enabled, FP8 input when FP8 is not enabled,
        or ``sf_vec_size`` other than 16 or 32).
    ValueError
        If the ``"cute-dsl"`` backend is requested for an unsupported
        parameter combination.

    Warnings
    --------
    The ``"cute-dsl"`` backend is **experimental** and not part of the
    stable API.  It may change or be removed in future versions without
    notice.
    """
    if sf_vec_size != 16 and sf_vec_size != 32:
        raise NotImplementedError("sf_vec_size can only be 16 or 32")

    # The quantize kernel reads global_scale as float32 on input's device. Normalize
    # so a bf16/fp16 or off-device scale isn't misread byte-wise / cross-device-read
    # into corrupt scale factors. This is a no-op when already float32 on input.device.
    if global_scale is not None:
        global_scale = global_scale.to(device=input.device, dtype=torch.float32)

    if backend == "cute-dsl":
        return _fp4_quantize_cute_dsl(
            input,
            global_scale,
            sf_vec_size,
            sf_use_ue8m0,
            is_sf_swizzled_layout,
            is_sf_8x4_layout,
            enable_pdl,
        )
    elif backend != "cuda":
        raise ValueError(f"Unknown backend: {backend}. Must be 'cuda' or 'cute-dsl'.")

    # for column major input, we need to transpose the input
    is_column_major = input.stride(-2) == 1
    if is_column_major:
        input = input.transpose(-2, -1)

    assert input.shape[-1] % sf_vec_size == 0
    if enable_pdl is None:
        enable_pdl = device_support_pdl(input.device)
    # get input device sm version
    major, minor = get_compute_capability(input.device)
    x_q, sf = get_fp4_quantization_module(f"{major}{minor}").fp4_quantize_sm100(
        input,
        global_scale,
        sf_vec_size,
        sf_use_ue8m0,
        is_sf_swizzled_layout,
        is_sf_8x4_layout,
        is_global_scale_inversed,
        enable_pdl,
    )
    # Swizzled sf includes row/column padding from block_scale_interleave
    # (rows to multiple of 128, cols to multiple of 4), so we use the padded
    # column count and let -1 absorb the padded row count.
    # Non-swizzled sf has exactly m * (k // sf_vec_size) elements, no padding.
    if is_sf_swizzled_layout:
        sf_cols = round_up(input.shape[-1] // sf_vec_size, 4)
        sf = sf.reshape((-1, sf_cols))
    else:
        sf = sf.reshape((-1, input.shape[-1] // sf_vec_size))
    if is_column_major:
        x_q = x_q.transpose(-2, -1)
        sf = sf.transpose(-2, -1)

    return x_q, sf


@torch.library.custom_op("flashinfer::fp4_quantize", mutates_args=())
def _fp4_quantize_custom_op(
    input: torch.Tensor,
    global_scale: Optional[torch.Tensor] = None,
    sf_vec_size: int = 16,
    sf_use_ue8m0: bool = False,
    is_sf_swizzled_layout: bool = True,
    is_sf_8x4_layout: bool = False,
    enable_pdl: Optional[bool] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    return fp4_quantize(
        input,
        global_scale,
        sf_vec_size,
        sf_use_ue8m0,
        is_sf_swizzled_layout,
        is_sf_8x4_layout,
        enable_pdl,
    )


@_fp4_quantize_custom_op.register_fake
def _fp4_quantize_fake(
    input: torch.Tensor,
    global_scale: Optional[torch.Tensor] = None,
    sf_vec_size: int = 16,
    sf_use_ue8m0: bool = False,
    is_sf_swizzled_layout: bool = True,
    is_sf_8x4_layout: bool = False,
    enable_pdl: Optional[bool] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    is_column_major = input.stride(-2) == 1
    if is_column_major:
        m = input.shape[-1]
        K = input.shape[-2]
    else:
        m = input.numel() // input.shape[-1]
        K = input.shape[-1]
    # Quantized output: 2 FP4 values packed per uint8
    if is_column_major:
        x_q = input.new_empty((*input.shape[:-2], K // 2, m), dtype=torch.uint8)
    else:
        x_q = input.new_empty((*input.shape[:-1], K // 2), dtype=torch.uint8)
    # Scale factors: shape depends on swizzled vs linear layout
    if is_sf_swizzled_layout:
        row_size = 8 if is_sf_8x4_layout else 128
        sf_rows = round_up(m, row_size)
        sf_cols = round_up(K // sf_vec_size, 4)
    else:
        sf_rows = m
        sf_cols = K // sf_vec_size
    if is_column_major:
        sf = input.new_empty((sf_cols, sf_rows), dtype=torch.uint8)
    else:
        sf = input.new_empty((sf_rows, sf_cols), dtype=torch.uint8)
    return x_q, sf


def _fp4_quantize_cute_dsl(
    input: torch.Tensor,
    global_scale: Optional[torch.Tensor],
    sf_vec_size: int,
    sf_use_ue8m0: bool,
    is_sf_swizzled_layout: bool,
    is_sf_8x4_layout: bool,
    enable_pdl: Optional[bool],
) -> Tuple[torch.Tensor, torch.Tensor]:
    """CuTe-DSL dispatch for fp4_quantize. Maps parameters to the appropriate kernel."""
    from ..cute_dsl import is_cute_dsl_available

    if not is_cute_dsl_available():
        raise RuntimeError(
            "CuTe-DSL backend requested but CuTe-DSL is not available. "
            "Please install the required dependencies."
        )

    if sf_vec_size == 16 and not sf_use_ue8m0:
        # NVFP4 path: E4M3 scale factors, sf_vec_size=16, all layouts. Import the
        # module (not the SF_LAYOUT_* names) so they are not redefined against the
        # MXFP4 branch below.
        from .kernels import nvfp4_quantize as _nvfp4_kernels

        if not is_sf_swizzled_layout:
            sf_layout = _nvfp4_kernels.SF_LAYOUT_LINEAR
        elif is_sf_8x4_layout:
            sf_layout = _nvfp4_kernels.SF_LAYOUT_8x4
        else:
            sf_layout = _nvfp4_kernels.SF_LAYOUT_128x4

        return _nvfp4_kernels.nvfp4_quantize_cute_dsl(
            input, global_scale, sf_layout=sf_layout, enable_pdl=enable_pdl
        )

    elif sf_vec_size == 32 and sf_use_ue8m0:
        # MXFP4 path: UE8M0 scale factors, sf_vec_size=32
        from .kernels.mxfp4_quantize import (
            SF_LAYOUT_128x4,
            SF_LAYOUT_8x4,
            SF_LAYOUT_LINEAR,
            mxfp4_quantize_cute_dsl,
        )

        if not is_sf_swizzled_layout:
            sf_layout = SF_LAYOUT_LINEAR
        elif is_sf_8x4_layout:
            sf_layout = SF_LAYOUT_8x4
        else:
            sf_layout = SF_LAYOUT_128x4
        return mxfp4_quantize_cute_dsl(
            input, sf_layout=sf_layout, enable_pdl=enable_pdl
        )

    else:
        raise ValueError(
            f"CuTe-DSL backend does not support sf_vec_size={sf_vec_size} with "
            f"sf_use_ue8m0={sf_use_ue8m0}. Supported: "
            f"(sf_vec_size=16, sf_use_ue8m0=False) for NVFP4, "
            f"(sf_vec_size=32, sf_use_ue8m0=True) for MXFP4."
        )


@flashinfer_api
def block_scale_interleave(unswizzled_sf: torch.Tensor) -> torch.Tensor:
    r"""Swizzle a block-scale tensor for FP4 layouts.

    Reorders an unswizzled FP4 block-scale tensor to optimize memory access
    patterns for FP4 GEMM/MoE kernels.  The output is padded in the ``m``
    dimension to a multiple of 128.

    Parameters
    ----------
    unswizzled_sf : torch.Tensor
        Input scale-factor tensor with dtype ``uint8`` or ``bfloat16``.

    Returns
    -------
    torch.Tensor
        1D flattened swizzled scale-factor buffer of shape
        ``(num_experts * expert_out_size,)`` where ``num_experts`` is
        ``unswizzled_sf.shape[0]`` for 3D inputs (and ``1`` otherwise)
        and ``expert_out_size`` is the padded swizzled size returned by
        ``_compute_swizzled_layout_sf_size``.  Note that this is *not*
        the same logical shape as ``unswizzled_sf``; downstream FP4
        GEMM/MoE kernels consume the flat buffer directly.

    Raises
    ------
    AssertionError
        If the input dtype is not ``uint8`` or ``bfloat16``.
    """
    # TODO(shuw): check input dtype is uint8
    assert (
        unswizzled_sf.dtype == torch.uint8 or unswizzled_sf.dtype == torch.bfloat16
    ), f"Input dtype must be uint8 or bfloat16, got {unswizzled_sf.dtype}"

    major, minor = get_compute_capability(unswizzled_sf.device)
    device_arch = f"{major * 10 + minor}"

    return get_fp4_quantization_module(device_arch).block_scale_interleave_sm100(
        unswizzled_sf,
    )


# Maintain compatibility with libraries using the old name
nvfp4_block_scale_interleave = block_scale_interleave


@flashinfer_api
def e2m1_and_ufp8sf_scale_to_float(
    e2m1_tensor: torch.Tensor,
    ufp8_scale_tensor: torch.Tensor,
    global_scale_tensor: Optional[torch.Tensor] = None,
    sf_vec_size: int = 16,
    ufp8_type: int = 1,
    is_sf_swizzled_layout: bool = True,
    is_sf_8x4_layout: bool = False,
) -> torch.Tensor:
    r"""Dequantize an E2M1 tensor with UFP8 scales back to float32.

    Performs dequantization by converting a packed FP4 tensor in E2M1
    format back to float values using the associated UFP8 scale factors
    and global scale.

    Parameters
    ----------
    e2m1_tensor : torch.Tensor
        Packed FP4 tensor in E2M1 format of shape ``[M, K/2]`` with dtype
        ``uint8``.
    ufp8_scale_tensor : torch.Tensor
        Scale-factor tensor in UFP8 format with dtype ``uint8``.
    global_scale_tensor : torch.Tensor, optional
        Global scale factor of shape ``[1]`` and dtype ``float32``.
    sf_vec_size : int
        Scale-factor vector size.  Defaults to ``16``.
    ufp8_type : int
        UFP8 scale-factor type (``0`` for UE8M0, ``1`` for E4M3).  Defaults
        to ``1``.
    is_sf_swizzled_layout : bool
        Whether the scale factors are stored in the swizzled layout.
        Defaults to ``True``.
    is_sf_8x4_layout : bool
        Whether swizzled scale factors use the 8x4 layout instead of the
        default 128x4 layout.  Must be ``False`` when
        ``is_sf_swizzled_layout=False``.  Defaults to ``False``.

    Returns
    -------
    torch.Tensor
        Dequantized float tensor of shape ``[M, K]`` with dtype
        ``float32``.
    """
    # NOTE(Zihao): this is another cpu op, should decouple it from cuda ops in the future
    major, minor = get_compute_capability(
        torch.device("cuda:0")
    )  # select any cuda device to get a compute capability
    if major * 10 + minor < 90:
        # No kernel available; use pure-PyTorch fallback
        return _e2m1_and_ufp8sf_scale_to_float_cpu(
            e2m1_tensor,
            ufp8_scale_tensor,
            global_scale_tensor,
            sf_vec_size,
            ufp8_type,
            is_sf_swizzled_layout,
        )
    device_arch = f"{major * 10 + minor}"
    return get_fp4_quantization_module(
        device_arch
    ).e2m1_and_ufp8sf_scale_to_float_sm100(
        e2m1_tensor,
        ufp8_scale_tensor,
        global_scale_tensor,
        sf_vec_size,
        ufp8_type,
        is_sf_swizzled_layout,
        is_sf_8x4_layout,
    )


@flashinfer_api
def shuffle_matrix_a(input_tensor: torch.Tensor, epilogue_tile_m: int) -> torch.Tensor:
    r"""PyTorch equivalent of TRT-LLM-gen ``shuffleMatrixA``.

    Parameters
    ----------
    input_tensor : torch.Tensor
        Row-major matrix to shuffle.
    epilogue_tile_m : int
        Epilogue tile size along the M dimension; determines the shuffle
        permutation.

    Returns
    -------
    torch.Tensor
        Row-shuffled copy of ``input_tensor``.
    """
    row_indices = get_shuffle_matrix_a_row_indices(input_tensor, epilogue_tile_m)

    return input_tensor[row_indices.to(input_tensor.device)]


@flashinfer_api
def shuffle_matrix_sf_a(
    input_tensor: torch.Tensor,
    epilogue_tile_m: int,
    num_elts_per_sf: int = 16,
):
    r"""CUDA implementation of TRT-LLM-gen ``shuffleMatrixSfA`` for linear-layout SF.

    Unlike upstream ``shuffleMatrixSfA`` (which both reads and writes the
    128x4 layout), this routine expects ``input_tensor`` in the *linear*
    layout that is used by quantized NVFP4 checkpoints.  No padding is
    added.

    Parameters
    ----------
    input_tensor : torch.Tensor
        Scale-factor tensor in linear layout.
    epilogue_tile_m : int
        Epilogue tile size along the M dimension; determines the row
        permutation.
    num_elts_per_sf : int
        Number of elements per scale-factor vector.  Defaults to ``16``.

    Returns
    -------
    torch.Tensor
        Row-shuffled scale-factor tensor, re-interleaved into the 128x4
        layout.
    """

    row_indices = get_shuffle_matrix_sf_a_row_indices(input_tensor, epilogue_tile_m)

    w_shuffled = input_tensor[row_indices.to(input_tensor.device)]

    # 128x4
    return block_scale_interleave(w_shuffled)


@flashinfer_api(trace=nvfp4_quantize_trace)
def nvfp4_quantize(
    a,
    a_global_sf,
    sfLayout=SfLayout.layout_128x4,
    do_shuffle=False,
    sf_vec_size=16,
    enable_pdl=None,
    backend: str = "cuda",
    per_token_activation: bool = False,
    expanded_idx_to_permuted_idx: Optional[torch.Tensor] = None,
):
    r"""Quantize input tensor to NVFP4 format.

    Parameters
    ----------
    a : torch.Tensor
        Input tensor of shape ``[M, K]`` with dtype fp16/bf16/float8_e4m3fn.
    a_global_sf : float or torch.Tensor
        Global scale factor. The CuTe-DSL backend accepts a host-side float
        or, for backward compatibility, a single-element tensor.
    sfLayout : SfLayout
        Scale-factor layout.  Defaults to ``SfLayout.layout_128x4``.
    do_shuffle : bool
        Whether to shuffle the scale factors.  Only the TRT-LLM backend
        needs to shuffle the tensor-B scale factors.  Defaults to
        ``False``.
    sf_vec_size : int
        Scale-factor vector size.  Defaults to ``16``.
    enable_pdl : bool, optional
        Whether to enable Programmatic Dependent Launch.  Auto-detected
        from device capability when ``None``.
    backend : str
        Backend to use for quantization:

        - ``"cuda"``: stable CUDA kernel (default).
        - ``"cute-dsl"``: CuTe-DSL kernel (SM100+, **experimental**);
          supports all ``sfLayout`` values
          (``layout_128x4`` / ``layout_8x4`` / ``layout_linear``)
          and input dtypes fp16/bf16/float8_e4m3fn, but only
          ``sf_vec_size == 16``.
    per_token_activation : bool
        Whether to use per-token NVFP4 activation scaling.  In this mode
        ``a_global_sf`` is the inverse base scale multiplier (typically
        ``1 / (448 * 6)``) and the function also returns per-token FP32
        scales.
    expanded_idx_to_permuted_idx : torch.Tensor, optional
        Optional row-remapping buffer for per-token activation
        quantization.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        ``(x_q, sf)`` where ``x_q`` has shape ``[M, K/2]`` with dtype
        ``FLOAT4_E2M1X2`` and ``sf`` is the scale-factor tensor (shape
        depends on the layout and ``sf_vec_size``).  When
        ``per_token_activation=True``, a third tensor containing per-token
        FP32 scales is also returned.

    Warnings
    --------
    The ``"cute-dsl"`` backend is **experimental** and not part of the
    stable API.  It may change or be removed in future versions without
    notice.
    """
    if per_token_activation:
        if sf_vec_size != 16:
            raise ValueError(
                "Per-token NVFP4 quantization only supports sf_vec_size=16"
            )

        sf_layout = SfLayout.layout_linear if do_shuffle else sfLayout
        if do_shuffle:
            assert sfLayout == SfLayout.layout_128x4

        if backend == "cuda":
            scale_inv = (
                float(a_global_sf.item())
                if isinstance(a_global_sf, torch.Tensor)
                else float(a_global_sf)
            )
            a_cuda = a.cuda()
            expanded_idx_to_permuted_idx_cuda = (
                expanded_idx_to_permuted_idx.cuda()
                if expanded_idx_to_permuted_idx is not None
                else None
            )
            major, minor = get_compute_capability(a_cuda.device)
            device_arch = f"{major * 10 + minor}"
            a_fp4, a_sf, per_token_scale = get_fp4_quantization_module(
                device_arch
            ).nvfp4_quant_and_per_token_scale_sm100(
                a_cuda,
                scale_inv,
                expanded_idx_to_permuted_idx_cuda,
                sf_layout.value,
            )
        elif backend == "cute-dsl":
            from ..cute_dsl import is_cute_dsl_available

            if expanded_idx_to_permuted_idx is not None:
                raise ValueError(
                    "CuTe-DSL per-token NVFP4 quantization does not support "
                    "expanded_idx_to_permuted_idx"
                )
            if not is_cute_dsl_available():
                raise RuntimeError(
                    "CuTe-DSL backend requested but CuTe-DSL is not available. "
                    "Please install the required dependencies."
                )
            from .kernels.nvfp4_quantize import (
                SF_LAYOUT_128x4,
                SF_LAYOUT_8x4,
                SF_LAYOUT_LINEAR,
                nvfp4_quantize_per_token_cute_dsl,
            )

            _sf_layout_map = {
                SfLayout.layout_128x4: SF_LAYOUT_128x4,
                SfLayout.layout_8x4: SF_LAYOUT_8x4,
                SfLayout.layout_linear: SF_LAYOUT_LINEAR,
            }
            a_fp4, a_sf, per_token_scale = nvfp4_quantize_per_token_cute_dsl(
                a.cuda(),
                (
                    a_global_sf.cuda()
                    if isinstance(a_global_sf, torch.Tensor)
                    else a_global_sf
                ),
                sf_layout=_sf_layout_map[sf_layout],
                enable_pdl=enable_pdl,
            )
        else:
            raise ValueError(
                f"Unknown backend: {backend}. Must be 'cuda' or 'cute-dsl'."
            )
    elif backend == "cuda":
        if expanded_idx_to_permuted_idx is not None:
            raise ValueError(
                "expanded_idx_to_permuted_idx is only supported with per_token_activation=True"
            )
        if do_shuffle:
            assert sfLayout == SfLayout.layout_128x4
            is_sf_swizzled_layout = False
            is_sf_8x4_layout = False
        else:
            is_sf_swizzled_layout = sfLayout != SfLayout.layout_linear
            is_sf_8x4_layout = sfLayout == SfLayout.layout_8x4

        a_fp4, a_sf = fp4_quantize(
            a.cuda(),
            a_global_sf.cuda(),
            sf_vec_size,
            sf_use_ue8m0=False,
            is_sf_swizzled_layout=is_sf_swizzled_layout,
            is_sf_8x4_layout=is_sf_8x4_layout,
            enable_pdl=enable_pdl,
        )
    elif backend == "cute-dsl":
        from ..cute_dsl import is_cute_dsl_available

        if not is_cute_dsl_available():
            raise RuntimeError(
                "CuTe-DSL backend requested but CuTe-DSL is not available. "
                "Please install the required dependencies."
            )
        if sf_vec_size != 16:
            raise ValueError(
                f"CuTe-DSL backend only supports sf_vec_size=16, got {sf_vec_size}"
            )
        from .kernels.nvfp4_quantize import (
            SF_LAYOUT_128x4,
            SF_LAYOUT_8x4,
            SF_LAYOUT_LINEAR,
            nvfp4_quantize_cute_dsl,
        )

        _sf_layout_map = {
            SfLayout.layout_128x4: SF_LAYOUT_128x4,
            SfLayout.layout_8x4: SF_LAYOUT_8x4,
            SfLayout.layout_linear: SF_LAYOUT_LINEAR,
        }
        if do_shuffle:
            assert sfLayout == SfLayout.layout_128x4
            sf_layout_int = SF_LAYOUT_LINEAR
        else:
            sf_layout_int = _sf_layout_map[sfLayout]

        a_fp4, a_sf = nvfp4_quantize_cute_dsl(
            a.cuda(), a_global_sf, sf_layout=sf_layout_int, enable_pdl=enable_pdl
        )
    else:
        raise ValueError(f"Unknown backend: {backend}. Must be 'cuda' or 'cute-dsl'.")

    if do_shuffle:
        epilogue_tile_m = 128
        if per_token_activation:
            row_indices = get_shuffle_matrix_a_row_indices(a_fp4, epilogue_tile_m).to(
                a_fp4.device
            )
            a_fp4 = a_fp4[row_indices]
            per_token_scale = per_token_scale[row_indices]
        else:
            a_fp4 = shuffle_matrix_a(a_fp4.view(torch.uint8), epilogue_tile_m)
        a_sf = shuffle_matrix_sf_a(a_sf.view(torch.uint8), epilogue_tile_m).reshape(
            a_sf.shape
        )

    if per_token_activation:
        return a_fp4, a_sf, per_token_scale
    return a_fp4, a_sf


@flashinfer_api(trace=mxfp4_quantize_trace)
def mxfp4_quantize(
    a: torch.Tensor,
    backend: str = "cuda",
    enable_pdl: Optional[bool] = None,
    sfLayout: SfLayout = SfLayout.layout_128x4,
) -> Tuple[torch.Tensor, torch.Tensor]:
    r"""Quantize input tensor to MXFP4 format.

    Parameters
    ----------
    a : torch.Tensor
        Input tensor of shape ``[M, K]`` with dtype fp16/bf16.
    backend : str
        Backend to use for quantization:

        - ``"cuda"``: stable CUDA kernel (default).
        - ``"cute-dsl"``: CuTe-DSL kernel (SM100+, **experimental**).
    enable_pdl : bool, optional
        Whether to enable Programmatic Dependent Launch.  Only used when
        ``backend == "cute-dsl"``.  Auto-detected from device capability
        when ``None``.
    sfLayout : SfLayout
        Scale-factor layout.  Defaults to ``SfLayout.layout_128x4``.
        Supported layouts are ``layout_128x4``, ``layout_8x4``, and
        ``layout_linear``.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        ``(x_q, sf)`` where ``x_q`` has shape ``[M, K/2]`` with dtype
        ``uint8`` (``FLOAT4_E2M1X2``) and ``sf`` is the UE8M0 scale-factor
        tensor (``uint8``) whose shape depends on the chosen layout and
        ``sf_vec_size`` (fixed at ``32`` here).

    Warnings
    --------
    The ``"cute-dsl"`` backend is **experimental** and not part of the
    stable API.  It may change or be removed in future versions without
    notice.  Use at your own risk for production workloads.
    """
    sfLayout = SfLayout(sfLayout)
    if backend == "cute-dsl":
        from ..cute_dsl import is_cute_dsl_available

        if not is_cute_dsl_available():
            raise RuntimeError(
                "CuTe-DSL backend requested but CuTe-DSL is not available. "
                "Please install the required dependencies."
            )
        from .kernels.mxfp4_quantize import (
            SF_LAYOUT_128x4,
            SF_LAYOUT_8x4,
            SF_LAYOUT_LINEAR,
            mxfp4_quantize_cute_dsl,
        )

        _sf_layout_map = {
            SfLayout.layout_128x4: SF_LAYOUT_128x4,
            SfLayout.layout_8x4: SF_LAYOUT_8x4,
            SfLayout.layout_linear: SF_LAYOUT_LINEAR,
        }

        return mxfp4_quantize_cute_dsl(
            a, sf_layout=_sf_layout_map[sfLayout], enable_pdl=enable_pdl
        )
    elif backend == "cuda":
        is_sf_swizzled_layout = sfLayout != SfLayout.layout_linear
        is_sf_8x4_layout = sfLayout == SfLayout.layout_8x4
        a_fp4, a_sf = fp4_quantize(
            a.cuda(),
            global_scale=None,
            sf_vec_size=32,
            sf_use_ue8m0=True,
            is_sf_swizzled_layout=is_sf_swizzled_layout,
            is_sf_8x4_layout=is_sf_8x4_layout,
        )
        return a_fp4, a_sf
    else:
        raise ValueError(f"Unknown backend: {backend}. Must be 'cuda' or 'cute-dsl'.")


@flashinfer_api
def mxfp4_dequantize(
    a_fp4,
    a_sf,
    sfLayout: SfLayout = SfLayout.layout_128x4,
):
    r"""Dequantize MXFP4 packed weights back to float32.

    Parameters
    ----------
    a_fp4 : torch.Tensor
        Quantized tensor of shape ``[M, K/2]`` with dtype ``uint8``
        (``FLOAT4_E2M1X2``).
    a_sf : torch.Tensor
        UE8M0 scale-factor tensor (``uint8``); shape depends on ``sfLayout``.
    sfLayout : SfLayout
        Scale-factor layout used by ``a_sf``.  Defaults to
        ``SfLayout.layout_128x4``.

    Returns
    -------
    torch.Tensor
        Dequantized tensor of shape ``[M, K]`` with dtype ``float32``.
    """
    sfLayout = SfLayout(sfLayout)
    return e2m1_and_ufp8sf_scale_to_float(
        a_fp4.cpu().view(torch.uint8),
        a_sf.cpu().view(torch.uint8).reshape(-1),
        torch.tensor([1.0], device=a_fp4.device),
        sf_vec_size=32,
        ufp8_type=0,
        is_sf_swizzled_layout=sfLayout != SfLayout.layout_linear,
        is_sf_8x4_layout=sfLayout == SfLayout.layout_8x4,
    )


@flashinfer_api
def mxfp4_dequantize_host(
    weight: torch.Tensor,
    scale: torch.Tensor,
    group_size: int = 32,
) -> torch.Tensor:
    r"""Host-side MXFP4 dequantization.

    Parameters
    ----------
    weight : torch.Tensor
        Quantized tensor of shape ``[M, K/2]`` with dtype ``uint8``
        (``FLOAT4_E2M1X2``).
    scale : torch.Tensor
        UE8M0 scale-factor tensor (``uint8``); shape depends on the
        layout and ``group_size`` / ``sf_vec_size`` (typically the
        swizzled buffer produced by :func:`mxfp4_quantize`).
    group_size : int
        Group size for dequantization.  Defaults to ``32``.

    Returns
    -------
    torch.Tensor
        Dequantized tensor of shape ``[M, K]`` with dtype ``float32``.
    """
    # NOTE(Zihao): the cpu op should be decouplied from cuda ops because it's device independent, should refactor this in the future
    major, minor = get_compute_capability(
        torch.device("cuda:0")
    )  # use any cuda device to get a compute capability
    device_arch = f"{major * 10 + minor}"
    return get_fp4_quantization_module(device_arch).mxfp4_dequantize_host(
        weight,
        scale,
        group_size,
    )


@flashinfer_api
def nvfp4_batched_quantize(
    a,
    a_global_sf,
    sf_vec_size=16,
):
    r"""Quantize batched input tensor to NVFP4 format.

    Parameters
    ----------
    a : torch.Tensor
        Input tensor of shape ``[B, M, K]`` with dtype fp16/bf16.
    a_global_sf : torch.Tensor
        Global scale factor of shape ``[1]`` with dtype ``float32``.
    sf_vec_size : int
        Scale-factor vector size.  Defaults to ``16``.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        ``(x_q, sf)`` where ``x_q`` has shape ``[B, M, K/2]`` with dtype
        ``FLOAT4_E2M1X2`` and ``sf`` is the per-batch swizzled
        scale-factor tensor of shape
        ``[B, ceil(M / 128) * 128 * ceil(K / sf_vec_size / 4) * 4]``
        (M is padded to a multiple of 128 and ``K / sf_vec_size`` is
        rounded up to a multiple of 4).
    """
    major, minor = get_compute_capability(a.device)
    device_arch = f"{major * 10 + minor}"
    a_fp4, a_sf = get_fp4_quantization_module(device_arch).fp4_batched_quantize_sm100(
        a,
        a_global_sf,
        sf_vec_size,
        False,
    )
    return a_fp4, a_sf


@flashinfer_api
def nvfp4_quantize_paged_kv_cache(
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    kv_layout: str = "HND",
    k_global_sf: Optional[torch.Tensor] = None,
    v_global_sf: Optional[torch.Tensor] = None,
) -> Tuple[
    Tuple[torch.Tensor, torch.Tensor],
    Tuple[torch.Tensor, torch.Tensor],
    float,
    float,
]:
    r"""Quantize a paged KV cache to NVFP4 for the trtllm-gen MHA kernel.

    Quantizes BF16/FP16 K/V caches to NVFP4 with two-level scaling (global
    FP32 + per-block FP8) and swizzles the scale factors for the SM100
    trtllm-gen MHA kernel layout.

    Parameters
    ----------
    k_cache : torch.Tensor
        Key cache tensor.  ``HND`` layout:
        ``[num_pages, num_kv_heads, page_size, head_dim]``;
        ``NHD`` layout: ``[num_pages, page_size, num_kv_heads, head_dim]``.
    v_cache : torch.Tensor
        Value cache tensor (same layout as ``k_cache``).
    kv_layout : str
        Layout of the input KV cache, either ``"HND"`` or ``"NHD"``.
    k_global_sf : torch.Tensor, optional
        Global scale factor for K (float32 scalar tensor).  When ``None``,
        auto-computed as ``FLOAT8_E4M3_MAX / k_amax``.
    v_global_sf : torch.Tensor, optional
        Global scale factor for V (float32 scalar tensor).  When ``None``,
        auto-computed as ``FLOAT8_E4M3_MAX / v_amax``.

    Returns
    -------
    Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor], float, float]
        ``(kv_cache_fp4, kv_cache_sf, k_global_scale, v_global_scale)``
        where ``kv_cache_fp4`` is ``(k_fp4, v_fp4)`` in the same layout as
        the input with ``head_dim`` replaced by ``head_dim // 2``, dtype
        ``uint8``; ``kv_cache_sf`` is ``(k_scales, v_scales)``
        (``k_scales`` keeps the linear input layout, ``v_scales`` uses
        TRT-LLM's 4-token interleaved layout, both with ``head_dim``
        replaced by ``head_dim // 16``, dtype ``float8_e4m3fn``); and the
        two trailing floats are ``1 / k_global_sf`` and ``1 / v_global_sf``.
    """
    _FLOAT8_E4M3_MAX = 448.0  # torch.finfo(torch.float8_e4m3fn).max
    # Extract dimensions based on layout
    if kv_layout == "NHD":
        num_pages, page_size, num_kv_heads, head_dim = k_cache.shape
    else:
        num_pages, num_kv_heads, page_size, head_dim = k_cache.shape

    device = k_cache.device
    scale_dim = head_dim // 16

    # Compute global scale factors if not provided.
    # global_sf = FLOAT8_E4M3_MAX / tensor_amax (no *E2M1_MAX factor).
    # This will let the fp4_quantize kernel output the scale factor in range [0, FLOAT8_E4M3_MAX/E2M1_MAX]
    if k_global_sf is None:
        k_amax = k_cache.float().abs().amax()
        k_global_sf = torch.tensor(
            [_FLOAT8_E4M3_MAX / max(k_amax.item(), 1e-12)],
            dtype=torch.float32,
            device=device,
        )
    if v_global_sf is None:
        v_amax = v_cache.float().abs().amax()
        v_global_sf = torch.tensor(
            [_FLOAT8_E4M3_MAX / max(v_amax.item(), 1e-12)],
            dtype=torch.float32,
            device=device,
        )

    # Flatten to 2D [total_tokens, head_dim] for fp4_quantize
    # Both layouts flatten identically since total elements are the same
    k_2d = k_cache.reshape(-1, head_dim)
    v_2d = v_cache.reshape(-1, head_dim)

    # Quantize using FlashInfer's GPU kernel with linear scale layout
    k_packed, k_sf = fp4_quantize(
        k_2d, k_global_sf, sf_vec_size=16, is_sf_swizzled_layout=False
    )
    v_packed, v_sf = fp4_quantize(
        v_2d, v_global_sf, sf_vec_size=16, is_sf_swizzled_layout=False
    )

    # fp4_quantize returns uint8 packed FP4 and uint8 scale factors (FP8 E4M3 encoded)
    # Reshape packed data and scale factors back to the original layout
    if kv_layout == "NHD":
        out_shape_fp4 = (num_pages, page_size, num_kv_heads, head_dim // 2)
        out_shape_sf = (num_pages, page_size, num_kv_heads, scale_dim)
    else:
        out_shape_fp4 = (num_pages, num_kv_heads, page_size, head_dim // 2)
        out_shape_sf = (num_pages, num_kv_heads, page_size, scale_dim)

    kv_cache_fp4 = (
        k_packed.view(torch.uint8).reshape(out_shape_fp4),
        v_packed.view(torch.uint8).reshape(out_shape_fp4),
    )

    # Reshape scale factors (FP8 E4M3 encoded as uint8)
    k_sf_fp8 = k_sf.view(torch.float8_e4m3fn).reshape(out_shape_sf)
    v_sf_fp8 = v_sf.view(torch.float8_e4m3fn).reshape(out_shape_sf)

    # Apply V scale factor swizzling for SM100 trtllm-gen MHA kernel.
    # The swizzle interleaves the token dimension by groups of 4 within each
    # [page_size, head_dim//16] tile per page/head:
    #   output[..., (t//4)*4*S + s*4 + t%4] = input[..., t*S + s]
    # This matches TRT-LLM's quantizeAndWriteFP4KVCache() V swizzle pattern.
    # K scale factors do NOT need swizzling — the kernel reads them with real strides.
    if page_size % 4 != 0 or head_dim % 64 != 0:
        raise ValueError(
            "V-scale swizzling requires page_size % 4 == 0 and head_dim % 64 == 0, "
            f"got page_size={page_size}, head_dim={head_dim}."
        )
    if kv_layout == "NHD":
        swizzle_shape = (
            num_pages,
            page_size // 4,
            4,
            num_kv_heads,
            4,
            scale_dim // 4,
        )
        swizzle_perm = (0, 1, 4, 3, 5, 2)
    else:
        swizzle_shape = (
            num_pages,
            num_kv_heads,
            page_size // 4,
            4,
            4,
            scale_dim // 4,
        )
        swizzle_perm = (0, 1, 2, 4, 5, 3)

    v_sf_fp8 = (
        v_sf_fp8.reshape(swizzle_shape)
        .permute(swizzle_perm)
        .reshape(out_shape_sf)
        .contiguous()
    )

    kv_cache_sf = (k_sf_fp8, v_sf_fp8)

    # Return the inverse of global_sf: global_scale = 1 / global_sf = amax / 448
    k_gs_ret = (1.0 / k_global_sf).item()
    v_gs_ret = (1.0 / v_global_sf).item()

    return kv_cache_fp4, kv_cache_sf, k_gs_ret, v_gs_ret


@flashinfer_api
def scaled_fp4_grouped_quantize(
    a,
    mask,
    a_global_sf,
):
    r"""Quantize a batched input tensor to NVFP4 with a per-row mask.

    Parameters
    ----------
    a : torch.Tensor
        Input tensor of shape ``[B, M, K]`` with dtype fp16/bf16.
    mask : torch.Tensor
        Mask tensor applied before quantization.
    a_global_sf : torch.Tensor
        Global scale factor of shape ``[1]`` with dtype ``float32``.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        ``(x_q, sf)`` where ``x_q`` has logical shape ``[M, K/2, B]``
        with dtype ``FLOAT4_E2M1X2`` (the implementation permutes the
        ``[B, M, K/2]`` physical layout so the batch dim is last, as
        required by FlashInfer's masked grouped GEMM), and ``sf`` is
        the 6D swizzled scale-factor tensor of logical shape
        ``[32, 4, padded_M // 128, 4, padded_K // 64, B]`` viewed as
        ``float8_e4m3fn``.  ``padded_M`` rounds ``M`` up to a multiple
        of 128 and ``padded_K`` rounds ``K // sf_vec_size`` (with
        ``sf_vec_size = 16``) up to a multiple of 4.
    """
    major, minor = get_compute_capability(a.device)
    device_arch = f"{major * 10 + minor}"
    a_fp4, a_sf = get_fp4_quantization_module(
        device_arch
    ).scaled_fp4_grouped_quant_sm100(
        a,
        a_global_sf,
        mask,
    )
    return a_fp4, a_sf


# ---------------------------------------------------------------------------
# NVFP4 KV cache quant/dequant with linear (non-swizzled) block scale layout
# ---------------------------------------------------------------------------


@functools.cache
def get_fp4_kv_dequantization_module():
    from ..jit.fp4_kv_dequantization import gen_fp4_kv_dequantization_module

    module = gen_fp4_kv_dequantization_module().build_and_load()

    @register_custom_op(
        "flashinfer::nvfp4_kv_dequant",
        mutates_args=("output",),
    )
    def nvfp4_kv_dequant(
        fp4_data: torch.Tensor,
        block_scales: torch.Tensor,
        global_scale: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        module.nvfp4_kv_dequant(fp4_data, block_scales, global_scale, output)

    @register_fake_op("flashinfer::nvfp4_kv_dequant")
    def _fake_nvfp4_kv_dequant(
        fp4_data: torch.Tensor,
        block_scales: torch.Tensor,
        global_scale: torch.Tensor,
        output: torch.Tensor,
    ) -> None:
        pass

    @register_custom_op(
        "flashinfer::nvfp4_paged_kv_dequant",
        mutates_args=("output_k", "output_v"),
    )
    def nvfp4_paged_kv_dequant(
        paged_k_cache: torch.Tensor,
        paged_v_cache: torch.Tensor,
        k_scales: torch.Tensor,
        v_scales: torch.Tensor,
        block_tables: torch.Tensor,
        seq_lens: torch.Tensor,
        k_global_scale: torch.Tensor,
        v_global_scale: torch.Tensor,
        output_k: torch.Tensor,
        output_v: torch.Tensor,
        kv_layout: int,
    ) -> None:
        module.nvfp4_paged_kv_dequant(
            paged_k_cache,
            paged_v_cache,
            k_scales,
            v_scales,
            block_tables,
            seq_lens,
            k_global_scale,
            v_global_scale,
            output_k,
            output_v,
            kv_layout,
        )

    @register_fake_op("flashinfer::nvfp4_paged_kv_dequant")
    def _fake_nvfp4_paged_kv_dequant(
        paged_k_cache: torch.Tensor,
        paged_v_cache: torch.Tensor,
        k_scales: torch.Tensor,
        v_scales: torch.Tensor,
        block_tables: torch.Tensor,
        seq_lens: torch.Tensor,
        k_global_scale: torch.Tensor,
        v_global_scale: torch.Tensor,
        output_k: torch.Tensor,
        output_v: torch.Tensor,
        kv_layout: int,
    ) -> None:
        pass

    return SimpleNamespace(
        nvfp4_kv_dequant=nvfp4_kv_dequant,
        nvfp4_paged_kv_dequant=nvfp4_paged_kv_dequant,
    )


@functools.cache
def get_fp4_kv_quantization_module():
    from ..jit.fp4_kv_quantization import gen_fp4_kv_quantization_module

    module = gen_fp4_kv_quantization_module().build_and_load()

    @register_custom_op(
        "flashinfer::nvfp4_kv_quant",
        mutates_args=("fp4_output", "block_scales"),
    )
    def nvfp4_kv_quant(
        input: torch.Tensor,
        global_scale: torch.Tensor,
        fp4_output: torch.Tensor,
        block_scales: torch.Tensor,
    ) -> None:
        module.nvfp4_kv_quant(input, global_scale, fp4_output, block_scales)

    @register_fake_op("flashinfer::nvfp4_kv_quant")
    def _fake_nvfp4_kv_quant(
        input: torch.Tensor,
        global_scale: torch.Tensor,
        fp4_output: torch.Tensor,
        block_scales: torch.Tensor,
    ) -> None:
        pass

    return SimpleNamespace(nvfp4_kv_quant=nvfp4_kv_quant)


_NVFP4_BLOCK_SIZE = 16


@supported_compute_capability([80, 86, 89, 90, 100, 103, 107, 110, 120, 121])
def _nvfp4_kv_dequant_check(fp4_data, block_scales, global_scale, output_dtype=None):
    return True


@backend_requirement({}, common_check=_nvfp4_kv_dequant_check)
@flashinfer_api
def nvfp4_kv_dequantize(
    fp4_data: torch.Tensor,
    block_scales: torch.Tensor,
    global_scale: torch.Tensor,
    output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    r"""GPU dequantization of an NVFP4 KV cache with linear block-scale layout.

    Requires SM80+.

    Parameters
    ----------
    fp4_data : torch.Tensor
        Packed FP4 data of shape ``[M, K/2]`` with dtype ``uint8``.
    block_scales : torch.Tensor
        Per-block FP8 E4M3 scales of shape ``[M, K/16]`` with dtype
        ``uint8``.
    global_scale : torch.Tensor
        Global scale factor of shape ``[1]`` with dtype ``float32``, on the
        same CUDA device as ``fp4_data``.
    output_dtype : torch.dtype
        Output dtype, either ``torch.bfloat16`` or ``torch.float16``.

    Returns
    -------
    torch.Tensor
        Dequantized tensor of shape ``[M, K]`` with the specified output
        dtype.
    """
    M = fp4_data.size(0)
    K = fp4_data.size(1) * 2
    if K % _NVFP4_BLOCK_SIZE != 0:
        raise ValueError(f"K dimension ({K}) must be divisible by {_NVFP4_BLOCK_SIZE}")
    output = torch.empty((M, K), dtype=output_dtype, device=fp4_data.device)
    get_fp4_kv_dequantization_module().nvfp4_kv_dequant(
        fp4_data, block_scales, global_scale, output
    )
    return output


@supported_compute_capability([80, 86, 89, 90, 100, 103, 107, 110, 120, 121])
def _nvfp4_paged_kv_dequant_check(*args, **kwargs):
    return True


@backend_requirement({}, common_check=_nvfp4_paged_kv_dequant_check)
@flashinfer_api(trace=nvfp4_kv_dequantize_paged_trace)
def nvfp4_kv_dequantize_paged(
    paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
    kv_cache_sf: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
    block_tables: torch.Tensor,
    seq_lens: torch.Tensor,
    k_scale: torch.Tensor,
    v_scale: torch.Tensor,
    output_k: torch.Tensor,
    output_v: torch.Tensor,
    kv_layout: str = "NHD",
) -> None:
    r"""Dequantize a paged NVFP4 KV cache into caller-owned contiguous outputs.

    Requires SM80+. This helper gathers pages through ``block_tables`` and
    writes dequantized K/V tensors in ``[batch, max_seq_len, num_heads,
    head_dim]`` layout. Tokens at positions ``>= seq_lens[batch]`` are left
    unchanged.

    Parameters
    ----------
    paged_kv_cache : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        The packed NVFP4 paged KV cache. Accepts the same tuple or stacked
        cache format as paged attention APIs. For tuple input, each tensor has
        shape ``[num_pages, page_size, num_kv_heads, head_dim // 2]`` when
        ``kv_layout="NHD"`` and ``[num_pages, num_kv_heads, page_size,
        head_dim // 2]`` when ``kv_layout="HND"``.
    kv_cache_sf : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
        Per-block FP8 E4M3 scales with the same tuple or stacked cache format
        as ``paged_kv_cache``, replacing ``head_dim // 2`` with
        ``head_dim // 16``.
    block_tables : torch.Tensor
        Physical page table of shape ``[batch, max_pages_per_request]`` with
        dtype ``int32`` or ``int64``.
    seq_lens : torch.Tensor
        Sequence lengths of shape ``[batch]`` with dtype ``int32``.
    k_scale, v_scale : torch.Tensor
        Global dequantization scale tensors of dtype ``float32`` on the same
        CUDA device as the cache.
    output_k, output_v : torch.Tensor
        Caller-owned output tensors in ``[batch, max_seq_len, num_heads,
        head_dim]`` layout. Each must be contiguous and have dtype
        ``torch.float16`` or ``torch.bfloat16``.
    kv_layout : str
        Layout of the paged input cache, either ``"NHD"`` or ``"HND"``.

    Returns
    -------
    None
        This function writes dequantized K/V values into ``output_k`` and
        ``output_v`` in place.
    """
    _check_kv_layout(kv_layout)
    paged_k_cache, paged_v_cache = _unpack_paged_kv_cache(paged_kv_cache, kv_layout)
    k_scales, v_scales = _unpack_paged_kv_cache(kv_cache_sf, kv_layout)

    if seq_lens.dtype != torch.int32:
        raise ValueError(f"seq_lens must have dtype torch.int32, got {seq_lens.dtype}")
    if block_tables.dtype not in (torch.int32, torch.int64):
        raise ValueError(
            f"block_tables must have dtype torch.int32 or torch.int64, got {block_tables.dtype}"
        )
    if k_scale.dtype != torch.float32 or v_scale.dtype != torch.float32:
        raise ValueError("k_scale and v_scale must have dtype torch.float32")
    if k_scale.numel() != 1 or v_scale.numel() != 1:
        raise ValueError("k_scale and v_scale must be scalar tensors")
    if output_k.dtype != output_v.dtype:
        raise ValueError("output_k and output_v must have the same dtype")
    if output_k.dtype not in (torch.float16, torch.bfloat16):
        raise ValueError(
            f"output dtype must be torch.float16 or torch.bfloat16, got {output_k.dtype}"
        )
    if output_k.ndim != 4 or output_v.ndim != 4:
        raise ValueError("output_k and output_v must be 4D tensors")
    if output_k.shape[:3] != output_v.shape[:3]:
        raise ValueError(
            "output_k and output_v must share batch, sequence, and head dims"
        )
    if output_k.shape[3] % _NVFP4_BLOCK_SIZE != 0:
        raise ValueError(
            f"output_k head_dim ({output_k.shape[3]}) must be divisible by {_NVFP4_BLOCK_SIZE}"
        )
    if output_v.shape[3] % _NVFP4_BLOCK_SIZE != 0:
        raise ValueError(
            f"output_v head_dim ({output_v.shape[3]}) must be divisible by {_NVFP4_BLOCK_SIZE}"
        )

    layout_code = 0 if kv_layout == "NHD" else 1
    get_fp4_kv_dequantization_module().nvfp4_paged_kv_dequant(
        paged_k_cache,
        paged_v_cache,
        k_scales,
        v_scales,
        block_tables,
        seq_lens,
        k_scale,
        v_scale,
        output_k,
        output_v,
        layout_code,
    )


@supported_compute_capability([100, 103, 107, 110, 120, 121])
def _nvfp4_kv_quant_check(input, global_scale):
    return True


@backend_requirement({}, common_check=_nvfp4_kv_quant_check)
@flashinfer_api(trace=nvfp4_kv_quantize_trace)
def nvfp4_kv_quantize(
    input: torch.Tensor,
    global_scale: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    r"""GPU quantization to the NVFP4 KV-cache format with linear block-scale layout.

    Requires SM100+ (Blackwell) for the ``cvt.rn.satfinite.e2m1x2.f32`` PTX
    instruction.

    Parameters
    ----------
    input : torch.Tensor
        Input tensor of shape ``[M, K]`` with dtype bf16 or fp16; ``K``
        must be divisible by 16.
    global_scale : torch.Tensor
        Global scale factor of shape ``[1]`` with dtype ``float32``, on the
        same CUDA device as ``input``.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor]
        ``(fp4_output, block_scales)`` where ``fp4_output`` is packed FP4
        data of shape ``[M, K/2]`` with dtype ``uint8`` and
        ``block_scales`` are per-block FP8 E4M3 scales of shape
        ``[M, K/16]`` with dtype ``uint8``.
    """
    M, K = input.shape
    if K % _NVFP4_BLOCK_SIZE != 0:
        raise ValueError(f"K dimension ({K}) must be divisible by {_NVFP4_BLOCK_SIZE}")
    fp4_output = torch.empty((M, K // 2), dtype=torch.uint8, device=input.device)
    block_scales = torch.empty(
        (M, K // _NVFP4_BLOCK_SIZE), dtype=torch.uint8, device=input.device
    )
    get_fp4_kv_quantization_module().nvfp4_kv_quant(
        input, global_scale, fp4_output, block_scales
    )
    return fp4_output, block_scales
