Repository navigation
[FEAT] Add CUDA Graphs for cross-round async self-speculation - #127
liuyao0322 wants to merge 21 commits into
Conversation
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
There was a problem hiding this comment.
Reviewed at head 9eb01bb (base 9d601ce). The runtime change is small and the ownership model looks right. I found no correctness blocker:
- Keying core graphs by row count is valid for Ouro. Depth only selects the block-table plane, which is copied in per replay, and Huginn is rejected for speculation.
- Warmup and capture KV writes are idempotent at real positions.
- The
finally: stream.synchronize()in_capturecloses the failed-warmup race. - Outputs are cloned before pool reuse.
_require_live_batchruns before both replay and fallback.
My main concerns are about what users get by default and under non-lockstep traffic:
1. Default graph limits disable most of the high-concurrency gain (ExecutionConfig.cuda_graph_max_batch_size=128, DeviceRecurrentGraphs.run).
The limit counts rows, not requests. With K=2, verify has 3C rows, so at C≥43 the verify core and coda silently run eager. Above C=128 the draft path does too. The native baseline also falls back above C=128.
The higher-concurrency table raises the limit to max(128, 3C) for all modes, so it does not show the default configuration. The doc example uses K=4, where verify has 5C rows and falls back from C≥26, but every measurement is K=2.
Could the limit be derived from max_num_seqs * (K+1), or could the engine at least warn at startup? Please also state this next to the tables.
2. The exact-row-count cache with 16 entries and no eviction fits only burst workloads (DeviceRecurrentGraphs.run, CodaGraphs.run).
The measurements are lockstep bursts: all requests arrive at once, use ignore_eos and have fixed output lengths. Under that workload shapes barely change, so "0 new captures / 0 fallbacks" is nearly guaranteed.
With staggered arrivals, mixed lengths and natural EOS, the per-offset draft count, verify rows and coda rows all vary. Once 16 shapes are cached, every new shape runs eager permanently. Each lazy capture also calls torch.cuda.synchronize mid-serving, which stalls both in-flight rounds.
Bucketing (padding to capture sizes, as vLLM does), LRU eviction, or pre-capture of expected sizes would make this robust. A test that fills the cache and then sees new shapes would also help.
3. The memory reserve does not scale with the number of graph pools (core/memory.py::plan_cache).
A single cuda_graph_memory_reserve_bytes (1 GiB) now covers four pools: native core/coda and spec core/coda. Each coda entry's output is the full bf16 logits, about 75 MB at 768 rows, so 16 large entries in the spec coda cache alone exceed the reserve. With gpu_memory_utilization planning, this can OOM at capture time.
4. Minor: CodaGraphs clones the full logits on every replay only for argmax. Capturing the argmax, or returning only token IDs on this path, would avoid rows × 49152 × 2 B of copying per call. That copying is likely part of the reported D2D growth from 48 to 198.
5. Test and evidence gaps:
- No test for new shapes after the cache is full.
- No test for high concurrency with default limits.
- No test for K=4.
- No test for a non-default CUDA device.
- All evidence is from H20 with torch 2.14.1 / triton 3.8.
I'm running supplementary tests on H100, Ouro-1.4B 574fa66c:
- PR GPU tests on head and base.
- C16–256 with K=2/4, default vs. raised graph limits, against native async Graph and spec eager.
- Staggered/ragged arrivals with natural EOS, counting captures, fallbacks and latency tails.
- A non-default device, two engines in one process, and DP4.
I'll post results here when they're in.
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
done |
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
Re-reviewed at head 73cfe4e. Thanks for the quick follow-up. The startup and cache-full warnings, the per-pool graph memory planning and the clarified docs address the "silent" part of my earlier points 1–3.
The measurements below were taken at 9eb01bb. The 73cfe4e delta does not touch replay or fallback selection: _cache_full is equivalent to the old condition, and planning is bypassed by the explicit KV sizes I used. So they apply to the current head.
I'm requesting changes for two reasons.
A. Coverage under non-lockstep traffic (blocking).
With exact-row-count keys and 16 entries, 77–93% of warmed core calls fall back to eager under staggered arrivals. With 64 entries, 44–58% still fall back. Graph then adds only 1–6% over spec eager. A warning documents this but does not fix it.
Please bucket row counts to a fixed set of capture sizes: pad inactive rows, and mask or skip their KV writes. Pre-capture those sizes, or at least use an LRU, so serving traffic reuses graphs. The bucketed path also needs a test that checks outputs stay bitwise equal to eager with padded rows.
B. Performance claims vs. native async Graph (blocking for the docs/README).
On H100 with non-repeated prompts, the gain over native is confined to the steady-state window. End to end, including the tail as requests finish at different times, short-output bursts reach only 0.92–1.00× native at K=2 and 0.71–0.81× at K=4. Longer outputs and inputs show real gains of 1.05–1.20×.
Please report end-to-end throughput next to the window/decode numbers and state where the feature wins. Either drop the K=4 example or benchmark it.
Non-blocking.
- The new planning reserves
3 × 1 GiBof pool headroom plus retained logits for async speculation. Your own measured graph overhead is 96–283 MiB, so memory-planned users lose ~2.5 GB of KV capacity. Consider sizing the per-pool headroom from rows/entries instead of a flat 1 GiB each. - Consider auto-sizing
cuda_graph_max_batch_sizefrommax_num_seqs * (K+1)instead of only warning. - The full-logits clone in
CodaGraphsforargmaxremains (point 4).
Supplementary H100 results. Raw JSON/logs are available on request.
Setup
- Hardware and software: one H100 per engine, torch 2.12.1 / triton 3.7.1. This is older than the PR's 2.14.1 / 3.8.0.
- Model: Ouro-1.4B
574fa66c, BF16, d=2/D=4, greedy, Triton attention. - Prompts: distinct natural-text slices, not 8 repeated texts. This gives lower acceptance than the PR: 72–91% instead of ~93%.
- Each arm ran two alternating blocks (ABCD/DCBA) on the same GPU. Each block was warmed once on the exact workload, then timed 3 times, so 6 timed trials per arm, reported as medians.
- Native async Graph uses
ouro_delayedatexit_threshold=1,max_loops=4. All early-exit checks are gated onexit_threshold < 1, so it is full depth, matching the PR's native control. graph defaultusescuda_graph_max_batch_size=128.graph raiseduses(K+1)*Crows. Native usesmax(128, C)rows.
Correctness
- The PR's GPU tests pass on H100: 212 passed at head, 183 at base, 16 deselected (FlashAttention), no skips.
- Speculative Graph vs. speculative eager is bitwise identical in every lockstep condition, for all requests:
- C16–256 at K=2 and K=4
- input/output 512/128
- input/output 64/256
- It is also bitwise identical under step-indexed ragged arrivals, 320/320 requests, with default limits at K=2/4 and with 64 graphs at K=4.
- The one exception is 64 graphs at K=2: 305/320 identical. Arrival steps are fixed, but ticket collection depends on
event.query()timing, so batch composition can still differ between runs. I attribute the divergence to that; it is not proven to be a graph bug.
- The one exception is 64 graphs at K=2: 305/320 identical. Arrival steps are fixed, but ticket collection depends on
- Speculation (eager or Graph) vs. native async Graph is not identical on these prompts, e.g. 52/64 at C64 and 230/256 at C256, with first divergences as early as output index 1–4. The eager path shows the same divergences, so this is not introduced by this PR. It does qualify the report's "all token IDs matched native" claim, which holds for its frozen inputs only.
Multi-GPU
- An engine on
cuda:1(with and without atorch.cuda.deviceguard), two Graph engines interleaved in one process, and 4 data-parallel replicas all produce identical tokens to thecuda:0reference: 32/32, and identical digests across all 4 GPUs.
Lockstep bursts, input/output 64/64 (tok/s)
"window" is the period in which every request has its first token and none has finished. "total" covers submit to last finish, including the tail as requests finish at different times.
| C / K | native window | spec eager | Graph default | Graph raised | total: native → Graph raised |
|---|---|---|---|---|---|
| 16 / 2 | 716 | 414 | 920 | 911 | 687 → 629 (0.92×) |
| 32 / 2 | 1347 | 830 | 1675 | 1675 | 1257 → 1257 (1.00×) |
| 64 / 2 | 2419 | 1567 | 2555 | 2814 | 2103 → 1989 (0.95×) |
| 128 / 2 | 3578 | 2978 | 4094 | 4295 | 2913 → 2892 (0.99×) |
| 192 / 2 | 4649 | 3993 | 3966 | 5029 | 3477 → 3333 (0.96×) |
| 256 / 2 | 5365 | 4802 | 4799 | 5670 | 3863 → 3606 (0.93×) |
| 16 / 4 | 716 | 374 | 873 | 873 | 686 → 484 (0.71×) |
| 64 / 4 | 2414 | 1422 | 2369 | 2572 | 2098 → 1621 (0.77×) |
| 128 / 4 | 3578 | 2744 | 3822 | 3930 | 2913 → 2371 (0.81×) |
| 256 / 4 | 5372 | 4166 | 4163 | 4934 | 3871 → 3075 (0.79×) |
Longer sequences, K=2, total:
- 32 × 512/128: 583 → 697 (1.20×)
- 64 × 512/128: 784 → 827 (1.05×)
- 64 × 64/256: 2092 → 2318 (1.11×)
DP4 at C128 64/64 K=2 total: native 2791–2814 vs. Graph raised 2697–2760 per GPU (0.97×).
Takeaways:
- Graph vs. eager is a real gain with raised limits, 1.1–2.2× in the window. With default limits it drops to 1.0× at C192/C256.
- Versus native async Graph, the gain is confined to the steady-state window. End to end, short-output bursts are 0.71–1.00× native at K=2/4. K=4, the doc's example, is 0.71–0.81×.
- With default limits, Graph gives no gain over eager at C≥192 (K=2) or C=256 (K=4). In those cells the window equals spec eager, and at C=192/256 it falls below native. This is concern 1 from my review, now measured.
- Warmed fallbacks are never zero, even in lockstep bursts with raised limits. Per timed trial, the core cache falls back 24–148 times and the coda cache 1–60 times, because the 16-entry cache is filled by tail shapes.
Staggered arrivals (320 requests at 24 req/s, input 32–1021, output 16–375, about half natural EOS, max_num_seqs=128; warm passes):
| K | arm | tok/s | TTFT p50/p99 (s) | TPOT p50/p99 (ms) | core replays / fallbacks |
|---|---|---|---|---|---|
| 2 | native async Graph | 614–641 | 7.6–8.4 / 26–27 | 112–119 / 273–301 | 190–450 / 2200–2380 |
| 2 | spec eager | 856–892 | 4.7 / 15.0–15.6 | 82–84 / 212–227 | – |
| 2 | spec Graph, 16 graphs | 897–900 | 4.5–4.6 / 15.0–15.2 | 82 / 206–207 | 142–206 / 1736–1800 |
| 2 | spec Graph, 64 graphs | 926–945 | 4.1–4.4 / 13.3–14.2 | 75–78 / 190–210 | 996–1080 / 854–978 |
| 4 | spec eager | 943–983 | 3.6–4.8 / 14.8–16.7 | 82–86 / 217–257 | – |
| 4 | spec Graph, 16 graphs | 1014–1028 | 3.5–4.1 / 14.8–15.0 | 81–82 / 209–219 | 330–410 / 1412–1492 |
| 4 | spec Graph, 64 graphs | 1032–1040 | 3.4–3.9 / 13.9–14.6 | 78–80 / 211–213 | 762–1014 / 844–1050 |
Under realistic arrival patterns, 89–93% (K=2) and 77–82% (K=4) of warmed core calls fall back to eager with the default 16 entries, and still 44–58% with 64 entries. Graph therefore adds only +1–6% over spec eager. That is concern 2 from my review. Speculation itself (eager or Graph) clearly beats native async Graph here. Native's own exact-row-count graph cache shows the same churn: 83–93% of core calls fall back. Speculative cold passes captured 16 or 64 shapes inside the serving window without step-time spikes beyond warm levels (max step 428–499 ms vs. 390–540 ms). Native's first cold pass hit a 1985 ms step.
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
|
Follow-up at head At high concurrency the head still loses to native async Graph end to end. Measured on H100 with Ouro-1.4B BF16, input/output 64/64, non-repeated prompts:
I implemented fixes on top of git fetch https://github.com/0z5a/vllm-rlt.git pr127-high-concurrency
git cherry-pick eba6b30..FETCH_HEADThe approach: 1. Derive the row limit when unset. 2. Exact full-batch buckets. I also removed the per-replay 3. Vectorize the per-round request loops. This was the main overhead. In the PR:
With this change, a forced K'=0 round is 1.05–1.16× native at C64–256 instead of 0.89–0.98×. 4. Adaptive drafts per round (optional, default in the PR).
Fixed K=4 still trails native at high concurrency even after step 3 (0.98–0.99×), because each extra verification row costs real compute once the GPU is saturated. The adaptive policy avoids that. Results on H100. End-to-end throughput, including the tail, divided by native async Graph at
Other fix-branch results:
Given these numbers, I'd suggest cherry-picking at least (1)–(3). (3) gives the largest gain and doesn't change semantics. (4) changes the default to adaptive drafts; if you'd rather keep fixed K as the default, drop the last commit ( |
…it is unset cuda_graph_max_batch_size now defaults to None. The engine resolves it once, before memory planning and graph construction: async speculation covers min(max_num_batched_tokens, max_num_seqs * (K + 1)) verification rows, never below the previous 128; native and synchronous paths keep 128. Explicit limits are unchanged and still warn when they cannot cover verification. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add SpeculativeConfig.draft_policy ('fixed' | 'adaptive') and a
draft_tokens_override diagnostic knob. The scheduler takes the round's draft
count from a callable over the eligible requests' token starts. The adaptive
policy picks K' in [0, K] maximizing expected committed tokens per predicted
round time: round time is fitted online by recursive least squares from
compute-stream events (prior from startup bucket replay timings), acceptance
by a Beta-prior truncated-geometric estimate, with periodic exploration.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… reset Async speculative buckets now keep max_num_seqs * (K' + 1) rows for every K' <= K exactly (bounded by the token budget), plus a halving ladder below max_num_seqs, falling back to the previous ladder when that exceeds cuda_graph_max_graphs. Full-batch draft and verification rounds no longer pad (e.g. 256 rows used the 384 bucket). Rows above the largest bucket run eager. Padding rows have zero context, so the per-replay block-table reset is unnecessary and removed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Each round previously issued several small kernels per request (state gathers, position adds, per-request acceptance and state writes): thousands of launches and copies at high concurrency, which made a forced K=0 round ~7% slower than native decode at 256 requests. Live request (token, position) pairs now sit in a device slot table; a round gathers and scatters them once, drafts and verifies with host-built index tensors uploaded in one copy, and accepts all requests with padded tensor ops. Row order and results are unchanged. The adaptive policy estimates acceptance per draft position instead of one geometric rate, which overstated deep drafts at K=4. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
SpeculativeConfig.draft_policy defaults to 'adaptive' (K is the per-round maximum); 'fixed' keeps K drafts per round. Mechanism tests pin 'fixed'. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
done |
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
LGTM. Approving, but the docs need to be updated before merge.
What I checked at c3a97d3
- The six cherry-picked commits match liuyao0322#1 one-to-one, and the final tree is identical to my validated branch (0z5a:pr127-high-concurrency). The H100 results from my Oct 17 comment therefore apply to this head as-is: 239 passed, 1 skipped (two-GPU case); FP32 Graph/eager token-identical; cuda:1, two-engine and DP4 checks match the reference.
- Ruff check/format are clean on the Python files, and there are no whitespace issues.
Blocking items from my Oct 15 review
- A, coverage under non-lockstep traffic: resolved. Row buckets are precaptured at startup, inactive rows are masked in the Triton KV scatter, and padded Graph matches padded eager 32/32 at K=2 and K=4. With the derived row limit and exact full-batch buckets, staggered arrivals ran with zero graph fallbacks.
- B, performance claims: resolved in substance, but the docs still describe an older head (see below).
Required before merge: docs
- README L107 still reports the pre-fix exact-row results (0.92–1.00× native at K=2, 0.71–0.81× at K=4). These no longer describe this head.
- docs/async-speculative.md L68–70 says the default 128-row limit stops covering verification at C43 (K=2) / C26 (K=4).
cuda_graph_max_batch_sizeis now unset by default and derived as min(max_num_batched_tokens, max_num_seqs × (K+1)), never below 128. The bucket description around L47 should also mention the exact full-batch sizes. - Add the end-to-end H100 results for this head (input/output 64/64, throughput including the tail ÷ native async Graph, medians of 6 trials):
| Config | C64 | C128 | C192 | C256 |
|---|---|---|---|---|
| K=2, default | 1.09 | 1.20 | 1.15 | 1.18 |
| K=4, default | 1.04 | 1.20 | 1.15 | 1.16 |
| K=4, fixed K | 0.99 | 1.08 | 0.98 | 0.98 |
- Document
draft_policy/--draft-policy. Async speculation now defaults to adaptive drafts, so K is the per-round maximum;fixedkeeps K drafts every round. Synchronous speculation is unchanged. - Update the PR description. The results and test plan still say the H100 numbers were not rerun after bucketing. The scope also changes: graphs remain opt-in, but the adaptive default applies to all async speculation, with or without graphs.
- Nit: the metadata audit mentions zeroed trailing pages. Block tables are no longer zeroed now that
block_tables.zero_()is gone.
Non-blocking follow-ups (in my commits, can go in a separate PR)
calibrated_priorcounts one coda call too many: the constant term should be-coda_call * verify_loops / draft_loops. The 5% tolerance intest_calibrated_prior_reproduces_bucket_round_costhides this. It only affects the prior, which RLS overrides quickly.- Per-position acceptance counts have no forgetting, so long-running servers adapt slowly to workload drift.
_retired_slot_valuesis never released after_grow_slots(small).
Disclosure and review request
6 of the 14 commits are mine, so this approval is not an independent review of them.
@bjf-frz PTAL, in particular:
- Are you OK with adaptive drafts becoming the default for all async speculation? If you'd rather keep fixed K, drop c3a97d3 or flip the default.
- This PR has no CI. All GPU evidence comes from manual runs (leo's H20 jobs and my H100 runs).
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Preserve bucketed capture and masked KV writes through the shared graph preparation/execution interface. Move adaptive policy and graph memory planning into StageEngine, retain scheduler-owned result application, and cover connector readiness in draft selection. Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
|
@0z5a Could you please re-review the latest head, The requested README/docs/description updates are complete, with the original H100 table embedded and the adaptive default made explicit. This head also merges main CPU CI is green. The description records the separate-file regression results and the remaining combined-process capture/GC errors and FA4/NIXL validation limits; the older H100 performance results are explicitly scoped to |
|
@bjf-frz hello, could you please review this pr? |
will do right now |
| sizes = preferred | set(ladder(min(preferred))) | ||
| if len(sizes) > execution.cuda_graph_max_graphs or 1 not in sizes: | ||
| return tuple(reversed(ladder(limit))) |
There was a problem hiding this comment.
The bucket-selection priority seems reversed here: we select preferred sizes, add smaller buckets, then discard all preferred sizes if the expanded set exceeds the budget.
For example, with at most 4 requests and K=2, verification has at most 12 rows. The preferred buckets (4, 8, 12) already fit a three-Graph budget. Adding (1, 2) exceeds that budget and triggers a fallback to (32, 64, 128) when the global row limit is 128. Consequently, the 64- and 128-row graphs are captured but never used, while every actual call is padded to 32 rows.
Could we prioritize preferred sizes and coverage of the maximum reachable row count, then add smaller buckets only if budget remains? If pruning is necessary, the resulting buckets should still stay within the speculative workload’s reachable range.
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
One startup preparation optimization:
| hidden = parameter.new_zeros((size, self.model.config.hidden_size)) | ||
| call = self.select(hidden) | ||
| self.prepare_inputs(call, hidden) | ||
| self.execute(call, hidden) |
There was a problem hiding this comment.
prepare() calls execute(), which clones the full logits output at line 459, but prepare() discards the returned tensor. This copy serves no purpose during startup and adds unnecessary GPU memory allocation and copying for every bucket. Could we skip the output clone during preparation while retaining any required warmup replay and stream dependencies? The clone in normal execution should remain to protect results from subsequent replays.
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
A non-blocking startup synchronization optimization to consider:
|
|
||
| def _capture(hidden, stream, pool, run): | ||
| """Warm and capture one stable-input graph on its own stream.""" | ||
| torch.cuda.synchronize(hidden.device) |
There was a problem hiding this comment.
For startup pre-capture, could we narrow this per-bucket device-wide synchronization to an explicit dependency on the input producer stream, once all input, KV, and graph-pool dependencies have been verified? For example, stream.wait_stream(torch.cuda.current_stream(hidden.device)) before warmup would order the capture stream after input preparation, while retaining the existing stream.synchronize() after warmup.
This suggestion is specific to startup pre-capture, not an unconditional replacement for runtime lazy capture. Simply synchronizing the capture stream would not establish input readiness, and the finally synchronization should remain to drain pending KV writes on failure. This is an optimization candidate rather than a confirmed bottleneck: please measure the startup cost, since the benefit may be small if no unrelated streams are active.
Keep preferred full-batch shapes within the graph budget before adding smaller buckets. Avoid discarded startup logits clones and order startup capture after the producer stream while preserving lazy-capture barriers and failure drains. Add regressions for tight budgets, startup dependencies and output ownership; document measured graph-preparation costs and their limits. Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
|
We benchmarked the CUDA Graph changes in isolation against main ( Setup: one B300, Ouro-1.4B BF16, FA4, fixed draft/target depths 2/4 and K=4. Both arms use async scheduling, CUDA Graphs, a 256-row limit, 256 graph entries, and 40,000 KV blocks. Following
Values are three-run medians, except main C32, which has two runs. All selected windows had zero request failures. W5/W6 had no new graph captures; main C32 included 20 and 12 captures, respectively, so its comparison is diagnostic rather than a steady-state speedup. The ShareGPT arms were measured separately, which also limits confidence in the smaller gains. Where the improvement comes from The change does more than improve graph coverage:
The 11.59 ms reduction in KV-write time closely tracks the 12.35 ms shorter GPU execution span. Almost all GPU work ran on the same compute stream. This supports fused KV writes as a major contributor, although profiled timings do not establish its exact contribution to end-to-end throughput. For fixed-length C8, throughput during the 60-second admission window improved by approximately 4.4%; including drain gives the reported 9.0%. Differences in the work remaining at cutoff and its completion timing therefore affect the overall result. We have not inspected per-round drain batch sizes and do not claim improved scheduling efficiency. Acceptance rates and verification rows per output token were nearly unchanged. Remaining validation A KV-only ablation is prepared but has not run because all GPUs are occupied. Also, only 227 of 244 common requests matched token-for-token in the fixed-length comparison. Both arms generated exactly 512 tokens per request, so shorter outputs do not explain that gain, but the mismatches remain unresolved. Overall, the fixed-length C8 measurements show a repeatable throughput benefit, with profiling pointing to cheaper KV writes as an important mechanism. These results remain diagnostic and do not establish correctness acceptance. |
|
Could we reuse the existing Triton These kernels already support GPU-resident slot indices and configurable row widths/strides. The I would prefer to build on the existing routing and buffer-management mechanisms. Please also explain which parts of the new slot allocation and retirement logic need to remain specific to speculative execution, versus what can be shared with Acceptance computation can remain speculation-specific; reusing the data-movement kernels does not require merging the entire speculative lifecycle into |
Pause automatic GC across capture so unreachable old graph executables cannot be destroyed while CUDA is capturing. Restore the caller's GC state after draining the capture stream, including failure paths, without adding per-bucket heap collections. Add a deterministic old-graph collection regression and GC-state failure coverage. The combined graph/speculation suite now passes in one process. Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
|
My main concern is that deciding the draft length based on position-level acceptance probabilities is difficult to justify. Why should a token's position in the draft determine whether we continue speculating? Historical acceptance rates at that position do not explain why continuing or stopping is appropriate for the current request. If we want to implement adaptive K, I would prefer a predictor conditioned on the current model state, such as a lightweight acceptance prediction head, instead of using position-level statistics as the decision basis. DSpark's confidence head is a relevant reference. EAGLE's use of hidden-state features is also worth exploring as a basis for such a predictor, although EAGLE itself is not an adaptive-K prediction head. Please also provide a dedicated benchmark isolating adaptive K, following the protocol in |
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Reviewed at e80f540884883cc9c1aac68fb3d6ac2212aa00a5; the remote head still matches. Request changes for the P2 adaptive-cost stability defect in the inline comment.
Reusing graph execution with device-owned positions is justified by the kernel-submission overhead. Eager asynchronous execution with fixed K remains a credible fallback. Making adaptive drafting the default adds the stability issue described below.
Checked the 22-file diff, prior reviews and subsequent fixes, device KV/slot ownership, output cloning, ticket retirement, and fallback. The earlier bucket-priority and startup-clone issues are addressed; I found no additional graph-buffer lifetime defect.
Validation: CI passed Ruff and 834 CPU tests, with 403 skipped, on synthetic merge 6a14f17 with newer main ff05fe1. PR metadata base and merge base are 4a3503440279. A reviewer-authored scalar arithmetic reproduction independently confirmed the finding in Python without importing or executing repository code.
No GPU execution or benchmark campaign was performed for this review. The reported H20 results remain manual evidence; the H100 performance results qualify c3a97d3, rather than this integration head. BF16 task accuracy and FA/NIXL qualification remain validation gaps.
| self.covariance = [ | ||
| [(p[r][c] - gain[r] * px[c]) / lam for c in range(NUM_FEATURES)] | ||
| for r in range(NUM_FEATURES) |
There was a problem hiding this comment.
[P2] Stabilize the adaptive cost covariance
Repeated full batches make several cost features linearly dependent. This update divides covariance by 0.99 on every round, so uncertainty in unobserved directions grows without a bound or recovery path. Eventually the covariance and fitted coefficients become non-finite. predict_ms() then returns the minimum cost for every K, causing maximum drafting even when drafts consistently fail. The adaptive policy persists across requests and is now the default for async speculation, with graphs enabled or disabled.
A deterministic reviewer-authored scalar Python reproduction used C=64, d/D=2/4, K=4, token budget 4096, the default prior/exploration schedule, contexts cycling through 64–127, and rejection of every first draft. Synthetic elapsed time was the feature dot product with [0.05, 0.002, 0.002, 0.002, 0.1]. The denominator first became negative at round 2,352; coefficients became NaN at round 69,856. Subsequent predictions were [0.001, 0.001, 0.001, 0.001, 0.001] and choices switched to K=4. This evaluates the scalar equations independently; it does not claim a GPU/runtime reproduction.
Please bound or recondition the covariance, recover from invalid fits, and add a repeated-full-batch regression that checks finite estimates and continued adaptation.
There was a problem hiding this comment.
Fixed in 05891dd and merged with main 66920da.
The fitter no longer updates an inverse covariance. Scaled features feed discounted normal equations with a persistent 1e-4 ridge; the information matrix stays bounded away from singularity in unobserved directions, and Cholesky solves avoid the subtractive covariance update. Non-finite observations are skipped; invalid fits recover to the calibrated prior and increment fit_resets. Non-finite predictions cannot silently become the minimum cost for every K.
The 80,000-round C64/K<=4 regression rejects every first draft and cycles contexts through 64–127. The exact test overlay fails on e80f540 with non-finite theta. The fixed version has zero resets, keeps selecting K=0 outside exploration, and learns a threefold timing change after round 75,000. Invalid coefficient/information/RHS recovery and invalid observations are covered too: python -m pytest -q tests/test_draft_policy.py (13 passed).
Full pre-commit passed. Merged-runtime single-GPU graph/speculation regressions: 231 passed, 2 skipped, 25 deselected. An additional merge suite has 201 passes and one eager split-prefill assertion failure independently reproduced on upstream main 66920da; details and the actual validation limits are in the description. No new adaptive-performance claim is made by this stability fix. Please recheck the numerical fix at this head.
Preserve staged-prefill planning and pre-capture cleanup from main. Replace covariance inflation with scaled, persistently regularized information fitting and Cholesky solves; recover invalid fits to the prior. Add an 80,000-round collinear full-batch regression with rejected drafts and a late timing change, plus invalid-state recovery checks. Clarify population acceptance estimates and the limits of combined performance evidence. Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
|
@bjf-frz Following the adaptive-cost stability fix and main integration at 05891dd, here is the response to the routing reuse and adaptive-K validation requests. Routing reuse. Agreed: device-owned positions do not prevent using Data movement can be shared; the speculation-specific ownership still needs to retain (token, position) per request generation, the host-built ragged round plan, two submitted rounds sharing a slot, copy/ticket retirement and delayed cancellation/reuse. What the current policy estimates. The position index is an offset within a draft, and p_j estimates P(candidate j accepted | all preceding candidates accepted) over prior observations. Products of these estimates provide a population estimate of committed tokens, which the policy divides by estimated round time. This is a workload-level cost policy, not a prediction conditioned on the current request/model hidden state, and does not establish the optimal K for that request. A state-conditioned confidence head would be a different learned predictor requiring training labels/weights, deployment compatibility, calibration and overhead/correctness qualification. It is not implemented by the current statistics. Dedicated validation contract. The existing combined Graph/vectorization tables do not establish adaptive-K gains. A separate comparison should use:
This dedicated adaptive-K campaign has not been run in this update; there is no new claim of adaptive benefit over the best fixed K. The current adaptive default remains a separate maintainer decision; the numerical fix addresses the confirmed long-running failure regardless of that decision. |
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Signed-off-by: leo <92263211+liuyao0322@users.noreply.github.com>
Required human contributor answers
The contributor explicitly supplied these answers for this implementation.
1. What outcome are you trying to achieve?
Improve low-concurrency Ouro asynchronous speculative inference latency and throughput while preserving correctness.
2. What is the highest-leverage variable for that outcome?
CPU kernel-launch overhead and GPU idle time.
3. How will this work move that variable rather than create busywork?
Use CUDA Graph replay to batch GPU submission and reduce per-operator CPU dispatch on the asynchronous speculative execution path.
4. Who is the customer or user, and what do they actually need?
People deploying Ouro inference services need lower latency and higher throughput at low concurrency, with correctness preserved.
PR necessity and bugfix evidence
The contributor explicitly requested a PR. The existing engine rejects asynchronous speculation with CUDA Graphs, and its synchronous graph wrapper requires host-owned positions. Supporting the combination therefore requires a runtime change. This is an opt-in feature, not a standalone bugfix claim. Refs #43; builds on merged #48, #66 and #115.
Purpose
Support fixed-depth greedy Ouro asynchronous speculation with optional recurrent-core/LM-head CUDA Graphs, startup row buckets, derived verification coverage, vectorized rounds and fixed draft counts. Device-owned positions and KV frontiers remain authoritative; the scheduler reserves an upper bound, and two rounds can be submitted before CPU delivery.
Default behavior: async and synchronous speculation use the configured fixed K, subject to ordinary output/context/token-budget limits. Async scheduling, speculation and Graphs remain opt-in. Adaptive K is removed from this PR, including its policy module/tests, configuration/CLI, scheduler callback, calibration and per-round timing feedback. Its latest numerically stabilized implementation is preserved on
codex/adaptive-k-followupfor separate review; it is not part of this PR.Routing reuse: the two-column
(token, position)device slot table now uses the existingkernels.routing.gather_kernelandscatter_kernelwithH=2. Acceptance remains request-local; compute/copy events, slot ownership, cancellation and deferred retirement retain their existing contracts. A GPU regression checks each gather/scatter against indexing, including untouched slots, ragged K=4 rounds and repeated request IDs/slot reuse.Reuse the synchronous graph capture/replay mechanism through a device-metadata entry point that validates allocation identity without consulting or updating CPU written-position state. Clone graph outputs before workspace reuse. Preserve compact round results, copy-stream ordering, cancellation and ticket retirement. Drain the capture stream when warmup fails before an ordinary compute event can be recorded.
The asynchronous speculative core/head pre-capture a finite bucket set during engine construction. Preferred sizes include
min(max_num_seqs * (k + 1), max_num_batched_tokens)for every k≤K, clipped to the row limit, plus a halving ladder below the smallest. Preferred full-batch shapes take priority; if they exceed the entry budget, keep the largest preferred sizes including the maximum reachable row count. Add smaller halving buckets only with remaining entries. Never expand beyond the speculative workload’s reachable range. Rows above the largest captured size run eager and never trigger lazy speculative capture. Native bootstrap and synchronous/Huginn graphs retain their existing behavior.The default row limit is now unset (
None). Before memory planning and capture, async speculation resolves it tomax(128, min(max_num_batched_tokens, max_num_seqs * (K + 1))); native and synchronous paths resolve to 128. Explicit limits are preserved and warn if they cannot cover configured verification capacity. The old C43/K=2 and C26/K=4 thresholds describe an explicit 128-row cap, not current async-speculative defaults.A Triton KV scatter checks the device active-row count before reading write addresses or writing KV. Padding has zero context; block tables are no longer zeroed per replay, since inactive rows and active trailing pages are not read. Live requests use a device
(token, position)slot table; each round gathers/scatters it in batches, uploads its host-built row indices in one copy, and vectorizes acceptance over[n, K+1]. Row/result layout and lifetime ordering remain intact. Prelude, metadata staging, argmax, acceptance and state updates stay outside capture. Random sampling, adaptive exit and preemption remain unsupported in async speculation. The merged main supports speculation on PD decode stages; this integration preserves that interface, but NIXL validation was skipped in this environment.The alternative is to retain eager asynchronous speculation or use synchronous graph speculation. Neither exposes cross-round device-state submission together with graph replay. Whole-round capture is deferred because it also needs replay-safe request-state updates and warmup isolation.
One pre-existing expression in
test_gsm8k_accuracy.pyis reformatted for the installed Ruff version; its behavior is unchanged.Test Plan
Current head
980c092is based on05891dd, which merged main66920da. It narrows the PR to fixed-K Graph/vectorized execution and existing Triton routing reuse. Historical H100 tables below remain embedded unchanged; adaptive rows describe the earlierc3a97d3runtime and do not qualify the current fixed-K default. Its fixed K=4 row reports 0.98–1.08× native for the tested C64–256 conditions; no new end-to-end speedup is claimed for this routing change.Test Result
python -m pytest --run-gpu -q tests/test_cuda_graph.py tests/test_async_speculative.py tests/test_async_speculative_graphs.py tests/test_async_speculative_lifecycle.py tests/test_speculative.py -k 'not flash and not fa2 and not fa4' --tb=short: 215 passed, 2 skipped, 25 deselected, 71.77 s. Skips: the two-device and NIXL cases. FlashAttention variants were excluded. The new production-path routing test checks complete slot-pool equality against indexing for each call, including untouched slots and ragged K=4 rounds with repeated IDs/slot reuse.Previous integration checks at
05891ddfoundtest_cuda_split_prefill_survives_decode_graph_capture_and_replayfailing in its eager arm (num_preluded_tokens: expected 2, actual 4). The same failure was independently reproduced on upstream main66920da; it remains a disclosed validation limitation, rather than a claim of all-green GPU qualification.Historical validation (results scoped to the stated revisions)
Capture/GC fix (
e80f540, 2026-10-23)The combined-process capture failure is now reproduced and fixed. Native debugging on torch 2.14.1+cu130 identified
gc_collect_main → CUDAGraph::~CUDAGraph → reset → cudaGraphExecDestroy, returning error 900 (cudaErrorStreamCaptureUnsupported) when an old graph was collected during a new capture. A pure-PyTorch two-graph reproduction and a closed-engine reproduction both triggered it; collecting only an output tensor or collecting graphs before capture succeeded. Engine/PreemptionManager reference cycles explain why old graphs can survive until a later automatic collection._capturenow pauses automatic cyclic GC across graph capture and restores the caller's prior state after the final stream drain, including errors. Warmup retains the caller's setting; no full-heap collection is added per bucket. Replay is unchanged. Explicit user-triggered collection inside capture is outside this automatic-GC protection.Validation at this head:
866985cfails with the original CUDA capture error; the fix passes, checks replay results and confirms collection occurs outside capture.The historical failed runs below are retained as evidence of the original defect. This correctness fix does not rerun or extend the H100 inference performance qualification; those tables retain their original revisions.
Reviewer follow-up (
866985c, 2026-10-23)All three comments on
805010dare addressed:(4, 8, 12)instead of(32, 64, 128). Tight budgets retain the largest preferred shapes, including maximum reachable rows; the memory planner reads the same bucket set.last_event. Warmup/failure drains remain. Runtime lazy capture retains its device-wide barrier. PyTorch's graph context still performs its own device synchronization.Validation: full pre-commit passed. On two scheduler-reserved H20 GPUs, 29 CUDA-graph, 9 graph-specific, 17 lifecycle and 16 targeted async regression cases passed (71 total); one NIXL case skipped and FlashAttention variants were deselected. Coverage includes startup clone avoidance and normal output ownership, readiness/failure drains, tight budgets, inactive-row KV sentinels, K=2/4 high concurrency, ragged arrivals, adaptive policy and non-default device. Against
805010d, the new test overlay reproduced four bucket-budget failures and the startup clone failure (5 failed, 1 passed). The full async test-file run again aborted during capture/GC; that failure is retained and this follow-up does not resolve or claim a clean full GPU suite.Startup measurement: H20-3e, Ouro-1.4B
574fa66c, BF16, torch 2.14.1 / Triton 3.8.0; C=4, K=2, token budget 16, row limit 128, three entries. Measure core+coda graph preparation only, excluding model loading and after kernel warmup, with no unrelated active streams. One warmup per arm, six timed trials per arm in three A/B/C/C/B/A blocks:805010d)The synchronization-only ranges overlap, so its benefit is small in this workload. These timings do not qualify full cold start or inference throughput; the original H100 tables below remain pinned to their original revisions.
Main integration validation (
805010d)2c77360also reproduced CUDA capture errors (67 passed, 1 failed); the affected high-concurrency pair passed separately on both heads. The successful separate-file runs do not establish that this capture issue is resolved. No new performance or numerical-accuracy qualification is claimed for the merge.Historical H100 results (
c3a97d3, reviewed 2026-10-19)Reviewer 0z5a's approval confirms that the six cherry-picked
commits at
c3a97d3match the validatedpr127-high-concurrencybranch one forone and have an identical final tree. The 2026-10-17 measurements
therefore apply to this inference runtime. These are reviewer-run manual results,
not new locally reproduced measurements or CI results.
Setup: one H100 per engine, Ouro-1.4B
574fa66c, BF16, d=2/D=4, greedy Tritonattention, torch 2.12.1 / Triton 3.7.1, distinct natural-text prompts and
input/output 64/64. Each entry is end-to-end output throughput including the
completion tail, divided by full-depth native async Graph throughput with
max(128, C)rows. Values are medians of six timed trials per arm, withalternating blocks. Absolute tok/s are not supplied in this table. The original
review table is reproduced below; values above 1 mean higher throughput.
In the historical
c3a97d3table, “default” means adaptive drafts plus the automatically derived row limit;K=2/4 is the maximum proposal count, not a fixed count on every round. The
fixed Krow used the then-availabledraft_policy="fixed"option. Adaptive K and that option are removed from this PR. Adaptive K=2/4 reaches 1.04–1.20×native in these C64–256 cases. Fixed K=4 ranges from 0.98–1.08× and still
regresses in some cells, so this is not a general benefit from enabling Graphs
or increasing K. The older short-output regressions below describe earlier
runtimes, not the defaults at
c3a97d3.The reviewer also reports C16/C32 default throughput of 1.10–1.13× native,
input/output 512/128 at 1.13–1.22×, and 64/256 at 1.24×. Four DP replicas at
C128/K=2 reach 1.16× native per GPU with identical cross-GPU token digests.
For 320 staggered requests at 24 req/s, mixed lengths and half natural EOS,
adaptive and fixed K are within 1–2% of each other, both 12–15% faster than
speculative eager, with zero Graph fallbacks. These bounded workloads do not
establish arbitrary-arrival SLOs or a universal saturation threshold.
Reviewer validation at the matching tree: 239 passed, 1 skipped (a two-GPU
case), CPU suite passed, FP32 Graph/eager tokens identical for 64/64 requests;
cuda:1, two live engines and DP4 checks match their references. BF16 paddingcan still change numerical results and greedy continuations, as the historical
investigation below demonstrates. The approval discloses that six commits are
by this reviewer, so it is not an independent review of those commits; manual
GPU evidence does not replace CI.
Historical bucketed-path validation (
dfbb0f7, 2026-10-16)At
dfbb0f7, the bucketed implementation passed 165 tests with 13 FlashAttention selectorsdeselected and no skips on a two-H20 Slurm allocation. This includes
FP32/BF16 KV sentinels during zero-active warmup/capture, shrinking bucket reuse,
new counts after every bucket is populated, stale allocations, failure cleanup,
K=2/4 ragged arrivals, non-default device and two live engines. With identical
padded inputs, core outputs and logits are bitwise equal to eager computation;
the frozen tiny-model cases also preserve unpadded eager token IDs. Full
pre-commit (Ruff lint/format and CPU regression) passes.
A separate H20 BF16 diagnostic uses Ouro-1.4B
574fa66c, torch 2.14.1+cu130 /Triton 3.8.0, 32 distinct topic prompts clipped/repeated to 32–128 input tokens,
24–44 output limits, natural EOS enabled on half, 16 concurrent slots, a
512-token budget, 128 graph rows and 16 entries. Eight requests arrive at each
of steps 0/12/24/36. One warm workload precedes three unprofiled observations
per arm and K, with explicit identical KV capacity. It deliberately synchronizes
the speculative runner at each step to fix batch composition; it is a controlled
correctness/coverage diagnostic, not a production-arrival latency or throughput
qualification. Arms ran sequentially, not interleaved.
73cfe4e)73cfe4e)The counts repeat in all three observations. Bucketed Graph and eager with
the same padded computation agree for 32/32 requests at K=2 and K=4 in all
three trials. Against the original unpadded eager, only 29/32 agree at each K;
the differing continuations are reproduced by padded eager with graph replay
replaced by ordinary model calls. Thus this validates replay of padded
computation, not universal bitwise equivalence to unpadded execution. First
unpadded-eager divergences are at output indices 10/26/35 (K=2) and 18/26/33
(K=4). No end-to-end task-accuracy qualification is inferred from this check.
A warmed complete-round ops-only/full pair for exact-row and bucketed paths
retains shapes, CPU/GPU stacks and memory timelines. Individual CPU kernel
submissions drop from 7,266 to 185, with Graph launches increasing from 3 to 11;
the bucketed trace contains the masked KV scatter kernel. Full capture exports
are nonempty. These instrumented windows diagnose cache coverage and submission,
not unprofiled speed or an end-to-end native-relative gain. Broader asynchronous
arrival performance was not established by that diagnostic. The current H100 results above validate the later vectorized/adaptive implementation at
c3a97d3.First-divergence investigation (2026-10-16)
The 29/32 unpadded-eager agreement was investigated at frozen runtime
dfbb0f7on one H20-3e, torch 2.14.1+cu130 / Triton 3.8.0, with the same BF16Ouro-1.4B checkpoint, inputs and controlled arrival/collection boundaries as
above. Six traced generations (unpadded eager, bucketed Graph, padded eager at
K=2/4) reproduced every frozen output exactly. Each first-divergence
comparison below has the same complete prompt and generated prefix, and
uses a verification row actually committed by acceptance, not a rejected draft
continuation. Indices are zero-based. These are correctness diagnostics, not
performance measurements.
For example, K=2/request 22/index 10 compares token 8102 at 16.0000 versus
3982 at 15.9375 in unpadded eager. Padding makes both 16.0000, so
argmaxselects the lower token index 3982. K=4/request 15/index 33 compares 314 and
2390 tied at 23.5000 in unpadded eager, versus 23.5000/23.6250 with padding.
The first actual output divergence occurs at arrival-loop step 59 for K=2
(request 22), and step 34 for K=4 (request 15). Padded Graph and padded eager
logit vectors are bitwise equal at all six first-divergence rows. Recomputing
only the LM head at exact/padded sizes on the same hidden states produces
identical valid logits at these rows: the differing states already originate
inside the recurrent core, rather than the head or token-ID copy.
The numerical sources were isolated with identical hidden inputs and a cloned
pre-call KV state, then reproduced without any KV cache or CUDA Graph:
step 8, recurrence depth 1, is
model.layers.12.input_layernorm_2(zero-basedlayer/depth), for 6 valid rows padded to 8. Its 12,288 input elements and
their squares match bitwise. The FP32 variance differs by at most
1.862645149e-9; reciprocal RMS by4.768371582e-7. Sixteen normalized BF16elements then differ, with maximum weighted output delta
0.0009765625.At one coordinate, normalized FP32 values
1.3554686307907104versus1.35546875cast to BF161.3515625versus1.359375across a roundingboundary. Standalone RMSNorm reproduces both outputs in 20/20 repetitions.
removes that local normalization difference, but is not a full fix: complete
unpadded/padded generations still agree for only 28/32 (K=2) and 30/32 (K=4),
and the changed precision also changes the original baseline. In this
diagnostic-only control, the next isolated first operator is
model.layers.0.mlp.down_proj, 36 valid rows padded to 64. All 202,752 validinput elements and weights are identical; 121 of 73,728 output elements
differ, by at most
0.000244140625. Standalone BF16 linear calls reproduceboth outputs in 20/20 repetitions. An FP64 GEMM counterfactual removes the
local difference. These precision controls were not applied to production
and do not establish a model-accuracy improvement.
The KV and acceptance checks distinguish numerical values from storage errors:
positions, block tables, write blocks/offsets and context lengths with the
original call arguments at
dfbb0f7, including the then-zeroed inactive rowsand trailing table pages. Current replay no longer zeroes block tables;
zero context and active context bounds exclude unused entries from reads.
padded computation, but replace the masked kernel with ordinary indexed
writes of only the active rows. Every key/value cache byte and active output
matches the real Graph path. This checks untouched storage as well as valid
writes; the reference never writes padded rows.
rejected proposal, emitted target prefix, next token and next device
position on CPU. All match GPU state and the compact result buffer. Final
delivered outputs match the reconstructed device histories, after ordinary
EOS/discard handling. All six first-divergence cases have identical proposal
IDs across arms; changed target logits correctly explain any changed accepted
length. For K=2/request 21,
[37, 32]is rejected immediately against target32 in unpadded eager, but both proposals are accepted when the padded target
starts with 37; the corresponding frontier increments are 1 and 3.
Within these frozen cases, there is no evidence of wrong KV addressing,
inactive-row writes, or a broken acceptance/frontier rule. Valid KV values
can differ numerically because their upstream computation differs, which is
not the same as overwriting another row's storage. The demonstrated causes are
batch-shape-sensitive reductions and GEMM followed by BF16 rounding and near-tie
greedy decisions. This explains why padding requires its own numerical scope;
it does not promise unpadded bitwise equivalence or qualify general task
accuracy. Fixing only the variance reduction is insufficient. Jobs 378/379 and
382/383 retain operator-level fixtures and standalone controls.
Previous head
73cfe4e: warning/memory-planning validationpre-commit run --all-filescuda:1while the current device remainscuda:0; outputs matchGPU command for the previous
73cfe4esnapshot:OMP_NUM_THREADS=1 python -m pytest -q --run-gpu \ tests/test_async_speculative.py tests/test_async_speculative_lifecycle.py \ tests/test_cuda_graph.py tests/test_huginn.py \ -k 'not flash and not fa2 and not fa4'These were tiny-model correctness regressions at
73cfe4e. The H100, natural-EOS, default-versus-raised-limit and DP4 results subsequently supplied by the reviewer are reproduced in full below; they are not tests of the bucketed implementation.Historical validation at performance-evidence runtime
08d3f99Logs, JUnit, source/model hashes and performance artifacts were retained for validation. Raw artifacts and diagnostic measurement scripts are not included in this PR; the linked report records the configuration, methodology, results and limits.
BF16 measurements and source attribution
One Slurm-reserved H20-3e, driver 580.105.08, PyTorch 2.14.1+cu130, Triton 3.8.0, Ouro-1.4B revision
574fa66c, BF16 weights/activations/KV, d=2/D=4/K=2, greedy LAST_EXITED, ignore EOS. Five warmed alternating paired observations per mode/condition, with profiling disabled. Decode excludes the first output token and ends after GPU completion. Inputs are eight distinct texts repeated/clipped to fixed lengths; output length is 64.Historical fixed K=2 at
08d3f99; C1/8 fits that runtime's 128-row default. Historical adaptive K=2/4 results are reported separately above.The versioned report and annotated timeline include observed variability, controls, reproduction and limits. All warmed speculative Graph trials reused captures without fallback. Isolated-engine peak allocated memory increased by 96.5/102.8 MiB at concurrency 1/8 (input/output 64/64).
In an independent C1 input/output 64/8 ops-only decode window, CPU runtime/driver kernel launches fell from 27,844 to 181 plus 27 Graph launches; GPU kernel count stayed approximately unchanged. GPU-active interval union was 65.19/67.48 ms, while the remaining decode interval fell from 372.34 to 14.98 ms. D2D copies increased 48→198. Host synchronization time rose because the CPU submitted faster and waited for GPU completion; this is dispatch reduction, not removal of all waiting. Profile durations are diagnostic, not the throughput evidence above.
C1 input/output 64/8 ops-only decode observations (profiled diagnostics):
Full profiles enable shapes, memory and stacks. Both bounded one-round captures have verified nonempty CPU/GPU stacks and memory timelines; the initial oversized whole-request full trace is preserved. Core/head capture removes most individual launches while preserving the existing state machine. Remaining metadata/acceptance launches, copies, cold capture costs and broader workloads remain explicit follow-up work.
A separate production-loader/create_app loopback HTTP test ran five warmed alternating paired SSE requests per mode, concurrency 1, the same seven-token prompt and 64-token output. All ten completed with
[DONE], identical text and exact usage. Median E2E fell from 2711.7 to 824.4 ms; TTFT was 55.4 vs 54.4 ms. E2E p95 was 2728.2/824.8 ms, descriptive with five samples per mode. Each response had 27 choice events/64 tokens; these intervals are not per-token ITL. This is a bounded loopback result, not a saturation or natural-EOS qualification.Code-quality self-review applied R02/R12/R16/R17/R21/R23/R25: GPU progress stays runner-owned, logical allocation identity stays cache-owned, failed capture drains its stream, cloned outputs outlive replay, and graph-off/fallback/lifecycle behavior is covered. No unresolved actionable ownership or correctness finding was identified within this scope.
Higher-concurrency qualification
The extended report and scaling plot add 150 unprofiled paired trials, concurrency 16/32/64 at input 64/512, concurrency 128/256 at input 64, and C64 longer-output cohorts (64/512 and 512/256 input/output). All output token IDs matched the matched native baseline, and speculative eager/Graph pairs matched exactly. All warmed speculative trials had zero new captures and zero fallback.
Historical tuned coverage at
08d3f99: fixed K=2 and a shared row limit ofmax(128, 3*C). At that runtime's 128-row default, verify core/coda fell back from C43, and drafting/native core from C129. Current async-speculative defaults derive the limit instead. This table does not measure the current adaptive policy or non-lockstep latency tails.Input lengths 64/512; output length 64:
Graph still accelerates speculation relative to eager by 1.16x at C256, but the native Graph baseline is faster there. At input 512/C64, Graph speculation reaches 2182.76 tok/s (1.42x over speculative eager, 1.12x over native Graph). C64 longer-output runs retain approximately 12% native-relative gains. Acceptance stays about 93.3–93.5%, with all counters retained per timed trial. These measurements do not define a universal enable/disable threshold or establish saturation.
All modes within each condition share GPU allocation, model, KV capacity and scheduling controls. The row limit is
max(128, 3*C); caches retain the default 16-entry limit. C128/256 raise the common admission scan limit to preserve the all-first-token boundary. An initial C256 warmup with the default scan failed that setup assertion before timing; its diagnostics are retained and both extension conditions were rerun. C1–64 use one physical H20; C128/256 use another H20 on the same host, so the figure separates those points and makes no strict cross-card scaling claim.The C64/input-512 complete-round ops/full pair retains nonempty shapes, CPU/GPU stacks and memory artifacts. Individual runtime/driver kernel submissions fall 12,040→981 plus 11 Graph launches; GPU-active time is 76.38/76.30 ms and the remaining profile interval 129.80/12.92 ms. This is diagnostic attribution, not timed throughput. Isolated warmed decode allocation overhead is 162.4–174.5 MiB for the two C64 long cohorts and 282.8 MiB at C256/input-output 64/64.
800 measured loopback SSE requests at C16/64 completed with valid
[DONE], length and usage. Each concurrency's eager/Graph text matches its native full-depth control. C64 p95 E2E falls 3191.4→1550.6 ms (speculative eager→Graph), with p95 TTFT 222.1→100.9 ms; five finite bursts per mode are not a steady-state capacity test. The C16/C64 continuation split is reproduced by native async Graph at identical input prefixes, first diverging at output index 29 (margin 0.125 versus a tied BF16 top pair); it is not a new speculative eager/Graph disagreement. Cross-batch bitwise invariance remains unqualified.Longer-output decode cohorts
These use the same K=2 and tuned graph coverage. Values are medians of five paired observations. Longer output also changes context lengths, so this is not a duration-only comparison with the 64-output table.
Isolated warmed decode memory
Per-engine peak allocated memory after a complete warm cohort; initialization/prefill temporary peaks are excluded. Graph extra MiB is relative to speculative eager. These use explicit KV capacities at runtime
08d3f99, before the automatic-planning change.Loopback HTTP bursts
Production
load_engine/create_app, seven-token prompt, 64 output tokens, ignore EOS; five warmed alternating paired bursts per mode at C16/64. Shared graph row limits are 128/192 respectively, with 16 entries per cache. All 800 measured requests completed with valid[DONE], length and usage, with zero measured speculative captures/fallbacks.Throughput includes prefill and the complete burst and is the median of five observations. Latency percentiles pool 80 requests per mode at C16 and 320 at C64. Choice events contain multiple tokens, so event gaps are not per-token ITL. These finite bursts do not establish steady-state capacity or arbitrary-arrival latency tails.
No runtime code changed for the higher-concurrency measurements: their runtime hashes match
08d3f99. Subsequent review changes have separate evidence: warnings/memory planning at73cfe4e, bucketing atdfbb0f7, and vectorized/adaptive rounds atc3a97d3. This documentation/plot delta was checked for whitespace, links, raw-result coverage, source identities and successful producer completion.Previous review follow-up at
73cfe4emax_num_seqs * (K + 1)exceeds coverage. State default-limit fallbacks beside the tables rather than suggesting the tuned results are defaults.cuda:1while the current device remainscuda:0. The existing ragged correctness matrix also covers K=4.Historical validation at
73cfe4e: full pre-commit passed (Ruff lint/format and CPU regression). An isolated source snapshot matching that revision passed 108 tests, 13 deselected on a Slurm two-H20 allocation: async speculation/lifecycle, CUDA Graph and Huginn tests excluding FlashAttention selectors. The non-default-device test ran with both allocated GPUs; it was not skipped. No new real-checkpoint throughput claim is made for this head.The exact-row cache limitation above describes
73cfe4e; it is superseded by the current bucketed implementation described in Purpose and Test Result. The full-logits clone and flat per-pool headroom remain follow-ups. Reviewer H100 results are reproduced below, attributed to their measured revision.Historical H100 exact-row review measurements (2026-10-16)
These results were supplied by reviewer 0z5a, measured at
9eb01bb. The reviewer confirmed applicability to73cfe4e: its warnings do not change replay/fallback selection, and the experiments use explicit KV sizes. These exact-row measurements are historical. Separate post-bucketing, vectorized/adaptive H100 results forc3a97d3are reproduced above; the two sets describe different runtimes. Raw JSON/logs are available from the reviewer on request; these are attributed external results, not a locally reproduced benchmark.Setup: one H100 per engine, torch 2.12.1 / Triton 3.7.1, Ouro-1.4B
574fa66c, BF16, d=2/D=4, greedy Triton attention. Distinct natural-text slices replace the eight repeated H20 prompts; observed acceptance is 72–91% rather than about 93%. Each arm runs two alternating blocks (ABCD/DCBA), warmed once per exact workload and timed three times per block: six timed observations per arm, reported as medians. Native async Graph is full depth (ouro_delayed, threshold 1, four loops).For this historical speculative Graph runtime, default coverage was 128 rows; raised coverage is
(K+1)*C. Native coverage ismax(128, C). In the original table below, “window” starts after all requests receive their first token and stops before any finishes. The speculative eager/default/raised columns are window throughput too. “Total” covers submit through last finish, including prefill and the uneven completion tail. It is not the H20 decode-to-GPU-completion window, so those absolute throughput values should not be combined.Lockstep bursts, input/output 64/64 (tok/s)
The original exact-row implementation wins over speculative eager in many window comparisons, but short-output end-to-end throughput is 0.92–1.00× native at K=2, and 0.71–0.81× at K=4. Even raised-limit lockstep trials have warmed fallbacks (core 24–148, coda 1–60 per trial) as tail shapes fill the caches. These results do not support a general latency/throughput improvement over native async Graph. K=2 is the usage example; K=4 remains supported and benchmarked here, with the regression disclosed.
Longer-sequence end-to-end results, K=2:
DP4, C128, input/output 64/64, K=2: native 2791–2814 versus speculative Graph raised 2697–2760 tok/s per GPU, approximately 0.97×.
Staggered arrivals: original exact-row caches
320 requests at 24 requests/s, input lengths 32–1021, output lengths 16–375, about half natural EOS,
max_num_seqs=128; warm passes. Values below are the reviewer's original ranges. TPOT uses the reviewer's reporting convention, not per-token SSE timestamps.With 16 entries, 89–93% (K=2) and 77–82% (K=4) of warmed speculative core calls fall back; with 64 entries, 44–58% still fall back. Graph adds only about 1–6% over speculative eager with 16 entries. Speculation itself beats native in this workload, so that gain must not be attributed entirely to Graph replay. Native's exact-row cache also falls back on 83–93% of its core calls. These measurements motivated the bucketed path; the current H100 staggered-arrival results above report zero fallbacks with derived coverage and exact full-batch buckets.
Correctness reported by the reviewer: H100 GPU selections passed (212 tests at head, 183 at base, 16 FlashAttention cases deselected, no skips). Speculative Graph/eager token IDs agree in all lockstep conditions, including K=2/4 and longer sequences. Under ragged arrivals, agreement is 320/320 at default limits for K=2/4 and with 64 entries for K=4, but 305/320 with 64 entries at K=2. Batch composition can vary with asynchronous event collection; this is a hypothesis for the mismatch, not a proven explanation. Speculation/native also diverges on these prompts (eager does too), so native equality in the H20 report is limited to its frozen inputs. Non-default device, interleaved two-engine and DP4 checks passed against their reference outputs.
Remaining review notes
Adaptive K is separate follow-up work, including acceptance estimation, cost calibration and comparison against the best fixed K. No adaptive default is proposed by this PR. The retained
_retired_slot_valuesgrowth storage remains a non-blocking follow-up. Earlier approvals and benchmark claims apply only to their recorded revisions; the routing/scope change needs review. GPU validation is manual, not official CI evidence.AI assistance: OpenAI Codex assisted with implementation, tests and documentation.