[RFC]: Adaptive early exit on Ouro-1.4B: accuracy–throughput evaluation and exact-policy async exit

[MinhaoLi0318]

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.2 with min_loops=2 and max_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.py raises if any exit depth is not 4. docs/cdb_runtime.md states that no task-accuracy result is claimed for ouro_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=1 in 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:

1. PR: extend benchmarks/gsm8k with 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 HF exit_threshold arm, and depth/compute statistics plus paired statistics in compare. Validated on CPU with the tiny model.
2. 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:

1. 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.
2. 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.
3. Finalize KV from depth k. finalize_token already 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 the last_exited state that sync ouro would have produced.
4. The extra round is still executed, so compute matches ouro_delayed. The gain is semantic: outputs and depths should match sync ouro, so async can be compared with sync without a policy confound, and it inherits ouro'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 ouro on 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_delayed async 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:

1. Is τ=0.2 with min_loops=2 the intended operating point, or should the sweep center elsewhere?
2. Is the full 1,319 split acceptable as the headline result, with GSM8K-87 as the screen?
3. Is Triton on H100/L40S acceptable for the accuracy arms, or do you need FA4 for comparability?
4. Should phase 2 be tracked here or in a separate issue?
5. 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.

@bjf-frz @hsliuustc0106

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.

[bjf-frz]

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]

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]

1. τ=0.2, min_loops=2 is fine as the primary operating point. Keep the sweep for reference.
2. GSM8K-87 for screening and the full 1,319 split for final validation sounds good.
3. No need to stick to FA4. Feel free to use or adapt the attention backend that best fits the available GPU; just keep it consistent within each comparison.
4. Please track Phase 2 in a separate issue. The BF16/multi-stream divergence is a standalone correctness/debugging problem.
5. Keep ouro_exact opt-in and data-gated. Only proceed if Phase 1 shows a meaningful B1-vs-B2 difference, and implement it after M3 settles the exit-policy interface.