Match Qwen head-chunked attention execution

src/engine/metal/qwen_mtplx/session_checkpoint.rs

@@ -23,6 +23,8 @@ pub(super) struct Identity {
     pub(super) context: u32,
     pub(super) vocabulary: u32,
     pub(super) mtp: bool,
+    #[serde(default)]
+    pub(super) head_chunk: bool,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -182,7 +184,11 @@ impl LiveSession {
         file.read_exact(&mut raw).map_err(|e| e.to_string())?;
         let mut manifest: Manifest = serde_json::from_slice(&raw).map_err(|e| e.to_string())?;
         progress(HEADER_BYTES + raw.len() as u64);
-        if manifest.identity != identity {
+        if (Identity {
+            head_chunk: identity.head_chunk,
+            ..manifest.identity
+        }) != identity
+        {
             return Err(
                 "Qwen checkpoint does not match the current model/context/MTP policy".into(),
             );
@@ -196,6 +202,14 @@ impl LiveSession {
                 "Qwen checkpoint uses obsolete MTP attention-gate math; rebuild the cache".into(),
             );
         }
+        // The route also changes short dense prefills, not just MTP verify.
+        // A <=3-token AR prefix cannot have executed a >=3-row body, whereas
+        // MTP can retain a short frontier from a wider verification window.
+        if manifest.identity.head_chunk != identity.head_chunk
+            && (identity.mtp || manifest.tokens.len() >= 4)
+        {
+            return Err("Qwen checkpoint SDPA head-chunk policy differs; rebuild the cache".into());
+        }
         // Old files do not record prefill chunk sizes or producer hardware.
         // At ten routes/token this length can contain an affected chunk;
         // shorter sessions cannot. Reject before allocating/restoring state.
@@ -312,6 +326,7 @@ fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() {
         context: 64,
         vocabulary: 256,
         mtp: true,
+        head_chunk: super::policy::sdpa_head_chunk_enabled(|key| std::env::var(key).ok()),
     };
     let directory =
         std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_CHECKPOINT_DIR").unwrap());
@@ -522,6 +537,87 @@ fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() {
             .contains("sorted-gather row guard")
     );
     assert_eq!(fingerprint(&live), ar_original);
+    // Policy changes in either direction invalidate affected checkpoints;
+    // missing metadata is the actual pre-port (unchunked) policy.
+    for mtp in [false, true] {
+        for count in [1, 3, 4] {
+            for produced in [false, true] {
+                let saved = live.current.as_mut().unwrap();
+                saved.token_ids = vec![1; count];
+                saved.mtp = mtp.then(|| vec![None]);
+                saved.boundaries.clear();
+                let producer = Identity {
+                    mtp,
+                    head_chunk: produced,
+                    ..identity
+                };
+                live.save_checkpoint(
+                    &path,
+                    producer,
+                    [113; 32],
+                    &streams,
+                    stream,
+                    |_| {},
+                    || false,
+                )
+                .unwrap();
+                let original = fingerprint(&live);
+                let bytes = fs::read(&path).unwrap();
+                let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+                for legacy in [false, true] {
+                    if legacy && produced {
+                        continue;
+                    }
+                    let mut manifest: serde_json::Value =
+                        serde_json::from_slice(&bytes[offset..]).unwrap();
+                    if legacy {
+                        manifest["identity"]
+                            .as_object_mut()
+                            .unwrap()
+                            .remove("head_chunk");
+                    }
+                    let raw = serde_json::to_vec(&manifest).unwrap();
+                    let mut file = bytes[..offset].to_vec();
+                    file.extend(&raw);
+                    fs::write(&broken, file).unwrap();
+                    for consumed in [false, true] {
+                        let consumer = Identity {
+                            head_chunk: consumed,
+                            ..producer
+                        };
+                        let mut read = 0;
+                        let result = live.load_checkpoint(
+                            &broken,
+                            consumer,
+                            &streams,
+                            stream,
+                            |n| read += n,
+                            || false,
+                        );
+                        if produced != consumed && (mtp || count >= 4) {
+                            assert!(
+                                result
+                                    .unwrap_err()
+                                    .contains("SDPA head-chunk policy differs")
+                            );
+                            assert_eq!(
+                                read,
+                                HEADER_BYTES + raw.len() as u64,
+                                "rejected before tensor reads"
+                            );
+                        } else {
+                            assert!(result.unwrap());
+                        }
+                        assert_eq!(fingerprint(&live), original);
+                        assert_eq!(live.current.as_ref().unwrap().token_ids, vec![1; count]);
+                    }
+                }
+            }
+        }
+    }
+    println!(
+        "#113 SDPA policy: legacy/current, both directions, AR/MTP, 1/3/4-token frontiers; exact compatible restore and incompatible rejection before tensor reads"
+    );
     println!(
         "Qwen sorted-gather checkpoint policy: old short state exact, possibly affected long state rejected before restore"
     );