Repository navigation
[FEAT] Add opt-in single-fallback proposals for greedy speculation - #59
prettygirlisnotme wants to merge 4 commits into
Conversation
prettygirlisnotme
left a comment
There was a problem hiding this comment.
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.
|
Remove the changes to the README and move this section to the speculative decoding documentation. |
| 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 |
There was a problem hiding this comment.
Rename fallback_margin to alternate_prob_gap_threshold to make its meaning explicit.
|
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:
2. Performance Use the current PR revision, comparing
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:
3. Numerical Accuracy and Output Quality Please investigate the existing BF16 output differences:
4. Profiling Follow the Focus on:
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. |
|
Moved the feature section from README to 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:
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:
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 |
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 defaultNonekeeps the plain runner. The implementation supports one branch per request per round withlast_exitedKV, 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 - 1verification rows for a primary chain ofnrows, so the packed verification batch fitsmax_num_batched_tokens.Validation
Based on main
ecb1f8b505b7e831815b40aec3b4598619cca23a, including #48/#51. Per-file SHA256 in the evidence identifies the tested patch.atol=rtol=1e-4checks in eight cases. BF16 fixed-input errors and natural-generation differences are reported separately, without an allclose or lossless claim.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:
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.