Update CUDA graph docs with current status, A/B split, buffer fixes, remaining blockers
GETTING_CUDAGRAPH_READY.md: - Updated architecture section for A/B split (Graph A + eager attention + Graph B) - Updated Section D integration order with current progress - Added all recent violation fix commits CUDA_GRAPH_SYNC_INVENTORY.md: - Added Category 6 fixes: _l1_out_buf 2x fix, GEMM output pre-allocation, swizzle CUDA kernel, gsa scalar assignment, router BF16 fix - Added remaining blockers for next session - Updated CUDAGraphDecoder architecture description for A/B split - Added capture/replay flow description
This commit is contained in:
@@ -1,37 +1,18 @@
|
||||
# CUDA Graph Readiness — Sync Violation Inventory
|
||||
|
||||
**Date:** 2026-06-03 (updated 23:40 UTC)
|
||||
**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
|
||||
|
||||
**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.
|
||||
**All sync violations in the compute forward path have been fixed.** The eager decode path works at 0.51-0.53s/token.
|
||||
|
||||
CUDA graph capture with A/B split architecture is partially working. Graph A/B capture has been attempted on B200 with `--cuda-graph` flag. Multiple per-step allocation issues have been found and fixed, but full 61-layer capture is NOT YET WORKING.
|
||||
|
||||
- **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
|
||||
|
||||
---
|
||||
|
||||
## B200 Detector Results
|
||||
|
||||
### Run 1 (commit 0ca7bed)
|
||||
- Method 1: 1 violation — `dec_tid_buf[0] = all_tokens[-1]` (CPU→GPU sync from Python int)
|
||||
- Method 2: FAIL — `expert_offsets[g] = (g + 1) * padded_rows_per_group` (CPU→GPU sync in Python loop)
|
||||
|
||||
### Run 2 (commit e07d798)
|
||||
- Method 1: 1 violation — same `dec_tid_buf` (test code not yet fixed)
|
||||
- Method 2: FAIL — `torch.bincount` in MoE (data-dependent shapes)
|
||||
|
||||
### Run 3 (commit 84655d0)
|
||||
- Method 1: 1 violation — same `dec_tid_buf`
|
||||
- Method 2: FAIL — illegal memory access from stride-0 gsa expand view
|
||||
|
||||
### Run 4 (commit 80bb27f) — CURRENT
|
||||
- Method 1: 0 violations in forward (only pinned memcpy flagged, which is graph-capturable)
|
||||
- Method 2: **PASS** ✅ — L0 graph capture succeeds
|
||||
- **Method 2** (L0 graph capture): **PASS** ✅ (from detector test, pre-A/B split)
|
||||
- **Multi-layer A/B capture**: 🔄 IN PROGRESS — multiple per-step allocation issues found and partially fixed
|
||||
|
||||
---
|
||||
|
||||
@@ -76,11 +57,13 @@ All VERBOSE-gated `.item()` calls (diagnostics) are safe at VERBOSE=0.
|
||||
|
||||
---
|
||||
|
||||
## CATEGORY 4: Cross-GPU transfers inside graph — NOT YET ADDRESSED ⏳
|
||||
## 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 |
|
||||
|
||||
---
|
||||
|
||||
@@ -93,82 +76,126 @@ All VERBOSE-gated `.item()` calls (diagnostics) are safe at VERBOSE=0.
|
||||
|
||||
---
|
||||
|
||||
## Section C Hazard Summary (from GETTING_CUDAGRAPH_READY.md)
|
||||
## CATEGORY 6: Per-step allocations inside CUDA graph capture — PARTIALLY FIXED 🔄
|
||||
|
||||
| # | Hazard | Status |
|
||||
|---|--------|--------|
|
||||
| 1 | Compressor returns None for most decode steps | ⏳ Phase 2 (eager-break) |
|
||||
| 2 | KV grows each step | ⏳ Phase 2 (eager-break) |
|
||||
| 3 | Indexer top-k → host reads count | ✅ Already fixed-shape |
|
||||
| 4 | MoE per-expert token counts | ✅ scatter_add_ with pre-allocated buffer |
|
||||
| 5 | Next token/positions on host | ✅ Pinned CPU buffers + copy_ |
|
||||
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 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 |
|
||||
## CATEGORY 8: Cross-GPU operations inside graph capture — FIXED ✅
|
||||
|
||||
## CUDAGraphDecoder Architecture (Current)
|
||||
| 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 |
|
||||
|
||||
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.
|
||||
---
|
||||
|
||||
## 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. 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
|
||||
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
|
||||
```
|
||||
|
||||
Commits: `486f74d` (initial), `92225b0` (simplified from A/B split)
|
||||
```
|
||||
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
|
||||
|
||||
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)
|
||||
1. **Fix Category 6 remaining allocations** — systematic audit of ALL per-step torch.zeros/empty/copy_ in forward path
|
||||
2. **Ensure swizzled buffers allocated before capture** — add explicit allocation in CUDAGraphDecoder.pre_allocate() or before capture
|
||||
3. **Extend capture to all 61 layers** — test on B200 with --cuda-graph
|
||||
4. **Replay verification** — bit-for-bit match with eager forward
|
||||
5. **Performance benchmark** — measure speedup from graph capture
|
||||
6. **Gate commits** on capture test
|
||||
7. **Phase 2**: Paged KV + device-side compressor for full vLLM graph capture
|
||||
|
||||
## Phase 2 (vLLM Integration)
|
||||
|
||||
|
||||
@@ -39,7 +39,7 @@ fire_b200_test tests/unit/test_cuda_graph_readiness.py kernel-test /tmp/kernel-t
|
||||
## SECTION B — The hidden-CPU checklist (grep the hot path for these) ✅ ADDRESSED
|
||||
|
||||
**Explicit device→host transfers** — All `.item()` calls on hot path eliminated:
|
||||
- mhc.py `post_block`: removed `X_next.abs().max().item()` (was 122 syncs/step across 61 layers × 2 mHC)
|
||||
- mhc.py `post_block`: removed `X_next.abs().max().item()` (122 syncs/step across 61 layers × 2 mHC)
|
||||
- All other `.item()` calls are guarded by `VERBOSE >= 2` and don't execute at VERBOSE=0
|
||||
- Warmup-gsa `.item()` calls run once at step 0, outside graph region
|
||||
|
||||
@@ -56,7 +56,7 @@ fire_b200_test tests/unit/test_cuda_graph_readiness.py kernel-test /tmp/kernel-t
|
||||
|
||||
**Host control flow on device values** — Eliminated:
|
||||
- `dec_tid_buf[0] = python_int` → pinned CPU buffer + `copy_` (async, graph-capturable)
|
||||
- `expert_offsets[g] = python_int * padded_rows` → element-wise GPU multiply with pre-allocated range tensor
|
||||
- `expert_offsets[g] = python_int` → element-wise GPU multiply with pre-allocated range tensor
|
||||
- `if group_offsets[0] != 0` → unconditional GPU-only update (no host read of GPU tensor)
|
||||
|
||||
**What is FINE (no sync, don't waste time on these)**
|
||||
@@ -83,13 +83,21 @@ Also confirmed:
|
||||
- **Router** is graph-safe: pre-allocated output buffers, GPU-only operations ✅
|
||||
- **mHC** is graph-safe: fixed-iteration Sinkhorn, no `.item()` on hot path ✅
|
||||
|
||||
### Architectural Decision: Eager-Break-at-Attention (Phase 1)
|
||||
### Architectural Decision: Eager-Break-at-Attention (Phase 1) — UPDATED 2026-06-04
|
||||
|
||||
The per-layer compute is split:
|
||||
- **Captured** (in CUDA graph): mHC pre_block → RMSNorm + quantize → attention projections → o_proj → mHC post_block → FFN mHC → Router → MoE → SE → mHC post_block
|
||||
- **Eager** (outside graph): Compressor → Indexer → KV gather → FMHA → inverse RoPE
|
||||
- **Rationale**: FMHA has dynamic sequence length; compressor/KV are data-dependent. Capturing the compute-heavy parts eliminates ~94ms of Python dispatch overhead per step.
|
||||
- **Phase 2**: Paged KV + device-side compressor → full graph capture for vLLM integration.
|
||||
The per-layer compute is split into **two graph-captured regions** with eager attention in between:
|
||||
- **Graph A** (captured): mHC pre_block(attn) + fused RMSNorm + quantize + q_a + q_a_norm + q_b + kv projections
|
||||
- Outputs written to pre-allocated buffers: x_normed, q_heads, kv_3d, ctx_a_B, ctx_a_C, X_mid
|
||||
- **Eager** (NOT captured): Compressor → Indexer → KV gather → FMHA → inverse RoPE → o_a + o_b → F_attn
|
||||
- Dynamic shapes (FMHA seq_len, compressor returns None) → cannot be captured
|
||||
- `forward_attention()` accepts optional `q_heads`/`kv_3d` to skip projections when called from graph replay
|
||||
- **Graph B** (captured): mHC post_block(attn) + FFN mHC + RMSNorm + quantize + Router + MoE + SE + mHC post_block(ffn)
|
||||
- Reads F_attn from pre-allocated buffer (written by eager attention)
|
||||
- Writes X_next to pre-allocated output buffer
|
||||
|
||||
**Rationale**: FMHA has dynamic sequence length; compressor/KV are data-dependent. Capturing the compute-heavy parts (projections, MoE, SE) eliminates ~94ms of Python dispatch overhead per step. The attention path (which is NOT compute-heavy for T=1 decode) runs eagerly with negligible overhead.
|
||||
|
||||
**Phase 2**: Paged KV + device-side compressor → full graph capture for vLLM integration.
|
||||
|
||||
---
|
||||
|
||||
@@ -97,16 +105,18 @@ The per-layer compute is split:
|
||||
|
||||
1. ✅ **Build Section A's detector and run it on the current forward** — DONE. `tests/unit/test_cuda_graph_readiness.py` on B200.
|
||||
2. ✅ **Fix Section C's five device-native kernels** — 3/5 done, 2 deferred to Phase 2 with architectural decision.
|
||||
3. 🔄 **Re-run capture-under-test until it captures clean** — L0 capture PASSES. Need to extend to all 61 layers + lm_head + replay verification.
|
||||
3. 🔄 **Re-run capture-under-test until it captures clean** — Graph A/B split architecture implemented. Graph capture attempted on B200. Multiple per-step allocation issues found and fixed (see CUDA_GRAPH_SYNC_INVENTORY.md). Still not fully capturing all 61 layers.
|
||||
4. ⬜ **Gate every commit on the capture test** — Not yet implemented.
|
||||
|
||||
### Next Steps
|
||||
1. Extend graph capture from L0 to all 61 layers
|
||||
2. Capture hc_head + norm + lm_head graph on cuda:0
|
||||
3. Implement replay loop and verify bit-for-bit match with eager
|
||||
4. Benchmark: measure speedup from graph capture vs eager decode
|
||||
5. Gate commits on capture test
|
||||
6. Phase 2: paged KV + device-side compressor for full vLLM graph capture
|
||||
### Next Steps (for next session)
|
||||
1. **Continue fixing per-step allocations in graph capture path** — the main blocker
|
||||
2. **Verify swizzled scale buffers are allocated before graph capture** — some paths hit None
|
||||
3. **Test graph capture on B200** — `fire_b200_test single_shot_inference.py kernel-test /tmp/kernel-test.log 1800 -- --max-tokens 30 --cuda-graph`
|
||||
4. **Extend capture to all 61 layers** once per-step allocation issues are resolved
|
||||
5. **Replay verification** — bit-for-bit match with eager forward
|
||||
6. **Benchmark** — measure speedup from graph capture vs eager decode (0.51-0.53s/token)
|
||||
7. **Gate commits** on capture test
|
||||
8. Phase 2: paged KV + device-side compressor for full vLLM graph capture
|
||||
|
||||
## Guardrails
|
||||
- Keep the stop-check, detokenize, and load-time BF16 dequant on the host — they're outside the captured region by design; don't contort them to be "graph-safe."
|
||||
@@ -126,3 +136,6 @@ The per-layer compute is split:
|
||||
| `f13a81d` | grouped_linear scale_a_buf pre-alloc, quantize zeros_like → scalar 0.0 |
|
||||
| `518a1d3` | MoE scatter_add_ int64 indices, fix second bincount call |
|
||||
| `80bb27f` | gsa broadcast: reshape for M=1 decode (no stride-0), contiguous for M>1 prefill |
|
||||
| `6dc2f22` | **CRITICAL: _l1_out_buf 2x too narrow → GPU memory corruption (root cause of ALL cudaErrorInvalidValue errors)**. Also: all GEMM output buffers pre-allocated, gsa copy_ → scalar assignment |
|
||||
| `69e15f1` | Blackwell swizzle CUDA kernel for graph capture, swizzled output buffers |
|
||||
| `ffa7842` | Dense router: BF16 GEMM instead of FP32 conversion during graph capture |
|
||||
|
||||
Reference in New Issue
Block a user