Repository navigation
[Perf] qwen3: Fuse SiLU*mul with the NVFP4 quantization of down_proj in Qwen2MLP - #44491
Conversation
|
@Fridge003 @ishandhanani @Qiaolin-Yu Could one of you take a look and add |
c9f0a75 to
1a613a0
Compare
|
Rebased onto the latest main so CI can run. Gentle ping @Fridge003 in case this got buried. Thanks! |
1a613a0 to
f91c752
Compare
…n2MLP Qwen3.5's dense MLP runs SiLU*mul and the down_proj NVFP4 input quantization as one FlashInfer kernel; Qwen3, whose decoder layers use Qwen2MLP, ran them as two. Move the fused call and its enabling check from Qwen2MoeMLP and qwen3_5.py to models/utils.py, give Qwen2MLP the same path, and enable it for Qwen3 layers whose MLP projections are both NVFP4 W4A4. Keep it off for NVFP4 AWQ checkpoints, whose per-input-channel scale the fused producer does not apply, and skip it in a forward where sequence parallelism is active or down_proj is wrapped by LoRA, since both read the activation tensor. Signed-off-by: Hanjun Cho <gkswns0531@gmail.com>
f91c752 to
644c16f
Compare
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
[BBuf Review Agent] One thing to fix before this can go in.
| prefix=add_prefix("mlp", prefix), | ||
| reduce_results=False, | ||
| ) | ||
| maybe_enable_silu_fp4_quant_fusion(self.mlp) |
There was a problem hiding this comment.
[BBuf Review Agent] Qwen3 NVFP4 with --fp4-gemm-backend marlin now sends an FP4 tuple to apply_fp4_marlin_linear, which requires a floating-point tensor, so the first forward crashes. Could we exclude Marlin in the shared enabling check and cover this MLP handoff?
There was a problem hiding this comment.
@BBuf
Thanks!
On SM80-SM90 auto also picks Marlin, so this would have broken the default path there. f5c2129 keeps the fusion off on the Marlin backend. Qwen3.5's dense MLP uses the same check, so it's covered too, and test_stays_off_for_the_marlin_backend fails without the change.
Checked with Qwen3-Reranker NVFP4 checkpoints:
- L4 (SM89,
autopicks Marlin): the previous head failed on the first forward. Greedy outputs and logprobs now match main exactly (0.6B and 4B). - RTX PRO 6000 with
--fp4-gemm-backend marlin: the previous head failed with'tuple' object has no attribute 'dtype'. It now serves normally. - RTX PRO 6000 with the default backend: the fusion stays on, and outputs match main exactly.
While re-checking the handoff I found one more case. When down_proj's K is not a multiple of 64, the FP4 scales are padded, and the fused kernel leaves the padding unwritten (fp4_quantize zeroes it). In a kernel test with K = 1376 and 608, a NaN byte left there turned the whole GEMM output into NaN. 47dcb3e keeps the fusion off for such K; every Qwen3 size at TP <= 8 is a multiple of 64. 8b51234 moves the kill switch into Envs and logs the enable message once.
There was a problem hiding this comment.
[BBuf Review Agent] The shared guard now excludes Marlin for both Qwen3 and Qwen3.5, and the regression test covers that check. This addresses the tuple handoff crash.
Marlin, the default FP4 GEMM backend on SM80-SM90, takes a floating-point activation, so it cannot consume the prequantized (fp4, scale) tuple and a Qwen3 NVFP4 model crashed on the first forward. Leave the fusion off when the FP4 GEMM backend is Marlin. The enabling-check tests now use a stand-in FlashInfer, so the AWQ and Marlin gates are exercised whether or not FlashInfer is installed, and a positive case shows the fusion turns on for W4A4 on a native FP4 GEMM backend. Signed-off-by: Hanjun Cho <gkswns0531@gmail.com>
The NVFP4 block scales of a K-wide input are padded to a multiple of four columns. fp4_quantize zeroes that padding, but the fused SiLU*mul+FP4 kernel never writes it, and the FP4 GEMM reads it. On an RTX PRO 6000, with K = 1376 or 608 and freed memory holding a NaN byte, every output of the down_proj GEMM became NaN (CUTLASS and cuDNN backends); K = 1216 was unaffected. Leave the fusion off unless down_proj's input width is a multiple of 64. Every Qwen3 intermediate size at TP <= 8 is. Signed-off-by: Hanjun Cho <gkswns0531@gmail.com>
Register SGLANG_DISABLE_SILU_FP4_QUANT_FUSION in Envs instead of reading os.environ in models/utils.py, log the enable message once per process rather than once per decoder layer, trim comments that restated the gates or named the caller, and state in each regression test's docstring the behavior it guards. Signed-off-by: Hanjun Cho <gkswns0531@gmail.com>
|
/tag-and-rerun-ci |
|
/rerun-failed-ci |
1 similar comment
|
/rerun-failed-ci |
…in Qwen2MLP (sgl-project#44491) Signed-off-by: Hanjun Cho <gkswns0531@gmail.com>
Motivation
Qwen3.5's dense MLP runs SiLU*mul and the NVFP4 quantization of the down_proj input as one FlashInfer kernel. Qwen3 uses
Qwen2MLP, which lacks that path, so its NVFP4 W4A4 MLPs write the activation in 16 bits (act_and_mul) and read it back (fp4_quantize).Modifications
Qwen2MoeMLP/qwen3_5.pytomodels/utils.pyand giveQwen2MLPthe same branch, enabled only byQwen3DecoderLayerwhen both projections use ModelOpt NVFP4 W4A4. OtherQwen2MLPusers are unchanged.pre_quant_scale([Bugfix] Preserve NVFP4 AWQ prescaling in MLP fusion and Marlin paths #38864 adds the same check to the Qwen3.5 copy).DualGemm.can_rundoes). Without the LoRA check, a Qwen3 NVFP4 server with a down_proj adapter fails at startup.Accuracy Tests
Qwen3-Embedding-4BNVFP4 on 60 questions each of Ko-StrategyQA / MIRACL-ko: Recall@10 0.876 / 0.767 and the distance to BF16 scores are identical to main.Qwen3-Reranker-4BNVFP4 as a causal LM (CUDA graphs on), andQwen3-Reranker-0.6BNVFP4 with a down_proj LoRA adapter: greedy outputs match main.autopicks Marlin), greedy outputs and logprobs match main (Qwen3-Reranker0.6B and 4B NVFP4).test_qwen3_5_modelopt_fp4.py: 15 passed.Speed Tests and Profiling
RTX PRO 6000 (SM120), FlashInfer 0.7.1.
Qwen3-Embedding-4BNVFP4, GPU time of oneencodecall on 128 sequences:Decode throughput of
Qwen3-Reranker-4BNVFP4 at batch 1 / 16 / 64 is unchanged.The kernels alone (CUPTI, cold L2), SiLU*mul + NVFP4 quantization of a [M, 2 x 9728] BF16 input
silu_and_mul+fp4_quantizeChecklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #41621617029
Latest PR Test (Extra): ❌ Run #41621617024
Latest PR Test (AMD ROCm 10): ✅ Run #41621617068
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: