Fix GPU kernel ordering and MXFP8 quantization dispatch (#37331)

Co-authored-by: Pranjal Shankhdhar <pranjalssh@meta.com>
Co-authored-by: Chengze Fan <fancz2002@gmail.com>

python/sglang/kernels/jit/include/sgl_kernel/utils.cuh

@@ -167,7 +167,10 @@ template <bool kUsePDL>
 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
 }

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()

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):