Match Qwen head-chunked attention execution

src/engine/metal/qwen_mtplx.rs

@@ -86,6 +86,7 @@ mod stream;
 #[cfg(test)]
 mod submission;
 mod unary;
+mod verify_sdpa;
 mod views;
 mod vision;
 mod vision_kernels;

src/engine/metal/qwen_mtplx/decoder.rs

@@ -376,6 +376,7 @@ fn mtplx_decoder_connected_attention_and_moe_advance_caches() {
         },
         qsa: qsa_attention::Options {
             verify_gate: false,
+            head_chunk: false, // Historical unchunked oracle fixture.
             indexer: QsaIndexerOptions {
                 fused: false,
                 compiled: false,

src/engine/metal/qwen_mtplx/execution.rs

@@ -186,6 +186,7 @@ impl Options {
             },
             qsa: qsa_attention::Options {
                 verify_gate: false,
+                head_chunk: super::policy::sdpa_head_chunk_enabled(&get),
                 indexer: QsaIndexerOptions {
                     fused: flag("MTPLX_FUSED_QSA_INDEXER"),
                     compiled: flag("MTPLX_COMPILED_QSA_INDEXER"),

src/engine/metal/qwen_mtplx/issue103_analysis.rs

@@ -391,6 +391,7 @@ fn issue103_vision_position_receipts() {
     .unwrap();
     let mut options = Options {
         verify_gate: false,
+        head_chunk: false, // Historical unchunked oracle fixture.
         indexer: super::QsaIndexerOptions {
             fused: false,
             compiled: false,

src/engine/metal/qwen_mtplx/issue106.rs

@@ -166,6 +166,7 @@ fn issue106_complete_state_and_checkpoint_recovery() {
         context: 4096,
         vocabulary: metadata.tokenizer.vocab_size() as u32,
         mtp: true,
+        head_chunk: super::policy::sdpa_head_chunk_enabled(|key| std::env::var(key).ok()),
     };
     let mut live = LiveSession::default();
     assert_eq!(&std::fs::read(&old).unwrap()[..8], b"QWMTKV01");

src/engine/metal/qwen_mtplx/issue109_product.rs

@@ -8,6 +8,16 @@ use crate::settings::{
 #[test]
 #[ignore = "installed model; serial external memory supervisor required"]
 fn issue109_verify_rows_and_warm_checkpoint() {
+    verify_rows_and_warm_checkpoint(true);
+}
+
+#[test]
+#[ignore = "installed model and actual pre-head-chunk checkpoint; serial external memory supervisor"]
+fn issue113_default_verify_and_checkpoint() {
+    verify_rows_and_warm_checkpoint(false);
+}
+
+fn verify_rows_and_warm_checkpoint(old_gate: bool) {
     assert_eq!(std::env::var("MTPLX_CONTEXT_COPY").as_deref(), Ok("0"));
     let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
     let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE109_DIR").unwrap());
@@ -151,13 +161,25 @@ fn issue109_verify_rows_and_warm_checkpoint() {
     owner
         .save_checkpoint(&directory.join("new.bin"), [109; 32], &mut |_| {})
         .unwrap();
-    let old = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE109_OLD").unwrap());
-    assert!(
-        owner
-            .load_checkpoint(&old, &mut |_| {})
-            .unwrap_err()
-            .contains("obsolete MTP attention-gate")
+    let old = std::path::PathBuf::from(
+        std::env::var_os(if old_gate {
+            "DS4_ISSUE109_OLD"
+        } else {
+            "DS4_ISSUE113_OLD"
+        })
+        .unwrap(),
     );
+    let reject_old = old_gate || owner.identity.head_chunk;
+    let loaded = owner.load_checkpoint(&old, &mut |_| {});
+    if reject_old {
+        assert!(loaded.unwrap_err().contains(if old_gate {
+            "obsolete MTP attention-gate"
+        } else {
+            "SDPA head-chunk policy differs"
+        }));
+    } else {
+        assert!(loaded.unwrap());
+    }
     assert_eq!(original, state(&owner));
     let mut next = messages.clone();
     next.push(first.message.clone());
@@ -179,6 +201,6 @@ fn issue109_verify_rows_and_warm_checkpoint() {
     assert_eq!(expected, state(&owner));
     println!(
         "{}",
-        serde_json::json!({"case":"checkpoint","actual_old_rejected":true,"warm_reopened_state_exact":true})
+        serde_json::json!({"case":"checkpoint","actual_old_rejected":reject_old,"warm_reopened_state_exact":true})
     );
 }

src/engine/metal/qwen_mtplx/owner.rs

@@ -42,6 +42,7 @@ impl QwenExecutor {
             vocabulary: u32::try_from(metadata.tokenizer.vocab_size())
                 .map_err(|e| e.to_string())?,
             mtp: settings.speculative.glm_mtp,
+            head_chunk: super::policy::sdpa_head_chunk_enabled(|key| std::env::var(key).ok()),
         };
         let (execution, policy) = Execution::open(settings)?;
         Ok(Self {

src/engine/metal/qwen_mtplx/policy.rs

@@ -25,6 +25,10 @@ fn falsey(value: Option<&str>) -> bool {
     })
 }
 
+pub(super) fn sdpa_head_chunk_enabled(get: impl Fn(&'static str) -> Option<String>) -> bool {
+    !falsey(get("MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK").as_deref())
+}
+
 pub(super) fn sorted_gather_enabled(get: impl Fn(&'static str) -> Option<String>) -> bool {
     !falsey(get("MTPLX_MOE_SORTED_GATHER_KERNEL").as_deref())
         && !truthy(get("MTPLX_FORCE_GPU_FAMILY_FALLBACK").as_deref())

src/engine/metal/qwen_mtplx/qsa_attention.rs

@@ -33,6 +33,7 @@ pub(super) struct Attention<'a> {
 /// Resolved original request-phase/environment switches, not new heuristics.
 pub(super) struct Options {
     pub(super) verify_gate: bool,
+    pub(super) head_chunk: bool,
     pub(super) indexer: QsaIndexerOptions,
     pub(super) flash_min_context: u32,
     pub(super) prefill_gather: bool,
@@ -351,8 +352,15 @@ impl Attention<'_> {
             }
             _ => return Err("unconsumed QSA attention selection".into()),
         };
-        let out =
-            attention::scaled_dot_product(&q, &k, &v, mask.as_ref(), self.scale as f32, stream)?;
+        let out = super::verify_sdpa::call(
+            &q,
+            &k,
+            &v,
+            mask.as_ref(),
+            self.scale as f32,
+            options.head_chunk,
+            stream,
+        )?;
         self.finish(
             &ops::transpose(&out, &[0, 2, 1, 3], stream)?,
             &gate,
@@ -546,6 +554,7 @@ fn mtplx_qsa_attention_module_matches_reference() {
             };
             let options = Options {
                 verify_gate: false,
+                head_chunk: false, // Historical unchunked oracle fixture.
                 indexer: qsa_indexer_options(&first["options"]),
                 flash_min_context: first["flash_min"].as_u64().unwrap() as u32,
                 prefill_gather: first["prefill_gather"].as_bool().unwrap(),

src/engine/metal/qwen_mtplx/session_cache.rs

@@ -532,6 +532,9 @@ fn mtplx_full_installed_chat_matches_reference() {
                     context: settings.context_tokens as u32,
                     vocabulary: metadata.tokenizer.vocab_size() as u32,
                     mtp,
+                    head_chunk: super::policy::sdpa_head_chunk_enabled(|key| {
+                        std::env::var(key).ok()
+                    }),
                 };
                 let key = crate::engine::conversation_key(
                     &settings.system_prompt,

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"
     );

src/engine/metal/qwen_mtplx/verify_sdpa.rs

@@ -0,0 +1,267 @@
+//! Pinned qwen4_exp._verify_sdpa; selected by geometry, not attention phase.
+use super::array::Array;
+use super::stream::Stream;
+use super::{attention, ops};
+
+pub(super) fn plan(rows: i32, heads: i32, kv_heads: i32) -> Option<(i32, i32)> {
+    if kv_heads <= 0 || heads % kv_heads != 0 {
+        return None;
+    }
+    let gqa = heads / kv_heads;
+    if gqa <= 0 || i64::from(rows) * i64::from(gqa) <= 32 || rows > 32 {
+        return None;
+    }
+    let width = (32 / rows.max(1)).min(gqa).max(1);
+    Some((width, kv_heads * ((gqa - 1) / width + 1)))
+}
+
+pub(super) fn call(
+    q: &Array,
+    k: &Array,
+    v: &Array,
+    mask: Option<&Array>,
+    scale: f32,
+    enabled: bool,
+    stream: Stream,
+) -> Result<Array, String> {
+    let qs = q.layout().shape().to_vec();
+    let ks = k.layout().shape().to_vec();
+    if qs.len() != 4 || ks.len() != 4 {
+        return Err("Qwen SDPA expects rank-four queries and keys".into());
+    }
+    let Some((width, chunks)) = enabled.then(|| plan(qs[2], qs[1], ks[1])).flatten() else {
+        return attention::scaled_dot_product(q, k, v, mask, scale, stream);
+    };
+    static ENGAGED: std::sync::Once = std::sync::Once::new();
+    ENGAGED.call_once(|| {
+        eprintln!(
+            "[mtplx] verify SDPA head-chunked: q_len={} heads_per_chunk={width} chunks={chunks}",
+            qs[2]
+        );
+    });
+    let gqa = qs[1] / ks[1];
+    let qg = ops::reshape(q, &[qs[0], ks[1], gqa, qs[2], qs[3]], stream)?;
+    let mut outputs = Vec::with_capacity(chunks as usize);
+    for group in 0..ks[1] {
+        let kv = |a: &Array| {
+            let mut end = a.layout().shape().to_vec();
+            end[1] = group + 1;
+            ops::slice(a, &[0, group, 0, 0], &end, &[1; 4], stream)
+        };
+        let kj = kv(k)?;
+        let vj = kv(v)?;
+        for start in (0..gqa).step_by(width as usize) {
+            let end = (start + width).min(gqa);
+            let qc = ops::squeeze(
+                &ops::slice(
+                    &qg,
+                    &[0, group, start, 0, 0],
+                    &[qs[0], group + 1, end, qs[2], qs[3]],
+                    &[1; 5],
+                    stream,
+                )?,
+                Some(&[1]),
+                stream,
+            )?;
+            // The pinned wrapper passes the original mask unchanged. The
+            // primitive broadcasts it or rejects incompatible head dimensions.
+            outputs.push(attention::scaled_dot_product(
+                &qc, &kj, &vj, mask, scale, stream,
+            )?);
+        }
+    }
+    ops::concatenate(&outputs, 1, stream)
+}
+
+#[test]
+fn issue113_plan_and_policy_match_pinned_oracle() {
+    let fixture: serde_json::Value = serde_json::from_str(include_str!(
+        "../../../../tests/fixtures/mtplx-head-chunk.json"
+    ))
+    .unwrap();
+    let plans = fixture["plans"].as_array().unwrap();
+    assert_eq!(plans.len(), 520);
+    for p in plans {
+        let n = |k| p[k].as_i64().unwrap() as i32;
+        assert_eq!(
+            serde_json::to_value(plan(n("rows"), n("heads"), n("kv"))).unwrap(),
+            p["plan"]
+        );
+    }
+    for (value, expected) in [
+        (None, true),
+        (Some(""), true),
+        (Some("unexpected"), true),
+        (Some("0"), false),
+        (Some(" FALSE "), false),
+        (Some("Off"), false),
+        (Some("NO"), false),
+    ] {
+        assert_eq!(
+            super::policy::sdpa_head_chunk_enabled(|_| value.map(str::to_owned)),
+            expected
+        );
+    }
+}
+
+#[test]
+#[ignore = "Metal, pinned head-chunk numerical and dispatch fixtures; serial supervisor"]
+fn issue113_head_chunk_outputs_and_dispatch() {
+    use super::super::{
+        gpu::Context,
+        qwen_mtplx_tests::{capture_dispatches, hex_digest, pattern},
+    };
+    use super::array::Dtype;
+    super::configure_sources().unwrap();
+    let _context = Context::open_qwen(0).unwrap();
+    let streams = super::stream::Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(super::stream::Device::Gpu).unwrap();
+    let fixture: serde_json::Value = serde_json::from_str(include_str!(
+        "../../../../tests/fixtures/mtplx-head-chunk.json"
+    ))
+    .unwrap();
+    let cases = fixture["cases"].as_array().unwrap();
+    assert_eq!(cases.len(), 216);
+    for (index, c) in cases.iter().enumerate() {
+        let n = |k| c[k].as_i64().unwrap() as i32;
+        let (b, h, kv, r, t) = (n("batch"), n("heads"), n("kv"), n("rows"), n("total"));
+        let bf = |shape: &[i32], salt| {
+            Array::new(
+                shape,
+                Dtype::BF16,
+                pattern(shape.iter().map(|&x| x as u32).product(), salt),
+            )
+            .unwrap()
+        };
+        let q = ops::transpose(&bf(&[b, r, h, 256], 93), &[0, 2, 1, 3], gpu).unwrap();
+        let source = |salt| {
+            ops::slice(
+                &bf(&[b, kv, t + 7, 256], salt),
+                &[0, 0, 2, 0],
+                &[b, kv, t + 2, 256],
+                &[1; 4],
+                gpu,
+            )
+            .unwrap()
+        };
+        let (k, v) = (source(94), source(95));
+        let mode = c["mode"].as_str().unwrap();
+        let mask = if mode == "none" {
+            None
+        } else {
+            let mh = if mode == "heads" {
+                h
+            } else if mode == "local" {
+                4
+            } else {
+                1
+            };
+            let mut bytes = Vec::new();
+            for head in 0..mh {
+                for row in 0..r {
+                    for col in 0..t {
+                        let on = col <= t - r + row && (col + head) % 3 != 1 && row != 1;
+                        if mode == "add" {
+                            bytes.extend((if on { 0u16 } else { 0xc120 }).to_le_bytes());
+                        } else {
+                            bytes.push(u8::from(on));
+                        }
+                    }
+                }
+            }
+            let shape = if mode == "two_dim" {
+                vec![r, t]
+            } else {
+                vec![1, mh, r, t]
+            };
+            Some(
+                Array::new(
+                    &shape,
+                    if mode == "add" {
+                        Dtype::BF16
+                    } else {
+                        Dtype::Bool
+                    },
+                    super::scalar_buffer(&bytes).unwrap(),
+                )
+                .unwrap(),
+            )
+        };
+        let result = call(
+            &q,
+            &k,
+            &v,
+            mask.as_ref(),
+            0.0625,
+            c["enabled"].as_bool().unwrap(),
+            gpu,
+        );
+        if c.get("error").is_some() {
+            assert!(
+                result.is_err(),
+                "case {index}: expected oracle mask/shape rejection"
+            );
+            continue;
+        }
+        let output = result.unwrap_or_else(|e| panic!("case {index}: {e}"));
+        let chunked = c["enabled"].as_bool().unwrap() && plan(r, h, kv).is_some();
+        let outputs = if chunked {
+            output.inputs().clone()
+        } else {
+            vec![output.clone()]
+        };
+        let calls = c["calls"].as_array().unwrap();
+        assert_eq!(outputs.len(), calls.len(), "case {index}");
+        for (part, expected) in outputs.iter().zip(calls) {
+            assert_eq!(
+                serde_json::json!(part.layout().shape()),
+                expected["q"],
+                "case {index}"
+            );
+            if matches!(
+                part.operation().as_deref(),
+                Some(super::ops::Operation::SdpaVector { .. })
+            ) {
+                for (input, name) in part.inputs().iter().take(3).zip(["q", "k", "v"]) {
+                    assert_eq!(
+                        serde_json::json!(input.layout().shape()),
+                        expected[name],
+                        "case {index}"
+                    );
+                }
+            }
+        }
+        drop(outputs);
+        let (_, dispatches) = capture_dispatches(|| {
+            ops::evaluate(&streams, std::slice::from_ref(&output), gpu, false).unwrap()
+        });
+        let fused = dispatches
+            .iter()
+            .filter(|d| d.0.starts_with("sdpa_vector") && !d.0.starts_with("sdpa_vector_2pass_2"))
+            .count();
+        assert_eq!(
+            fused,
+            c["fused_primitives"].as_u64().unwrap() as usize,
+            "case {index}: {dispatches:?}"
+        );
+        assert_eq!(
+            dispatches
+                .iter()
+                .filter(|d| d.0.contains("softmax"))
+                .count(),
+            calls.len() - fused,
+            "case {index}"
+        );
+        let mut bytes = vec![0; output.layout().nbytes()];
+        output.buffer().read(output.offset(), &mut bytes).unwrap();
+        assert_eq!(
+            hex_digest(&bytes),
+            c["sha256"].as_str().unwrap(),
+            "case {index}: {c}"
+        );
+        println!(
+            "{}",
+            serde_json::json!({"case":"head_chunk","index":index,"exact":true,"fused":fused,"subcalls":calls.len()})
+        );
+    }
+}