diff --git a/python/sglang/kernels/ops/quantization/mxfp8_aiter_gfx95.py b/python/sglang/kernels/ops/quantization/mxfp8_aiter_gfx95.py new file mode 100644 index 000000000000..b56f09eae26e --- /dev/null +++ b/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) diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py index 15b6a831f20a..bfcb3926ae84 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/fp8_hip.py b/python/sglang/srt/layers/quantization/fp8_hip.py index 9e5ea1e2d0de..5223e3c170c9 100644 --- a/python/sglang/srt/layers/quantization/fp8_hip.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 04b542a54955..d69e691b84e4 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_common/amd/deepseek_v2_hip_act.py b/python/sglang/srt/models/deepseek_common/amd/deepseek_v2_hip_act.py index ab48d486cb19..9b62db3f8f07 100644 --- a/python/sglang/srt/models/deepseek_common/amd/deepseek_v2_hip_act.py +++ b/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, ) diff --git a/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_gfx95_dense.py b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_gfx95_dense.py index 898bb9450914..360c10afa1ba 100644 --- a/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_gfx95_dense.py +++ b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_gfx95_dense.py @@ -59,22 +59,23 @@ 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 @@ -82,8 +83,13 @@ def _native_mxfp8_consumer(linear: Optional[nn.Module]) -> bool: 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 diff --git a/test/registered/kernels/ops/quantization/test_mxfp8_amd_gfx95.py b/test/registered/kernels/ops/quantization/test_mxfp8_amd_gfx95.py index 09f2b8d9a881..f49f322789cc 100644 --- a/test/registered/kernels/ops/quantization/test_mxfp8_amd_gfx95.py +++ b/test/registered/kernels/ops/quantization/test_mxfp8_amd_gfx95.py @@ -1,10 +1,13 @@ -"""The gfx950 native MXFP8 GEMV and dense route against fp64 and the bf16-dequant route: within one bf16 ulp, repeatable, batch-invariant.""" +"""The gfx950 native MXFP8 GEMV and the native / aiter dense routes against fp64 and the bf16-dequant route: within one bf16 ulp, repeatable, batch-invariant.""" +import types import unittest import torch from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import ( + Fp8GridActivation, + Mxfp8Activation, bf16_dequant_blockscaled_linear, dequant_block_fp8_weight_to_bf16, fake_quant_fp8_activation, @@ -15,6 +18,11 @@ mxfp8_native_blockscaled_linear, prepare_mxfp8_native_weight, ) +from sglang.srt.layers.quantization import fp8_hip +from sglang.srt.layers.quantization.fp8_utils import ( + Mxfp8DenseGemmBackend, + dispatch_block_fp8_mxfp8_linear, +) from sglang.srt.utils import is_gfx95_supported, is_hip from sglang.test.ci.ci_register import register_amd_ci from sglang.test.test_utils import CustomTestCase @@ -140,5 +148,65 @@ def test_within_one_bf16_ulp_of_the_bf16_route(self): self.assertTrue(torch.equal(out_q, out), (n, k, m)) +def _aiter_mxfp8_gemm_available() -> bool: + try: + from aiter.ops.triton.gemm.basic import ( # noqa: F401 + gemm_a8w8_blockscale_group32, + ) + except ImportError: + return False + return True + + +# decode / verify rows and the dot_scaled buckets the native test covers +AITER_MS = (1, 6, 33, 1025) + + +@unittest.skipUnless( + is_hip() and is_gfx95_supported() and _aiter_mxfp8_gemm_available(), + "gfx950 aiter MXFP8 route (ROCm/aiter#5750)", +) +class TestMxfp8AiterRouteGfx95(CustomTestCase): + def _route(self, wq, ws): + """apply_dense of a layer on the aiter route, its weights processed as at load time.""" + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(wq.clone(), requires_grad=False) + layer.weight_scale_inv = torch.nn.Parameter(ws.clone(), requires_grad=False) + layer.block_fp8_mxfp8_ready = True + backend = Mxfp8DenseGemmBackend.GFX95_MXFP8_AITER + method = types.SimpleNamespace( + mxfp8_dense_backend=backend, + weight_block_size=[32, 32], + w8a8_mxfp8_linear=dispatch_block_fp8_mxfp8_linear(backend), + ) + fp8_hip.process_dense_weights(method, layer, None) + return lambda x: fp8_hip.apply_dense(method, layer, x, None) + + def test_within_one_bf16_ulp_of_the_bf16_route(self): + for n, k in ROUTE_SHAPES: + torch.manual_seed(0) + w = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) + wq, ws = _quant_weight_block32(w) + w_bf16 = dequant_block_fp8_weight_to_bf16(wq, ws, [32, 32]) + route = self._route(wq, ws) + for m in AITER_MS: + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + ref = bf16_dequant_blockscaled_linear(x, w_bf16) + out = route(x) + row_max = ref.float().abs().amax(dim=1, keepdim=True).clamp(min=1.0) + ulp_of_row_max = torch.exp2(torch.floor(torch.log2(row_max)) - 7) + diff = (out.float() - ref.float()).abs() + self.assertTrue( + bool((diff <= ulp_of_row_max).all()), + (n, k, m, (diff / ulp_of_row_max).max().item()), + ) + # the producers' fp8-grid and fp8 + scales operands quantize to the same codes + out_grid = route(Fp8GridActivation(fake_quant_fp8_activation(x))) + self.assertTrue(torch.equal(out_grid, out), (n, k, m)) + xq, xs = fp8_grid_quantize(x) + out_q = route(Mxfp8Activation(xq, xs)) + self.assertTrue(torch.equal(out_q, out), (n, k, m)) + + if __name__ == "__main__": unittest.main()