diff --git a/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh index 28b9245e3..53356ff89 100644 --- a/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh +++ b/python/sglang/kernels/jit/include/sgl_kernel/utils.cuh @@ -167,7 +167,10 @@ template SGL_DEVICE void PDLTriggerSecondary() { #if SGL_ARCH_HOPPER_OR_GREATER if constexpr (kUsePDL) { - asm volatile("griddepcontrol.launch_dependents;" :::); + // The "memory" clobber is load-bearing: without it the compiler may sink + // this kernel's stores past the trigger, and the dependent grid's + // griddepcontrol.wait only covers writes issued BEFORE launch_dependents. + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); } #endif } diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 2e79c953e..7f4145353 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1293,6 +1293,12 @@ def mxfp8_group_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: 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" + if _is_hip and _is_gfx95_supported: + from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import ( + mxfp8_e4m3_quantize, + ) + + return mxfp8_e4m3_quantize(x) 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() diff --git a/test/registered/kernels/ops/moe/test_minimax_m3_mxfp8.py b/test/registered/kernels/ops/moe/test_minimax_m3_mxfp8.py index bc9f18d94..1b3a4f567 100644 --- a/test/registered/kernels/ops/moe/test_minimax_m3_mxfp8.py +++ b/test/registered/kernels/ops/moe/test_minimax_m3_mxfp8.py @@ -33,6 +33,9 @@ from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import ( # noqa: E402 _mxfp8_e4m3_quantize_triton, dequant_mxfp8_to_bf16, ) +from sglang.srt.layers.quantization.fp8_utils import ( # noqa: E402 + mxfp8_group_quantize, +) from sglang.test.ci.ci_register import register_amd_ci register_amd_ci(est_time=20, stage="jit-kernel-unit", runner_config="amd") @@ -81,6 +84,24 @@ def test_mxfp8_quant_triton_matches_torch(shape, dtype): assert _relerr(deq_k, deq_t) < 1e-2 +@requires_gfx950 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_mxfp8_group_quantize_uses_gfx950_quantizer(dtype): + torch.manual_seed(0) + x = torch.randn(64, 128, device=DEVICE, dtype=dtype) + + q, scale = mxfp8_group_quantize(x) + expected_q, expected_scale = _mxfp8_e4m3_quantize_triton(x) + + assert q.dtype == torch.float8_e4m3fn + assert scale.dtype == torch.uint8 + assert q.shape == x.shape + assert scale.shape == (64, 4) + torch.testing.assert_close(q.float(), expected_q.float(), rtol=0, atol=0) + torch.testing.assert_close(scale, expected_scale, rtol=0, atol=0) + + @pytest.mark.parametrize("m,inter", [(8, 512), (65, 2048)]) @torch.inference_mode() def test_minimax_swiglu_mxfp8_quant_matches_unfused_fp32(m, inter):