perf(models): project only the sampled row at prefill for 8 families

docs/benchmark_results/last-row-prefill-families-m1ultra-2026-09-24.md

@@ -0,0 +1,55 @@
+# Last-row prefill for eight more families (issue #1970)
+
+Qwen3, Qwen3-MoE, Phi-3, Mixtral, Gemma 2, Gemma 3, Jamba, and Nemotron-H used
+the `LanguageModel::forward_last_logits` default: a single-sequence prefill
+computed `[1, L, vocab]` logits through the LM head and kept one row. Each
+model's forward is now split into a hidden-state path ending at the final norm
+and a head (Gemma 2 keeps its final-logit softcap on the head side), and the
+three `forward_last_logits*` entry points slice the hidden state at the sampled
+position before the head, as Cohere2 (#1955) and Llama3 (#1969) already do.
+Models with per-sequence state (the Gemma 3 wrapper, Jamba, Nemotron-H) use the
+same state closure and mask rules as their existing `forward` and
+`forward_with_sequence_id`. Models without an embeddings input route the
+embeddings variant to the plain path, as the trait default did.
+
+- **Date:** 2026-09-24
+- **Hardware:** Apple M1 Ultra 128 GB, macOS 27.0 (26A428)
+- **Build:** `cargo build --release --features metal,accelerate`
+- **Arms:** main `13b6d33a` binary against the change, ABBA order
+- **Condition:** `mlxcel-bench-decode` with `bench_decode.sh`'s arguments
+  (warmup 20, tg128, `--ignore-eos`), `--prompt-tokens` 512, prefill tok/s
+- **Machine load:** indexers suspended; load average 3.7 to 4.3
+
+## Results (prefill tok/s, two runs per arm)
+
+| Model | vocab | main | last row | Change |
+|---|---|---|---|---|
+| Gemma 2 2B | 256k | 2106.1, 2103.6 | 2659.0, 2664.7 | +26.5% |
+| Gemma 3 4B | 262k | 1025.6, 1025.6 | 1222.1, 1222.4 | +19.2% |
+| Qwen3 1.7B | 152k | 2175.2, 2165.4 | 2566.7, 2582.4 | +18.6% |
+| Qwen3-MoE 30B-A3B | 152k | 887.5, 886.1 | 947.4, 945.1 | +6.7% |
+| Phi-3 mini | 32k | 1520.8, 1476.9 | 1555.9, 1557.4 | about +3.8% |
+| Nemotron 3 Nano 30B | 131k | 354.4, 354.5 | 365.5, 355.1 | +0.2% to +3% |
+| Jamba reasoning 3B | 65k | 217.1, 217.0 | 218.8, 218.6 | +0.8% |
+| Mixtral 8x7B | 32k | 337.0, 336.9 | 339.3, 339.3 | +0.7% |
+
+The gain tracks the head's share of prefill work, which is the vocabulary size
+relative to the model size: small models with 150k to 262k vocabularies gain
+the most, large MoE and hybrid models with small vocabularies are near noise.
+The Phi-3 and Nemotron ranges are wide because one run in one arm moved.
+
+## Correctness
+
+- Greedy output byte-identical to main on all eight, 100 tokens each for a
+  short prompt and a 512-token prompt.
+- Unit tests (Qwen3, Gemma 2, alongside the existing Llama3 and Cohere2 tests):
+  a shared helper, `test_support::last_logits::assert_last_logits_match_full_forward`,
+  checks all three entry points against the sliced full forward at the last and
+  an interior position. Shifting the slice by one row fails both new tests.
+
+## Not measured
+
+CUDA and ROCm; pp2048 for these families; DeepSeek V3, which also still uses
+the trait default and has no local checkpoint. Jamba, Nemotron-H, Gemma 3,
+Mixtral, Qwen3-MoE, and Phi-3 are covered by the greedy comparison but have no
+dedicated unit test.

src/models/gemma2.rs

@@ -419,33 +419,63 @@ impl Gemma2Model {
         caches: &mut [KVCache],
         mask: Option<&MlxArray>,
     ) -> UniquePtr<MlxArray> {
-        // Embed tokens and multiply by sqrt(hidden_size)
-        let mut h = self.embed_tokens.forward(input_ids);
-        let hidden_size = mlxcel_core::array_shape(&h)[2]; // [..., hidden_size]
+        let h = self.hidden_states(input_ids, None, caches, mask);
+        self.logits_from_hidden(&h)
+    }
+
+    /// Embeddings (the caller's or looked up) scaled by sqrt(hidden_size),
+    /// every layer and the final norm. The scale applies to all embeddings,
+    /// matching Python GemmaModel.__call__.
+    fn hidden_states(
+        &self,
+        input_ids: &MlxArray,
+        input_embeddings: Option<&MlxArray>,
+        caches: &mut [KVCache],
+        mask: Option<&MlxArray>,
+    ) -> UniquePtr<MlxArray> {
+        let mut h = if let Some(embeds) = input_embeddings {
+            mlxcel_core::copy(embeds)
+        } else {
+            self.embed_tokens.forward(input_ids)
+        };
+        let hidden_size = mlxcel_core::array_shape(&h)[2];
         let scale = (hidden_size as f32).sqrt();
         h = mlxcel_core::multiply_scalar(&h, scale);
 
-        // Pass through transformer layers
         let n = self.layers.len();
         for (i, layer) in self.layers.iter().enumerate() {
             h = layer.forward(&h, &mut caches[i], mask);
             pipeline_hint(&h, i, n);
         }
 
-        // Final norm
-        let h = self.norm.forward(&h);
+        self.norm.forward(&h)
+    }
 
-        // LM head
-        let mut logits = if let Some(head) = &self.lm_head {
-            head.forward(&h)
+    /// LM head (or the tied embedding) and the final logit softcap, both
+    /// per position.
+    fn logits_from_hidden(&self, h: &MlxArray) -> UniquePtr<MlxArray> {
+        let logits = if let Some(head) = &self.lm_head {
+            head.forward(h)
         } else {
-            self.embed_tokens.as_linear(&h)
+            self.embed_tokens.as_linear(h)
         };
+        mlxcel_core::compiled_softcap(&logits, self.final_logit_softcapping)
+    }
 
-        // Apply final logit softcapping
-        logits = mlxcel_core::compiled_softcap(&logits, self.final_logit_softcapping);
-
-        logits
+    /// Logits `[B, 1, vocab]` for position `last_pos` only: the hidden state
+    /// is sliced before the LM head and softcap, which act per position, so a
+    /// prefill does not project every prompt row through a 256k vocabulary.
+    fn last_logits(
+        &self,
+        input_ids: &MlxArray,
+        input_embeddings: Option<&MlxArray>,
+        caches: &mut [KVCache],
+        mask: Option<&MlxArray>,
+        last_pos: usize,
+    ) -> UniquePtr<MlxArray> {
+        let h = self.hidden_states(input_ids, input_embeddings, caches, mask);
+        let row = mlxcel_core::generate::logits_at_position(&h, last_pos);
+        self.logits_from_hidden(&row)
     }
 
     /// Create KV caches for all layers
@@ -475,36 +505,8 @@ impl Gemma2Model {
         caches: &mut [KVCache],
         mask: Option<&MlxArray>,
     ) -> UniquePtr<MlxArray> {
-        let mut h = if let Some(embeds) = input_embeddings {
-            mlxcel_core::copy(embeds)
-        } else {
-            self.embed_tokens.forward(input_ids)
-        };
-
-        // Apply sqrt(hidden_size) normalization to ALL embeddings.
-        // Python: normalizer = sqrt(config.hidden_size); h = h * normalizer
-        // Used by: PaliGemma VLM (Gemma2 backbone)
-        let hidden_size = mlxcel_core::array_shape(&h)[2];
-        let scale = (hidden_size as f32).sqrt();
-        h = mlxcel_core::multiply_scalar(&h, scale);
-
-        let n = self.layers.len();
-        for (i, layer) in self.layers.iter().enumerate() {
-            h = layer.forward(&h, &mut caches[i], mask);
-            pipeline_hint(&h, i, n);
-        }
-
-        let h = self.norm.forward(&h);
-
-        let mut logits = if let Some(head) = &self.lm_head {
-            head.forward(&h)
-        } else {
-            self.embed_tokens.as_linear(&h)
-        };
-
-        logits = mlxcel_core::compiled_softcap(&logits, self.final_logit_softcapping);
-
-        logits
+        let h = self.hidden_states(input_ids, input_embeddings, caches, mask);
+        self.logits_from_hidden(&h)
     }
 
     /// Load model from directory
@@ -578,6 +580,39 @@ fn get_weight_copy(weights: &WeightMap, name: &str) -> Result<UniquePtr<MlxArray
 
 // LanguageModel trait implementation.
 impl LanguageModel for Gemma2Model {
+    fn forward_last_logits(
+        &self,
+        input_ids: &MlxArray,
+        caches: &mut [KVCache],
+        mask: Option<&MlxArray>,
+        last_pos: usize,
+    ) -> UniquePtr<MlxArray> {
+        self.last_logits(input_ids, None, caches, mask, last_pos)
+    }
+
+    fn forward_last_logits_with_sequence_id(
+        &self,
+        input_ids: &MlxArray,
+        _seq_id: Option<mlxcel_core::cache::SequenceId>,
+        caches: &mut [KVCache],
+        mask: Option<&MlxArray>,
+        last_pos: usize,
+    ) -> UniquePtr<MlxArray> {
+        self.last_logits(input_ids, None, caches, mask, last_pos)
+    }
+
+    fn forward_last_logits_with_embeddings_and_sequence_id(
+        &self,
+        input_ids: &MlxArray,
+        input_embeddings: Option<&MlxArray>,
+        _seq_id: Option<mlxcel_core::cache::SequenceId>,
+        caches: &mut [KVCache],
+        mask: Option<&MlxArray>,
+        last_pos: usize,
+    ) -> UniquePtr<MlxArray> {
+        self.last_logits(input_ids, input_embeddings, caches, mask, last_pos)
+    }
+
     fn forward(
         &self,
         input_ids: &MlxArray,
@@ -613,3 +648,7 @@ impl LanguageModel for Gemma2Model {
         vec![1, 107] // Gemma2 EOS tokens: <eos> (1) and <end_of_turn> (107)
     }
 }
+
+#[cfg(test)]
+#[path = "gemma2_tests.rs"]
+mod tests;

src/models/gemma2_tests.rs

@@ -0,0 +1,90 @@
+// Copyright 2025-2026 Lablup Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Unit tests for Gemma 2 on a tiny synthetic model.
+
+use super::{Gemma2Model, ModelArgs};
+use mlxcel_core::weights::WeightMap;
+use mlxcel_core::{MlxArray, UniquePtr};
+
+fn filled(shape: &[i32], phase: usize) -> UniquePtr<MlxArray> {
+    let n: i32 = shape.iter().product();
+    let data: Vec<f32> = (0..n)
+        .map(|i| ((i as usize + phase) as f32).sin() * 0.3)
+        .collect();
+    mlxcel_core::from_slice_f32(&data, shape)
+}
+
+/// Two layers, tied embeddings (the `as_linear` head) and the default final
+/// logit softcap.
+fn tiny_model() -> Gemma2Model {
+    let args: ModelArgs = serde_json::from_value(serde_json::json!({
+        "model_type": "gemma2",
+        "hidden_size": 8,
+        "num_hidden_layers": 2,
+        "intermediate_size": 16,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 1,
+        "rms_norm_eps": 1e-6,
+        "vocab_size": 12,
+        "head_dim": 4,
+        "tie_word_embeddings": true
+    }))
+    .expect("tiny gemma2 config parses");
+    let (h, q, kv, ff) = (8, 8, 4, 16);
+    let mut w = WeightMap::new();
+    let mut phase = 1usize;
+    let mut next = || {
+        phase += 5;
+        phase
+    };
+    w.insert("model.embed_tokens.weight".into(), filled(&[12, h], next()));
+    w.insert("model.norm.weight".into(), filled(&[h], next()));
+    for l in 0..2 {
+        let p = format!("model.layers.{l}");
+        for (name, shape) in [
+            ("self_attn.q_proj.weight", [q, h]),
+            ("self_attn.k_proj.weight", [kv, h]),
+            ("self_attn.v_proj.weight", [kv, h]),
+            ("self_attn.o_proj.weight", [h, q]),
+            ("mlp.gate_proj.weight", [ff, h]),
+            ("mlp.up_proj.weight", [ff, h]),
+            ("mlp.down_proj.weight", [h, ff]),
+        ] {
+            w.insert(format!("{p}.{name}"), filled(&shape, next()));
+        }
+        for norm in [
+            "input_layernorm",
+            "post_attention_layernorm",
+            "pre_feedforward_layernorm",
+            "post_feedforward_layernorm",
+        ] {
+            w.insert(format!("{p}.{norm}.weight"), filled(&[h], next()));
+        }
+    }
+    Gemma2Model::from_weights(&w, &args).expect("tiny gemma2 builds")
+}
+
+/// The last-logits overrides slice the hidden state before the LM head and
+/// the softcap, and must return the full forward's row (#1968 follow-up).
+#[test]
+fn last_logits_match_the_sliced_full_forward() {
+    let model = tiny_model();
+    crate::test_support::last_logits::assert_last_logits_match_full_forward(
+        &model,
+        &[1, 3, 5, 7, 2, 9],
+        true,
+        1e-4,
+    );
+}