biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update CUDA_GRAPH_SYNC_INVENTORY.md with session 2 progress

Browse Source 
- Category 6: Per-step allocations (partially fixed, 6 done, ~6 blocking)
- Category 7: CuTeDSL from_dlpack fix (v3 works, v1/v2 failed)
- Category 8: Cross-GPU operations in graph capture (fixed)
- CUDAGraphDecoder architecture: single-graph-per-layer (simplified from A/B split)
- Multi-layer capture still blocked by Category 6 allocations
This commit is contained in:
biondizzle
2026-06-03 23:41:42 +00:00
parent a468f72a0e
commit 800e974d20
1 changed files with 64 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
CUDA_GRAPH_SYNC_INVENTORY.md
View File

@@ -1,15 +1,17 @@
# CUDA Graph Readiness — Sync Violation Inventory
**Date:** 2026-06-03 (updated 19:12 UTC)
**Source:** Section A detector runs on B200 + manual code grep (Section B checklist)
**Date:** 2026-06-03 (updated 23:40 UTC)
**Source:** Section A detector runs on B200 + manual code grep (Section B checklist) + graph capture attempts
**Target:** single_shot_inference.py decode forward (1 token step, T=1)
## Summary
**ALL sync violations in the compute forward path have been fixed.** Layer 0 CUDA graph capture PASSES on B200.
**All sync violations in the compute forward path have been fixed.** Layer 0 CUDA graph capture PASSES on B200.
Multi-layer capture is blocked by per-step tensor allocations inside the forward path.
- **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** ✅
- **Multi-layer capture**: BLOCKED — per-step allocations in CuTeDSL GEMM runners + MoE/SE paths
---
@@ -103,11 +105,67 @@ All VERBOSE-gated `.item()` calls (diagnostics) are safe at VERBOSE=0.
---
## CATEGORY 6: Per-step allocations inside CUDA graph capture — PARTIALLY FIXED ⏳
These are `torch.zeros()`, `torch.empty()`, and `.copy_()` calls inside the forward
path that work fine in eager mode but are disallowed during `torch.cuda.graph()` capture.
Each was discovered during attempted multi-layer capture on B200.
| File | Issue | Status | Fix Commit |
|------|-------|--------|------------|
| `dsv4/ops/gemm_runner.py:189` | `torch.zeros()` in `run_nvfp4_grouped_gemm` | ✅ FIXED | `188ecae` |
| `dsv4/ops/gemm_runner.py:433` | `torch.zeros()` in `run_fused_swiglu_grouped_gemm` | ✅ FIXED | `188ecae` |
| `dsv4/layers/grouped_linear.py` | No pre-allocated GEMM output buffer | ✅ FIXED | `b32713c`, `f57de06` |
| `dsv4/layers/moe.py` | No pre-allocated L1 output buffer | ✅ FIXED | `a468f72` |
| `dsv4/layers/shared_expert.py` | No pre-allocated L1 output buffer | ✅ FIXED | `a468f72` |
| `dsv4/ops/quantize.py:147-150` | `torch.zeros_like()` in mhc_rmsnorm_quantize | ✅ FIXED | `188ecae` |
| `dsv4/layers/shared_expert.py:367` | `_l2_gsa_buf.copy_()` during capture | ⏳ BLOCKING | Not yet fixed |
| `dsv4/layers/moe.py` | `_l2_gsa_buf.copy_()` and similar | ⏳ NOT YET CHECKED | |
| `dsv4/layers/linear.py` | Any per-step allocations in run/run_from_quantized | ⏳ NOT YET CHECKED | |
| `dsv4/ops/quantize.py` | Any remaining torch.zeros/empty in mhc_rmsnorm_quantize | ⏳ NOT YET CHECKED | |
| `dsv4/kernels/attention/` | FMHA output allocations | ⏳ NOT YET CHECKED | |
## CATEGORY 7: CuTeDSL from_dlpack device mismatch in graph capture — FIXED ✅
When capturing on non-default GPUs (cuda:1-7), `cutlass_torch.from_dlpack(t)` fails
because PyTorch's `__dlpack__` checks `torch.cuda.current_device()` against the
tensor's device, and inside graph capture on cuda:1, `current_device()` may return 0.
| 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` |
## CATEGORY 8: Cross-GPU operations inside graph capture — PARTIALLY FIXED ⏳
| Issue | Status | Fix |
|-------|--------|-----|
| `positions.to(rope_cos.device)` inside forward_layer during capture | ✅ FIXED | Per-GPU `dec_pos_per_gpu`/`dec_tid32_per_gpu` buffers (`56b816a`) |
| `X.to(f"cuda:{gpu}")` in layer loop | ✅ AVOIDED | Graph uses per-layer x_in_bufs, copy_ before replay |
| `token_id.to(x.device)` in moe_forward | ✅ AVOIDED | Per-GPU dec_tid32_per_gpu buffers |
## CUDAGraphDecoder Architecture (Current)
The decoder captures the ENTIRE `forward_layer` as a single graph per layer (not the A/B split).
L0 capture passed. Multi-layer capture blocked by Category 6 allocations.
```
Capture flow:
1. Step 0: warmup (eager) + warmup_gsa (fix gsa values)
2. Capture: for each layer li, capture forward_layer(x_in_bufs[li], ...) → x_out_bufs[li]
3. Capture: hc_head + norm + lm_head on cuda:0
4. Replay: copy X → x_in_bufs[li] → replay graph → read x_out_bufs[li] → next layer
5. Replay: copy X → x_lm_in → replay lm_graph → read logits_buf
```
Commits: `486f74d` (initial), `92225b0` (simplified from A/B split)
## Remaining Work for Full Graph Capture
1. **Extend capture to all 61 layers** — L0 passes, need L1-L60
2. **Capture hc_head + norm + lm_head** on cuda:0
3. **Cross-GPU transfers** — per-GPU X buffers, or per-GPU subgraphs
1. **Fix Category 6 allocations** — systematic audit of ALL per-step torch.zeros/empty/copy_ in forward path
2. **Extend capture to all 61 layers** — L0 passes, L1+ blocked by allocations
3. **Capture hc_head + norm + lm_head** on cuda:0 (code written, untested beyond L0)
4. **Replay verification** — bit-for-bit match with eager forward
5. **Performance benchmark** — measure speedup from graph capture
6. **Gate commits** on capture test