From 5b0440878431ebc7b621009ed568125cfcdd361d Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Tue, 1 Sep 2026 11:04:41 +0800 Subject: [PATCH] [MoE] Add FlashInfer SM90 MXFP4 W4A8 CUTLASS MoE (#34967) Co-authored-by: luoyuan.luo --- .../autoregressive/DeepSeek/DeepSeek-V4.mdx | 2 +- .../advanced_features/server_arguments.mdx | 4 +- .../reference/support_features.mdx | 2 +- .../configs/deepseek-ai/deepseek-v4.jsx | 18 +- .../srt/layers/moe/fused_moe_triton/layer.py | 5 + .../moe/moe_runner/flashinfer_cutlass.py | 44 ++- .../sglang/srt/layers/quantization/mxfp4.py | 138 +++++-- .../mxfp4_flashinfer_cutlass_moe.py | 85 +++-- .../mxfp4_flashinfer_trtllm_moe.py | 3 + python/sglang/srt/server_args.py | 6 +- .../layers/moe/bench_mxfp4_sm90_kernels.py | 88 ++++- .../quantization/test_mxfp4_sm120_cutlass.py | 12 +- .../quantization/test_mxfp4_sm90_cutlass.py | 336 +++++++++++++++++- 13 files changed, 662 insertions(+), 81 deletions(-) diff --git a/docs/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx b/docs/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx index 90d7007c9..3ada73b40 100644 --- a/docs/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx +++ b/docs/cookbook/autoregressive/DeepSeek/DeepSeek-V4.mdx @@ -337,7 +337,7 @@ automatically. Two options are available for running DeepSeek-V4 on Hopper: -- **Original FP4 checkpoints** — apply the W4A16 MoE kernels (Marlin) as the command generator picks for Hopper cells. This path works on both H100 and H200 and is the only option for H100 (no FP8 path). It is TP-only; on H200 the Pro variant fits on a single 8-GPU node, while H100 Pro needs 2 nodes (TP=16). +- **Original FP4 checkpoints** — run the MoE experts with W4A16 kernels (Marlin or the FlashInfer SM90 CUTLASS runner) as the command generator picks for Hopper cells. With FlashInfer >= 0.6.18 you can instead select the **W4A8** path — MXFP4 weights with FP8 activations via FlashInfer's Humming kernels — by adding `--flashinfer-mxfp4-moe-precision fp8` to `--moe-runner-backend flashinfer_mxfp4`; the low-latency Hopper cells now generate this form. Both work on H100 and H200; FP4 is the only option for H100 (no FP8 path). It is TP-only; on H200 the Pro variant fits on a single 8-GPU node, while H100 Pro needs 2 nodes (TP=16). - **Converted FP8 checkpoints** (H100 and H200 only) — pre-repackaged FP8 weights at [`sgl-project/DeepSeek-V4-Flash-FP8`](https://huggingface.co/sgl-project/DeepSeek-V4-Flash-FP8) and [`sgl-project/DeepSeek-V4-Pro-FP8`](https://huggingface.co/sgl-project/DeepSeek-V4-Pro-FP8) unlock DP-attention + DeepEP and richer parallelism (e.g. Pro TP=16 across 2 nodes). On these FP8 checkpoints you can additionally enable the all-FP8 **MegaMoE** path on SM90 for higher long-context / large-decode throughput — see the **SM90 (Hopper) FP8 MegaMoE** note in Configuration Tips below. diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 496000c89..143c986ea 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -1808,9 +1808,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--flashinfer-mxfp4-moe-precision` - Choose the computation precision of flashinfer mxfp4 moe + Choose the computation precision of flashinfer mxfp4 moe. On SM90, `fp8` selects the Humming-style MXFP4-weight x FP8-activation path introduced by FlashInfer #3738 and requires FlashInfer >= 0.6.18. `default` - default, bf16 + default, bf16, fp8 `--enable-flashinfer-allreduce-fusion` diff --git a/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx b/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx index 177a14351..2dff30551 100644 --- a/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx +++ b/docs/docs/hardware-platforms/ascend-npus/reference/support_features.mdx @@ -1638,7 +1638,7 @@ click [Server Arguments](../../../advanced_features/server_arguments). `--flashinfer-mxfp4-moe-precision` `default` - `default`,
`bf16` + `default`,
`bf16`,
`fp8` Special for GPU diff --git a/docs/src/snippets/configs/deepseek-ai/deepseek-v4.jsx b/docs/src/snippets/configs/deepseek-ai/deepseek-v4.jsx index 535e2ff49..f00dd8f2a 100644 --- a/docs/src/snippets/configs/deepseek-ai/deepseek-v4.jsx +++ b/docs/src/snippets/configs/deepseek-ai/deepseek-v4.jsx @@ -1904,13 +1904,17 @@ sgl-eval run mmmu_pro \\ { match: { hw: "h200", variant: "flash-official", quant: "fp4", strategy: "low-latency", nodes: "single" }, - verified: true, + // W4A8 (MXFP4 weights x FP8 activations, FlashInfer Humming kernels); + // requires FlashInfer >= 0.6.18. Falls back: drop the precision flag + // for the W4A16 path, or use --moe-runner-backend marlin. + verificationStatus: "in-progress", env: [], flags: [ "--trust-remote-code", "--model-path {{MODEL_NAME}}", "--tp 4", - "--moe-runner-backend marlin", + "--moe-runner-backend flashinfer_mxfp4", + "--flashinfer-mxfp4-moe-precision fp8", "--speculative-algorithm DSPARK", "--host {{HOST_IP}}", "--port {{PORT}}", @@ -1918,13 +1922,15 @@ sgl-eval run mmmu_pro \\ }, { match: { hw: "h200", variant: "flash", quant: "fp4", strategy: "low-latency", nodes: "single" }, - verified: true, + // W4A8 Humming path -- see the flash-official cell above. + verificationStatus: "in-progress", env: [], flags: [ "--trust-remote-code", "--model-path {{MODEL_NAME}}", "--tp 4", - "--moe-runner-backend marlin", + "--moe-runner-backend flashinfer_mxfp4", + "--flashinfer-mxfp4-moe-precision fp8", "--speculative-algorithm EAGLE", "--speculative-num-steps 3", "--speculative-eagle-topk 1", @@ -1993,13 +1999,15 @@ sgl-eval run mmmu_pro \\ }, { match: { hw: "h200", variant: "pro", quant: "fp4", strategy: "low-latency", nodes: "single" }, - verified: true, + // W4A8 Humming path -- see the flash-official cell above. + verificationStatus: "in-progress", env: [], flags: [ "--trust-remote-code", "--model-path {{MODEL_NAME}}", "--tp 8", "--moe-runner-backend flashinfer_mxfp4", + "--flashinfer-mxfp4-moe-precision fp8", "--speculative-algorithm EAGLE", "--speculative-num-steps 3", "--speculative-eagle-topk 1", diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 56a8d4156..a1f0e7f0f 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -401,6 +401,11 @@ class FusedMoE(torch.nn.Module): self.quant_config = quant_config self.use_flashinfer_mxfp4_moe = get_moe_runner_backend().is_flashinfer_mxfp4() # TODO maybe we should remove this `if`, since `Mxfp4MoEMethod` does another round-up logic + # Keep the pre-round value: the mxfp4 SM90 CUTLASS post-load processor + # needs to know which trailing K columns are padding, because the + # loader never writes them and their scale entries keep the buffer's + # _UE8M0_ONE fill. + self.hidden_size_unpadded = hidden_size if ( self.quant_config is not None and self.quant_config.get_name() == "mxfp4" diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py index 10fff8245..ac5d1c78b 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py @@ -65,8 +65,9 @@ class FlashInferCutlassMoeQuantInfo(MoeQuantInfo): class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo): """Quantization payload for CUTLASS MXFP4 MoE. - SM90 consumes W4A16-interleaved weights and scales. SM120 consumes packed - MXFP4 weights and block-interleaved scales with MXFP8 activations. + SM90 consumes either W4A16-interleaved weights/scales or the Humming-style + W4A8 layouts. SM120 consumes packed MXFP4 weights and block-interleaved + scales with MXFP8 activations. """ # SM90 weights are interleaved; SM120 weights remain checkpoint-packed. @@ -80,6 +81,13 @@ class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo): # A non-None global scale selects the SM120 MXFP8 activation path. mxfp4_weight_global_scale: Optional[torch.Tensor] = None + # A complete non-None triplet selects the SM90 Humming W4A8 path. The + # residuals are FP32 [num_local_experts] and already include the fixed 2^6 + # compensation required by FlashInfer's epilogue. + w13_humming_residual_scale: Optional[torch.Tensor] = None + w2_humming_residual_scale: Optional[torch.Tensor] = None + humming_fc2_act_scale: Optional[torch.Tensor] = None + # Per-expert bias. GPT-OSS has both; DSv4 leaves both None. w13_bias: Optional[torch.Tensor] = None # bf16 [E, 2*N] w2_bias: Optional[torch.Tensor] = None # bf16 [E, K] @@ -340,6 +348,22 @@ def fused_experts_none_to_flashinfer_mxfp4( weight_global_scale = quant_info.mxfp4_weight_global_scale use_mxfp8_act_scaling = weight_global_scale is not None + w13_humming_residual_scale = quant_info.w13_humming_residual_scale + w2_humming_residual_scale = quant_info.w2_humming_residual_scale + humming_fc2_act_scale = quant_info.humming_fc2_act_scale + humming_scales = ( + w13_humming_residual_scale, + w2_humming_residual_scale, + humming_fc2_act_scale, + ) + use_wfp4afp8_humming = any(scale is not None for scale in humming_scales) + if use_wfp4afp8_humming and not all(scale is not None for scale in humming_scales): + raise ValueError( + "SM90 Humming MXFP4 MoE requires both expert residual scales " + "and the FC2 activation scale." + ) + if use_wfp4afp8_humming and use_mxfp8_act_scaling: + raise ValueError("SM90 Humming and SM120 MXFP8 scaling are mutually exclusive.") input_sf = None fc1_expert_weights = quant_info.w13_weight fc2_expert_weights = quant_info.w2_weight @@ -359,6 +383,17 @@ def fused_experts_none_to_flashinfer_mxfp4( quant_info.w2_weight_scale.view(torch.int32), weight_global_scale, ] + elif use_wfp4afp8_humming: + assert w13_humming_residual_scale is not None + assert w2_humming_residual_scale is not None + assert humming_fc2_act_scale is not None + quant_scales = [ + quant_info.w13_weight_scale.view(torch.int32), + w13_humming_residual_scale, + humming_fc2_act_scale, + quant_info.w2_weight_scale.view(torch.int32), + w2_humming_residual_scale, + ] else: quant_scales = [ quant_info.w13_weight_scale.view(torch.int32), @@ -367,6 +402,10 @@ def fused_experts_none_to_flashinfer_mxfp4( out_hidden = padded_hidden if do_pad else origin_hidden output_dtype = torch.bfloat16 + # FlashInfer 0.6.17 intentionally reverted the Humming API. Do not pass the + # new keyword at all on the existing W4A16/MXFP8 paths, so those paths keep + # working with SGLang's currently pinned release. + humming_kwargs = {"use_wfp4afp8_humming": True} if use_wfp4afp8_humming else {} with use_symmetric_memory(get_tp_group(), disabled=not is_allocation_symmetric()): out = torch.empty(x.shape[0], out_hidden, dtype=output_dtype, device=x.device) @@ -398,6 +437,7 @@ def fused_experts_none_to_flashinfer_mxfp4( tune_max_num_tokens=next_power_of_2(x.shape[0]), output=out, use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(), + **humming_kwargs, ) if do_pad: diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 27dd4d7b4..33a8b6194 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -405,14 +405,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.flashinfer_mxfp4_moe_precision = ( get_exec().moe.flashinfer_mxfp4_moe_precision ) + self._use_sm90_humming = False # When `flashinfer_mxfp4` is enabled, dispatch to one of three FlashInfer # entry points depending on the GPU: # - SM100 (Blackwell) -> trtllm_fp4_block_scale_moe (existing) # - SM120 (Blackwell) -> cutlass_fused_moe(MXFP8 x MXFP4) - # - SM90 (Hopper) -> cutlass_fused_moe(use_w4_group_scaling=True) - # (FlashInfer PR #3084, post-0.6.10) + # - SM90 (Hopper) -> cutlass_fused_moe(use_w4_group_scaling=True), + # W4A16 by default (PR #3084) or opt-in + # Humming W4A8 (PR #3738/#4431) self._fi_kernel: Optional[str] = None if self.use_flashinfer: + # precision=fp8 is an SM90 knob (Humming W4A8). The Blackwell + # paths already run MXFP8 activations, so the flag is inert there + # rather than an error -- one config can move across hardware. if get_platform().is_sm100: self._fi_kernel = "trtllm_sm100" elif get_platform().is_sm120: @@ -425,6 +430,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): "from FlashInfer PR #3084 (>= 0.6.11). Upgrade flashinfer-python " "or pick a different backend (e.g. marlin / triton_kernel)." ) + self._use_sm90_humming = self.flashinfer_mxfp4_moe_precision == "fp8" self._fi_kernel = "cutlass_sm90" else: raise NotImplementedError( @@ -497,6 +503,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # CUTLASS post-load processor after the load completes. self._padded_intermediate = round_up(intermediate_size_per_partition, 128) self._padded_hidden = round_up(hidden_size, 128) + # `hidden_size` here may ALREADY be FusedMoE's rounded value (GPT-OSS + # 2880 -> 3072). Remember the checkpoint's K so the post-load + # processor can exclude the never-written tail from Humming's + # per-expert scale range. + self._unpadded_hidden = getattr(layer, "hidden_size_unpadded", hidden_size) # create_weights below uses the *unpadded* sizes so the loader's # naive-copy fast path is correct. intermediate_size_per_partition_after_pad = intermediate_size_per_partition @@ -1113,7 +1124,22 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): _interleaved = getattr(layer.moe_runner_config, "gate_up_interleaved", True) - def _stack_up_gate_w13(unpadded_w13, last_pad, last_un): + # FusedMoE may have rounded hidden up (GPT-OSS 2880 -> 3072) BEFORE + # create_weights, so K_un above is that rounded value, not the + # checkpoint's K. The loader never writes the trailing columns, so they + # keep the scale buffer's _UE8M0_ONE (2^0) fill -- far above a real + # per-expert max. Humming derives its residual from each expert's + # min/max E8M0 exponent, so letting those columns through would shift + # the residual and perturb the REAL weights. Copy only the checkpoint's + # columns and let the preserve_expert_range fill cover the rest; the + # packed weights there are zero, so the scale is numerically inert. + K_real = min(getattr(self, "_unpadded_hidden", None) or K_un, K_un) + # ceil: a partial trailing group is still real and must be kept. + w13_scale_real = -(-K_real // sf_block_size) + + def _stack_up_gate_w13( + unpadded_w13, last_pad, last_un, preserve_expert_range=False, last_real=None + ): # unpadded_w13: [E, 2*N_un, last_un] # Returns: [E, 2*N_pad, last_pad] in [up_padded; gate_padded] order. if _interleaved: @@ -1126,10 +1152,22 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): out = torch.zeros( E, 2 * N_pad, last_pad, dtype=unpadded_w13.dtype, device=device ) + if preserve_expert_range: + # Humming derives one residual from each expert's min/max E8M0 + # exponents. Fill padding with an existing expert value so + # padding cannot change that range. + out.copy_(unpadded_w13[:, :1, :1]) + # When protecting the expert range, stop at the checkpoint's K so + # the stale tail is covered by the fill instead of copied through. + copy_un = ( + min(last_real, last_un) + if (preserve_expert_range and last_real is not None) + else last_un + ) # First half: up (with row + col padding zeros). - out[:, :N_un, :last_un] = up_rows + out[:, :N_un, :copy_un] = up_rows[:, :, :copy_un] # Second half: gate. - out[:, N_pad : N_pad + N_un, :last_un] = gate_rows + out[:, N_pad : N_pad + N_un, :copy_un] = gate_rows[:, :, :copy_un] return out w13_padded = _stack_up_gate_w13( @@ -1139,6 +1177,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): layer.w13_weight_scale.data, K_pad // sf_block_size, K_un // sf_block_size, + preserve_expert_range=self._use_sm90_humming, + last_real=w13_scale_real, ) # Bias: same de-interleave on dim=-1. if _interleaved: @@ -1151,9 +1191,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w13_bias_padded[:, :N_un] = w13_bias_up w13_bias_padded[:, N_pad : N_pad + N_un] = w13_bias_gate - def _pad_w2_3d(unpadded, last_pad, last_un): + def _pad_w2_3d( + unpadded, last_pad, last_un, preserve_expert_range=False, k_real=None + ): out = torch.zeros(E, K_pad, last_pad, dtype=unpadded.dtype, device=device) - out[:, :K_un, :last_un] = unpadded[:, :K_un, :] + if preserve_expert_range: + out.copy_(unpadded[:, :1, :1]) + # Same stale-tail exclusion as _stack_up_gate_w13, on w2's K rows. + k_copy = ( + min(k_real, K_un) + if (preserve_expert_range and k_real is not None) + else K_un + ) + out[:, :k_copy, :last_un] = unpadded[:, :k_copy, :] return out # ---- w2 (no halving, just pad to [E, K_pad, N_pad/2]) ---------------- @@ -1164,6 +1214,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): layer.w2_weight_scale.data, N_pad // sf_block_size, N_un // sf_block_size, + preserve_expert_range=self._use_sm90_humming, + k_real=K_real, ) w2_bias_padded = torch.zeros(E, K_pad, dtype=bias_dtype, device=device) w2_bias_padded[:, :K_un] = layer.w2_weight_bias.data @@ -1187,26 +1239,47 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # ---- FlashInfer SM90 byte / scale interleave ----------------------- # The padded buffers above are contiguous by construction (allocated # via torch.zeros + slice assignment), so we feed them straight in. - layer.w13_weight = Parameter( - interleave_moe_weights_for_sm90_mixed_gemm(w13_padded, "fp4"), - requires_grad=False, - ) - layer.w2_weight = Parameter( - interleave_moe_weights_for_sm90_mixed_gemm(w2_padded, "fp4"), - requires_grad=False, - ) - layer.w13_weight_scale = Parameter( - interleave_moe_scales_for_sm90_mixed_gemm( + if self._use_sm90_humming: + from flashinfer.fused_moe import ( + preprocess_moe_weights_for_sm90_mixed_gemm_humming, + ) + + w13_il, w13_scale_il, w13_residual = ( + preprocess_moe_weights_for_sm90_mixed_gemm_humming( + w13_padded, w13_scale_padded + ) + ) + w2_il, w2_scale_il, w2_residual = ( + preprocess_moe_weights_for_sm90_mixed_gemm_humming( + w2_padded, w2_scale_padded + ) + ) + # Humming keeps the FP4->FP8 exponent-bias compensation in the + # epilogue. FlashInfer #4431 consumes these in local expert order. + layer.w13_humming_residual_scale = Parameter( + (w13_residual * 64.0).contiguous(), requires_grad=False + ) + layer.w2_humming_residual_scale = Parameter( + (w2_residual * 64.0).contiguous(), requires_grad=False + ) + layer.humming_fc2_act_scale = Parameter( + torch.ones((), dtype=torch.float32, device=device), + requires_grad=False, + ) + else: + w13_il = interleave_moe_weights_for_sm90_mixed_gemm(w13_padded, "fp4") + w2_il = interleave_moe_weights_for_sm90_mixed_gemm(w2_padded, "fp4") + w13_scale_il = interleave_moe_scales_for_sm90_mixed_gemm( w13_scale_padded, group_size=sf_block_size - ), - requires_grad=False, - ) - layer.w2_weight_scale = Parameter( - interleave_moe_scales_for_sm90_mixed_gemm( + ) + w2_scale_il = interleave_moe_scales_for_sm90_mixed_gemm( w2_scale_padded, group_size=sf_block_size - ), - requires_grad=False, - ) + ) + + layer.w13_weight = Parameter(w13_il, requires_grad=False) + layer.w2_weight = Parameter(w2_il, requires_grad=False) + layer.w13_weight_scale = Parameter(w13_scale_il, requires_grad=False) + layer.w2_weight_scale = Parameter(w2_scale_il, requires_grad=False) layer.w13_weight_bias = Parameter(w13_bias_padded, requires_grad=False) layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False) @@ -1349,10 +1422,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): pass def _apply_sm90_cutlass(self, layer, dispatch_output): - """SM90 (Hopper) MXFP4 x BF16 MoE via FlashInfer's cutlass mixed-input - path (PR #3084). Routed through the unified ``MoeRunner`` -- this - helper only builds the quant_info; the actual kernel call lives in - :mod:`sglang.srt.layers.moe.moe_runner.flashinfer_cutlass`.""" + """SM90 MXFP4 x BF16/FP8 MoE via FlashInfer's mixed-input kernels. + + Routed through the unified ``MoeRunner``; this helper only builds the + quant_info. The actual kernel call lives in + :mod:`sglang.srt.layers.moe.moe_runner.flashinfer_cutlass`. + """ from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import ( FlashInferCutlassMxfp4MoeQuantInfo, ) @@ -1362,6 +1437,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w2_weight=layer.w2_weight, w13_weight_scale=layer.w13_weight_scale, w2_weight_scale=layer.w2_weight_scale, + w13_humming_residual_scale=getattr( + layer, "w13_humming_residual_scale", None + ), + w2_humming_residual_scale=getattr(layer, "w2_humming_residual_scale", None), + humming_fc2_act_scale=getattr(layer, "humming_fc2_act_scale", None), w13_bias=layer.w13_weight_bias, w2_bias=layer.w2_weight_bias, swiglu_alpha=layer.swiglu_alpha, diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py index 8fbdc4aed..62c9aa9e0 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py @@ -14,7 +14,7 @@ import torch from torch.nn import Module from torch.nn.parameter import Parameter -from sglang.srt.runtime_context import get_platform +from sglang.srt.runtime_context import get_exec, get_platform from sglang.srt.utils import is_flashinfer_available, log_info_on_rank0 # Suppress TRT-LLM CUTLASS trace logs without overriding user configuration. @@ -30,7 +30,7 @@ _GROUP_SIZE = 32 class Mxfp4FlashinferCutlassMoEMethod: - """FlashInfer MXFP4 MoE: W4A16 on SM90 and W4A8 on SM120.""" + """FlashInfer MXFP4 MoE: W4A16/W4A8 on SM90 and W4A8 on SM120.""" fuse_routed_scaling_factor_in_topk = True @@ -38,6 +38,11 @@ class Mxfp4FlashinferCutlassMoEMethod: if not is_flashinfer_available(): raise RuntimeError("Mxfp4FlashinferCutlassMoEMethod requires FlashInfer.") self._use_mxfp8_act_scaling = get_platform().is_sm120 + precision = get_exec().moe.flashinfer_mxfp4_moe_precision + # precision=fp8 is an SM90 knob (Humming W4A8); on SM120 the MXFP8 + # activation path already computes in FP8, so the flag is simply inert + # there rather than an error -- one config can move across hardware. + self._use_sm90_humming = not self._use_mxfp8_act_scaling and precision == "fp8" self._fp8 = fp8_method self.prefix = prefix self._swiglu_limit_tensor: torch.Tensor | None = None @@ -124,9 +129,12 @@ class Mxfp4FlashinferCutlassMoEMethod: return arch = "SM120" if self._use_mxfp8_act_scaling else "SM90" + precision = ( + "W4A8" if self._use_sm90_humming or self._use_mxfp8_act_scaling else "W4A16" + ) log_info_on_rank0( logger, - f"Preparing DSv4 MXFP4 experts for FlashInfer {arch} CUTLASS " + f"Preparing DSv4 MXFP4 experts for FlashInfer {arch} CUTLASS {precision} " f"(layer: {self.prefix})...", ) @@ -152,23 +160,51 @@ class Mxfp4FlashinferCutlassMoEMethod: for scale_u8 in (w13_scale_u8, w2_scale_u8): scale_u8.copy_(block_scale_interleave(scale_u8).reshape_as(scale_u8)) else: - from flashinfer.fused_moe import ( - interleave_moe_scales_for_sm90_mixed_gemm, - interleave_moe_weights_for_sm90_mixed_gemm, - ) + if self._use_sm90_humming: + from flashinfer.fused_moe import ( + preprocess_moe_weights_for_sm90_mixed_gemm_humming, + ) - w13_il = interleave_moe_weights_for_sm90_mixed_gemm( - layer.w13_weight.data.view(torch.uint8).contiguous(), "fp4" - ) - w2_il = interleave_moe_weights_for_sm90_mixed_gemm( - layer.w2_weight.data.view(torch.uint8).contiguous(), "fp4" - ) - w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm( - w13_scale_u8, group_size=_GROUP_SIZE - ) - w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm( - w2_scale_u8, group_size=_GROUP_SIZE - ) + w13_il, w13_s_il, w13_residual = ( + preprocess_moe_weights_for_sm90_mixed_gemm_humming( + layer.w13_weight.data.view(torch.uint8).contiguous(), + w13_scale_u8, + ) + ) + w2_il, w2_s_il, w2_residual = ( + preprocess_moe_weights_for_sm90_mixed_gemm_humming( + layer.w2_weight.data.view(torch.uint8).contiguous(), + w2_scale_u8, + ) + ) + layer.w13_humming_residual_scale = Parameter( + (w13_residual * 64.0).contiguous(), requires_grad=False + ) + layer.w2_humming_residual_scale = Parameter( + (w2_residual * 64.0).contiguous(), requires_grad=False + ) + layer.humming_fc2_act_scale = Parameter( + torch.ones((), dtype=torch.float32, device=w13_scale_u8.device), + requires_grad=False, + ) + else: + from flashinfer.fused_moe import ( + interleave_moe_scales_for_sm90_mixed_gemm, + interleave_moe_weights_for_sm90_mixed_gemm, + ) + + w13_il = interleave_moe_weights_for_sm90_mixed_gemm( + layer.w13_weight.data.view(torch.uint8).contiguous(), "fp4" + ) + w2_il = interleave_moe_weights_for_sm90_mixed_gemm( + layer.w2_weight.data.view(torch.uint8).contiguous(), "fp4" + ) + w13_s_il = interleave_moe_scales_for_sm90_mixed_gemm( + w13_scale_u8, group_size=_GROUP_SIZE + ) + w2_s_il = interleave_moe_scales_for_sm90_mixed_gemm( + w2_scale_u8, group_size=_GROUP_SIZE + ) layer.w13_weight = Parameter(w13_il, requires_grad=False) layer.w2_weight = Parameter(w2_il, requires_grad=False) layer.w13_weight_scale_inv = Parameter(w13_s_il, requires_grad=False) @@ -177,7 +213,11 @@ class Mxfp4FlashinferCutlassMoEMethod: layer._dsv4_mxfp4_backend = ( "flashinfer_cutlass_sm120" if self._use_mxfp8_act_scaling - else "flashinfer_cutlass_sm90" + else ( + "flashinfer_cutlass_sm90_fp8" + if self._use_sm90_humming + else "flashinfer_cutlass_sm90" + ) ) # SM90 creates full-size interleaved copies; release old layouts per layer. if not self._use_mxfp8_act_scaling: @@ -198,6 +238,11 @@ class Mxfp4FlashinferCutlassMoEMethod: w13_weight_scale=layer.w13_weight_scale_inv, w2_weight_scale=layer.w2_weight_scale_inv, mxfp4_weight_global_scale=self._mxfp4_weight_global_scale_tensor, + w13_humming_residual_scale=getattr( + layer, "w13_humming_residual_scale", None + ), + w2_humming_residual_scale=getattr(layer, "w2_humming_residual_scale", None), + humming_fc2_act_scale=getattr(layer, "humming_fc2_act_scale", None), w13_bias=None, w2_bias=None, swiglu_alpha=None, diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index 3d0add5e3..64b786375 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -54,6 +54,9 @@ class Mxfp4FlashinferTrtllmMoEMethod: def __init__(self, fp8_method, prefix: str): self._fp8 = fp8_method self.prefix = prefix + # precision=fp8 is an SM90 knob (Humming W4A8); this SM100 trtllm path + # already runs MXFP8 activations, so the flag is inert here rather than + # an error -- one config can move across hardware. self.flashinfer_mxfp4_moe_precision = ( get_exec().moe.flashinfer_mxfp4_moe_precision ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 49ba9ee32..fa4b80ed3 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2422,8 +2422,10 @@ class ServerArgs: NS("exec.moe"), ] = "auto" flashinfer_mxfp4_moe_precision: A[ - Literal["default", "bf16"], - "Choose the computation precision of flashinfer mxfp4 moe", + Literal["default", "bf16", "fp8"], + "Choose the computation precision of flashinfer mxfp4 moe. " + "On SM90, `fp8` selects the Humming-style MXFP4-weight x FP8-activation " + "path introduced by FlashInfer #3738 and requires FlashInfer >= 0.6.18.", NS("exec.moe"), ] = "default" deepep_mode: A[ diff --git a/test/manual/layers/moe/bench_mxfp4_sm90_kernels.py b/test/manual/layers/moe/bench_mxfp4_sm90_kernels.py index b835612f7..eece05df9 100644 --- a/test/manual/layers/moe/bench_mxfp4_sm90_kernels.py +++ b/test/manual/layers/moe/bench_mxfp4_sm90_kernels.py @@ -1,10 +1,10 @@ -"""Benchmark MXFP4 MoE kernels on H100/H200: SGLang Marlin vs FlashInfer cutlass. +"""Benchmark MXFP4 MoE kernels on H100/H200: Marlin vs FlashInfer CUTLASS. Compares per-call latency of: - * Marlin path : ``fused_marlin_moe(...)`` after Marlin weight repack - * FlashInfer : ``cutlass_fused_moe(use_w4_group_scaling=True, ...)`` - (PR #3084's SM90 mixed-input path) + * Marlin path : ``fused_marlin_moe(...)`` after Marlin weight repack + * FlashInfer W4A16 : PR #3084's SM90 mixed-input path + * FlashInfer W4A8 : PR #3738/#4431's corrected Humming path, when available Both run on the same random MXFP4 weights/scales (semantics differ slightly -- Marlin uses a scalar swiglu clamp + no bias, FlashInfer fuses per-expert @@ -15,7 +15,7 @@ Run on H100/H200: cd /sgl-workspace/sglang_dev3 && \\ PYTHONPATH=python:/sgl-workspace/flashinfer FLASHINFER_DISABLE_VERSION_CHECK=1 \\ - python python/sglang/test/bench_mxfp4_sm90_kernels.py + python test/manual/layers/moe/bench_mxfp4_sm90_kernels.py """ from __future__ import annotations @@ -25,15 +25,29 @@ from dataclasses import dataclass from typing import Callable, List, Tuple import torch -from flashinfer.autotuner import autotune # ---- FlashInfer ---- +from flashinfer import __version__ as flashinfer_version +from flashinfer.autotuner import autotune from flashinfer.fused_moe import ( cutlass_fused_moe, interleave_moe_scales_for_sm90_mixed_gemm, interleave_moe_weights_for_sm90_mixed_gemm, ) from flashinfer.fused_moe.core import ActivationType +from packaging.version import Version + +try: + from flashinfer.fused_moe import ( + preprocess_moe_weights_for_sm90_mixed_gemm_humming, + ) +except ImportError: + preprocess_moe_weights_for_sm90_mixed_gemm_humming = None + +_fi_release = Version(flashinfer_version).release +_fi_release = _fi_release + (0,) * (3 - len(_fi_release)) +if _fi_release[:3] < (0, 6, 18): + preprocess_moe_weights_for_sm90_mixed_gemm_humming = None # ---- SGLang Marlin ---- from sglang.kernels.ops.quantization.gptq_marlin_repack import gptq_marlin_repack @@ -162,14 +176,50 @@ def build_flashinfer_inputs(shape: Shape, w13, w2, w13_s, w2_s, w13_b, w2_b): } +def build_flashinfer_humming_inputs(shape: Shape, w13, w2, w13_s, w2_s, w13_b, w2_b): + if preprocess_moe_weights_for_sm90_mixed_gemm_humming is None: + raise RuntimeError("FlashInfer does not provide the corrected Humming API.") + w13_il, w13_s_il, w13_residual = preprocess_moe_weights_for_sm90_mixed_gemm_humming( + w13, w13_s + ) + w2_il, w2_s_il, w2_residual = preprocess_moe_weights_for_sm90_mixed_gemm_humming( + w2, w2_s + ) + e = shape.num_experts + return { + "w13": w13_il, + "w2": w2_il, + "quant_scales": [ + w13_s_il.view(torch.int32), + (w13_residual * 64.0).contiguous(), + torch.ones((), dtype=torch.float32, device="cuda"), + w2_s_il.view(torch.int32), + (w2_residual * 64.0).contiguous(), + ], + "w13_b": w13_b, + "w2_b": w2_b, + "swiglu_alpha": torch.full((e,), 1.702, dtype=torch.float32, device="cuda"), + "swiglu_beta": torch.full((e,), 1.0, dtype=torch.float32, device="cuda"), + "swiglu_limit": torch.full((e,), 7.0, dtype=torch.float32, device="cuda"), + } + + def make_flashinfer_runner( - shape: Shape, prep, x, topk_w, topk_i, autotuned: bool, with_bias: bool = True + shape: Shape, + prep, + x, + topk_w, + topk_i, + autotuned: bool, + with_bias: bool = True, + use_humming: bool = False, ): out = torch.empty(shape.tokens, shape.hidden, dtype=torch.bfloat16, device="cuda") fc1_b = prep["w13_b"] if with_bias else None fc2_b = prep["w2_b"] if with_bias else None def _call(): + humming_kwargs = {"use_wfp4afp8_humming": True} if use_humming else {} cutlass_fused_moe( input=x, token_selected_experts=topk_i, @@ -186,6 +236,7 @@ def make_flashinfer_runner( use_w4_group_scaling=True, activation_type=ActivationType.Swiglu, output=out, + **humming_kwargs, ) if autotuned: @@ -328,6 +379,29 @@ def run_one_shape(shape: Shape, run_marlin: bool): ) fi_med = fi_at_med # alias for downstream speedup print + if preprocess_moe_weights_for_sm90_mixed_gemm_humming is not None: + humming_prep = build_flashinfer_humming_inputs( + shape, w13, w2, w13_s, w2_s, w13_b, w2_b + ) + humming_call = make_flashinfer_runner( + shape, + humming_prep, + x, + topk_w, + topk_i, + autotuned=True, + with_bias=True, + use_humming=True, + ) + humming_med, humming_min = time_call(humming_call) + print( + f" FlashInfer Humming W4A8: median={humming_med:.3f} ms " + f"min={humming_min:.3f} ms" + ) + print(f" speedup (FI W4A16 / FI W4A8): {fi_med / humming_med:.2f}x") + else: + print(" FlashInfer Humming W4A8: SKIPPED (requires >= 0.6.18)") + # Marlin if run_marlin: try: diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py index cd23043a6..b24090f5d 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py @@ -84,6 +84,7 @@ def test_cutlass_adapter_import_does_not_require_flashinfer(monkeypatch): def test_dsv4_sm120_load_contract(monkeypatch, request): import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as adapter_module + from sglang.srt.runtime_context import get_context platform = override_platform(is_sm120=True) platform.install() @@ -95,7 +96,8 @@ def test_dsv4_sm120_load_contract(monkeypatch, request): def create_weights(self, *args, **kwargs): captured.update(kwargs) - method = adapter_module.Mxfp4FlashinferCutlassMoEMethod(_Fp8Method(), "test") + with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"): + method = adapter_module.Mxfp4FlashinferCutlassMoEMethod(_Fp8Method(), "test") method.create_weights( SimpleNamespace(), num_experts=4, @@ -126,6 +128,7 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch): from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import ( Mxfp4FlashinferCutlassMoEMethod, ) + from sglang.srt.runtime_context import get_context monkeypatch.setattr( runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext() @@ -155,9 +158,10 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch): moe_ep_rank=0, ) - method = Mxfp4FlashinferCutlassMoEMethod( - SimpleNamespace(process_weights_after_loading=lambda layer: None), "test" - ) + with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"): + method = Mxfp4FlashinferCutlassMoEMethod( + SimpleNamespace(process_weights_after_loading=lambda layer: None), "test" + ) config = MoeRunnerConfig( num_experts=num_experts, num_local_experts=num_experts, diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py index 4d1c0ae7e..8d440b2ce 100644 --- a/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py +++ b/test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py @@ -2,9 +2,9 @@ Builds a single-layer GPT-OSS-style MoE with random MXFP4 weights, drives the SGLang plumbing (``_process_weights_for_sm90_cutlass`` + ``_apply_sm90_cutlass``) -and compares against a direct FlashInfer ``cutlass_fused_moe`` call with the -same inputs. Both paths invoke the same SM90 kernel from FlashInfer PR #3084, -so outputs must be bit-exact. +and compares against direct FlashInfer ``cutlass_fused_moe`` calls. It covers +both PR #3084's W4A16 path and PR #3738/#4431's corrected Humming W4A8 path; +outputs must be bit-exact within each path. Run on H100/H200: @@ -24,6 +24,16 @@ register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large") flashinfer_fused_moe = pytest.importorskip("flashinfer.fused_moe") +HAS_CORRECTED_HUMMING_API = hasattr( + flashinfer_fused_moe, + "preprocess_moe_weights_for_sm90_mixed_gemm_humming", +) +preprocess_humming = getattr( + flashinfer_fused_moe, + "preprocess_moe_weights_for_sm90_mixed_gemm_humming", + None, +) + if not hasattr(flashinfer_fused_moe, "interleave_moe_weights_for_sm90_mixed_gemm"): pytest.skip( "FlashInfer build does not include PR #3084 SM90 mixed-input helpers", @@ -154,11 +164,12 @@ def _round_up(x, base): return ((x + base - 1) // base) * base -def _build_method(num_experts, hidden, inter): +def _build_method(num_experts, hidden, inter, *, use_humming=False): from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod) method._fi_kernel = "cutlass_sm90" + method._use_sm90_humming = use_humming method.num_experts = num_experts # The new SM90 cutlass path tracks padded sizes in dedicated attrs; # ``hidden_size`` / ``intermediate_size_per_partition`` keep the unpadded @@ -348,6 +359,11 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct( ) monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False) monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None) + monkeypatch.setattr( + fi_cutlass_mod.envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE, + "get", + lambda: False, + ) w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter) x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1 @@ -392,6 +408,7 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct( swiglu_limit=layer.swiglu_limit, use_w4_group_scaling=True, activation_type=ActivationType.Swiglu, + use_fused_finalize=False, output=out_ref_padded, ) out_ref = ( @@ -404,6 +421,258 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct( ) +@pytest.mark.skipif( + not HAS_CORRECTED_HUMMING_API, + reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18", +) +def test_process_weights_humming_matches_flashinfer_direct(): + """The SM90 fp8 option must use #3738's preprocessing and retain #4431's + per-local-expert residual contract.""" + num_experts, hidden, inter = 4, 256, 256 + w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter) + + # Make each expert's residual distinct so an accidental scalar/broadcast + # contract cannot pass this check. + expert_offsets = torch.arange(num_experts, dtype=torch.uint8, device="cuda").view( + -1, 1, 1 + ) + w13_s = w13_s + expert_offsets + w2_s = w2_s + expert_offsets + 1 + + layer = _build_mock_layer( + num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b + ) + method = _build_method(num_experts, hidden, inter, use_humming=True) + method._process_weights_for_sm90_cutlass(layer) + + # GPT-OSS loads pair-wise [gate, up]; FlashInfer consumes halved [up; gate]. + ref_w13 = torch.cat((w13[:, 1::2], w13[:, 0::2]), dim=1).contiguous() + ref_w13_s = torch.cat((w13_s[:, 1::2], w13_s[:, 0::2]), dim=1).contiguous() + expected_w13, expected_w13_s, expected_w13_residual = preprocess_humming( + ref_w13, ref_w13_s + ) + expected_w2, expected_w2_s, expected_w2_residual = preprocess_humming(w2, w2_s) + + assert torch.equal(layer.w13_weight, expected_w13) + assert torch.equal(layer.w2_weight, expected_w2) + assert torch.equal(layer.w13_weight_scale, expected_w13_s) + assert torch.equal(layer.w2_weight_scale, expected_w2_s) + assert torch.equal(layer.w13_humming_residual_scale, expected_w13_residual * 64.0) + assert torch.equal(layer.w2_humming_residual_scale, expected_w2_residual * 64.0) + assert layer.w13_humming_residual_scale.shape == (num_experts,) + assert layer.w2_humming_residual_scale.shape == (num_experts,) + assert layer.humming_fc2_act_scale.shape == () + + +@pytest.mark.skipif( + not HAS_CORRECTED_HUMMING_API, + reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18", +) +def test_humming_padding_preserves_per_expert_residual(): + """Synthetic alignment padding must not change an expert's E8M0 range.""" + num_experts, hidden, inter = 4, 192, 192 + w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter) + ref_w13 = torch.cat((w13[:, 1::2], w13[:, 0::2]), dim=1).contiguous() + ref_w13_s = torch.cat((w13_s[:, 1::2], w13_s[:, 0::2]), dim=1).contiguous() + _, _, expected_w13_residual = preprocess_humming( + ref_w13, ref_w13_s, interleave=False + ) + _, _, expected_w2_residual = preprocess_humming(w2, w2_s, interleave=False) + + layer = _build_mock_layer( + num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b + ) + method = _build_method(num_experts, hidden, inter, use_humming=True) + method._process_weights_for_sm90_cutlass(layer) + + assert torch.equal(layer.w13_humming_residual_scale, expected_w13_residual * 64.0) + assert torch.equal(layer.w2_humming_residual_scale, expected_w2_residual * 64.0) + + +def _build_prerounded_case(E, hidden_real, hidden_rounded, inter, tail_fill, seed=0): + """Weights as ``create_weights`` leaves them when FusedMoE pre-rounds hidden. + + Buffers are allocated at ``hidden_rounded``; the loader only ever writes the + first ``hidden_real`` columns, so the tail keeps whatever the buffer was + filled with (``_UE8M0_ONE`` in production). Real scales sit well BELOW 2^0 + so a leaked tail moves the per-expert max. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + kr_bytes = hidden_real // 2 + kr_grp = hidden_real // GROUP_SIZE + + w13 = torch.zeros( + (E, 2 * inter, hidden_rounded // 2), dtype=torch.uint8, device="cuda" + ) + w13[:, :, :kr_bytes] = torch.randint( + 0, 256, (E, 2 * inter, kr_bytes), dtype=torch.uint8, device="cuda", generator=g + ) + w2 = torch.zeros((E, hidden_rounded, inter // 2), dtype=torch.uint8, device="cuda") + w2[:, :hidden_real, :] = torch.randint( + 0, + 256, + (E, hidden_real, inter // 2), + dtype=torch.uint8, + device="cuda", + generator=g, + ) + + w13_s = torch.full( + (E, 2 * inter, hidden_rounded // GROUP_SIZE), + tail_fill, + dtype=torch.uint8, + device="cuda", + ) + w13_s[:, :, :kr_grp] = torch.randint( + 100, 110, (E, 2 * inter, kr_grp), dtype=torch.uint8, device="cuda", generator=g + ) + w2_s = torch.full( + (E, hidden_rounded, inter // GROUP_SIZE), + tail_fill, + dtype=torch.uint8, + device="cuda", + ) + w2_s[:, :hidden_real, :] = torch.randint( + 100, + 110, + (E, hidden_real, inter // GROUP_SIZE), + dtype=torch.uint8, + device="cuda", + generator=g, + ) + + w13_b = torch.zeros((E, 2 * inter), dtype=torch.bfloat16, device="cuda") + w2_b = torch.zeros((E, hidden_rounded), dtype=torch.bfloat16, device="cuda") + return w13, w2, w13_s, w2_s, w13_b, w2_b + + +@pytest.mark.skipif( + not HAS_CORRECTED_HUMMING_API, + reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18", +) +def test_humming_range_ignores_prerounded_hidden_tail(): + """FusedMoE rounds GPT-OSS hidden 2880 -> 3072 BEFORE ``create_weights``, so + the trailing scale columns keep the ``_UE8M0_ONE`` buffer fill. Those bytes + are 2^0 -- above any real per-expert max -- and must not reach Humming's + min/max, or the residual shifts and perturbs the real weights. + + Invariant: the residual must not depend on what the never-written tail holds. + """ + from sglang.srt.layers.quantization.mxfp4 import _UE8M0_ONE + + E, hidden_real, hidden_rounded, inter = 4, 2880, 3072, 256 + + residuals = [] + for tail_fill in (_UE8M0_ONE, 105): # 105 sits inside the real 100..110 band + w13, w2, w13_s, w2_s, w13_b, w2_b = _build_prerounded_case( + E, hidden_real, hidden_rounded, inter, tail_fill + ) + layer = _build_mock_layer( + E, hidden_rounded, inter, w13, w2, w13_s, w2_s, w13_b, w2_b + ) + method = _build_method(E, hidden_rounded, inter, use_humming=True) + # What create_weights records from layer.hidden_size_unpadded. + method._unpadded_hidden = hidden_real + method._process_weights_for_sm90_cutlass(layer) + residuals.append( + ( + layer.w13_humming_residual_scale.clone(), + layer.w2_humming_residual_scale.clone(), + ) + ) + + assert torch.equal(residuals[0][0], residuals[1][0]), ( + "w13 Humming residual changed with the never-written hidden tail; " + "the _UE8M0_ONE fill leaked into the per-expert E8M0 range" + ) + assert torch.equal( + residuals[0][1], residuals[1][1] + ), "w2 Humming residual changed with the never-written hidden tail" + + +@pytest.mark.skipif( + not HAS_CORRECTED_HUMMING_API, + reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18", +) +@pytest.mark.parametrize( + "tokens,hidden,inter,ep_size,ep_rank", + [(8, 256, 256, 1, 0), (8, 192, 192, 1, 0), (8, 256, 256, 2, 1)], +) +def test_apply_sm90_humming_matches_flashinfer_direct( + tokens, hidden, inter, ep_size, ep_rank, monkeypatch +): + """SGLang must forward the five Humming scales and enable the new kernel.""" + import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as fi_cutlass_mod + + monkeypatch.setattr( + fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext() + ) + monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False) + monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None) + monkeypatch.setattr( + fi_cutlass_mod.envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE, + "get", + lambda: False, + ) + + num_experts, top_k = 4, 2 + w13, w2, w13_s, w2_s, w13_b, w2_b = _make_random_mxfp4(num_experts, hidden, inter) + x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1 + topk_w, topk_i = _make_topk(tokens, num_experts, top_k) + topk_i = topk_i + ep_rank * num_experts + + layer = _build_mock_layer( + num_experts, hidden, inter, w13, w2, w13_s, w2_s, w13_b, w2_b + ) + layer.moe_ep_size = ep_size + layer.moe_ep_rank = ep_rank + method = _build_method(num_experts, hidden, inter, use_humming=True) + method._process_weights_for_sm90_cutlass(layer) + out_sglang = method._apply_sm90_cutlass( + layer, _MockDispatchOutput(x.clone(), topk_w, topk_i) + ).hidden_states + + padded_hidden = method._padded_hidden + x_ref = ( + torch.nn.functional.pad(x, (0, padded_hidden - hidden)) + if padded_hidden != hidden + else x + ) + out_ref_padded = torch.empty( + tokens, padded_hidden, dtype=torch.bfloat16, device="cuda" + ) + cutlass_fused_moe( + input=x_ref, + token_selected_experts=topk_i, + token_final_scales=topk_w, + fc1_expert_weights=layer.w13_weight, + fc2_expert_weights=layer.w2_weight, + output_dtype=torch.bfloat16, + quant_scales=[ + layer.w13_weight_scale.view(torch.int32), + layer.w13_humming_residual_scale, + layer.humming_fc2_act_scale, + layer.w2_weight_scale.view(torch.int32), + layer.w2_humming_residual_scale, + ], + fc1_expert_biases=layer.w13_weight_bias, + fc2_expert_biases=layer.w2_weight_bias, + swiglu_alpha=layer.swiglu_alpha, + swiglu_beta=layer.swiglu_beta, + swiglu_limit=layer.swiglu_limit, + ep_size=ep_size, + ep_rank=ep_rank, + use_w4_group_scaling=True, + use_wfp4afp8_humming=True, + activation_type=ActivationType.Swiglu, + tune_max_num_tokens=tokens, + use_fused_finalize=False, + output=out_ref_padded, + ) + out_ref = out_ref_padded[:, :hidden].contiguous() + assert torch.equal(out_sglang, out_ref) + + # ============================================================================= # DeepSeek-V4 path: Mxfp4FlashinferCutlassMoEMethod (sibling of Marlin / # trtllm-gen). Wired into fp8.py's get_quant_method when SM90 + @@ -494,10 +763,13 @@ def test_dsv4_apply_matches_flashinfer_direct( # ---- SGLang DSv4 path ---- # plain SiLU * up — all three SwiGLU scalars None (no clamp configured). - method = ds_mod.Mxfp4FlashinferCutlassMoEMethod( - SimpleNamespace(process_weights_after_loading=lambda layer: None), - "test", - ) + from sglang.srt.runtime_context import get_context + + with get_context().override_server_args(flashinfer_mxfp4_moe_precision="default"): + method = ds_mod.Mxfp4FlashinferCutlassMoEMethod( + SimpleNamespace(process_weights_after_loading=lambda layer: None), + "test", + ) # Wire the unified MoeRunner -> flashinfer_mxfp4 fused func that # ``apply`` now dispatches through. method.runner = _build_flashinfer_mxfp4_runner(num_experts, hidden, inter) @@ -559,6 +831,54 @@ def test_dsv4_apply_matches_flashinfer_direct( ) +@pytest.mark.skipif( + not HAS_CORRECTED_HUMMING_API, + reason="requires corrected per-expert Humming API from FlashInfer >= 0.6.18", +) +def test_dsv4_process_weights_humming_matches_flashinfer_direct(): + """DSv4's native [up; gate] layout must use the same #3738 transform.""" + from types import SimpleNamespace + + import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as ds_mod + from sglang.srt.runtime_context import get_context + + num_experts, hidden, inter = 4, 256, 256 + w13, w2, w13_s, w2_s = _make_random_dsv4_mxfp4(num_experts, hidden, inter) + w1, w3 = w13.chunk(2, dim=1) + w1_s, w3_s = w13_s.chunk(2, dim=1) + w31 = torch.cat((w3, w1), dim=1).contiguous() + w31_s = torch.cat((w3_s.view(torch.uint8), w1_s.view(torch.uint8)), dim=1).view( + torch.float8_e8m0fnu + ) + + with get_context().override_server_args(flashinfer_mxfp4_moe_precision="fp8"): + method = ds_mod.Mxfp4FlashinferCutlassMoEMethod( + SimpleNamespace(process_weights_after_loading=lambda layer: None), + "test", + ) + + layer = _MockLayer() + layer.w13_weight = torch.nn.Parameter(w31.clone(), requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2.clone(), requires_grad=False) + layer.w13_weight_scale_inv = torch.nn.Parameter(w31_s.clone(), requires_grad=False) + layer.w2_weight_scale_inv = torch.nn.Parameter(w2_s.clone(), requires_grad=False) + layer.num_local_experts = num_experts + method.process_weights_after_loading(layer) + + ref_w13, ref_w13_s, ref_w13_residual = preprocess_humming( + w31.view(torch.uint8), w31_s.view(torch.uint8) + ) + ref_w2, ref_w2_s, ref_w2_residual = preprocess_humming( + w2.view(torch.uint8), w2_s.view(torch.uint8) + ) + assert torch.equal(layer.w13_weight, ref_w13) + assert torch.equal(layer.w2_weight, ref_w2) + assert torch.equal(layer.w13_weight_scale_inv, ref_w13_s) + assert torch.equal(layer.w2_weight_scale_inv, ref_w2_s) + assert torch.equal(layer.w13_humming_residual_scale, ref_w13_residual * 64.0) + assert torch.equal(layer.w2_humming_residual_scale, ref_w2_residual * 64.0) + + class _MockDispatchOutput: """Stand-in for StandardDispatchOutput. ``topk_output`` is a real ``StandardTopKOutput`` so ``TopKOutputChecker.format_is_standard``