docs: update README with full NVFP4 coverage, dequant anti-pattern, v2 status

- Added NVFP4 coverage table (what's native, what's converted, why) - Documented the dequant→requant anti-pattern that caused vLLM hangs - Updated plan: Phase 2 done, Phase 3 targets remaining conversions - Removed stale REWRITE_PLAN reference - Updated project structure (nvfp4_cutedsl.py, removed old refs)
2026-05-16 05:43:33 +00:00
parent 3445bd24c1
commit e5370140cb
1 changed files with 49 additions and 19 deletions
--- a/README.md
+++ b/README.md
@@ -1,11 +1,12 @@
 # NVFP4 MegaMoE Kernel

-NVFP4 block-scaled Mixture-of-Experts kernel for DeepSeek-V4 on NVIDIA Blackwell (SM100). Uses CuTeDSL — NVIDIA's Python-based CUTLASS DSL — for a native NVFP4 pipeline that takes full advantage of Blackwell's TMA, MMA, and epilogue overlap.
+Full NVFP4 inference pipeline for DeepSeek-V4 on NVIDIA Blackwell (SM100). The entire model — MoE experts, shared experts, and attention projections — runs in native NVFP4 with zero dequantization overhead.

 ## What This Is

-A fused MoE FFN kernel that runs the entire expert forward pass in NVFP4:
+A native NVFP4 inference stack for DeepSeek-V4:

+**MoE Experts** — CuTeDSL ScaledGroupedGemmKernel (our work):
 ```
 BF16 input → quantize to NVFP4
  L1 GEMM: NVFP4 × NVFP4 → BF16 (gate + up)
@@ -15,7 +16,15 @@ BF16 input → quantize to NVFP4
  Scatter with routing weights → BF16 output
 ```

-Both GEMMs are fully NVFP4: A and B in `float4_e2m1fn_x2`, block scales in `float8_e4m3fn`, global scales in `float32`. BF16 is used only for the SiLU activation and the final scatter — the minimum possible.
+**Attention Projections** — FlashInferCutlassNvFp4LinearKernel (vLLM built-in):
+- `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):
+- `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.

 ## How We Got Here

@@ -42,12 +51,22 @@ NVIDIA's CuTeDSL approach (Python-based CUTLASS kernels compiled via MLIR → PT

 The 0.989 cosine is entirely from activation quantization. The weights are bit-identical to the checkpoint — no BF16 round-trip, no precision loss.

+### The Dequant→Requant Anti-Pattern
+
+Early versions dequantized all NVFP4 weights to BF16, then let vLLM's `FlashInferCutlassNvFp4LinearKernel` requantize them back to NVFP4 at inference time. This:
+- Wasted 5 minutes on load doing NVFP4→BF16 conversion
+- Lost precision on the double round-trip
+- Caused vLLM to hang — the NVFP4 attention kernel expects native NVFP4 weights, not BF16 weights with an NVFP4 quant_method attached
+
+The fix: **keep everything in NVFP4**. The checkpoint stores NVFP4. The kernels consume NVFP4. No conversion needed.
+
 ### Key Lessons

 1. **A wrong reference is worse than no reference** — the 0.2 cosine against a broken BF16 dequant sent us chasing SF remap bugs for weeks
 2. **The C++ CUTLASS API is a footgun for FP4** — CuTeDSL handles tensor layouts, tiling, and SF construction correctly by construction
 3. **Test with real data early** — uniform tests pass even with broken kernels; random data reveals real bugs
 4. **Separate the GEMM from the pipeline** — our `layertest.py` runs without vLLM, Docker, or tensor parallelism. It caught the kernel bug that vLLM's integration layers masked.
+5. **Don't dequant what's already quantized** — if the kernel expects NVFP4 and the checkpoint is NVFP4, leave it alone. No BF16 round-trips.

 ## Project Structure

@@ -59,22 +78,21 @@ nvfp4-megamoe-kernel/
 │   └── kernel/moe/                   # NVIDIA's ScaledGroupedGemmKernel (untouched)
 │       ├── torch_scaled_grouped_mm.py   # The working kernel (3900 lines)
 │       ├── moe_utils.py
-│       ├── moe_persistent_scheduler.py
+│       moe_persistent_scheduler.py
 │       └── moe_sched_extension.py
-├── src/nvfp4_megamoe_kernel/         # OLD Python pipeline (being replaced)
-│   ├── nvfp4_mega_moe.py            # Old pipeline — calls broken C++ kernel
-│   └── cutlass_nvfp4_gemm/          # OLD C++ CUTLASS extension (BROKEN)
 ├── vllm/                             # vLLM integration
+│   ├── nvfp4_cutedsl.py             # CuTeDSLMoERunner — MoE kernel interface
 │   └── patches/
-│       └── deepseek_v4.py           # DeepSeek-V4 model patch
+│       ├── deepseek_v4.py           # DeepSeek-V4 model patch (NVFP4 native)
+│       └── deepseek_v4_attention.py # Attention patch (NVFP4 native)
+├── src/nvfp4_megamoe_kernel/         # OLD Python pipeline (tagged the-last-of-cutlass)
 ├── tests/
 │   ├── layertest.py                 # Layer 0 comparison: CuTeDSL vs BF16 (✅ cosine 0.989)
 │   ├── test_cutedsl.py              # Small standalone CuTeDSL test (✅ cosine 0.991)
 │   ├── test_uniform_fp4.py          # Uniform data GEMM test
 │   ├── test_b_layout.py             # B matrix column layout test
 │   └── test_quick_rand.py           # Quick random GEMM sanity check
-├── reference/                        # Reference files for study
-└── REWRITE_PLAN.md                  # Original rewrite plan
+└── reference/                        # Reference files for study
 ```

 ## The Bridge Layer (`cutedsl/bridge.py`)
@@ -106,6 +124,17 @@ python3 test_cutedsl.py
 python3 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 |
+
 ## Plan

 ### Phase 1: Kernel ✅ DONE
@@ -113,21 +142,22 @@ python3 layertest.py
 - Bridge layer handles all tensor layout conversion
 - Full MoE pipeline (L1→SiLU→L2→scatter) produces cosine 0.989 vs BF16

-### Phase 2: vLLM Integration (IN PROGRESS)
- Wire `cutedsl/moe_pipeline.py` into the vLLM DeepSeek-V4 model
- Replace `nvfp4_mega_moe_full()` call with `CuTeDSLMoERunner.run()`
- Weight loading: checkpoint uint8 → float4_e2m1fn_x2 view-cast (bit-preserving, no BF16 round-trip)
- Block scales (float8_e4m3fn) and global scales (float32) pass through directly from checkpoint
+### Phase 2: vLLM Integration ✅ DONE
+- CuTeDSLMoERunner wires CuTeDSL kernel into vLLM
+- Weight loading: checkpoint uint8 → float4_e2m1fn_x2 view-cast (bit-preserving)
+- Block scales (float8_e4m3fn) and global scales (float32) pass through directly
 - L1 dual global scale handling: normalize to max(gate_gs, up_gs), fold ratio into block scales
- Remove C++ CUTLASS extension build from Dockerfile
- Add CuTeDSL dependency to the Docker build
+- Attention projections stay native NVFP4 (FlashInferCutlassNvFp4LinearKernel)
+- CuTeDSL kernel warmup during model load (prevents RPC timeout)
+- Removed all debug prints and env var gates from vLLM serving path

 ### Phase 3: Optimization
+- Replace wo_a FP8 conversion with native NVFP4 GEMM (eliminate last dequant)
+- Fix compressor weight_loader so it stays NVFP4 native
 - Explore larger tile sizes for better occupancy
 - Profile end-to-end inference on full model

 ### Phase 4: Production
- Clean up debug artifacts
- Remove old C++ kernel code
+- Clean up old C++ kernel code (tagged `the-last-of-cutlass`)
 - Add proper error handling and logging
 - Benchmark vs BF16 baseline