Match Qwen compiled verification gate math

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(())
+    })
+}