Repository navigation
[quantization] NVFP4 MoE: split fused w13 gate/up global scales - #27588
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
/tag-and-rerun-ci |
ec68e52 to
be819e5
Compare
The fused w13 GEMM in an NVFP4 MoE layer packs the gate (w1) and up (w3) projections together. A checkpoint may carry a distinct NVFP4 weight global scale per half, stored as [num_experts, 2] (col 0 = gate, col 1 = up). Previously the fused-MoE path collapsed these to the gate scale (col 0) and warned when the two columns differed, so the up half of GEMM1's output was dequantized with the wrong scale. Compute both g1_alphas (gate) and g1_alphas_up (up), and let the FlashInfer TRT-LLM kernel dequantize each half of GEMM1's output with its own scale via g1_scale_c. Shared-scale checkpoints (1-D or single-column) collapse back to the original single-scale path, and non-gated (Relu2) layers are unchanged. The Marlin path continues to use a single shared scale. Add CPU unit tests pinning the gate/up scale split (_compute_gemm1_alphas) and the derived TRT-LLM up-half output scalar (_compute_g1_scale_c).
be819e5 to
ba6bac9
Compare
|
/tag-and-rerun-ci |
|
@merrymercy can you please tell me what are the next steps for landing this ? |
…it tests Rewrite test_modelopt_nvfp4_moe_scales.py to parametrize the gate/up global-scale split and g1_scale_c contracts over representative NVFP4 ModelOpt MoE shapes (DeepSeek-R1, Llama-4, Qwen3, Mixtral, and a non-gated Relu2 config) with realistic small fp32 global-scale magnitudes, instead of toy 2-3 expert sizes. Fix the g1_scale_c tests to use the real TRT-LLM scalar w2_input_scale_quant (1 / w2_input_scale.max()) rather than a per-expert tensor, and assert the per-expert output shape plus the non-gated .expand().contiguous() contract; keep a per-expert broadcast case for robustness. Document that the surrounding wiring (Marlin single-scale collapse, the w1/w3 mismatch warning, and the layer-level g1_alphas / g1_alphas_up registration in process_weights_after_loading) is Blackwell-only and is exercised on-device, not by these CPU tests.
Add a test for the per-expert w13_input_scale shape the non-TRT-LLM backends pass (the existing tests only exercised the scalar TRT-LLM shape), drop the over-specific model-dim claim from the module docstring, and tighten the comments.
|
/rerun-test test_modelopt_nvfp4_moe_scales.py test_qwen3_fp4_trtllm_gen_moe.py::TestFlashinferTrtllmGenMoeBackend test_nvidia_nemotron_3_super_nvfp4.py::TestNvidiaNemotron3SuperNVFP4 test_deepseek_v3_fp4.py::TestDeepseekV3FP4SymmetricMemory |
|
⛔ Please ask a maintainer to run this command, or use the normal CI flow. |
|
/tag-and-rerun-ci |
|
/rerun-test test_modelopt_nvfp4_moe_scales.py test_qwen3_fp4_trtllm_gen_moe.py::TestFlashinferTrtllmGenMoeBackend test_nvidia_nemotron_3_super_nvfp4.py::TestNvidiaNemotron3SuperNVFP4 test_deepseek_v3_fp4.py::TestDeepseekV3FP4SymmetricMemory |
|
⛔ Please ask a maintainer to run this command, or use the normal CI flow. |
|
/tag-and-rerun-ci extra |
github-actions
Bot
added
the
run-ci-extra
|
/rerun-test test_modelopt_nvfp4_moe_scales.py test_qwen3_fp4_trtllm_gen_moe.py::TestFlashinferTrtllmGenMoeBackend test_nvidia_nemotron_3_super_nvfp4.py::TestNvidiaNemotron3SuperNVFP4 test_deepseek_v3_fp4.py::TestDeepseekV3FP4SymmetricMemory |
|
Results for 🚀 🚀 🚀 |
|
/rerun-test test_modelopt_nvfp4_moe_scales.py test_qwen3_fp4_trtllm_gen_moe.py::TestFlashinferTrtllmGenMoeBackend test_nvidia_nemotron_3_super_nvfp4.py::TestNvidiaNemotron3SuperNVFP4 test_deepseek_v3_fp4.py::TestDeepseekV3FP4SymmetricMemory |
|
Results for 🚀 🚀 🚀 |
merrymercy
removed
the
run-ci-extra
…project#27588) Co-authored-by: prajjwal1 <prajjwal1@protonmail.com>
Motivation
In an NVFP4 MoE layer the
w13GEMM fuses the gate (w1) and up (w3)projections. Some checkpoints quantize the two halves with separate NVFP4 weight
global scales, stored as
[num_experts, 2](column 0 for gate, column 1 for up).Today the fused-MoE path drops column 1: it keeps only the gate scale and logs a
w1_weight_scale_2 must match w3_weight_scale_2warning when the columns differ.The up half of the GEMM1 output then gets dequantized with the gate scale, which
is wrong whenever the two halves were quantized separately.
Modifications
modelopt_quant.py: add_compute_gemm1_alphas, which folds the GEMM1 inputscale into the per-half weight scales and returns both
g1_alphas(gate) andg1_alphas_up(up). Column 1 is read only when it exists, so 1-D,single-column, and non-gated checkpoints still collapse to one shared scale.
Both alphas are registered on the layer; single-alpha backends such as CUTLASS
keep reading
g1_alphas.flashinfer_trtllm.py: add_compute_g1_scale_c. The FlashInfer TRT-LLMkernel dequantizes the two GEMM1 halves separately, so
g1_scale_cnow usesthe up alpha. With a shared scale it reduces to the previous single-scale
value.
into the Marlin branch, which is now its only consumer.
g1_scale_c.Checkpoints with a shared (or near-equal) w1/w3 scale behave exactly as before.
Only checkpoints with genuinely different gate/up scales change, and they now
dequantize each half correctly.
Accuracy Tests
test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_scales.py(CPU, no GPU) pins the contract of both helpers: the gate/up split and that the
two columns stay decoupled, a simulated GEMM1 checking each half reconstructs
from its own alpha, the shared-scale fallbacks (1-D / single-column / equal
columns), and the non-gated case. All pass locally.
The TRT-LLM MoE path only runs on Blackwell, and most checkpoints ship
near-equal gate/up scales, so this regression wouldn't surface in a normal
accuracy eval. The unit tests cover it directly instead.
Speed Tests and Profiling
No effect on the forward path. The change runs once in
process_weights_after_loadingand only adds one small per-expert tensor(
g1_alphas_up) at load time.Checklist
F401,F821).Original commits
ac341a49eCI States
Latest PR Test (Base): ✅ Run #27511566795
Latest PR Test (Extra): ❌ Run #27511566758