[Discussion] CUDA graphs for the decode path: staged plan + full shim audit

[riccardo-galbani]

In the recent video you mentioned CUDA graphs as one of the directions worth
pursuing after the GB10 decode work, since the CUDA backend currently launches
every kernel individually. I'd like to work on this and, given the impact
radius, I'm opening a discussion before writing code, with a staged plan and a
complete audit of the decode path. Line numbers below refer to main @ 80ebbc3.

Why the seam already exists

The architecture is already designed for batched submission:
ds4_gpu_begin_commands / ds4_gpu_end_commands. On Metal, begin creates a
command buffer and end commits it — one submission per token. On CUDA, begin
is a no-op stub and end is cudaDeviceSynchronize() (ds4_cuda.cu:2479). This
asymmetry is a large part of the Metal-vs-CUDA overhead gap: a decode token is
~350 individual kernel launches at ~5-10µs each, plus the idle gaps between
them.

Because the Metal programming model forbids host readbacks inside a command
buffer, the token body (metal_graph_encode_token_raw_swa) is by construction
free of mid-token host decisions, and MoE expert selection already flows
through a device buffer (selected tensor in ds4_gpu_routed_moe_one_tensor).
Both are exactly the invariants CUDA graph capture needs — the design work is
mostly already done.

Prior art: the Entrpi fork (decode-perf-tuning branch) implemented per-layer
graph capture and measured +5-10% decode on GB10, growing with context, with
bit-exact output vs eager. I've studied that implementation in depth and the
plan below builds on its validated patterns (and its documented failure modes).

Staged plan

Every stage is env-gated, falls back to eager, and must be bit-exact
(md5 of output, graphs ON vs OFF) before moving on. Flag off = byte-identical
behavior to today.

• Stage 0 — infrastructure, zero behavior change. A thread-local stream
  override (ds4_current_stream() / ds4_capture_set_stream()) so that all
  ~100 ds4_gpu_* shims can be redirected to a capture stream without
  touching their ABI. Small, self-contained; I'd propose this as the first PR.
• Stage 1 — dense-FFN layer cluster. First layers have no router: minimal
  capture, validates the mechanics (warm-up before capture since cudaMalloc is
  forbidden inside one, sync brackets, replay, cache keyed on stable bits).
• Stage 2 — MoE cluster. Adds nothing conceptually: expert selection is
  already device-side.
• Stage 3 — whole layer, then whole token. begin/end_commands become
  capture/replay with a graph cache — the CUDA implementation of what Metal
  already does, with one structural difference: Metal cheaply re-encodes every
  token, a graph replays a frozen recording, so all per-token scalars must
  move from kernel arguments to a small device-side struct (~40B, memcpy'd
  once per token) that kernels dereference at execution time.

Audit of the decode path (complete)

All shims reachable from metal_graph_eval_token_raw_swa (19461), including
static helpers (metal_graph_decode_kv_store, 13463):

• 23 green — no per-token arguments, capturable as-is: all matmuls, norms,
  swiglu, residual adds, MoE dispatch, attention output.
• 11 red (scalar) — carry per-token scalars (pos0, raw_row, n_raw,
  raw_start, n_comp, emit counters, token id in embed and hash-mode
  routing) that must migrate to the device-side struct.
• 1 red (pointer) — dsv4_indexer_qat at 15230 receives a host-computed
  view at offset index_row * row_bytes: the pointer changes every emit.
  Needs a row-view variant (stable base + row index read device-side). The
  same 8 lines also contain a host if (emit) branch and a host counter
  increment — a compact example of everything replay must handle.
• 4 structural — the expert-streaming readback family
  (signal/commit_and_wait_selected_readback,
  stream_expert_cache_begin_selected_load, guarded by
  overlap_selected_shared, 15780) performs a mid-token host readback and
  cannot live inside a graph. Proposed scope cut: v1 graphs only when weights
  are resident; streaming keeps the eager path. Note this doesn't affect the
  primary target (GB10 + Flash q2 = resident).

Questions

1. Is this a direction you want in the tree at all, and at which granularity
   would you want it to stop (per-layer is where the measured gain lives;
   whole-token is architecturally cleaner but marginal, ~1%)?
2. Would you take Stage 0 as a standalone PR? It's behavior-invariant and
   reviewable in isolation.
3. Preference on reusing/adapting the Entrpi implementation (with credit)
   vs. a minimal reimplementation shaped on ds4's simplicity philosophy?

Happy to adjust scope. Development starts on an RTX 4070 (8GB); a GB10 arrives
this week for the remaining stages and the final numbers.

[riccardo-galbani]

Stage 0 is up as #566 , verified bit-exact on sm_89 and sm_121

[riccardo-galbani]

One correction to the staged plan above, caught while scoping stage 1
against the source.

Stage 1 was described as "dense-FFN cluster — first layers without
router". That premise is wrong for V4 Flash: all 43 layers are MoE
(DS4_SHAPE_FLASH: n_expert = 256, no dense-layer field), and the
first 3 layers are hash-routing (n_hash_layer = 3) — still MoE, and
actually the one place where router_select takes the token as a
per-token scalar. The wording was a mental leftover from V3-style
architectures.

Stage 1 is therefore retargeted to the shared-expert FFN cluster,
identical in every layer: shared_gate_up_swiglu_q8_0 → swiglu →
shared_down_hc_expand_q8_0. All three shims are in the capturable set
from the audit (no per-token scalars, no routing indirection, stable
buffer pointers), which makes it the smallest cluster that still
exercises the full capture mechanics (non-blocking stream, scoped
override, warm-up, sync brackets, per-layer graph cache, eager
fallback).

Terminology note for anyone cross-reading Entrpi's decode-perf-tuning
fork: dense_graph there means dense matmul (single attention-side
projections at n_tok=1), not "dense layer". Unrelated to this cluster.

Stage 2 (routed MoE cluster with the device-side selected
indirection) and stage 3 are unchanged.

[riccardo-galbani]

Stage 1 complete: shared-expert FFN cluster capture/replay — correctness validated, performance neutral as expected

Stage 1 is done on riccardo-galbani/ds4, branch cuda-graphs-stage1 (stacked on cuda-graphs-stage0, PR #566). Full stage 1 diff (excluding stage 0): compare view. Five commits:

C1 — process-global NonBlocking capture stream, created in ds4_gpu_init under the DS4_CUDA_GRAPHS gate (value, not presence).
C2 — scoped thread-local stream override (enter at BeginCapture, exit at EndCapture).
C3 — per-cluster state machine (EMPTY → WARMED → CAPTURED), key on 8 fields, bidirectional sync brackets, invalidation on arena generation.
C4 — wiring at the resident call site of metal_graph_encode_decode_layer. Two findings corrected the plan:
the default fused cluster is TWO shims (gate_up_swiglu → down_hc), and the output pointer alternates per token parity (43 layers, odd), so
slots are 4-way set-associative per layer, full-key match, deterministic policy, CAPTURED ways never evicted.
C5 — arena pre-warm: cuda_tmp_alloc sized to the computed session maximum at both graph-creation points (session and sessionless
paths), so the arena never grows after t0 — no cudaMalloc can fall in a capture bracket and the generation stays stable for the whole
session. The dominant prefill site verified exactly linear at 96 KiB/token across two decades of prefill_cap.

Correctness (GB10, sm_121): logits bit-identical to upstream in all graphs × MTP combinations; counter signatures stable across -n 40, ~900 and ~2000-token runs (non-MTP: 86 captured / 86 eager, one capture per way, zero re-captures; MTP: 88/88); arena generation 1 for the whole session. Gate off: byte-identical output to upstream behavior.

Performance (GB10, interleaved ABBA, 5+5 runs, -n 2000, temp 0, identical token streams by bit-exactness): generation 14.004 ± 0.14
t/s (on) vs 14.002 ± 0.16 t/s (off) — neutral within noise, as predicted: the cluster spans 2 shims/layer, so the ~43 saved launches
per token are offset by the 86 sync-bracket events per token that the small capture perimeter requires. This stage was about proving the
mechanics (state machine, pointer-parity slots, arena stability) at zero regression; the measurable gains live in the coverage extension
(the Entrpi fork's +5.2–5.6% comes from a much larger captured perimeter, which is exactly stage 2+ territory: bigger clusters, fewer
brackets per token).

A question for you before I open the next PR: would you rather review stage 1 now as its own PR (infrastructure, neutral perf, gated off by
default), or wait until the stack shows a measurable gain and review a consolidated one? Happy either way — the branches are public and
stacked either way.