nvfp4-megamoe-kernel/CURRENT_BUG.md

Current State: Building our own NVFP4 kernels

Status: WIP — shared expert CuTeDSL kernel in progress Date: 2026-05-18

What happened today

Spent the day debugging why vLLM produces empty/garbage output. The journey:

1. NaN from layer 0 — diagnostic prints showed NaN from the very first layer
2. MoE kernel is fine — standalone test: cosine 0.988, no NaN
3. Root cause: FlashInferCutlassNvFp4LinearKernel uses broken input_scale — checkpoint values cause 3977x amplification during activation quantization → NaN
4. BF16 dequant fix — dequantize NVFP4 weights to BF16, replace quant method
5. process_weights_after_loading timing — our fix runs inside load_weights(), but vLLM's quant method runs AFTER. Fix gets overwritten.
6. Post-quant hook approach — forward pre-hooks don't fire (torch.compile + model wrappers bypass them)
7. Patched utils.py — added _post_quant_fix() call at end of process_weights_after_loading. This works — 305 projections dequantized to BF16.
8. Still garbage — even with 183 attention + 122 shared expert projections in BF16, output is still empty.
9. Conclusion: vLLM's pipeline has deeper issues. The FlashInferCutlassNvFp4LinearKernel is untrustworthy on B200 (same class of C++ CUTLASS FP4 bugs we hit with MoE). BF16 dequant doesn't fix it because something else is broken in vLLM's execution path.

Decision: Build our own NVFP4 kernels for shared experts and attention. Same CuTeDSL approach that works for MoE. Stop fighting vLLM's broken kernels.

Confirmed Working

Component	Kernel	Status
MoE experts (384/layer)	CuTeDSL ScaledGroupedGemm	✅ cosine 0.988, cudagraph-safe
All NVFP4 weights	Dequant to BF16	✅ Valid output in standalone test
Full attention weight chain	BF16 matmul	✅ No NaN, no zeros

In Progress

Component	Kernel	Status
Shared experts	CuTeDSL GEMM (1 group)	🔧 Runner WIP, scale assembly needs fixing
Attention projections	CuTeDSL GEMM	📋 Next after shared experts

WIP: Shared Expert CuTeDSL Kernel

Files:

• cutedsl/shared_expert_pipeline.py — dedicated runner (needs scale assembly fix)
• tests/test_shared_expert.py — standalone test

Issue: Tried reusing MoE runner with num_experts=1 — fails because MoE runner's scatter assumes hidden_size != HC_DIM. The MoE runner does output.scatter_add_ which expects expert output shape [tokens, hidden_size] but shared expert operates on HC_DIM (28672).

Fix needed: Dedicated runner with correct scale assembly for num_groups=1. The MoE runner's _assemble_scales_cudagraph_safe is the template. For a single group:

• No expert offsets needed
• No scatter needed (all tokens go to the same expert)
• Scale assembly is just: quantize activation → pad to 128-row alignment → Blackwell swizzle
• Simpler than the MoE case

Plan

Phase 1: Shared Expert Kernel (WIP)

1. Fix shared_expert_pipeline.py — implement scale assembly for num_groups=1
2. Test with test_shared_expert.py — target cosine ≥ 0.98 vs BF16 reference
3. Add cudagraph test
4. Wire into vLLM via DeepseekV4MoE.forward()

Phase 2: Attention NVFP4 Kernel

• Each attention projection is a standard NVFP4 GEMM
• fused_wqa_wkv has dual weight_scale_2 (same as MoE gate+up)
• Handle wo_a — currently FP8, could stay FP8 or go native NVFP4
• Test each projection individually, then integrate

Phase 3: Clean Up

• Remove all BF16 dequant code
• Remove vllm/patches/utils.py patch
• Remove _post_quant_fix()
• All NVFP4 through CuTeDSL, no vLLM FlashInfer kernels

Memory Layout

Component	NVFP4 Size	BF16 Size	Notes
Shared expert (per layer)	33MB	66MB	Small, 2GB total
Attention (per layer)	~TBD	~TBD	5 projections
MoE experts (per layer)	~TBD	~TBD	48 experts, stays NVFP4