diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index 3aa8591d8bad..0e2836ba35f6 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -1932,6 +1932,11 @@ SGLang supports various environment variables that can be used to configure its Token count above which SGLANG_ROCM_K3_FUSE_KDA_INPROJ stops applying and the split projection runs instead. The merged shape is only faster while the projection is bandwidth bound. 256 + + SGLANG_ROCM_GDN_ONLINE_MXFP4 + On gfx95x with SGLANG_USE_AITER=1, quantize the BF16 GDN in_proj_qkvz and out_proj 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 gemm_a4w4. Set to 0 to keep the BF16 projections. + true + diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 10eee809c188..0c2e51def227 100644 --- a/python/sglang/srt/environ.py +++ b/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") diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py index 91689ca7a700..0dc1544ec103 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py +++ b/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 diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 5e452cb6b256..dea5385e54d0 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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 diff --git a/test/registered/amd/test_qwen35_gdn_online_mxfp4.py b/test/registered/amd/test_qwen35_gdn_online_mxfp4.py new file mode 100644 index 000000000000..26264acd9678 --- /dev/null +++ b/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()