nvfp4-megamoe-kernel/CUDA_GRAPH_SYNC_INVENTORY.md

# CUDA Graph Readiness — Sync Violation Inventory

**Date:** 2026-06-06 (updated 09:15 UTC)
**Source:** Section A detector runs on B200 + manual code grep (Section B checklist) + graph capture attempts + full 61-layer replay verification
**Target:** single_shot_inference.py decode forward (1 token step, T=1)

## Summary

**CUDA graph capture WORKS on all 8 GPUs as of 2026-06-06!** Decode speed: 0.28-0.30s/token (2x faster than eager 0.55s/token).

**ROOT CAUSE of all-zeros replay bug (FIXED)**: PyTorch CUDA graphs on non-default GPUs require explicit `torch.cuda.Stream(device=device)` for capture and replay. Using `torch.cuda.set_device()` alone causes empty graphs (GPU 0) or stale data replay (GPU 1+). See `tests/unit/test_cuda_graph_stream.py` for the minimal reproduction.

The eager decode path works at 0.51-0.53s/token.

- **Method 1** (sync debug): 0 violations in forward compute. The `dec_tid_buf.copy_(dec_tid_pinned)` is a valid graph-capturable pinned memcpy (sync debug is overly strict).
- **Method 2** (L0 graph capture): **PASS** ✅ (from detector test, pre-A/B split)
- **Multi-layer A/B capture**: ✅ WORKING on all 8 GPUs (with explicit stream fix)

---

## CATEGORY 1: Explicit `.item()` syncs on hot path — ALL FIXED ✅

| File | Line | Fix | Commit |
|------|------|-----|--------|
| `dsv4/layers/mhc.py` | 422 | Removed `X_next.abs().max().item()` (122 syncs/step) | `a9ea303` |
| `single_shot_inference.py` | ~1600 | Warmup-gsa `.item()` — one-time, outside graph | OK (by design) |
| `single_shot_inference.py` | ~1642 | `argmax(logits).item()` — outside graph (sampling) | OK (by design) |

All VERBOSE-gated `.item()` calls (diagnostics) are safe at VERBOSE=0.

---

## CATEGORY 2: Per-step tensor allocations — ALL FIXED ✅

| File | Line | Fix | Commit |
|------|------|-----|--------|
| `dsv4/layers/linear.py` | 128 | Pre-allocated `_scale_a_buf` | `a9ea303` |
| `dsv4/layers/shared_expert.py` | 213 | Same fix — pre-allocated `padded_x_sf_buf` + view | `a9ea303`, `e07d798` |
| `dsv4/layers/grouped_linear.py` | 240 | Pre-allocated `_scale_a_buf` | `f13a81d` |
| `dsv4/layers/grouped_linear.py` | ~374 | Pre-allocated `_output_buf` | `0ca7bed` |
| `dsv4/layers/moe.py` | ~508 | `torch.full` → `self._l1_gsa_buf.fill_()` | `84655d0` |
| `dsv4/ops/quantize.py` | 84,88 | `torch.zeros_like` → scalar `0.0` | `f13a81d` |
| `dsv4/ops/quantize.py` | 327-329 | gsa: reshape for M=1, contiguous for M>1 | `80bb27f` |
| `dsv4/layers/mhc.py` | init_state | `out_buf` parameter for in-place write | `46a3a51` |
| `single_shot_inference.py` | ~1600 | Pre-allocated `dec_X_buf` | `46a3a51` |

---

## CATEGORY 3: Data-dependent control flow — FIXED / DEFERRED

| File | Issue | Status | Fix |
|------|-------|--------|-----|
| `single_shot_inference.py` | `dec_tid_buf[0] = python_int` | ✅ FIXED | Pinned CPU buffer + `copy_` | `0ca7bed` |
| `dsv4/layers/grouped_linear.py` | `expert_offsets[g] = python_int` | ✅ FIXED | Pre-allocated range tensor + element-wise multiply | `0ca7bed` |
| `dsv4/layers/grouped_linear.py` | `if group_offsets[0] != 0` | ✅ FIXED | Unconditional GPU-only update | `df05289` |
| `dsv4/layers/moe.py` | `torch.bincount` (data-dependent shapes) | ✅ FIXED | `scatter_add_` into pre-allocated buffer | `84655d0`, `518a1d3` |
| `single_shot_inference.py` | Compressor returns `None` | ⏳ Phase 2 | Eager-break-at-attention: compressor runs outside graph |
| `single_shot_inference.py` | KV `n_comp` Python int | ⏳ Phase 2 | Eager-break: attention runs outside graph |

---

## CATEGORY 4: Cross-GPU transfers inside graph — ADDRESSED ✅

| File | Issue | Fix |
|------|-------|-----|
| `single_shot_inference.py` | `X.to(f"cuda:{gpu}")` in layer loop | Per-GPU X buffers + cross-GPU memcpy outside graph, or capture per-GPU subgraphs |
| `single_shot_inference.py` | `positions.to(rope_cos.device)` | Per-GPU `dec_pos_per_gpu`/`dec_tid32_per_gpu` buffers | `56b816a` |
| `single_shot_inference.py` | `token_id.to(x.device)` in moe_forward | Per-GPU dec_tid32_per_gpu buffers |

---

## CATEGORY 5: torch.cuda.synchronize() on hot path — ALL CONDITIONAL ✅

| File | Line | Guard |
|------|-------|-------|
| `single_shot_inference.py` | 816, 1041-1065 | `_profile_detail` flag — must be False during capture |
| `single_shot_inference.py` | 1088 | Profile flag |

---

## CATEGORY 6: Per-step allocations inside CUDA graph capture — ALL FIXED ✅

### FIXED — GEMM output buffers

| File | Issue | Fix | Commit |
|------|-------|-----|--------|
| `dsv4/ops/gemm_runner.py:189` | `torch.zeros()` in `run_nvfp4_grouped_gemm` | Pre-allocated `out` parameter | `188ecae` |
| `dsv4/ops/gemm_runner.py:433` | `torch.zeros()` in `run_fused_swiglu_grouped_gemm` | Pre-allocated `out` parameter | `188ecae` |
| `dsv4/layers/grouped_linear.py` | No pre-allocated GEMM output buffer | Pre-allocated `_output_buf` | `b32713c`, `f57de06` |
| `dsv4/layers/moe.py` | No pre-allocated L1 output buffer | Pre-allocated `_l1_out_buf` (2*intermediate_size) | `6dc2f22` |
| `dsv4/layers/shared_expert.py` | No pre-allocated L1 output buffer | Pre-allocated `_l1_out_buf` (2*intermediate_size) | `6dc2f22` |
| `dsv4/layers/moe.py` | No pre-allocated L2 output buffer | Pre-allocated `_l2_out_buf` | `6dc2f22` |
| `dsv4/layers/shared_expert.py` | No pre-allocated L2 output buffer | Pre-allocated `_l2_out_buf` | `6dc2f22` |
| `dsv4/layers/linear.py` | No pre-allocated GEMM output buffer | Pre-allocated `_gemm_out_buf` | `6dc2f22` |

### FIXED — Blackwell 32_4_4 scale swizzle

| File | Issue | Fix | Commit |
|------|-------|-----|--------|
| `dsv4/kernels/gemm/grouped.py` | `to_blocked()` uses Python view ops (reshape, transpose, permute) — not graph-capturable | CUDA kernel `blackwell_swizzle.cu` during graph capture, Python fallback for eager | `69e15f1` |
| `dsv4/layers/moe.py` | `_assemble_scales_cudagraph_safe` uses Python view ops | Same CUDA kernel treatment + pre-allocated `_padded_x_sf_swizzled_buf_l1/l2` | `69e15f1` |
| `dsv4/layers/shared_expert.py` | `_assemble_scales_single_group` calls `pad_and_swizzle_single` | Same CUDA kernel treatment + pre-allocated `_padded_x_sf_swizzled_buf_l1/l2` | `69e15f1`, `f259d63` |

**CRITICAL BUG FIXED (2026-06-06)**: In shared_expert.py, `_padded_x_sf_swizzled_buf_l1/l2` were allocated at line 183-184 but then **overwritten with None** at line 190-191. This meant that during graph capture, `_assemble_scales_single_group` would find the swizzled buffer is None and fall through to the Python path, which FAILS during graph capture (Python view ops like reshape/transpose can't be recorded). Fixed by removing the None overwrite.

### FIXED — gsa copy_ from view

| File | Issue | Fix | Commit |
|------|-------|-----|--------|
| `dsv4/layers/shared_expert.py` | `_l1_gsa_buf.copy_(gsa_l1_gpu[:1].reshape(1))` | `self._l1_gsa_buf[0] = gsa_l1_gpu[0]` | `6dc2f22` |
| `dsv4/layers/shared_expert.py` | `_l2_gsa_buf.copy_(gsa_l2_gpu[:1].reshape(1))` | `self._l2_gsa_buf[0] = gsa_l2_gpu[0]` | `6dc2f22` |
| `dsv4/layers/moe.py` | Same pattern for L1 and L2 gsa | Same scalar assignment fix | `6dc2f22` |
| `dsv4/layers/linear.py` | `_gsa_buf.copy_(gsa[:1].reshape(1))` and `gsa.max().reshape(1)` | `self._gsa_buf[0] = gsa_gpu[0]` / `self._gsa_buf[0] = quant.gsa.max()` | `6dc2f22` |
| `dsv4/layers/grouped_linear.py` | `_gsa_buf[:1].copy_()` + `_gsa_buf[1:].copy_(expand(...))` | `self._gsa_buf[0] = gsa_gpu[0]` + `self._gsa_buf[1:] = self._gsa_buf[0]` | `6dc2f22` |

### FIXED — Router gate FP32 conversion

| File | Issue | Fix | Commit |
|------|-------|-----|--------|
| `dsv4/kernels/router/dense_router_decode.py` | `hidden_states.float() @ gate_bf16.T.float()` creates new FP32 tensors during capture | Run GEMM in BF16, convert only logits output to FP32 for sqrt(softplus) | `ffa7842` |

### FIXED — Norm weight pre-caching (2026-06-06)

| File | Issue | Fix | Commit |
|------|-------|-----|--------|
| `single_shot_inference.py` CUDAGraphDecoder | `attn_norm_w.to(dev, torch.float32)` creates new tensor during capture | Pre-cache norm weights on correct device in FP32 before capture; store on `self` to prevent GC | `32902d1`, `5a98cc6` |

### Known allocations inside graph capture that are FINE (recorded and replayed correctly)

| File | Issue | Notes |
|------|-------|-------|
| `dsv4/layers/mhc.py` | `_dynamic_params` does `X_flat.float()` → new FP32 tensor | Captured and replayed. Should be fine. |
| `dsv4/layers/mhc.py` | `sinkhorn_knopp` CUDA kernel returns new tensor | Captured and replayed. Should be fine. |
| `dsv4/layers/moe.py` | `l1_out[padded_dst]` — advanced indexing creates new tensor | Captured and replayed. Should be fine. |
| `dsv4/layers/moe.py` | `deinterleave_l1_weights` — creates new tensor (non-fused path only) | Not used with fused_swiglu=True. |
| `dsv4/ops/quantize.py` | `quantize_nvfp4_gpu_fused` returns new tensors from CUDA kernels | Captured and replayed (kernel output is recorded). Should be fine. |
| Various layers | `.contiguous()` calls on non-contiguous tensors | Allocates new tensor during capture; recorded and replayed. Fine. |

---

## CATEGORY 7: CuTeDSL from_dlpack device mismatch in graph capture — FIXED ✅

| Attempt | Fix | Result | Commit |
|---------|-----|--------|--------|
| v1 | `torch.cuda.set_device(t.device.index)` before from_dlpack | ❌ 'Capture must end on the same stream it began on' | `87b6c99` (reverted) |
| v2 | `_DLPatchTensor` wrapper forcing `dl_device` in `__dlpack__` | ❌ 'Cannot copy between CPU and CUDA tensors' | `5c94dbb` (reverted) |
| v3 | Patch `torch.cuda.current_device` lambda to return tensor's device index | ✅ WORKS | `91c3703` |

**NOTE**: The from_dlpack patch is still needed during CAPTURE (Python-side). During REPLAY, the GPU kernel arguments are replayed directly — no from_dlpack call. The patch does not interfere with explicit stream management.

---

## CATEGORY 8: Cross-GPU operations inside graph capture — FIXED ✅

| Issue | Fix |
|-------|-----|
| `positions.to(rope_cos.device)` inside forward_layer during capture | Per-GPU `dec_pos_per_gpu`/`dec_tid32_per_gpu` buffers (`56b816a`) |
| `X.to(f"cuda:{gpu}")` in layer loop | Graph uses per-layer x_in_bufs, copy_ before replay |
| `token_id.to(x.device)` in moe_forward | Per-GPU dec_tid32_per_gpu buffers |

---

## CATEGORY 9: Multi-GPU CUDA graph stream issue — FIXED ✅

**THIS WAS THE ROOT CAUSE OF THE ALL-ZEROS REPLAY BUG.**

| Issue | Fix |
|-------|-----|
| Graph capture on non-default GPUs (cuda:1-7) produces all-zero output during replay | Use explicit `torch.cuda.Stream(device=device)` per layer for capture AND replay |
| GPU 0: Empty graph with `torch.cuda.set_device()` | Same fix — explicit stream |
| No sync between graph streams and default stream (eager attention) | `torch.cuda.Event` + `record()` + `wait_event()` |

**Minimal reproduction**: `tests/unit/test_cuda_graph_stream.py`

**Implementation in CUDAGraphDecoder**:
- `self.streams[li] = torch.cuda.Stream(device=dev)` — per-layer stream
- Capture: `with torch.cuda.graph(graph_a, stream=s):`
- Replay: `with torch.cuda.stream(s): graph_a.replay()`
- Sync: Event between graph stream and default stream for eager attention

---

## CUDAGraphDecoder Architecture (Current — A/B Split with Explicit Streams)

The decoder captures the compute-heavy path as two graphs per layer, with eager attention in between:

```
Capture flow:
1. Step 0: warmup (eager) + warmup_gsa (fix gsa values)
2. For each layer li:
   a. Create per-device stream: s = torch.cuda.Stream(device=dev)
   b. Capture Graph A (on stream s): mHC pre_block(attn) + RMSNorm + quantize + q_a + q_b + kv projections
      → writes to x_normed_bufs[li], q_heads_bufs[li], kv_3d_bufs[li], ctx_a_B/C_bufs[li], X_mid_bufs[li], q_a_bufs[li]
   c. Capture Graph B (on stream s): mHC post_block(attn) + FFN + Router + MoE + SE + mHC post_block(ffn)
      → reads F_attn_bufs[li], X_mid_bufs[li]; writes x_out_bufs[li]
3. Capture hc_head + norm + lm_head on cuda:0 (on lm_stream)
```

```
Replay flow:
1. For each layer li:
   a. Copy X → x_in_bufs[li] (handles cross-GPU transfer)
   b. Replay Graph A on stream s:
      with torch.cuda.stream(s): graphs_a[li].replay()
   c. Sync: graph stream → default stream (Event + wait_event)
   d. Eager attention: forward_attention(q_heads=q_heads, kv_3d=kv_3d, ...)
   e. Copy F_attn → F_attn_bufs[li]
   f. Sync: default stream → graph stream (Event + synchronize)
   g. Replay Graph B on stream s:
      with torch.cuda.stream(s): graphs_b[li].replay()
   h. X = x_out_bufs[li]
2. Copy X → x_lm_in → replay lm_graph on lm_stream
3. Read logits_buf
```

Key commits: `6dc2f22` (initial A/B split + critical buffer fixes), `69e15f1` (swizzle kernel), `ffa7842` (router fix), `f259d63` (SE swizzle bug), `6650f06` (explicit stream fix — THE critical fix)

---

## Performance

| Mode | Decode Speed | Notes |
|------|-------------|-------|
| Eager (no --cuda-graph) | 0.51-0.53s/token | Baseline, stable |
| CUDA Graph (--cuda-graph) | 0.28-0.30s/token | ~2x faster, matching numerical output |

**Decode degeneration**: Model generates repetition loop (`psych` ↔ `istically`) in BOTH modes. This is NOT caused by CUDA graph capture — it's a model-level issue. Root cause still UNKNOWN. Components exonerated: mHC, FMHA, compression.

---

## Remaining Work

### Phase 1 (current — nearly complete)
1. ⬜ **Gate commits on capture test** — implement CI check
2. ⬜ **Optimize stream sync** — pre-create events, reduce per-step overhead
3. ⬜ **Long-run stability test** — --max-tokens 512+ with --cuda-graph
4. ⬜ **Memory leak check** — ensure no growing GPU usage over many steps
5. ⬜ **Numerical drift check** — verify logit range stays stable over 512+ steps

### Phase 2 (vLLM Integration — future)
- Paged KV cache (fixed blocks + block table)
- Device-side compressor boundary detection + fixed-shape output
- Full graph capture including FMHA
- Bucket-by-shape for variable sequence lengths