docs: update README with cudagraph compatibility work and decisions

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