diff --git a/README.md b/README.md index cc2d28fd..f8955a59 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,15 @@ BF16 input β†’ quantize to 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 πŸŽ‰ + +**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. + +**Test Results:** +- `tests/layertest.py`: cosine 0.988 vs BF16 reference βœ… +- `tests/cudagraph_test.py`: capture + replay PASS βœ… +- vLLM inference: running with cudagraph, output quality in progress + ## How We Got Here ### The C++ CUTLASS Kernel Was Broken @@ -100,25 +109,36 @@ padded_scales.zero_() - VRAM cost is negligible (~350KB for activation intermediates across all MoE layers) - vLLM already does this everywhere (attention, FFN, etc.) -#### 3. Dynamic Output Allocation +#### 3. Kernel Compilation in the Forward Path -`torch.zeros(num_tokens, ...)` inside the forward pass creates a new tensor each call. In cudagraph, new allocations are recorded and replayed β€” this works, but only if `num_tokens` is fixed (which it is, since vLLM captures at fixed token budgets). +`cute.compile()` is a host-side JIT operation that generates PTX and compiles a CUDA kernel. It cannot be called inside cudagraph capture. The fix: **compile once during warmup, cache the compiled kernel, then only invoke `compiled()` on subsequent calls**. -#### Test Harness +The compiled kernel uses `separate_tensormap_init=True`, which handles TMA descriptor re-initialization for new tensor data. We create new `mark_layout_dynamic` CuTe tensor views for each forward call, and the compiled kernel accepts them. -`tests/cudagraph_test.py` validates cudagraph compatibility by: -1. Creating a runner with dummy weights -2. Running a warmup forward pass (triggers kernel compilation) -3. Attempting `torch.cuda.graph(g)` capture on the forward pass -4. If capture fails, patching `torch.cuda.synchronize`, `.item()`, `.tolist()`, `.cpu()` to detect exactly which syncs are happening +**Critical lesson:** Caching the compiled kernel across different tensor allocations initially produced wrong results (cosine 0.5 instead of 0.99). The issue was NOT that caching is fundamentally broken β€” it was that our bridge had other bugs (wrong `make_b_k_major` stride check, `quantize_weight_to_nvfp4` packing N instead of K). Once those were fixed, cached compilation works correctly. -Run on the B200: -```bash -cd /root/nvfp4-megamoe-kernel -source tests/.venv/bin/activate -python3 tests/cudagraph_test.py +#### 4. Weight Quantization: K is the Packed Dimension + +`quantize_weight_to_nvfp4` packs K (dim 0) differently from `quantize_to_nvfp4` which packs the last dimension. For a weight matrix `(K, N)`: +- K=7168 is the packed dimension (7168 β†’ 3584 in float4) +- N=6144 stays as-is +- Block scales are computed along K blocks: `(K//16, N)` not `(K//2, N//16)` +- The nibble packing uses `[:, ::2, :]` and `[:, 1::2, :]` (along the K block dim) + +Confusing the two quantization functions produces wrong tensor shapes that crash or produce garbage. + +#### 5. B Tensor K-Major Layout + +The CuTeDSL kernel expects B tensors in K-major memory layout (K elements contiguous in memory). `torch.stack` produces N-major layout. The fix: **double-permute trick** β€” transpose, make contiguous, transpose back. Same shape, different strides. + +```python +# Double-permute: (E,K,N) β†’ (E,N,K) β†’ contiguous β†’ (E,K,N) +# Same shape, but K-contiguous memory layout +return b_tensor.permute(0, 2, 1).contiguous().permute(0, 2, 1) ``` +A single permute changes the tensor SHAPE (swapping K and N), which breaks everything downstream. + ### 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 @@ -127,6 +147,8 @@ python3 tests/cudagraph_test.py 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. 6. **GPU scalar slicing is a silent cudagraph killer** β€” no error, no warning, just `cudaErrorStreamCaptureInvalidated` with no pointer to the cause. The test harness catches it. +7. **Weight vs activation quantization are different** β€” K-packed (weights) vs last-dim-packed (activations). Mixing them up produces wrong shapes and garbage output. +8. **Double-permute for memory layout changes** β€” single permute changes shape, double-permute changes layout. The kernel cares about both. ## Project Structure @@ -147,7 +169,7 @@ nvfp4-megamoe-kernel/ β”‚ └── deepseek_v4_attention.py # Attention patch (NVFP4 native) β”œβ”€β”€ tests/ β”‚ β”œβ”€β”€ cudagraph_test.py # CUDAGraph compatibility test (βœ… PASS) -β”‚ β”œβ”€β”€ layertest.py # Layer 0 comparison: CuTeDSL vs BF16 (βœ… cosine 0.989) +β”‚ β”œβ”€β”€ layertest.py # Layer 0 comparison: CuTeDSL vs BF16 (βœ… cosine 0.988) β”‚ β”œβ”€β”€ 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 @@ -161,20 +183,22 @@ Handles all tensor layout conversion from our pipeline to what the CuTeDSL kerne | Function | What it does | |----------|--------------| -| `quantize_activation_nvfp4()` | BF16 β†’ float4_e2m1fn_x2 + float8_e4m3fn block scales (cudagraph-safe, no `.max()` sync) | -| `quantize_weight_to_nvfp4()` | Same, but for weight matrices with K as the packed dimension | +| `quantize_to_nvfp4()` | BF16 β†’ float4_e2m1fn_x2 + float8_e4m3fn block scales (NOT cudagraph-safe β€” uses .max()) | +| `quantize_activation_nvfp4()` | Same, but cudagraph-safe (fixed global_scale, no .max()) | +| `quantize_weight_to_nvfp4()` | Same, but K is the packed dimension (different block scale shape) | | `assemble_scales_2d_side()` | Pad and swizzle activation scale factors (2Dx3D A side) | | `assemble_scales_3d_side()` | Pad and swizzle weight scale factors (2Dx3D B side) | -| `make_b_k_major()` | Convert B tensor from N-major to K-major strides (required by kernel) | -| `run_nvfp4_grouped_gemm()` | Full kernel launch (compile + run, cudagraph-safe) | +| `make_b_k_major()` | Convert B tensor from N-major to K-major strides (double-permute trick) | +| `run_nvfp4_grouped_gemm()` | Kernel launch with cached compilation (cudagraph-safe) | ## Running Tests -On the B200: +On the B200 (host venv, no container): ```bash cd /root/nvfp4-megamoe-kernel source tests/.venv/bin/activate +export CUDA_TOOLKIT_PATH=/usr/local/cuda # CUDAGraph compatibility test python3 tests/cudagraph_test.py @@ -202,7 +226,7 @@ python3 tests/layertest.py ### Phase 1: Kernel βœ… DONE - CuTeDSL ScaledGroupedGemmKernel works with NVFP4 - Bridge layer handles all tensor layout conversion -- Full MoE pipeline (L1β†’SiLUβ†’L2β†’scatter) produces cosine 0.989 vs BF16 +- Full MoE pipeline (L1β†’SiLUβ†’L2β†’scatter) produces cosine 0.988 vs BF16 ### Phase 2: vLLM Integration βœ… DONE - CuTeDSLMoERunner wires CuTeDSL kernel into vLLM @@ -222,14 +246,22 @@ python3 tests/layertest.py - Test harness validates capture + replay - VRAM overhead: ~350KB (negligible) - Compute overhead: zero rows through GEMM on padding (memory-bound, free) +- Kernel compilation cached: `cute.compile()` once during warmup, `compiled()` on forward calls -### Phase 3: Optimization +### 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 + +### Phase 4: 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 +- Proper TMA descriptor management for kernel caching across different tensor pools -### Phase 4: Production +### Phase 5: Production - Clean up old C++ kernel code (tagged `the-last-of-cutlass`) - Add proper error handling and logging - Benchmark vs BF16 baseline