diff --git a/README.md b/README.md index b27af4b4..cc2d28fd 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ The original kernel was a C++ `.cu` file using CUTLASS's C++ API directly. It pa ### The CuTeDSL Kernel Works -NVIDIA's CuTeDSL approach (Python-based CUTLASS kernels compiled via MLIR → PTX) is what the CUTLASS team recommends for Blackwell. Their official MoE scaled grouped GEMM example (`torch_scaled_grouped_mm.py`) supports NVFP4 out of the box. We adapted it. +NVIDIA's CuTeDSL approach (Python-based CUTLASS kernels compiled via MLML → PTX) is what the CUTLASS team recommends for Blackwell. Their official MoE scaled grouped GEMM example (`torch_scaled_grouped_mm.py`) supports NVFP4 out of the box. We adapted it. **Results with real DeepSeek-V4 layer 0 weights:** - L1 GEMM alone: cosine 0.995 @@ -60,6 +60,65 @@ Early versions dequantized all NVFP4 weights to BF16, then let vLLM's `FlashInfe The fix: **keep everything in NVFP4**. The checkpoint stores NVFP4. The kernels consume NVFP4. No conversion needed. +### CUDAGraph Compatibility + +vLLM uses CUDA graphs to eliminate kernel launch overhead in the decode path. CUDA graphs record the entire forward pass once, then replay it — but they require **fixed tensor shapes, fixed memory addresses, and zero CPU-GPU syncs**. + +Our original runner was not cudagraph-safe. We had to fix several classes of issues: + +#### 1. CPU↔CUDA Tensor Copies + +`torch.tensor([0,1,...], device=x.device)` creates the tensor on **CPU first**, then copies to CUDA. This copy is forbidden during graph capture. The fix: **cache tensors per device** on first use, outside the graph. + +```python +# BAD — CPU→CUDA copy inside graph +step_to_idx = torch.tensor([0,1,2,3,4,4,5,5,6,6,6,7,7], device=x.device) + +# GOOD — cached on first use, reused in graph +step_to_idx = _get_step_to_idx_lut(x.device) # returns cached CUDA tensor +``` + +Similarly, `torch.zeros` and `torch.rand` don't support `float4_e2m1fn_x2` or `float8_e4m3fn` dtypes. The fix: create as `uint8` or `float16`, then `.view()` or `.to()` the target dtype. + +#### 2. GPU Scalar Slicing + +`buf[:gpu_scalar, :]` requires the runtime to query the GPU scalar's value to determine the output shape. This triggers an implicit CPU-GPU sync, which invalidates the graph. The fix: **always use full pre-allocated buffers**. Extra rows are zeros that contribute nothing to the computation. + +```python +# BAD — GPU scalar as slice index (implicit sync) +total_padded_rows = padded_expert_offsets[-1] # GPU scalar +padded_scales = buf[:total_padded_rows, :padded_cols] # sync! + +# GOOD — full pre-allocated buffer, zero out before use +padded_scales = self._padded_scales_buf # always max size +padded_scales.zero_() +``` + +**Design decision:** Padding to max size wastes a few rows of compute on zero data, but: +- The extra rows are zeros → zero GEMM output → no accuracy impact +- GEMMs are memory-bandwidth bound → multiplying zeros is nearly free +- VRAM cost is negligible (~350KB for activation intermediates across all MoE layers) +- vLLM already does this everywhere (attention, FFN, etc.) + +#### 3. Dynamic Output Allocation + +`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). + +#### Test Harness + +`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 + +Run on the B200: +```bash +cd /root/nvfp4-megamoe-kernel +source tests/.venv/bin/activate +python3 tests/cudagraph_test.py +``` + ### 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 @@ -67,6 +126,7 @@ The fix: **keep everything in NVFP4**. The checkpoint stores NVFP4. The kernels 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. +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. ## Project Structure @@ -81,12 +141,12 @@ nvfp4-megamoe-kernel/ │ moe_persistent_scheduler.py │ └── moe_sched_extension.py ├── vllm/ # vLLM integration -│ ├── nvfp4_cutedsl.py # CuTeDSLMoERunner — MoE kernel interface +│ ├── nvfp4_cutedsl.py # CuTeDSLMoERunner — cudagraph-safe MoE kernel interface │ └── patches/ │ ├── 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/ +│ ├── cudagraph_test.py # CUDAGraph compatibility test (✅ PASS) │ ├── 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 @@ -101,27 +161,29 @@ Handles all tensor layout conversion from our pipeline to what the CuTeDSL kerne | Function | What it does | |----------|--------------| -| `quantize_to_nvfp4()` | BF16 → float4_e2m1fn_x2 + float8_e4m3fn block scales + float32 global scale | +| `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 | | `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) | -| `compute_expert_offsets()` | Compute cumulative token offsets for grouped GEMM | -| `run_nvfp4_grouped_gemm()` | Full kernel launch (compile + run) | +| `run_nvfp4_grouped_gemm()` | Full kernel launch (compile + run, cudagraph-safe) | ## Running Tests On the B200: ```bash -cd /root/nvfp4-megamoe-kernel/tests -source .venv/bin/activate +cd /root/nvfp4-megamoe-kernel +source tests/.venv/bin/activate + +# CUDAGraph compatibility test +python3 tests/cudagraph_test.py # Small standalone test -python3 test_cutedsl.py +python3 tests/test_cutedsl.py # Full layer 0 comparison with real weights -python3 layertest.py +python3 tests/layertest.py ``` ## NVFP4 Coverage @@ -151,6 +213,16 @@ python3 layertest.py - CuTeDSL kernel warmup during model load (prevents RPC timeout) - Removed all debug prints and env var gates from vLLM serving path +### Phase 2.5: CUDAGraph Compatibility ✅ DONE +- CuTeDSLMoERunner is fully cudagraph-safe +- Zero CPU-GPU syncs, zero dynamic shapes, zero GPU scalar slicing +- All intermediate buffers pre-allocated at max_num_tokens * top_k +- `quantize_activation_nvfp4` uses cached LUT (no CPU→CUDA copy) +- `torch.zeros/rand` for float4/float8 → uint8→view or float16→cast +- Test harness validates capture + replay +- VRAM overhead: ~350KB (negligible) +- Compute overhead: zero rows through GEMM on padding (memory-bound, free) + ### Phase 3: Optimization - Replace wo_a FP8 conversion with native NVFP4 GEMM (eliminate last dequant) - Fix compressor weight_loader so it stays NVFP4 native