[void-mckenzie]

Towards #828. This captures the Gemma 4 assistant draft in CUDA graphs; today it runs eager. The draft keeps no KV of its own, so every one of its k steps reads the same [0, cached_len) window from the target's cache. The existing chain graphs need one set of attention metadata per step because each step writes KV. The shared draft does not, so one graph per batch size holds all k steps with a single metadata set.

• Per call, three things move: the page table row, the cached length and the pending token, copied into fixed buffers before replay.
• The page table walk, the KV index list, the sliding window translation and the argmax all run inside the graph, so replay holds at any sequence length.
• Capture runs after every attach, so a runtime cache rebuild re-captures against the new pool and page table. FREETOKEN_SHARED_DRAFT_GRAPHS=0 keeps the draft eager.

Results

E4B BF16, 25 greedy prompts. Decode tokens/s is sum(predicted_n) / sum(predicted_ms) * 1000.

	decode tokens/s	25-prompt wall
batch 1	280.2 -> 416.2 (+48.5%)	17.8 -> 13.5 s
batch 4	202.1 -> 257.2 (+27.3%)	8.26 -> 7.19 s

Verified

• Outputs identical to the eager draft: 25/25 at batch 1, 25/25 at batch 4.
• Tiny contract model: 336 draft ids compared step by step between the eager and graphs runs, 0 differ.
• The stored batch 4 pair differs at one step, prompt lon4 at cached_len 257. The eager path alone drafted 2856 or 2859 there across boots, and inside one boot the graph and eager drafts match id for id, so the difference is boot to boot target numerics, not the graphs.

Tests

The contract test keeps its strict per-step checks on the eager draft, and now records the graphs run's draft ids too. A graph draft may differ from the eager one only if both are in the eager head's top 3 at that step and within two fp16 rounding steps of each other. Three targeted cases pass; the full file is 6 passed, 5 failed, the same five qwen tie cases as before this branch.

Questions

• The sliding window translation runs over the whole fixed index buffer (sized for the largest batch and the longest sequence), not just the rows in use. ChainGraphs sizes its buffer the same way, so the two paths match. If you want it trimmed, I think both should change.
• A batch larger than the captured sizes falls back to the eager draft and logs once. I did not make it an error because --cuda-graph-bs can legitimately sit below --max-running-requests. Say if you would rather it refuse.

@jason-fxz, which way do you want the first one?

[jason-fxz]

Checked on sm_120 (gemma-4-12B-it + 12B assistant, k=3): outputs and accept len match the eager draft, and decode goes from 117 to 131 tok/s at bs 1.

Request: the graphs-vs-eager check cannot catch wrong attention inputs. The tiny assistant has zero attention/FFN weights and its pre_projection drops the token embedding, so the drafts depend only on the seed hidden. The test still passes with the SWA translate removed, a one-key window, or a zeroed pending token. Please also run it with a random-weight assistant, comparing each replay against eager on the same inputs.

Q1: please trim it in this PR; ChainGraphs needs no change, since its one buffer is shared by every bs and only the gather writes it, touching live entries only. Here each replay translates all max_bs * width entries, and each bs keeps its own swa_indices: at bs <= 64 with a 128K table, that is ~195 µs vs ~10 µs per bs-1 propose, and ~370 MB. Doing the lookup inside gather_kv_indices into one shared buffer fixes both; the per-bs decode scratch can be shared too.

Q2: keep the eager fallback, but drop the counter and the warning. GraphRunner and ChainGraphs fall back silently.

Nits:

• Drop FREETOKEN_SHARED_DRAFT_GRAPHS; a popped draft_graphs test kwarg can no-op capture_graphs.
• _batch duplicates _TritonSteps.metadata.
• The chain_graph.py docstring is stale, and the detach comment's reason is wrong: graphs hold no tensor refs.
• draft["top3"] in the test is unused.