[AMD][V4.1][*/N] Switch the fp8 dense GEMMs on gfx950 to aiter's MXFP8 GEMM

python/sglang/kernels/ops/quantization/mxfp8_aiter_gfx95.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+"""gfx950 aiter dense route for 32x32-block ue8m0 fp8 checkpoints: aiter's MXFP8 A8W8 GEMM on the
+plain [N, K] fp8 weight and its compact [N/32, K/32] ue8m0 block scales, with the activation as fp8 +
+ue8m0 [M, K/32]; aiter picks the kernel config from its tuned MXFP8 tables."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from aiter import gemm_a8w8_blockscale
+
+from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import mxfp8_e4m3_quantize
+
+
+def aiter_mxfp8_blockscaled_linear(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale_ue8m0: torch.Tensor,
+    input_scale: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Dense linear of the aiter route. input is bf16 (plain or on the fp8 grid, where the per-32
+    requant is exact), or fp8 e4m3 with input_scale ue8m0 [M, K/32]."""
+    input_2d = input.view(-1, input.shape[-1])
+    n = weight.shape[0]
+    if input_2d.shape[0] == 0:
+        out = input_2d.new_empty((0, n), dtype=torch.bfloat16)
+    else:
+        if input_scale is None:
+            xq, xs = mxfp8_e4m3_quantize(input_2d.to(torch.bfloat16))
+        else:
+            assert input_2d.dtype == torch.float8_e4m3fn, input_2d.dtype
+            xq, xs = input_2d.contiguous(), input_scale.contiguous()
+        out = gemm_a8w8_blockscale(
+            xq, weight, xs, weight_scale_ue8m0, dtype=torch.bfloat16
+        )
+    if bias is not None:
+        out = out + bias
+    return out.view(*input.shape[:-1], n)

python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py

@@ -436,13 +436,23 @@ def commit_kv_proj_fused(
     stacked = _stacked_wkv_weight(wkv_linears=wkv_linears)
 
     if stacked.mxfp8_scale is not None:
-        kv_all = wkv_linears[0].quant_method.w8a8_mxfp8_linear(
-            input=main_x,
-            weight=stacked.weight,
-            weight_scale=stacked.mxfp8_scale,
-            input_scale=None,
-            bias=None,
-        )
+        quant_method = wkv_linears[0].quant_method
+        if quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_aiter():
+            kv_all = quant_method.w8a8_mxfp8_linear(
+                input=main_x,
+                weight=stacked.weight,
+                weight_scale_ue8m0=stacked.mxfp8_scale,
+                input_scale=None,
+                bias=None,
+            )
+        else:
+            kv_all = quant_method.w8a8_mxfp8_linear(
+                input=main_x,
+                weight=stacked.weight,
+                weight_scale=stacked.mxfp8_scale,
+                input_scale=None,
+                bias=None,
+            )
     elif stacked.fp8_scale is not None:
         quant_method = wkv_linears[0].quant_method
         kv_all = quant_method.w8a8_block_fp8_linear(
@@ -481,6 +491,18 @@ def _block_quant_stack_applies(*, wkv_linears: list[torch.nn.Module]) -> bool:
     block_quant = hasattr(quant_method, "block_quant") and quant_method.block_quant
     if not (block_quant and hasattr(quant_method, "w8a8_block_fp8_linear")):
         return False
+    # on gfx950 only the aiter MXFP8 route stacks (the Triton block kernel's ue8m0 activation
+    # quant is CUDA-only); the others serve 32x32-block weights through apply()
+    backend = getattr(quant_method, "mxfp8_dense_backend", None)
+    if backend is not None and backend.is_gfx95():
+        if not (
+            backend.is_gfx95_mxfp8_aiter()
+            and all(
+                getattr(linear, "block_fp8_mxfp8_ready", False)
+                for linear in wkv_linears
+            )
+        ):
+            return False
     block_out = quant_method.quant_config.weight_block_size[0]
     # the gfx950 native MXFP8 route keeps the weight in a 3-D lane-order layout
     return all(
@@ -536,6 +558,12 @@ def _build_stacked_wkv_weight(
                 [linear.weight_scale_inv_swizzled.reshape(-1) for linear in wkv_linears]
             )
             return _StackedWkvWeight(weight=weight, fp8_scale=None, mxfp8_scale=scale)
+        if backend is not None and backend.is_gfx95_mxfp8_aiter():
+            # compact ue8m0 block scales, one row per 32 output rows: whole blocks stack
+            scale = torch.cat(
+                [linear.weight_scale_mx_e8m0 for linear in wkv_linears], dim=0
+            )
+            return _StackedWkvWeight(weight=weight, fp8_scale=None, mxfp8_scale=scale)
         if wkv_linears[0].weight_scale_inv.dtype == torch.int32:
             from sglang.srt.layers.quantization.fp8_utils import (
                 inverse_transform_scale_ue8m0,

python/sglang/srt/layers/quantization/fp8_hip.py

@@ -1,7 +1,8 @@
 """gfx950 dense routes of Fp8LinearMethod for 32x32-block fp8 checkpoints served as
-MXFP8 (block_fp8_as_mxfp8): the Triton dot_scaled kernel, or the native scaled-MFMA
-kernels on a lane-ordered weight. The fused gfx950 producers hand these routes their
-operand as an Fp8GridActivation (bf16 already on the fp8 grid) or an Mxfp8Activation."""
+MXFP8 (block_fp8_as_mxfp8): the Triton dot_scaled kernel, the native scaled-MFMA
+kernels on a lane-ordered weight, or aiter's MXFP8 GEMM on the plain one. The fused gfx950
+producers hand these routes their operand as an Fp8GridActivation (bf16 already on the fp8
+grid) or an Mxfp8Activation."""
 
 from __future__ import annotations
 
@@ -19,6 +20,7 @@
 from sglang.kernels.ops.quantization.mxfp8_native_amd_gfx95 import (
     native_route_supports,
     prepare_mxfp8_native_weight,
+    ue8m0_weight_scale,
 )
 from sglang.srt.layers.utils import copy_or_rebind_param
 
@@ -31,6 +33,14 @@ def process_dense_weights(method, layer: torch.nn.Module, scale_u8) -> None:
         if scale_u8 is not None:
             copy_or_rebind_param(layer, "weight_scale_inv_mx", scale_u8.contiguous())
         return
+    if backend.is_gfx95_mxfp8_aiter():
+        # the weight stays [N, K]; aiter reads the block scales compact, one row per 32 output rows
+        copy_or_rebind_param(
+            layer,
+            "weight_scale_mx_e8m0",
+            ue8m0_weight_scale(layer.weight_scale_inv.data),
+        )
+        return
     assert backend.is_gfx95_mxfp8_native()
     n, k = layer.weight.shape
     layer.mxfp8_native_ready = False
@@ -65,12 +75,13 @@ def apply_dense(
     backend = method.mxfp8_dense_backend
     mxfp8_ready = layer.block_fp8_mxfp8_ready
     native_route = mxfp8_ready and backend.is_gfx95_mxfp8_native()
-    # Unwrap the producer's operand: the native route takes fp8 + scales or the fp8-grid
-    # bf16 directly; the dot_scaled route quantizes the plain tensor itself (per-32
-    # rounding is idempotent, so the wrapper's rounding is exact for it).
+    aiter_route = mxfp8_ready and backend.is_gfx95_mxfp8_aiter()
+    # Unwrap the producer's operand: the native and aiter routes take fp8 + scales directly,
+    # the native one also the fp8-grid bf16; the other routes quantize the plain tensor
+    # themselves (per-32 rounding is idempotent, so the wrapper's rounding is exact for them).
     input_scale, on_fp8_grid = None, False
     if isinstance(x, Mxfp8Activation):
-        if native_route:
+        if native_route or aiter_route:
             x, input_scale = x.q, x.scale
         else:
             x = dequant_mxfp8_to_bf16(x.q, x.scale)
@@ -80,6 +91,14 @@ def apply_dense(
         x, input_scale = x
     if native_route:
         return _apply_native(method, layer, x, bias, input_scale, on_fp8_grid)
+    if aiter_route:
+        return method.w8a8_mxfp8_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale_ue8m0=layer.weight_scale_mx_e8m0,
+            input_scale=input_scale,
+            bias=bias,
+        )
     if mxfp8_ready and input_scale is None:
         return method.w8a8_mxfp8_linear(
             input=x,

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -383,6 +383,8 @@ class Mxfp8DenseGemmBackend(Enum):
     GFX95_DOT_SCALED = "gfx95_dot_scaled"
     # gfx950 native MXFP8: lane-ordered fp8 weight + ue8m0 scale bytes; untiled shapes keep bf16
     GFX95_MXFP8_NATIVE = "gfx95_mxfp8_native"
+    # gfx950 aiter MXFP8 GEMM: plain fp8 weight + compact 32x32 ue8m0 block scale bytes
+    GFX95_MXFP8_AITER = "gfx95_mxfp8_aiter"
     UNSUPPORTED = "unsupported"
 
     def is_flashinfer_cutlass(self) -> bool:
@@ -406,10 +408,14 @@ def is_gfx95_dot_scaled(self) -> bool:
     def is_gfx95_mxfp8_native(self) -> bool:
         return self == Mxfp8DenseGemmBackend.GFX95_MXFP8_NATIVE
 
+    def is_gfx95_mxfp8_aiter(self) -> bool:
+        return self == Mxfp8DenseGemmBackend.GFX95_MXFP8_AITER
+
     def is_gfx95(self) -> bool:
         return self in (
             Mxfp8DenseGemmBackend.GFX95_DOT_SCALED,
             Mxfp8DenseGemmBackend.GFX95_MXFP8_NATIVE,
+            Mxfp8DenseGemmBackend.GFX95_MXFP8_AITER,
         )
 
     def is_unsupported(self) -> bool:
@@ -742,6 +748,8 @@ def resolve_block_fp8_mxfp8_backend() -> Mxfp8DenseGemmBackend:
     if _is_hip and _is_gfx95_supported:
         if backend.is_triton():
             return Mxfp8DenseGemmBackend.GFX95_DOT_SCALED
+        if backend.is_aiter():
+            return Mxfp8DenseGemmBackend.GFX95_MXFP8_AITER
         return Mxfp8DenseGemmBackend.GFX95_MXFP8_NATIVE
     # Explicit CUTLASS / CuTe-DSL only: they leave the weight untouched and store
     # the swizzled scale separately, so the block layout stays readable by Triton.
@@ -787,6 +795,12 @@ def dispatch_block_fp8_mxfp8_linear(backend: Mxfp8DenseGemmBackend) -> Callable:
         )
 
         return mxfp8_native_blockscaled_linear
+    if backend.is_gfx95_mxfp8_aiter():
+        from sglang.kernels.ops.quantization.mxfp8_aiter_gfx95 import (
+            aiter_mxfp8_blockscaled_linear,
+        )
+
+        return aiter_mxfp8_blockscaled_linear
     return _unsupported_mxfp8_linear
 
 

python/sglang/srt/models/deepseek_common/amd/deepseek_v2_hip_act.py

@@ -23,19 +23,28 @@ def resolve_fused_clamp_route(mlp, half_width: int) -> None:
         and quant_method.block_quant
         and quant_method.weight_block_size == [128, 128]
     )
-    # gfx950 32-block route: the activation lands on the fp8 grid, so down_proj skips its fake-quant
+    # gfx950 32-block native or aiter route: the activation lands on the fp8 grid, so down_proj
+    # skips its fake-quant
     mlp._hip_act_fp8_grid = bool(
         isinstance(quant_method, Fp8LinearMethod)
         and quant_method.block_fp8_as_mxfp8
         and mlp.down_proj.block_fp8_mxfp8_ready
-        and quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_native()
+        and (
+            quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_native()
+            or quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_aiter()
+        )
         and silu_and_mul_clamp_fp8_grid_supported(half_width)
     )
-    # the native MXFP8 route takes fp8 + ue8m0 straight from the epilogue at decode token counts
+    # the native and aiter MXFP8 routes take fp8 + ue8m0 straight from the epilogue
     mlp._hip_act_native_consumer = bool(
         mlp._hip_act_fp8_grid
-        and quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_native()
-        and mlp.down_proj.mxfp8_native_ready
+        and (
+            quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_aiter()
+            or (
+                quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_native()
+                and mlp.down_proj.mxfp8_native_ready
+            )
+        )
     )
     mlp._fused_clamp_fp8_checked = True
 
@@ -48,6 +57,6 @@ def silu_and_mul_clamp(mlp, gate_up: torch.Tensor):
         gate_up,
         float(mlp.swiglu_limit),
         fp8_grid=mlp._hip_act_fp8_grid,
-        # the native down_proj runs on fp8 + ue8m0 at every M, so hand it those directly
+        # the native / aiter down_proj runs on fp8 + ue8m0 at every M, so hand it those directly
         emit_fp8=mlp._hip_act_native_consumer,
     )

python/sglang/srt/models/deepseek_common/amd/deepseek_v4_gfx95_dense.py

@@ -59,31 +59,37 @@ def fused_rmsnorm_fake_quant_eligible(
     quant_config: Optional[QuantizationConfig],
 ) -> bool:
     """Whether rmsnorm_fake_quant_fp8 applies: gfx950 with a 32-wide-block checkpoint
-    (V4.1), whose dense route takes the norm output already on the fp8 grid as an
-    Fp8GridActivation."""
+    (V4.1), whose native or aiter dense route takes the norm output already on the fp8
+    grid (Fp8GridActivation) or as fp8 + ue8m0 (Mxfp8Activation)."""
     if not (_is_hip and _is_gfx95_supported and isinstance(quant_config, Fp8Config)):
         return False
     block = quant_config.weight_block_size
+    backend = resolve_block_fp8_mxfp8_backend()
     return (
         block is not None
         and block[1] == 32
         and quant_config.scale_fmt == "ue8m0"
-        and resolve_block_fp8_mxfp8_backend().is_gfx95_mxfp8_native()
+        and (backend.is_gfx95_mxfp8_native() or backend.is_gfx95_mxfp8_aiter())
     )
 
 
-def _native_mxfp8_consumer(linear: Optional[nn.Module]) -> bool:
-    """Whether linear runs the gfx950 native MXFP8 route, which consumes fp8 + ue8m0
-    scales directly at every M."""
+def _mxfp8_consumer(linear: Optional[nn.Module]) -> bool:
+    """Whether linear runs a gfx950 MXFP8 route that consumes fp8 + ue8m0 scales directly
+    at every M: aiter's MXFP8 GEMM, or the native kernels on a tiled weight."""
     if linear is None:
         return False
     quant_method = linear.quant_method
     return (
         isinstance(quant_method, Fp8LinearMethod)
         and quant_method.block_fp8_as_mxfp8
         and linear.block_fp8_mxfp8_ready
-        and quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_native()
-        and linear.mxfp8_native_ready
+        and (
+            quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_aiter()
+            or (
+                quant_method.mxfp8_dense_backend.is_gfx95_mxfp8_native()
+                and linear.mxfp8_native_ready
+            )
+        )
     )
 
 
@@ -99,7 +105,7 @@ def q_norm_fake_quant(attn, q_lora: torch.Tensor) -> Tuple[torch.Tensor, object]
         q_lora = attn.q_norm(q_lora)
         return q_lora, q_lora
     if not attn._wq_b_native_consumer_checked:
-        attn._wq_b_native_consumer = _native_mxfp8_consumer(attn.wq_b)
+        attn._wq_b_native_consumer = _mxfp8_consumer(attn.wq_b)
         attn._wq_b_native_consumer_checked = True
     q_for_wq_b, q_lora = rmsnorm_fake_quant_fp8(
         q_lora,
@@ -123,7 +129,7 @@ def input_norm_fake_quant(
         return norm(hidden_states), None
     if not layer._wqkv_a_native_consumer_checked:
         # wqkv_a exists only when the q / kv projections are fused
-        layer._wqkv_a_native_consumer = _native_mxfp8_consumer(
+        layer._wqkv_a_native_consumer = _mxfp8_consumer(
             layer.self_attn.wqkv_a if layer.self_attn.fuse_wqa_wkv else None
         )
         layer._wqkv_a_native_consumer_checked = True