Repository navigation
perf(gemma4): project only the sampled row at prefill on every path - #1973
Merged
Merged
Conversation
Gemma 4 prefill still built the full [1, L, 262144] logits plus a softcap copy on the paths the shipped checkpoints take. The text wrapper kept full logits up to 4096 prompt tokens (#672), had no sequence-id overrides so server prefill used the trait default, and the published checkpoints load as Gemma4VLModel or Gemma4UnifiedModel, which forwarded none of the forward_last_logits entry points. The text wrapper now overrides all three entry points at every length through a shared logits_at helper, and both VLM wrappers forward them, keeping their per-layer-input and vision block-id selection. Against main on M1 Ultra (bench_decode.sh arguments, ABBA): E4B prefill +15.5% at pp512 and +15.3% at pp2048, 12B +8.8% and +8.6%. Greedy output identical on both for a short, a 512-token and an image prompt. New unit test covers the three entry points at two positions for both attention families and fails when the row is shifted. Closes #1972
inureyes
added
status:review
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.
Gemma 4 prefill still computed full 262k-vocab logits (plus a softcap copy) and kept one row on the paths the shipped checkpoints take:
Gemma4Wrapper::forward_last_logitskept full logits up to 4096 prompt tokens (fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672) to stay byte-identical with the batched head kernel.Gemma4Wrapperhad no sequence-id overrides, so server prefill used the trait default at every length.Gemma4VLModelorGemma4UnifiedModel, neither of which forwarded anyforward_last_logits*entry point. The fix: gemma-4 long-prompt prefill OOMs on CUDA because attention materializes O(heads*L^2) scores instead of using flash SDPA #672 long-prompt saving never reached them.This PR drops the 4096 gate, overrides all three entry points on the text wrapper through one
logits_athelper (head + softcap on the sliced row), and forwards them from both VLM wrappers with the same per-layer-input and vision block-id selection as their forward paths. The gate's premise was byte-identity; 11 families moved to the single-row kernel in #1955, #1969 and #1971 with identical greedy output, and so does Gemma 4 here.Measurements
M1 Ultra,
mlxcel-bench-decodewithbench_decode.sharguments, main2dfc33e5against this branch, ABBA:Decode unchanged. Record:
docs/benchmark_results/gemma4-last-row-prefill-m1ultra-2026-09-25.md.Validation
gemma4_last_logits_match_the_sliced_full_forwardchecks the three entry points at the last and an interior position for sliding and full attention; shifting the row by one fails it.make verify-testpass locally.Not covered: CUDA/ROCm, the MTP target path (keeps its own full forward).
Closes #1972