[FEAT] Capture fixed-depth speculative decode with CUDA Graphs

[0z5a]

Scope

Implements the CUDA Graph feature claimed by @0z5a in #43 on top of #44. Fixed-depth synchronous speculation captures recurrent loops and the LM head. Shape/position changes use matching graph keys or eager fallback; sampling and request bookkeeping remain outside capture.

RTX 5090 E2E

Ouro-1.4B, BF16/Triton, d=2/D=4, K=4, 64 output tokens, one resident model, five warmups per arm and five alternating paired trials. Both arms used one OMP/MKL/OpenBLAS CPU thread. Every output token and exit depth matched.

Workload	Eager median	Graph median	Speedup
Repeated prompt, 1 request	3.028 s	0.763 s	3.969×
Repeated prompt, 4 requests	2.223 s	1.209 s	1.838×
Prose prompt, 1 request	1.865 s	0.753 s	2.477×

Recurrent graphs: 11 captures / 5,400 replays. LM-head graphs: 11 captures / 2,240 replays. Eager fallbacks: zero. The single-request repeated prompt is sensitive to host scheduling; a separate run during host load around 162/128 logical CPUs had substantial timing drift. The table above summarizes the recorded trials.

Verification and limits

• RTX 5090 tiny-model eager/graph exact output and ragged replay tests passed (2 GPU cases); graph fallback/regression test passed.
• Full CPU suite: 286 passed, 122 GPU tests skipped. Ruff check/format passed.
• GPU tests and E2E used the existing /home/gongji/0z5a environment inside Docker with a verified 96 GiB RAM cap.
• Packed FA4 capture remains unverified: the installed paged FA4 backend does not support SM12 on RTX 5090.

Built on merged #44. This PR contains only the CUDA Graph work claimed in #43.

CUDA Graph capture cleanup

Recurrent and coda graphs now share one warmup/capture helper. Their graph keys, eager fallbacks, and steady-state replay behavior are unchanged. Post-change verification: CPU suite 286 passed, 11 skipped, 111 GPU cases deselected; RTX 5090 Triton CUDA Graph suite 12 passed, 14 deselected; Ruff check/format passed. The latest stacked scheduler run also exercised this Graph parent with exact output matches and zero eager fallbacks.

[bjf-frz]

Tested speculative decoding with CUDA Graphs enabled/disabled on a B300 using Ouro-1.4B, BF16 and FA4.

Following the README workload methodology: 1,024 input tokens, 128 output tokens, 2× concurrency requests with closed-loop replacement. Timing includes prefill and drain, excluding loading and warmup. Both arms use synchronous speculation with draft depth 2, target depth 4 and K=4.

Throughput (output tokens/s; higher is better):

Concurrency	Percentile	Graph off (tokens/s)	Graph on (tokens/s)	Improvement
1	P50	38.76	90.83	+134.4%
1	P95	39.05	90.90	+132.8%
1	P99	39.07	90.91	+132.7%
32	P50	539.35	792.52	+46.9%
32	P95	540.94	801.20	+48.1%
32	P99	541.08	801.97	+48.2%
128	P50	862.63	972.11	+12.7%
128	P95	863.35	972.91	+12.7%
128	P99	863.42	972.98	+12.7%

Improvement = (Graph on / Graph off − 1) × 100%, computed from the corresponding throughput percentiles.

Sample-size note: these are percentiles across only three measured runs per configuration after two warmups, calculated with linear interpolation. P95/P99 are descriptive of these three runs, not reliable tail estimates or request-latency percentiles. C32 eager had one slow run (292.39 vs. 539.35–541.12 tokens/s), so its improvement remains preliminary.

Generated token IDs and exit depths matched across all measured pairs. Three additional FA4 correctness tests passed, covering packed metadata updates, graph reuse, and fallback behavior.

Configuration details: graph limits were raised to 640 rows / 64 shapes, rather than the README defaults. Capture occurred during warmup; caches were frozen for measured runs with eager fallback allowed. At C128, recurrent graphs recorded 888 replays and 1,076 fallbacks per run.

Tested commit: b574eb3f78281ac96cfb35e898ac6c166f390357, plus the local _capture change replacing the explicit device-wide synchronization with wait_stream. This comparison does not isolate the effect of that change.

One potential follow-up optimization is improving graph coverage at high concurrency: at C128, only 45.2% of recurrent calls used graph replay (888 replays vs. 1,076 eager fallbacks per run). Tracking fallback reasons and shape frequencies could help identify whether a different shape-cache policy or shape bucketing would improve coverage. This can be considered future work.

[bjf-frz]

plz remove benchmarks & results dirs.

[0z5a]

plz remove benchmarks & results dirs.

done