Repository navigation
perf(llama3): project only the sampled row through the LM head at prefill - #1969
Merged
Merged
Conversation
…fill Llama3Model, which also serves Mistral and Qwen2/Qwen2.5, used the forward_last_logits default, so every single-sequence prefill computed [1, L, vocab] logits through the LM head and kept one row (about 0.5 TFLOP at 512 tokens for a 128k-152k vocabulary). Its two duplicated forward bodies become one hidden_states path, and the three forward_last_logits entry points slice the final-normed hidden state at last_pos before the head; the final norm and head act per position, so the row is the one the full forward produces. generate::logits_at_position becomes public for that. VisionLanguageModel now forwards the three entry points to its text model, so text-model overrides are no longer skipped behind the vision wrapper. M1 Ultra, bench_decode.sh arguments, ABBA: Llama 3.1 8B 4-bit prefill 771 to 826 tok/s at pp512 and 765 to 819 at pp2048 (+7.1%); Qwen2.5 7B 4-bit 817 to 881 at both lengths (+7.8%); decode unchanged. Greedy output identical on three prompts per model and on LLaVA 1.6 Mistral with an image. A new test checks all three entry points against the sliced full forward; shifting the slice by one row fails it. Closes #1968
inureyes
added
status:review
This was referenced Sep 24, 2026
Closed
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.
Llama3-family prefill (Llama 3.x, Mistral, Qwen2/2.5) now projects only the sampled row through the LM head, as Cohere2 does since #1955.
Llama3Model's two duplicated forward bodies become onehidden_statespath (embeddings, layers, final norm), andforward_last_logits,forward_last_logits_with_sequence_idandforward_last_logits_with_embeddings_and_sequence_idslice the hidden state atlast_posbefore the head. The final norm and head act per position, so the row equals the full forward's.generate::logits_at_positionis now public so overrides can slice hidden states with it.VisionLanguageModelforwards the three entry points to its text model; before, the trait default ran behind the wrapper and skipped any text-model override (Gemma 3 VLM, LLaVA, FastVLM, Llama 4 VLM).Measurement
M1 Ultra,
mlxcel-bench-decodewithbench_decode.sh's arguments, the old path reproduced through a measurement-only switch, ABBA:Decode unchanged. Record:
docs/benchmark_results/llama3-last-row-prefill-m1ultra-2026-09-24.md.Validation
make verify-test: 11433 passed, 0 failed. Workspace clippy and fmt pass.Closes #1968