[jthomson04]

Why and what changed

Speculative replay (ngram, MTP) spent about 40% of replay in VllmCore::emit_speculative_ready_tokens. Most of that was repeated FxHashMap<Uuid, _> lookups: about 8 per ready request per pass, plus 3 per emitted token. The ChaCha12 acceptance sampler accounted for another 4–6%.

• Commit 1 (identical output). Build one burst record per ready request from one request lookup and one scheduled-work lookup. Price the forward from the records that survive reservation, then sample and emit each burst under a single get_mut. RNG draw order and draw count are unchanged.
• Commit 2 (changes sampled streams). Sample acceptance with rand's SmallRng (Xoshiro256++) instead of StdRng (ChaCha12). Fractional acceptance rates now produce different seeded streams for vLLM and SGLang. Seeds, per-rank offsets, and rates of 0 and 1 are unchanged.

Review map

• Risk level: medium
• Start with: speculative_bursts and emit_speculative_ready_tokens in crates/core/src/engine/scheduler/vllm/core.rs, especially decode pricing after reservation drops or preempts a request.
• Public or serialized contract changed: no. Seeded speculative outputs at fractional acceptance rates change with commit 2.
• Compatibility or rollback concern: SmallRng's algorithm, and so the stream for a given seed, may change on a rand upgrade. This is noted on the sampler.

Evidence

• Tests: cargo test -p aisimulate-core --lib --tests; pytest -c pytest.ini tests/test_ngram_cli.py tests/test_runner.py and the other speculative replay suites; SDK speculation tests.
• Negative or boundary cases:
  • A fractional-acceptance pass-sequence pin covers budget-capped bursts, max_model_len, and preemption. It was recorded on main, passes unchanged on commit 1, and is re-recorded by commit 2.
  • Two survivor tests check that the timing model is charged only the survivor after a preemption or an impossible lookahead.
• Before/after: Tyche Grace node, 131,072 requests, median of 3 rotated rounds of wall_time_ms.
  • ngram (Qwen3-32B TP4): 11.65 s → 10.58 s → 9.79 s, −15.9%.
  • MTP (DeepSeek-V3.2 DP8/EP8): 14.05 s → 13.01 s → 12.53 s, −10.9%.
  • No-speculation control: within 0.6%.
• Identity for commit 1:
  • prediction.json matches main for both workloads and for SGLang MTP.
  • Per-request rows match by session at 8,192 requests with --capture-per-request.
• Commit 2:
  • Request and output-token totals are unchanged.
  • Mean TPOT moves by at most 1.1%. That matches the shift from changing only the seed on main (MTP: +1.01% at seed 43, +0.99% for commit 2).

Modeling or data provenance

N/A

🤖 Generated with Claude Code

[coderabbitai]

📝 Summary

Risk: Medium. The three areas that need the most human attention are:

1. Burst accounting after reservation drops or preemption. Confirm timing charges and acceptance statistics use only surviving bursts.
2. Seeded RNG compatibility. SmallRng replaces StdRng, so the sampled sequence can change for the same seed.
3. Performance and workload coverage. The reported gains need validation against the target workloads and environments.

Behavior and contracts

• The vLLM scheduler records speculative burst data per ready request, then uses that data for reservation, timing, and token emission. Completion handoff delay is limited to completed requests.
• SpeculativeDecodeSampler now uses SmallRng with the supplied seed. The seeded output sequence is not preserved as a compatibility contract.
• No public declarations are reported as changed. The RNG stream change remains observable behavior for users relying on exact seeded output.
• The supplied changes concern vLLM. The objective mentions SGLang behavior, but the supplied change summary does not identify a SGLang file change.

Evidence and readiness

• The supplied summary describes tests for capped bursts, context limits, preemption, and timing charges for surviving requests. It also reports broader test runs and benchmark comparisons. These results were not independently verified here.
• The author-reported benchmarks show ngram replay improving from 11.65 s to 9.79 s and MTP from 14.05 s to 12.53 s. The no-speculation control reportedly stayed within 0.6%; mean TPOT reportedly changed by at most 1.1%.
• Test results and benchmark claims are reported evidence, not independently confirmed results. Current review finding counts are unavailable. No bot review constitutes approval.
• Technical quality: The summarized implementation and focused regression tests address the key risk of charging timing for requests removed during reservation.
• Merge readiness: Cannot be established from the available evidence. The repository inspection did not provide a diff, and the reported validation was not independently run.

The speculative sampler now uses SmallRng. The vLLM scheduler records per-request burst data for reservation, decode timing, acceptance accounting, and output emission.

Changes: Speculative decode and vLLM scheduling

Layer / File(s)	Summary
Sampler RNG crates/core/src/engine/common/speculative.rs	SpeculativeDecodeSampler now uses SmallRng with the existing seed argument. Its documentation notes that seeded streams may change with rand.
Burst reservation and surviving decode work crates/core/src/engine/scheduler/vllm/core.rs, crates/core/src/engine/scheduler/vllm/tests.rs	Per-request burst records capture reservation and decode-pricing data. Decode timing uses surviving records after reservation and preemption. Tests record batch size and active KV tokens, and check timing for surviving requests and impossible lookahead.
Acceptance and burst emission crates/core/src/engine/scheduler/vllm/core.rs, crates/core/src/engine/scheduler/vllm/tests.rs	Acceptance accounting and output construction use saved burst records. Tests check emitted bursts under fractional acceptance, token-budget, context-limit, and preemption constraints.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• ai-dynamo/aisimulate#519: Adds native vLLM MTP scheduling, reservations, preemption, and burst emission in VllmCore, the same speculative emission flow updated here.
• ai-dynamo/aisimulate#612: Adds mixed prefill/decode timing passed through speculative token emission, which is used by the vLLM emission path updated here.

Merge Risk: Minimal · up to 093fa

No actionable regression was identified in speculative scheduling or emission; the PR is mergeable after normal checks.

Pre-merge checks: 8

Check name	Status	Explanation
Title check		The title precisely describes the behavioral changes: reduced per-token lookup overhead and RNG cost in vLLM speculative emission.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Cross-Layer Contract		Cross-layer trace of the changed inputs and outputs. Status: PASS, with non-blocking gaps listed below. Changed files (base..head): speculative.rs (sampler RNG type), vllm/core.rs (emission re…
Modeling And Data Evidence		The check applies. The PR changes selection logic (speculative burst sampling and emission) and predicted outputs (seeded acceptance streams at fractional rates). The PR supplies provenance and before…
Compatibility Boundaries		The PR changes only three Rust source files (speculative.rs, vllm/core.rs, vllm/tests.rs). It touches no Cargo.toml, pyproject.toml, Python package, bindings, schema, CLI, or docs. The SpeculativeDeco…
Review Evidence		The PR description meets the check's criteria. - Commands and results are named. It lists cargo test -p aisimulate-core --lib --tests and `pytest -c pytest.ini tests/test_ngram_cli.py tests/test…
Description check		The description is mostly complete. It covers the behavior change, risk, compatibility impact, tests, boundary cases, performance results, and data provenance. The Evidence section omits Fast CI, Full…

Full details: Cross-Layer Contract

Explanation

Cross-layer trace of the changed inputs and outputs. Status: PASS, with non-blocking gaps listed below. Changed files (base..head): speculative.rs (sampler RNG type), vllm/core.rs (emission refactor), vllm/tests.rs (new tests). No Python, schema, CLI, or docs files changed. Traced consumers: - seed / aic_mtp_seed plumbing is unchanged. Python config/engine.py (seed field, default 42) → runner.py (aic_mtp_seed) → rank.rs / config.rs → SpeculativeDecodeSampler::new(rates, seed + offset) in both vllm/core.rs:847 and sglang/core.rs:226. The seed argument and per-rank offsets are unchanged. Only the RNG type changed (StdRng → SmallRng). - rand 0.9.5 enables small_rng by default (rand-0.9.5/Cargo.toml), so SmallRng compiles without a Cargo change. - Sampler unit tests (speculative.rs tests) cover exact 0/1 rates, first-rejection stop, RNG non-advance for unscheduled drafts, seed determinism, seed-offset distinctness, and an empirical-mean check at fractional rates with 200k samples. These exercise the new RNG directly. - vLLM emission: new pass-sequence pin, reservation-preemption survivor test, and impossible-lookahead survivor test cover the changed emission and timing paths. - Python: test_ngram_cli.py checks same-seed reproducibility and seed-difference. Neither pins fractional-rate numeric values, so the stream change does not stale any fixture. - Docs: docs/replay/engine/speculation.md describes seed as an unsigned 64-bit sampling seed and does not promise a stable stream. Its statements remain accurate after the change. Gaps (not FAIL conditions on their own): 1. The PR description says SGLang's per-token handoff delay "now gets the same completion guard." No SGLang file is in the diff. The guard exists only in vllm/core.rs. SGLang's call sites in sglang/decode.rs already pass a literal true, so behavior is not affected. The description overstates the SGLang change. 2. SGLang shares the changed sampler, and the PR says its fractional-rate streams change. No test in the diff pins SGLang output at fractional rates. sglang/tests.rs only uses rates 1,1. The new stream therefore has no SGLang-specific pin. The PR gives no SGLang before/after evidence for commit 2. 3. The PR states "Public or serialized contract changed: no." Seeded fractional-rate outputs do change. The PR discloses this in the review map. The speculation.md seed row does not mention the stream change. The sampler doc comment does. Verification limits: I could not run the Rust or Python test suites. Reported performance and prediction.json identity figures come from the PR author and are not independently checked.

Full details: Modeling And Data Evidence

Explanation

The check applies. The PR changes selection logic (speculative burst sampling and emission) and predicted outputs (seeded acceptance streams at fractional rates). The PR supplies provenance and before/after evidence that meets the stated bar. Evidence found: - Golden pin in the diff: fractional_acceptance_follows_a_fixed_burst_sequence in crates/core/src/engine/scheduler/vllm/tests.rs asserts an exact per-pass burst sequence under budget caps, max_model_len, and KV preemption. It also asserts that preemptions occur, so the workload is not trivial. - Survivor tests in the same file (reservation_preemption_prices_only_the_survivor, impossible_lookahead_prices_only_the_survivor) assert the exact timing-model charge (batch, active KV tokens). These are explicit oracles for the timing change. - Commit 1 identity claims: prediction.json matches main for both workloads and SGLang MTP. Per-request rows match at 8,192 requests. - Commit 2 explanation: the seed-stream change is stated. Request and output-token totals are unchanged. The mean TPOT drift (≤1.1%) is compared with a seed-only control on main (MTP +1.01% vs +0.99%), which explains the drift as RNG variance. - Performance provenance: hardware (Tyche Grace node), workload size (131,072 requests), models and parallelism, median of 3 rotated rounds, and per-round wall_time_ms values. A no-speculation control stays within 0.6%. - Commands are listed for reproduction. Gap (not disqualifying): the raw benchmark artifacts (prediction.json, per-request captures, timing rounds) are not in the repository. The performance numbers rest on a summary in the PR description. I could not verify the commit 1 vs commit 2 split from the single reviewed range, so I rely on the author's description for that point. The golden pin, identity checks, and seed-only control together give an explained diff. No unexplained refresh is present.

Full details: Compatibility Boundaries

Explanation

The PR changes only three Rust source files (speculative.rs, vllm/core.rs, vllm/tests.rs). It touches no Cargo.toml, pyproject.toml, Python package, bindings, schema, CLI, or docs. The SpeculativeDecodeSampler change is pub(crate); the rand dependency is unchanged (0.9.2 in workspace, 0.9.5 in Cargo.lock). The documented engine.speculation.seed default (42) is unchanged, and the docs make no cross-version stream-stability promise. The changed files contain no Dynamo references. The wheel and crate artifact contract is untouched. Residual notes: seeded fractional-acceptance streams change (disclosed in the PR and noted in the code comment), and SmallRng feature availability under the current rand feature set was not verified by a build.

Full details: Review Evidence

Explanation

The PR description meets the check's criteria. - Commands and results are named. It lists cargo test -p aisimulate-core --lib --tests and pytest -c pytest.ini tests/test_ngram_cli.py tests/test_runner.py, plus other speculative replay suites and SDK speculation tests. It also reports numeric results: wall-time changes, TPOT drift, and token totals. - Negative and boundary validation fits the risk. The description cites a fractional-acceptance pass-sequence pin covering budget-capped bursts, max_model_len, and preemption. It also cites two survivor tests for reservation preemption and impossible lookahead. The diff confirms both survivor tests exist in crates/core/src/engine/scheduler/vllm/tests.rs (reservation_preemption_prices_only_the_survivor, impossible_lookahead_prices_only_the_survivor) and that the pass-sequence test exists. - Evidence types are distinguished. Local tests are separated from parity checks (prediction.json identity and per-request row matching at 8,192 requests). Production-like measurements are reported separately as a Tyche Grace node benchmark with rotated rounds and a no-speculation control. - No inactive workflow is cited as CI evidence. The description does not mention hosted CI or any .github/workflows file, so the failure condition does not apply. Minor gap: the description does not use the term "fake-runner" or label any run as fake-runner. The check does not require that label to pass, because the other evidence classes are clearly distinguished. The author may add the label for clarity.

Full details: Description check

Explanation

The description is mostly complete. It covers the behavior change, risk, compatibility impact, tests, boundary cases, performance results, and data provenance. The Evidence section omits Fast CI, Full CI, CodeRabbit reviewed commit, Codex reviewed commit, and expected-value derivation. The Tracking section is also missing.

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[tedzhouhk]

Reviewed speculative reservation, survivor-only timing, acceptance accounting, and emission/completion boundaries. Local Rust unit/integration suites passed 2,579 tests (one data-dependent test intentionally ignored). An independent 168-case replay matrix matched the parent exactly on the refactor-only commit; the final head differed only in the 41 fractional-acceptance cases, with completed-request and output-token totals preserved. The disclosed RNG stream change also reaches the TensorRT-LLM simulation through the shared core. No actionable correctness findings; current-head Full CI is green. Author-reported performance benchmarks were not rerun.