[Perf] [AMD] qwen3_5: Quantize the BF16 GDN projections to MXFP4 online

docs/docs/references/environment_variables.mdx

@@ -1932,6 +1932,11 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Token count above which <code>SGLANG_ROCM_K3_FUSE_KDA_INPROJ</code> stops applying and the split projection runs instead. The merged shape is only faster while the projection is bandwidth bound.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>256</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ROCM_GDN_ONLINE_MXFP4</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>On gfx95x with <code>SGLANG_USE_AITER=1</code>, quantize the BF16 GDN <code>in_proj_qkvz</code> and <code>out_proj</code> of Qwen3.5 / Qwen3.8 to MXFP4 after load. Projections with N or K not a multiple of 256 stay BF16. Batches of up to 64 rows run the AITER preshuffled MXFP4 GEMM if AITER has a tuned config for the shape. All other batches run the ASM <code>gemm_a4w4</code>. Set to <code>0</code> to keep the BF16 projections.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
+    </tr>
   </tbody>
 </table>
 

python/sglang/srt/environ.py

@@ -1046,6 +1046,9 @@ class Envs:
     # output columns ride along nearly free.
     SGLANG_ROCM_K3_FUSE_KDA_INPROJ = EnvBool(True)
     SGLANG_ROCM_K3_FUSE_KDA_INPROJ_MAX_TOKENS = EnvInt(256)
+    # On gfx95x with SGLANG_USE_AITER=1, quantize the BF16 Qwen3.5/3.8 GDN
+    # in_proj_qkvz and out_proj to MXFP4 after loading. 0 keeps them BF16.
+    SGLANG_ROCM_GDN_ONLINE_MXFP4 = EnvBool(True)
     # ROCm decode attention kernel: auto (aiter_sparse on gfx950, tilelang elsewhere) |
     # aiter_sparse | tilelang | triton | torch | comparison | unified_kv_triton
     SGLANG_HACK_FLASHMLA_BACKEND = EnvStr("auto")

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py

@@ -14,6 +14,7 @@
     PerTensorScaleParameter,
 )
 from sglang.srt.layers.quantization import QuantizationConfig
+from sglang.srt.layers.quantization.base_config import LinearMethodBase
 from sglang.srt.layers.quantization.dequantization import (
     copy_missing_attrs,
     dequantize_fp8,
@@ -367,7 +368,7 @@ def fused_gemm_afp4wfp4_split_cat(x, w, y, x_scale, w_scale, S1, S2, dtype):
         )
 
 
-__all__ = ["QuarkW4A4MXFP4"]
+__all__ = ["QuarkW4A4MXFP4", "AiterMxfp4OnlineLinearMethod"]
 logger = logging.getLogger(__name__)
 
 OCP_MX_BLOCK_SIZE = 32
@@ -968,3 +969,86 @@ def apply_weights(
             return y.view(*output_shape)
         else:
             return y
+
+
+class AiterMxfp4OnlineLinearMethod(LinearMethodBase):
+    """MXFP4 GEMM for a loaded BF16 linear that ``quantize_`` converts once.
+
+    A model sets this method on a BF16 linear after ``create_weights``. The
+    loader then calls ``process_weights_after_loading``, which runs ``quantize_``.
+    ``quantize_`` stores the weight and its E8M0 scales preshuffled. Batches of
+    at most ``PRESHUFFLE_MAX_M`` rows run the Triton preshuffled GEMM if AITER
+    has a tuned config for (N, K). All other batches run the ASM ``gemm_a4w4``.
+    """
+
+    PRESHUFFLE_MAX_M = 64
+
+    @classmethod
+    def quantize_(cls, layer: torch.nn.Module) -> None:
+        from aiter.ops.shuffle import shuffle_scale, shuffle_weight
+        from aiter.ops.triton._triton_kernels.gemm.basic.gemm_afp4wfp4 import (
+            _get_config,
+        )
+
+        w = layer.weight.data
+        n, k = w.shape
+        # gemm_a4w4 and the preshuffled scales read whole 256x256 tiles.
+        assert w.dtype == torch.bfloat16 and n % 256 == 0 and k % 256 == 0, (
+            f"online MXFP4 needs a BF16 weight with N, K % 256 == 0, got {w.dtype} {n}x{k}"
+        )
+        w_q, w_s = _dynamic_mxfp4_quant_orig(w)
+        layer.weight = torch.nn.Parameter(
+            shuffle_weight(w_q.view(torch.uint8), layout=(16, 16)),
+            requires_grad=False,
+        )
+        layer.weight_scale = torch.nn.Parameter(
+            shuffle_scale(w_s.view(torch.uint8)), requires_grad=False
+        )
+        layer.mxfp4_preshuffle_tuned = _get_config(
+            cls.PRESHUFFLE_MAX_M, n, k // 2, True
+        )[1]
+        layer.quant_method = cls()
+
+    def create_weights(self, *args, **kwargs):
+        raise NotImplementedError("use quantize_() on a loaded BF16 linear")
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if layer.weight.dtype == torch.bfloat16:
+            self.quantize_(layer)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        from aiter import gemm_a4w4, per_1x32_f4_quant_hip
+        from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4_preshuffle
+
+        m, n = x.shape[0], layer.weight.shape[0]
+        w_s = layer.weight_scale
+        if m <= self.PRESHUFFLE_MAX_M and layer.mxfp4_preshuffle_tuned:
+            # Below 32 rows the kernel takes un-shuffled activation scales.
+            x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=m >= 32)
+            x_s = x_s.view(torch.uint8)
+            x_s = x_s.view(x_s.shape[0] // 32, -1) if m >= 32 else x_s[:m]
+            y = torch.empty(m, n, dtype=x.dtype, device=x.device)
+            gemm_afp4wfp4_preshuffle(
+                x_q.view(torch.uint8),
+                layer.weight.view(n // 16, -1),
+                x_s,
+                w_s.view(w_s.shape[0] // 32, -1),
+                x.dtype,
+                y,
+            )
+        else:
+            x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
+            y = gemm_a4w4(
+                x_q,
+                layer.weight.view(x_q.dtype),
+                x_s,
+                w_s.view(x_s.dtype),
+                dtype=x.dtype,
+                bpreshuffle=True,
+            )[:m]
+        return y if bias is None else y + bias

python/sglang/srt/models/qwen3_5.py

@@ -184,6 +184,11 @@
 _gdn_decode_fused_proj_conv = (
     _is_cuda and envs.SGLANG_ENABLE_GDN_DECODE_FUSED_PROJ_CONV.get()
 )
+_gdn_online_mxfp4 = (
+    ("in_proj_qkvz", "out_proj")
+    if envs.SGLANG_ROCM_GDN_ONLINE_MXFP4.get() and _use_aiter and _is_gfx95
+    else ()
+)
 _is_amx_available = cpu_has_amx_support()
 _is_xpu = is_xpu()
 
@@ -528,6 +533,22 @@ def __init__(
             prefix=add_prefix("out_proj", prefix),
         )
 
+        if _gdn_online_mxfp4 and config.model_type in _QWEN3_5_ROCM_PACKED_MODEL_TYPES:
+            from sglang.srt.layers.quantization.quark.schemes.quark_w4a4_mxfp4 import (
+                AiterMxfp4OnlineLinearMethod,
+            )
+
+            # The loader quantizes these in process_weights_after_loading.
+            # gemm_a4w4 needs N, K % 256 == 0; other shapes stay BF16.
+            for name in _gdn_online_mxfp4:
+                proj = getattr(self, name)
+                if (
+                    isinstance(proj.quant_method, UnquantizedLinearMethod)
+                    and proj.weight.dtype == torch.bfloat16
+                    and all(d % 256 == 0 for d in proj.weight.shape)
+                ):
+                    proj.quant_method = AiterMxfp4OnlineLinearMethod()
+
     @staticmethod
     def _override_weight_loader(param, loader):
         """Robustly override loader for:
@@ -1878,13 +1899,24 @@ def get_input_embeddings(self):
     def prepare_before_cuda_graph_capture(self, model_runner) -> None:
         if _use_aiter and self.config.model_type in _QWEN3_5_ROCM_PACKED_MODEL_TYPES:
             packed = 0
+            mxfp4 = 0
             for module in self.modules():
                 if isinstance(module, Qwen3_5GatedDeltaNet):
+                    mxfp4 += sum(
+                        hasattr(getattr(module, name, None), "mxfp4_preshuffle_tuned")
+                        for name in _gdn_online_mxfp4
+                    )
                     module.finalize_fused_in_proj()
                     packed += int(module._fused_in_proj_weight is not None)
             logger.info(
                 "Packed BF16/FP8 GDN input projection enabled for %d layers", packed
             )
+            if _gdn_online_mxfp4:
+                logger.info(
+                    "Online MXFP4 GDN %s: %d projections",
+                    ",".join(_gdn_online_mxfp4),
+                    mxfp4,
+                )
 
     def set_dflash_layers_to_capture(self, layers_to_capture: list[int]):
         self.layers_to_capture = layers_to_capture

test/registered/amd/test_qwen35_gdn_online_mxfp4.py

@@ -0,0 +1,62 @@
+"""ROCm coverage for the online MXFP4 Qwen3.5/3.8 GDN projections."""
+
+import itertools
+import unittest
+
+import torch
+
+from sglang.srt.utils import is_gfx95_supported
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+# Qwen3.8-2.4T TP8: in_proj_qkvz and out_proj.
+SHAPES = ((4608, 8192), (8192, 2048))
+# Rows on both sides of the 32-row scale-layout split and of the 64-row kernel cut.
+ROWS = (1, 4, 16, 31, 32, 64, 65, 256)
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available() and torch.version.hip and is_gfx95_supported(),
+    "requires AMD gfx95",
+)
+class TestQwen35GdnOnlineMxfp4(CustomTestCase):
+    def test_matches_dequantized_reference(self):
+        from aiter import per_1x32_f4_quant_hip
+
+        from sglang.srt.layers.quantization.quark.schemes.quark_w4a4_mxfp4 import (
+            AiterMxfp4OnlineLinearMethod,
+            _dequant_mxfp4_to_bf16,
+            _dynamic_mxfp4_quant_orig,
+        )
+
+        for (n, k), m in itertools.product(SHAPES, ROWS):
+            with self.subTest(n=n, k=k, m=m):
+                torch.manual_seed(m)
+                w = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.02
+                x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
+                w_q, w_s = _dynamic_mxfp4_quant_orig(w)
+                x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False)
+                ref = torch.nn.functional.linear(
+                    _dequant_mxfp4_to_bf16(
+                        x_q.view(torch.uint8), x_s.view(torch.uint8)[:m]
+                    ).float(),
+                    _dequant_mxfp4_to_bf16(w_q, w_s.contiguous()).float(),
+                )
+
+                layer = torch.nn.Module()
+                layer.weight = torch.nn.Parameter(w, requires_grad=False)
+                layer.quant_method = AiterMxfp4OnlineLinearMethod()
+                # The second call must not quantize the MXFP4 weight again.
+                layer.quant_method.process_weights_after_loading(layer)
+                layer.quant_method.process_weights_after_loading(layer)
+                y = layer.quant_method.apply(layer, x)
+
+                self.assertEqual(y.shape, (m, n))
+                rel = ((y.float() - ref).norm() / ref.norm()).item()
+                self.assertLess(rel, 1e-2)
+
+
+if __name__ == "__main__":
+    unittest.main()