nvfp4-megamoe-kernel/CURRENT_BUG.md

Current Bug: vLLM produces empty/garbage output

Status: Weights confirmed good — bug is in vLLM's quant pipeline for attention Date: 2026-05-18

Symptom

• vLLM server starts, loads model, processes requests (200 OK)
• Chat completions return content: "" with finish_reason: "length"
• 20 completion tokens generated but all produce empty/NaN logits
• With enforce-eager + diagnostics: NaN from layer 0 onward on real requests

✅ Confirmed: Weights produce valid output

Standalone test (test_attn_moe_chain.py) running directly on B200:

Step	Operation	amax	NaN?
1	Embed tokens	1.27	No
2	hc_mult expansion	1.27	No
3	RMSNorm	0.20	No
4	q_a_proj (NVFP4→BF16 dequant + matmul)	0.50	No
5	kv_proj (NVFP4→BF16 dequant + matmul)	1.30	No
6	q_norm + kv_norm	0.11 / 1.87	No
7	q_b_proj (NVFP4→BF16 dequant + matmul)	1.10	No
8	MoE CuTeDSL runner (with warmup gs)	cosine 0.988	No

Every step produces valid, non-NaN, non-zero output. The problem is NOT the weights.

❌ Root Cause: vLLM's process_weights_after_loading breaks attention

The timeline

1. load_weights() → our _convert_nvfp4_post_load() runs
2. process_weights_after_loading() → vLLM's quant method runs AFTER, overwriting our fixes
3. FlashInferCutlassNvFp4LinearKernel gets set up with broken input_global_scale_inv

What the quant method does

CompressedTensorsW4A4Fp4.process_weights_after_loading():

input_global_scale_inv = layer.input_scale.max()  # = 0.00025141 (WRONG)
layer.input_global_scale = 1.0 / input_global_scale_inv  # = 3977.6
layer.input_global_scale_inv = input_global_scale_inv   # = 0.00025141
layer.alpha = input_global_scale * weight_global_scale

At runtime: scaled_fp4_quant(x, input_global_scale_inv=0.00025141) divides by 0.00025141 → multiplies by 3977.6 → massive overflow → NaN.

Why our fixes didn't work

Attempt	Why it failed
BF16 dequant + UnquantizedLinearMethod	process_weights_after_loading overwrites quant_method back to FlashInferCutlassNvFp4LinearKernel
Fix input_scale before quant method	Runs too early — quant method reads input_scale and overwrites our value
Fix input_global_scale_inv directly	Attribute doesn't exist yet when our code runs — it's set BY the quant method

The key insight

Our code runs inside load_weights(). The quant method's process_weights_after_loading() runs after load_weights() returns. Any changes we make get overwritten.

Config values (corrected)

Parameter	Value
head_dim	512 (NOT 56)
num_attention_heads	128
num_key_value_heads	1
q_lora_rank	1536
qk_rope_head_dim	64
o_lora_rank	1024
hc_mult	4
n_routed_experts	384 (48 per EP rank)

Next step: Post-init hook

The fix must run AFTER process_weights_after_loading and BEFORE the first inference. Options:

Option A: Override input_global_scale_inv post-init

• Add a _fix_nvfp4_activation_scales() method
• Call it from the right hook point (after quant method setup, before inference)
• Compute correct input_global_scale_inv from BF16 warmup
• Override the Parameter on each attention module

Option B: Replace quant_method with UnquantizedLinearMethod post-init

• After process_weights_after_loading, dequant weights to BF16
• Swap quant_method on attention modules to UnquantizedLinearMethod
• This time the quant method won't overwrite us (it already ran)

Option C: Override the quant config to skip attention modules

• Tell CompressedTensorsW4A4Fp4 to skip attention projections
• Then dequantize to BF16 ourselves
• Cleanest but requires modifying the quant config

Option B is most straightforward. The quant method already ran and set up its attributes. We can then come in and replace everything with BF16.

Architecture notes

• Attention uses MLA (Multi-head Latent Attention) with 2-step Q projection (q_a → q_b)
• fused_wqa_wkv = MergedColumnParallelLinear(q_a + kv fused)
• wo_a = FP8 via fp8_einsum (no input_scale, weight-only)
• wo_b = standard ColumnParallelLinear
• hc_pre / hc_post = Head-Conditioned mixing (tilelang custom ops)
• Dummy run zeros attention output by design (out.zero_(); return)
• FlashMLA handles the actual MLA attention kernel