Skip to content

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

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

prettygirlisnotme wants to merge 4 commits into
ThinkFlowLab: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?

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.

2 participants