Skip to content

[RFC] Speculative Decoding for rlt #43

Open

Description

@bjf-frz

1. Motivation

Native RLT decoding executes all recurrent loops before producing each token. At low request concurrency, the small token batch can leave GPU capacity underused.

This proposal uses shallow loops to draft tokens, reuses their hidden states and KV, and batches the remaining deep loops for verification. The initial configuration is fixed d=2, D=4, with last_exited KV semantics. K is selected through measurement. The goal is faster decoding while preserving the full-depth target's generation semantics.

LoopSpec-style pipelining across token positions is a future direction.

2. Algorithm

At the start of a round, the last committed token A has not been forwarded. KV contains the prefix before A.

  1. Autoregressively generate K candidates using loops 1–2, retaining shallow hidden states and KV.
  2. Run loops 1–2 on the final candidate as well, producing states for K+1 input positions.
  3. Reuse these states and batch loops 3–4, producing K verification distributions and one bonus distribution.
  4. Accept candidates in order. At the first rejection, emit a correction and discard the suffix. If all candidates are accepted, emit a bonus token.
  5. Commit output and truncate KV. The correction/bonus token becomes the next round's input, not yet forwarded.

Example with K=2:

flowchart TD
    A["A: loops 1–2 → propose B"] --> B["B: loops 1–2 → propose C"]
    B --> C["C: complete loops 1–2"]
    C --> V["Reuse shallow states and KV<br/>Batch loops 3–4 for all three positions"]
    V --> H["A row verifies B; B row verifies C<br/>C row provides the bonus distribution"]
    H --> R1["Reject B: commit Z<br/>Keep KV through A"]
    H --> R2["Accept B, reject C: commit B, Z<br/>Keep KV through B"]
    H --> OK["Accept both: commit B, C, E<br/>Keep KV through C"]
Loading

Reuse condition: with fixed depth and last_exited, each loop reads historical KV at the same depth. Deeper computation does not modify shallow states. Computing the candidate chain's shallow loops first and then batching its deep loops preserves the dependencies, provided positions, normalization, and causal masks match. Shallow KV must not be copied into uncomputed deeper loops.

Verification: greedy compares each candidate with the target argmax and corrects the first mismatch with that argmax. For sampling, let q be the actual proposal distribution and p the target distribution:

Acceptance probability: min(1, p(y) / q(y))
On rejection:           z ~ normalize(max(p - q, 0))
If all accepted:        e ~ p_bonus

Both p and q include their actual temperature/top-k/top-p processing. This rule preserves the target distribution, not identical text under the same seed. Numerical differences between batched and serial BF16 computation require separate validation.

State boundary: retain KV only for the accepted prefix and the round's initial input. The correction/bonus token has no KV yet. Truncate every layer and depth. Apply EOS when committing output, and shorten K to respect output, context, and resource limits.

3. Scheduling and Execution

3.1 Synchronous

Synchronous CPU/GPU timeline

3.2 Asynchronous

Asynchronous CPU/GPU timeline

4. Modules

Module Responsibility
Scheduling Select requests and effective K; reserve verification rows and KV; manage prefill fairness and in-flight rounds
Draft / Target Generate candidates with shallow loops; retain states and reuse them for batched deep-loop execution
Verification and correction Greedy checks, sampling rejection, correction and bonus generation; return valid output and accepted length
State commit Apply EOS/length limits, commit output, truncate KV, and establish the next input
Asynchronous execution Manage round tickets, events, and buffers; preserve result ordering and resource lifetimes

5. Accuracy and Performance Evaluation

5.1 Matched Baselines

Use the same Ouro checkpoint, fixed D=4, last_exited, prefill policy, sampling parameters, and inputs. GPU evaluation uses BF16 + FA4 with fixed hardware/software configuration.

Compare native synchronous vs speculative synchronous, and native asynchronous fixed-depth vs speculative asynchronous separately. Comparing across execution modes measures combined gains, not the contribution of speculation alone.

5.2 Accuracy

Fixed inputs: validate reuse. Compare native serial full-depth execution with the reuse verifier on the same token chain, including positions that will eventually be rejected. Record maximum absolute and RMS errors for loop-2 hidden states, KV at each depth, and final logits, plus the argmax mismatch rate. Use a small FP32 reference to diagnose dependency/indexing errors and real Ouro BF16 to assess numerical behavior.

Free generation: validate output. Run native greedy and speculative greedy with K=1/2/4/8 on the same real prompts. Report exact sequence agreement, the first divergent position, and the native top-1/top-2 logit margin there. Replay a divergent prefix with fixed inputs to distinguish numerical differences from KV/rollback errors. Do not compare logits on different prefixes after divergence.

Sampling: validate distributions. Use analytically tractable small-vocabulary p/q distributions to verify acceptance plus residual correction. Across multiple seeds, report empirical deviations from p and sampling uncertainty. Then cover actual temperature/top-k/top-p settings. Same-seed text equality is not a sampling correctness criterion.

Task evaluations such as GSM8K may additionally compare native and speculative scores using the same dataset and evaluator, with uncertainty reported. Similar task scores do not establish losslessness. Agree on BF16 tolerances and acceptable greedy divergence before acceptance testing.

5.3 Performance

Metric Method
Sweep d=2, D=4; K=1/2/4/8; concurrency 1 and 8 up to near saturation; short/long contexts and outputs
Decode throughput Complete matched prefill for all requests before timing; generate equal numbers of new tokens; divide committed tokens by elapsed time
Decode time Match concurrency and generated work; speedup = T_native / T_spec
Serving latency Under matched arrival workloads, measure end-to-end latency, TTFT, streaming chunk intervals and tokens per chunk; report p50/p95
Cost analysis Record accepted length per round, verification rows, draft/target/verification/commit times, and peak memory; profile bandwidth and MFU separately

Match output length and EOS policy in fixed-work tests; evaluate natural termination separately in serving tests. Include draft, the final shallow fill, target execution, verification, sampling, synchronization, and rollback. Synchronize the GPU at timing boundaries, warm up, repeat at least five times, and report median and variability. Use an exclusive GPU, interleave baseline and speculative trials, and keep profiling outside timed runs.

Prioritize effective throughput and latency when interpreting MFU. Computation on rejected candidates increases executed FLOPs without necessarily improving useful performance. Report results by concurrency, length, and K rather than only the best speedup.

6. Open Questions

  • Keep synchronous rounds atomic, or allow other work between phases?
  • Start asynchronous execution with one in-flight round per request, a CPU commit barrier, and one GPU stream?
  • What sampling scope, default-K policy, and BF16 agreement criteria should the first release adopt?

7. Feature Tracking

⬜ Unclaimed · 🚧 In progress / partial · ✅ Merged for the stated scope · 🔬 Exploratory

Progress refreshed on 2026-10-22 against main at 4a350344027960a48384626aa646f1765fda6d4f. Merged scope is separate from hardware/workload qualification. Existing claims are retained.

Feature Scope Owner Status / evidence
Synchronous self-speculative decoding Fixed loops, compute reuse, greedy/sampling rejection, KV rollback and output integration @bjf-frz ✅ #44; execution shared with the stage runtime through #118.
Asynchronous self-speculative decoding Device-resident candidate chain and cross-round submission/collection @liuyao0322 ✅ #66, cleanup #115 and stage integration #118. Fixed-depth greedy CUDA, LAST_EXITED, refill; compatible Triton/paged FlashAttention backends.
Accuracy and performance evaluation Numerical agreement, serving latency, K sweeps and profiling @Levius-Fubuki 🚧 #47 remains open; retarget/rebase reported complete on Oct 21. #118 adds scoped serving/quality evidence, not universal equivalence.
🔬 Adaptive speculation Select K / draft depth based on load, acceptance and cost @Dmaner 🚧 #84 proposal; open #127 also implements adaptive K. No merged adaptive policy at this baseline; coordinate these scopes. Dynamic draft depth remains separate.
🔬 Loop-level pipelining Batch shallow work for later tokens with deep work for earlier tokens @YuDeng0102 🚧 Draft #133 now exists. #121 is another open pipelined implementation; compare against a shared current-main baseline before choosing integration.
🔬 Device-side execution loop Reduce host barriers in acceptance, stopping and KV management @IDEA-V 🚧 #51 readback reductions and #66/#118 GPU-owned speculative state merged; a fully device-side execution loop is not complete.
CUDA Graph support Capture speculative core/CODA and manage variable shapes @0z5a ✅ #48 synchronous graphs; #118 includes cross-round async core/CODA graphs. Open #127 adds bucketing, derived limits, vectorization and adaptive drafts beyond that foundation.
🔬 Scheduling / preemption compatibility Suspend/resume at safe boundaries @0z5a 🚧 #52 remains open for synchronous committed-round boundaries. Async speculative preemption remains rejected.
PD integration Speculative decode after P→D transfer @jyizheng ✅ Bounded D-side speculation integrated through #118; draft #119 adds real-checkpoint handoff agreement validation. This does not qualify cross-node RDMA or arbitrary feature combinations.
🔬 Proposal improvements Alternate proposals/candidate trees @prettygirlisnotme #59 closed unmerged after the measured negative result. #120 remains open to record the findings; no alternate-branch runtime was merged.

Remaining integration and validation

Versioned RFC document

Activity

  1. changed the title [-][RFC] Fixed-loop self-speculative decoding for RLT[/-] [+][RFC] Self-Speculative Decoding for rlt[/+] on Sep 23, 2026
  2. changed the title [-][RFC] Self-Speculative Decoding for rlt[/-] [+][RFC] Speculative Decoding for rlt[/+] on Sep 23, 2026
  3. liuyao0322 commented on Sep 23, 2026

    @liuyao0322
    Contributor

    I’d like to take ownership of asynchronous self-speculative decoding. could you please assign it to me ?

  4. Levius-Fubuki commented on Sep 23, 2026

    @Levius-Fubuki
    Contributor

    Hi @bjf-frz, I'd like to take ownership of the Accuracy and performance evaluation feature in this RFC.

    My related contributions include merged Triton attention work with numerical and performance validation (vllm#55629) and merged fixes for benchmark clients counting failed streams as successful (vllm-omni#8000, #8015).

    I'll start with the synchronous implementation in #44 and extend the same evaluation contracts to the asynchronous implementation with @liuyao0322 when it is ready:

    • Pin implementation/baseline SHAs, checkpoint, environment and workloads; agree on BF16 acceptance criteria before running acceptance tests.
    • Validate reuse on fixed token chains (shallow states, per-depth KV and logits), greedy divergence with prefix replay, and sampling distributions with small-vocabulary references. Cover rejection/rollback and EOS/length boundaries through the production path.
    • Extend existing tests/benchmarks with matched native-vs-speculative runs: K=1/2/4/8, concurrency and length sweeps, committed-token throughput, serving latency/chunk sizes, acceptance length and peak memory. Profile phase costs and MFU separately from timing; retain raw repeated measurements and report regressions as well as gains.

    My first deliverable will be a gap analysis against #44's existing coverage and a reproducible evaluation matrix. GPU/backend availability will be recorded explicitly, with unrun cases left pending. Please record me as the owner of this feature if this scope fits; I'll coordinate with the implementation owners to avoid duplicating their tests.

  5. bjf-frz commented on Sep 23, 2026

    @bjf-frz
    CollaboratorAuthor

    I’d like to take ownership of asynchronous self-speculative decoding. could you please assign it to me ?

    Of course, thanks for volunteering! RFC list has been updated.

  6. bjf-frz commented on Sep 23, 2026

    @bjf-frz
    CollaboratorAuthor

    Hi @bjf-frz, I'd like to take ownership of the Accuracy and performance evaluation feature in this RFC.

    My related contributions include merged Triton attention work with numerical and performance validation (vllm#55629) and merged fixes for benchmark clients counting failed streams as successful (vllm-omni#8000, #8015).

    I'll start with the synchronous implementation in #44 and extend the same evaluation contracts to the asynchronous implementation with @liuyao0322 when it is ready:

    • Pin implementation/baseline SHAs, checkpoint, environment and workloads; agree on BF16 acceptance criteria before running acceptance tests.
    • Validate reuse on fixed token chains (shallow states, per-depth KV and logits), greedy divergence with prefix replay, and sampling distributions with small-vocabulary references. Cover rejection/rollback and EOS/length boundaries through the production path.
    • Extend existing tests/benchmarks with matched native-vs-speculative runs: K=1/2/4/8, concurrency and length sweeps, committed-token throughput, serving latency/chunk sizes, acceptance length and peak memory. Profile phase costs and MFU separately from timing; retain raw repeated measurements and report regressions as well as gains.

    My first deliverable will be a gap analysis against #44's existing coverage and a reproducible evaluation matrix. GPU/backend availability will be recorded explicitly, with unrun cases left pending. Please record me as the owner of this feature if this scope fits; I'll coordinate with the implementation owners to avoid duplicating their tests.

    Thx for your volunteering! RFC list has been updated. Feel free to share your test plan, btw, no need to prioritize testing the Triton backend for now. You can test directly with FA4.

  7. Dmaner commented on Sep 23, 2026

    @Dmaner

    @bjf-frz I’d like to take Adaptive speculation .could you please assign it to me ?

  8. YuDeng0102 commented on Sep 23, 2026

    @YuDeng0102

    hi @bjf-frz I’d like to take Loop-level pipelining .could you please assign it to me ?

  9. bjf-frz commented on Sep 23, 2026

    @bjf-frz
    CollaboratorAuthor

    @bjf-frz I’d like to take Adaptive speculation .could you please assign it to me ?

    Of course yes, thx for your volunteering! Feel free to share your implementation plan, we could discuss how to achieve Adaptive speculation a better way.

  10. bjf-frz commented on Sep 23, 2026

    @bjf-frz
    CollaboratorAuthor

    hi @bjf-frz I’d like to take Loop-level pipelining .could you please assign it to me ?

    Of course yes, RFC list has been updated.

  11. 0z5a commented on Sep 23, 2026

    @0z5a
    Contributor

    I’d like to take ownership of coda graph support and scheduling and distributed compatibility, please assign them to me if they are still available.

  12. Dmaner commented on Sep 23, 2026

    @Dmaner

    @bjf-frz I’d like to take Adaptive speculation .could you please assign it to me ?

    Of course yes, thx for your volunteering! Feel free to share your implementation plan, we could discuss how to achieve Adaptive speculation a better way.

    ​thx, I'll draft an implementation plan asap for feedback

  13. bjf-frz commented on Sep 23, 2026

    @bjf-frz
    CollaboratorAuthor

    I’d like to take ownership of coda graph support and scheduling and distributed compatibility, please assign them to me if they are still available.
    Thx for your volunteering! Cuda graph has been proved to be a valid method to improve performance in previous work. Happy to see what will happen on your job.

  14. 0z5a commented on Sep 23, 2026

    @0z5a
    Contributor

    I’d like to take ownership of coda graph support and scheduling and distributed compatibility, please assign them to me if they are still available.
    Thx for your volunteering! Cuda graph has been proved to be a valid method to improve performance in previous work. Happy to see what will happen on your job.

    ok i see

  15. Levius-Fubuki commented on Sep 23, 2026

    @Levius-Fubuki
    Contributor

    Thanks! I’ll focus on BF16 + FA4, review #44’s existing coverage, and share a concise test plan.

  16. IDEA-V commented on Sep 23, 2026

    @IDEA-V

    Hi @bjf-frz, I’d like to take ownership of Device-side execution loop, if it’s still available.

    My background is in GPU kernel optimization and efficient inference. I’ll start by profiling the CPU/GPU synchronization overhead in #44 and assessing the benefit of moving acceptance, stopping, and KV updates onto the GPU, then share an implementation plan. I’ll coordinate with @liuyao0322 and @0z5a on the asynchronous execution and CUDA Graph interfaces.

    I’ll primarily develop on an A100, with access to a B200 for subsequent BF16 + FA4 validation and performance measurements.

    Could you please record me as the owner if this scope fits? Thanks!

  17. prettygirlisnotme commented on Sep 23, 2026

    @prettygirlisnotme

    Hi @bjf-frz, I'd like to take on the Proposal improvements direction, starting with the feasibility and benefit assessment.

    I'll start from the current self-speculative implementation and characterize where the shallow proposals disagree with the full-depth target. Then I'll test lightweight ways to improve proposal quality while keeping K and draft depth fixed initially, so this stays separate from adaptive speculation.

    I'll evaluate both acceptance improvement and end-to-end decoding cost, since a proposal change is only useful if the added overhead translates into real latency or throughput gains.

    If the feasibility results are promising, I'll narrow this into a concrete first implementation. Please let me know if someone is already working on this direction.

  18. bjf-frz commented on Sep 23, 2026

    @bjf-frz
    CollaboratorAuthor

    Hi @bjf-frz, I’d like to take ownership of Device-side execution loop, if it’s still available.

    My background is in GPU kernel optimization and efficient inference. I’ll start by profiling the CPU/GPU synchronization overhead in #44 and assessing the benefit of moving acceptance, stopping, and KV updates onto the GPU, then share an implementation plan. I’ll coordinate with @liuyao0322 and @0z5a on the asynchronous execution and CUDA Graph interfaces.

    I’ll primarily develop on an A100, with access to a B200 for subsequent BF16 + FA4 validation and performance measurements.

    Could you please record me as the owner if this scope fits? Thanks!

    Sounds good. RFC list has been updated.

  19. bjf-frz commented on Sep 23, 2026

    @bjf-frz
    CollaboratorAuthor

    Hi @bjf-frz, I'd like to take on the Proposal improvements direction, starting with the feasibility and benefit assessment.

    I'll start from the current self-speculative implementation and characterize where the shallow proposals disagree with the full-depth target. Then I'll test lightweight ways to improve proposal quality while keeping K and draft depth fixed initially, so this stays separate from adaptive speculation.

    I'll evaluate both acceptance improvement and end-to-end decoding cost, since a proposal change is only useful if the added overhead translates into real latency or throughput gains.

    If the feasibility results are promising, I'll narrow this into a concrete first implementation. Please let me know if someone is already working on this direction.

    You can check this paper “ LoopSpec: Pipelined Self-Speculative Decoding for Looped Transformers”, this is an open problem. Our goal is to reduce the average depth and further improve throughput. RFC has been updated.

  20. IDEA-V commented on Sep 24, 2026

    @IDEA-V

    Initial profiling for the Device-side execution loop is up in #51, a draft PR that includes two small sync reductions and the profiling harness. In short:

    • Eager speculation is launch-bound: about 1,400 kernels per core pass and 13–16% GPU busy. Per-candidate syncs matter little there.
    • With #48's CUDA Graphs, host syncs are 5–9% of decode time (B≤8, B200 MIG). The PR removes most of them: 73–110 → 5 syncs per round, +2.5% at B=1 and +5.3% at B=8 on top of #48, with outputs token-identical.
    • Its remaining upper bound in this setting is only about 3%; The share would grow if graph execution gets faster.
  21. prettygirlisnotme commented on Sep 24, 2026

    @prettygirlisnotme

    I opened a draft PR at #59.
    On the frozen prototype, the matched-depth throughput changes were +5.8%/+5.3%/−1.1% for d1 at C1/C4/C8, and +4.1%/+4.2%/+6.5% for d2. The current port is opt-in, greedy/eager only, with focused correctness and KV lifecycle tests.
    Still open are the BF16 divergence, FA4 and natural-EOS/online-serving validation, and a fresh performance rerun on the final port.
    Would appreciate feedback on the scope and KV integration, especially wrt #48/#51.

  22. jyizheng commented on Sep 24, 2026

    @jyizheng

    Hi @bjf-frz @0z5a, I'd like to take the PD integration part of Scheduling and distributed compatibility, if @0z5a is OK with splitting it off. It looks like #52 now covers scheduling and preemption only, since #49's PD decode handoff was dropped.

    Scope: speculative decode on the D side after NIXL handoff (#31). That means defining which KV depths and hidden state P must hand over so D can start round 1 with the fixed d/D, and handling handoff alongside in-flight rounds, rollback, and cancellation. I'll build on #52's scheduling interfaces rather than duplicate them.

    First deliverable: a short design note on the handoff contract (valid KV range per depth, first-token/bonus semantics), plus a P→D greedy agreement test against colocated speculative decode. I have H100/B200 nodes with RDMA for the disagg runs.

    Please record me as owner if this split works.

  23. 0z5a commented on Sep 25, 2026

    @0z5a
    Contributor

    Hi @bjf-frz @0z5a, I'd like to take the PD integration part of Scheduling and distributed compatibility, if @0z5a is OK with splitting it off. It looks like #52 now covers scheduling and preemption only, since #49's PD decode handoff was dropped.

    Scope: speculative decode on the D side after NIXL handoff (#31). That means defining which KV depths and hidden state P must hand over so D can start round 1 with the fixed d/D, and handling handoff alongside in-flight rounds, rollback, and cancellation. I'll build on #52's scheduling interfaces rather than duplicate them.

    First deliverable: a short design note on the handoff contract (valid KV range per depth, first-token/bonus semantics), plus a P→D greedy agreement test against colocated speculative decode. I have H100/B200 nodes with RDMA for the disagg runs.

    Please record me as owner if this split works.

    Yes feel free to cherry-pick if you find something useful.

  24. Dmaner commented on Oct 4, 2026

    @Dmaner

    @bjf-frz I’d like to take Adaptive speculation .could you please assign it to me ?

    Of course yes, thx for your volunteering! Feel free to share your implementation plan, we could discuss how to achieve Adaptive speculation a better way.

    Hi @bjf-frz, I opened a RFC on adaptive speculation in #84. Any feedback would be greatly appreciated whenever you have time!

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions