Reach DS4 Metal performance parity

src/engine/metal/glm.rs

@@ -8,6 +8,11 @@ const CACHE_F16: bool = true;
 const DECODE_FLUSH_LAYERS: usize = 4;
 const AUTO_CACHE_BYTES: u64 = 12 * 1024 * 1024 * 1024;
 const STREAMING_TOKEN_PREFILL_MAX: u32 = 64;
+
+fn live_prefix_rewind_target(live: &[i32], incoming: &[i32]) -> Option<usize> {
+    (incoming.len() > 1 && incoming.len() < live.len() && live.starts_with(incoming))
+        .then_some(incoming.len() - 1)
+}
 const STREAMING_FULL_ATTN_CONTEXT: u32 = 8192;
 const LONG_CONTEXT_THRESHOLD: u32 = 65_536;
 const LONG_CONTEXT_FULL_ATTN_CONTEXT: u32 = 4096;
@@ -2365,6 +2370,14 @@ impl GlmExecutor {
     }
 
     pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result<usize, String> {
+        if let Some(rewind) = live_prefix_rewind_target(&self.tokens, tokens) {
+            self.tokens.truncate(rewind);
+            if let Some(mtp) = &mut self.mtp {
+                mtp.pending = None;
+                mtp.min_pos = None;
+            }
+            return Ok(rewind);
+        }
         if !tokens.starts_with(&self.tokens) {
             self.reset()?;
         }
@@ -3217,7 +3230,7 @@ fn f32_project_rows(
 mod tests {
     use super::{
         GlmExecutor, argmax, dynamic_expert_budget, full_indexer_layer, indexed_prefill_rows,
-        streaming_token_prefill_eligible,
+        live_prefix_rewind_target, streaming_token_prefill_eligible,
     };
     use crate::engine::{GLM, Model, ReasoningMode};
     use crate::model::ModelChoice;
@@ -3250,6 +3263,15 @@ mod tests {
         assert_eq!(indexed_prefill_rows(0, 17, 2048), 17);
     }
 
+    #[test]
+    fn repeated_glm_prompt_rewinds_one_token_for_logits() {
+        let live = [1, 2, 3, 4, 5, 6];
+        assert_eq!(live_prefix_rewind_target(&live, &[1, 2, 3, 4]), Some(3));
+        assert_eq!(live_prefix_rewind_target(&live, &[1]), None);
+        assert_eq!(live_prefix_rewind_target(&live, &live), None);
+        assert_eq!(live_prefix_rewind_target(&live, &[1, 2, 9]), None);
+    }
+
     #[test]
     fn glm_byte_budget_reserves_prefill_before_dynamic_experts() {
         let mib = 1024 * 1024;