Repository navigation
perf(eagle3): reuse the argmax the forward already computes - #1295
Merged
Merged
Conversation
Contributor
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the
 |
| Layer / File(s) | Summary |
|---|---|
Add optional precomputed IDs to metrics src/speculators/models/eagle3/metrics.py, tests/unit/models/test_eagle3_metrics.py |
compute_metrics accepts optional prediction and target IDs, aligns supplied IDs for accuracy calculation, and computes missing IDs internally. Tests compare loss and metrics across TTT steps and mask settings, including the case where only target IDs are supplied. |
Pass precomputed IDs from the forward path src/speculators/models/eagle3/core.py |
Eagle3DraftModel.forward computes verifier target IDs before the TTT loop, computes draft prediction IDs before metric calculation, and passes both to compute_metrics. |
Priority:
Merge Risk: 
57cb8
-
The argmax reuse has no identified correctness issue and is ready to merge after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Title check The title clearly and concisely describes the main performance change: reusing the draft argmax computed during Eagle3 forward processing. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Description check The description directly matches the changes. It explains the argmax reuse, target ID computation outside the TTT loop, metric alignment, performance measurements, and unit tests.
✨ Finishing Touches: 💡 2
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
Contributor
Merge Protections🟢 Merge protection satisfied — ready to merge. Show 1 satisfied protection🟢 Require approval from approved reviewers listAll pull requests must have at least one approving review from a member of the approved reviewers list before merging.
|
This was referenced Oct 21, 2026
orestis-z
marked this pull request as ready for review
October 22, 2026 13:25
orestis-z
requested review from
fynnsu,
rahul-tuli and
shanjiaz
as code owners
October 22, 2026 13:25
orestis-z
force-pushed
the
perf/argmax-reuse
branch
2 times, most recently
from
October 22, 2026 14:32
57cb865 to
7614c0f
Compare
orestis-z
force-pushed
the
perf/argmax-reuse
branch
2 times, most recently
from
October 22, 2026 14:34
53e343f to
2dc252f
Compare
fynnsu
approved these changes
Oct 26, 2026
fynnsu
left a comment
fynnsu
left a comment
Collaborator
There was a problem hiding this comment.
Generally approve but could use some cleanup.
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
orestis-z
force-pushed
the
perf/argmax-reuse
branch
from
October 27, 2026 15:22
2dc252f to
a9818dd
Compare
Each TTT step took three argmax passes over a ``[total_seq_len, draft_vocab_size]`` tensor: one over the draft logits for the accuracy metrics, one over the targets for the same, and one over the draft logits again to pick the next step's input ids. At a ~32k draft vocab that is three full streams of a tensor that does not fit in cache, for two distinct results. Take the draft argmax once, before the metrics, and hand it to ``compute_metrics``. The targets do not change across TTT steps, so their argmax moves out of the loop entirely and is computed once per forward. ``compute_metrics`` takes the ids unaligned and applies the same slicing ``align_for_step`` applies to the tensors they came from, so the reported per-step accuracies are unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>
orestis-z
added a commit
that referenced
this pull request
Oct 27, 2026
<!-- markdownlint-disable --> ## Purpose Split out of the measurements in #1259 / #1262 / #1263, one change per PR. Every training step built its metric scalars with `torch.tensor(<python float>, device=cuda)`, which copies from pageable host memory and so cannot be issued asynchronously — the host blocks until the work already queued on the device drains. The same step ran `torch.masked_select` on the accuracy mask, whose output size depends on the mask's values, forcing another device-to-host sync before the count could be read. Building those scalars with `torch.full`/`torch.zeros`/`torch.ones` and counting masked positions with a reduction gives identical numbers without either sync. `full_total` and `cond_total` stay separate tensors so a caller reducing the metrics dict in place cannot count one of them twice. ### Measured Single H100, eagle3, synthetic batches, full 151936-token vocab, `total_seq_len` 2048, 10 warmup + 50 measured steps, 3 reps, median of the per-run means, `log_freq=1` (the default), `scripts/benchmark.py`. This PR against main: | Qwen3-0.6B | Qwen3-1.7B | Qwen3-8B | Qwen3-32B | |---|---|---|---| | **1.111x** | **1.110x** | **1.030x** | **1.013x** | <details> <summary>Where the gain comes from, change by change</summary> Same harness and settings, each change measured against its parent commit, so every cell is one PR's marginal contribution at that verifier size. Run-to-run spread within a cell is under 1% everywhere except main at 0.6B. | PR | | Qwen3-0.6B | Qwen3-1.7B | Qwen3-8B | Qwen3-32B | |---|---|---|---|---|---| | #1294 | metrics off host | **1.111x** | **1.110x** | **1.030x** | **1.013x** | | #1295 | reuse the argmax | **1.076x** | **1.026x** | **1.013x** | 1.007x | | #1296 | single-pass fused KL | **1.047x** | **1.017x** | **1.012x** | 1.003x | | #1297 | logs off the thread | **1.029x** | **1.042x** | **1.015x** | 1.006x | | #1298 | one metric readback | 0.980x | 1.004x | 0.997x | 0.999x | | #1299 | H2D prefetch † | 1.024x | 1.013x | 1.029x | 1.021x | | #1300 | skip dropped positions | 0.997x | 0.988x | 0.980x | 0.988x | | #1301 | precast frozen heads | **1.055x** | **1.096x** | **1.059x** | **1.036x** | | | **cumulative** | **1.358x** | **1.329x** | **1.141x** | **1.077x** | † #1299's row is not the prefetch. That commit also pins the synthetic batch in `scripts/benchmark.py`, and the real `DataLoader` already sets `pin_memory=True`, so it is a benchmark-fidelity fix rather than a training speedup. Applying only the pinning hunk on top of #1298 gives 1.017x / 1.017x / 1.030x / 1.022x — the whole row. Discounting it, the cumulative row is **1.335x / 1.307x / 1.108x / 1.054x**. Reading it: - **#1294 and #1301 carry most of the gain.** They are the two largest at every size, and #1301 is the only one that also moves peak memory (−1.5 GB at 0.6B to −7.4 GB at 32B). - **#1295, #1296 and #1297 fade with verifier size** but stay positive — 1–2% at 8B, 0.3–0.7% at 32B. - **#1298 is a tie at every size.** One sync per step instead of a few dozen is the right shape, but it does not show up here. - **#1299's prefetch is inside noise** in this regime once the harness fix is separated out, though it does pay once the metric log emit is deferred — see #1299. - **#1300 is negative at `seq_len` 2048 at every size** and costs 0.3–1.5 GB of peak; it needs both long sequences and a small hidden size to pay. Only #1300 and #1301 change peak allocated memory at all; the other six are flat to the megabyte. Most of these changes remove fixed per-step host and metric overhead, so the relative win shrinks as the verifier grows even though the absolute saving does not. At 8B and 32B the step is dominated by the optimizer (42% and 58%, all Muon Newton–Schulz), which nothing here touches — the same point #1111's side note made. </details> ## Tests ```bash pytest tests/unit/models/test_metrics.py ``` New cases cover that the metric tensors are built without a host-to-device copy and that the masked counts match `masked_select`. Passes. ## Checklist I have filled in: - [x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [x] The test plan/results, such as providing test command and pasting the results. - [ ] (Optional) The necessary documentation update. - [x] I (a human) have written or reviewed the code in this pr to the best of my ability. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com> Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
orestis-z
force-pushed
the
perf/argmax-reuse
branch
from
October 27, 2026 15:23
a9818dd to
b8dcc74
Compare
orestis-z
enabled stacked auto-merge (squash)
October 27, 2026 15:29
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Each TTT step took three argmax passes over a
[total_seq_len, draft_vocab_size]tensor for two distinct results: one over the draft logits for the accuracy metrics, one over the targets for the same, and one over the draft logits again to pick the next step's input ids. At a ~32k draft vocab that tensor does not fit in cache, so each pass is a full stream of it. This takes the draft argmax once and hands it tocompute_metrics, and moves the target argmax out of the TTT loop entirely since the targets do not change across steps.compute_metricstakes the ids unaligned and applies the same slicingalign_for_stepapplies to the tensors they came from, so the reported per-step accuracies are unchanged.Measured
Single H100, eagle3, synthetic batches, full 151936-token vocab,
total_seq_len2048, 10 warmup + 50 measured steps, 3 reps, median of the per-run means,log_freq=1,scripts/benchmark.py. This PR against #1294:The argmax it removes is over the draft vocabulary, which does not grow with the verifier, so the relative win shrinks as the rest of the step gets heavier — but it is positive at every size.
Tests
New cases cover that the accuracy counts match the argmax of the tensors
align_for_stepproduces, for severalttt_stepswith and without a loss mask, plus a deterministic case that only passes ifpred_idsandtarget_idsget the opposite slices. Passes.Checklist
I have filled in:
🤖 Generated with Claude Code