Repository navigation
[feat] [speculative] implement pipelined self-speculative decoding for looped transformers - #121
Conversation
…integration tests
|
If the review is approved, I will proceed with the implementation. Please also let me know if you have any suggestions, and I will implement them. @hsliuustc0106 @bjf-frz |
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
Review: PR #121 — Pipelined Self-Speculative Decoding
Review status: Request changes (correctness)
Thanks for implementing the pipelined scheduler and CDB executor. The direction is interesting, and the initial single-request speedup is encouraging. My main concern, however, is branch/KV correctness in the actual execution path, not merely API contracts or the breadth of the benchmark matrix.
Running longer contexts or higher concurrency would not, by itself, prove that the underlying issue is fixed. We need a branch-aware cache invariant, an explanation of how the implementation enforces it, and deterministic tests that exercise the conflicting histories.
P0 — Speculative branches do not appear to have isolated KV histories
The scheduler tracks branch / parent_branch, but PipelinedSpeculativeRunner.forward() builds KV batches using the same request_id and a position derived only from lane.index and recurrent depth. The KV address preparation has no branch identity. Checking for duplicate write addresses within one batch does not protect a live branch's history from writes by another branch, either in that batch or in an earlier cycle.
A concrete path to examine with S=4, K=1, second_step=2:
- At position
p,q1(A)proposesBforp+1;q1(B)later proposesDforp+2. q2(A)proposes a distinct candidateC != B, also forp+1.- In the following cycle, the lanes for
C(depth=0, position=p+1) andD(depth=0, position=p+2) may execute in the same CDB batch. - The model writes all rows' KV before attending. Therefore,
Dcan readC's newly written KV atp+1, whereasDbelongs to theBancestry and requiresB's KV. The duplicate-write guard would not detect this because the rows write different positions.
This is a possible read-after-overwrite across branch histories, not simply an allocation or API-semantics issue. Even if the final decision selects B, a descendant may already have consumed C's KV, so pruning the losing branch afterward cannot repair its hidden state.
Could you either implement branch-aware/versioned speculative KV (with explicit commit/rollback semantics), or document and test an equivalent mechanism that prevents conflicting writes/reads and safely recomputes affected descendants? In particular, please state the invariant connecting a live lane's ancestry to every KV position and depth it may attend.
P0 — Pruning states is not sufficient to reconcile the KV cache
retain_branches() currently retains/recycles hidden-state slots only. truncate_suffix() is called at generation completion (and in the first-token stop path), not when a speculative branch is rejected. Thus the advertised rollback behavior is not yet evident in the implementation.
The key requirement is stronger than eventually truncating the suffix: after each acceptance/rejection, no surviving lane may read KV or written-position metadata left by an incompatible branch. Please show how this is enforced before the next attention operation, including when the selected branch changes after one or more speculative descendants have already run.
P0 — Add directed executor-level tests, not just longer random generations
The current end-to-end test uses only S=4, K=1, second_step=None, three short CPU/FP32 generations, and compares final token IDs. The logical mock tests exercise q2/rejection decisions but cannot expose shared physical KV corruption. The reported real-model example has 72 accepted / 0 rejected, so it does not validate rejection handling either.
I would consider the following the minimum correctness evidence before calling this lossless:
- Force
q1 != q2and keep both branches live, including theB/C/Dinterleaving above; test q1 acceptance, q2 acceptance, and rejection of both proposals using the real CDB executor/cache, not only a mock scheduler. - Inject deterministic rejection patterns (all reject, alternating accept/reject, consecutive rollbacks), including a KV-block boundary and EOS/length termination with outstanding descendants.
- Compare at the first affected cycle, not only at the end: branch ancestry, KV ownership/written metadata, per-depth KV and hidden states where comparable, and final logits against a teacher-forced serial reference on the same token prefix. In FP32, use this to diagnose structural errors; in BF16, report tolerances, top-1/top-2 margins, and the first divergent token.
- Assert that pruning/slot reuse never leaves a surviving lane dependent on another branch's KV, and that a corrected branch can continue generating consistently for several more tokens.
A targeted small test that deterministically exercises the conflicting branch histories is more convincing here than a large benchmark that happens not to trigger them.
P1 — Separate correctness claims from scalability/performance evidence
Long-context and high-concurrency benchmarks are useful after the cache lineage/rollback problem is addressed, for capacity, latency, and performance regressions. They are not a substitute for the P0 invariant and adversarial tests; high concurrency is also outside the current single-request runner's serving integration scope.
Separately, DecisionReadback.read() synchronizes on the CUDA event immediately after staging in the present forward path. This can still be a compact asynchronous copy, but does not demonstrate useful host/device overlap. It would be helpful to profile the synchronization cost and distinguish copy asynchrony from actual pipeline overlap when reporting the speedup.
Suggested merge bar: first establish branch-aware KV correctness and deterministic q1/q2/rejection coverage; then use real-model BF16 validation to support the scope of the lossless claim. Broader context/concurrency sweeps are follow-up performance evidence, not proof of the fix.
If larger-scale GPU validation becomes useful later, I can help run verification on an H100. Feel free to ping me.
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
Thanks for the PR. The single-proposal pipeline looks right, but tree branching via second_step is not lossless yet, so I'm requesting changes for that part. Below: a reproducer, then supplementary H100 measurements of the runner on Ouro-1.4B.
1. second_step branching breaks greedy equivalence
Everything here runs on the PR head (8961c6a), on CPU in FP32. It reuses the PR's own test helpers (tests/test_pipelined_speculative.py: make_cache, native_autoregressive_decode) with tiny_ouro_config(total_ut_steps=4).
| schedule | seeds tried | result |
|---|---|---|
PipelineSchedule(4, 1, None) |
40 seeds, 40 tokens | 40/40 token-identical to AR |
PipelineSchedule(4, 1, 2) |
42, 999, 71 (the seeds in the PR's test), 0–3; prompt [1, 2, 3], 12 tokens |
every run raises ValueError: a write batch cannot contain duplicate request/depth/position addresses |
PipelineSchedule(4, 1, 3) |
40 seeds, 40 tokens | no error, but 2/40 silently differ from AR. Example: seed 9, prompt [3, 2, 3, 5], first divergence at output index 22 |
Minimal repro, run from the repository root on the PR branch:
import torch
from tests.helpers import tiny_ouro_config
from tests.test_pipelined_speculative import make_cache, native_autoregressive_decode
from vllm_rlt.models import OuroForCausalLM
from vllm_rlt.worker.pipelined_speculative import PipelinedSpeculativeRunner
from vllm_rlt.worker.speculative_pipeline import PipelineSchedule
for second, seed, prompt, n in ((2, 42, [1, 2, 3], 12), (3, 9, [3, 2, 3, 5], 40)):
torch.manual_seed(seed)
cfg = tiny_ouro_config(total_ut_steps=4)
model = OuroForCausalLM(cfg)
ref = native_autoregressive_decode(model, make_cache(cfg, blocks=256), prompt, n)
runner = PipelinedSpeculativeRunner(model, make_cache(cfg, blocks=256), PipelineSchedule(4, 1, second))
try:
print(second, runner.generate("r", prompt, max_new_tokens=n) == ref)
except ValueError as e:
print(second, e)Here is my reading of the cause; please correct me if the tree logic intends something else. The paged cache holds exactly one KV slot per (request, position, depth). The q1 and q2 branches both spawn a lane for the same next position, so they share those slots.
- X=2: both lanes sit at stage 0 in the same cycle, so the batch writes the same address twice.
- X=3: the q2 lane writes the depth-0 slot of position p+1 two cycles later. By then the q1 branch's descendants (p+2, p+3) still read that slot, and they read the alternative token's KV.
Branches would need their own KV storage (for example, copy-on-branch into separate pages), or second_step should be rejected until they have it. It would also help to add second_step cases and a longer random-model stress test to test_pipelined_speculative.py. The current lossless test only covers second_step=None with 5–8 tokens, and the 40-token stress above is what exposes the silent case.
A smaller point: retain_branches says it rolls back the uncommitted KV suffix, but it only recycles state slots. For second_step=None, correctness then depends on rejected positions being rewritten before any later row reads them. As far as I can tell that holds today, but it should be either documented or asserted.
2. H100 measurements (Ouro-1.4B, BF16, Triton, eager, C=1, greedy)
Setup:
- 16 GSM8K 3-shot prompts of 294–354 tokens, 256 new tokens each,
ignore_eos. - Prefill is untimed. Decode tokens/s is the median of 3 interleaved reps.
- Every mode runs through the same
KVCacheManagerandmodel.recurrent_preparedAPI, so the AR baseline pays the same per-call overhead. - Token equality is checked against AR on the same prompts, in both FP32 and BF16.
| mode | GSM8K BF16 tok/s (×AR) | MATH-500 BF16 tok/s (×AR) | core calls / token | FP32: token-identical to AR |
|---|---|---|---|---|
| AR, same cache API | 22.3 (1.00) | 22.4 (1.00) | 4.00 | — |
PipelineSchedule(4, 1) |
66.6 (2.99) | 65.9 (2.94) | 1.28 | 16/16 |
PipelineSchedule(4, 2) |
41.5 (1.86) | 42.0 (1.87) | 2.10 | 16/16 |
PipelineSchedule(4, 1, 3) |
69.3 (3.11) | 70.4 (3.14) | 1.23 | 8/16; first divergence as early as token 9 |
PipelineSchedule(4, 1, 2) |
— | — | — | raises the duplicate-address ValueError on the real model too |
The GSM8K rows use the 16 prompts above (FP32 runs use the same prompts). The MATH-500 rows use 8 Minerva 4-shot prompts of 721–899 tokens, also with 256 new tokens. Rep-to-rep spread is ≤3%, except for one MATH warm-up rep at −8.5%.
For context, the native engine on the same prompts gives:
| native engine (same 16 GSM8K prompts × 256) | tok/s | ×eager AR | greedy-identical to eager AR (BF16) |
|---|---|---|---|
| AR, eager | 21.9 | 1.00 | — |
| AR, CUDA graphs | 45.5 | 2.07 | 16/16 |
| sync self-spec d=1, K=4, eager | 45.8 | 2.09 | 7/16 |
| async self-spec d=1, K=4 | 46.0 | 2.10 | 7/16 |
| sync self-spec d=1, K=4, CUDA graphs | 92.9 | 4.23 | 7/16 |
Notes:
- The FP32 column is the correctness signal. In BF16, even the exact modes (
step_size=1/2here and the natived1k4) match AR on only 5–7 of 16 prompts because of batch-shape numerics; the FP32 runs remove that noise. In FP32,step_size=1andstep_size=2are exact, whilesecond_step=3diverges on half of the prompts. That matches the CPU reproducer above. - Throughput: the eager pipelined runner is about 3× eager AR, about 1.46× CUDA-graph AR, and about 1.45× the existing eager/async
d1k4self-speculation. However, the existingd1k4with CUDA graphs is already at 92.9 tok/s on the same prompts, which is faster than the eager pipeline (66.6). The comparison that will decide the value of this PR is therefore Phase 3 (graphs) against graph-captured self-spec. - An independent prototype of the same scheme (one mixed-depth core call per step, draft from depth 1) reaches the same 1.28 core calls per committed token. In FP32 it is token-identical to AR on all 16 prompts. That independently supports the
second_step=Nonedesign.
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Maintainer self-review of head 8961c6ae9084f42f4d4ff5f22cb69217b44586f5. Verdict: FAIL — blocking correctness defect, not ready to merge.
Blocker (reproduced): the documented two-question (second_step) path crashes on a real KV cache. Minimal reproduction: PipelineSchedule(stages=4, step_size=1, second_step=2) + tiny_ouro_config(total_ut_steps=4) + runner.generate("r", [1,2,3], 8) raises ValueError: duplicate KV write addresses for request/depth/position rows, stack vllm_rlt/worker/pipelined_speculative.py:128 → vllm_rlt/core/kv_cache_manager.py:587 → vllm_rlt/worker/metadata.py:153. The failing batch has depths [3,2,1,1,0,0,0], positions [3,4,4,5,5,5,6], duplicate triple (req,0,5): q1 and q2 create two sibling lanes that write KV at the same depth/position. The q2 path is only covered by a mock test (tests/test_speculative_pipeline.py:50-124, no KV); the only real-model test uses second_step=None (tests/test_pipelined_speculative.py:73). The schedule's own reproducer is the one path that fails. Fix q2 (single lane per position) and add a real-cache test, or remove q2 from the documented/claimed surface and make the unsupported configuration raise.
The change is not integrated with the engine. The PR adds six files and modifies none: vllm_rlt/worker/pipelined_speculative.py, speculative_pipeline.py, speculative_decisions.py plus three test files. Nothing in StageEngine, ModelRunner, scheduling, KV leases, prefix pinning, CUDA Graphs or PD imports or calls them, so the "natively integrates with the CDB engine" statement is not supported by this diff, and the 2.62× serving number cannot be produced by merged code.
Other correctness gaps found:
PipelineSchedule.stagesis not validated againstmodel.config.total_ut_steps(pipelined_speculative.py:39-58);stages=2on a 4-loop model is accepted and returns wrong tokens (native[9,9,9,9,9]vs[20,41,37,37,41]).- The acceptance counter is inconsistent with the claim:
stats.accepted_proposalscan increment at most once per committed token (pipelined_speculative.py:169-173), so "72 accepted / 0 rejected" for 64 output tokens is impossible. retain_branchesdocstring says it rolls back KV (:195-202); the code reclaims slots only (:229-231) and truncates KV at the end.speculative_pipeline.py:60claims Huginn/Raven support, but the runner callsprelude/prepared/coda(:84/:92/:95) while Huginn exposes only*_prepared(models/huginn.py:251/278/295).- Greedy-only, no RNG, no default EOS, no chunked prefill; single-request scheduling only, so "continuous batching" is not delivered either.
Tests run (CPU, torch 2.8.0, no CUDA): full pytest -q -m "not gpu" on the merged tree (main + this PR) → 793 passed, 35 skipped, 322 deselected; the three new test files → 13 passed, 1 skipped (CUDA readback). ruff check clean. Not run: any GPU path, the RTX 5090 benchmark (its script is not in the PR), real-model q2/step_size>1 correctness.
Path to merge: fix the q2 KV collision with a real-cache regression test, validate stages against the model, correct the counter/docstrings, and either wire the runner into the engine or retitle/descope the PR to an experimental standalone prototype with the claims removed.
|
Follow-up: the maintainer's re-verification criteria for this PR. I ran four read-only probes against the unfixed head 1. q2 real-KV correctness. Currently 18 failing cases: 2. 3. Acceptance counter and performance claims. Required invariants: 4. Huginn claim. Currently construction succeeds and 5. Claims and limits to converge in the same PR. Without an engine/CLI integration path, do not claim "natively integrates with the CDB engine" or continuous batching; either wire one real path or label the code an experimental building block. Fix the CPU-side gates for the new head: full |
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Changes requested: the q2 pipeline branch is not correct against a real KV cache
The idea (overlap the next recurrent step with reading back the previous decision) is worth pursuing, but in its current form this PR cannot be merged: one of its documented code paths crashes on a real KV cache, and the module is not reachable from the engine. Details below, all reproduced on this branch.
Blocking: second_step (q2) produces duplicate KV writes
Minimal reproduction with the repository's tiny-model helpers:
schedule = PipelineSchedule(stages=4, step_size=1, second_step=2)
runner = <PipelineRunner built on tiny_ouro_config(total_ut_steps=4)>
runner.generate("r", [1, 2, 3], 8)ValueError: duplicate KV write addresses for request/depth/position rows
at vllm_rlt/worker/pipelined_speculative.py:128 (cache._prepare_batch)
-> vllm_rlt/core/kv_cache_manager.py:587
-> vllm_rlt/worker/metadata.py:153
Failing batch arguments: ids all "r", depths=[3,2,1,1,0,0,0], positions=[3,4,4,5,5,5,6]; the duplicate tuple is (request="r", depth=0, position=5). q2 spawns a sibling lane at the same index, and step 3 of the cycle calls cache._prepare_batch(ids, depths, positions) for all lanes at once, so both siblings write the same depth/position in the same step.
This is not a small patch: the current design writes KV for every live lane on every cycle, and q2 therefore needs either per-branch KV staging or a deferred write until the surviving branch is committed. The existing q2 coverage is mock-only (tests/test_speculative_pipeline.py:50-124 uses a MockExecutor that never touches KV), and the real-model integration test only runs second_step=None (tests/test_pipelined_speculative.py:73). Please add a real-cache test that crosses the q2 path before re-requesting review.
Blocking: the module is not integrated with the engine
git grep finds the three new modules referenced only from tests; the PR itself lists CLI/serving integration as future work. As a result the reported 2.62x serving gain cannot be produced by the code being merged, and the claim of native CDB-engine integration does not hold: the implementation bypasses StageEngine, ModelRunner, the scheduler, KV leases/pins, CUDA graphs and PD, and holds a single self.request_id (no continuous batching).
Also:
PipelinedSpeculativeRunnerdoes not validateschedule.stagesagainstmodel.config.total_ut_steps(pipelined_speculative.py:39-58); building it withstages=2for a 4-loop model does not raise, and the output is wrong ([9,9,9,9,9]versus[20,41,37,37,41]from the engine path).StageEnginedoes validate this class of invariant, so the module should too.- The 2.62x measurement references a script that is not in the PR, and the reported "72 accepted / 0 rejected" is not reproducible from the counters in this diff (
accepted_proposalscan increase by at most one per committed token,pipelined_speculative.py:169-173; 64 committed tokens cannot reach 72). retain_branchesdocstring says it rolls back the uncommitted KV suffix (:195-202), but the implementation only recycles the state slot; KV is truncated at the end ofgenerate(:229-231).speculative_pipeline.py:60claims Huginn/Raven support, but the runner callsmodel.prelude/recurrent_prepared/coda(:84/:92/:95) while Huginn only exposes the*_preparedvariants (models/huginn.py:251/278/295).
What I could verify
On a synthesized merge tree (main + the six new files), the CPU suite is green: 793 passed / 35 skipped (780 baseline + 13 new cases), and ruff check passes. The failure above is on the GPU-free path as well; there is nothing environment-specific about it.
Recommendation: keep this as a prototype branch. Re-request review after (a) the q2 KV design plus a real-cache test, (b) validation of stages against the model, and (c) either engine integration or a clearly reduced claim with a reproducible benchmark script in the PR.
with regression test
…_ut_steps at construction
cs-ai-li
left a comment
cs-ai-li
left a comment
There was a problem hiding this comment.
1. Commit Message
fix(speculative): reject unsupported second_step (q2) at construction with regression
test
PipelinedSpeculativeRunner currently performs synchronous physical KV writes
for all active lanes on every cycle. Multi-proposal schedules with second_step
(q2) spawn sibling lanes at the same token index, causing collision error:
'ValueError: duplicate KV write addresses for request/depth/position rows'.
To prevent unhandled runtime crashes:
- Explicitly raise NotImplementedError in PipelinedSpeculativeRunner.init when schedule.second_step is not None.
- Clarify in docstrings that second_step is currently logical-only and not supported with physical KV cache execution.
- Add regression test test_second_step_rejected_at_construction across second_step=(2, 3) ensuring clean rejection at construction time.
──────
2. PR Description
Summary
Fixes the blocking KV write collision crash reported on real models when using
second_step (q2).
Root Cause
PipelinedSpeculativeRunner.forward() writes KV cache for all live lanes across recurrent steps in a single cache._prepare_batch() call with for_write=True. When second_step (q2) spawns sibling lanes for the same token position, both lanes attempt to write to (request_id, depth=0, position) concurrently, raising ValueError: duplicate KV write addresses for request/depth/position rows.
Changes
1. Explicit Construction Guard: In PipelinedSpeculativeRunner.__init__, raise
NotImplementedError if schedule.second_step is not None, directing users to
second_step=None.
2. Docstring Clarification: Updated docstrings in pipelined_speculative.py and
speculative_pipeline.py to clarify that second_step is only supported in the pure-
logical scheduler, not the physical runner.
3. Regression Tests: Added test_second_step_rejected_at_construction in
tests/test_pipelined_speculative.py to verify construction-time rejection.
Validation
- pytest tests/test_pipelined_speculative.py: 5 passed in 0.32s
- Existing single-proposal (second_step=None) lossless decode tests remain 100% green.
- ruff check and ruff format clean.
──────
cs-ai-li
left a comment
cs-ai-li
left a comment
There was a problem hiding this comment.
1. Commit Message
fix(speculative): validate schedule.stages against model.config.total_ut_steps at
construction
Previously, passing an incompatible schedule.stages (e.g. stages=2 on a 4-loop model) was accepted without warning, producing silently wrong tokens ([20,41,37,37,41] vs native [9,9,9,9,9]), while stages=8 failed late during execution.
Following the validation pattern in StageEngine:
- Validate schedule.stages == model.config.total_ut_steps in PipelinedSpeculativeRunner.init.
- Raise ValueError naming both fields if a mismatch is detected.
- Add regression test test_stages_mismatch_rejected_at_construction covering
both under-staged (stages=2) and over-staged (stages=8) configurations.
──────
2. PR Description
Summary
Validates schedule.stages against model.config.total_ut_steps at construction time to prevent silent corruption of decoded tokens.
Problem
When PipelineSchedule.stages did not match the recurrent model's expected loop count:
- stages=2 on a 4-loop model ran to completion but silently output incorrect tokens.
- stages=8 failed during execution after wasting compute.
Changes
1. Construction Validation: Added strict validation inPipelinedSpeculativeRunner.__init__ verifying schedule.stages ==model.config.total_ut_steps. Raises ValueError:schedule.stages (X) must match model.config.total_ut_steps (Y).
2. Regression Tests: Added test_stages_mismatch_rejected_at_construction in
tests/test_pipelined_speculative.py testing stages=2 and stages=8 on a 4-loop Ouro
configuration.
Validation
- pytest tests/test_pipelined_speculative.py: 7 passed in 0.34s
- Valid schedules (stages=4) remain 100% lossless against native decoding.
- Code style: clean on ruff check and ruff format.
cs-ai-li
left a comment
cs-ai-li
left a comment
There was a problem hiding this comment.
- Commit Message
fix(speculative): validate model interface at construction and remove Huginn claim
PipelinedSpeculativeRunner requires model methods 'prelude', 'recurrent_prepared' , and 'coda' (used by Ouro-class models). HuginnForCausalLM exposes only *_prepared variants, causing generate() to raise an unhandled AttributeError at runtime,while PipelineSchedule docstring claimed Huginn/Raven support.
To guard against runtime crashes and align documentation with implementation:
- Validate that model implements 'prelude', 'recurrent_prepared', and 'coda' in
PipelinedSpeculativeRunner.init, raising NotImplementedError at construction.
- Remove unsupported '32 for Huginn/Raven' claim from PipelineSchedule docstring.
- Add regression test test_huginn_rejected_at_construction verifying clean
construction-time rejection when a Huginn model is passed.
──────
2. PR Description
Summary
Guards against unsupported model architectures at construction time and removes misleading claims regarding Huginn/Raven support.
Problem
- speculative_pipeline.py claimed support for Huginn/Raven models ("32 for
Huginn/Raven").
- However, PipelinedSpeculativeRunner relies on model.prelude and model.coda,
which HuginnForCausalLM does not implement (it exposes prelude_prepared and
coda_prepared).
- Instantiating with Huginn succeeded, but calling generate() crashed with
AttributeError: 'HuginnForCausalLM' object has no attribute 'prelude'.
Changes
1. Model Interface Validation: Checked for required methods (prelude,
recurrent_prepared, coda) in PipelinedSpeculativeRunner.__init__, raising
NotImplementedError with actionable explanation at construction time.
2. Docstring Correction: Removed "32 for Huginn/Raven" from PipelineSchedule
docstring.
3. Regression Tests: Added test_huginn_rejected_at_construction in
tests/test_pipelined_speculative.py using a real HuginnForCausalLM instance.
Validation
- pytest tests/test_pipelined_speculative.py: 8 passed in 0.32s
- Full speculative test suite (19 tests) passed cleanly.
- Code style: clean on ruff check and ruff format.
…docstring, and verify counter invariants
cs-ai-li
left a comment
cs-ai-li
left a comment
There was a problem hiding this comment.
Commit Message
fix(speculative): document prototype limits, correct retain_branches docstring, and
verify counter invariants
To address maintainer review feedback on claims, docstrings, and counter invariants:
- Label PipelinedSpeculativeRunner as an experimental standalone prototype in module and class docstrings, documenting known limits (single-request, greedy-only, no default EOS, not yet integrated into StageEngine/serving stack).
- Correct retain_branches docstring: clarifies that it reclaims state slots only,while KV truncation is deferred to generate() via cache.truncate_suffix().
- Fix stats.committed_tokens accounting on early prefill stopping.
- Add test_speculative_stats_invariants to verify required counter invariants: committed_tokens == len(tokens), accepted + rejected == len(tokens) - 1, 0 <= accepted <= len(tokens) - 1, and cycles >= 1.
Future Work & Follow-up RoadmapTo maintain clean responsibility boundaries and reviewable diffs, the remaining
|
Hi @0z5a , Thank you for the exceptionally thorough review, the concrete minimal reproducer, and Your analysis of the cross-branch KV history pollution is spot-on: because physical KV At the same time, your H100 results provide invaluable independent validation for the Actions Taken in Response to Your Feedback Following your recommendation ("Branches would need their own KV storage... or
All 22 speculative unit & integration tests are passing cleanly on the updated head, We would appreciate your re-review of the updated PR. Thanks again for your insightful |
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Re-review of head 155d878cc444130a72e0f5b7962aecd1eefa75c3 against base a32c89a757d42a5a1702440e27425d000186068b. Recommendation: request changes.
The new construction guards address the earlier unsupported-q2, model-depth mismatch, and Huginn-interface findings. The single-proposal design has a plausible correctness argument with LAST-EXITED KV and full-depth verification: predecessors populate the depth-specific history needed by later lanes, and rejection prunes descendants before corrected tokens advance. The three inline comments identify remaining cache-layout, CUDA synchronization, and test-selection defects.
The performance and integration claims also need to match the reduced scope:
- This is a standalone single-request, greedy prototype; it is not reachable from the engine/serving path.
forward()immediately callsread()afterqueue(), and that read synchronizes the event. The current implementation therefore still waits for decisions each cycle; pinned nonblocking copies alone do not demonstrate computation/readback overlap.- The reported 2.62x eager-AR comparison does not establish an improvement over the existing graph-enabled block speculation. That is a relevant alternative baseline for qualification.
- For one generation producing 64 tokens, including the prefill token, the counters can record only 63 final proposal decisions. The reported "72 accepted / 0 rejected" needs a reproducible script and an explanation of any warmup or cumulative-counter scope.
Validation: local Ruff lint, formatting, and diff checks passed. The recorded CI merge build 981f15d reported 821 passed / 357 skipped; this is existing CI evidence, not a local runtime rerun. Runtime failures in the inline comments were not reproduced locally because user namespaces are disabled and the isolated execution attempt could not start. No GPU execution, BF16 losslessness validation, or performance measurements were performed for this review.
| ) | ||
|
|
||
| self.model = model | ||
| self.cache = cache |
There was a problem hiding this comment.
[P1] Reject SHARED KV at construction
This standalone runner accepts KVCacheManager(layout="shared"), but its execution requires depth-separated KV. SHARED maps every depth to plane 0 (KVCacheManager._plane), while prefill() advances the whole prompt one depth at a time and forward() advances multiple positions at different depths. For example, the second prompt token's first loop reads the first token's first-loop KV here; sequential SHARED execution requires the first token's final-loop KV. Overlapping decode lanes have the same history problem, so accepting this configuration can silently change greedy outputs. Please require cache.layout == "last_exited" in this constructor, matching the existing speculative guard in LLMEngine, and add a construction-rejection regression test.
There was a problem hiding this comment.
Resolved in the shared runtime integration (49829d4 / 1ca951e): the standalone PipelinedSpeculativeDecoder was replaced by SpeculativeRunner (vllm_rlt/worker/speculative.py), StageEngine.__init__ (vllm_rlt/engine/stage_engine.py) enforces cache_config.layout == "last_exited" at construction, and SpeculativeConfig.__post_init__ (vllm_rlt/config.py) enforces target_loops % draft_loops == 0 when mode == "pipeline".
| target = self.host[: rows * columns].view(rows, columns) | ||
| if self.is_cuda: | ||
| target.copy_(decisions, non_blocking=True) | ||
| self.ready.record() |
There was a problem hiding this comment.
[P1] Record the completion event on the decision tensor's device
When decisions is on cuda:1 but the process's current device remains cuda:0, the nonblocking D2H copy is submitted on GPU 1, whereas bare self.ready.record() uses the current stream of GPU 0. read() can therefore finish synchronizing an unrelated event and consume the pinned buffer before its copy completes, yielding stale or uninitialized token decisions. Please record on torch.cuda.current_stream(self.device) (or the actual copy stream), and cover a tensor device that differs from the current device. This is a source-level finding; no multi-GPU reproducer was run in this review. PyTorch documents the default stream selection here: https://docs.pytorch.org/docs/2.14/generated/torch.cuda.Event.html#torch.cuda.Event.record
There was a problem hiding this comment.
Resolved in the shared runtime integration (49829d4 / 1ca951e): SpeculativeRunner._submit_pipeline (vllm_rlt/worker/speculative.py) copies endpoint choices into ticket-owned pinned host memory (decisions) and records the completion event on torch.cuda.current_stream(self.device) within self.compute_stream.
| assert recovered == [[101], [202], [303]] | ||
|
|
||
|
|
||
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
There was a problem hiding this comment.
[P2] Mark the CUDA readback test as gpu
This test is gated only by CUDA availability and lacks @pytest.mark.gpu. Consequently, both plain pytest and pytest -m "not gpu" can execute its CUDA allocations on a GPU host without --run-gpu; the opt-in hook in tests/conftest.py only skips tests carrying the gpu marker. Please add that marker so the test follows the repository's reservation and execution gate.
There was a problem hiding this comment.
Resolved in the shared runtime integration (49829d4 / 1ca951e): the standalone prototype tests were replaced by the unified pipeline test suite in tests/test_speculative.py, with all CUDA tests marked @pytest.mark.gpu.
bjf-frz
changed the title
Review Response & Real-Hardware Benchmark Evidence (
|
| Input len / batch | Block median tok/s | Pipeline median tok/s | Change | Block range | Pipeline range | Matched requests |
|---|---|---|---|---|---|---|
128 / 8 |
704.70 | 887.42 | +25.9% | 704.61-705.09 | 802.90-888.22 | 8/8 |
4096 / 1 |
39.43 | 35.62 | -9.7% | 38.05-39.56 | 35.55-35.63 | 0/1 |
4096 / 4 |
47.35 | 46.98 | -0.8% | 47.35-47.37 | 46.98-47.05 | 3/4 |
- Block 2-loop vs Pipeline 2-loop (
draft_loops=2, target_loops=4):
| Input len / batch | Block median tok/s | Pipeline median tok/s | Change | Block range | Pipeline range | Matched requests |
|---|---|---|---|---|---|---|
128 / 8 |
495.38 | 641.86 | +29.6% | 495.38-495.43 | 574.07-643.45 | 8/8 |
4096 / 1 |
33.50 | 36.33 | +8.5% | 33.50-33.50 | 36.29-36.38 | 0/1 |
4096 / 4 |
46.98 | 45.98 | -2.1% | 46.98-47.00 | 45.95-45.98 | 4/4 |
4. Real-Checkpoint GSM8K Benchmarks (Ouro-1.4B, BF16 & FP32, triton, max_tokens=64, K=4)
Single-Request Latency & Throughput (Concurrency = 1, GSM8K prompt = 50 tokens, 3 reps after warmup)
| Mode | Draft Loops | Compilation | Throughput (tok/s) | TPOT (ms/tok) | E2E Latency (s) | Speedup vs AR Graph | Speedup vs Block Graph | Exact Match vs AR |
|---|---|---|---|---|---|---|---|---|
Standard AR |
— | Eager | 40.86 | 24.47 | 1.5663 | 0.59x | 0.31x | Baseline |
Standard AR |
— | CUDA Graph | 68.90 | 14.51 | 0.9289 | 1.00x | 0.53x | 1/1 (100%) |
Block Spec (K=4) |
1 | Eager | 88.91 | 11.25 | 0.7198 | 1.29x | 0.68x | 1/1 (100%) |
Block Spec (K=4) |
1 | CUDA Graph | 129.84 | 7.71 | 0.4936 | 1.88x | 1.00x | 1/1 (100%) |
Pipeline Spec (K=4) |
1 | Eager | 114.57 | 8.73 | 0.5587 | 1.66x | 0.88x | 1/1 (100%) |
Pipeline Spec (K=4) |
1 | CUDA Graph | 156.56 | 6.40 | 0.4095 | 2.27x | 1.21x (+20.6%) | 1/1 (100%) |
Pipeline Spec (K=4) |
2 | CUDA Graph | 106.54 | 9.39 | 0.6007 | 1.55x | 0.82x | 1/1 (100%) |
Multi-Request Continuous Batching (Concurrency = 8, 8 ragged GSM8K prompts [38..114] tokens, 512 output tokens)
| Mode | Draft Loops | Compilation | Async Sched | Total Throughput (tok/s) | Wall Time (s) | Speedup vs AR Graph | Speedup vs Block Graph | Deterministic Equivalence |
|---|---|---|---|---|---|---|---|---|
Standard AR |
— | Eager | Off | 291.08 | 1.7605 | 0.74x | 0.39x | Baseline |
Standard AR |
— | CUDA Graph | Off | 393.27 | 1.3019 | 1.00x | 0.53x | 8/8 vs AR Eager |
Block Spec (K=4) |
1 | CUDA Graph | Off | 738.55 | 0.6932 | 1.88x | 1.00x | 6/8 BF16 (8/8 FP32) |
Pipeline Spec (K=4) |
1 | Eager | Off | 725.80 | 0.7054 | 1.85x | 0.98x | pipe_eager == pipe_graph (8/8) |
Pipeline Spec (K=4) |
1 | CUDA Graph | Off | 954.02 | 0.5367 | 2.43x | 1.29x (+29.2%) | pipe_graph == pipe_eager (8/8) |
Pipeline Spec (K=4) |
1 | CUDA Graph | On | 957.14 | 0.5349 | 2.43x | 1.30x (+29.6%) | pipe_async == pipe_graph (8/8) |
Pipeline Spec (K=4) |
2 | CUDA Graph | On | 619.28 | 0.8268 | 1.57x | 0.84x | pipe_async == pipe_graph (8/8) |
Numerical Precision Verification (BF16 vs FP32 on 8-Request Ragged Batch)
BF16 8-req: pipe_eager == pipe_graph: True (8/8 exact match)
BF16 8-req: pipe_graph == pipe_async: True (8/8 exact match)
BF16 8-req: exact sequence match vs AR: 6/8
FP32 8-req: pipe_graph_async == ar_eager: True (8/8 exact match)
Separate synchronous and asynchronous pipeline execution, preserve lifecycle and KV masking, and document the reverted boundary row-selection experiment.
|
Updated the implementation and PR description in 483c10e.
|
|
lgtm |
Resolve conflicts with pipelined speculative decoding (InferMatrixLab#121): - SpeculativeConfig keeps main's mode and WISE's draft_attention; a WISE draft only exists in block rounds, so draft_attention='wise' rejects mode='pipeline' (new test). - Recurrent CUDA graph entries carry both main's write_mask buffer and WISE's static inputs; the graph key separates WISE from dense decode and masked from unmasked writes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Summary
Add
--speculative-mode pipelineto the shared self-speculative runtime. Block mode serializes autoregressive drafting and multi-token verification. Pipeline mode batches tokens at different recurrent depths, advancing drafting and verification together; rejection discards dependent descendants and restarts from the corrected token. Block remains the default.The current implementation separates synchronous execution in
speculative.pyfrom genuinely asynchronous execution inasync_speculative.py. The asynchronous path advances fixed-capacity device state before CPU collection, supports two outstanding steps for the same request, and uses the existing copy stream and leased pinned buffers for compact D2H results. Result decoding is independent of whether statistics collection is enabled.CUDA Graph replay supports active-row attention metadata and masked KV writes. Cancellation, request-ID reuse, failed submissions, and resident/CPU-offloaded preemption have explicit lifetime handling; snapshots require drained pending work. Native Linear/RMSNorm are retained. The experimental PRELUDE/CODA row selection has been reverted.
Pipeline mode requires greedy decoding,
last_exitedKV layout, and compatible draft/target loop counts. Mixed prefill/decode is not supported for this mode. Pipeline plus PD has not been validated end to end.Token-level execution
The diagrams show PRELUDE, RECURRENT and CODA, token-level batches, and corresponding acceptance/rejection events. Device decisions become available before asynchronous host delivery.
See the self-speculation documentation.
Current validation
Results below are reviewer-run on NVIDIA B300. Performance and quality were measured on 2026-10-26; all 65 runtime Python files were SHA-256 checked against publication commit
483c10e. The final working tree was regression-tested again on 2026-10-27: 327 passed, 1 skipped in 168.74 s. The skip is the unsupported FA4/FP32 combination, not an unrun GPU suite.This covers the shared runtime and pipeline tests, including asynchronous ticket/state handling, cancellation and ID reuse, preemption/restore, and graph execution. The mixed-execution tests do not establish pipeline/mixed support.
Matched performance comparison
Ouro-1.4B, BF16, one B300 per experiment, FA4 4.0.0b30, PyTorch 2.13.0+cu130, Triton 3.7.1, Transformers 4.55.0. Both block and pipeline use asynchronous scheduling and CUDA Graphs. Target depth 4, draft depth 1 or 2, K=4; max sequences 32, token budget 256, prefill chunk 1024, graph row/entry limits 256, one CPU thread. Graph limits are benchmark settings, not changed defaults.
Fixed ShareGPT-derived token IDs, input lengths 128/4096, 128 generated tokens per request, greedy, ignore EOS, no chat template. Each case submits C requests together. Two engines share the same model in one process, warm up fully, then run three alternating A/B repetitions. 48 timed runs, 588 completed requests, 75,264 output tokens, no failed measured requests; zero timed graph captures or fallbacks.
Throughput is median output tokens / complete batch wall time, including prefill and final synchronization. E2E and TTFT run from local request arrival to host-visible outputs; HTTP/IPC is excluded. Latency percentiles pool requests across three repetitions. TPOT is
(last output - first output) / 127, not individual inter-token latency. These are fixed-burst comparisons, not a saturation sweep; small-sample p95 values are descriptive.All arrows below are block → pipeline. Relative changes use unrounded measurements; negative latency changes are improvements.
Short prompts and low concurrency benefit most. At 4096/32, throughput changes of +0.22%/+0.80% do not establish a meaningful gain; E2E p50 worsens by 2.34%/1.88%. TPOT also rises substantially. Earlier first-token delivery and later completion both increase that interval, so it is not a direct measure of recurrent-kernel slowdown.
Pipelining combines shallow draft and deeper verification work into fewer submissions; it does not remove all corresponding token-loop computation. At high concurrency, block verification already has substantial batching, while pipeline state updates and more frequent output commits remain. Long-prompt prefill/admission time also dominates much of the end-to-end interval. These explain the limited opportunity, rather than proving a single isolated source for every timing difference.
Full throughput ranges and TTFT/E2E/TPOT p50/p95
Acceptance and rejection
Collected in one separate, untimed replay per case. Rates divide by actually tested candidates, excluding untested descendants discarded after rejection. They are not throughput proxies. Counts differ because block and pipeline commit tokens differently.
Output agreement and task accuracy
Timed repeats were deterministic within each configuration. Block/pipeline outputs matched exactly for 288/294 paired requests across repetitions. The six differences are the same short-prompt request at two draft depths across three repetitions. All long-prompt pairs matched. This is output agreement, not a task-accuracy metric.
Separately evaluated 256 fixed GSM8K test questions in each of five modes: existing raw 3-shot protocol and strict answer extraction, no chat template, greedy generation, natural EOS/protocol stop, 1024-token cap, concurrency 8, fixed target depth 4, BF16/FA4/async/graphs. The subset was selected before evaluating outputs by sorting SHA-256(
pr121-native-quality-20261026:ID), taking 256 IDs, then restoring ID order. Protocol hash:f63b08813b0703f57ca3c7f740dcef46db18284f19078d03f3ca4a25ce685544.Pipeline and block have equal aggregate scores at each draft depth, but individual answers change: draft 1 has five correct→wrong and five wrong→correct; draft 2 has four each. Exact full-sequence matches are 192/256 and 189/256. This slice shows no aggregate quality loss; it does not prove improvement or sub-1-percentage-point non-inferiority. A separate 64-question pipeline draft-1 batch-sensitivity check scored 42/41/42 at concurrency 1/8/32, with no monotonic accuracy trend.
Numerical cause: an untimed same-input layer trace found the first divergence in
model.layers.0.mlp.down_proj, with 8 versus 32 rows: three BF16 elements differed, maximum absolute difference 0.0001220703125. Later recurrent layers amplify small rounding differences. For one short request, output token 93 changed from0to:: the block logits tied at 21.25/21.25, whereas pipeline favored:at 21.375 versus 21.25. This directly demonstrates a native linear-algebra batch-shape effect; it should not be attributed solely to FA4 or called evidence of a state-machine bug.Differences can grow with changing batch shapes, repeated recurrent computation, near-tied logits, and longer autoregressive continuations after the first changed token. They need not grow monotonically with concurrency or sequence length. Native Linear/RMSNorm remain in use; there is no custom batch-invariant math file in this version.
Evaluated and reverted: PRELUDE/CODA row selection
The experiment selected active newborn rows for PRELUDE and root/newborn endpoints for CODA on GPU, ran the existing native modules on compact buffers, then scattered results back. Static capacities preserved graph replay and avoided CPU-side device decisions.
Matched B300/FA4/async/graph comparisons, draft 1, three alternating repetitions:
A separate short-prompt ablation found PRELUDE-only -0.56%, CODA-only -0.55%, and both -0.46%. Absolute timing drifted between phases, so only within-phase comparisons are used. Timed outputs matched; all 512 quality comparisons matched exactly. Paired ops/full profiles over eight steps showed CODA projection rows falling from 32 to 16, but projection time only changed from 0.297 to 0.293 ms while ordinary launches rose from 4,552 to 4,672. Indexing, scatter and selection overhead outweighed the measured saving. The experiment is reverted; active/write_mask correctness handling remains.
Implementation details, ablation and profiling evidence are recorded in the existing self-speculation chapter. Raw local reports remain under
.reviews/async-pipeline-native-perf,.reviews/async-pipeline-native-quality,.reviews/async-pipeline-selective-perf, and.reviews/async-pipeline-selective-quality; benchmark dumps are not committed.Historical author-reported RTX 5090 results (previous implementation)
Preserved from the previous PR description for provenance. These results concern the earlier shared-runtime revision through
df3c17e, not the current asynchronous implementation, and have not been rerun by us on RTX 5090. The author reported 31 GPU tests passing, six deselected, using the Triton backend. The synthetic inputs and older execution path differ from the current B300 measurements; do not compare the absolute numbers across these campaigns. The prior inference that FA4 was the primary source of disagreement is not established by these cross-device tests; the current same-input trace above demonstrates a Linear batch-shape contribution. The small GSM8K prompt runs below measure speed/output agreement, not GSM8K task accuracy.2. Matched A/B Matrix vs. B300 Setup (
Ouro-1.4B, BF16,triton,max_tokens=128,num_speculative_tokens=4)Using the same prompt synthesis (
(x + i*17) % 1000 + 100),max_tokens=128,async_scheduling=True,cuda_graphs=True,cuda_graph_max_batch_size=256,cuda_graph_max_graphs=256, 1 warmup + 2 timed iterations on 1× NVIDIA GeForce RTX 5090 (32 GB,tritonbackend) (4096 / 4used instead of4096 / 8due to 32 GB VRAM capacity):draft_loops=1, target_loops=4):128 / 84096 / 14096 / 4draft_loops=2, target_loops=4):128 / 84096 / 14096 / 43. Real-Prompt GSM8K Benchmark (
Ouro-1.4B, BF16 & FP32,triton,max_tokens=64,K=4)Single-Request Latency & Throughput (
Concurrency = 1, GSM8K prompt = 50 tokens, 3 reps after warmup)Standard ARStandard ARBlock Spec(K=4)Block Spec(K=4)Pipeline Spec(K=4)Pipeline Spec(K=4)Pipeline Spec(K=4)Multi-Request Continuous Batching (
Concurrency = 8, 8 ragged GSM8K prompts[38..114]tokens, 512 output tokens)Standard ARStandard ARBlock Spec(K=4)Pipeline Spec(K=4)pipe_eager == pipe_graph(8/8)Pipeline Spec(K=4)pipe_graph == pipe_eager(8/8)Pipeline Spec(K=4)pipe_async == pipe_graph(8/8)Pipeline Spec(K=4)pipe_async == pipe_graph(8/8)Numerical Precision Verification (
BF16vsFP32on 8-Request Ragged GSM8K Batch)Required human contributor answers