Skip to content

[FEAT] Add opt-in single-fallback proposals for greedy speculation - #59

Closed
prettygirlisnotme wants to merge 4 commits into
InferMatrixLab:mainfrom
prettygirlisnotme:proposal/single-fallback
Closed

prettygirlisnotme wants to merge 4 commits into
InferMatrixLab:mainfrom
prettygirlisnotme:proposal/single-fallback

Conversation

@prettygirlisnotme

@prettygirlisnotme prettygirlisnotme commented Sep 24, 2026 •

Copy link
Copy Markdown

Purpose

Refs #43. Add an opt-in alternate suffix for fixed-depth greedy speculation. When the shallow top-1/top-2 full-softmax probability gap is small, draft one alternate suffix. Reuse it only when the target verifier confirms the alternate token at the first rejection.

SpeculativeConfig(alternate_prob_gap_threshold=0.2, ...) enables this path. The default None keeps the plain runner. The implementation supports one branch per request per round with last_exited KV, synchronous eager execution and refill scheduling.

The runner shares committed pages, copies the private stem, and commits or releases alternate KV within the round. A failed spare-capacity allocation continues on the primary path without evicting cached prefixes. The scheduler reserves up to 2*n - 1 verification rows for a primary chain of n rows, so the packed verification batch fits max_num_batched_tokens.

Validation

Based on main ecb1f8b505b7e831815b40aec3b4598619cca23a, including #48/#51. Per-file SHA256 in the evidence identifies the tested patch.

  • CPU: 461 passed / 153 skipped; Ruff lint and format passed.
  • L40S/Triton speculative regression: 104 passed / 7 FA4 deselected. That job later failed in the separate old diagnostic entrypoint; the diagnostic was repaired and run separately.
  • FP32 fixed-input branch/plain/native comparisons passed the existing atol=rtol=1e-4 checks in eight cases. BF16 fixed-input errors and natural-generation differences are reported separately, without an allclose or lossless claim.
  • Coverage includes selected/discarded alternate suffixes, partial suffix acceptance, fork misses, rollback, ragged verification budgets, cross-page copies, prefix preservation under capacity pressure, EOS, output limits and cancellation.

Current performance and quality

Ouro-1.4B 574fa66cb8bf5abdc979642d01cf2b79b16bfab1, L40S, BF16/Triton, greedy synchronous eager, D4/K4. The performance workload uses 64 fixed prompts, 128 output tokens with EOS ignored, and five alternating repetitions at C1/C4/C8. Comparisons hold draft depth fixed.

C1 E2E throughput improved by 4.46% at d1 and 2.64% at d2. At C8 it regressed by 8.44% and 2.11%, respectively. The d2/C8 decode-only gain did not carry over to E2E. This is an experimental opt-in mechanism with measured regressions, not a general throughput improvement.

Full quality runs use C1, natural EOS/task stops and a 1024-token output cap:

Task / draft depth Baseline Alternate Paired wins / losses
GSM8K / d1 814/1319 (61.71%) 814/1319 (61.71%) 0 / 0
GSM8K / d2 819/1319 (62.09%) 819/1319 (62.09%) 0 / 0
HumanEval / d1 114/164 (69.51%) 114/164 (69.51%) 0 / 0
HumanEval / d2 114/164 (69.51%) 114/164 (69.51%) 0 / 0

HumanEval uses EvalPlus 0.3.1's standard chat prompt and sanitization with the original HumanEval tests. Earlier invalid-format runs are retained in the evidence. Complete sequences matched at each depth on both datasets, but alternate activity was limited: HumanEval d1 forked on 5/164 tasks with three alternate selections; d2 forked on 1/164 with no alternate selection. Four/six HumanEval tasks reached the output cap at d1/d2, identically in each pair. GSM8K did not record per-example fork counters.

Observed score changes were 0 pp under the predeclared proposed 1 pp drop check. This finite-dataset result does not establish statistical noninferiority or BF16 losslessness; BF16 sequence differences remain in the separate numerical diagnostic.

Evidence: results and reproduction inputs, eight original profiles.

Limits

BF16 greedy sequences can differ. The branch retains host decisions and extra stem copies, and does not inherit the plain runner's device-side draft-ID optimization. FA4, branch CUDA Graph/async execution and online-serving performance are unvalidated. E2E here is a simultaneous 64-request workload including prefill and queuing.

@prettygirlisnotme
prettygirlisnotme marked this pull request as ready for review September 27, 2026 17:11

@prettygirlisnotme prettygirlisnotme left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review summary for e05f453

Checked the default runner selection, greedy/eager guards, alternate KV commit/discard/exception cleanup, and fixed-input branch comparisons.

The integration patch adapts the runner constructor and test helpers to main ecb1f8b, including the merged CUDA Graph and device-readback changes. It also fixes one stale OuroConfig.tiny() call in the baseline GPU test. The branch remains eager-only; the default plain runner retains upstream execution support.

Cluster validation of that patch:

  • Full pre-commit and CPU regression suite: 448 passed, 153 skipped.
  • GPU-enabled branch and baseline speculative suite on L40/Triton: 91 passed, 7 FA4 cases deselected. The count includes CPU cases.

BF16 free-generation differences remain. FA4 and online-serving performance are unvalidated. The reported throughput measurements predate the merged #48/#51 optimizations; gains against the current optimized baseline have not been established.

This is an author self-review, not an independent approval.

@bjf-frz bjf-frz changed the title Add opt-in single-fallback proposals for greedy speculation [FEAT] Add opt-in single-fallback proposals for greedy speculation Oct 8, 2026
@bjf-frz

bjf-frz commented Oct 8, 2026

Copy link
Copy Markdown
Collaborator

Remove the changes to the README and move this section to the speculative decoding documentation.

Comment thread vllm_rlt/config.py Outdated
target_loops: int = 4
# None keeps the plain runner. A finite margin in [0, 1] enables the optional
# single-fallback branch runner, which is greedy-only.
fallback_margin: float | None = None

@bjf-frz bjf-frz Oct 8, 2026 •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rename fallback_margin to alternate_prob_gap_threshold to make its meaning explicit.

@bjf-frz

bjf-frz commented Oct 8, 2026

Copy link
Copy Markdown
Collaborator

Please validate this feature in the following four areas before we proceed with the code review.

1. Correctness

Cover the branch logic and resource lifecycle:

  • The primary chain is fully accepted, and the alternate branch is discarded.
  • The first rejection occurs at the fork and matches the alternate candidate; the alternate KV is committed correctly, and its suffix is verified.
  • Rejection occurs elsewhere, or the alternate candidate does not match; correction, truncation, and rollback remain correct.
  • Insufficient spare KV capacity falls back to the primary path without evicting cached prefixes or leaking resources.
  • Cover cross-page forks, ragged batches, EOS, length limits, and cancellation.
  • After fixing the scheduling-budget issue, ensure actual verification rows never exceed max_num_batched_tokens.

2. Performance

Use the current PR revision, comparing fallback_margin=None against 0.2, with all other settings held constant.

Item Configuration
Model and execution Ouro-1.4B, pinned checkpoint; same GPU, BF16, Triton, synchronous eager, greedy
Speculation D=4, K=4, draft depths 1 and 2
Concurrency 1, 4, 8
Workload Fixed 64 prompts: 32 GSM8K and 32 HumanEval; retain example IDs, templates, and token IDs
Output length 128 generated tokens, ignoring EOS for throughput measurements
Repetitions Warm up each configuration, alternate baseline/candidate runs, and collect at least five repetitions

Additionally, use one representative configuration to test short/long contexts and KV pressure, checking whether copying costs or frequent allocation fallback cause regressions.

Measure both decode-only performance and E2E performance, including prefill and queuing. Report:

  • Absolute throughput, relative changes, and run-to-run variability.
  • Request completion latency and TTFT, including p50/p95.
  • Committed tokens per round, fork rate, alternate-selection rate, and capacity-fallback rate.
  • Primary/alternate computation rows, peak GPU memory, and additional KV blocks.
  • Failures, timeouts, and unfinished requests.

3. Numerical Accuracy and Output Quality

Please investigate the existing BF16 output differences:

  • On identical fixed input chains, compare native serial execution, plain speculation, and fallback speculation. Report logits and KV errors across layers/depths, plus argmax mismatch rates.
  • For natural generation with EOS enabled, report complete-sequence agreement, the first divergent position, and the top-1/top-2 margin at that position.
  • Replay the same divergent prefix to distinguish numerical differences from state-management errors.
  • If sequence differences are acceptable, define the acceptance criteria explicitly and compare GSM8K accuracy and HumanEval pass@1 using consistent evaluation settings. The 64-prompt performance workload does not replace a full quality evaluation.

4. Profiling

Follow the rlt-perf-opt skill referenced in the developer must-read documentation. Select one low-concurrency configuration and one at concurrency 8, and capture separate ops-only and full profiles for both baseline and fallback, covering complete speculative rounds after warmup.

Focus on:

  • Whether fewer decoding rounds offset the additional draft/verification work.
  • CPU/GPU waits caused by argmax readbacks, .tolist(), and gate decisions.
  • KV stem copying and alternate-page allocation/release costs.
  • Whether metadata preparation loses optimizations already present in the plain runner.

Measure throughput in separate runs with the profiler disabled. Provide code/model revisions, commands, configurations, per-run raw results, and profile artifacts, clearly identifying both beneficial and regressing scenarios.

@prettygirlisnotme

Copy link
Copy Markdown
Author

Moved the feature section from README to docs/cdb_runtime.md and renamed the option to alternate_prob_gap_threshold. I also fixed the scheduler budget: a primary chain of n rows now reserves up to 2*n - 1 verification rows, including the alternate suffix.

The regression checks cover commit/discard/rollback, partial alternate-suffix acceptance, fork misses, capacity fallback without prefix eviction, ragged verification budgets, cross-page copies, EOS, length limits and cancellation. CPU: 461 passed / 153 skipped. L40S/Triton speculative suite: 104 passed / 7 FA4 deselected.

I reran the requested performance matrix against the current patched revision: pinned Ouro-1.4B, L40S, BF16/Triton synchronous eager, greedy, D4/K4, 64 fixed prompts, 128 outputs ignoring EOS, five alternating repetitions. These are matched-depth comparisons:

Draft / concurrency Decode change, cohort-bootstrap 95% CI E2E tokens/s, baseline → branch E2E change
d1 / 1 +3.45% [1.64, 5.38] 38.06 → 39.75 +4.46%
d1 / 4 −0.02% [−2.43, 2.99] 138.54 → 134.34 −3.03%
d1 / 8 −7.37% [−13.30, −1.67] 252.38 → 231.09 −8.44%
d2 / 1 +1.72% [0.51, 3.12] 29.83 → 30.61 +2.64%
d2 / 4 +0.64% [−0.84, 2.12] 109.46 → 109.23 −0.21%
d2 / 8 +6.76% [1.20, 14.55] 195.85 → 191.72 −2.11%

The gain is limited to some configurations; C8 E2E regresses. E2E includes prefill and queuing for 64 simultaneously submitted requests. The main matrix completed all 1820 runs; a harness argument error prevented the final pressure section from starting, so I repaired the harness and ran only those 12 cases separately. All 12 completed. The long-context KV-pressure cases were slower by about 2–5% with repeated failed branch-allocation attempts.

FP32 fixed-input comparisons passed the existing tolerance. BF16 differences remain: branch/plain full-sequence agreement was 4/8 at each draft depth in the natural-generation diagnostic. Replaying one divergent prefix with fresh serial prefill removed the branch/plain logits/KV difference for that case; this does not establish the cause of every divergence.

Full quality results, with natural EOS/task stops and a 1024-token cap:

Task / draft depth Baseline Alternate Paired wins / losses
GSM8K / d1 814/1319 (61.71%) 814/1319 (61.71%) 0 / 0
GSM8K / d2 819/1319 (62.09%) 819/1319 (62.09%) 0 / 0
HumanEval / d1 114/164 (69.51%) 114/164 (69.51%) 0 / 0
HumanEval / d2 114/164 (69.51%) 114/164 (69.51%) 0 / 0

HumanEval uses EvalPlus 0.3.1's standard chat prompt and sanitization, scored against the original HumanEval tests. The earlier raw/fenced formats produced invalid completions; those records remain in the evidence. The first full CPU scoring attempt hit a filelock/fork compatibility error; rerunning the unchanged scorer with spawned processes completed on the same generated outputs.

Complete sequences matched within each depth on both full datasets. Activity was sparse: HumanEval d1 forked on 5/164 tasks (10 rounds, three alternate selections); d2 forked on 1/164 (one round, no alternate selection). Four/six tasks hit the output cap at d1/d2, identically in each pair. GSM8K lacks per-example fork counters. The proposed acceptance check was at most a 1 pp observed drop per task/depth; all four measured changes were 0 pp. Limited branch activity and the separate BF16 divergences prevent a general losslessness or statistical noninferiority claim.

Eight profiles cover baseline/branch at C1/C8, separately ops-only/full, after warmup. At C8 the ops-only window had 261 stream synchronizations versus 2, and about 59.4 MB of DtoD copies versus zero. The branch still makes host decisions and does not use the plain runner's device-side draft-ID path. Profile timings are excluded from the throughput comparisons.

Results and reproduction inputs and eight original profiles include raw runs, latency/variability/capacity/memory counters, numerical checks, quality protocols and source/model identities. I could not find rlt-perf-opt in the developer documentation or skills at the pinned main revision; I used docs/profiling.md and the current profiler API. Could you point me to that skill?

@0z5a

0z5a commented Oct 12, 2026

Copy link
Copy Markdown
Contributor

High-concurrency e2e result seems not very stable. I will take a look at this and give some idea on how to improve.

bjf-frz commented Oct 13, 2026

Copy link
Copy Markdown
Collaborator

Based on the measurements, I don’t think this multi-branch approach is worth pursuing further in its current form.

We tested three variants:

  • Standard speculative decoding without an alternate branch.
  • The original policy: fork when p1 - p2 <= 0.2.
  • A revised policy: fork only when p1 + p2 >= 0.4 and p1 - p2 <= 0.2.

The benchmark used Ouro-1.4B on a single B300 per configuration, BF16, target depth 4, draft depths 1/2, and 4 draft tokens. Each run submitted 64 requests—32 GSM8K and 32 HumanEval—with 128 output tokens per request. Results below are medians of five runs.

Draft depth / concurrency No branch (tokens/s) Original policy Revised policy
d1 / 1 65.46 67.91 66.36
d1 / 4 234.49 225.29 226.66
d1 / 8 393.05 361.17 364.84
d2 / 1 49.48 50.95 50.46
d2 / 4 185.05 184.02 181.43
d2 / 8 322.17 317.81 310.24

The revised policy also provided no consistent latency improvement: request completion p95 increased in five of the six configurations compared with the original policy. These measurements include queueing, prefill, and decoding, but exclude HTTP/network overhead.

Our review of the current vLLM implementation did not identify an equivalent multi-branch verification path; the tree-based implementation previously referenced was from an older version.

The main reasons for stopping this exploration are:

  1. Fewer speculative rounds do not mean less total work. At d1/concurrency 8, the original policy reduced rounds from 2,109 to 1,983, but increased verification token rows from 10,396 to 12,568. Throughput consequently fell by approximately 8.1%. Extra drafting, verification, and KV operations outweighed the saved rounds.

  2. The gate does not reliably predict a useful branch. A small probability gap does not guarantee that the preceding draft tokens will be accepted, that the target will reject the primary candidate specifically at the fork, or that it will select the second candidate. The alternate branch was selected in only approximately 21–26% of forks under the original policy.

  3. A higher branch hit rate did not translate into meaningful acceleration. At d1/concurrency 8, the revised policy increased the hit rate from 21.1% to 29.7%, but reduced successful branch selections from 193 to 131 and increased speculative rounds from 1,983 to 2,041. Throughput improved by only approximately 1% over the original policy and remained 7.2% below the no-branch baseline.

  4. The decision logic adds overhead of its own. The current implementation reads probabilities back to the CPU to decide whether to fork. Stricter filtering can require checking more positions: gate evaluations increased from 5,936 to 7,332 at d1/concurrency 8. This is an additional potential cost, although we have not separately quantified its runtime contribution.

  5. The benefit is limited to low concurrency. Both branching policies showed small throughput gains at concurrency 1, but both lost to the no-branch baseline at concurrency 4 and 8.

Given these results, the current implementation is unlikely to deliver enough practical benefit to justify its complexity. I suggest stopping this multi-branch exploration and retaining the standard speculative decoding path.

@prettygirlisnotme

Copy link
Copy Markdown
Author

Got it, thanks for the extra benchmarks. Agreed, let's stop here given the results.

@bjf-frz

bjf-frz commented Oct 13, 2026

Copy link
Copy Markdown
Collaborator

Got it, thanks for the extra benchmarks. Agreed, let's stop here given the results.

Although this approach did not deliver consistent practical performance gains, understanding the underlying mechanisms is still valuable. Please document the strategies we evaluated, the benchmark results, and why the additional branching costs outweighed the benefits in the speculative decoding documentation, so future work can build on these findings.

@prettygirlisnotme

Copy link
Copy Markdown
Author

Sure, I'll add the tested approaches, results, and trade-offs to the speculative decoding docs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants