Repository navigation
perf(models): project only the sampled row at prefill for 8 families - #1971
Merged
Merged
Conversation
Qwen3, Qwen3-MoE, Phi-3, Mixtral, Gemma 2, Gemma 3, Jamba and Nemotron-H used the trait default for forward_last_logits, which runs the full [1, L, vocab] LM head at prefill and keeps one row. Each model now splits its forward into a hidden-state path ending at the final norm and a head (Gemma 2 keeps its logit softcap on the head side), and the three forward_last_logits entry points slice the hidden state at the sampled position first. Gemma 3, Jamba and Nemotron-H keep the same per-sequence state and mask rules as their existing forward paths. pp512 prefill on M1 Ultra against main: Gemma 2 2B +26.5%, Gemma 3 4B +19.2%, Qwen3 1.7B +18.6%, Qwen3-30B-A3B +6.7%, Phi-3 mini about +3.8%, Nemotron, Jamba and Mixtral within 0.2% to 3%. Greedy output is identical to main on all eight. A shared test helper checks all three entry points against the sliced full forward; Qwen3 and Gemma 2 tests use it and fail when the slice is shifted by one row. Closes #1970
inureyes
added
status:review
This was referenced Sep 24, 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
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.
Extends the prefill last-row LM head projection from Cohere2 (#1955) and Llama3 (#1969) to Qwen3, Qwen3-MoE, Phi-3, Mixtral, Gemma 2, Gemma 3, Jamba and Nemotron-H. These models used the
forward_last_logitstrait default, which projects every prompt position through the head and keeps one row.Each model now splits its forward into a hidden-state path ending at the final norm and a head (Gemma 2's final-logit softcap stays on the head side, since it is per position). The three
forward_last_logits*entry points slice the hidden state atlast_posbefore the head. Gemma 3, Jamba and Nemotron-H reuse the same sequence-state closures and mask handling as their existingforward/forward_with_sequence_id. Models without an embeddings input keep routing the embeddings variant to the plain path, as the default did.Measurements
M1 Ultra,
mlxcel-bench-decodewithbench_decode.sharguments, pp512 prefill, main13b6d33aagainst this branch, ABBA:Full record:
docs/benchmark_results/last-row-prefill-families-m1ultra-2026-09-24.md.Validation
test_support::last_logits::assert_last_logits_match_full_forward, used by new Qwen3 and Gemma 2 tests; shifting the slice by one row fails both.make verify-testpass locally.Not covered: CUDA/ROCm, and DeepSeek V3 (still on the trait default, no local checkpoint).
Closes #1970