Match Qwen compiled verification gate math

src/engine/metal/qwen_mtplx.rs

@@ -37,6 +37,8 @@ mod issue105;
 mod issue106;
 #[cfg(test)]
 mod issue108;
+#[cfg(test)]
+mod issue109;
 mod linear;
 mod load;
 mod mlp;

src/engine/metal/qwen_mtplx/decoder.rs

@@ -146,6 +146,7 @@ impl DecoderLayer<'_> {
             &hidden,
             options.hyper_fused,
             options.hyper_v3,
+            options.qsa.verify_gate,
             streams,
             stream,
         )?;
@@ -176,6 +177,7 @@ impl DecoderLayer<'_> {
             &hidden,
             options.hyper_fused,
             options.hyper_v3,
+            options.qsa.verify_gate,
             streams,
             stream,
         )?;
@@ -372,6 +374,7 @@ fn mtplx_decoder_connected_attention_and_moe_advance_caches() {
             ..Default::default()
         },
         qsa: qsa_attention::Options {
+            verify_gate: false,
             indexer: QsaIndexerOptions {
                 fused: false,
                 compiled: false,

src/engine/metal/qwen_mtplx/execution.rs

@@ -185,6 +185,7 @@ impl Options {
                 capture,
             },
             qsa: qsa_attention::Options {
+                verify_gate: false,
                 indexer: QsaIndexerOptions {
                     fused: flag("MTPLX_FUSED_QSA_INDEXER"),
                     compiled: flag("MTPLX_COMPILED_QSA_INDEXER"),
@@ -359,6 +360,19 @@ impl Execution {
     ) -> Result<Output, String> {
         let ids = self.input(ids)?;
         let options = self.graph_options(phase_prefill, capture)?;
+        self.forward_with_options(ids, &options)
+    }
+
+    /// Every target forward in the MTP loop (including repairs and final
+    /// commits) uses MTPLX's decode_verify scope. Draft/history calls do not.
+    pub(super) fn forward_verify(&mut self, ids: &[i32], capture: bool) -> Result<Output, String> {
+        let ids = self.input(ids)?;
+        let mut options = self.graph_options(false, capture)?;
+        options.qsa.verify_gate = true;
+        self.forward_with_options(ids, &options)
+    }
+
+    fn forward_with_options(&mut self, ids: Array, options: &Options) -> Result<Output, String> {
         let output = self.model.with_dependent_mut(|_, model| {
             model.call(
                 &ids,
@@ -369,13 +383,15 @@ impl Execution {
                     logits_keep: 0,
                     return_hidden: true,
                 },
-                &options,
+                options,
                 None,
                 &self.streams,
                 self.stream,
             )
         })?;
         #[cfg(test)]
+        super::issue109::record(self, &ids, &output, options.qsa.verify_gate)?;
+        #[cfg(test)]
         let output = self.inject_output(super::issue105::Stage::Target, output)?;
         #[cfg(test)]
         let output = self.script_output(false, output)?;

src/engine/metal/qwen_mtplx/hyper.rs

@@ -23,6 +23,7 @@ impl GatedResidual<'_> {
         input: &Array,
         fused: bool,
         v3: bool,
+        verify: bool,
         streams: &Streams,
         stream: Stream,
     ) -> Result<Read, String> {
@@ -110,18 +111,8 @@ impl GatedResidual<'_> {
         let sum = super::reduce::sum(&product, &[-2], false, stream)?;
         let mixed = binary(&sum, &scalar(0x3e80)?, Binary::Multiply, stream)?;
         let combine = if let Some(inject) = &self.inject {
-            let gate = binary(
-                &inject.apply(&normed, stream)?,
-                &scalar(0x4080)?,
-                Binary::Divide,
-                stream,
-            )?;
-            let gate = unary(&gate, Unary::Sigmoid, stream)?;
-            Some((
-                input.clone(),
-                // Python __rmul__ keeps the array first for 2.0 * gate.
-                binary(&gate, &scalar(0x4000)?, Binary::Multiply, stream)?,
-            ))
+            let gate = inject_gate(&inject.apply(&normed, stream)?, verify, stream)?;
+            Some((input.clone(), gate))
         } else {
             None
         };
@@ -129,6 +120,34 @@ impl GatedResidual<'_> {
     }
 }
 
+/// The eager HC inject site follows MTPLX _verify_inject on decode_verify.
+/// Other phases retain the original standalone divide/sigmoid/multiply.
+pub(super) fn inject_gate(logits: &Array, verify: bool, stream: Stream) -> Result<Array, String> {
+    use super::{
+        binary::{Binary, binary},
+        unary::{Unary, unary},
+    };
+    let body = |x: &Array| {
+        let scalar =
+            |bits: u16| Array::new(&[], Dtype::BF16, super::scalar_buffer(&bits.to_le_bytes())?);
+        let divided = binary(x, &scalar(0x4080)?, Binary::Divide, stream)?;
+        binary(
+            &unary(&divided, Unary::Sigmoid, stream)?,
+            &scalar(0x4000)?,
+            Binary::Multiply,
+            stream,
+        )
+    };
+    if !verify {
+        return body(logits);
+    }
+    let inputs = std::slice::from_ref(logits);
+    super::compiled::Replay::trace_compiled(inputs, |a| Ok(vec![body(&a[0])?]))?
+        .run(inputs)?
+        .pop()
+        .ok_or_else(|| "compiled HC inject omitted output".into())
+}
+
 pub(super) fn fused_read(
     x: &Array,
     norm: &Array,
@@ -379,7 +398,7 @@ fn mtplx_graph_hyper_forward_matches_reference() {
                     let mode = receipt["mode"].as_u64().unwrap();
                     let input = bf16(&[1, rows, 10240], 1);
                     let result = layer
-                        .forward(&input, mode & 1 != 0, mode & 2 != 0, &streams, gpu)
+                        .forward(&input, mode & 1 != 0, mode & 2 != 0, false, &streams, gpu)
                         .unwrap();
                     assert_eq!(
                         layer.pack.is_some(),

src/engine/metal/qwen_mtplx/issue103_analysis.rs

@@ -175,7 +175,16 @@ fn issue103_numeric_receipts() {
     ops::evaluate(&streams, &[gate.clone(), value.clone()], gpu, false).unwrap();
     let body = |a: &[Array]| {
         binary(
-            &unary::unary(&a[0], unary::Unary::Sigmoid, gpu)?,
+            &unary::unary(
+                &ops::astype(
+                    &ops::astype(&a[0], Dtype::F32, false, gpu)?,
+                    Dtype::BF16,
+                    false,
+                    gpu,
+                )?,
+                unary::Unary::Sigmoid,
+                gpu,
+            )?,
             &a[1],
             Binary::Multiply,
             gpu,
@@ -185,7 +194,7 @@ fn issue103_numeric_receipts() {
     let eager = body(&[gate.clone(), value.clone()]).unwrap();
     let compiled = Replay::trace_compiled(&[gate.clone(), value.clone()], body)
         .unwrap()
-        .run(&[gate, value])
+        .run(&[gate.clone(), value.clone()])
         .unwrap();
     let read = |a: &Array| {
         ops::read_f32_as_f64(
@@ -201,7 +210,14 @@ fn issue103_numeric_receipts() {
         "{}",
         serde_json::json!({"case":"bf16_gate","gate":-6.84375,"eager":e,"compiled":c})
     );
-    assert_ne!(e, c, "diagnosis no longer reproduces");
+    assert_ne!(
+        e, c,
+        "the two reference phase laws must remain distinguishable"
+    );
+    let actual = super::qsa_attention::attention_gate(&value, &gate, true, gpu).unwrap();
+    assert_eq!(read(&actual), c);
+    assert_eq!(e, vec![0.00106048583984375; 4]);
+    assert_eq!(c, vec![0.001068115234375; 4]);
     for (name, operand, mode) in [
         ("qsa_divisor", 128_f32.sqrt(), Binary::Divide),
         (
@@ -367,7 +383,8 @@ fn issue103_vision_position_receipts() {
         .unwrap(),
     )
     .unwrap();
-    let options = Options {
+    let mut options = Options {
+        verify_gate: false,
         indexer: super::QsaIndexerOptions {
             fused: false,
             compiled: false,
@@ -647,6 +664,7 @@ fn issue103_vision_position_receipts() {
             .map(|i| f32::from(i / 8 < 8 && i / 8 == i % 8))
             .collect::<Vec<_>>(),
     );
+    options.verify_gate = true;
     let mut gate_outputs = Vec::new();
     for gate in [0., -6.84375] {
         let bias = bf(&[16], &[vec![0.; 8], vec![gate; 8]].concat());
@@ -724,5 +742,5 @@ fn issue103_vision_position_receipts() {
         "{}",
         serde_json::json!({"case":"active_attention_gate","gate":-6.84375,"actual":actual,"expected_verify":expected})
     );
-    assert_ne!(actual, expected, "diagnosis no longer reproduces");
+    assert_eq!(actual, expected, "verify must use the compiled gate law");
 }

src/engine/metal/qwen_mtplx/issue106.rs

@@ -142,7 +142,7 @@ fn advance(execution: &mut Execution, saved: &SessionSnapshot, token: i32) -> Se
         .append_mtp_history(hidden.as_ref().unwrap(), &[token], false)
         .unwrap();
     ops::evaluate(&execution.streams, &[updated], execution.stream, false).unwrap();
-    let output = execution.forward(&[token], false, false).unwrap();
+    let output = execution.forward_verify(&[token], false).unwrap();
     let logits = ar::last_row(&output.logits.unwrap(), execution.stream).unwrap();
     let mut key = saved.token_ids.clone();
     key.push(token);

src/engine/metal/qwen_mtplx/issue107_product.rs

@@ -6,7 +6,7 @@ use crate::settings::{
 };
 use base64::Engine as _;
 
-fn state(owner: &QwenExecutor) -> String {
+pub(super) fn state(owner: &QwenExecutor) -> String {
     let saved = owner.live.current.as_ref().unwrap();
     let mut hashes = Vec::new();
     let payload = saved

src/engine/metal/qwen_mtplx/issue109.rs

@@ -0,0 +1,119 @@
+use super::array::Dtype;
+use super::{ops, qsa_attention};
+
+#[test]
+#[ignore = "Metal and exhaustive pinned gate fixtures; external supervisor required"]
+fn issue109_gate_phase_and_dtype_laws() {
+    super::configure_sources().unwrap();
+    let _context = crate::engine::metal::gpu::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 root = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE109_GATE").unwrap());
+    for (name, dtype) in [
+        ("bf16", Dtype::BF16),
+        ("f16", Dtype::F16),
+        ("f32", Dtype::F32),
+    ] {
+        let load = |part: &str| {
+            let bytes = std::fs::read(root.join(format!("{name}-{part}.f32"))).unwrap();
+            let array = super::array::Array::new(
+                &[(bytes.len() / 4) as i32],
+                Dtype::F32,
+                super::scalar_buffer(&bytes).unwrap(),
+            )
+            .unwrap();
+            ops::astype(&array, dtype, false, gpu).unwrap()
+        };
+        let gate = load("gate");
+        let output = load("output");
+        for phase in ["prefill", "ar_decode", "decode_verify"] {
+            let actual =
+                qsa_attention::attention_gate(&output, &gate, phase == "decode_verify", gpu)
+                    .unwrap();
+            let actual = ops::read_f32_as_f64(
+                &ops::astype(&actual, Dtype::F32, false, gpu).unwrap(),
+                &streams,
+                gpu,
+            )
+            .unwrap()
+            .into_iter()
+            .flat_map(|v| (v as f32).to_le_bytes())
+            .collect::<Vec<_>>();
+            let expected = std::fs::read(root.join(format!("{name}-{phase}.f32"))).unwrap();
+            assert_eq!(actual.len(), expected.len());
+            let different = actual
+                .chunks_exact(4)
+                .zip(expected.chunks_exact(4))
+                .enumerate()
+                .filter_map(|(i, (a, b))| (a != b).then_some(i))
+                .collect::<Vec<_>>();
+            assert!(
+                different.is_empty(),
+                "{name}/{phase}: {} different rows; first {:?}",
+                different.len(),
+                different.first()
+            );
+            println!(
+                "{}",
+                serde_json::json!({"case":"gate","dtype":name,"phase":phase,"values":actual.len()/4,"exact":true})
+            );
+            if dtype == Dtype::BF16 {
+                let inject =
+                    super::hyper::inject_gate(&gate, phase == "decode_verify", gpu).unwrap();
+                let actual = ops::read_f32_as_f64(
+                    &ops::astype(&inject, Dtype::F32, false, gpu).unwrap(),
+                    &streams,
+                    gpu,
+                )
+                .unwrap()
+                .into_iter()
+                .flat_map(|v| (v as f32).to_le_bytes())
+                .collect::<Vec<_>>();
+                assert!(
+                    actual == std::fs::read(root.join(format!("inject-{phase}.f32"))).unwrap(),
+                    "inject/{phase}"
+                );
+                println!(
+                    "{}",
+                    serde_json::json!({"case":"inject","phase":phase,"values":actual.len()/4,"exact":true})
+                );
+            }
+        }
+    }
+}
+
+thread_local! {
+    static RECEIPTS: std::cell::RefCell<Option<(std::path::PathBuf, usize)>> = const { std::cell::RefCell::new(None) };
+}
+
+pub(super) fn capture(directory: Option<std::path::PathBuf>) {
+    RECEIPTS.with_borrow_mut(|slot| {
+        *slot = directory.map(|d| {
+            std::fs::create_dir_all(&d).unwrap();
+            (d, 0)
+        })
+    });
+}
+
+pub(super) fn record(
+    execution: &super::execution::Execution,
+    ids: &super::array::Array,
+    output: &super::model::Output,
+    verify: bool,
+) -> Result<(), String> {
+    RECEIPTS.with_borrow_mut(|slot| {
+        let Some((directory, index)) = slot else { return Ok(()); };
+        let ids = ops::read_i64(ids, &execution.streams, execution.stream)?;
+        let mut record=serde_json::json!({"index":*index,"ids":ids,"verify":verify});
+        for (part, a) in [("logits",&output.logits),("hidden",&output.hidden)] {
+            let a=a.as_ref().ok_or("missing #109 output")?;
+            let bytes=execution.values_f32(a)?.into_iter().flat_map(f32::to_le_bytes).collect::<Vec<_>>();
+            std::fs::write(directory.join(format!("{index}-{part}.f32")),&bytes).map_err(|e|e.to_string())?;
+            record[part]=serde_json::json!({"shape":a.layout().shape(),"sha256":crate::engine::metal::qwen_mtplx_tests::hex_digest(&bytes)});
+        }
+        std::fs::write(directory.join(format!("{index}.json")),serde_json::to_vec(&record).unwrap()).map_err(|e|e.to_string())?;
+        println!("{}",serde_json::json!({"case":"verify_forward","record":record}));
+        *index+=1;
+        Ok(())
+    })
+}

src/engine/metal/qwen_mtplx/issue109_product.rs

@@ -0,0 +1,184 @@
+use super::issue107_product::state;
+use super::*;
+use crate::settings::{
+    GenerationPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences,
+    effective_settings,
+};
+
+#[test]
+#[ignore = "installed model; serial external memory supervisor required"]
+fn issue109_verify_rows_and_warm_checkpoint() {
+    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());
+    std::fs::create_dir_all(&directory).unwrap();
+    let settings = effective_settings(
+        crate::model::ModelChoice::Qwen38FlashNext,
+        &GenerationPreferences {
+            context_tokens: 8192,
+            max_generated_tokens: 12,
+            system_prompt: String::new(),
+            temperature: Some(0.),
+            top_p: Some(0.95),
+            min_p: Some(0.),
+            seed: Some(1),
+            reasoning_mode: ReasoningMode::Low,
+        },
+        &RuntimePreferences {
+            speculative: SpeculativePreferences {
+                glm_mtp: true,
+                ..Default::default()
+            },
+            ..Default::default()
+        },
+        root.parent().unwrap(),
+    )
+    .unwrap();
+    let mut owner = QwenExecutor::open(&settings.engine).unwrap();
+    owner.policy.store_min_suffix = Some(0);
+    let metrics = Metrics::new(&directory.join("metrics"));
+    let messages = vec![ChatTurn {
+        user: true,
+        tool: false,
+        system: false,
+        skip_previous_eos: false,
+        reasoning: None,
+        reasoning_complete: true,
+        content: "Count from one to ten.".into(),
+    }];
+    let run = |owner: &mut QwenExecutor, messages: &[ChatTurn], label: &str| {
+        super::super::issue109::capture(Some(directory.join(label)));
+        let (out, complete) = owner
+            .generate(
+                messages,
+                &settings.turn,
+                &AtomicBool::new(false),
+                &metrics,
+                |_, _, _| {},
+                |_, _| {},
+            )
+            .unwrap();
+        super::super::issue109::capture(None);
+        assert!(complete);
+        if let Some(prefix) = &owner.live.prefix {
+            if let Some(super::super::session_cache::State::Qsa(a)) = &prefix.trunk[35] {
+                let bytes = owner
+                    .execution
+                    .values_f32(a[0].as_ref().unwrap())
+                    .unwrap()
+                    .into_iter()
+                    .flat_map(f32::to_le_bytes)
+                    .collect::<Vec<_>>();
+                std::fs::write(directory.join(format!("{label}-prefill-key35.f32")), bytes)
+                    .unwrap();
+            }
+            for (part, a) in [
+                ("logits", Some(&prefix.logits)),
+                ("hidden", prefix.hidden.as_ref()),
+            ] {
+                let bytes = owner
+                    .execution
+                    .values_f32(a.unwrap())
+                    .unwrap()
+                    .into_iter()
+                    .flat_map(f32::to_le_bytes)
+                    .collect::<Vec<_>>();
+                std::fs::write(directory.join(format!("{label}-prefill-{part}.f32")), bytes)
+                    .unwrap();
+            }
+        }
+        let saved = owner.live.current.as_ref().unwrap();
+        let states = |states: &[Option<super::super::session_cache::State>]| {
+            states.iter().map(|s| {
+                s.as_ref().map(|s| {
+                    let arrays: &[Option<super::super::array::Array>] = match s {
+                        super::super::session_cache::State::Recurrent(a) => a,
+                        super::super::session_cache::State::Qsa(a) => a,
+                    };
+                    arrays.iter().map(|a| {
+                        a.as_ref().map(|a| {
+                            let values = owner.execution.values_f32(a).unwrap();
+                            let bytes = values.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<_>>();
+                            let canonical_zero = values.iter().flat_map(|&v| (if v == 0. { 0. } else { v }).to_le_bytes()).collect::<Vec<_>>();
+                            let hash = crate::engine::metal::qwen_mtplx_tests::hex_digest;
+                            serde_json::json!({"shape":a.layout().shape(),"sha256":hash(&bytes),"value_sha256":hash(&canonical_zero)})
+                        })
+                    }).collect::<Vec<_>>()
+                })
+            }).collect::<Vec<_>>()
+        };
+        let arrays = serde_json::json!({"trunk":states(&saved.trunk),"mtp":states(saved.mtp.as_ref().unwrap())});
+        std::fs::write(
+            directory.join(format!("{label}-arrays.json")),
+            serde_json::to_vec(&arrays).unwrap(),
+        )
+        .unwrap();
+        if let Some(super::super::session_cache::State::Qsa(a)) = &saved.trunk[35] {
+            let bytes = owner
+                .execution
+                .values_f32(a[0].as_ref().unwrap())
+                .unwrap()
+                .into_iter()
+                .flat_map(f32::to_le_bytes)
+                .collect::<Vec<_>>();
+            std::fs::write(directory.join(format!("{label}-key35.f32")), bytes).unwrap();
+        }
+        let record = serde_json::json!({"case":label,"prompt":&saved.token_ids[..out.prompt_tokens as usize],"tokens":&saved.token_ids[out.prompt_tokens as usize..],"cached":out.cached_tokens,"state_sha256":state(owner)});
+        std::fs::write(
+            directory.join(format!("{label}.json")),
+            serde_json::to_vec(&record).unwrap(),
+        )
+        .unwrap();
+        for (part, a) in [
+            ("logits", Some(&saved.logits)),
+            ("hidden", saved.hidden.as_ref()),
+        ] {
+            let bytes = owner
+                .execution
+                .values_f32(a.unwrap())
+                .unwrap()
+                .into_iter()
+                .flat_map(f32::to_le_bytes)
+                .collect::<Vec<_>>();
+            std::fs::write(directory.join(format!("{label}-{part}.f32")), bytes).unwrap();
+        }
+        println!("{record}");
+        out
+    };
+    let first = run(&mut owner, &messages, "cold");
+    assert_eq!(first.cached_tokens, 0);
+    let original = state(&owner);
+    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")
+    );
+    assert_eq!(original, state(&owner));
+    let mut next = messages.clone();
+    next.push(first.message.clone());
+    next.push(ChatTurn {
+        content: "Continue counting.".into(),
+        ..messages[0].clone()
+    });
+    let warm = run(&mut owner, &next, "warm");
+    assert!(warm.cached_tokens >= first.prompt_tokens);
+    let expected = state(&owner);
+    owner.reset().unwrap();
+    assert!(
+        owner
+            .load_checkpoint(&directory.join("new.bin"), &mut |_| {})
+            .unwrap()
+    );
+    assert_eq!(original, state(&owner));
+    run(&mut owner, &next, "reopened");
+    assert_eq!(expected, state(&owner));
+    println!(
+        "{}",
+        serde_json::json!({"case":"checkpoint","actual_old_rejected":true,"warm_reopened_state_exact":true})
+    );
+}

src/engine/metal/qwen_mtplx/model.rs

@@ -75,6 +75,7 @@ fn runs(layers: impl IntoIterator<Item = (bool, bool)>) -> Vec<Run> {
 struct RunKey {
     range: Range<usize>,
     capture: bool,
+    verify_gate: bool,
     stream: Stream,
     signature: Vec<(Vec<i32>, Dtype)>,
 }
@@ -223,7 +224,14 @@ impl<'a> TextModel<'a> {
         let h = mtp.fuse_and_run(hidden, &emb, cache, options, vision, streams, stream)?;
         let mixed = mtp
             .mixer
-            .forward(&h, options.hyper_fused, options.hyper_v3, streams, stream)?
+            .forward(
+                &h,
+                options.hyper_fused,
+                options.hyper_v3,
+                options.qsa.verify_gate,
+                streams,
+                stream,
+            )?
             .mixed;
         let logits = self.head.unwrap_or(self.embedding).apply(&mixed, stream)?;
         Ok(Output {
@@ -516,7 +524,14 @@ impl<'a> TextModel<'a> {
         self.last_widened = Some(h.clone());
         let mixed = self
             .mixer
-            .forward(&h, options.hyper_fused, options.hyper_v3, streams, stream)?
+            .forward(
+                &h,
+                options.hyper_fused,
+                options.hyper_v3,
+                options.qsa.verify_gate,
+                streams,
+                stream,
+            )?
             .mixed;
         let logits = if output.emit_logits {
             let mixed = if output.logits_keep != 0 {
@@ -610,6 +625,7 @@ impl<'a> TextModel<'a> {
             let key = RunKey {
                 range: range.clone(),
                 capture: options.gdn.capture,
+                verify_gate: options.qsa.verify_gate,
                 stream,
                 signature: flat
                     .iter()

src/engine/metal/qwen_mtplx/owner.rs

@@ -1072,3 +1072,7 @@ mod issue106_product;
 #[cfg(test)]
 #[path = "issue107_product.rs"]
 mod issue107_product;
+
+#[cfg(test)]
+#[path = "issue109_product.rs"]
+mod issue109_product;

src/engine/metal/qwen_mtplx/qsa_attention.rs

@@ -1,7 +1,5 @@
 //! qwen4_exp.Attention: projection, indexer, RoPE, KV and consumer branch order.
-use super::array::Array;
-#[cfg(test)]
-use super::array::Dtype;
+use super::array::{Array, Dtype};
 use super::attention::{self, EagerSelection};
 use super::binary::{Binary, binary};
 use super::cache::QsaCache;
@@ -34,12 +32,54 @@ 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) indexer: QsaIndexerOptions,
     pub(super) flash_min_context: u32,
     pub(super) prefill_gather: bool,
     pub(super) gather_tile_rows: i32,
 }
 
+/// Only decode_verify uses the fused lowering from MTPLX attention_math.py.
+/// In BF16 the value-preserving cast round trip selects that lowering without
+/// changing the sigmoid's dtype. All other phases and FP16 keep the stock graph.
+pub(super) fn attention_gate(
+    output: &Array,
+    gate: &Array,
+    verify: bool,
+    stream: Stream,
+) -> Result<Array, String> {
+    let dtype = gate.layout().dtype();
+    let body = |output: &Array, gate: &Array| {
+        binary(
+            output,
+            &unary::unary(gate, unary::Unary::Sigmoid, stream)?,
+            Binary::Multiply,
+            stream,
+        )
+    };
+    if !verify || !matches!(dtype, Dtype::BF16 | Dtype::F32) {
+        return body(output, gate);
+    }
+    let inputs = [output.clone(), gate.clone()];
+    let replay = super::compiled::Replay::trace_compiled(&inputs, |args| {
+        let gate = if dtype == Dtype::BF16 {
+            ops::astype(
+                &ops::astype(&args[1], Dtype::F32, false, stream)?,
+                Dtype::BF16,
+                false,
+                stream,
+            )?
+        } else {
+            args[1].clone()
+        };
+        Ok(vec![body(&args[0], &gate)?])
+    })?;
+    replay
+        .run(&inputs)?
+        .pop()
+        .ok_or_else(|| "compiled attention gate omitted output".into())
+}
+
 impl Attention<'_> {
     fn finish(
         &self,
@@ -47,13 +87,12 @@ impl Attention<'_> {
         gate: &Array,
         batch: i32,
         rows: i32,
+        verify: bool,
         stream: Stream,
     ) -> Result<Array, String> {
         let out = ops::reshape(out, &[batch, rows, -1], stream)?;
-        let gate = unary::unary(gate, unary::Unary::Sigmoid, stream)?;
-        let out = self
-            .output
-            .apply(&binary(&out, &gate, Binary::Multiply, stream)?, stream)?;
+        let out = attention_gate(&out, gate, verify, stream)?;
+        let out = self.output.apply(&out, stream)?;
         match self.output_bias {
             Some(bias) => binary(&out, bias, Binary::Add, stream),
             None => Ok(out),
@@ -211,7 +250,7 @@ impl Attention<'_> {
                 self.scale,
                 stream,
             )?;
-            return self.finish(&out, &gate, batch, rows, stream);
+            return self.finish(&out, &gate, batch, rows, options.verify_gate, stream);
         }
         if let Some(EagerSelection::PrefillBlocks(ids, valid)) = &selected {
             let keys = cache.kv.keys.as_ref().unwrap().borrow();
@@ -266,6 +305,7 @@ impl Attention<'_> {
                     &gate,
                     batch,
                     rows,
+                    options.verify_gate,
                     stream,
                 );
             }
@@ -279,6 +319,7 @@ impl Attention<'_> {
                 &gate,
                 batch,
                 rows,
+                options.verify_gate,
                 stream,
             );
         }
@@ -317,6 +358,7 @@ impl Attention<'_> {
             &gate,
             batch,
             rows,
+            options.verify_gate,
             stream,
         )
     }
@@ -503,6 +545,7 @@ fn mtplx_qsa_attention_module_matches_reference() {
                 }),
             };
             let options = Options {
+                verify_gate: false,
                 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_checkpoint.rs

@@ -30,6 +30,8 @@ struct Manifest {
     identity: Identity,
     tag: [u8; 32],
     #[serde(default)]
+    verify_gate: bool,
+    #[serde(default)]
     images: Vec<crate::engine::qwen_vision::Identity>,
     tokens: Vec<i32>,
     lazy_kv: bool,
@@ -116,6 +118,7 @@ impl LiveSession {
             check_abort(&abort)?;
             let offset = file.stream_position().map_err(|e| e.to_string())?;
             let manifest = serde_json::to_vec(&Manifest {
+                verify_gate: true,
                 identity,
                 tag,
                 images: self.current_images.clone(),
@@ -185,6 +188,11 @@ impl LiveSession {
         if !crate::engine::qwen_vision::compatible(&manifest.images) {
             return Err("Qwen vision cache policy differs; rebuild the cache".into());
         }
+        if identity.mtp && !manifest.verify_gate {
+            return Err(
+                "Qwen checkpoint uses obsolete MTP attention-gate math; rebuild the cache".into(),
+            );
+        }
         if !manifest.payload["mtp_history_snapshot"].is_null() != identity.mtp {
             return Err("Qwen checkpoint MTP payload differs".into());
         }
@@ -433,7 +441,46 @@ fn qwen_checkpoint_roundtrip_and_failed_restore_preserve_live_state() {
     }
     assert_eq!(live.tag, [9; 32]);
     assert_eq!(fingerprint(&live), original);
+    // Files predating the verify-phase law lack this field. Only MTP state
+    // changed numerically; old AR files must remain exactly restorable.
+    let without_gate_policy = |bytes: &[u8]| {
+        let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+        let mut manifest: serde_json::Value = serde_json::from_slice(&bytes[offset..]).unwrap();
+        manifest.as_object_mut().unwrap().remove("verify_gate");
+        let mut result = bytes[..offset].to_vec();
+        result.extend(serde_json::to_vec(&manifest).unwrap());
+        result
+    };
+    fs::write(&broken, without_gate_policy(&valid)).unwrap();
+    assert!(
+        live.load_checkpoint(&broken, identity, &streams, stream, |_| {}, || false)
+            .unwrap_err()
+            .contains("obsolete MTP attention-gate")
+    );
+    assert_eq!(fingerprint(&live), original);
+    live.current.as_mut().unwrap().mtp = None;
+    let ar_identity = Identity {
+        mtp: false,
+        ..identity
+    };
+    let ar_original = fingerprint(&live);
+    live.save_checkpoint(
+        &path,
+        ar_identity,
+        [10; 32],
+        &streams,
+        stream,
+        |_| {},
+        || false,
+    )
+    .unwrap();
+    fs::write(&broken, without_gate_policy(&fs::read(&path).unwrap())).unwrap();
+    assert!(
+        live.load_checkpoint(&broken, ar_identity, &streams, stream, |_| {}, || false)
+            .unwrap()
+    );
+    assert_eq!(fingerprint(&live), ar_original);
     println!(
-        "Qwen checkpoint: exact payload roundtrip, interrupted save, incompatible/corrupt restore passed"
+        "Qwen checkpoint: exact roundtrip, cancellation/corruption controls, old MTP rejected and old AR restored"
     );
 }

src/engine/metal/qwen_mtplx/speculative.rs

@@ -436,7 +436,7 @@ impl Execution {
             .collect::<Result<Vec<_>, _>>()
             .map_err(|_| "context-copy token exceeds int32")?;
         let started = Instant::now();
-        let verified = self.forward(&input, false, options.family_capture)?;
+        let verified = self.forward_verify(&input, options.family_capture)?;
         let logits = verified.logits.ok_or("copy verify omitted logits")?;
         let hidden = verified.hidden.ok_or("copy verify omitted hidden")?;
         let greedy = if sampler.temperature <= 0.0 {
@@ -492,7 +492,8 @@ impl Execution {
         }
         let keep = accepted + 1;
         let committed = if options.family_capture {
-            let graph = self.graph_options(false, false)?;
+            let mut graph = self.graph_options(false, false)?;
+            graph.qsa.verify_gate = true;
             self.model.with_dependent_mut(|_, model| {
                 model.commit_verified_window(
                     &mut self.cache,
@@ -513,7 +514,7 @@ impl Execution {
         if !committed {
             self.rollback_mtp_verify(&snapshots, input.len() as i32)?;
             let started = Instant::now();
-            let repaired = self.forward(&input[..1], false, false)?;
+            let repaired = self.forward_verify(&input[..1], false)?;
             let logits = repaired.logits.ok_or("copy repair omitted logits")?;
             let hidden = repaired.hidden.ok_or("copy repair omitted hidden")?;
             ops::evaluate(
@@ -810,7 +811,7 @@ impl Execution {
                     .map_err(|_| "MTP token exceeds int32")?,
             );
             let started = Instant::now();
-            let verified = self.forward(&input, false, options.family_capture)?;
+            let verified = self.forward_verify(&input, options.family_capture)?;
             let verify_logits = verified.logits.ok_or("verify omitted logits")?;
             let verify_hidden = verified.hidden.ok_or("verify omitted hidden")?;
             ops::evaluate(
@@ -911,7 +912,8 @@ impl Execution {
                 // MTPLX exits verify_capture_scope before replaying the kept
                 // prefix. Keeping capture armed here disables the fused GDN
                 // single-row path on the PLE layer's repair forward.
-                let graph = self.graph_options(false, false)?;
+                let mut graph = self.graph_options(false, false)?;
+                graph.qsa.verify_gate = true;
                 if options.family_capture {
                     cycle.capture_committed = self.model.with_dependent_mut(|_, model| {
                         model.commit_verified_window(
@@ -940,7 +942,7 @@ impl Execution {
                         .map(|&id| i32::try_from(id))
                         .collect::<Result<Vec<_>, _>>()
                         .map_err(|_| "MTP token exceeds int32")?;
-                    let repaired = self.forward(&ids, false, false)?;
+                    let repaired = self.forward_verify(&ids, false)?;
                     let logits = repaired.logits.ok_or("repair omitted logits")?;
                     let hidden = repaired.hidden.ok_or("repair omitted hidden")?;
                     if !options.defer_repair_eval {
@@ -999,10 +1001,9 @@ impl Execution {
                     result.history_seconds +=
                         self.append_committed(&result.hidden, &[token], options.lazy_history)?;
                 }
-                let output = self.forward(
+                let output = self.forward_verify(
                     &[i32::try_from(token).map_err(|_| "MTP token exceeds int32")?],
                     false,
-                    false,
                 )?;
                 let logits = output.logits.ok_or("final commit omitted logits")?;
                 let hidden = output.hidden.ok_or("final commit omitted hidden")?;