Files

biondizzle ac213bdee8 Update docs: CUDA graph capture WORKING on all 8 GPUs, 0.28s/token (2x eager)

2026-06-06 08:29:40 +00:00

12 KiB

Raw Blame History

CUDA Graph Readiness — Sync Violation Inventory

Date: 2026-06-04 (updated 05:10 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

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: 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+).

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 — PARTIALLY FIXED 🔄

These are torch.zeros(), torch.empty(), and Python view operations that work fine in eager mode but are disallowed during torch.cuda.graph() capture.

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`
`dsv4/layers/linear.py`	`_assemble_scales_single_group` calls `pad_and_swizzle_single`	Same CUDA kernel treatment + pre-allocated `_padded_x_sf_swizzled_buf`	`69e15f1`

IMPORTANT: The swizzled buffers are allocated in _allocate_buffers() / _ensure_buffer_size(). If these haven't been called before graph capture, the buffers will be None. A safety fallback falls through to the Python path (which will fail during graph capture). Ensure all layer buffers are allocated before calling graph_decoder.capture().

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`

STILL BLOCKING ⏳ — Known remaining issues for next session

File	Issue	Notes
Various layers	`.contiguous()` calls inside graph capture may allocate new tensors	Need systematic audit. During graph capture, `.contiguous()` on a non-contiguous tensor allocates. Pre-ensure tensors are contiguous before capture.
`dsv4/layers/mhc.py`	`_dynamic_params` does `X_flat.float()` → new FP32 tensor	This IS captured (new allocation inside graph is recorded and replayed). But need to verify no issues.
`dsv4/layers/mhc.py`	`sinkhorn_knopp` CUDA kernel returns new tensor	Same — allocation is recorded and replayed. Should be fine.
`dsv4/layers/moe.py`	`l1_out[padded_dst]` — advanced indexing creates new tensor	This IS captured and replayed. Should be fine.
`dsv4/layers/moe.py`	`deinterleave_l1_weights` — creates new tensor	Need to verify graph-capturable
`dsv4/layers/moe.py`	`sorted_token_ids` from `argsort` — creates new tensor	Captured and replayed. Should be fine.
`dsv4/ops/quantize.py`	`quantize_nvfp4_gpu_fused` returns new tensors from CUDA kernels	Captured and replayed (kernel output is recorded). Should be fine.
Shared expert / linear	Swizzled buffers may be None if `_allocate_buffers()` not called before capture	Safety fallback to Python path will FAIL during graph capture. Must ensure buffers allocated.

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`

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

CUDAGraphDecoder Architecture (Current — A/B Split)

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. Capture Graph A: 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_bufs[li], ctx_a_C_bufs[li], X_mid_bufs[li]
   b. Capture Graph B: 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

Replay flow:
1. For each layer li:
   a. Copy X → x_in_bufs[li] (handles cross-GPU transfer)
   b. Replay Graph A → read q_heads_bufs[li], kv_3d_bufs[li], x_normed_bufs[li]
   c. Run eager attention: forward_attention(... q_heads=q_heads, kv_3d=kv_3d ...)
   d. Copy F_attn → F_attn_bufs[li]
   e. Replay Graph B → read x_out_bufs[li]
   f. X = x_out_bufs[li]
2. Copy X → x_lm_in → replay lm_graph → read logits_buf

Commits: 6dc2f22 (initial A/B split + critical buffer fixes), 69e15f1 (swizzle kernel), ffa7842 (router fix)

Remaining Work for Full Graph Capture

Fix Category 6 remaining allocations — systematic audit of ALL per-step torch.zeros/empty/copy_ in forward path
Ensure swizzled buffers allocated before capture — add explicit allocation in CUDAGraphDecoder.pre_allocate() or before capture
Extend capture to all 61 layers — test on B200 with --cuda-graph
Replay verification — bit-for-bit match with eager forward
Performance benchmark — measure speedup from graph capture
Gate commits on capture test
Phase 2: Paged KV + device-side compressor for full vLLM graph capture

Phase 2 (vLLM Integration)

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

12 KiB Raw Blame History