Repository navigation
[RFC]: Adaptive early exit on Ouro-1.4B: accuracy–throughput evaluation and exact-policy async exit #56
Description
Motivation.
Adaptive early exit is the main reason to run a looped model through vllm-rlt, but its throughput and its accuracy have been measured separately, under different settings. Baseline: main @ 3314c1b3.
Throughput is measured; accuracy at those settings is not.
- PR [feat]Implement depth-aware KV and asynchronous decode scheduling #30 measures early exit at
exit_threshold=0.2withmin_loops=2andmax_loops=4. At the highest concurrency tested, E2E throughput improves by about +11.9% / +15.9% / +18.7% / +14.2% (P2/P1) for 1K / 4K / 16K / 64K contexts, with about 2.4–2.5 mean decode rounds. - The same PR notes that the delayed variant (P3, required by async scheduling) runs about 3.3–3.4 mean rounds, so the extra round "still consumes a substantial part of the original early-exit benefit". It also says that "policy-changing comparisons do not establish equal model quality".
- The only task-accuracy result covers fixed four-loop greedy decoding: GSM8K 62.02% native vs. 61.64% Transformers on all 1,319 test questions (
docs/accuracy.md). The harness enforces this:benchmarks/gsm8k_backends.pyraises if any exit depth is not 4.docs/cdb_runtime.mdstates that no task-accuracy result is claimed forouro_delayed.
So we cannot yet say how much accuracy the measured speedups cost, or which threshold is a reasonable default for users.
Output differences under the same exit policy are still unresolved. In the PR #30 campaign, 196 of 612 trials differed from their own policy reference (182 in tokens, 168 in exit depths). docs/cdb_runtime.md also records a BF16 multi-stream token divergence whose root cause is unconfirmed: it disappeared when diagnostic logit copies were added. The PR itself lists investigating same-policy differences as a step to take before expanding the matrix. Until that is resolved, accuracy differences under async execution cannot be cleanly attributed to the exit policy.
Why this is separable. The released Transformers model already supports exit_threshold. It runs all four loops, keeps full-depth KV for every token, and only selects the exited loop's hidden state (modeling_ouro.py @ 574fa66c). The native engine really skips loops and uses last_exited KV. Comparing the two at the same threshold therefore isolates the effect of the KV approximation from the effect of the early-exit decision itself.
Proposed Change.
The work has three phases. Phases 1–2 measure and diagnose, and leave runtime behavior, exit-policy defaults and public APIs unchanged. Phase 3 is a runtime change that is proposed now but only implemented if the phase-1 data supports it.
Phase 1: accuracy–compute trade-off (GSM8K, BF16, greedy)
Arms. Each step changes exactly one factor:
| Arm | Backend | Exit policy | KV | Isolates |
|---|---|---|---|---|
| A0 | HF | fixed 4 loops | full depth | reference (existing baseline) |
| A1 | HF | exit_threshold=τ |
full depth | early-exit decision alone |
| B1 | native sync | ouro, τ |
last_exited |
+ KV approximation (B1 vs A1) |
| B2 | native sync | ouro_delayed, τ |
last_exited |
+ one-round delay (B2 vs B1) |
| B3 | native async | ouro_delayed, τ |
last_exited |
+ async scheduling (B3 vs B2, expected to be 0) |
- Thresholds: τ ∈ {1.0 (control), 0.9, 0.7, 0.5, 0.2}. 0.2 is the PR [feat]Implement depth-aware KV and asynchronous decode scheduling #30 operating point.
min_loops=1in the parity arms, because HF has no minimum. The project default,min_loops=2, is reported as an extra native row.- Dataset and prompts: the existing pinned protocol (lm-eval
gsm8k_cot, 3-shot, strict match, 1,024 new tokens, 2,048 context). First screen every arm on the fixed GSM8K-87 case, then run the full 1,319-question test split for the thresholds that matter. - Metrics per arm:
- strict accuracy, plus paired gains and losses against A0 on the same questions;
- unparseable-answer and length-termination counts;
- mean exit depth and depth histogram;
- generated tokens per question;
- decode loops per question, the sum of per-token depths. This is a hardware-independent compute measure; it also catches the case where early exit makes answers longer and cancels out the per-token savings.
- Measured throughput: questions/s and output tokens/s at concurrency 1 / 8 / 32 on the same GPU, for B1–B3 against the τ=1.0 control, interleaved in shuffled order with three repetitions (min–max reported). Absolute numbers are not comparable with PR [feat]Implement depth-aware KV and asynchronous decode scheduling #30 (different GPU and backend); only within-run ratios are reported.
Deliverables:
- PR: extend
benchmarks/gsm8kwith exit controls: exit mode, threshold, min/max loops, and async scheduling. These go into the protocol fingerprint so baselines cannot be mixed. The PR also adds per-token exit-depth recording, the HFexit_thresholdarm, and depth/compute statistics plus paired statistics incompare. Validated on CPU with the tiny model. - Report (
docs/benchmarks/…): the accuracy vs. decode-loops Pareto curve per arm, with pinned SHAs, environment, raw outputs and checksums, following the existing accuracy-doc conventions. It will propose, not decide, a default threshold.
Phase 2: same-policy output differences
- Reproduce the BF16 multi-stream divergence on the real checkpoint with a fixed prompt set.
- At the first divergent position, record the reference top-1/top-2 logit margin:
- near-ties point to reduced-precision or schedule-dependent numerics;
- large margins point to a state, ordering or stream-dependency problem.
- Bisect across the P4–P8 features (async single-stream → multi-stream → static buffers → padding → graphs), holding the exit policy fixed.
- Deliverable: a diagnosis report. If a real ordering bug is found, a fix PR with a regression test, kept separate from any structural change.
Phase 3 (data-gated): exact-policy exit under async scheduling
Problem. Async scheduling currently requires ouro_delayed. By the time round k's gate score reaches the host, round k+1 has already been submitted, so every exiting token runs one extra round. Its output also follows a different policy from ouro. This is the P2→P3 confound in PR #30: an async-vs-sync comparison also changes model semantics.
Candidate design, an opt-in exit mode (working name ouro_exact) that keeps the async pipeline but restores the original policy's outputs:
- After each recurrent round, snapshot the per-request hidden state on the device, into one extra
[max_num_seqs, hidden]pool indexed by the existing async request slot. No host transfer is involved. - When round k's signal triggers an exit, run coda from the round-k snapshot instead of the round-(k+1) hidden state, and record exit depth k.
- Finalize KV from depth k.
finalize_tokenalready copies the exit depth's K/V into every deeper plane, so the K/V that round k+1 wrote at this position is overwritten, and later tokens see exactly thelast_exitedstate that syncourowould have produced. - The extra round is still executed, so compute matches
ouro_delayed. The gain is semantic: outputs and depths should match syncouro, so async can be compared with sync without a policy confound, and it inheritsouro's measured accuracy.
When to implement:
- Go ahead if phase 1 shows a material B2-vs-B1 difference in accuracy or outputs.
- If the difference is negligible, report that and drop or deprioritize this phase.
Open points to settle in design review: where the snapshot lives under CUDA Graph capture; the finalize_many async path; padding rows; and preemption snapshots, which would need to include the extra hidden state.
Acceptance:
- Token and depth agreement with sync
ouroon the phase-1 prompt set, under FP32 on the tiny model and the real checkpoint, with BF16 differences reported against the phase-2 findings. - GSM8K accuracy matching arm B1 within the agreed tolerance.
- Throughput and peak memory against
ouro_delayedasync under matched conditions. - CPU and GPU regression tests for exit, EOS, abort and preemption.
This touches exit handling that M3 (#32) is moving from the engine into the scheduler. I would implement it against whatever interface M3 settles on, coordinating with @bjf-frz on ordering. The implementation would not be written in parallel with that migration.
Non-goals. No trained lookahead gate, no paper reproduction, and no change to the existing ouro / ouro_delayed semantics or defaults (phase 3 adds an opt-in mode). No overlap with the speculative-decoding evaluation in #43; the per-token depth and acceptance data may still be useful to the adaptive-speculation work there. A fully device-side exit decision is out of scope and relates to #43's device-side execution loop.
Hardware. Phase 1 uses the Triton backend on a rented single GPU (L40S or H100 class). The exact GPU, driver, package versions and CPU affinity are recorded, and cases that were not run are marked as such. FA4 arms can be added if maintainers want them for comparability with PR #30.
Feedback Period.
One week. I will start the CPU-only harness PR (deliverable 1) in the meantime, since it doesn't depend on the answers below.
Questions for maintainers:
- Is τ=0.2 with
min_loops=2the intended operating point, or should the sweep center elsewhere? - Is the full 1,319 split acceptable as the headline result, with GSM8K-87 as the screen?
- Is Triton on H100/L40S acceptable for the accuracy arms, or do you need FA4 for comparability?
- Should phase 2 be tracked here or in a separate issue?
- For phase 3: is an opt-in exit mode the right shape, or would you prefer this as the default async behavior for
ouro? And should it wait for the M3 exit-policy migration, or land first behind the current engine interface?
CC List.
Any Other Things.
I'm happy to own this end to end. If someone is already looking at same-policy differences or adaptive-exit accuracy, I'll coordinate rather than duplicate.
Activity
bjf-frz commentedon Sep 24, 2026
ouro_delayed is only a workaround for testing async scheduling, since we do not have weights trained for a true delayed-exit policy. Its accuracy is therefore expected to be unreliable and does not need to be evaluated as a meaningful model-quality baseline.
The other evaluations are still useful, especially scheduling correctness, sync/async consistency, exit depth, compute, and throughput.
MinhaoLi0318 commentedon Sep 25, 2026
Thanks @bjf-frz, that makes sense. I'll treat ouro_delayed as a scheduling-test policy, not a quality baseline. Revised phase 1:
Accuracy (GSM8K): only the fixed-depth HF reference, HF with exit_threshold, and native sync ouro, across thresholds. ouro_delayed gets no accuracy claim.
ouro_delayed, sync vs. async: token and exit-depth agreement on the same prompts, recording the first divergent position and the logit margin there whenever the two differ. Also exit-depth distribution, decode loops per token and per question, and throughput. This is also where phase 2 starts.
This changes the case for phase 3. Async scheduling currently accepts only ouro_delayed, random_lookahead or trace, so no async exit policy uses the trained gate with its original semantics. ouro_exact would fill that gap: the trained gate's ouro decision under async scheduling, with outputs expected to match sync ouro. The extra round would still run, so compute stays at the ouro_delayed level. My original trigger for phase 3 (an accuracy gap between ouro_delayed and ouro) no longer applies. Would you be open to this as the phase-3 target on its own merits?
Two questions before I rent a GPU:
Is τ=0.2 with min_loops=2 the operating point you care about, or should I sweep {0.9, 0.7, 0.5, 0.2}?
Is Triton on an L40S/H100 acceptable, or do you want FA4 for comparability with PR #30?
The harness support for both parts is in #58 (validated on CPU).
bjf-frz commentedon Sep 25, 2026
bench: add adaptive-exit options and depth statistics to the GSM8K ha…
MinhaoLi0318 commentedon Oct 7, 2026
@bjf-frz thanks, all five points are adopted. Status:
config.json), so the comparison runs on main.Meanwhile my vllm-rlt time goes to the #70 Phase 1 rollout outputs (#83).
AI assistance: I used Claude Code to check the thread and the harness state and to draft this update; I reviewed it before posting.
MinhaoLi0318 commentedon Oct 22, 2026
@bjf-frz Before the Phase 1 run (one A100, 10-24), two things fixed in advance:
The per-token logit margins I mentioned earlier move to the Phase 2 issue, which I'll open with this run's outputs. Tables here by 10-25.
AI assistance: Claude Code helped check the harness and draft this comment; I reviewed it before posting.
MinhaoLi0318 commentedon Oct 25, 2026
Phase 1 results: the accuracy, exit-depth, compute and throughput tables promised on 10-07.
Setup. vllm-rlt
09f229a(the #58 merge, pre-#118 engine),ByteDance/Ouro-1.4B@574fa66c, 1× A100-SXM4-80GB, torch 2.13.0+cu130, Triton 3.7.1, BF16, greedy. Native arms use the Triton attention backend, HF arms the release model (eager). Pinnedgsm8k_cotprotocol (3-shot, strict match); B0 passed the stored GSM8K-87 gate (60/87 vs 59/87). Depth and loops count decode tokens only.Full 1,319 split
A0 / B0: HF / native, fixed 4 loops. B1m2_0.2: native sync
ouro, τ=0.2, min_loops=2 (primary point). Dsync_0.2 (B2): native syncouro_delayed, same τ and min_loops; its accuracy is used only in the rule below, not as a quality baseline.† The HF release computes all 4 loops per token. Unparseable: no strict-match answer (e.g.
The answer is $18.).At the primary point, early exit runs 44.6% fewer decode loops per question than fixed depth, with answers about as long (111.1 vs 109.9 tokens/q), and costs 7.35 pp. B2's extra round leaves a 20.4% saving.
ouro_exact rule (pre-registered 10-22): met
B2 − B1 = +5.76 pp [95% CI +3.7, +7.8], lost/gained 56/132, exact two-sided McNemar p = 2.95e-08 (needed: p < 0.05 and |Δ| ≥ 1 pp). Loops/q 241.5 (B1) vs 346.8 (B2), +0.99 loops/token on shared prefixes; outputs differ on 1,076 of 1,319 questions.
This is not a quality claim for
ouro_delayed. It shows thatouro_delayed, which async scheduling requires today, changes what the trained gate produces. As designed,ouro_exactwould restore B1's outputs (55.42% here) while still running the extra round, so its case is sync/async equivalence, not accuracy.Throughput (Run 1, matched work)
Requests cycle through the GSM8K-87 prompts, and each decodes exactly screen B0's natural length for its prompt (
ignore_eos), so every policy does the same token work. 16/32/128 requests at C=1/8/32, fresh engine per cell, 3 reps with the cell order shuffled; median output tok/s (×fixed; req/s ratios are identical). Every cell's 3-rep spread is within 5% of its median. Policies with τ<1 use min_loops=2.Gains at C=1 / 8 / 32 (median ratios):
fixed_async/fixed): 0.997 / 0.992 / 0.995The speedup comes from depth alone; async and multi-stream overlap add nothing on this setup.
ouroτ=0.2 gets 1.757× at C=1 (depth ratio 4/2.27 = 1.76×) and 1.478× at C=32 (depth ratio 1.79×).Consistency: 0 divergent requests at every C for delayed τ=0.2 sync vs async (multi- and single-stream),
fixedvsfixed_async, and rep 0 vs rep 2 of every policy; in the harness, Dsync and Dasync match on 87/87 at τ=0.5 and 0.2. Divergence appears only where batch composition changes (the BF16 batch-variance baseline): underfixed, C=1 vs C=8 9/48 requests, C=1 vs C=32 18/48, C=8 vs C=32 27/96, repeated prompts within a C=32 cell 36/123.GSM8K-87 screen (informational)
min_loopsmatters only at τ=0.2 (min 1 and min 2 give identical tokens and depths on 87/87 at τ ≥ 0.5). There, min_loops=1 loses 16.09 pp vs B0 (p = 0.00936) and answers grow to 129.9 tokens/q (B0: 103.8).last_exitedKV effect: it also includes BF16-vs-FP32 exit arithmetic and eager-vs-Triton kernels.All 18 arms
Δ: A1 vs A0, B1/B1m2 vs B0. "=": identical outputs on 87/87. † HF: selected exit loop; all 4 loops still computed. B2 − B1: τ=0.5 +4.60 pp (p 0.344), τ=0.2 +9.20 pp (p 0.0574).
Deviations
Not claimed: carry-over to current main (these are pre-#118 numbers), and logit margins, which go to the Phase 2 issue I'll open next with this run's per-request outputs and batch logs.
@bjf-frz The pre-registered rule is met, with B2 above B1. Should
ouro_exactgo ahead as the opt-in async mode once M3 settles the exit-policy interface, or, since it would return async to B1's outputs at about B2's compute, would you rather deprioritize Phase 3?AI assistance: I used Claude Code to drive the pod run, recompute the statistics from the per-question samples and draft this comment; I checked the numbers before posting.