Apply Qwen vision RoPE per row across table boundaries

src/engine/metal/qwen_mtplx.rs

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

src/engine/metal/qwen_mtplx/issue103_analysis.rs

@@ -430,7 +430,7 @@ fn issue103_vision_position_receipts() {
                 gpu,
             )
             .unwrap();
-        ops::evaluate(&streams, &[output], gpu, false).unwrap();
+        ops::evaluate(&streams, std::slice::from_ref(&output), gpu, false).unwrap();
         let projected = ops::reshape(
             &ops::rms_norm(&x, Some(&norm), 1e-6, gpu).unwrap(),
             &[1, rows, 1, 8],
@@ -490,11 +490,81 @@ fn issue103_vision_position_receipts() {
             "{}",
             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());
+        assert!(bad.is_empty(), "each row must use its own vision position");
+        // An equivalent split preserves outputs, offsets and the next call.
+        let mut split = QsaCache::new(4);
+        if pos > 0 {
+            let prefix = bf(&[1, 1, pos, 8], &vec![0.; pos as usize * 8]);
+            split.kv.update_and_fetch(&prefix, &prefix, gpu).unwrap();
         }
+        let boundary = (4 - pos).clamp(0, rows);
+        let mut pieces = Vec::new();
+        for (begin, end) in [(0, boundary), (boundary, rows)] {
+            if begin == end {
+                continue;
+            }
+            let piece = ops::slice(&x, &[0, begin, 0], &[1, end, 8], &[1; 3], gpu).unwrap();
+            pieces.push(
+                attention
+                    .call(
+                        &piece,
+                        &mut split,
+                        &options,
+                        Some((Some(&table), -1)),
+                        &streams,
+                        gpu,
+                    )
+                    .unwrap(),
+            );
+        }
+        let joined = ops::concatenate(&pieces, 1, gpu).unwrap();
+        assert_eq!(read(&output), read(&joined));
+        assert_eq!(cache.kv.offset, pos + rows);
+        assert_eq!(split.kv.offset, cache.kv.offset);
+        let next = bf(&[1, 1, 8], &[0.3, -0.7, 0.2, 0.9, -0.2, 0.4, 0.6, -0.1]);
+        let next_whole = attention
+            .call(
+                &next,
+                &mut cache,
+                &options,
+                Some((Some(&table), -1)),
+                &streams,
+                gpu,
+            )
+            .unwrap();
+        let next_split = attention
+            .call(
+                &next,
+                &mut split,
+                &options,
+                Some((Some(&table), -1)),
+                &streams,
+                gpu,
+            )
+            .unwrap();
+        assert_eq!(read(&next_whole), read(&next_split));
+        assert_eq!(cache.kv.offset, pos + rows + 1);
+        assert_eq!(split.kv.offset, cache.kv.offset);
+        for (actual, expected) in [
+            (&cache.kv.keys, &split.kv.keys),
+            (&cache.kv.values, &split.kv.values),
+        ] {
+            let valid = |array: &Option<super::cache::ArrayObject>| {
+                ops::slice(
+                    &array.as_ref().unwrap().borrow(),
+                    &[0, 0, 0, 0],
+                    &[1, 1, pos + rows + 1, 8],
+                    &[1; 4],
+                    gpu,
+                )
+                .unwrap()
+            };
+            assert_eq!(read(&valid(actual)), read(&valid(expected)));
+        }
+        println!(
+            "{}",
+            serde_json::json!({"case":"vision_chunk_next_state","start":pos,"rows":rows,"offset":cache.kv.offset,"exact":true})
+        );
     }
     let prefix = bf(
         &[1, 12, 8],

src/engine/metal/qwen_mtplx/issue108.rs

@@ -0,0 +1,90 @@
+use super::array::{Array, Dtype};
+use super::{ops, rope};
+
+#[test]
+#[ignore = "Metal and pinned MTPLX chunk fixtures; external supervisor required"]
+fn issue108_attention_chunk_positions() {
+    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 directory = std::path::PathBuf::from(std::env::var_os("DS4_ISSUE108_ORACLE").unwrap());
+    let cases: Vec<serde_json::Value> =
+        serde_json::from_slice(&std::fs::read(directory.join("manifest.json")).unwrap()).unwrap();
+    let ints = |shape: &[i32], values: &[i32]| {
+        Array::new(
+            shape,
+            Dtype::I32,
+            super::scalar_buffer(
+                &values
+                    .iter()
+                    .flat_map(|n| n.to_le_bytes())
+                    .collect::<Vec<_>>(),
+            )
+            .unwrap(),
+        )
+        .unwrap()
+    };
+    for case in &cases {
+        let n = |name: &str| case[name].as_i64().unwrap() as i32;
+        let key = case["key"].as_str().unwrap();
+        let sections: Vec<i32> = serde_json::from_value(case["sections"].clone()).unwrap();
+        let axes = ints(&[n("rotary") / 2], &rope::build_mrope_axes(&sections, true));
+        let table: Option<Vec<Vec<i32>>> = serde_json::from_value(case["table"].clone()).unwrap();
+        let table = table.map(|t| ints(&[3, t[0].len() as i32], &t.concat()));
+        let (freq, _) = rope::inv_freq_and_scaling(n("rotary"), 10000000., None, gpu).unwrap();
+        let scaling = case["scaling"].as_f64().unwrap();
+        let (cos, sin) = rope::vision_chunk_cos_sin(
+            (table.as_ref(), n("delta")),
+            n("start"),
+            n("rows"),
+            &freq,
+            &axes,
+            scaling,
+            gpu,
+        )
+        .unwrap();
+        let check = |part: &str, value: &Array| {
+            let actual = ops::read_f32_as_f64(
+                &ops::astype(value, Dtype::F32, false, gpu).unwrap(),
+                &streams,
+                gpu,
+            )
+            .unwrap()
+            .into_iter()
+            .flat_map(|v| (v as f32).to_le_bytes())
+            .collect::<Vec<_>>();
+            assert_eq!(
+                actual,
+                std::fs::read(directory.join(format!("{key}-{part}.f32"))).unwrap(),
+                "{key} {part}"
+            );
+        };
+        check("cos", &cos);
+        check("sin", &sin);
+        for (part, heads, seed) in [("q", n("heads"), 3), ("k", 1, 11)] {
+            let shape = [1, n("rows"), heads, n("dim")];
+            let values = (0..shape.iter().product::<i32>())
+                .map(|i| ((i * 17 + seed) % 97 - 48) as f32 / 32.)
+                .collect::<Vec<_>>();
+            let dtype = if case["dtype"] == "bf16" {
+                Dtype::BF16
+            } else {
+                Dtype::F32
+            };
+            let raw = ops::astype(
+                &super::vision::f32_array(&shape, &values).unwrap(),
+                dtype,
+                false,
+                gpu,
+            )
+            .unwrap();
+            check(part, &rope::apply_partial(&raw, &cos, &sin, gpu).unwrap());
+        }
+        println!(
+            "{}",
+            serde_json::json!({"case":"attention_chunk","key":key,"exact_q_k_and_tables":true})
+        );
+    }
+    assert_eq!(cases.len(), 56);
+}

src/engine/metal/qwen_mtplx/qsa_attention.rs

@@ -171,22 +171,15 @@ impl Attention<'_> {
             rope::cos_sin(&positions, self.inv_freq, self.rope_scaling, stream)
         };
         let (cos, sin) = match (vision, self.mrope_axes) {
-            (Some((table, delta)), Some(axes)) => {
-                if let Some(table) =
-                    table.filter(|t| t.layout().shape().get(1).is_some_and(|&n| end <= n))
-                {
-                    let table = ops::slice(
-                        table,
-                        &[0, pos],
-                        &[table.layout().shape()[0], end],
-                        &[1, 1],
-                        stream,
-                    )?;
-                    rope::mrope_cos_sin(&table, self.inv_freq, axes, stream)?
-                } else {
-                    plain_rope(delta)?
-                }
-            }
+            (Some((table, delta)), Some(axes)) => rope::vision_chunk_cos_sin(
+                (table, delta),
+                pos,
+                rows,
+                self.inv_freq,
+                axes,
+                self.rope_scaling,
+                stream,
+            )?,
             _ => plain_rope(0)?,
         };
         let q = ops::transpose(

src/engine/metal/qwen_mtplx/rope.rs

@@ -246,6 +246,70 @@ pub(super) fn mrope_cos_sin(
     ))
 }
 
+/// MTPLX _vision_chunk_cos_sin: preserve the inside/outside graphs and
+/// concatenate their tables when one attention chunk contains both kinds of row.
+pub(super) fn vision_chunk_cos_sin(
+    (table, delta): (Option<&Array>, i32),
+    start: i32,
+    rows: i32,
+    inv_freq: &Array,
+    axes: &Array,
+    scaling: f64,
+    stream: Stream,
+) -> Result<(Array, Array), String> {
+    let end = start
+        .checked_add(rows)
+        .ok_or("QSA vision position overflow")?;
+    let table_len = match table {
+        Some(table) => {
+            let layout = table.layout();
+            let shape = layout.shape();
+            if shape.len() != 2 || shape[0] != 3 {
+                return Err("vision positions require a three-axis table".into());
+            }
+            shape[1]
+        }
+        None => 0,
+    };
+    let inside = (end.min(table_len) - start).max(0);
+    let covered = if inside > 0 {
+        let positions = ops::slice(
+            table.unwrap(),
+            &[0, start],
+            &[3, start + inside],
+            &[1, 1],
+            stream,
+        )?;
+        let (mut cos, mut sin) = mrope_cos_sin(&positions, inv_freq, axes, stream)?;
+        if scaling != 1.0 {
+            cos = binary(&cos, &scalar(scaling)?, Binary::Multiply, stream)?;
+            sin = binary(&sin, &scalar(scaling)?, Binary::Multiply, stream)?;
+        }
+        if inside == rows {
+            return Ok((cos, sin));
+        }
+        Some((cos, sin))
+    } else {
+        None
+    };
+    let first = start
+        .checked_add(inside)
+        .and_then(|p| p.checked_add(delta))
+        .ok_or("QSA vision position overflow")?;
+    let stop = first
+        .checked_add(rows - inside)
+        .ok_or("QSA vision position overflow")?;
+    let positions = indexing::arange_range_i32(first, stop, stream)?;
+    let (cos, sin) = cos_sin(&positions, inv_freq, scaling, stream)?;
+    match covered {
+        Some((cos_in, sin_in)) => Ok((
+            ops::concatenate(&[cos_in, cos], 0, stream)?,
+            ops::concatenate(&[sin_in, sin], 0, stream)?,
+        )),
+        None => Ok((cos, sin)),
+    }
+}
+
 /// MTPLX _vision_position_cos_sin: choose table versus text+delta per row.
 pub(super) fn vision_position_cos_sin(
     positions: &Array,