sgl-project--sglang/python/sglang/srt/layers/quantization/fp8_utils.py

from __future__ import annotations

import logging
from enum import Enum
from functools import lru_cache
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Union

import torch

from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.quantization.fp8_kernel import (
    sglang_per_token_group_quant_fp8,
    sglang_per_token_group_quant_fp8_row_padded,
)
from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils.common import torch_release

if TYPE_CHECKING:
    from sglang.srt.server_args import ServerArgs

from sglang.srt.layers.quantization.fp8_kernel import (
    fp8_dtype,
    fp8_max,
    fp8_min,
    is_fp8_fnuz,
    mxfp8_block_scaled_matmul_triton,
    pack_mxfp8_scales_triton,
    per_token_group_quant_fp8,
    scaled_fp8_quant,
    sglang_per_token_quant_fp8,
    static_quant_fp8,
    triton_scaled_mm,
    w8a8_block_fp8_matmul_deepgemm,
    w8a8_block_fp8_matmul_triton,
)
from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils import (
    ceil_align,
    ceil_div,
    get_bool_env_var,
    get_cuda_version,
    get_device_capability,
    get_hip_version,
    is_blackwell_supported,
    is_cuda,
    is_flashinfer_available,
    is_gfx95_supported,
    is_hip,
    is_musa,
    is_sm90_supported,
    is_sm100_supported,
    is_sm120_supported,
    offloader,
)
from sglang.srt.utils.custom_op import register_custom_op

logger = logging.getLogger(__name__)

_is_hip = is_hip()
_is_cuda = is_cuda()
_is_fp8_fnuz = is_fp8_fnuz()
_is_sm100_supported = is_sm100_supported()
_is_sm120_supported = is_sm120_supported()
_is_gfx95_supported = is_gfx95_supported()
_is_musa = is_musa()

_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_use_aiter_gfx95 = _use_aiter and _is_gfx95_supported
# ROCm 7.0 hipcc miscompiles gemm_a8w8_blockscale_bpreshuffle on gfx95 (#23319).
_use_aiter_bpreshuffle_gfx95 = _use_aiter_gfx95 and get_hip_version() >= (7, 2, 0)
# gfx95 + ROCm < 7.2: bpreshuffle CK is disabled (above), and the non-bpreshuffle
# fallback ck_gemm_a8w8_blockscale returns NaN above a per-shape M for some shapes
# (measured NaN onset: (2560,4096)@M>=4096, (4096,1024)@M>=8192), corrupting prefill.
# Map each affected (n, k) to the largest M for which CK is confirmed correct
# (conservative = last verified-safe M). Keep the faster CK path at/below that M and
# fall back to the numerically-correct Triton FP8 GEMM above it. Fixed in ROCm 7.2.
_AITER_GFX95_CK_W8A8_MAX_SAFE_M = {
    (2560, 4096): 2048,
    (4096, 1024): 4096,
}


class _MXFP4QuantizedData(MXFP4QuantizeUtil):
    def __init__(
        self,
        original_shape: torch.Size,
        original_dtype: torch.dtype,
        quantized_data: torch.Tensor,
    ):
        self.original_shape = original_shape
        self.original_dtype = original_dtype
        self.quantized_data = quantized_data


# Force CK bpreshuffle (not Triton) for the dense w8a8-block GEMMs (MLA q/kv/o
# projections), to match ATOM (CK preshuffle; Triton FP8 blockscale is slower).
# Default OFF; DeepseekV4 enables it via set_force_ck_w8a8(True). The env var
# SGLANG_FORCE_CK_W8A8=1 still works as an override.
_FORCE_CK_W8A8: bool = False


def set_force_ck_w8a8(enabled: bool = True) -> None:
    global _FORCE_CK_W8A8
    _FORCE_CK_W8A8 = enabled


def materialize_bpreshuffle_fp8_scale(scale: torch.Tensor) -> torch.Tensor:
    """Materialize the physical scale layout consumed by gfx95 bpreshuffle GEMM."""
    return scale.t().contiguous().t() if scale.dim() == 2 else scale


def materialize_bpreshuffle_fp8_scale_tuple(
    value: Tuple[torch.Tensor, ...],
) -> Tuple[torch.Tensor, ...]:
    """Materialize the scale slot in FP8 ``(q_input, x_scale, ...)`` tuples."""
    return (
        value[0],
        materialize_bpreshuffle_fp8_scale(value[1]),
        *value[2:],
    )


def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool:
    if _FORCE_CK_W8A8:
        return False
    return (n, k) in [
        (1024, 8192),
        (16384, 1536),
        (2112, 7168),
        (3072, 1536),
        (32768, 8192),
        (4096, 7168),
        (4608, 7168),
        (512, 7168),
        (7168, 2048),
        (7168, 2304),
        (7168, 16384),
        (7168, 256),
        (8192, 1024),
        (8192, 32768),
    ]


if _use_aiter:
    import aiter
    from aiter import gemm_a8w8_blockscale as ck_gemm_a8w8_blockscale
    from aiter import (
        gemm_a8w8_blockscale_bpreshuffle,
        gemm_a8w8_bpreshuffle,
        get_hip_quant,
    )
    from aiter.ops.triton.gemm_a8w8_blockscale import (
        gemm_a8w8_blockscale as triton_gemm_a8w8_blockscale,
    )

    aiter_per1x128_quant = get_hip_quant(aiter.QuantType.per_1x128)


if _is_cuda:
    from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm

    from sglang.srt.utils.patch_torch import register_fake_if_exists

    @register_fake_if_exists("sgl_kernel::fp8_scaled_mm")
    def _fp8_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
        # mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N].
        M = mat_a.shape[-2]
        N = mat_b.shape[-1]
        return mat_a.new_empty((M, N), dtype=out_dtype)

    @register_fake_if_exists("sgl_kernel::fp8_blockwise_scaled_mm")
    def _fp8_blockwise_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype):
        # mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N].
        M = mat_a.shape[-2]
        N = mat_b.shape[-1]
        return mat_a.new_empty((M, N), dtype=out_dtype)


use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL")

# Input scaling factors are no longer optional in _scaled_mm starting
# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale
TORCH_DEVICE_IDENTITY = None


def use_rowwise_torch_scaled_mm():
    if _is_hip:
        # The condition to determine if it is on a platform that supports
        # torch._scaled_mm rowwise feature.
        # The condition is determined once as the operations
        # are time consuming.
        return get_device_capability() >= (9, 4) and torch_release >= (2, 7)
    return False


USE_ROWWISE_TORCH_SCALED_MM = use_rowwise_torch_scaled_mm()


@lru_cache(maxsize=1)
def cutlass_fp8_supported():
    if not _is_cuda:
        return False
    major, minor = get_device_capability()
    cuda_version = get_cuda_version()
    if major >= 9:
        return cuda_version >= (12, 0)
    elif major == 8 and minor == 9:
        return cuda_version >= (12, 4)
    return False


def normalize_e4m3fn_to_e4m3fnuz(
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
    assert weight.dtype == torch.float8_e4m3fn
    # The bits pattern 10000000(-128) represents zero in e4m3fn
    # but NaN in e4m3fnuz. So here we set it to 0.
    # https://onnx.ai/onnx/technical/float8.html
    weight_as_int8 = weight.view(torch.int8)
    ROCM_FP8_NAN_AS_INT = -128
    weight_as_int8[weight_as_int8 == ROCM_FP8_NAN_AS_INT] = 0
    weight = weight_as_int8.view(torch.float8_e4m3fnuz)

    # For the same bits representation, e4m3fnuz value is half of
    # the e4m3fn value, so we should double the scaling factor to
    # get the same dequantized value.
    # https://onnx.ai/onnx/technical/float8.html
    weight_scale = weight_scale * 2.0
    if input_scale is not None:
        input_scale = input_scale * 2.0
    return weight, weight_scale, input_scale


class Fp8GemmRunnerBackend(Enum):
    """Enum for FP8 GEMM runner backend selection."""

    AUTO = "auto"
    FLASHINFER_TRTLLM = "flashinfer_trtllm"
    FLASHINFER_CUTLASS = "flashinfer_cutlass"
    FLASHINFER_DEEPGEMM = "flashinfer_deepgemm"
    CUTLASS = "cutlass"
    DEEP_GEMM = "deep_gemm"
    TRITON = "triton"
    AITER = "aiter"

    def is_auto(self) -> bool:
        return self == Fp8GemmRunnerBackend.AUTO

    def is_flashinfer_trtllm(self) -> bool:
        return self == Fp8GemmRunnerBackend.FLASHINFER_TRTLLM

    def is_flashinfer_cutlass(self) -> bool:
        return self == Fp8GemmRunnerBackend.FLASHINFER_CUTLASS

    def is_flashinfer_deepgemm(self) -> bool:
        return self == Fp8GemmRunnerBackend.FLASHINFER_DEEPGEMM

    def is_cutlass(self) -> bool:
        return self == Fp8GemmRunnerBackend.CUTLASS

    def is_deep_gemm(self) -> bool:
        return self == Fp8GemmRunnerBackend.DEEP_GEMM

    def is_triton(self) -> bool:
        return self == Fp8GemmRunnerBackend.TRITON

    def is_aiter(self) -> bool:
        return self == Fp8GemmRunnerBackend.AITER


FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None


def _check_cutlass_block_fp8_hardware_support() -> bool:
    """Return True if CUTLASS block FP8 is supported (Hopper or newer with CUDA 12.0+)."""
    return is_sm90_supported() or is_blackwell_supported()


if is_blackwell_supported() and is_flashinfer_available():
    from flashinfer import SfLayout
    from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8
    from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8
    from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize
    from flashinfer.gemm import gemm_fp8_nt_groupwise as _raw_gemm_fp8_nt_groupwise

    from sglang.srt.utils.custom_op import register_custom_op

    @register_custom_op(
        op_name="flashinfer_bmm_fp8",
        mutates_args=[],
        fake_impl=lambda q_input, weight, x_scale, weight_scale, out_dtype: (
            q_input.new_empty((q_input.shape[0], weight.shape[1]), dtype=out_dtype)
        ),
    )
    def flashinfer_bmm_fp8(
        q_input: torch.Tensor,  # [M, K] fp8 e4m3
        weight: torch.Tensor,  # [K, N] fp8 e4m3, column-major
        x_scale: torch.Tensor,  # per-tensor scalar
        weight_scale: torch.Tensor,  # per-tensor scalar
        out_dtype: torch.dtype,
    ) -> torch.Tensor:
        m, n = q_input.shape[0], weight.shape[1]
        return _raw_flashinfer_bmm_fp8(
            q_input.unsqueeze(0),
            weight.unsqueeze(0),
            x_scale.reshape(1),
            weight_scale.reshape(1),
            out_dtype,
            backend="cublas",
        ).view(m, n)

    @lru_cache(maxsize=1)
    def _get_flashinfer_groupwise_backend() -> str:
        if get_fp8_gemm_runner_backend().is_flashinfer_cutlass():
            return "cutlass"
        if get_fp8_gemm_runner_backend().is_flashinfer_trtllm():
            return "trtllm"

        major, minor = get_device_capability()
        # SM120/121: CUTLASS only.
        # SM100/103: TRTLLM only.
        if major >= 12:
            return "cutlass"
        return "trtllm"

    # Wrap gemm_fp8_nt_groupwise as a custom op so torch.compile does not trace
    # into flashinfer's JIT compilation code (pathlib/cubin_loader ops).
    @register_custom_op(
        op_name="flashinfer_gemm_fp8_nt_groupwise",
        mutates_args=[],
        fake_impl=lambda q_input, weight, x_scale, weight_scale, out_dtype: (
            q_input.new_empty((q_input.shape[0], weight.shape[0]), dtype=out_dtype)
        ),
    )
    def gemm_fp8_nt_groupwise(
        q_input: torch.Tensor,
        weight: torch.Tensor,
        x_scale: torch.Tensor,
        weight_scale: torch.Tensor,
        out_dtype: torch.dtype,
    ) -> torch.Tensor:
        backend = _get_flashinfer_groupwise_backend()
        if backend == "cutlass":
            # FlashInfer CUTLASS groupwise kernel requires contiguous scale tensors
            x_scale = x_scale.contiguous()
            weight_scale = weight_scale.contiguous()
            return _raw_gemm_fp8_nt_groupwise(
                q_input,
                weight,
                x_scale,
                weight_scale,
                out_dtype=out_dtype,
                backend="cutlass",
                scale_major_mode="MN",
            )
        return _raw_gemm_fp8_nt_groupwise(
            q_input,
            weight,
            x_scale,
            weight_scale,
            out_dtype=out_dtype,
            backend=backend,
        )

    # Wrap MXFP8 ops as custom ops so torch.compile does not trace into
    # flashinfer's JIT compilation path (filesystem checks/cubin loader).
    def _fake_flashinfer_mxfp8_quantize(
        input: torch.Tensor,
        _is_sf_swizzled_layout: bool = True,
        alignment: int = 32,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Fake mode only needs dtypes and output rank to propagate compile graph.
        # The scale tensor shape is not consumed before the following fake mm op.
        k_aligned = ((input.shape[1] + alignment - 1) // alignment) * alignment
        q_input = input.new_empty(
            (input.shape[0], k_aligned), dtype=torch.float8_e4m3fn
        )
        scale = input.new_empty((1,), dtype=torch.uint8)
        return q_input, scale

    @register_custom_op(
        op_name="flashinfer_mxfp8_quantize",
        mutates_args=[],
        fake_impl=_fake_flashinfer_mxfp8_quantize,
    )
    def flashinfer_mxfp8_quantize(
        input: torch.Tensor,
        is_sf_swizzled_layout: bool = True,
        alignment: int = 32,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        return _raw_flashinfer_mxfp8_quantize(
            input,
            is_sf_swizzled_layout=is_sf_swizzled_layout,
            alignment=alignment,
            sf_swizzle_layout=SfLayout.layout_128x4,
        )

    @register_custom_op(
        op_name="flashinfer_mm_mxfp8",
        mutates_args=[],
        fake_impl=lambda q_input, weight_t, x_scale_u8, weight_scale_t, out_dtype, use_8x4_sf_layout=False, backend="auto": (
            q_input.new_empty((q_input.shape[0], weight_t.shape[1]), dtype=out_dtype)
        ),
    )
    def flashinfer_mm_mxfp8(
        q_input: torch.Tensor,
        weight_t: torch.Tensor,
        x_scale_u8: torch.Tensor,
        weight_scale_t: torch.Tensor,
        out_dtype: torch.dtype,
        use_8x4_sf_layout: bool = False,
        backend: str = "auto",
    ) -> torch.Tensor:
        return _raw_flashinfer_mm_mxfp8(
            q_input,
            weight_t,
            x_scale_u8,
            weight_scale_t,
            out_dtype=out_dtype,
            use_8x4_sf_layout=use_8x4_sf_layout,
            backend=backend,
        )


if is_sm90_supported() and is_flashinfer_available():
    # FlashInfer SM90 DeepGEMM with automatic swapAB optimization for small M
    from flashinfer.gemm import fp8_blockscale_gemm_sm90


def dispatch_w8a8_block_fp8_linear() -> Callable:
    """
    Dispatch to the appropriate FP8 block linear implementation.

    This function selects the backend based on:
    1. The --fp8-gemm-backend server argument (preferred)
    2. Auto-detection based on hardware capabilities
    """
    backend = get_fp8_gemm_runner_backend()

    # Handle explicit backend selection via --fp8-gemm-backend
    if not backend.is_auto():
        return _dispatch_explicit_backend(backend)

    # Auto mode: Select based purely on hardware/backend availability
    return _dispatch_auto_backend()


def dispatch_w8a8_mxfp8_linear() -> Callable:
    backend = get_fp8_gemm_runner_backend()
    if backend.is_deep_gemm():
        return _deepgemm_w8a8_mxfp8_linear_with_fallback
    elif backend.is_flashinfer_cutlass() or backend.is_flashinfer_trtllm():
        return flashinfer_mxfp8_blockscaled_linear
    elif backend.is_triton():
        return triton_mxfp8_blockscaled_linear
    elif _is_hip and _is_gfx95_supported:
        from sglang.srt.layers.quantization.mxfp8_amd_gfx95 import (
            dot_scaled_mxfp8_blockscaled_linear,
        )

        return dot_scaled_mxfp8_blockscaled_linear
    return triton_mxfp8_blockscaled_linear


def _deepgemm_w8a8_mxfp8_linear_with_fallback(
    input: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
    weight_scale_fallback: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    from sglang.srt.layers.quantization.fp8_kernel import (
        sglang_per_token_group_quant_fp8,
        w8a8_mxfp8_matmul_deepgemm,
    )

    assert input_scale is None
    output_dtype = input.dtype

    shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0
    dtype_supported = output_dtype == torch.bfloat16

    if not (shape_supported and dtype_supported):
        return triton_mxfp8_blockscaled_linear(
            input, weight, weight_scale_fallback, input_scale, bias
        )

    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]

    q_input, x_scale = sglang_per_token_group_quant_fp8(
        input_2d,
        32,
        column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
        scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
        scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
    )

    # weight_scale format is set per-backend by _process_mxfp8_linear_weight_scale
    # (int32 packed TMA-aligned on Blackwell, float32 on Hopper); NOT uint8 — Triton form is routed to the fallback above.

    output = w8a8_mxfp8_matmul_deepgemm(
        q_input, weight, x_scale, weight_scale, output_dtype=output_dtype
    )
    if bias is not None:
        output += bias
    return output.to(dtype=output_dtype).view(*output_shape)


def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable:
    """Dispatch based on explicitly selected backend."""
    if backend.is_flashinfer_trtllm():
        if not (is_sm100_supported() and is_flashinfer_available()):
            raise RuntimeError(
                "FlashInfer FP8 GEMM requested via --fp8-gemm-backend=flashinfer_trtllm, "
                "but FlashInfer is not available or not supported on this hardware. "
                "FlashInfer TRTLLM FP8 GEMM requires SM100/SM103 GPUs and FlashInfer."
            )
        return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback

    elif backend.is_flashinfer_cutlass():
        if not (is_blackwell_supported() and is_flashinfer_available()):
            raise RuntimeError(
                "FlashInfer FP8 GEMM requested via --fp8-gemm-backend=flashinfer_cutlass, "
                "but FlashInfer is not available or not supported on this hardware. "
                "FlashInfer CUTLASS FP8 GEMM requires Blackwell GPUs and FlashInfer."
            )
        return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback

    elif backend.is_flashinfer_deepgemm():
        if not (is_sm90_supported() and is_flashinfer_available()):
            raise RuntimeError(
                "FlashInfer DeepGEMM with swapAB requested via --fp8-gemm-backend=flashinfer_deepgemm, "
                "but it's not available. This backend requires Hopper (SM90) GPUs and FlashInfer "
                "to be installed."
            )
        return flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback

    elif backend.is_cutlass():
        if not _check_cutlass_block_fp8_hardware_support():
            raise RuntimeError(
                "CUTLASS block FP8 requested via --fp8-gemm-backend=cutlass, "
                "but hardware does not support it. CUTLASS block FP8 requires "
                "Hopper (SM90+) GPUs with CUDA 12.0+."
            )
        return cutlass_w8a8_block_fp8_linear_with_fallback

    elif backend.is_aiter():
        if not _use_aiter:
            raise RuntimeError(
                "AITER backend requested via --fp8-gemm-backend=aiter, "
                "but AITER is not available. AITER requires AMD GPUs with "
                "SGLANG_USE_AITER=1 environment variable set."
            )
        return aiter_w8a8_block_fp8_linear

    elif backend.is_deep_gemm():
        if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
            raise RuntimeError(
                "DeepGEMM backend requested via --fp8-gemm-backend=deep_gemm, "
                "but DeepGEMM is not available. This usually means the deep_gemm package "
                "is not installed or has been disabled via SGLANG_ENABLE_JIT_DEEPGEMM=0."
            )
        return deepgemm_w8a8_block_fp8_linear_with_fallback

    elif backend.is_triton():
        return triton_w8a8_block_fp8_linear

    else:
        raise ValueError(f"Unknown FP8 GEMM backend: {backend}")


def _dispatch_auto_backend() -> Callable:
    """Auto-select the best backend based on hardware capabilities."""
    # Priority order for auto selection:
    # 1. DeepGEMM (if enabled and available)
    # 2. FlashInfer TRTLLM (if Blackwell GPU and FlashInfer available)
    # 3. CUTLASS (if Hopper+ GPU and CUDA 12.0+)
    # 4. AITER (if AMD GPU with AITER enabled)
    # 5. Triton (fallback)

    if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
        return deepgemm_w8a8_block_fp8_linear_with_fallback
    elif is_blackwell_supported() and is_flashinfer_available():
        return flashinfer_gemm_w8a8_block_fp8_linear_with_fallback
    elif _check_cutlass_block_fp8_hardware_support():
        return cutlass_w8a8_block_fp8_linear_with_fallback
    elif _use_aiter:
        return aiter_w8a8_block_fp8_linear
    else:
        return triton_w8a8_block_fp8_linear


def initialize_fp8_gemm_config(server_args: ServerArgs) -> None:
    """Initialize FP8 GEMM configuration."""
    global FP8_GEMM_RUNNER_BACKEND

    backend = server_args.fp8_gemm_runner_backend
    if backend == "auto" and is_sm120_supported():
        # TODO(brayden): Verify if CUTLASS can be set by default once SwapAB is supported
        backend = "triton"

    backend = Fp8GemmRunnerBackend(backend)

    if (
        backend.is_auto()
        and server_args.quantization == "mxfp8"
        and _is_sm100_supported
        and is_flashinfer_available()
    ):
        backend = Fp8GemmRunnerBackend.FLASHINFER_CUTLASS

    FP8_GEMM_RUNNER_BACKEND = backend


def get_fp8_gemm_runner_backend() -> Fp8GemmRunnerBackend:
    """Get the current FP8 GEMM runner backend."""
    global FP8_GEMM_RUNNER_BACKEND
    if FP8_GEMM_RUNNER_BACKEND is None:
        FP8_GEMM_RUNNER_BACKEND = Fp8GemmRunnerBackend.AUTO
    return FP8_GEMM_RUNNER_BACKEND


def flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
    input: torch.Tensor,
    weight: torch.Tensor,
    block_size: List[int],
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    assert input_scale is None

    input_2d = input.view(-1, input.shape[-1])
    backend = _get_flashinfer_groupwise_backend()
    # Fall back to triton for non-supported formats.
    # TODO: Check if flashinfer supports other output dtypes besides bf16.
    if backend == "trtllm" and (
        input_2d.shape[1] < 256 or input_2d.dtype != torch.bfloat16
    ):
        return triton_w8a8_block_fp8_linear(
            input, weight, block_size, weight_scale, input_scale, bias
        )

    output_shape = [*input.shape[:-1], weight.shape[0]]

    # TRTLLM uses the existing SGLang column-major scale layout.
    # CUTLASS with scale_major_mode="MN" expects (k//block_k, m), so we normalize below.
    q_input, x_scale = sglang_per_token_group_quant_fp8(
        input_2d, block_size[1], column_major_scales=(backend == "trtllm")
    )
    if backend == "cutlass":
        block_n, block_k = block_size
        m, k = input_2d.shape
        n = weight.shape[0]
        expected_x_scale_shape = (k // block_k, m)
        expected_weight_scale_shape = (k // block_k, n // block_n)
        if x_scale.shape == (m, k // block_k):
            x_scale = x_scale.transpose(-1, -2).contiguous()
        if weight_scale.shape == (n // block_n, k // block_k):
            weight_scale = weight_scale.transpose(-1, -2).contiguous()
        assert x_scale.shape == expected_x_scale_shape, (
            "FlashInfer CUTLASS groupwise FP8 expects A scale layout "
            f"(k//block_k, m) for scale_major_mode='MN', got {tuple(x_scale.shape)}; "
            f"expected {expected_x_scale_shape}. "
            f"strides={x_scale.stride()} is_contiguous={x_scale.is_contiguous()} "
            f"m={m} n={n} k={k} block_size={block_size}"
        )
        assert weight_scale.shape == expected_weight_scale_shape, (
            "FlashInfer CUTLASS groupwise FP8 expects B scale layout "
            f"(k//block_k, n//block_n) for scale_major_mode='MN', got {tuple(weight_scale.shape)}; "
            f"expected {expected_weight_scale_shape}. "
            f"strides={weight_scale.stride()} is_contiguous={weight_scale.is_contiguous()} "
            f"m={m} n={n} k={k} block_size={block_size}"
        )
        assert x_scale.dtype == torch.float32, (
            "FlashInfer CUTLASS groupwise FP8 expects x_scale dtype float32, "
            f"got {x_scale.dtype}."
        )
        assert weight_scale.dtype == torch.float32, (
            "FlashInfer CUTLASS groupwise FP8 expects weight_scale dtype float32, "
            f"got {weight_scale.dtype}."
        )
    # TRTLLM path continues using the original quantized scale layout.
    output = gemm_fp8_nt_groupwise(
        q_input,
        weight,
        x_scale,
        weight_scale,
        out_dtype=input_2d.dtype,
    )

    if bias is not None:
        output += bias

    return output.to(dtype=input_2d.dtype).view(*output_shape)


def flashinfer_deepgemm_w8a8_block_fp8_linear_with_fallback(
    input: torch.Tensor,
    weight: torch.Tensor,
    block_size: List[int],
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    FlashInfer DeepGEMM backend for SM90 (Hopper) with swapAB optimization.

    Uses flashinfer.gemm.fp8_blockscale_gemm_sm90 which automatically selects
    the swapAB kernel for small M dimensions (M < 32) for better performance
    during decoding/low batch size scenarios.

    For SM90 (Hopper), this uses the DeepGEMM JIT with automatic swapAB selection.
    """
    assert input_scale is None

    output_dtype = input.dtype
    dtype_supported = output_dtype == torch.bfloat16

    # fp8_blockscale_gemm_sm90 requires: N % 64 == 0, K % 128 == 0
    shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0

    if not (shape_supported and dtype_supported):
        if weight_scale.dtype == torch.int32:
            weight_scale = _unpack_ue8m0_scale_for_triton(
                weight_scale, weight.shape, block_size
            )
        return triton_w8a8_block_fp8_linear(
            input, weight, block_size, weight_scale, input_scale, bias
        )

    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]

    # - input: (M, K) BF16 or FP8
    # - weight: (N, K) FP8 with weight_scale
    # - weight_scale: (N, K//128) for per-token or (N//128, K//128) for per-block

    output = fp8_blockscale_gemm_sm90(
        input_2d,
        weight,
        input_scale=None,  # BF16 input, internal quantization
        weight_scale=weight_scale,
        out_dtype=output_dtype,
    )

    if bias is not None:
        output += bias
    return output.view(*output_shape)


def cutlass_w8a8_block_fp8_linear_with_fallback(
    input: torch.Tensor,
    weight: torch.Tensor,
    block_size: List[int],
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    assert input_scale is None

    # TODO: add more robust shape check here
    shape_supported = weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0

    if not shape_supported:
        # fallback to triton
        return triton_w8a8_block_fp8_linear(
            input, weight, block_size, weight_scale, input_scale, bias
        )

    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]

    # Quantize into row-padded buffers so the sgl-kernel wrapper's per-call
    # pad_tensor() on mat_a / scales_a short-circuits (saves 2x fill + 2x cat
    # kernels per GEMM). weight_scale.T is left as a K-major view because the
    # kernel requires scales_b.stride(0) == 1 and materializes it internally.
    q_input, x_scale = sglang_per_token_group_quant_fp8_row_padded(
        input_2d, block_size[1]
    )
    output = fp8_blockwise_scaled_mm(
        q_input, weight.T, x_scale, weight_scale.T, out_dtype=input_2d.dtype
    )
    if output.shape[0] != input_2d.shape[0]:
        # GEMM ran on the row-padded buffer; drop the padding rows.
        output = output[: input_2d.shape[0]]
    if bias is not None:
        output += bias
    return output.to(dtype=input_2d.dtype).view(*output_shape)


def deepgemm_w8a8_block_fp8_linear_with_fallback(
    input: torch.Tensor,
    weight: torch.Tensor,
    block_size: List[int],
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    assert input_scale is None

    output_dtype = input.dtype
    dtype_supported = output_dtype == torch.bfloat16

    # TODO: https://github.com/sgl-project/sglang/pull/6890#issuecomment-2943395737
    shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0

    if not (shape_supported and dtype_supported):
        # fall back to triton
        # If weight_scale is in UE8M0 packed format (int32), convert back to float32
        # UE8M0 format has shape (N, K//block_k//4) with dtype int32
        # Triton expects shape (N//block_n, K//block_k) with dtype float32
        if weight_scale.dtype == torch.int32:
            weight_scale = _unpack_ue8m0_scale_for_triton(
                weight_scale, weight.shape, block_size
            )
        return triton_w8a8_block_fp8_linear(
            input, weight, block_size, weight_scale, input_scale, bias
        )

    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]

    if not _is_musa:
        q_input, x_scale = sglang_per_token_group_quant_fp8(
            input_2d,
            block_size[1],
            column_major_scales=True,
            scale_tma_aligned=True,
            scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
        )
    else:
        q_input, x_scale = sglang_per_token_group_quant_fp8(
            input_2d,
            block_size[1],
        )

    output = w8a8_block_fp8_matmul_deepgemm(
        q_input, weight, x_scale, weight_scale, block_size, output_dtype=output_dtype
    )
    if bias is not None:
        output += bias
    return output.to(dtype=output_dtype).view(*output_shape)


def _unpack_ue8m0_scale_for_triton(
    sf_packed: torch.Tensor,
    weight_shape: Tuple[int, int],
    block_size: List[int],
) -> torch.Tensor:
    """
    Unpack UE8M0 packed scale tensor back to float32 format for triton kernel.

    The UE8M0 format packs scales as:
    - Shape: (N, K//block_k//4) with dtype int32
    - Each int32 contains 4 uint8 scale values

    Triton expects:
    - Shape: (N//block_n, K//block_k) with dtype float32

    Args:
        sf_packed: Packed scale tensor with shape (N, packed_k_groups) and dtype int32
        weight_shape: (N, K) shape of the weight tensor
        block_size: [block_n, block_k] quantization block size

    Returns:
        Unpacked scale tensor with shape (n_groups, k_groups) and dtype float32
    """
    assert sf_packed.dtype == torch.int32
    assert len(sf_packed.shape) == 2

    N, K = weight_shape
    block_n, block_k = block_size
    n_groups = ceil_div(N, block_n)
    k_groups = ceil_div(K, block_k)

    mn_repeat, k_div_4 = sf_packed.shape
    k_packed = k_div_4 * 4

    # Unpack int32 -> 4x uint8 -> float32
    # Each uint8 represents an exponent in UE8M0 format
    sf_u8 = sf_packed.contiguous().view(torch.uint8).view(mn_repeat, k_packed)
    sf_fp32 = (sf_u8.to(torch.int32) << 23).view(torch.float32)

    # Handle row dimension - may have 128x replication or direct mapping
    if mn_repeat == N:
        # Rows are replicated 128 times, take every 128th row
        # sf_fp32 shape: (N, k_packed) -> (n_groups, k_packed)
        # Select representative rows at indices 0, 128, 256, ...
        indices = torch.arange(0, N, block_n, device=sf_packed.device)
        sf_fp32 = sf_fp32.index_select(0, indices)
    elif mn_repeat == n_groups:
        # Already in the correct n_groups format
        pass
    else:
        raise ValueError(
            f"Unexpected scale shape: sf_packed.shape={sf_packed.shape}, "
            f"weight_shape={weight_shape}, block_size={block_size}"
        )

    # Crop k dimension to expected size (remove padding if any)
    sf_fp32 = sf_fp32[:, :k_groups].contiguous()

    return sf_fp32


def aiter_w8a8_block_fp8_linear(
    input: torch.Tensor,
    weight: torch.Tensor,
    block_size: List[int],
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    # assert input_scale is None
    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]

    n, k = weight.shape

    if _use_aiter_bpreshuffle_gfx95:
        use_triton = use_aiter_triton_gemm_w8a8_tuned_gfx950(n, k)
    elif _use_aiter_gfx95:
        # gfx95 on ROCm < 7.2: keep the (faster) CK path at/below the per-shape
        # CK-safe M bound; above it, ck_gemm_a8w8_blockscale returns NaN, so use
        # Triton. Unlisted shapes keep their original decision. Fixed in ROCm 7.2.
        _ck_safe_m = _AITER_GFX95_CK_W8A8_MAX_SAFE_M.get((n, k))
        use_triton = use_aiter_triton_gemm_w8a8_tuned_gfx950(n, k) or (
            _ck_safe_m is not None and input_2d.shape[0] > _ck_safe_m
        )
    else:
        use_triton = True

    # if input_scale not None, input is quanted
    if input_scale is not None:
        q_input = input_2d
        x_scale = input_scale
        if _use_aiter_bpreshuffle_gfx95 and not use_triton:
            x_scale = materialize_bpreshuffle_fp8_scale(x_scale)
        # On ROCm >= 7.2, scale is in bpreshuffle's transposed layout.
        # Triton needs a row-major view, so adjust strides only. No copy.
        elif use_triton and _use_aiter_bpreshuffle_gfx95:
            x_scale = torch.as_strided(x_scale, x_scale.shape, (1, x_scale.shape[0]))
    else:
        materialize_bpreshuffle_scale = _use_aiter_bpreshuffle_gfx95 and not use_triton
        q_input, x_scale = aiter_per1x128_quant(
            input_2d,
            quant_dtype=aiter.dtypes.fp8,
            transpose_scale=False,
        )
        if materialize_bpreshuffle_scale:
            x_scale = materialize_bpreshuffle_fp8_scale(x_scale)

    if use_triton:
        gemm_a8w8_blockscale_op = triton_gemm_a8w8_blockscale
    elif _use_aiter_bpreshuffle_gfx95:
        gemm_a8w8_blockscale_op = gemm_a8w8_blockscale_bpreshuffle
    else:
        gemm_a8w8_blockscale_op = ck_gemm_a8w8_blockscale

    output = gemm_a8w8_blockscale_op(
        q_input,
        weight,
        x_scale,
        weight_scale,
        dtype=torch.bfloat16 if input_scale is not None else input.dtype,
    )

    if bias is not None:
        output += bias

    return output.to(
        dtype=torch.bfloat16 if input_scale is not None else input_2d.dtype
    ).view(*output_shape)


def triton_w8a8_block_fp8_linear(
    input: torch.Tensor,
    weight: torch.Tensor,
    block_size: List[int],
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    assert input_scale is None
    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]

    q_input, x_scale = per_token_group_quant_fp8(
        input_2d, block_size[1], column_major_scales=False
    )
    output = w8a8_block_fp8_matmul_triton(
        q_input, weight, x_scale, weight_scale, block_size, output_dtype=input_2d.dtype
    )
    if bias is not None:
        output += bias
    return output.to(dtype=input_2d.dtype).view(*output_shape)


@lru_cache(maxsize=1)
def _get_triton_mxfp8_downcast():
    try:
        from triton_kernels.numerics_details.mxfp import downcast_to_mxfp
    except Exception as err:
        raise RuntimeError(
            "MXFP8 quantization requires triton_kernels with MXFP8 support."
        ) from err
    return downcast_to_mxfp


def mxfp8_group_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """Quantize a 2D contiguous tensor to MXFP8 with UE8M0 scales per group (32)."""
    assert x.dim() == 2, f"Expected 2D input, got {x.dim()}D"
    assert x.is_contiguous(), "MXFP8 quantization requires a contiguous 2D tensor."
    _, k = x.shape
    assert k % 32 == 0, f"{k=} must be divisible by 32"
    downcast_to_mxfp = _get_triton_mxfp8_downcast()
    q_input, scale_u8 = downcast_to_mxfp(x, torch.float8_e4m3fn, axis=1)
    return q_input.contiguous(), scale_u8.contiguous()


def _pack_mxfp8_scales(scale_u8: torch.Tensor) -> torch.Tensor:
    if (
        _is_hip
        and _is_gfx95_supported
        and scale_u8.is_cuda
        and scale_u8.shape[0] % 128 != 0
    ):
        return pack_mxfp8_scales_triton(scale_u8)

    # Pack (M, K//32) UE8M0 scales into the layout expected by tl.dot_scaled.
    assert scale_u8.dim() == 2, f"Expected 2D scale tensor, got {scale_u8.dim()}D"
    scale_u8 = scale_u8.contiguous()
    m, k_groups = scale_u8.shape
    assert (
        k_groups % 4 == 0
    ), f"{k_groups=} must be divisible by 4 (K must be multiple of 128)"

    scale_m = ceil_div(m, 128)
    if m % 128 != 0:
        pad_rows = scale_m * 128 - m
        pad = torch.full(
            (pad_rows, k_groups),
            127,
            dtype=scale_u8.dtype,
            device=scale_u8.device,
        )
        scale_u8 = torch.cat([scale_u8, pad], dim=0)

    scale_k = k_groups // 4
    scale_u8 = scale_u8.view(scale_m, 128, scale_k, 4)
    scale_u8 = scale_u8.view(scale_m, 4, 32, scale_k, 4)
    packed = scale_u8.permute(0, 3, 2, 1, 4).contiguous()
    return packed.view(1, scale_m, scale_k, 2, 256)


@register_custom_op(
    op_name="triton_mxfp8_block_scaled_matmul",
    mutates_args=[],
    fake_impl=lambda a, a_scale, b, b_scale, output_dtype, block_m=128, block_n=256, block_k=128, num_stages=None: (  # noqa: E501
        a.new_empty((a.shape[0], b.shape[0]), dtype=output_dtype)
    ),
)
def triton_mxfp8_block_scaled_matmul(
    a: torch.Tensor,
    a_scale: torch.Tensor,
    b: torch.Tensor,
    b_scale: torch.Tensor,
    output_dtype: torch.dtype,
    *,
    block_m: int = 128,
    block_n: int = 256,
    block_k: int = 128,
    num_stages: Optional[int] = None,
) -> torch.Tensor:
    """Opaque custom op wrapper to prevent Dynamo tracing Triton grid math."""
    return mxfp8_block_scaled_matmul_triton(
        a,
        a_scale,
        b,
        b_scale,
        output_dtype=output_dtype,
        block_m=block_m,
        block_n=block_n,
        block_k=block_k,
        num_stages=num_stages,
    )


def _raw_triton_mxfp8_blockscaled_linear(
    input: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
    output_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
    if not (
        (_is_cuda and (_is_sm100_supported or _is_sm120_supported))
        or (_is_hip and _is_gfx95_supported)
    ):
        raise RuntimeError(
            "MXFP8 dense linear requires Blackwell GPUs (SM100/SM120) or ROCm gfx95."
        )

    input_2d = input.view(-1, input.shape[-1]).contiguous()
    output_shape = [*input.shape[:-1], weight.shape[0]]

    block_m = 128
    block_n = 256 if weight.shape[0] % 256 == 0 else 128
    block_k = 128

    m, k = input_2d.shape
    n, k_w = weight.shape
    assert k == k_w, f"{k=} does not match {k_w=}"
    assert k % 128 == 0, f"{k=} must be divisible by 128 for MXFP8"
    assert n % block_n == 0, f"{n=} must be divisible by {block_n}"
    assert weight.dtype == torch.float8_e4m3fn, "MXFP8 weight must be FP8 E4M3."
    assert weight_scale.dtype == torch.uint8, "MXFP8 weight_scale must be UE8M0 uint8."
    assert weight_scale.dim() in (
        2,
        5,
    ), (
        "MXFP8 weight_scale must be canonical 2D or packed 5D, "
        f"got {weight_scale.dim()}D."
    )

    if input_scale is None:
        q_input, x_scale_u8 = mxfp8_group_quantize(input_2d)
    else:
        q_input = input_2d
        x_scale_u8 = input_scale
        assert x_scale_u8.dtype == torch.uint8, "MXFP8 input_scale must be UE8M0 uint8."
        assert x_scale_u8.shape == (m, k // 32)

    if output_dtype is None:
        if input_2d.dtype in (torch.float16, torch.bfloat16, torch.float32):
            output_dtype = input_2d.dtype
        else:
            output_dtype = torch.bfloat16

    if m % block_m != 0:
        pad_rows = ceil_div(m, block_m) * block_m - m
        q_input = torch.cat(
            [
                q_input,
                torch.zeros((pad_rows, k), device=q_input.device, dtype=q_input.dtype),
            ],
            dim=0,
        )
        pad_scale = torch.full(
            (pad_rows, k // 32),
            127,
            device=x_scale_u8.device,
            dtype=x_scale_u8.dtype,
        )
        x_scale_u8 = torch.cat([x_scale_u8, pad_scale], dim=0)

    a_scale_packed = _pack_mxfp8_scales(x_scale_u8)
    b_scale_packed = (
        weight_scale.contiguous()
        if weight_scale.dim() == 5
        else _pack_mxfp8_scales(weight_scale)
    )

    num_stages = 1 if _is_sm120_supported else (4 if _is_sm100_supported else 1)
    output = triton_mxfp8_block_scaled_matmul(
        q_input,
        a_scale_packed,
        weight.contiguous(),
        b_scale_packed,
        output_dtype=output_dtype,
        block_m=block_m,
        block_n=block_n,
        block_k=block_k,
        num_stages=num_stages,
    )
    output = output[:m, :]
    if bias is not None:
        output += bias
    return output.to(dtype=output_dtype).view(*output_shape)


@register_custom_op(
    op_name="triton_mxfp8_blockscaled_linear",
    mutates_args=[],
    fake_impl=lambda input, weight, weight_scale, input_scale=None, bias=None, output_dtype=None: (
        input.new_empty(
            (*input.shape[:-1], weight.shape[0]),
            dtype=(output_dtype if output_dtype is not None else input.dtype),
        )
    ),
)
def triton_mxfp8_blockscaled_linear(
    input: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
    output_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
    """Opaque custom-op wrapper to prevent Dynamo guards on MXFP8 padding branches."""
    return _raw_triton_mxfp8_blockscaled_linear(
        input=input,
        weight=weight,
        weight_scale=weight_scale,
        input_scale=input_scale,
        bias=bias,
        output_dtype=output_dtype,
    )


def flashinfer_mxfp8_blockscaled_linear(
    input: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
    output_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
    """MXFP8 dense linear via FlashInfer mm_mxfp8."""
    input_2d = input.view(-1, input.shape[-1]).contiguous()
    output_shape = [*input.shape[:-1], weight.shape[0]]

    m, k = input_2d.shape
    n, k_w = weight.shape
    if k != k_w:
        raise ValueError(f"Input K={k} does not match weight K={k_w}.")
    if k % 32 != 0:
        raise ValueError(f"K={k} must be divisible by 32 for MXFP8.")
    if weight.dtype != torch.float8_e4m3fn:
        raise TypeError("MXFP8 weight must be FP8 E4M3.")

    if input_scale is None:
        q_input, x_scale_u8 = flashinfer_mxfp8_quantize(
            input_2d, is_sf_swizzled_layout=True, alignment=32
        )
    else:
        q_input = input_2d
        x_scale_u8 = input_scale.contiguous()

    if output_dtype is None:
        if input_2d.dtype in (torch.float16, torch.bfloat16, torch.float32):
            output_dtype = input_2d.dtype
        else:
            output_dtype = torch.bfloat16

    # Ensure transposed tensors are contiguous for FlashInfer's internal runner.
    weight_t = weight.contiguous().t()

    if get_fp8_gemm_runner_backend().is_flashinfer_trtllm():
        weight_scale_t = weight_scale.contiguous().view(-1)
        output = flashinfer_mm_mxfp8(
            q_input,
            weight_t,
            x_scale_u8,
            weight_scale_t,
            out_dtype=output_dtype,
            use_8x4_sf_layout=False,
            backend="trtllm",
        )
    elif get_fp8_gemm_runner_backend().is_flashinfer_cutlass():
        weight_scale_t = (
            weight_scale.contiguous().t()
            if weight_scale.ndim == 2
            else weight_scale.contiguous()
        )
        output = flashinfer_mm_mxfp8(
            q_input,
            weight_t,
            x_scale_u8,
            weight_scale_t,
            out_dtype=output_dtype,
            use_8x4_sf_layout=False,
            backend="cutlass",
        )

    if bias is not None:
        output += bias
    return output.to(dtype=output_dtype).view(*output_shape)


def dequant_mxfp4(
    w_block: torch.Tensor,
    w_scale: torch.Tensor,
    out_dtype,
) -> torch.Tensor:
    """
    :param w_block: (batch, n, k, 16), uint8, pack two mxfp4 into one byte
    :param w_scale: (batch, n, k), uint8
    :return: (batch, n, k * 32), float32
    """

    assert w_block.dtype == torch.uint8
    assert w_scale.dtype == torch.uint8

    batch, n, k, pack_dim = w_block.shape
    batch_, n_, k_ = w_scale.shape
    assert pack_dim == 16
    assert batch == batch_
    assert n == n_
    assert k == k_

    out_raw = MXFP4QuantizeUtil.dequantize(
        quantized_data=w_block, scale=w_scale, dtype=out_dtype, block_sizes=[32]
    )
    return out_raw.reshape(batch, n, k * 32)


def input_to_float8(
    x: torch.Tensor, dtype: torch.dtype = fp8_dtype
) -> Tuple[torch.Tensor, torch.Tensor]:
    """This function quantizes input values to float8 values with tensor-wise quantization."""
    min_val, max_val = x.aminmax()
    amax = torch.maximum(min_val.abs(), max_val.abs()).float().clamp(min=1e-12)

    if _is_fp8_fnuz:
        dtype = fp8_dtype
        fp_max = fp8_max
    else:
        finfo = torch.finfo(dtype)
        fp_max = finfo.max

    scale = fp_max / amax
    x_scl_sat = (x.float() * scale).clamp(min=-fp_max, max=fp_max)
    return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()


def block_quant_to_tensor_quant(
    x_q_block: torch.Tensor,
    x_s: torch.Tensor,
    block_size: List[int],
) -> Tuple[torch.Tensor, torch.Tensor]:
    """This function converts block-wise quantization to tensor-wise quantization.
    The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale
    and the block size.
    The outputs are tensor-wise quantization tensor and tensor-wise quantization scale.
    Note only float8 is supported for now.
    """
    block_n, block_k = block_size[0], block_size[1]
    n, k = x_q_block.shape
    n_tiles = (n + block_n - 1) // block_n
    k_tiles = (k + block_k - 1) // block_k
    assert n_tiles == x_s.shape[0]
    assert k_tiles == x_s.shape[1]

    x_dq_block = x_q_block.to(torch.float32)

    x_dq_block_tiles = [
        [
            x_dq_block[
                j * block_n : min((j + 1) * block_n, n),
                i * block_k : min((i + 1) * block_k, k),
            ]
            for i in range(k_tiles)
        ]
        for j in range(n_tiles)
    ]

    for i in range(k_tiles):
        for j in range(n_tiles):
            x_dq_block_tiles[j][i][:, :] = x_dq_block_tiles[j][i] * x_s[j][i]

    x_q_tensor, scale = (
        scaled_fp8_quant(x_dq_block)
        if _is_cuda
        else input_to_float8(x_dq_block, dtype=x_q_block.dtype)
    )
    return x_q_tensor, scale


def block_quant_dequant(
    x_q_block: torch.Tensor,
    x_s: torch.Tensor,
    block_size: List[int],
    dtype: torch.dtype,
) -> torch.Tensor:
    """This function converts block-wise quantization to unquantized.
    The inputs are block-wise quantization tensor `x_q_block`, block-wise quantization scale
    and the block size.
    The output is an unquantized tensor with dtype.
    """
    block_n, block_k = block_size[0], block_size[1]
    *_, n, k = x_q_block.shape

    # ... n_scale k_scale -> ... (n_scale block_n) (k_scale block_k)
    x_scale_repeat = x_s.repeat_interleave(block_n, dim=-2).repeat_interleave(
        block_k, dim=-1
    )
    x_scale_repeat = x_scale_repeat[..., :n, :k]

    return (x_q_block.to(torch.float32) * x_scale_repeat).to(dtype)


def quantize_block_fp8_weight_to_mxfp4(
    fp8_weight: torch.Tensor,
    fp8_scale: torch.Tensor,
    weight_block_size: List[int],
    mxfp4_block_size: int = 32,
) -> Tuple[torch.Tensor, torch.Tensor]:
    fp8_weight_dequant = block_quant_dequant(
        fp8_weight,
        fp8_scale.to(torch.float32),
        weight_block_size,
        torch.bfloat16,
    )
    fp4_weight, fp4_scale = _MXFP4QuantizedData.quantize(
        fp8_weight_dequant, block_size=mxfp4_block_size
    )
    fp4_weight = fp4_weight.quantized_data
    fp4_weight = fp4_weight.contiguous().view(torch.int8)
    fp4_scale = fp4_scale.view(
        *fp8_weight_dequant.shape[:-1],
        fp8_weight_dequant.shape[-1] // mxfp4_block_size,
    )
    return fp4_weight, fp4_scale.contiguous().view(torch.float8_e8m0fnu)


def requant_weight_ue8m0_inplace(weight, weight_scale_inv, weight_block_size):
    assert isinstance(weight, torch.nn.Parameter)
    assert isinstance(weight_scale_inv, torch.nn.Parameter)

    new_weight, new_weight_scale_inv = requant_weight_ue8m0(
        weight.to(weight_scale_inv.device), weight_scale_inv, weight_block_size
    )

    offloader.update_param(weight, new_weight)
    weight_scale_inv.data = new_weight_scale_inv


def requant_block_scale_ue8m0_for_deepgemm(
    weight: torch.nn.Parameter,
    weight_scale: torch.nn.Parameter,
    weight_block_size: Optional[List[int]],
    use_deepgemm_runner: bool,
    output_dtype: Optional[torch.dtype] = None,
    weight_shape=None,
) -> bool:
    """Requantize block-FP8 weight scales to UE8M0 in place for DeepGEMM.

    No-op (returns False) unless the caller selected the DeepGEMM runner, the
    block size is 128x128 (the only layout the requant kernel supports), the
    scales are not already UE8M0, and DeepGEMM can run the layer (bf16 output,
    aligned shape). Returns True when it requantizes.
    """
    from sglang.srt.model_loader.utils import (
        should_deepgemm_weight_requant_ue8m0,
    )

    if (
        not use_deepgemm_runner
        or weight_block_size != [128, 128]
        or getattr(weight_scale, "format_ue8m0", False)
        or not should_deepgemm_weight_requant_ue8m0(
            weight_block_size=weight_block_size,
            output_dtype=output_dtype,
            weight_shape=weight_shape,
        )
    ):
        return False

    requant_weight_ue8m0_inplace(weight, weight_scale, weight_block_size)
    weight_scale.format_ue8m0 = True
    return True


def requant_weight_ue8m0(
    weight: torch.Tensor,
    weight_scale_inv: torch.Tensor,
    weight_block_size: List[int],
):
    assert weight_block_size == [128, 128]

    *_, n, k = weight.shape

    weight_dequant = block_quant_dequant(
        weight,
        weight_scale_inv,
        weight_block_size,
        torch.bfloat16,
    )

    out_w, out_s = quant_weight_ue8m0(
        weight_dequant=weight_dequant,
        weight_block_size=weight_block_size,
    )

    out_s = transform_scale_ue8m0(out_s, mn=out_w.shape[-2])

    return out_w, out_s


def quant_weight_ue8m0(
    weight_dequant: torch.Tensor,
    weight_block_size: List[int],
):
    assert weight_block_size == [128, 128]
    assert (
        weight_dequant.dtype == torch.bfloat16
    ), f"{weight_dequant.dtype=} {weight_dequant.shape=}"

    *batch_dims, n, k = weight_dequant.shape

    weight_dequant_flat = weight_dequant.view((-1, k))
    out_w_flat, out_s_flat = per_block_cast_to_fp8(weight_dequant_flat)

    out_w = out_w_flat.view((*batch_dims, n, k))
    out_s = out_s_flat.view(
        (
            *batch_dims,
            ceil_div(n, weight_block_size[0]),
            ceil_div(k, weight_block_size[1]),
        )
    )

    return out_w, out_s


def transform_scale_ue8m0_inplace(param, mn):
    param.data = transform_scale_ue8m0(param.data, mn=mn)


# NOTE copy and modified from DeepGEMM
def transform_scale_ue8m0(sf, mn, use_torch_impl: bool = False):
    import deep_gemm.utils.layout

    get_mn_major_tma_aligned_packed_ue8m0_tensor = (
        _get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl
        if use_torch_impl
        else deep_gemm.utils.layout.get_mn_major_tma_aligned_packed_ue8m0_tensor
    )

    sf = sf.index_select(-2, torch.arange(mn, device=sf.device) // 128)
    sf = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf)

    # In sgl-deep-gemm, the C++ deepgemm path returns through DLPack which collapses the stride
    # of size-1 trailing dims to 1 (happens when packed_sf_k == 1, i.e.
    # K <= block_k * 4). Restore the TMA-aligned stride so the deepgemm
    # assertion sf.stride(-1) == get_tma_aligned_size(mn, element_size) holds.
    if not use_torch_impl and sf.shape[-1] == 1:
        from deep_gemm.utils import get_tma_aligned_size

        aligned_mn = get_tma_aligned_size(sf.shape[-2], sf.element_size())
        if sf.stride(-1) != aligned_mn:
            new_stride = list(sf.stride())
            new_stride[-1] = aligned_mn
            sf = sf.as_strided(sf.shape, tuple(new_stride))
    return sf


# Copied from DeepGEMM tests
def _get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(
    x: torch.Tensor,
) -> torch.Tensor:
    from deep_gemm.utils import align, get_tma_aligned_size

    assert x.dtype == torch.float and x.dim() in (2, 3)

    # First, convert into UE8M0 `uint8_t`
    ue8m0_tensor = (x.view(torch.int) >> 23).to(torch.uint8)

    # Second, make padded packed tensors
    mn, k = x.shape[-2], x.shape[-1]
    remove_dim = False
    if x.dim() == 2:
        x, remove_dim = x.unsqueeze(0), True
    b = x.shape[0]
    aligned_mn = get_tma_aligned_size(mn, 4)
    aligned_k = align(k, 4)
    padded = torch.zeros((b, aligned_mn, aligned_k), device=x.device, dtype=torch.uint8)
    padded[:, :mn, :k] = ue8m0_tensor
    padded = padded.view(-1).view(dtype=torch.int).view(b, aligned_mn, aligned_k // 4)

    # Finally, transpose
    transposed = torch.zeros(
        (b, aligned_k // 4, aligned_mn), device=x.device, dtype=torch.int
    ).mT
    transposed[:, :, :] = padded
    aligned_x = transposed[:, :mn, :]
    return aligned_x.squeeze(0) if remove_dim else aligned_x


def inverse_transform_scale_ue8m0(sf_packed, mn):
    sf_fp32 = _inverse_transform_scale_ue8m0_impl(sf_packed)
    # Can call consistency check every time since this is only called on startup
    sf_packed_recreated = transform_scale_ue8m0(sf_fp32, mn=mn, use_torch_impl=True)
    assert torch.all(
        sf_packed == sf_packed_recreated
    ), f"{sf_packed=} {sf_packed_recreated=} {sf_fp32=}"
    return sf_fp32


# Inverse impl can refer to DeepGEMM's torch impl in get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl
def _inverse_transform_scale_ue8m0_impl(sf_packed):
    """
    NOTE: We assume k is aligned
    :param sf_packed: (scale_mn, scale_k/4) int32
    :return: (scale_mn, scale_k), float32
    """
    if len(sf_packed.shape) == 3:
        return torch.stack(
            [_inverse_transform_scale_ue8m0_impl(x) for x in sf_packed], dim=0
        )

    block_size = 128
    assert len(sf_packed.shape) == 2, f"{sf_packed.shape=}"
    assert sf_packed.dtype == torch.int32

    mn_repeat_128, k_div_4 = sf_packed.shape
    mn = mn_repeat_128 // block_size
    k = k_div_4 * 4

    # packed u8 -> fp32
    sf_u8 = sf_packed.contiguous().flatten().view(torch.uint8).view(mn_repeat_128, k)
    sf_fp32 = (sf_u8.to(torch.int32) << 23).view(torch.float32)

    # remove repeat
    sf_reshaped = sf_fp32.view(mn, block_size, k)
    sf_unrepeated = sf_reshaped[:, 0:1, :]
    if not torch.all(sf_unrepeated == sf_reshaped):
        from sglang.srt.debug_utils.dumper import get_tensor_info

        raise AssertionError(
            f"sf_unrepeated != sf_reshaped ({get_tensor_info(sf_unrepeated)=} {get_tensor_info(sf_reshaped)=})"
        )
    sf_unrepeated = sf_unrepeated.squeeze(1).contiguous()

    assert sf_unrepeated.shape == (mn, k)
    return sf_unrepeated


# COPIED FROM DeepGEMM
def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    assert x.dim() == 2
    m, n = x.shape
    x_padded = torch.zeros(
        (ceil_align(m, 128), ceil_align(n, 128)), dtype=x.dtype, device=x.device
    )
    x_padded[:m, :n] = x
    x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)
    x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
    sf = ceil_to_ue8m0(x_amax / 448.0)
    x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn)
    return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
        x_view.size(0), x_view.size(2)
    )


# COPIED FROM DeepGEMM
def ceil_to_ue8m0(x: torch.Tensor):
    bits = x.abs().float().view(torch.int32)
    exp = (bits >> 23) & 0xFF
    mantissa = bits & 0x7FFFFF
    exp = exp + (mantissa != 0).to(torch.int32)
    exp = exp.clamp(1, 254)
    return (exp << 23).view(torch.float32)


def channel_quant_to_tensor_quant(
    x_q_channel: torch.Tensor,
    x_s: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    x_dq_channel = x_q_channel.to(torch.float32) * x_s
    x_q_tensor, scale = (
        scaled_fp8_quant(x_dq_channel)
        if _is_cuda
        else input_to_float8(x_dq_channel, dtype=x_q_channel.dtype)
    )
    return x_q_tensor, scale


def _process_scaled_mm_output(output, input_2d_shape, output_shape):
    if type(output) is tuple and len(output) == 2:
        output = output[0]
    return torch.narrow(output, 0, 0, input_2d_shape[0]).view(*output_shape)


def _apply_fallback_scaled_mm(
    qinput,
    weight,
    x_scale,
    weight_scale,
    input_2d_shape,
    output_shape,
    bias,
    input_dtype,
):
    global TORCH_DEVICE_IDENTITY
    if TORCH_DEVICE_IDENTITY is None:
        TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32, device=weight.device)

    output = torch._scaled_mm(
        qinput,
        weight,
        scale_a=TORCH_DEVICE_IDENTITY,
        scale_b=TORCH_DEVICE_IDENTITY,
        out_dtype=torch.float32,
    )

    output = _process_scaled_mm_output(output, input_2d_shape, output_shape)
    x_scale = torch.narrow(x_scale, 0, 0, input_2d_shape[0])

    output = output * x_scale * weight_scale.t()
    if bias is not None:
        output = output + bias
    return output.to(dtype=input_dtype)


def apply_fp8_linear_bmm_flashinfer(
    input: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Per-tensor static fp8 linear via flashinfer bmm_fp8 (SM10X only)."""
    output_shape = [*input.shape[:-1], weight.shape[1]]
    input_2d = input.view(-1, input.shape[-1])
    qinput, x_scale = static_quant_fp8(input_2d, input_scale, repeat_scale=False)
    output = flashinfer_bmm_fp8(qinput, weight, x_scale, weight_scale, input.dtype)
    if bias is not None:
        output = output + bias
    return output.view(*output_shape)


def apply_fp8_linear(
    input: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    input_scale_ub: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
    cutlass_fp8_supported: bool = cutlass_fp8_supported(),
    use_per_token_if_dynamic: bool = False,
    pad_output: Optional[bool] = None,
    compressed_tensor_quant: bool = False,
) -> torch.Tensor:
    # Note: we pad the input because torch._scaled_mm is more performant
    # for matrices with batch dimension > 16.
    # This could change in the future.
    # We also don't pad when using torch.compile,
    # as it breaks with dynamic shapes.
    if pad_output is None:
        pad_output = not cutlass_fp8_supported and not get_bool_env_var(
            "SGLANG_ENABLE_TORCH_COMPILE"
        )
    output_padding = 17 if pad_output else None

    # View input as 2D matrix for fp8 methods
    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[1]]

    if compressed_tensor_quant:
        # Maybe apply padding to output, see comment in __init__
        num_token_padding = output_padding
        if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
            num_token_padding = None
        # For static per-tensor activation scales when using inductor compiler,
        # use pure PyTorch ops instead of the opaque sgl_kernel quant kernel.
        # Inductor fuses these with surrounding ops (RMSNorm, residual add),
        # eliminating a separate kernel launch per linear layer.
        # weight_scale shape does not matter here -- it is only used in the
        # GEMM epilogue, not in the activation quant fusion. Only activates when
        # cuda_graph_config[prefill].tc_compiler=inductor; eager PCG and
        # decode both use the faster custom kernel.

        if (
            input_scale is not None
            and input_scale.numel() == 1
            and get_server_args().cuda_graph_config.prefill.tc_compiler == "inductor"
        ):
            qinput = (
                (input_2d * input_scale.reciprocal())
                .clamp(min=fp8_min, max=fp8_max)
                .to(fp8_dtype)
            )
            x_scale = input_scale
        else:
            qinput, x_scale = scaled_fp8_quant(
                input_2d,
                input_scale,
                num_token_padding=num_token_padding,
                use_per_token_if_dynamic=use_per_token_if_dynamic,
            )
    else:
        # cutlass w8a8 fp8 sgl-kernel only supports per-token scale
        if input_scale is not None:
            assert input_scale.numel() == 1
            # broadcast per-tensor scale to per-token scale when supporting cutlass
            qinput, x_scale = static_quant_fp8(
                input_2d, input_scale, repeat_scale=cutlass_fp8_supported
            )
        else:
            # default use per-token quantization if dynamic
            if _is_cuda:
                qinput, x_scale = sglang_per_token_quant_fp8(input_2d)
            else:
                # TODO(kkhuang): temporarily enforce per-tensor activation scaling if weight is per-tensor scaling
                # final solution should be: 1. add support to per-tensor activation scaling.
                # 2. solve the torch.compile error from weight_scale.numel() == 1 and x_scale.numel() > 1 (below line#308)
                if _is_hip and weight_scale.numel() == 1:
                    qinput, x_scale = scaled_fp8_quant(
                        input_2d,
                        input_scale,
                        use_per_token_if_dynamic=use_per_token_if_dynamic,
                    )
                else:
                    qinput, x_scale = per_token_group_quant_fp8(
                        input_2d, group_size=input_2d.shape[1]
                    )

    if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
        cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
        if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
            # Massage the input to be 2D
            qinput = qinput.view(-1, qinput.shape[-1])
            output = triton_scaled_mm(
                qinput, weight, x_scale, weight_scale, input.dtype, bias
            )
        else:
            output = fp8_scaled_mm(
                qinput,
                weight,
                x_scale,
                weight_scale,
                out_dtype=input.dtype,
                bias=bias,
            )
        return output.view(*output_shape)

    # torch.scaled_mm supports per tensor weights + activations only
    # so fallback to naive if per channel or per token
    per_tensor_weights = weight_scale.numel() == 1
    # When the number of token is 1,
    # per-token scale has shape (1, 1), per-tensor scale has shape (1) or ().
    per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2

    if (
        use_per_token_if_dynamic
        and not per_tensor_weights
        and not per_tensor_activations
        and (USE_ROWWISE_TORCH_SCALED_MM or _use_aiter)
    ):
        # into this sector means use dynamic per-token-per-channel quant
        # per-token scale quant for input matrix, every row(one token) have one scale factor
        # per-channel scale quant for weight matrix, every col(one channel) have one scale factor
        if _use_aiter:
            # gemm_a8w8_bpreshuffle(XQ, WQ, x_scale, w_scale, dtype)
            # XQ -> input tensor, shape = (m, k)
            # WQ -> weight tensor, shape = (n, k), with preshuffe get better perf
            # x_scale -> input scale tensor, shape = (m, 1)
            # w_scale -> weight scale tensor, shape = (n ,1)
            # dtype -> output dtype
            output = gemm_a8w8_bpreshuffle(
                XQ=qinput,
                WQ=weight.T,
                x_scale=x_scale,
                w_scale=weight_scale,
                dtype=input.dtype,
            )
            if bias is not None:
                output += bias
            return _process_scaled_mm_output(output, input_2d.shape, output_shape)
        else:
            # For now validated on ROCm platform
            # fp8 rowwise scaling in torch._scaled_mm is introduced in
            # https://github.com/pytorch/pytorch/pull/144432 using hipBLASLt
            # and ROCm 6.3, which only exists in torch 2.7 and above.
            # For CUDA platform please validate if the
            # torch._scaled_mm support rowwise scaled GEMM
            # Fused GEMM_DQ Rowwise GEMM
            output = torch._scaled_mm(
                qinput,
                weight,
                out_dtype=input.dtype,
                scale_a=x_scale,
                scale_b=weight_scale.t(),
                bias=bias,
            )
            return _process_scaled_mm_output(output, input_2d.shape, output_shape)

    if per_tensor_weights and per_tensor_activations:
        # Fused GEMM_DQ; _scaled_mm with torch.compile requires len(weight_scale.shape) == len(x_scale.shape)
        if weight_scale.ndim == 0 and x_scale.ndim == 1:
            weight_scale = weight_scale.unsqueeze(0)
        output = torch._scaled_mm(
            qinput,
            weight,
            out_dtype=input.dtype,
            scale_a=x_scale,
            scale_b=weight_scale,
            bias=bias,
        )
        return _process_scaled_mm_output(output, input_2d.shape, output_shape)

    # Fallback for channelwise case, where we use unfused DQ
    # due to limitations with scaled_mm

    # Symmetric quantized GEMM by definition computes the following:
    #   C = (s_x * X) (s_w * W) + bias
    # This is equivalent to dequantizing the weights and activations
    # before applying a GEMM.
    #
    # In order to compute quantized operands, a quantized kernel
    # will rewrite the above like so:
    #   C = s_w * s_x * (X * W) + bias
    #
    # For the scaled_mm fallback case, we break this down, since it
    # does not support s_w being a vector.
    return _apply_fallback_scaled_mm(
        qinput,
        weight,
        x_scale,
        weight_scale,
        input_2d.shape,
        output_shape,
        bias,
        input.dtype,
    )


def can_auto_enable_marlin_fp8() -> bool:
    try:
        major, minor = get_device_capability()
        sm = major * 10 + minor
        return 80 <= sm < 89
    except Exception:
        return False


def apply_fp8_ptpc_linear(
    input: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    input_scale_ub: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
    cutlass_fp8_supported: bool = cutlass_fp8_supported(),
    use_per_token_if_dynamic: bool = False,
    pad_output: Optional[bool] = None,
    compressed_tensor_quant: bool = False,
) -> torch.Tensor:
    """FP8 per-token per-channel linear. Only used with the aiter (ROCm) backend."""
    # Handle pre-quantized (fp8_tensor, scale) tuple from fused RMSNorm+Quant
    if isinstance(input, tuple):
        q_input, x_scale = input
        q_input = q_input.view(-1, q_input.shape[-1])
        output_shape = [*q_input.shape[:-1], weight.shape[0]]
        output = aiter.gemm_a8w8_bpreshuffle(
            q_input, weight, x_scale, weight_scale, None, torch.bfloat16
        )
        if bias is not None:
            output = output + bias
        return output.view(*output_shape)

    # View input as 2D matrix for fp8 methods
    input_2d = input.view(-1, input.shape[-1])

    # weight is transposed (K, N)
    output_shape = [*input.shape[:-1], weight.shape[1]]

    q_input, x_scale = aiter.per_token_quant_hip(input_2d, quant_dtype=aiter.dtypes.fp8)

    per_tensor_weights = (weight_scale.numel() == 1) and weight_scale.dim() < 2
    per_tensor_activations = (x_scale.numel() == 1) and x_scale.dim() < 2

    if not (per_tensor_weights and per_tensor_activations):
        # weight is in (N, K)
        output_shape = [*input.shape[:-1], weight.shape[0]]

    output = aiter.gemm_a8w8_bpreshuffle(
        q_input, weight, x_scale, weight_scale, None, input.dtype
    )
    if bias is not None:
        output = output + bias
    return output.view(*output_shape)


def validate_fp8_block_shape(
    layer: torch.nn.Module,
    input_size: int,
    output_size: int,
    input_size_per_partition: int,
    output_partition_sizes: list[int],
    block_size: list[int],
) -> None:
    """Validate block quantization shapes for tensor parallelism."""

    tp_size = getattr(layer, "tp_size", get_parallel().tp_size)
    block_n, block_k = block_size[0], block_size[1]

    # Required by row parallel
    if (
        tp_size > 1
        and input_size // input_size_per_partition == tp_size
        and input_size_per_partition % block_k != 0
    ):
        raise ValueError(
            f"Weight input_size_per_partition = {input_size_per_partition} "
            f"is not divisible by weight quantization block_k = {block_k}."
        )

    # Required by column parallel or enabling merged weights
    is_tp_split = tp_size > 1 and output_size // sum(output_partition_sizes) == tp_size
    is_merged_gemm = len(output_partition_sizes) > 1
    if is_tp_split or is_merged_gemm:
        sizes_to_check = output_partition_sizes
        if not is_tp_split and is_merged_gemm:
            # In case of merged matrices, we allow the last
            # matrix to not be a multiple of block size
            sizes_to_check = output_partition_sizes[:-1]
        for output_partition_size in sizes_to_check:
            if output_partition_size % block_n != 0:
                raise ValueError(
                    f"Weight output_partition_size = "
                    f"{output_partition_size} is not divisible by "
                    f"weight quantization block_n = {block_n}."
                )