[RFC]: Support Adaptive Speculation for Ouro #84

Progress update — 2026-10-22

Baseline: main 4a350344027960a48384626aa646f1765fda6d4f.

• Main has fixed-depth cross-round async speculation and recurrent/CODA graphs through [FEAT] add cross-round async self-speculative decoding #66/[Architecture] Unify stage execution with loop-level scheduling and mixed batching #118, but no adaptive draft-policy option at this baseline.
• Open [FEAT] Add CUDA Graphs for cross-round async self-speculation #127 now includes adaptive per-round K (including zero drafts), cost/acceptance feedback, graph bucketing and batch-operation optimizations. Its documented runtime changes adaptive drafting to the async default; that policy choice still needs maintainer confirmation.
• Coordinate this issue's K/controller scope with [FEAT] Add CUDA Graphs for cross-round async self-speculation #127 before implementing a parallel controller. Dynamic draft depth, workload-drift handling and broader model coverage remain distinct proposals.
• [FEAT] Add CUDA Graphs for cross-round async self-speculation #127's manual H100 evidence applies to its recorded branch; it needs reconciliation with [Architecture] Unify stage execution with loop-level scheduling and mixed batching #118/[Refactor] Move runner result application to Scheduler (M2/R1) #107 and does not constitute current-main qualification.
• The candidate-selection/EMA plan below remains a proposal, not a merged or selected design. Existing ownership is unchanged.

Motivation.

Related #43

Support Adaptive Speculation for Ouro, Select K / draft depth d based on load and acceptance

Proposed Change.

A few prior observations
Sampling experiments use H100, Ouro-1.4B, BF16, FA4, CUDA Graph ON, target D=4, L=128/512, and N=64.

Speculative Decoding Speedup Heatmap (K/d)

Performance varies with workload: low concurrency can benefit from a larger K, while a smaller K is more suitable at high concurrency. Candidate configurations should be selected separately for each B range.

Draft Acceptance Heatmap (K/d)

The effects of K/d on acceptance rate: with d fixed, increasing K generally lowers the pure draft acceptance rate; with K fixed, increasing d generally raises acceptance.

Speedup Ratio vs. Acceptance Rate

The effects of K and d on throughput depend on the workload. At low concurrency, a moderate increase in K can increase the number of tokens effectively committed in one verification round. At high concurrency, the additional draft and verification work from a large K is more likely to offset the benefit of fewer decoding rounds. The tradeoff for d also depends on workload: in these experiments, B1 favors a shallower draft, while some B8/B32 combinations perform better at d=2 and lose benefit as d increases further.
A high acceptance rate does not guarantee high throughput. For example, at L512/B8/K1, increasing d from 2 to 3 raises acceptance from 90.8% to 94.9%, but the decode performance ratio drops from 1.20 to 1.09. A deeper draft increases loop cost as well as acceptance, so at high concurrency, higher acceptance can still accompany a reduction in speculative benefit.

Idealized model of TPS and acceptance rate

Theoretical execution time for an idealized round of speculative decoding

$$ T_{\mathrm{spec}}(B,L,K,d)\approx d(K+1)c_{\mathrm{shallow}}(B,L)+(D-d)c_{\mathrm{verify}}\left(B(K+1),L\right)+T_{\mathrm{other}} $$

For a simple idealized calculation, assume $c_{\mathrm{shallow}}=c_{\mathrm{verify}}=c$ and ignore $T_{\mathrm{other}}$.

$$T_{\mathrm{spec}}=cd(K+1)+c(D-d)=c[dK+d+D-d]=c(D+dK)$$

The corresponding idealized speedup and predicted throughput are:

$$ S_{\mathrm{AR}}=\frac{D(1+K\alpha)}{D+dK} $$

$$ \widehat{\mathrm{TPS}}=\frac{B(1+K\alpha)}{T_{\mathrm{spec}}(B,L,K,d)} $$

Here, $\alpha$ is the pure draft acceptance rate.

The measurements, interpreted through the idealized cost model, suggest two conclusions:
1. Higher concurrency reduces speculative gains and favors smaller speculative budgets.
The observed trend is consistent with reduced batching benefits and higher costs from expanded verification. Smaller $K$ limits wasted candidate computation. A shallower $d$ can also help when its time savings outweigh the loss in acceptance; the best depth remains workload-dependent.
2. The correlation between acceptance rate and throughput depends on workload.
Across different $(K,d)$ configurations, higher acceptance can correlate with either higher or lower throughput. At fixed $B,L,K$, the cost model gives $\widehat{\mathrm{TPS}}_1/\widehat{\mathrm{TPS}}_0=\frac{1+K\alpha_1}{1+K\alpha_0}\cdot\frac{T_0}{T_1}$: throughput falls when round time grows faster than useful output.
At $B=8,\ L=512,\ K=1$, increasing $d$ from 2 to 3 raised acceptance from 90.8% to 94.9%, while decode speedup fell from 1.20× to 1.09× ( - 0.91). With $D=4$, the idealized model predicts $\widehat{\mathrm{TPS}}_1/\widehat{\mathrm{TPS}}_0=\frac{1+0.949}{1+0.908}\cdot\frac{4+2}{4+3}\approx0.88<1$. The 2.1% increase in useful output is outweighed by the 16.7% increase in modeled round time, explaining why the acceptance gain can fail to improve throughput.

Implementation Plan

Phase 1

Based on the pre-observation data above, we plan to use the following approach:

1. Offline sampling

Perform offline sampling for different workloads to obtain multiple candidate (K,d) ranges. Use them as the default configuration for Ouro in each batch-size range and as the candidate range for subsequent tuning.

2. Online Optimization: dynamically adjust based on candidate configurations

Use EMA (Exponential Moving Average) to dynamically adjust within the candidate range, with the draft acceptance rate as the feedback metric:

$$ AR_{raw}=\frac{\sum_i A_i}{\sum_i K_{actual,i}} $$

$$ E^{AR}_t=(1-\beta)E^{AR}_{t-1}+\beta AR_{raw,t} $$

Adjust dynamically based on E, using thresholds to determine whether to move to a higher or lower configuration. Here, j denotes the candidate position.

$$ j_{\text{next}} = \begin{cases} j + 1, & E_j > \theta_{B, j}^{up} \\ j - 1, & E_j < \theta_{B, j}^{down} \\ j, & \text{others} \end{cases} $$

For small batches, consider increasing the number of candidates and the draft depth to improve acceptance while keeping throughput unchanged.

For large batches, consider reducing the number of candidates and the draft depth to reduce the extra workload introduced by speculation. Primarily reduce K and d, and consider increasing K or d only when acceptance falls below a certain threshold.

Example parameter configuration based on the current experimental data:

Phase 2

Support adaptively disabling Speculative Decoding. Combine offline-measured benefit boundaries with the online batch size and acceptance EMA to switch to ordinary decoding when speculative benefit is insufficient. After disabling it, use exponential backoff to reduce probe frequency. Re-enable speculative decoding when the workload changes or a probe finds that benefit has recovered.

Phase 3

Extract a ModelAdapter to support Looped Transformers such as Ouro, Nanbeige4.2, and Huginn. Encapsulate loop-depth control and early-exit behavior for different models, reuse the existing candidate-selection and EMA-feedback logic, and provide each model with default candidate configurations generated through offline sampling.

Architecture and Implementation Outline

Overall Architecture

Core Structures

Overall flow: select a configuration by batch size → execute one speculative-decoding round → collect acceptance feedback → adjust the configuration for the next round. Offline experiments first determine the available (K,d) candidates for different batch sizes. Online adjustment then uses the draft acceptance rate to move within the corresponding candidate list.

1. Select the configuration for this round. AdaptiveController obtains the number of requests currently waiting for speculative scheduling from Scheduler, finds the corresponding batch-size range, and retrieves its current (K,d). The candidates and initial selections for each range are stored in AdaptiveSpecConfig. The selected configuration is written to SpeculationPlan and remains fixed during this round.

2. Generate candidates and verify them. Scheduler allocates the token budget according to the plan and passes the batch and plan to SpeculativeRunner (the runner). Draft first generates candidate tokens, and Target then verifies which candidates can be accepted. The runner returns the actual proposed and accepted counts for each request.

3. Commit the results and return feedback. LLMEngine first saves the verification counts, then commits the output and updates the KV cache. After all processing succeeds, it combines the counts with the current plan identifier into RoundFeedback and sends it to the controller. The controller uses this information to identify the range and candidate used in this round and passes the feedback to the policy module.

4. Calculate and smooth the acceptance rate. CandidatePolicy (the policy module) uses the draft acceptance rate. It then smooths this ratio with EMA (Exponential Moving Average) to reduce fluctuations between rounds. Each range and candidate maintains its own EMA, with the state stored in BucketState.

$$ r_t=\frac{\sum_i A_i}{\sum_i K_{\mathrm{actual},i}},\qquad E_t=(1-\beta)E_{t-1}+\beta r_t $$

The formula above applies to the same candidate within the same range: A is the number of accepted draft tokens, K_actual is the actual proposed count, r is the acceptance rate for this round, and E is the smoothed acceptance rate. β is the weight of new feedback; a smaller β produces a smoother result but a slower response.

5. Decide the configuration for the next round. After enough valid feedback has accumulated, the policy module checks the EMA at the configured interval: above the upper threshold, move to the next candidate; below the lower threshold, move to the previous candidate; between the two thresholds, keep the current candidate. Each adjustment moves one position and stays within the current range’s candidate list. The controller writes this selection into the next batch’s plan.

Test Plan

A single H100 SXM 80GB, Ouro-1.4B (BF16, D=4, greedy, CUDA Graph). Use GSM8K, HumanEval, and WikiText as performance-test inputs.

Build two non-overlapping groups of performance-test inputs from GSM8K, HumanEval, and WikiText: one for offline candidate-configuration selection, and the other for comparative evaluation after the configuration is frozen.

• Offline sampling: Sample different (K,d) configurations at B=1/8/32, screen them with short outputs, confirm them with long outputs, and generate and freeze the candidate configurations and EMA parameters for each batch-size range.
• Online evaluation: Use 512 input tokens and 2048 new output tokens per request. Test fixed concurrency B=1/8/32 and a changing workload of 1→8→32→8→1.

Feedback Period.

one week.

CC List.

@bjf-frz @hsliuustc0106 when you have a moment, please let me know if this approach makes sense or if there is room for improvement. If it looks good, I will proceed with the implementation quickly. Any alternative ideas or suggestions are also welcome!

Anything else.

Why use draft acceptance rate as the feedback?

My rationale mainly consists of the following three points:

1. The metric is closely related to speculative decoding.
2. It is relatively simple to collect accurately.
3. The experiments also indicate that a good acceptance rate corresponds to a better speedup in most batch-size ranges.

However, a higher acceptance rate does not necessarily mean faster execution. For example, a deeper draft can improve acceptance but also takes more compute time. At high concurrency, this extra overhead may instead reduce speculative benefit. We therefore first test (K,d) at different batch sizes and select candidates with better throughput, maintaining a good acceptance rate while ensuring that throughput is not affected too much by speculative decoding.

If parameters are adjusted directly according to throughput increases or decreases, changes caused by other factors can easily be attributed to speculation. For example, more long requests, a higher prefill share, slower CPU scheduling, or GPU contention from other tasks can all reduce throughput even when the current speculative configuration is still appropriate.

Limitations

• Hardware dependence: The default Ouro adaptive-speculation preset is based on empirical measurements from a single H100 GPU. For multi-GPU deployments or GPUs with different specifications, we recommend running offline sampling on the target hardware to determine suitable candidate (K,d) ranges, or providing custom candidate ranges.
• Stable workloads: When the workload remains stable, we recommend using a fixed (K,d) configuration to avoid the additional overhead of adaptive control.

Before submitting a new issue...

• Make sure you already searched for relevant issues in the issue tracker, and looked for the answer in the documentation.