Take argmax only over the last position - #248
Merged
JakeStevens merged 1 commit intoAug 11, 2026
Merged
Conversation
Parallel prefill returns logits for the whole prompt, and the sampling step
argmaxed that entire tensor before indexing the last row:
next_token = torch.argmax(logits, dim=-1)[0, -1].item()
For Llama 3.2 1B at a 1984-token prompt that scans [1, 1984, 128256] to
obtain a single token -- seq_len times the necessary work. Slice first
instead. Same result, and the cost stops scaling with prompt length.
Measured on an AMD EPYC Genoa devserver, 16 cores, MKL/OMP=1:
prompt dtype before after
1984 float32 267.89 ms 0.14 ms
1984 bfloat16 1150.24 ms 0.59 ms
8128 float32 1094.65 ms 0.15 ms
8128 bfloat16 4718.06 ms 0.59 ms
It lands hardest on reduced precision: torch's CPU argmax is roughly 4x
slower on bfloat16 than float32, so a bf16 model paid ~1.15 s of the 1.42 s
it took to pick one token at a 1984-token prompt.
This also distorts reported prefill throughput, because on_prompt_eval_end()
is recorded after the sampling step. Splitting forward() from the argmax on
Llama 3.2 1B showed the ExecuTorch forward gap between bf16 and fp32 was
-19 ms (bf16 faster) at a 1984-token prompt while the argmax gap was +886 ms;
at 8128 tokens, forward +856 ms against argmax +3651 ms. Nearly all of an
apparent bf16 regression was this line.
Only the parallel-prefill path is affected. The sequential-prefill fallback
and the decode loop call argmax on single-position logits, where the two
forms are equivalent, and the other call sites in this file already slice.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
digantdesai
approved these changes
Aug 11, 2026
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
3 participants
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.
Parallel prefill returns logits for the whole prompt, and the sampling step argmaxed that entire tensor before indexing the last row:
This means there is a bunch of unnecessary results.
Instead, slice first then argmax on that.
Only the parallel prefill path is affected. The sequential prefill fallback and the decode loop call argmax on single position logits, where the two forms are equivalent, and the other call sites in this file already slice.