Preserve image-aware sparse QSA for Qwen vision

src/engine/metal/qwen_mtplx/binding.rs

@@ -847,6 +847,7 @@ impl<'a> Bind<'a> {
                 q_norm: self.array(&format!("{ip}.q_layernorm.weight"))?,
                 k_norm: self.array(&format!("{ip}.k_layernorm.weight"))?,
                 inv_freq: &r.frequency,
+                mrope_axes: r.axes.as_ref(),
                 eps: c.rms_norm_eps,
                 scaling: r.scaling,
                 scratch: 32 * 1024 * 1024,

src/engine/metal/qwen_mtplx/indexing.rs

@@ -1,5 +1,5 @@
 //! Select, take/Gather, GatherAxis, Arange and Scan from the pinned runtime.
-//! take exports BF16/F16/F32/U32 with I64 and I32 indices.
+//! take exports BF16/F16/F32/I32/U32 with I64 and I32 indices.
 use super::super::gpu::Buffer;
 use super::array::{Array, Dtype, Layout};
 use super::stream::Stream;
@@ -503,7 +503,7 @@ pub(super) fn take(table: &Array, ids: &Array, axis: i32, stream: Stream) -> Res
         || !(il.dtype() == Dtype::I64 || il.dtype() == Dtype::U32 || il.dtype() == Dtype::I32)
         || !matches!(
             sl.dtype(),
-            Dtype::BF16 | Dtype::F16 | Dtype::F32 | Dtype::U32
+            Dtype::BF16 | Dtype::F16 | Dtype::F32 | Dtype::I32 | Dtype::U32
         )
         || (sl.shape()[axis] == 0 && il.size() != 0)
     {
@@ -551,7 +551,8 @@ fn gather_take_evaluate(
         Dtype::BF16 => "bf16",
         Dtype::F16 => "f16",
         Dtype::F32 => "f32",
-        Dtype::U32 => "u32",
+        // Gather copies 32-bit payloads without value arithmetic.
+        Dtype::I32 | Dtype::U32 => "u32",
         _ => unreachable!(),
     };
     let large = sl.size() > i32::MAX as usize

src/engine/metal/qwen_mtplx/issue103_analysis.rs

@@ -532,6 +532,7 @@ fn issue103_vision_position_receipts() {
                 q_norm: &norm,
                 k_norm: &norm,
                 inv_freq: &freq,
+                mrope_axes: Some(&axes),
                 eps: 1e-6,
                 scaling: 1.,
                 scratch: 32 * 1024 * 1024,
@@ -566,7 +567,10 @@ fn issue103_vision_position_receipts() {
         "{}",
         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");
+    assert_eq!(
+        outputs[0], outputs[1],
+        "equal vision positions must preserve sparse QSA"
+    );
     let qweight = bf(
         &[16, 8],
         &(0..128)

src/engine/metal/qwen_mtplx/issue107_indexer.rs

@@ -0,0 +1,180 @@
+use super::*;
+
+#[test]
+#[ignore = "Metal and independently verified pinned indexer fixtures; external supervisor required"]
+fn issue107_real_geometry_query_and_block_start_positions() {
+    super::super::configure_sources().unwrap();
+    let _context = crate::engine::metal::gpu::Context::open_qwen(0).unwrap();
+    let streams = super::super::stream::Streams::new(Some(
+        super::super::allocator::Allocator::new().unwrap(),
+    ));
+    let stream = streams
+        .default_stream(super::super::stream::Device::Gpu)
+        .unwrap();
+    let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_INDEXER").unwrap());
+    let ints = |shape: &[i32], data: &[i32]| {
+        Array::new(
+            shape,
+            Dtype::I32,
+            super::super::scalar_buffer(
+                &data
+                    .iter()
+                    .flat_map(|i| i.to_le_bytes())
+                    .collect::<Vec<_>>(),
+            )
+            .unwrap(),
+        )
+        .unwrap()
+    };
+    let axes = ints(&[32], &rope::build_mrope_axes(&[11, 11, 10], true));
+    // I32 Gather reuses the bit-preserving U32 copy kernel, including signed
+    // extrema that a detour through F32 could not preserve.
+    let signed = ints(&[2, 3], &[i32::MIN, -7, i32::MAX, 31, 0, -1]);
+    for (axis, indices, expected) in [
+        (1, vec![2, 0], vec![i32::MAX, i32::MIN, -1, 31]),
+        (0, vec![1, 0], vec![31, 0, -1, i32::MIN, -7, i32::MAX]),
+    ] {
+        let copied = indexing::take(&signed, &ints(&[2], &indices), axis, stream).unwrap();
+        assert_eq!(
+            ops::read_i64(&copied, &streams, stream).unwrap(),
+            expected.into_iter().map(i64::from).collect::<Vec<_>>()
+        );
+    }
+    let (freq, scaling) = rope::inv_freq_and_scaling(64, 10000000., None, stream).unwrap();
+    for (name, dtype) in [("f32", Dtype::F32), ("bf16", Dtype::BF16)] {
+        let pattern = |shape: &[i32], seed: usize| {
+            let values = (0..shape.iter().map(|&n| n as usize).product())
+                .map(|i| (((i * 17 + seed) % 97) as f32 - 48.) / 32.)
+                .collect::<Vec<_>>();
+            ops::astype(
+                &super::super::vision::f32_array(shape, &values).unwrap(),
+                dtype,
+                false,
+                stream,
+            )
+            .unwrap()
+        };
+        let norm = ops::astype(
+            &super::super::vision::f32_array(&[128], &[1.; 128]).unwrap(),
+            dtype,
+            false,
+            stream,
+        )
+        .unwrap();
+        let query = pattern(&[1, 16, 4, 128], 3);
+        let raw = pattern(&[1, 2072, 128], 11);
+        let mut indexer = Indexer {
+            projection: None,
+            heads: 4,
+            kv_heads: 1,
+            head_dim: 128,
+            block_topk: 512,
+            ratio: 4,
+            q_norm: &norm,
+            k_norm: &norm,
+            inv_freq: &freq,
+            mrope_axes: Some(&axes),
+            eps: 1e-6,
+            scaling,
+            scratch: 32 * 1024 * 1024,
+            compiled: None,
+        };
+        for mode in ["table", "equal", "delta", "empty"] {
+            let delta = if mode == "equal" { 0 } else { -7 };
+            let table = match mode {
+                "table" => Some(ints(
+                    &[3, 2060],
+                    &(0..3)
+                        .flat_map(|axis| (0..2060).map(move |p| p / (axis + 1)))
+                        .collect::<Vec<_>>(),
+                )),
+                "equal" => Some(ints(
+                    &[3, 2060],
+                    &(0..3).flat_map(|_| 0..2060).collect::<Vec<_>>(),
+                )),
+                "empty" => Some(ints(&[3, 0], &[])),
+                _ => None,
+            };
+            let vision = Some((table.as_ref(), delta));
+            let q = indexer
+                .prepare_queries(&query, 2056, true, vision, stream)
+                .unwrap();
+            let k = indexer
+                .pool_keys(&raw, 0, 518, true, vision, stream)
+                .unwrap();
+            for (part, array) in [("query", q), ("pooled", k)] {
+                let got = ops::read_f32_as_f64(
+                    &ops::astype(&array, Dtype::F32, false, stream).unwrap(),
+                    &streams,
+                    stream,
+                )
+                .unwrap();
+                let expected =
+                    std::fs::read(directory.join(format!("{name}-{mode}-{part}.f32"))).unwrap();
+                let actual = got
+                    .iter()
+                    .flat_map(|v| (*v as f32).to_le_bytes())
+                    .collect::<Vec<_>>();
+                assert_eq!(
+                    crate::engine::metal::qwen_mtplx_tests::hex_digest(&actual),
+                    crate::engine::metal::qwen_mtplx_tests::hex_digest(&expected),
+                    "{name}/{mode}/{part}"
+                );
+                println!(
+                    "{}",
+                    serde_json::json!({"case":"indexer","dtype":name,"mode":mode,"part":part,"exact_oracle_match":true})
+                );
+            }
+            let mut cache = QsaCache::new(4);
+            cache.kv.offset = 2056;
+            cache.raw_keys = Some(super::super::cache::object(
+                ops::slice(&raw, &[0, 0, 0], &[1, 2056, 128], &[1; 3], stream).unwrap(),
+            ));
+            let qk = ops::concatenate(
+                &[
+                    ops::reshape(&query, &[1, 16, 512], stream).unwrap(),
+                    ops::slice(&raw, &[0, 2056, 0], &[1, 2072, 128], &[1; 3], stream).unwrap(),
+                ],
+                -1,
+                stream,
+            )
+            .unwrap();
+            let options = super::super::QsaIndexerOptions {
+                fused: true,
+                compiled: true,
+                prefill: false,
+                phase_prefill: false,
+                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,
+            };
+            let (selection, lane) = indexer
+                .call(
+                    &qk,
+                    2056,
+                    &mut cache,
+                    Some(&qk),
+                    &options,
+                    &streams,
+                    vision,
+                    stream,
+                )
+                .unwrap();
+            assert_ne!(lane, "compiled");
+            assert_eq!(cache.pooled_len, 518);
+            let super::super::attention::EagerSelection::Dense(mask) = selection.unwrap() else {
+                panic!("expected sparse block selection mask")
+            };
+            ops::evaluate(&streams, &[mask], stream, false).unwrap();
+        }
+    }
+    streams.clear_streams().unwrap();
+}

src/engine/metal/qwen_mtplx/issue107_product.rs

@@ -0,0 +1,328 @@
+//! Long multi-image request and cache compatibility through the product owner.
+use super::*;
+use crate::settings::{
+    GenerationPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences,
+    effective_settings,
+};
+use base64::Engine as _;
+
+fn state(owner: &QwenExecutor) -> String {
+    let saved = owner.live.current.as_ref().unwrap();
+    let mut hashes = Vec::new();
+    let payload = saved
+        .encode_payload(
+            256,
+            &owner.execution.streams,
+            owner.execution.stream,
+            |name, bytes| {
+                hashes.push((
+                    name,
+                    crate::engine::metal::qwen_mtplx_tests::hex_digest(&bytes),
+                ));
+                Ok(())
+            },
+            || false,
+        )
+        .unwrap();
+    crate::engine::metal::qwen_mtplx_tests::hex_digest(&serde_json::to_vec(&serde_json::json!({"ids":saved.token_ids,"images":owner.live.current_images,"payload":payload,"hashes":hashes})).unwrap())
+}
+
+fn legacy_live(owner: &QwenExecutor, path: &Path) -> LiveSession {
+    // Import the actual baseline tensors solely to exercise an old live/resident
+    // object. The normal loader must reject this file, tested separately below.
+    let bytes = std::fs::read(path).unwrap();
+    let offset = u64::from_le_bytes(bytes[8..16].try_into().unwrap()) as usize;
+    let manifest: serde_json::Value = serde_json::from_slice(&bytes[offset..]).unwrap();
+    let saved = super::super::session_codec::decode_payload(
+        &manifest["payload"],
+        serde_json::from_value(manifest["tokens"].clone()).unwrap(),
+        manifest["lazy_kv"].as_bool().unwrap(),
+        true,
+        &owner.execution.streams,
+        owner.execution.stream,
+        |name| {
+            let span = &manifest["blobs"][name];
+            let start = span[0].as_u64().unwrap() as usize;
+            Ok(bytes[start..start + span[1].as_u64().unwrap() as usize].to_vec())
+        },
+    )
+    .unwrap();
+    LiveSession {
+        current: Some(saved),
+        current_images: serde_json::from_value(manifest["images"].clone()).unwrap(),
+        ..Default::default()
+    }
+}
+
+#[test]
+#[ignore = "installed model plus vision tower; serial external memory supervisor required"]
+fn issue107_vision_product() {
+    let root = std::path::PathBuf::from(std::env::var_os("DS4_QWEN38_ARTIFACTS").unwrap());
+    let image = std::fs::read(std::env::var_os("DS4_QWEN_VISION_IMAGE").unwrap()).unwrap();
+    let directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_DIR").unwrap());
+    std::fs::create_dir_all(&directory).unwrap();
+    let settings = effective_settings(
+        crate::model::ModelChoice::Qwen38FlashNext,
+        &GenerationPreferences {
+            context_tokens: 8192,
+            max_generated_tokens: 1,
+            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 marker = crate::engine::vision_data_marker(&format!(
+        "data:image/png;base64,{}",
+        base64::engine::general_purpose::STANDARD.encode(&image)
+    ));
+    let messages = [ChatTurn {
+        user: true,
+        tool: false,
+        system: false,
+        skip_previous_eos: false,
+        reasoning: None,
+        reasoning_complete: true,
+        content: format!("Describe these images{marker}{marker}"),
+    }];
+    let mut owner = QwenExecutor::open(&settings.engine).unwrap();
+    owner.policy.store_min_suffix = Some(0);
+    let metrics = Metrics::new(&directory.join("metrics"));
+    let old = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_BAD_CHECKPOINT").unwrap());
+    assert!(
+        owner
+            .load_checkpoint(&old, &mut |_| {})
+            .unwrap_err()
+            .contains("vision cache policy")
+    );
+    let mut resident = Some(legacy_live(&owner, &old));
+    owner.swap_resident_state(&mut resident).unwrap();
+    assert!(owner.live.current.is_none());
+    drop(resident);
+    owner.live = legacy_live(&owner, &old);
+    assert!(
+        owner
+            .save_checkpoint(&directory.join("stale-save.bin"), [0; 32], &mut |_| {})
+            .unwrap_err()
+            .contains("vision cache policy")
+    );
+    let (out, complete) = owner
+        .generate(
+            &messages,
+            &settings.turn,
+            &AtomicBool::new(false),
+            &metrics,
+            |p, _, _| {
+                eprintln!("#107 prefill {p}");
+            },
+            |_, _| {},
+        )
+        .unwrap();
+    assert!(complete);
+    assert!(out.prompt_tokens > 2048);
+    assert_eq!(out.cached_tokens, 0);
+    let saved = owner.live.current.as_ref().unwrap();
+    let logits = owner.execution.values_f32(&saved.logits).unwrap();
+    let oracle = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE107_ORACLE_DIR").unwrap());
+    assert_eq!(
+        logits
+            .iter()
+            .flat_map(|v| v.to_le_bytes())
+            .collect::<Vec<_>>(),
+        std::fs::read(oracle.join("final.f32")).unwrap()
+    );
+    let oracle_ids: Vec<i32> =
+        serde_json::from_slice(&std::fs::read(oracle.join("input.json")).unwrap()).unwrap();
+    assert_eq!(&saved.token_ids[..out.prompt_tokens as usize], oracle_ids);
+    std::fs::write(
+        directory.join("long-final.f32"),
+        logits
+            .iter()
+            .flat_map(|v| v.to_le_bytes())
+            .collect::<Vec<_>>(),
+    )
+    .unwrap();
+    std::fs::write(directory.join("long-input.json"),serde_json::to_vec(&serde_json::json!({"ids":saved.token_ids,"images":owner.live.current_images,"prompt_tokens":out.prompt_tokens})).unwrap()).unwrap();
+    owner
+        .save_checkpoint(&directory.join("long.bin"), [107; 32], &mut |_| {})
+        .unwrap();
+    println!(
+        "{}",
+        serde_json::json!({"case":"vision_long","prompt_tokens":out.prompt_tokens,"cached_tokens":out.cached_tokens,
+        "tokens":owner.live.current.as_ref().unwrap().token_ids,"text":out.message.content,"reasoning":out.message.reasoning,
+        "logits_sha256":crate::engine::metal::qwen_mtplx_tests::hex_digest(&logits.iter().flat_map(|v|v.to_le_bytes()).collect::<Vec<_>>())})
+    );
+    let original = state(&owner);
+    assert!(
+        owner
+            .load_checkpoint(&old, &mut |_| {})
+            .unwrap_err()
+            .contains("vision cache policy")
+    );
+    assert_eq!(
+        state(&owner),
+        original,
+        "rejected disk cache preserves healthy current state"
+    );
+    let run = |owner: &mut QwenExecutor, messages: &[ChatTurn], label: &str| {
+        let (out, complete) = owner
+            .generate(
+                messages,
+                &settings.turn,
+                &AtomicBool::new(false),
+                &metrics,
+                |_, _, _| {},
+                |_, _| {},
+            )
+            .unwrap();
+        assert!(complete);
+        assert!(
+            !owner
+                .live
+                .current_images
+                .iter()
+                .any(|i| i.start < out.cached_tokens as usize
+                    && (out.cached_tokens as usize) < i.end)
+        );
+        println!(
+            "{}",
+            serde_json::json!({"case":label,"cached_tokens":out.cached_tokens,"prompt_tokens":out.prompt_tokens,"state_sha256":state(owner)})
+        );
+        out
+    };
+    let repeated = run(&mut owner, &messages, "unchanged");
+    assert_eq!(repeated.cached_tokens, out.prompt_tokens);
+    assert_eq!(state(&owner), original);
+    let mut followup = messages.to_vec();
+    followup.push(out.message.clone());
+    followup.push(ChatTurn {
+        content: "Continue.".into(),
+        ..messages[0].clone()
+    });
+    let warm = run(&mut owner, &followup, "resident_followup");
+    assert!(warm.cached_tokens >= out.prompt_tokens);
+    let warm_state = state(&owner);
+    owner.reset().unwrap();
+    assert!(
+        owner
+            .load_checkpoint(&directory.join("long.bin"), &mut |_| {})
+            .unwrap()
+    );
+    assert_eq!(state(&owner), original);
+    let reopened = run(&mut owner, &followup, "reopened_followup");
+    assert_eq!(reopened.cached_tokens, warm.cached_tokens);
+    assert_eq!(state(&owner), warm_state);
+    let mut other = std::io::Cursor::new(Vec::new());
+    image::RgbImage::from_pixel(1024, 1024, image::Rgb([0, 0, 255]))
+        .write_to(&mut other, image::ImageFormat::Png)
+        .unwrap();
+    let changed_marker = crate::engine::vision_data_marker(&format!(
+        "data:image/png;base64,{}",
+        base64::engine::general_purpose::STANDARD.encode(other.get_ref())
+    ));
+    for (label, content) in [
+        (
+            "text_before",
+            format!("Please describe these images{marker}{marker}"),
+        ),
+        (
+            "text_between",
+            format!("Describe these images{marker}Compare them.{marker}"),
+        ),
+        (
+            "text_after",
+            format!("Describe these images{marker}{marker}Be brief."),
+        ),
+        (
+            "changed_pixels",
+            format!("Describe these images{changed_marker}{changed_marker}"),
+        ),
+    ] {
+        owner.reset().unwrap();
+        owner
+            .load_checkpoint(&directory.join("long.bin"), &mut |_| {})
+            .unwrap();
+        let edited = [ChatTurn {
+            content,
+            ..messages[0].clone()
+        }];
+        let changed = run(&mut owner, &edited, label);
+        assert_eq!(changed.cached_tokens, 0);
+        if label == "changed_pixels" {
+            assert_eq!(
+                &owner.live.current.as_ref().unwrap().token_ids[..changed.prompt_tokens as usize],
+                oracle_ids
+            );
+            assert_ne!(state(&owner), original);
+        }
+        let expected = state(&owner);
+        owner.reset().unwrap();
+        run(&mut owner, &edited, &format!("{label}_cold_control"));
+        assert_eq!(state(&owner), expected);
+    }
+    owner.reset().unwrap();
+    let cancelled = AtomicBool::new(false);
+    let interrupted = owner.generate(
+        &messages,
+        &settings.turn,
+        &cancelled,
+        &metrics,
+        |p, _, _| {
+            if p >= 2048 {
+                cancelled.store(true, std::sync::atomic::Ordering::Relaxed);
+            }
+        },
+        |_, _| {},
+    );
+    assert!(cancelled.load(std::sync::atomic::Ordering::Relaxed));
+    assert!(interrupted.is_err() || interrupted.is_ok_and(|(_, complete)| !complete));
+    assert!(owner.live.current.is_none());
+    let resent = run(&mut owner, &messages, "cancel_resend");
+    assert_eq!(resent.cached_tokens, 0);
+    assert_eq!(state(&owner), original);
+    // An artificially supplied prefix ending within the second image must be
+    // refused before any of its tensors are restored.
+    owner
+        .live
+        .current
+        .as_mut()
+        .unwrap()
+        .token_ids
+        .truncate(2048);
+    owner.live.prefix = None;
+    let guarded = run(&mut owner, &messages, "inside_image_guard");
+    assert_eq!(guarded.cached_tokens, 0);
+    assert_eq!(state(&owner), original);
+    let text = [ChatTurn {
+        content: "Say hello.".into(),
+        ..messages[0].clone()
+    }];
+    run(&mut owner, &text, "text_control");
+    let text_state = state(&owner);
+    assert!(owner.live.current_images.is_empty());
+    owner
+        .save_checkpoint(&directory.join("text.bin"), [0; 32], &mut |_| {})
+        .unwrap();
+    owner.reset().unwrap();
+    assert!(
+        owner
+            .load_checkpoint(&directory.join("text.bin"), &mut |_| {})
+            .unwrap()
+    );
+    assert_eq!(state(&owner), text_state);
+    println!(
+        "{}",
+        serde_json::json!({"case":"compatibility","old_disk_rejected":true,"old_live_rejected":true,"old_resident_rejected":true,"stale_save_rejected":true,"new_vision_and_text_reopen_exactly":true})
+    );
+}

src/engine/metal/qwen_mtplx/owner.rs

@@ -253,7 +253,8 @@ impl QwenExecutor {
         incoming: &mut Option<LiveSession>,
     ) -> Result<(), String> {
         self.clear_execution_cache()?;
-        let next = incoming.take().unwrap_or_default();
+        let mut next = incoming.take().unwrap_or_default();
+        next.discard_incompatible_vision();
         *incoming = Some(std::mem::replace(&mut self.live, next));
         Ok(())
     }
@@ -1068,3 +1069,6 @@ mod issue105_product;
 #[cfg(test)]
 #[path = "issue106_product.rs"]
 mod issue106_product;
+#[cfg(test)]
+#[path = "issue107_product.rs"]
+mod issue107_product;

src/engine/metal/qwen_mtplx/qsa_attention.rs

@@ -106,6 +106,7 @@ impl Attention<'_> {
                                 idx.as_ref(),
                                 &options.indexer,
                                 streams,
+                                vision,
                                 stream,
                             )?
                             .0
@@ -119,7 +120,16 @@ impl Attention<'_> {
                 let selected = match &mut self.indexer {
                     Some(indexer) => {
                         indexer
-                            .call(x, pos, cache, None, &options.indexer, streams, stream)?
+                            .call(
+                                x,
+                                pos,
+                                cache,
+                                None,
+                                &options.indexer,
+                                streams,
+                                vision,
+                                stream,
+                            )?
                             .0
                     }
                     None => None,
@@ -192,8 +202,8 @@ impl Attention<'_> {
         let v = ops::transpose(&v, &[0, 2, 1, 3], stream)?;
         let (k, v) = cache.kv.update_and_fetch(&k, &v, stream)?;
         let total = k.borrow().layout().shape()[2];
-        // Vision bypasses selection, not the indexer/cache work above.
-        if vision.is_some() {
+        // Match the oracle's explicit diagnostic rollback; normal vision keeps QSA.
+        if vision.is_some() && crate::engine::qwen_vision::cache_policy() == 0 {
             selected = None;
         }
         if let Some(EagerSelection::FlashBlocks(ids, tail)) = selected {
@@ -492,6 +502,7 @@ fn mtplx_qsa_attention_module_matches_reference() {
                     q_norm: &iqnorm,
                     k_norm: &iknorm,
                     inv_freq: &freq,
+                    mrope_axes: (vision == "table").then_some(&axes),
                     eps: 1e-6,
                     scaling: 1.,
                     scratch: 32 * 1024 * 1024,
@@ -512,7 +523,7 @@ fn mtplx_qsa_attention_module_matches_reference() {
                 attn.indexer
                     .as_ref()
                     .unwrap()
-                    .extend_pooled(&mut cache, prefix, options.indexer.fused, gpu)
+                    .extend_pooled(&mut cache, prefix, options.indexer.fused, None, gpu)
                     .unwrap();
                 let kv =
                     ops::transpose(&source(&[1, prefix, 2, 256], 89), &[0, 2, 1, 3], gpu).unwrap();

src/engine/metal/qwen_mtplx/qsa_indexer.rs

@@ -18,6 +18,7 @@ pub(super) struct Indexer<'a> {
     pub(super) q_norm: &'a Array,
     pub(super) k_norm: &'a Array,
     pub(super) inv_freq: &'a Array,
+    pub(super) mrope_axes: Option<&'a Array>,
     pub(super) eps: f64,
     pub(super) scaling: f64,
     pub(super) scratch: u64,
@@ -31,6 +32,10 @@ pub(super) struct PreparedRows {
     pub(super) total: i32,
 }
 
+#[cfg(test)]
+#[path = "issue107_indexer.rs"]
+mod issue107;
+
 impl<'a> Indexer<'a> {
     #[expect(
         clippy::too_many_arguments,
@@ -195,6 +200,7 @@ impl<'a> Indexer<'a> {
         qk_rows: Option<&Array>,
         options: &super::QsaIndexerOptions,
         streams: &super::stream::Streams,
+        vision: Option<(Option<&Array>, i32)>,
         stream: Stream,
     ) -> Result<(Option<super::attention::EagerSelection>, &'static str), String> {
         use super::{
@@ -221,16 +227,21 @@ impl<'a> Indexer<'a> {
             self.ratio as u32,
             self.block_topk as u32,
         );
-        if self.compiled_supported(
-            qk_rows.unwrap_or(hidden),
-            cache,
-            pos,
-            qk_rows.is_some(),
-            decode,
-            mode,
-            options,
-            stream,
-        ) {
+        let vision = vision.filter(|_| {
+            self.mrope_axes.is_some() && crate::engine::qwen_vision::cache_policy() != 0
+        });
+        if vision.is_none()
+            && self.compiled_supported(
+                qk_rows.unwrap_or(hidden),
+                cache,
+                pos,
+                qk_rows.is_some(),
+                decode,
+                mode,
+                options,
+                stream,
+            )
+        {
             return self
                 .call_compiled(
                     hidden,
@@ -244,7 +255,7 @@ impl<'a> Indexer<'a> {
                 .map(|a| (a, "compiled"));
         }
         let Some(prepared) =
-            self.prepare_rows(hidden, pos, cache, qk_rows, options.fused, stream)?
+            self.prepare_rows(hidden, pos, cache, qk_rows, options.fused, vision, stream)?
         else {
             return Ok((None, "update_only"));
         };
@@ -573,9 +584,13 @@ impl<'a> Indexer<'a> {
         q: &Array,
         pos: i32,
         fused: bool,
+        vision: Option<(Option<&Array>, i32)>,
         stream: Stream,
     ) -> Result<Array, String> {
-        if fused && qsa_prepare::supported(q, self.q_norm, self.inv_freq, 4, stream) {
+        if vision.is_none()
+            && fused
+            && qsa_prepare::supported(q, self.q_norm, self.inv_freq, 4, stream)
+        {
             return qsa_prepare::queries(
                 q,
                 self.q_norm,
@@ -588,7 +603,18 @@ impl<'a> Indexer<'a> {
         }
         let q = ops::rms_norm(q, Some(self.q_norm), self.eps as f32, stream)?;
         let positions = indexing::arange_range_i32(pos, pos + q.layout().shape()[1], stream)?;
-        let (cos, sin) = rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?;
+        let (cos, sin) = match (vision, self.mrope_axes) {
+            (Some((table, delta)), Some(axes)) => rope::vision_position_cos_sin(
+                &positions,
+                table,
+                delta,
+                self.inv_freq,
+                axes,
+                self.scaling,
+                stream,
+            )?,
+            _ => rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?,
+        };
         rope::apply_partial(&q, &cos, &sin, stream)
     }
 
@@ -598,9 +624,13 @@ impl<'a> Indexer<'a> {
         old: i32,
         total: i32,
         fused: bool,
+        vision: Option<(Option<&Array>, i32)>,
         stream: Stream,
     ) -> Result<Array, String> {
-        if fused && qsa_prepare::supported(fresh, self.k_norm, self.inv_freq, 3, stream) {
+        if vision.is_none()
+            && fused
+            && qsa_prepare::supported(fresh, self.k_norm, self.inv_freq, 3, stream)
+        {
             return qsa_prepare::pool(
                 fresh,
                 self.k_norm,
@@ -632,7 +662,18 @@ impl<'a> Indexer<'a> {
             Binary::Multiply,
             stream,
         )?;
-        let (cos, sin) = rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?;
+        let (cos, sin) = match (vision, self.mrope_axes) {
+            (Some((table, delta)), Some(axes)) => rope::vision_position_cos_sin(
+                &positions,
+                table,
+                delta,
+                self.inv_freq,
+                axes,
+                self.scaling,
+                stream,
+            )?,
+            _ => rope::cos_sin(&positions, self.inv_freq, self.scaling, stream)?,
+        };
         ops::squeeze(
             &rope::apply_partial(
                 &ops::expand_dims(&pooled, &[2], stream)?,
@@ -650,6 +691,7 @@ impl<'a> Indexer<'a> {
         cache: &mut QsaCache,
         total: i32,
         fused: bool,
+        vision: Option<(Option<&Array>, i32)>,
         stream: Stream,
     ) -> Result<Option<Array>, String> {
         let nb = total / self.ratio;
@@ -669,7 +711,7 @@ impl<'a> Indexer<'a> {
             )?;
             drop(raw);
             cache.write_pooled(
-                &self.pool_keys(&fresh, old, nb, fused, stream)?,
+                &self.pool_keys(&fresh, old, nb, fused, vision, stream)?,
                 old,
                 nb,
                 stream,
@@ -689,6 +731,10 @@ impl<'a> Indexer<'a> {
     /// Common noncompiled _call_rows prefix. None means the complete-block
     /// budget fits: cache updates still happen, but lazy query work is unused.
     /// Cache.kv.offset is advanced by Attention, never by this indexer.
+    #[expect(
+        clippy::too_many_arguments,
+        reason = "original preparation inputs and request positions"
+    )]
     pub(super) fn prepare_rows(
         &self,
         hidden: &Array,
@@ -696,6 +742,7 @@ impl<'a> Indexer<'a> {
         cache: &mut QsaCache,
         qk_rows: Option<&Array>,
         fused: bool,
+        vision: Option<(Option<&Array>, i32)>,
         stream: Stream,
     ) -> Result<Option<PreparedRows>, String> {
         let shape = hidden.layout().shape().to_vec();
@@ -733,9 +780,9 @@ impl<'a> Indexer<'a> {
         let parts = ops::split(&qk, &[width], -1, stream)?;
         let query = ops::reshape(&parts[0], &[1, shape[1], self.heads, self.head_dim], stream)?;
         let keys = ops::reshape(&parts[1], &[1, shape[1], self.head_dim], stream)?;
-        let query = self.prepare_queries(&query, pos, fused, stream)?;
+        let query = self.prepare_queries(&query, pos, fused, vision, stream)?;
         cache.write_raw(&keys, stream)?;
-        let pooled = self.extend_pooled(cache, total, fused, stream)?;
+        let pooled = self.extend_pooled(cache, total, fused, vision, stream)?;
         if total / self.ratio <= self.block_topk {
             return Ok(None);
         }
@@ -793,6 +840,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() {
         q_norm: &q_norm,
         k_norm: &k_norm,
         inv_freq: &inv_freq,
+        mrope_axes: None,
         eps: 1e-6,
         scaling: 1.,
         scratch: 32 * 1024 * 1024,
@@ -835,7 +883,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() {
                         .write_raw(&source(&[1, prefix, 128], 88), gpu)
                         .unwrap();
                     indexer
-                        .extend_pooled(&mut cache, prefix, options.fused, gpu)
+                        .extend_pooled(&mut cache, prefix, options.fused, None, gpu)
                         .unwrap();
                     let kv = ops::transpose(&source(&[1, prefix, 2, 256], 89), &[0, 2, 1, 3], gpu)
                         .unwrap();
@@ -875,6 +923,7 @@ fn mtplx_qsa_canonical_indexer_routes_match_reference() {
                         qk.as_ref(),
                         &options,
                         &streams,
+                        None,
                         gpu,
                     )
                     .unwrap()
@@ -1111,6 +1160,7 @@ fn mtplx_qsa_indexer_compiled_commit_matches_reference() {
                 q_norm: &q_norm[0],
                 k_norm: &k_norm[0],
                 inv_freq: &freq[0],
+                mrope_axes: None,
                 eps: 1e-6,
                 scaling: 1.,
                 scratch: 32 * 1024 * 1024,
@@ -1399,6 +1449,7 @@ fn mtplx_qsa_indexer_prefix_matches_reference() {
                     q_norm: &q_norm,
                     k_norm: &k_norm,
                     inv_freq: &inv_freq,
+                    mrope_axes: None,
                     eps: 1e-6,
                     scaling: 1.125,
                     scratch: 32 * 1024 * 1024,
@@ -1454,7 +1505,15 @@ fn mtplx_qsa_indexer_prefix_matches_reference() {
                         }
                         let (prefix, calls) = capture_dispatches(|| {
                             indexer
-                                .prepare_rows(&hidden, pos, &mut cache, qk.as_ref(), fused, gpu)
+                                .prepare_rows(
+                                    &hidden,
+                                    pos,
+                                    &mut cache,
+                                    qk.as_ref(),
+                                    fused,
+                                    None,
+                                    gpu,
+                                )
                                 .unwrap()
                         });
                         assert!(calls.is_empty(), "QSA prefix must remain lazy");

src/engine/metal/qwen_mtplx/qsa_select.rs

@@ -391,6 +391,7 @@ fn mtplx_qsa_fused_graph_matches_reference() {
                                 q_norm: &q,
                                 k_norm: &q,
                                 inv_freq: &q,
+                                mrope_axes: None,
                                 eps: 1e-6,
                                 scaling: 1.,
                                 scratch: 32 * 1024 * 1024,

src/engine/metal/qwen_mtplx/request.rs

@@ -100,6 +100,14 @@ pub(in crate::engine) struct LiveSession {
 }
 
 impl LiveSession {
+    pub(super) fn discard_incompatible_vision(&mut self) {
+        if !crate::engine::qwen_vision::compatible(&self.current_images)
+            || !crate::engine::qwen_vision::compatible(&self.prefix_images)
+        {
+            *self = Self::default();
+        }
+    }
+
     pub(super) fn prompt(
         &self,
         model: &QwenMetadata,
@@ -131,6 +139,7 @@ impl LiveSession {
         emit: impl FnMut(bool, String),
         observe: impl FnMut(TurnProgress) -> Result<(), String>,
     ) -> Result<TextTurn, String> {
+        self.discard_incompatible_vision();
         let (prompt, images) = execution
             .vision_prompt
             .clone()

src/engine/metal/qwen_mtplx/rope.rs

@@ -246,6 +246,53 @@ pub(super) fn mrope_cos_sin(
     ))
 }
 
+/// MTPLX _vision_position_cos_sin: choose table versus text+delta per row.
+pub(super) fn vision_position_cos_sin(
+    positions: &Array,
+    table: Option<&Array>,
+    delta: i32,
+    inv_freq: &Array,
+    axes: &Array,
+    scaling: f64,
+    stream: Stream,
+) -> Result<(Array, Array), String> {
+    let integer = |n: i32| Array::new(&[], Dtype::I32, super::scalar_buffer(&n.to_le_bytes())?);
+    let shifted = binary(positions, &integer(delta)?, Binary::Add, stream)?;
+    let mut positions3 = ops::broadcast_to(
+        &ops::expand_dims(&shifted, &[0], stream)?,
+        &[3, positions.layout().size() as i32],
+        stream,
+    )?;
+    if let Some(table) = table {
+        let shape = table.layout().shape().to_vec();
+        if shape.len() != 2 || shape[0] != 3 {
+            return Err("vision positions require a three-axis table".into());
+        }
+        if shape[1] > 0 {
+            let clipped = binary(
+                &binary(positions, &integer(0)?, Binary::Maximum, stream)?,
+                &integer(shape[1] - 1)?,
+                Binary::Minimum,
+                stream,
+            )?;
+            let prompt = indexing::take(table, &clipped, 1, stream)?;
+            let covered = binary(positions, &integer(shape[1])?, Binary::Less, stream)?;
+            positions3 = indexing::select(
+                &ops::expand_dims(&covered, &[0], stream)?,
+                &prompt,
+                &positions3,
+                stream,
+            )?;
+        }
+    }
+    let (mut cos, mut sin) = mrope_cos_sin(&positions3, inv_freq, axes, stream)?;
+    if scaling != 1.0 {
+        cos = binary(&cos, &scalar(scaling)?, Binary::Multiply, stream)?;
+        sin = binary(&sin, &scalar(scaling)?, Binary::Multiply, stream)?;
+    }
+    Ok((cos, sin))
+}
+
 pub(super) fn apply_partial(
     x: &Array,
     cos: &Array,

src/engine/metal/qwen_mtplx/session_checkpoint.rs

@@ -78,6 +78,9 @@ impl LiveSession {
             .as_ref()
             .ok_or("Qwen has no committed checkpoint")?;
         validate_tokens(&saved.token_ids, identity)?;
+        if !crate::engine::qwen_vision::compatible(&self.current_images) {
+            return Err("Qwen vision cache policy differs; rebuild the cache".into());
+        }
         if saved.mtp.is_some() != identity.mtp {
             return Err("Qwen checkpoint MTP configuration differs".into());
         }
@@ -179,6 +182,9 @@ impl LiveSession {
             );
         }
         validate_tokens(&manifest.tokens, identity)?;
+        if !crate::engine::qwen_vision::compatible(&manifest.images) {
+            return Err("Qwen vision cache policy differs; rebuild the cache".into());
+        }
         if !manifest.payload["mtp_history_snapshot"].is_null() != identity.mtp {
             return Err("Qwen checkpoint MTP payload differs".into());
         }

src/engine/metal/qwen_mtplx/vision.rs

@@ -510,6 +510,7 @@ pub(super) fn prepare(
                 start,
                 end: ids.len(),
                 digest: *digest,
+                policy: crate::engine::qwen_vision::cache_policy(),
             });
             index += 1;
         } else {

src/engine/qwen_vision.rs

@@ -207,12 +207,27 @@ pub(super) struct Identity {
     pub(super) start: usize,
     pub(super) end: usize,
     pub(super) digest: [u8; 32],
+    /// Missing on checkpoints written before image-aware sparse QSA.
+    #[serde(default)]
+    pub(super) policy: u32,
+}
+
+pub(super) fn cache_policy() -> u32 {
+    u32::from(std::env::var("MTPLX_QWEN4_VISION_QSA").as_deref() != Ok("0"))
+}
+
+pub(super) fn compatible(images: &[Identity]) -> bool {
+    images.iter().all(|image| image.policy == cache_policy())
 }
 
 pub(super) fn same_prefix(a: &[Identity], b: &[Identity], tokens: usize) -> bool {
-    a.iter()
-        .filter(|i| i.start < tokens)
-        .eq(b.iter().filter(|i| i.start < tokens))
+    // Image pads do not identify partial image embeddings. MTPLX only restores
+    // before an image or after its complete span, never into the middle.
+    !b.iter()
+        .any(|image| image.start < tokens && tokens < image.end)
+        && a.iter()
+            .filter(|i| i.start < tokens)
+            .eq(b.iter().filter(|i| i.start < tokens))
 }
 
 #[test]
@@ -271,14 +286,19 @@ fn qwen_image_cache_keys_distinguish_pixels_and_preserve_prior_text() {
         start: 10,
         end: 14,
         digest: [1; 32],
+        policy: cache_policy(),
     }];
     let changed = vec![Identity {
         start: 10,
         end: 14,
         digest: [2; 32],
+        policy: cache_policy(),
     }];
     assert!(same_prefix(&old, &old, 20));
     assert!(!same_prefix(&old, &changed, 20));
     assert!(same_prefix(&old, &changed, 10));
     assert!(!same_prefix(&old, &[], 20));
+    assert!(!same_prefix(&old, &old, 11));
+    assert!(!same_prefix(&old, &old, 13));
+    assert!(same_prefix(&old, &old, 14));
 }