From e5370140cb0280801a4c9e38209ed96543db94cb Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 05:43:33 +0000 Subject: [PATCH] docs: update README with full NVFP4 coverage, dequant anti-pattern, v2 status MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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) --- README.md | 68 +++++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 49 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index c8edece3..b27af4b4 100644 --- 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