Record MTPLX upgrade defect reproductions

src/engine/metal/qwen_mtplx.rs

@@ -27,6 +27,8 @@ mod gdn;
 mod graph_sync;
 mod hyper;
 mod indexing;
+#[cfg(test)]
+mod issue103_analysis;
 mod linear;
 mod load;
 mod mlp;

src/engine/metal/qwen_mtplx/issue103_analysis.rs

@@ -0,0 +1,759 @@
+//! Isolated diagnostics for #103. Assertions pin observed defects, not fixes.
+//! Run serially with --ignored --nocapture; no production option is changed.
+use super::array::{Array, Dtype};
+use super::binary::{Binary, binary};
+use super::compiled::Replay;
+use super::sampling::{self, SamplerConfig};
+use super::stream::{Device, Streams};
+use super::{ops, random, unary};
+use crate::engine::Rng;
+
+fn f32_array(shape: &[i32], values: &[f32]) -> Array {
+    let bytes = values
+        .iter()
+        .flat_map(|v| v.to_le_bytes())
+        .collect::<Vec<_>>();
+    Array::new(shape, Dtype::F32, super::scalar_buffer(&bytes).unwrap()).unwrap()
+}
+
+#[test]
+#[ignore = "#103 diagnostic; Metal required, pins currently broken sampling law and invalid rows"]
+fn issue103_sampling_receipts() {
+    super::configure_sources().unwrap();
+    let _context = super::super::gpu::Context::open_qwen(0).unwrap();
+    let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(Device::Gpu).unwrap();
+    let values = (0..64)
+        .map(|i| match i {
+            0 => 0.4_f32.ln(),
+            1 => 0.3_f32.ln(),
+            _ => (0.3_f32 / 62.).ln(),
+        })
+        .collect::<Vec<_>>();
+    let row = f32_array(&[64], &values);
+    let config = SamplerConfig {
+        temperature: 1.,
+        top_p: 0.5,
+        top_k: 2,
+        ..Default::default()
+    };
+    let distribution =
+        sampling::distribution_from_logits(&row, &config, &[], &[], &streams, gpu).unwrap();
+    let mut ordinary = [0; 64];
+    let mut rng = Rng::new_qwen(0);
+    for _ in 0..512 {
+        ordinary[distribution.sample(&mut rng).unwrap() as usize] += 1;
+    }
+    let mut lazy = [0; 64];
+    for seed in 0..128 {
+        let token = sampling::lazy_sample(&row, &config, &random::key(seed).unwrap(), gpu).unwrap();
+        let id = ops::read_i64(&token, &streams, gpu).unwrap()[0];
+        lazy[id as usize] += 1;
+    }
+    println!(
+        "{}",
+        serde_json::json!({"case":"sampling_law", "temperature":1.,"top_k":2,"top_p":0.5,"ordinary_probabilities":[distribution.probability(0).unwrap(),distribution.probability(1).unwrap()],"ordinary_counts":&ordinary[..2],"lazy_counts":&lazy[..2],"expected_probabilities":[4./7.,3./7.]})
+    );
+    assert!(ordinary[0] > 0 && ordinary[1] > 0);
+    assert_eq!(
+        &lazy[..2],
+        &[128, 0],
+        "diagnosis no longer reproduces; update child regression"
+    );
+    // Official sampler, with logits scaled so its post-temperature mass is known.
+    for (temperature, top_k, top_p) in [(0.6, 20, 0.95), (0.2, 8, 0.8), (1.0, 2, 1.0)] {
+        let row = f32_array(
+            &[64],
+            &values
+                .iter()
+                .map(|v| v * temperature as f32)
+                .collect::<Vec<_>>(),
+        );
+        let config = SamplerConfig {
+            temperature,
+            top_k,
+            top_p,
+            ..Default::default()
+        };
+        let distribution =
+            sampling::distribution_from_logits(&row, &config, &[], &[], &streams, gpu).unwrap();
+        let probabilities = (0..64)
+            .map(|i| distribution.probability(i).unwrap())
+            .collect::<Vec<_>>();
+        let mut counts = [0; 64];
+        for seed in 0..512 {
+            let token =
+                sampling::lazy_sample(&row, &config, &random::key(seed).unwrap(), gpu).unwrap();
+            counts[ops::read_i64(&token, &streams, gpu).unwrap()[0] as usize] += 1;
+        }
+        println!(
+            "{}",
+            serde_json::json!({"case":"sampling_matrix","temperature":temperature,"top_k":top_k,"top_p":top_p,"ordinary_probabilities":probabilities,"lazy_counts":counts.to_vec()})
+        );
+    }
+    for (name, values) in [
+        ("nan", vec![f32::NAN; 64]),
+        ("positive_inf", [vec![f32::INFINITY], vec![0.; 63]].concat()),
+        ("all_negative_inf", vec![f32::NEG_INFINITY; 64]),
+        (
+            "valid_mask",
+            [
+                vec![f32::NEG_INFINITY, 2., f32::NEG_INFINITY, 1.],
+                vec![f32::NEG_INFINITY; 60],
+            ]
+            .concat(),
+        ),
+    ] {
+        let row = f32_array(&[64], &values);
+        for temperature in [0., 0.6, 1.] {
+            let config = SamplerConfig {
+                temperature,
+                ..Default::default()
+            };
+            let result = sampling::sample_from_logits(
+                &row,
+                &config,
+                &mut Rng::new_qwen(0),
+                &[],
+                &[],
+                &streams,
+                gpu,
+            );
+            println!(
+                "{}",
+                serde_json::json!({"case":"invalid_logits","row":name,"temperature":temperature,"ordinary":result.as_ref().map(|r|r.0).map_err(Clone::clone)})
+            );
+            if name == "valid_mask" {
+                assert!([1, 3].contains(&result.unwrap().0));
+            } else {
+                assert!(result.is_ok(), "diagnosis no longer reproduces");
+            }
+        }
+        let config = SamplerConfig {
+            temperature: 1.,
+            ..Default::default()
+        };
+        let result = sampling::lazy_sample(&row, &config, &random::key(7).unwrap(), gpu)
+            .and_then(|a| ops::read_i64(&a, &streams, gpu));
+        println!(
+            "{}",
+            serde_json::json!({"case":"invalid_logits_lazy","row":name,"result":result})
+        );
+    }
+}
+
+#[test]
+#[ignore = "#103 diagnostic; Metal required, eager versus compiled numeric operands"]
+fn issue103_numeric_receipts() {
+    super::configure_sources().unwrap();
+    let _context = super::super::gpu::Context::open_qwen(0).unwrap();
+    let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(Device::Gpu).unwrap();
+    let gate = ops::astype(&f32_array(&[4], &[-6.84375; 4]), Dtype::BF16, false, gpu).unwrap();
+    let value = ops::astype(&f32_array(&[4], &[1.; 4]), Dtype::BF16, false, gpu).unwrap();
+    ops::evaluate(&streams, &[gate.clone(), value.clone()], gpu, false).unwrap();
+    let body = |a: &[Array]| {
+        binary(
+            &unary::unary(&a[0], unary::Unary::Sigmoid, gpu)?,
+            &a[1],
+            Binary::Multiply,
+            gpu,
+        )
+        .map(|o| vec![o])
+    };
+    let eager = body(&[gate.clone(), value.clone()]).unwrap();
+    let compiled = Replay::trace_compiled(&[gate.clone(), value.clone()], body)
+        .unwrap()
+        .run(&[gate, value])
+        .unwrap();
+    let read = |a: &Array| {
+        ops::read_f32_as_f64(
+            &ops::astype(a, Dtype::F32, false, gpu).unwrap(),
+            &streams,
+            gpu,
+        )
+        .unwrap()
+    };
+    let e = read(&eager[0]);
+    let c = read(&compiled[0]);
+    println!(
+        "{}",
+        serde_json::json!({"case":"bf16_gate","gate":-6.84375,"eager":e,"compiled":c})
+    );
+    assert_ne!(e, c, "diagnosis no longer reproduces");
+    for (name, operand, mode) in [
+        ("qsa_divisor", 128_f32.sqrt(), Binary::Divide),
+        (
+            "yarn_amplitude",
+            (0.1 * 4_f64.ln() + 1.) as f32,
+            Binary::Multiply,
+        ),
+    ] {
+        let x = f32_array(&[4], &[1.5000001, 1.5000002, 1., 0.]);
+        let scalar = f32_array(&[], &[operand]);
+        let body = |a: &[Array]| binary(&a[0], &scalar, mode, gpu).map(|o| vec![o]);
+        let eager = body(std::slice::from_ref(&x)).unwrap();
+        let compiled = Replay::trace_compiled(std::slice::from_ref(&x), body)
+            .unwrap()
+            .run(&[x])
+            .unwrap();
+        let e = read(&eager[0]);
+        let c = read(&compiled[0]);
+        println!(
+            "{}",
+            serde_json::json!({"case":name,"operand_bits":operand.to_bits(),"eager":e,"compiled":c})
+        );
+        assert_ne!(e, c, "diagnosis no longer reproduces");
+    }
+}
+
+#[test]
+#[ignore = "#103 diagnostic; M5 and pinned runtime shader required; initialized sentinels"]
+fn issue103_sorted_gather_receipts() {
+    super::configure_sources().unwrap();
+    let _context = super::super::gpu::Context::open_qwen(0).unwrap();
+    assert!(super::sparse_attention::nax_available());
+    let scope = super::encoder::Scope::new().unwrap();
+    let buffer =
+        |count: usize, bits: u16| super::scalar_buffer(&bits.to_le_bytes().repeat(count)).unwrap();
+    let w = super::scalar_buffer(&vec![0; 64 * 256 * 256 / 8 * 4]).unwrap();
+    let scales = buffer(64 * 256 * 8, 0);
+    let biases = buffer(64 * 256 * 8, 0x3f80); // all dequantized weights exactly 1
+    for tokens in [3276_u32, 3277, 3404, 4095, 4096] {
+        let rows = tokens * 10;
+        let x = buffer(rows as usize * 256, 0x3f80);
+        let lhs = super::scalar_buffer(&(0..rows).flat_map(u32::to_le_bytes).collect::<Vec<_>>())
+            .unwrap();
+        let rhs = super::scalar_buffer(
+            &(0..rows)
+                .flat_map(|r| (r * 64 / rows).to_le_bytes())
+                .collect::<Vec<_>>(),
+        )
+        .unwrap();
+        let y = buffer(rows as usize * 256, 0x7fc1); // initialized NaN, independent of allocator reuse
+        let (_, dispatches) = super::super::qwen_mtplx_tests::capture_dispatches(|| {
+            super::gather_bf16(
+                [&w, &scales, &biases, &x, &lhs, &rhs, &y],
+                [rows, 1, 256, 256, 64, rows],
+                [4, 32],
+                true,
+                true,
+            )
+            .unwrap();
+        });
+        scope.synchronize().unwrap();
+        let mut data = vec![0; rows as usize * 256 * 2];
+        y.read(0, &mut data).unwrap();
+        let bad = data
+            .chunks_exact(512)
+            .enumerate()
+            .filter_map(|(r, bytes)| {
+                bytes
+                    .chunks_exact(2)
+                    .any(|b| u16::from_le_bytes(b.try_into().unwrap()) != 0x4380)
+                    .then_some(r)
+            })
+            .collect::<Vec<_>>();
+        println!(
+            "{}",
+            serde_json::json!({"case":"sorted_gather","tokens":tokens,"routed_rows":rows,"bad_rows":bad.len(),"first_bad":bad.first(),"last_bad":bad.last(),"dispatches":format!("{dispatches:?}")})
+        );
+        if tokens == 3276 || tokens == 4096 {
+            assert!(bad.is_empty());
+        } else {
+            assert!(!bad.is_empty(), "diagnosis no longer reproduces");
+        }
+    }
+}
+
+#[test]
+#[ignore = "#103 diagnostic; installed tokenizer and latest oracle JSONL; no GPU"]
+fn issue103_tokenizer_receipts() {
+    let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
+    let tokenizer =
+        crate::engine::tokenizer::Tokenizer::load_qwen(&root.join("tokenizer.json")).unwrap();
+    let input = std::fs::read_to_string(std::env::var_os("DS4_ISSUE103_ORACLE").unwrap()).unwrap();
+    let mut mismatches = 0;
+    for case in input
+        .lines()
+        .filter_map(|l| serde_json::from_str::<serde_json::Value>(l).ok())
+        .filter(|c| c["case"] == "tokenizer")
+    {
+        let text = case["text"].as_str().unwrap();
+        let ids = tokenizer.tokenize(text);
+        let matches = serde_json::json!(ids) == case["ids"];
+        mismatches += usize::from(!matches);
+        println!(
+            "{}",
+            serde_json::json!({"case":"native_tokenizer","text":text,"actual_ids":ids,"expected_ids":case["ids"],"matches":matches})
+        );
+    }
+    println!(
+        "{}",
+        serde_json::json!({"case":"tokenizer_summary","mismatches":mismatches})
+    );
+    assert_eq!(mismatches, 5, "diagnosis changed; inspect exact token IDs");
+}
+
+#[test]
+#[ignore = "#103 diagnostic; full pinned model, SSD PLE and external memory supervisor"]
+fn issue103_mtp_final_receipts() {
+    use super::decoder::AttentionCache;
+    use super::execution::Execution;
+    use super::speculative::Options;
+    let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
+    assert!(!super::ngram::resident_policy(), "use SSD PLE");
+    eprintln!("#103 loading unchanged Bare-Speed model");
+    let mut execution = Execution::load(&root, true, false, 100).unwrap();
+    let ids = [9707, 11, 1075]; // fixed valid IDs; no template is involved in this model-level repro
+    let prepared = execution
+        .prefill_mtp_history(
+            &ids,
+            &super::prefill::Options {
+                chunk_size: Some(2048),
+                final_logits_only: true,
+                external_cache_only: false,
+                external_emit_logits: true,
+                return_hidden: true,
+            },
+            None,
+            false,
+            |_| Ok(()),
+            |_| Ok(()),
+        )
+        .unwrap();
+    let baseline = execution
+        .capture_session(
+            &ids,
+            &prepared.target.logits,
+            prepared.target.hidden.as_ref(),
+            true,
+            false,
+        )
+        .unwrap();
+    let offset = |e: &Execution| {
+        e.cache
+            .iter()
+            .filter_map(|c| match &c.as_ref()?.attention {
+                AttentionCache::Qsa(c) => Some(c.kv.offset),
+                _ => None,
+            })
+            .collect::<Vec<_>>()
+    };
+    let mut final_token = 0;
+    for depth in [1, 2, 3] {
+        for temperature in [0., 0.6] {
+            for stop_exit in [false, true] {
+                let (logits, hidden, _) = execution.restore_session(&baseline, &ids).unwrap();
+                let sampler = SamplerConfig {
+                    temperature,
+                    ..Default::default()
+                };
+                let result = execution
+                    .decode_mtp(
+                        &logits,
+                        hidden.as_ref().unwrap(),
+                        &sampler,
+                        &sampler,
+                        &Options {
+                            prompt_ids: &ids,
+                            depth,
+                            max_tokens: if stop_exit { 8 } else { 1 },
+                            seed: 0,
+                            stops: if stop_exit {
+                                if temperature == 0. { &[11] } else { &[314] }
+                            } else {
+                                &[]
+                            },
+                            committed_history: true,
+                            family_capture: true,
+                            lazy_history: false,
+                            defer_repair_eval: false,
+                            capture_final_state: true,
+                        },
+                        |_| Ok(()),
+                        |_| Ok(()),
+                    )
+                    .unwrap();
+                let mut key = ids.to_vec();
+                key.push(result.tokens[0] as i32);
+                let snapshot = execution
+                    .capture_session(&key, &result.logits, Some(&result.hidden), true, false)
+                    .unwrap();
+                let (.., reused) = execution.restore_session(&snapshot, &key).unwrap();
+                let offsets = offset(&execution);
+                println!(
+                    "{}",
+                    serde_json::json!({"case":"mtp_final","exit":if stop_exit {"first_primary_stop"} else {"max_tokens"},"depth":depth,"temperature":temperature,"prompt_ids":ids,"tokens":result.tokens,"key_length":key.len(),"restored_cached_tokens":reused,"trunk_offsets":offsets,"mtp_offset":execution.mtp_offset().unwrap(),"pending":result.pending,"final_state_error":result.final_state_error})
+                );
+                assert_eq!(offsets, vec![3; 12]);
+                assert_eq!(reused, 4);
+                assert!(result.pending.is_none() && result.final_state_error.is_none());
+                final_token = result.tokens[0] as i32;
+            }
+        }
+    }
+    execution.restore_session(&baseline, &ids).unwrap();
+    let bad = execution
+        .forward(&[1075], false, false)
+        .unwrap()
+        .logits
+        .unwrap();
+    let bad = execution.values_f32(&bad).unwrap();
+    execution.restore_session(&baseline, &ids).unwrap();
+    let good = execution
+        .forward(&[final_token, 1075], false, false)
+        .unwrap()
+        .logits
+        .unwrap();
+    let good = execution
+        .values_f32(&super::ar::last_row(&good, execution.stream).unwrap())
+        .unwrap();
+    let different = bad
+        .iter()
+        .zip(&good)
+        .filter(|(a, b)| a.to_bits() != b.to_bits())
+        .count();
+    println!(
+        "{}",
+        serde_json::json!({"case":"mtp_next_turn","missing_token":final_token,"suffix_ids":[1075],"different_logits":different,"bad_sha256":super::super::qwen_mtplx_tests::hex_digest(&bad.iter().flat_map(|v|v.to_le_bytes()).collect::<Vec<_>>()),"good_sha256":super::super::qwen_mtplx_tests::hex_digest(&good.iter().flat_map(|v|v.to_le_bytes()).collect::<Vec<_>>())})
+    );
+    assert!(different > 0);
+}
+
+#[test]
+#[ignore = "#103 diagnostic; small active attention module on Metal; mixed position boundary"]
+fn issue103_vision_position_receipts() {
+    use super::cache::QsaCache;
+    use super::mlp::Linear;
+    use super::qsa_attention::{Attention, Input, Options};
+    super::configure_sources().unwrap();
+    let _context = super::super::gpu::Context::open_qwen(0).unwrap();
+    let streams = Streams::new(Some(super::allocator::Allocator::new().unwrap()));
+    let gpu = streams.default_stream(Device::Gpu).unwrap();
+    let bf = |shape: &[i32], v: &[f32]| {
+        ops::astype(&f32_array(shape, v), Dtype::BF16, false, gpu).unwrap()
+    };
+    let matrix = |n: usize| {
+        bf(
+            &[n as i32, 8],
+            &(0..n * 8)
+                .map(|i| f32::from(i / 8 % 8 == i % 8))
+                .collect::<Vec<_>>(),
+        )
+    };
+    let q = matrix(16);
+    let k = matrix(8);
+    let v = matrix(8);
+    let out = matrix(8);
+    let norm = bf(&[8], &[1.; 8]);
+    let freq = f32_array(&[4], &[1., 0.1, 0.01, 0.001]);
+    let axes = Array::new(
+        &[4],
+        Dtype::I32,
+        super::scalar_buffer(
+            &[0_i32, 1, 2, 0]
+                .into_iter()
+                .flat_map(i32::to_le_bytes)
+                .collect::<Vec<_>>(),
+        )
+        .unwrap(),
+    )
+    .unwrap();
+    let table = Array::new(
+        &[3, 4],
+        Dtype::I32,
+        super::scalar_buffer(
+            &[0_i32, 1, 2, 2, 0, 1, 3, 3, 0, 1, 4, 4]
+                .into_iter()
+                .flat_map(i32::to_le_bytes)
+                .collect::<Vec<_>>(),
+        )
+        .unwrap(),
+    )
+    .unwrap();
+    let options = Options {
+        indexer: super::QsaIndexerOptions {
+            fused: false,
+            compiled: false,
+            prefill: false,
+            phase_prefill: true,
+            prefill_min_rows: 128,
+            prefill_compile_rows: 2048,
+            prefill_min_context: 16384,
+            prefill_workspace: 0,
+            tile_rows: 64,
+            flash: false,
+            gather_decode: false,
+            gather: false,
+            gather_max_rows: 4,
+            gather_min_context: 0,
+            tf32: false,
+        },
+        flash_min_context: 0,
+        prefill_gather: false,
+        gather_tile_rows: 64,
+    };
+    for (pos, rows) in [(0, 4), (4, 2), (2, 4)] {
+        let mut attention = Attention {
+            input: Input::Separate(
+                [Linear::Dense(&q), Linear::Dense(&k), Linear::Dense(&v)],
+                [None; 3],
+            ),
+            output: Linear::Dense(&out),
+            output_bias: None,
+            heads: 1,
+            kv_heads: 1,
+            dim: 8,
+            eps: 1e-6,
+            scale: 1. / 8_f64.sqrt(),
+            q_norm: &norm,
+            k_norm: &norm,
+            inv_freq: &freq,
+            rope_scaling: 1.,
+            mrope_axes: Some(&axes),
+            indexer: None,
+        };
+        let mut cache = QsaCache::new(4);
+        if pos > 0 {
+            let prefix = bf(&[1, 1, pos, 8], &vec![0.; pos as usize * 8]);
+            cache.kv.update_and_fetch(&prefix, &prefix, gpu).unwrap();
+        }
+        let x = bf(
+            &[1, rows, 8],
+            &(0..rows * 8)
+                .map(|i| (i + 1) as f32 / 32.)
+                .collect::<Vec<_>>(),
+        );
+        let output = attention
+            .call(
+                &x,
+                &mut cache,
+                &options,
+                Some((Some(&table), -1)),
+                &streams,
+                gpu,
+            )
+            .unwrap();
+        ops::evaluate(&streams, &[output], gpu, false).unwrap();
+        let projected = ops::reshape(
+            &ops::rms_norm(&x, Some(&norm), 1e-6, gpu).unwrap(),
+            &[1, rows, 1, 8],
+            gpu,
+        )
+        .unwrap();
+        let positions = (0..3)
+            .flat_map(|axis| {
+                (pos..pos + rows).map(move |p| {
+                    if p < 4 {
+                        [vec![0, 1, 2, 2], vec![0, 1, 3, 3], vec![0, 1, 4, 4]][axis][p as usize]
+                    } else {
+                        p - 1
+                    }
+                })
+            })
+            .collect::<Vec<i32>>();
+        let positions = Array::new(
+            &[3, rows],
+            Dtype::I32,
+            super::scalar_buffer(
+                &positions
+                    .iter()
+                    .flat_map(|p| p.to_le_bytes())
+                    .collect::<Vec<_>>(),
+            )
+            .unwrap(),
+        )
+        .unwrap();
+        let (cos, sin) = super::rope::mrope_cos_sin(&positions, &freq, &axes, gpu).unwrap();
+        let expected = super::rope::apply_partial(&projected, &cos, &sin, gpu).unwrap();
+        let actual = ops::slice(
+            &cache.kv.keys.as_ref().unwrap().borrow(),
+            &[0, 0, pos, 0],
+            &[1, 1, pos + rows, 8],
+            &[1; 4],
+            gpu,
+        )
+        .unwrap();
+        let read = |a: &Array| {
+            ops::read_f32_as_f64(
+                &ops::astype(a, Dtype::F32, false, gpu).unwrap(),
+                &streams,
+                gpu,
+            )
+            .unwrap()
+        };
+        let a = read(&actual);
+        let e = read(&expected);
+        let bad = a
+            .chunks_exact(8)
+            .zip(e.chunks_exact(8))
+            .enumerate()
+            .filter_map(|(i, (a, e))| (a != e).then_some(i))
+            .collect::<Vec<_>>();
+        println!(
+            "{}",
+            serde_json::json!({"case":"vision_positions","start":pos,"rows":rows,"table_end":4,"delta":-1,"different_rows":bad,"actual_keys":a,"expected_keys":e})
+        );
+        if pos == 2 {
+            assert_eq!(bad, vec![0, 1]);
+        } else {
+            assert!(bad.is_empty());
+        }
+    }
+    let prefix = bf(
+        &[1, 12, 8],
+        &(0..96)
+            .map(|i| ((i * 17 % 31) as f32 - 15.) / 16.)
+            .collect::<Vec<_>>(),
+    );
+    let current = bf(&[1, 1, 8], &[0.3, -0.7, 0.2, 0.9, -0.2, 0.4, 0.6, -0.1]);
+    let mut outputs = Vec::new();
+    for vision in [false, true] {
+        let mut attention = Attention {
+            input: Input::Separate(
+                [Linear::Dense(&q), Linear::Dense(&k), Linear::Dense(&v)],
+                [None; 3],
+            ),
+            output: Linear::Dense(&out),
+            output_bias: None,
+            heads: 1,
+            kv_heads: 1,
+            dim: 8,
+            eps: 1e-6,
+            scale: 1. / 8_f64.sqrt(),
+            q_norm: &norm,
+            k_norm: &norm,
+            inv_freq: &freq,
+            rope_scaling: 1.,
+            mrope_axes: Some(&axes),
+            indexer: Some(super::qsa_indexer::Indexer {
+                projection: Some(Linear::Dense(&q)),
+                heads: 1,
+                kv_heads: 1,
+                head_dim: 8,
+                block_topk: 1,
+                ratio: 4,
+                q_norm: &norm,
+                k_norm: &norm,
+                inv_freq: &freq,
+                eps: 1e-6,
+                scaling: 1.,
+                scratch: 32 * 1024 * 1024,
+                compiled: None,
+            }),
+        };
+        let mut cache = QsaCache::new(4);
+        let initial = attention
+            .call(&prefix, &mut cache, &options, None, &streams, gpu)
+            .unwrap();
+        ops::evaluate(&streams, &[initial], gpu, false).unwrap();
+        let output = attention
+            .call(
+                &current,
+                &mut cache,
+                &options,
+                vision.then_some((None, 0)),
+                &streams,
+                gpu,
+            )
+            .unwrap();
+        outputs.push(
+            ops::read_f32_as_f64(
+                &ops::astype(&output, Dtype::F32, false, gpu).unwrap(),
+                &streams,
+                gpu,
+            )
+            .unwrap(),
+        );
+    }
+    println!(
+        "{}",
+        serde_json::json!({"case":"vision_sparse_qsa","prefix_rows":12,"topk_blocks":1,"ratio":4,"position_delta":0,"text_output":outputs[0],"vision_output":outputs[1]})
+    );
+    assert_ne!(outputs[0], outputs[1], "diagnosis no longer reproduces");
+    let qweight = bf(
+        &[16, 8],
+        &(0..128)
+            .map(|i| f32::from(i / 8 < 8 && i / 8 == i % 8))
+            .collect::<Vec<_>>(),
+    );
+    let mut gate_outputs = Vec::new();
+    for gate in [0., -6.84375] {
+        let bias = bf(&[16], &[vec![0.; 8], vec![gate; 8]].concat());
+        let mut attention = Attention {
+            input: Input::Separate(
+                [
+                    Linear::Dense(&qweight),
+                    Linear::Dense(&k),
+                    Linear::Dense(&v),
+                ],
+                [Some(&bias), None, None],
+            ),
+            output: Linear::Dense(&out),
+            output_bias: None,
+            heads: 1,
+            kv_heads: 1,
+            dim: 8,
+            eps: 1e-6,
+            scale: 1. / 8_f64.sqrt(),
+            q_norm: &norm,
+            k_norm: &norm,
+            inv_freq: &freq,
+            rope_scaling: 1.,
+            mrope_axes: Some(&axes),
+            indexer: None,
+        };
+        let output = attention
+            .call(
+                &current,
+                &mut QsaCache::new(4),
+                &options,
+                None,
+                &streams,
+                gpu,
+            )
+            .unwrap();
+        ops::evaluate(&streams, std::slice::from_ref(&output), gpu, false).unwrap();
+        gate_outputs.push(output);
+    }
+    let ungated = binary(&gate_outputs[0], &bf(&[], &[2.]), Binary::Multiply, gpu).unwrap();
+    let gate = bf(&[1, 1, 8], &[-6.84375; 8]);
+    ops::evaluate(&streams, &[ungated.clone(), gate.clone()], gpu, false).unwrap();
+    let args = [ungated, gate];
+    let replay = Replay::trace_compiled(&args, |a| {
+        binary(
+            &a[0],
+            &unary::unary(
+                &ops::astype(
+                    &ops::astype(&a[1], Dtype::F32, false, gpu)?,
+                    Dtype::BF16,
+                    false,
+                    gpu,
+                )?,
+                unary::Unary::Sigmoid,
+                gpu,
+            )?,
+            Binary::Multiply,
+            gpu,
+        )
+        .map(|a| vec![a])
+    })
+    .unwrap();
+    let expected = replay.run(&args).unwrap();
+    let read = |a: &Array| {
+        ops::read_f32_as_f64(
+            &ops::astype(a, Dtype::F32, false, gpu).unwrap(),
+            &streams,
+            gpu,
+        )
+        .unwrap()
+    };
+    let actual = read(&gate_outputs[1]);
+    let expected = read(&expected[0]);
+    println!(
+        "{}",
+        serde_json::json!({"case":"active_attention_gate","gate":-6.84375,"actual":actual,"expected_verify":expected})
+    );
+    assert_ne!(actual, expected, "diagnosis no longer reproduces");
+}