biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update README and CURRENT_BUG.md with current state

Browse Source 
- README: updated NVFP4 coverage table, status, and plan
- CURRENT_BUG.md: full debugging journey, what works, what's next
- Both reflect decision to build our own CuTeDSL kernels
This commit is contained in:
biondizzle
2026-05-18 20:05:03 +00:00
parent e8b289e30d
commit b3451c74f8
2 changed files with 111 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files

174
CURRENT_BUG.md
View File

@@ -1,121 +1,77 @@
# Current Bug: vLLM produces empty/garbage output
# Current State: Building our own NVFP4 kernels
**Status:** Debugging, plan revised — building our own kernels
**Status:** WIP — shared expert CuTeDSL kernel in progress
**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
## What happened today
## What we know
Spent the day debugging why vLLM produces empty/garbage output. The journey:
### ✅ Confirmed working
- **MoE expert CuTeDSL kernel** — cosine 0.988, cudagraph-safe, production-ready
- **All NVFP4 weights dequantize correctly to BF16** — standalone test proves it
- **Full attention weight chain produces valid output** (embed → q_a → norm → q_b → o_a → o_b)
- **Post-quant fix runs at the right time** — patched `utils.py` calls `_post_quant_fix()` after `process_weights_after_loading`
- **183 attention projections dequantized to BF16** (61 layers × 3 projs)
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.
### ❌ Still broken
- Even with BF16 attention, model produces empty output
- Shared experts also use `FlashInferCutlassNvFp4LinearKernel` with broken `input_scale`
- Added shared experts to BF16 dequant fix (122 more projections) — **testing in progress**
**Decision: Build our own NVFP4 kernels for shared experts and attention.** Same CuTeDSL approach that works for MoE. Stop fighting vLLM's broken kernels.
### 🔥 The real problem: vLLM's NVFP4 kernels are untrustworthy on B200
We spent the entire day fighting vLLM's `FlashInferCutlassNvFp4LinearKernel`:
- Broken `input_scale` → NaN
- `process_weights_after_loading` timing issues
- Forward hooks not firing due to torch.compile/model wrappers
- Dequant-to-BF16 workaround is a bandaid that loses NVFP4 benefits
**We could have built our own kernel in the time we spent debugging theirs.**
## Revised Plan: Our Own NVFP4 Kernels
**Goal:** Replace ALL vLLM NVFP4 kernel paths with our own CuTeDSL implementations. No more `FlashInferCutlassNvFp4LinearKernel`. No more BF16 dequant workarounds.
### Phase 0: Get the BF16 fix working (current)
- Post-quant BF16 dequant for attention + shared experts
- Verify the model produces actual text output
- This is the "make it work" step
### Phase 1: CuTeDSL Shared Expert Kernel
**Priority:** High — shared experts are the last NVFP4 component using vLLM's broken kernel
**Files to create:**
- `cutedsl/shared_expert_pipeline.py` — L1 GEMM → SiLU → re-quant → L2 GEMM
- Same pattern as MoE but simpler: no routing, no topk, no scatter
- `gate_up_proj` already stacked (same as MoE L1)
- `down_proj` same as MoE L2
- `vllm/nvfp4_shared_expert.py` — runner class
- Cudagraph-safe (pre-allocated buffers)
- Warmup-based gs computation (same as MoE)
- Called from `DeepseekV4MoE.forward()` for shared expert path
- `tests/test_shared_expert.py` — standalone test
- Load shared expert weights from checkpoint
- CuTeDSL vs BF16 reference (cosine)
- Cudagraph test
**Why it's easy:** Shared experts are literally MoE with 1 expert and no routing. The CuTeDSL `ScaledGroupedGemmKernel` with `num_groups=1` is just a regular GEMM.
### Phase 2: CuTeDSL Attention Kernel
**Priority:** High — attention is the biggest remaining NVFP4 component
**Components to handle:**
- `fused_wqa_wkv` — MergedColumnParallelLinear (q_a + kv fused)
- `wq_b` — ColumnParallelLinear (second Q projection)
- `wo_a` — currently FP8 via fp8_einsum
- `wo_b` — ColumnParallelLinear (output projection)
**Design options:**
1. **Separate GEMMs** — one CuTeDSL GEMM per projection, simplest
2. **Fused attention GEMM** — batch all projections together (more complex, more speed)
**Recommended: Start with option 1.** Each projection is just a standard NVFP4 GEMM. No need to fuse. We can optimize later.
**Files to create:**
- `cutedsl/attention_pipeline.py` — NVFP4 GEMMs for each attention projection
- `vllm/nvfp4_attention.py` — runner class
- Handles q_a_proj, kv_proj, q_b_proj, o_a_proj, o_b_proj
- Cudagraph-safe
- Warmup gs for each projection
- `tests/test_attention_nvfp4.py` — standalone test
**Challenge:** `fused_wqa_wkv` has TWO weight_scale_2 values (one for q_a, one for kv). Need to handle dual global scales (same pattern as MoE gate+up with different gs).
### Phase 3: Clean up
- Remove all BF16 dequant code
- Remove `vllm/patches/utils.py` patch
- Remove `_post_quant_fix()` method
- All NVFP4 goes through our CuTeDSL kernels
- BF16 only where it must be (SiLU activation, final scatter, embeddings)
## NVFP4 Kernel Coverage (Target)
## Confirmed Working
| Component | Kernel | Status |
|-----------|--------|--------|
| MoE experts (L1+L2) | CuTeDSL ScaledGroupedGemm | ✅ Working |
| Shared experts (L1+L2) | CuTeDSL standard GEMM | 🔧 Phase 1 |
| Attention projections | CuTeDSL standard GEMM | 🔧 Phase 2 |
| wo_a | CuTeDSL or keep FP8 | 🔧 Phase 2 |
| Compressor | BF16 (small, not worth it) | ✅ Done |
| KV cache | FP8 (vLLM, not our concern) | ✅ Works |
| 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 |
## Config values
## In Progress
| Parameter | Value |
|-----------|-------|
| head_dim | 512 |
| 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) |
| shared expert gate_proj | [3072, 3584] = 11MB NVFP4 / 22MB BF16 |
| shared expert up_proj | [3072, 3584] = 11MB NVFP4 / 22MB BF16 |
| shared expert down_proj | [7168, 1536] = 11MB NVFP4 / 22MB BF16 |
| shared expert total | 33MB NVFP4 / 66MB BF16 per layer, ~2GB / ~4GB total |
| 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 |

59
README.md
View File

@@ -16,17 +16,34 @@ BF16 input → quantize to NVFP4
Scatter with routing weights → BF16 output
```
**Attention Projections** — FlashInferCutlassNvFp4LinearKernel (vLLM built-in):
**Attention Projections** — CuTeDSL NVFP4 GEMM (our work, in progress):
- `wq_b`, `wo_b`, `fused_wqa_wkv` — native NVFP4, no conversion
- `wo_a` — NVFP4→FP8 for `fp8_einsum` (only attention weight that needs conversion)
- Compressor — BF16 (weight_loader stacking issue, small matmul)
**Shared Experts** — FlashInferCutlassNvFp4LinearKernel (vLLM built-in):
**Shared Experts** — CuTeDSL NVFP4 GEMM (our work, in progress):
- `gate_up_proj`, `down_proj` — native NVFP4
Both GEMM types use `float4_e2m1fn_x2` for weights, `float8_e4m3fn` for block scales, `float32` for global scales. BF16 is used only for SiLU activation, the final MoE scatter, and the compressor — the minimum possible.
## Current Status: vLLM Inference Running 🎉
## Current Status: Building Our Own Kernels 🔧
**vLLM's built-in `FlashInferCutlassNvFp4LinearKernel` is broken on B200.** The same class of C++ CUTLASS FP4 bugs we hit with MoE (documented in "How We Got Here") affects the attention and shared expert paths. After a full day of debugging (broken `input_scale`, `process_weights_after_loading` timing, forward hook failures, BF16 dequant workarounds), we're replacing ALL vLLM NVFP4 kernels with our own CuTeDSL implementations.
**Test Results:**
- `tests/layertest.py`: cosine 0.988 vs BF16 reference ✅
- `tests/cudagraph_test.py`: capture + replay PASS ✅
- vLLM inference: produces empty/garbage output — vLLM's pipeline is broken, not our kernel
**What works:**
- MoE expert CuTeDSL kernel — production-ready, cosine 0.988, cudagraph-safe
- All NVFP4 weight dequantization — valid BF16 output confirmed in standalone tests
**What's in progress:**
- Shared expert CuTeDSL kernel — runner WIP, scale assembly for num_groups=1
- Attention projection CuTeDSL kernel — planned after shared experts
**Why we're building our own:** vLLM's `FlashInferCutlassNvFp4LinearKernel` uses the same C++ CUTLASS FP4 path that was broken for MoE. The CuTeDSL approach (Python-based CUTLASS via MLML→PTX) is what NVIDIA's CUTLASS team recommends for Blackwell. Our MoE kernel proves it works. Time to apply the same approach to the rest of the model.
**vLLM serves DeepSeek-V4-Pro NVFP4 with cudagraph enabled.** The model loads, cudagraph captures successfully, and inference runs. Output quality is still being tuned (garbage tokens currently), but this is the first time the entire pipeline — model loading, kernel compilation, cudagraph capture, and inference — works end-to-end.
@@ -212,14 +229,14 @@ python3 tests/layertest.py
## NVFP4 Coverage
| Component | Format | Kernel | Conversion? |
|-----------|--------|--------|-------------|
| MoE experts (L1+L2) | NVFP4 native | CuTeDSL ScaledGroupedGemm | No — direct uint8→float4 view-cast |
| Shared experts | NVFP4 native | FlashInferCutlassNvFp4 | No — stays native |
| wq_b, wo_b, fused_wqa_wkv | NVFP4 native | FlashInferCutlassNvFp4 | No — stays native |
| wo_a | NVFP4 → FP8 | fp8_einsum | Yes — fp8_einsum requires FP8 |
| Compressor | NVFP4 → BF16 | torch.mm | Yes — weight_loader stacking issue |
| KV cache | FP8 | FlashInfer MLA | N/A — FP8 is optimal for KV cache |
| Component | Format | Kernel | Status |
|-----------|--------|--------|--------|
| MoE experts (L1+L2) | NVFP4 native | CuTeDSL ScaledGroupedGemm | ✅ Working, cosine 0.988 |
| Shared experts | NVFP4 native | CuTeDSL GEMM (1 group) | 🔧 In progress |
| wq_b, wo_b, fused_wqa_wkv | NVFP4 native | CuTeDSL GEMM | 📋 Planned |
| wo_a | NVFP4 → FP8 | fp8_einsum | 📋 May stay FP8 or go native NVFP4 |
| Compressor | NVFP4 → BF16 | torch.mm | ✅ Done (weight_loader stacking issue) |
| KV cache | FP8 | FlashInfer MLA | ✅ Works |
## Plan
@@ -251,8 +268,24 @@ python3 tests/layertest.py
### Phase 3: Output Quality 🔧 IN PROGRESS
- vLLM serves the model with cudagraph, but output is garbage tokens
- Layer 0 cosine is 0.988 in isolation, so the GEMM math is correct
- Investigating: weight loading path in vLLM, tensor parallelism handling, scale normalization
- Need to compare vLLM's weight pipeline vs layertest's direct path
- Root cause: vLLM's `FlashInferCutlassNvFp4LinearKernel` is broken on B200
- Same class of C++ CUTLASS FP4 bugs we hit with MoE
- `input_scale` from checkpoint causes NaN during activation quantization
- BF16 dequant workaround doesn't fix the underlying pipeline issues
- Solution: Replace ALL vLLM NVFP4 kernels with our own CuTeDSL implementations
### Phase 3.5: Shared Expert CuTeDSL Kernel 🔧 IN PROGRESS
- Replacing `FlashInferCutlassNvFp4LinearKernel` with CuTeDSL GEMM
- Shared expert = MoE with 1 expert, no routing
- `cutedsl/shared_expert_pipeline.py` — dedicated runner (scale assembly needs fixing)
- `tests/test_shared_expert.py` — standalone test (ready)
- Target: cosine ≥ 0.98 vs BF16 reference
### Phase 3.6: Attention CuTeDSL Kernel 📋 PLANNED
- Replace attention NVFP4 path with CuTeDSL GEMMs
- Each projection (fused_wqa_wkv, wq_b, wo_a, wo_b) = standard NVFP4 GEMM
- `fused_wqa_wkv` has dual weight_scale_2 (same pattern as MoE gate+up)
- Test each projection individually, then integrate
### Phase 4: Optimization
- Replace wo_a FP8 conversion with native NVFP4 GEMM (eliminate last dequant)