Files

biondizzle 55f1ddd502 Update GETTING_CUDAGRAPH_READY.md and CUDA_GRAPH_SYNC_INVENTORY.md with full current status, multi-GPU stream fix, and next steps

2026-06-06 09:17:49 +00:00

14 KiB

Raw Permalink Blame History

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_`
`dsv4/layers/grouped_linear.py`	`expert_offsets[g] = python_int`	✅ FIXED	Pre-allocated range tensor + element-wise multiply
`dsv4/layers/grouped_linear.py`	`if group_offsets[0] != 0`	✅ FIXED	Unconditional GPU-only update
`dsv4/layers/moe.py`	`torch.bincount` (data-dependent shapes)	✅ FIXED	`scatter_add_` into pre-allocated buffer
`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
`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)

⬜ Gate commits on capture test — implement CI check
⬜ Optimize stream sync — pre-create events, reduce per-step overhead
⬜ Long-run stability test — --max-tokens 512+ with --cuda-graph
⬜ Memory leak check — ensure no growing GPU usage over many steps
⬜ 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

14 KiB Raw Permalink Blame History