feat(neuron): M-RoPE Stage 4 — wire interleaved M-RoPE into the TP path

Mirror Stage 3 into the tensor-parallel Qwen3.6 model:

- TpQwen3_5Attention / DecoderLayer take (cos, sin) instead of a scalar
  offset and apply via apply_cos_sin.
- TpQwen3_5Model gains the replicated rotary + rope_delta (reset in
  clear_kv_cache, settable). forward_inner builds the cos/sin once —
  interleaved M-RoPE from explicit position_ids (vision) or plain at
  offset+rope_delta (text/decode). forward() and forward_with_positions()
  delegate; the old single-shot forward_with_vision is gone.
- prefill_with_images_chunked now computes get_rope_index over the whole
  prompt once, stores rope_delta on the base model, and slices the
  (3, prompt_len) position tensor per chunk — so every rank assigns image
  tokens their 14×14 grid coordinates and steps in lockstep (every chunk,
  text or image, carries the M-RoPE slice because the image shifts the
  surrounding text positions).

Also build the position-id tensor as f32 directly (positions are small
integers, exact in f32) to avoid an i64→f32 cast on the GPU.

The TP forward is cuda-gated — CI CUDA type-check is the compile gate.
Non-cuda build + clippy + full workspace tests green; rope math + the
plain-RoPE-reduction invariant covered by unit tests.

Completes the interleaved-M-RoPE work for the vision spatial misread.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

crates/neuron/src/harness/arch/qwen3_5/decoder.rs

@@ -93,12 +93,13 @@ impl Qwen3_5DecoderLayer {
         &mut self,
         x: &Tensor,
         attn_mask: Option<&Tensor>,
-        offset: usize,
+        cos: &Tensor,
+        sin: &Tensor,
     ) -> candle_core::Result<Tensor> {
         let h = self.input_layernorm.forward(x)?;
         let attn_out = match &mut self.attention {
-            AttentionKind::Full(attn) => attn.forward(&h, attn_mask, offset)?,
-            // Linear attention ignores attn_mask + offset; its causal
+            AttentionKind::Full(attn) => attn.forward(&h, attn_mask, cos, sin)?,
+            // Linear attention ignores attn_mask + rope; its causal
             // structure is baked into the recurrent state lifecycle.
             AttentionKind::Linear(net) => net.forward(&h)?,
         };

crates/neuron/src/harness/arch/qwen3_5/full_attn.rs

@@ -96,7 +96,8 @@ impl Qwen3_5Attention {
         &mut self,
         x: &Tensor,
         attn_mask: Option<&Tensor>,
-        offset: usize,
+        cos: &Tensor,
+        sin: &Tensor,
     ) -> candle_core::Result<Tensor> {
         let (b, l, _) = x.dims3()?;
 
@@ -131,8 +132,9 @@ impl Qwen3_5Attention {
             .transpose(1, 2)?
             .contiguous()?;
 
-        // 3. RoPE on q, k.
-        let (q, k) = self.rotary.apply(&q, &k, offset)?;
+        // 3. RoPE on q, k (cos/sin built once per forward by the model —
+        //    interleaved M-RoPE for image tokens, plain for text).
+        let (q, k) = self.rotary.apply_cos_sin(&q, &k, cos, sin)?;
 
         // 4. KV cache.
         let (k, v) = self.kv_cache.append(&k, &v)?;

crates/neuron/src/harness/arch/qwen3_5/linear_attn.rs

@@ -737,6 +737,8 @@ mod tests {
                 rope_theta: 10000.0,
                 partial_rotary_factor: 1.0,
                 rope_type: None,
+                mrope_section: Vec::new(),
+                mrope_interleaved: false,
             },
             rms_norm_eps: 1e-6,
             tie_word_embeddings: false,

crates/neuron/src/harness/arch/qwen3_5/mod.rs

@@ -191,11 +191,12 @@ fn default_hidden_act() -> String {
 }
 
 /// Nested `rope_parameters` block from a Qwen3-Next `config.json`.
-/// `mrope_section` and `mrope_interleaved` are accepted via the
-/// `#[serde(default)]` flatten-tolerance below but ignored — we treat
-/// MRoPE as plain RoPE for text-only inference (the three position
-/// grids carry identical ids when there's no vision input, so the
-/// interleaving is a no-op).
+///
+/// For text-only inference the three MRoPE position grids carry
+/// identical ids, so the interleave is a no-op and plain RoPE applies.
+/// For vision inputs `mrope_section` + `mrope_interleaved` drive the
+/// per-axis (text/height/width) rotary used by image tokens — see
+/// `rope.rs`.
 #[derive(Debug, Clone, Deserialize)]
 pub struct RopeParameters {
     /// Base for the inverse-frequency computation. Qwen3.6: 10_000_000.
@@ -211,6 +212,16 @@ pub struct RopeParameters {
     /// implemented here.
     #[serde(default)]
     pub rope_type: Option<String>,
+    /// MRoPE per-axis section sizes `[text, height, width]` — e.g.
+    /// `[11, 11, 10]` for Qwen3.6, summing to the rotary half-dim.
+    /// Empty for models that don't declare MRoPE (→ plain RoPE).
+    #[serde(default)]
+    pub mrope_section: Vec<usize>,
+    /// Whether the three MRoPE axes are interleaved per-frequency
+    /// (Qwen3-VL / Qwen3.6 style, `true`) rather than block-concatenated
+    /// (Qwen2-VL style, `false`).
+    #[serde(default)]
+    pub mrope_interleaved: bool,
 }
 
 fn default_rope_theta() -> f64 {
@@ -303,6 +314,16 @@ pub struct Qwen3_5Model {
     embed_tokens: Embedding,
     layers: Vec<Qwen3_5DecoderLayer>,
     norm: Qwen3_5RmsNorm,
+    /// Shared with every full-attention layer; the model uses it to
+    /// build the per-forward cos/sin (interleaved M-RoPE for image
+    /// tokens, plain for text) once, which the layers then apply.
+    rotary: Arc<RotaryEmbedding>,
+    /// `offset + rope_delta` is the text-axis position during decode.
+    /// 0 for text-only; set from `get_rope_index` during a vision
+    /// prefill (image tokens compress the position space, so text after
+    /// the image resumes from a smaller counter than the sequence
+    /// index). Reset in `clear_kv_cache`.
+    rope_delta: i64,
     device: Device,
     dtype: DType,
 }
@@ -354,6 +375,8 @@ impl Qwen3_5Model {
             embed_tokens,
             layers,
             norm,
+            rotary,
+            rope_delta: 0,
             device,
             dtype,
         })
@@ -367,6 +390,9 @@ impl Qwen3_5Model {
         for l in &mut self.layers {
             l.clear_kv_cache();
         }
+        // New request → no image-compressed position offset until the
+        // next vision prefill sets one.
+        self.rope_delta = 0;
     }
 
     fn causal_mask(&self, b: usize, tgt: usize, offset: usize) -> candle_core::Result<Tensor> {
@@ -424,16 +450,16 @@ impl Qwen3_5Model {
     ) -> candle_core::Result<Tensor> {
         let (b, l) = input.dims2()?;
         let mut h = self.embed_tokens.forward(input)?;
-        // Splice image embeddings at `image_token_id` positions. The
-        // caller pre-expanded the prompt so every patch token in the
-        // image_embeds tensor has a matching position in `input`. We
-        // index_put the rows in place.
-        if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
-            // Locate image-token positions in input_ids. Operate on
-            // CPU since the input ids are tiny (max ~10k entries
-            // including the patch expansion) and the comparison is
-            // not in the per-step hot path.
+
+        // Vision path: splice image embeddings at `image_token_id`
+        // positions and build interleaved M-RoPE cos/sin so image tokens
+        // carry their 14×14 grid coordinates. Text / decode skip the
+        // device→host id copy entirely and take the plain-RoPE fast path
+        // — bit-for-bit the pre-M-RoPE behaviour when `rope_delta == 0`.
+        let (cos, sin) = if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
+            // Token ids on CPU — reused for the splice + position ids.
             let ids: Vec<u32> = input.flatten_all()?.to_vec1()?;
+
             let mut positions: Vec<u32> = Vec::with_capacity(img.dim(0)?);
             for (idx, id) in ids.iter().enumerate() {
                 if *id == tok_id {
@@ -451,22 +477,22 @@ impl Qwen3_5Model {
                 );
             }
             if !positions.is_empty() {
-                // Cast image_embeds to the LM's dtype so the splice
-                // produces a uniform tensor for the decoder stack.
+                // Cast image_embeds to the LM's dtype, then splice the
+                // contiguous `<|image_pad|>` runs in place.
                 let img = img.to_dtype(self.dtype)?;
-                // index_select would return the rows; we want to put.
-                // candle's slice_assign with explicit positions ranges
-                // doesn't exist; use scatter via index_select + an
-                // accumulator: build a `(B, L, hidden)` zero tensor,
-                // scatter the image rows in, then add to a masked
-                // version of `h`. Simpler approach: walk positions
-                // and use `slice_assign` for contiguous runs. Since
-                // image_pad runs are contiguous (template emits
-                // `<|vision_start|><|image_pad|>×N<|vision_end|>`),
-                // we group positions and assign per run.
                 h = splice_runs(&h, &img, &positions)?;
             }
-        }
+
+            let (text, height, width, delta) = rope::get_rope_index(&ids, tok_id)
+                .map_err(|e| candle_core::Error::Msg(format!("get_rope_index: {e}")))?;
+            self.rope_delta = delta;
+            let pos = rope::mrope_position_tensor(&text, &height, &width, &self.device)?;
+            self.rotary.mrope_cos_sin(&pos)?
+        } else {
+            let base = (offset as i64 + self.rope_delta).max(0) as usize;
+            self.rotary.plain_cos_sin(base, l)?
+        };
+
         // Causal mask only needed for L > 1 prefill; full-attention
         // layers consume it via broadcast_add. Linear-attention layers
         // ignore the mask.
@@ -476,7 +502,7 @@ impl Qwen3_5Model {
             Some(self.causal_mask(b, l, offset)?)
         };
         for layer in &mut self.layers {
-            h = layer.forward(&h, causal.as_ref(), offset)?;
+            h = layer.forward(&h, causal.as_ref(), &cos, &sin)?;
         }
         self.norm.forward(&h)
     }

crates/neuron/src/harness/arch/qwen3_5/rope.rs

@@ -1,19 +1,27 @@
 //! Rotary position embedding for Qwen3-Next's full-attention layers.
 //!
-//! Qwen3.6 ships with MRoPE (multimodal RoPE) machinery in the
-//! reference Python — three position grids interleaved per
-//! `mrope_section`. For text-only inference all three grids carry the
-//! same position ids and the interleave is a no-op, so this module
-//! implements the plain (non-mrope) flavour: the standard inv_freq
-//! cosine/sine tables driven by `rope_theta` and `head_dim`.
+//! Qwen3.6 declares **interleaved M-RoPE** (multimodal RoPE): the
+//! rotary half-dimension is split across three position axes —
+//! `[text, height, width]` per `mrope_section` (`[11,11,10]` for
+//! Qwen3.6) — interleaved per-frequency. For **text** every token's
+//! three axes carry the same position id, so the interleave is a no-op
+//! and this reduces exactly to plain RoPE. For **image** tokens the
+//! height/width axes carry the patch's 2D grid coordinates, which is
+//! how the model reads the 14×14 patch layout (without it, all patches
+//! share a height position and the image reads as vertical repetition).
 //!
-//! Rotation flavour: **GLM-style** rotate-half (the second half of the
-//! head dim is negated and swapped into the first). The reference
-//! Python uses `apply_rotary_pos_emb` with `rotate_half`; candle's
-//! `rope_slow` is the matching helper.
+//! Two cos/sin builders feed a shared [`RotaryEmbedding::apply`]:
+//! - [`RotaryEmbedding::plain_cos_sin`] narrows the precomputed tables
+//!   at a scalar position — the text / decode fast path.
+//! - [`RotaryEmbedding::mrope_cos_sin`] builds per-token cos/sin from a
+//!   `(3, seq)` position-id tensor, blending the three axes' frequencies
+//!   at the interleave index sets — the vision-prefill path.
+//!
+//! Rotation flavour: **GLM-style** rotate-half (candle's `rope_slow`),
+//! matching the reference Python's `apply_rotary_pos_emb` + `rotate_half`.
 
 use anyhow::Result;
-use candle_core::{DType, Device, Tensor};
+use candle_core::{DType, Device, IndexOp, Tensor};
 
 use super::TextConfig;
 
@@ -21,6 +29,18 @@ use super::TextConfig;
 pub struct RotaryEmbedding {
     sin: Tensor,
     cos: Tensor,
+    /// Inverse frequencies, shape `(1, rotary_dim/2)`. Retained (beyond
+    /// the precomputed `sin`/`cos` tables) so [`Self::mrope_cos_sin`] can
+    /// build cos/sin from arbitrary per-axis position ids.
+    inv_freq: Tensor,
+    /// Per-axis column masks over the rotary half-dim, shape `(1, half)`,
+    /// f32 0/1. `mask_t + mask_h + mask_w` partitions the columns; a
+    /// column belongs to exactly one axis. For a non-MRoPE config
+    /// `mask_t` is all-ones and the others all-zero (→ plain RoPE).
+    mask_t: Tensor,
+    mask_h: Tensor,
+    mask_w: Tensor,
+    dtype: DType,
     /// Number of dims at the head's leading edge that the rotation
     /// covers. The remaining `head_dim - rotary_dim` dims pass through
     /// unchanged. Qwen3-Next uses `partial_rotary_factor = 0.25`, so
@@ -29,6 +49,52 @@ pub struct RotaryEmbedding {
     head_dim: usize,
 }
 
+/// Build the per-axis 0/1 column masks over the rotary half-dim from
+/// `mrope_section`. Returns `(temporal, height, width)` each length
+/// `half`. Temporal is the complement of height ∪ width, so the three
+/// masks always partition `0..half` and reduce to all-temporal (plain
+/// RoPE) when no usable section is given.
+fn mrope_masks(
+    half: usize,
+    section: &[usize],
+    interleaved: bool,
+) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
+    let mut mh = vec![0f32; half];
+    let mut mw = vec![0f32; half];
+    if section.len() == 3 {
+        if interleaved {
+            // Qwen3-VL: height at columns 1,4,7,… ; width at 2,5,8,… ;
+            // temporal keeps 0,3,6,… — each `take`n from `mrope_section`.
+            for i in (1..half).step_by(3).take(section[1]) {
+                mh[i] = 1.0;
+            }
+            for i in (2..half).step_by(3).take(section[2]) {
+                mw[i] = 1.0;
+            }
+        } else {
+            // Qwen2-VL: contiguous blocks [text | height | width].
+            let h_start = section[0].min(half);
+            let h_end = (section[0] + section[1]).min(half);
+            for m in mh.iter_mut().take(h_end).skip(h_start) {
+                *m = 1.0;
+            }
+            for m in mw.iter_mut().take(half).skip(h_end) {
+                *m = 1.0;
+            }
+        }
+    }
+    let mt: Vec<f32> = (0..half)
+        .map(|i| {
+            if mh[i] == 0.0 && mw[i] == 0.0 {
+                1.0
+            } else {
+                0.0
+            }
+        })
+        .collect();
+    (mt, mh, mw)
+}
+
 impl RotaryEmbedding {
     pub fn new(dtype: DType, cfg: &TextConfig, dev: &Device) -> Result<Self> {
         let head_dim = cfg.head_dim;
@@ -52,44 +118,88 @@ impl RotaryEmbedding {
             .step_by(2)
             .map(|i| 1f32 / rope.rope_theta.powf(i as f64 / rotary_dim as f64) as f32)
             .collect();
-        let n = inv_freq.len();
-        let inv_freq = Tensor::from_vec(inv_freq, (1, n), dev)?.to_dtype(DType::F32)?;
+        let half = inv_freq.len();
+        let inv_freq = Tensor::from_vec(inv_freq, (1, half), dev)?.to_dtype(DType::F32)?;
         let t = Tensor::arange(0u32, max_seq_len as u32, dev)?
             .to_dtype(DType::F32)?
             .reshape((max_seq_len, 1))?;
         let freqs = t.matmul(&inv_freq)?;
+
+        // MRoPE axis masks. `sum(mrope_section)` should equal `half`;
+        // warn-tolerant: any shortfall just stays on the temporal axis.
+        let (mt, mh, mw) = mrope_masks(half, &rope.mrope_section, rope.mrope_interleaved);
+        let mask_t = Tensor::from_vec(mt, (1, half), dev)?;
+        let mask_h = Tensor::from_vec(mh, (1, half), dev)?;
+        let mask_w = Tensor::from_vec(mw, (1, half), dev)?;
+
         Ok(Self {
             sin: freqs.sin()?.to_dtype(dtype)?,
             cos: freqs.cos()?.to_dtype(dtype)?,
+            inv_freq,
+            mask_t,
+            mask_h,
+            mask_w,
+            dtype,
             rotary_dim,
             head_dim,
         })
     }
 
-    /// Apply RoPE to q, k.
-    ///
-    /// `q`, `k` shape: `(B, H, L, head_dim)`. `offset` is the index
-    /// into the cached cos/sin table — the position of the first token
-    /// in the current step.
-    ///
-    /// When `rotary_dim < head_dim` the rotation is applied only to the
-    /// first `rotary_dim` dims of each head; the tail passes through
-    /// unchanged (matches the reference Python's
-    /// `apply_rotary_pos_emb` with non-trivial `partial_rotary_factor`).
-    pub fn apply(
+    /// cos/sin for a contiguous run of `seq_len` positions starting at
+    /// `pos`, by narrowing the precomputed tables. The text / decode
+    /// path (all three MRoPE axes equal → plain RoPE). Shape
+    /// `(seq_len, rotary_dim/2)`.
+    pub fn plain_cos_sin(
+        &self,
+        pos: usize,
+        seq_len: usize,
+    ) -> candle_core::Result<(Tensor, Tensor)> {
+        let cos = self.cos.narrow(0, pos, seq_len)?;
+        let sin = self.sin.narrow(0, pos, seq_len)?;
+        Ok((cos, sin))
+    }
+
+    /// cos/sin from explicit per-token 3D position ids, shape
+    /// `(3, seq_len)` (axes: text, height, width). Builds each axis's
+    /// frequencies and blends them at the interleave index sets, so
+    /// every rotary frequency slot is driven by exactly one axis.
+    /// Reduces exactly to [`Self::plain_cos_sin`] when the three axes are
+    /// equal. Returns cos/sin of shape `(seq_len, rotary_dim/2)`.
+    pub fn mrope_cos_sin(&self, position_ids: &Tensor) -> candle_core::Result<(Tensor, Tensor)> {
+        let pos = position_ids.to_dtype(DType::F32)?;
+        let (axes, seq_len) = pos.dims2()?;
+        debug_assert_eq!(axes, 3, "mrope position_ids must have 3 axes");
+        // Per-axis freqs: pos[a] (seq,1) @ inv_freq (1,half) → (seq,half).
+        let ft = pos.i(0)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
+        let fh = pos.i(1)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
+        let fw = pos.i(2)?.reshape((seq_len, 1))?.matmul(&self.inv_freq)?;
+        // Blend: each column belongs to exactly one axis (masks partition
+        // the half-dim), so this picks the right axis per frequency slot.
+        let blended = ft
+            .broadcast_mul(&self.mask_t)?
+            .add(&fh.broadcast_mul(&self.mask_h)?)?
+            .add(&fw.broadcast_mul(&self.mask_w)?)?;
+        let cos = blended.cos()?.to_dtype(self.dtype)?;
+        let sin = blended.sin()?.to_dtype(self.dtype)?;
+        Ok((cos, sin))
+    }
+
+    /// Apply rotary to `q`, `k` (shape `(B, H, L, head_dim)`) using
+    /// precomputed `cos`/`sin` of shape `(L, rotary_dim/2)`. Partial
+    /// rotary: only the first `rotary_dim` dims rotate; the tail passes
+    /// through unchanged.
+    pub fn apply_cos_sin(
         &self,
         q: &Tensor,
         k: &Tensor,
-        offset: usize,
+        cos: &Tensor,
+        sin: &Tensor,
     ) -> candle_core::Result<(Tensor, Tensor)> {
-        let (_, _, seq_len, head_dim_in) = q.dims4()?;
+        let (_, _, _seq_len, head_dim_in) = q.dims4()?;
         debug_assert_eq!(head_dim_in, self.head_dim, "q head_dim mismatch");
-        let cos = self.cos.narrow(0, offset, seq_len)?;
-        let sin = self.sin.narrow(0, offset, seq_len)?;
         if self.rotary_dim == self.head_dim {
-            // Full rotation.
-            let q_embed = candle_nn::rotary_emb::rope_slow(&q.contiguous()?, &cos, &sin)?;
-            let k_embed = candle_nn::rotary_emb::rope_slow(&k.contiguous()?, &cos, &sin)?;
+            let q_embed = candle_nn::rotary_emb::rope_slow(&q.contiguous()?, cos, sin)?;
+            let k_embed = candle_nn::rotary_emb::rope_slow(&k.contiguous()?, cos, sin)?;
             Ok((q_embed, k_embed))
         } else {
             // Partial rotation: narrow → rotate → cat the untouched tail.
@@ -102,8 +212,8 @@ impl RotaryEmbedding {
                 .narrow(candle_core::D::Minus1, 0, self.rotary_dim)?
                 .contiguous()?;
             let k_pass = k.narrow(candle_core::D::Minus1, self.rotary_dim, tail)?;
-            let q_rotated = candle_nn::rotary_emb::rope_slow(&q_rot, &cos, &sin)?;
-            let k_rotated = candle_nn::rotary_emb::rope_slow(&k_rot, &cos, &sin)?;
+            let q_rotated = candle_nn::rotary_emb::rope_slow(&q_rot, cos, sin)?;
+            let k_rotated = candle_nn::rotary_emb::rope_slow(&k_rot, cos, sin)?;
             let q_embed =
                 Tensor::cat(&[&q_rotated, &q_pass.contiguous()?], candle_core::D::Minus1)?;
             let k_embed =
@@ -112,3 +222,262 @@ impl RotaryEmbedding {
         }
     }
 }
+
+/// Compute interleaved-M-RoPE 3D position ids for a full prompt that may
+/// contain image-placeholder runs, plus the decode `rope_delta`.
+///
+/// Mirrors the reference `get_rope_index`:
+/// - text tokens advance a single running counter `c`, all three axes
+///   equal (`[c, c, c]`);
+/// - each contiguous run of `image_token_id` is one image; its tokens get
+///   `[base + t, base + h, base + w]` in row-major (t outer, h, w inner),
+///   where `base` is the counter at the run's start; after the run the
+///   counter resumes from `base + max(grid_t, grid_h, grid_w)`.
+///
+/// Returns `(text_pos, height_pos, width_pos, rope_delta)`, each pos `Vec`
+/// length `input_ids.len()`. `rope_delta = final_counter - seq_len`: add it
+/// to a plain decode offset so text resumes from the counter after the
+/// (position-compressed) image blocks.
+///
+/// Fixed-resolution assumption (Stage C): each image run is a perfect
+/// square with `grid_t = 1` (still image) and `grid_h = grid_w =
+/// isqrt(run_len)` — 196 → 14×14. Dynamic resolution (#14) would thread
+/// real per-image grids instead.
+pub(crate) fn get_rope_index(input_ids: &[u32], image_token_id: u32) -> Result<MRopeIndex> {
+    let n = input_ids.len();
+    let mut text = Vec::with_capacity(n);
+    let mut height = Vec::with_capacity(n);
+    let mut width = Vec::with_capacity(n);
+    let mut counter: i64 = 0;
+    let mut i = 0;
+    while i < n {
+        if input_ids[i] == image_token_id {
+            let start = i;
+            while i < n && input_ids[i] == image_token_id {
+                i += 1;
+            }
+            let run = i - start;
+            let g = run.isqrt();
+            if g * g != run {
+                anyhow::bail!(
+                    "get_rope_index: image run length {run} is not a perfect square \
+                     (fixed-resolution Stage C assumes a square grid; dynamic resolution is #14)"
+                );
+            }
+            let (grid_t, grid_h, grid_w) = (1usize, g, g);
+            let base = counter;
+            for tt in 0..grid_t {
+                for hh in 0..grid_h {
+                    for ww in 0..grid_w {
+                        text.push(base + tt as i64);
+                        height.push(base + hh as i64);
+                        width.push(base + ww as i64);
+                    }
+                }
+            }
+            counter = base + grid_t.max(grid_h).max(grid_w) as i64;
+        } else {
+            text.push(counter);
+            height.push(counter);
+            width.push(counter);
+            counter += 1;
+            i += 1;
+        }
+    }
+    let delta = counter - n as i64;
+    Ok((text, height, width, delta))
+}
+
+/// `(text_pos, height_pos, width_pos, rope_delta)` returned by
+/// [`get_rope_index`]; the three vectors combine into the `(3, seq)`
+/// MRoPE position-id tensor.
+pub(crate) type MRopeIndex = (Vec<i64>, Vec<i64>, Vec<i64>, i64);
+
+/// Build the `(3, seq)` position-id tensor consumed by
+/// [`RotaryEmbedding::mrope_cos_sin`] from the three axis vectors.
+///
+/// Built directly as **f32** (positions are small integers, exact in
+/// f32 well past any context length): the freqs matmul needs float
+/// anyway, and this avoids an i64 tensor / i64→f32 cast on the GPU.
+pub(crate) fn mrope_position_tensor(
+    text: &[i64],
+    height: &[i64],
+    width: &[i64],
+    dev: &Device,
+) -> candle_core::Result<Tensor> {
+    let seq = text.len();
+    let mut flat = Vec::with_capacity(3 * seq);
+    flat.extend(text.iter().map(|&x| x as f32));
+    flat.extend(height.iter().map(|&x| x as f32));
+    flat.extend(width.iter().map(|&x| x as f32));
+    Tensor::from_vec(flat, (3, seq), dev)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::IndexOp;
+
+    /// A TextConfig stub with Qwen3.6's rope params (head_dim 256,
+    /// partial 0.25 → rotary_dim 64 → half 32; section [11,11,10]).
+    fn qwen36_cfg() -> TextConfig {
+        serde_json::from_value(serde_json::json!({
+            "hidden_size": 5120,
+            "num_hidden_layers": 1,
+            "num_attention_heads": 64,
+            "num_key_value_heads": 8,
+            "head_dim": 256,
+            "intermediate_size": 1,
+            "vocab_size": 10,
+            "rms_norm_eps": 1e-6,
+            "max_position_embeddings": 64,
+            "layer_types": ["full_attention"],
+            "rope_parameters": {
+                "rope_theta": 10000000.0,
+                "partial_rotary_factor": 0.25,
+                "mrope_section": [11, 11, 10],
+                "mrope_interleaved": true
+            }
+        }))
+        .expect("cfg")
+    }
+
+    #[test]
+    fn mrope_masks_partition_the_half_dim() {
+        let (mt, mh, mw) = mrope_masks(32, &[11, 11, 10], true);
+        // Each column belongs to exactly one axis.
+        for i in 0..32 {
+            let s = mt[i] + mh[i] + mw[i];
+            assert_eq!(s, 1.0, "column {i} covered {s} times");
+        }
+        assert_eq!(mt.iter().sum::<f32>(), 11.0);
+        assert_eq!(mh.iter().sum::<f32>(), 11.0);
+        assert_eq!(mw.iter().sum::<f32>(), 10.0);
+        // Interleave: temporal 0,3,…; height 1,4,…; width 2,5,…
+        assert_eq!(mt[0], 1.0);
+        assert_eq!(mh[1], 1.0);
+        assert_eq!(mw[2], 1.0);
+        assert_eq!(mt[3], 1.0);
+    }
+
+    /// The load-bearing invariant: when all three position axes are
+    /// equal (text), `mrope_cos_sin` must reproduce `plain_cos_sin`
+    /// bit-for-bit — i.e. M-RoPE is a no-op for text, so text inference
+    /// is unchanged.
+    #[test]
+    fn mrope_reduces_to_plain_for_equal_axes() {
+        let dev = Device::Cpu;
+        let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
+
+        // positions 5,6,7 on all three axes.
+        let base: Vec<i64> = vec![5, 6, 7];
+        let pos =
+            Tensor::from_vec([base.clone(), base.clone(), base].concat(), (3, 3), &dev).unwrap();
+
+        let (mc, ms) = rope.mrope_cos_sin(&pos).unwrap();
+        let (pc, ps) = rope.plain_cos_sin(5, 3).unwrap();
+
+        let dcos = (mc - pc).unwrap().abs().unwrap().max_all().unwrap();
+        let dsin = (ms - ps).unwrap().abs().unwrap().max_all().unwrap();
+        assert!(
+            dcos.to_scalar::<f32>().unwrap() < 1e-6,
+            "cos mismatch {dcos:?}"
+        );
+        assert!(
+            dsin.to_scalar::<f32>().unwrap() < 1e-6,
+            "sin mismatch {dsin:?}"
+        );
+    }
+
+    /// Hand-checked interleave: a width-axis column (index 2) must track
+    /// the WIDTH position, while a temporal column (index 0) tracks the
+    /// TEXT position, even when the axes differ.
+    #[test]
+    fn mrope_blends_axes_at_interleave_columns() {
+        let dev = Device::Cpu;
+        let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
+        let half = rope.inv_freq.dim(1).unwrap();
+        let inv: Vec<f32> = rope.inv_freq.i(0).unwrap().to_vec1().unwrap();
+
+        // One token: text=10, height=3, width=7 — all distinct.
+        let pos = Tensor::from_vec(vec![10i64, 3, 7], (3, 1), &dev).unwrap();
+        let (cos, _sin) = rope.mrope_cos_sin(&pos).unwrap();
+        let cos_row: Vec<f32> = cos.i(0).unwrap().to_vec1().unwrap();
+        assert_eq!(cos_row.len(), half);
+
+        // Column 0 (temporal) → text pos 10. Column 1 (height) → 3.
+        // Column 2 (width) → 7.
+        assert!((cos_row[0] - (10.0 * inv[0]).cos()).abs() < 1e-5);
+        assert!((cos_row[1] - (3.0 * inv[1]).cos()).abs() < 1e-5);
+        assert!((cos_row[2] - (7.0 * inv[2]).cos()).abs() < 1e-5);
+        assert!((cos_row[3] - (10.0 * inv[3]).cos()).abs() < 1e-5);
+    }
+
+    #[test]
+    fn get_rope_index_text_only_is_sequential() {
+        let (t, h, w, delta) = get_rope_index(&[1, 2, 3, 4], 99).unwrap();
+        assert_eq!(t, vec![0, 1, 2, 3]);
+        assert_eq!(h, vec![0, 1, 2, 3]);
+        assert_eq!(w, vec![0, 1, 2, 3]);
+        assert_eq!(delta, 0, "no image → delta 0 → plain decode positions");
+    }
+
+    #[test]
+    fn get_rope_index_text_image_text() {
+        // [text, image(2x2 run of 4), text]. image_token = 99.
+        let ids = [1u32, 99, 99, 99, 99, 2];
+        let (t, h, w, delta) = get_rope_index(&ids, 99).unwrap();
+        // token 0: text → 0. image base=1, grid 1x2x2:
+        //   t all = 1; h = base+row = [1,1,2,2]; w = base+col = [1,2,1,2].
+        // resume from base + max(1,2,2) = 3. trailing text → 3.
+        assert_eq!(t, vec![0, 1, 1, 1, 1, 3]);
+        assert_eq!(h, vec![0, 1, 1, 2, 2, 3]);
+        assert_eq!(w, vec![0, 1, 2, 1, 2, 3]);
+        // final counter = 4, seq_len = 6 → delta = -2 (the 4 image tokens
+        // advanced the counter by only 2).
+        assert_eq!(delta, -2);
+        // Decode after the prompt (offset = 6) → text position 6 + (-2) = 4.
+        assert_eq!(6 + delta, 4);
+    }
+
+    #[test]
+    fn get_rope_index_rejects_non_square_image_run() {
+        // 196 is square (14x14) — ok. 195 is not.
+        assert!(get_rope_index(&[99u32; 196], 99).is_ok());
+        assert!(get_rope_index(&[99u32; 195], 99).is_err());
+    }
+
+    #[test]
+    fn position_tensor_round_trips_through_mrope_cos_sin() {
+        // get_rope_index → (3,seq) tensor → mrope_cos_sin, and confirm an
+        // image token's height column tracks its grid row (not the text
+        // counter), i.e. the end-to-end position plumbing is wired right.
+        let dev = Device::Cpu;
+        let rope = RotaryEmbedding::new(DType::F32, &qwen36_cfg(), &dev).unwrap();
+        let ids = [1u32, 99, 99, 99, 99]; // text + 2x2 image
+        let (t, h, w, _d) = get_rope_index(&ids, 99).unwrap();
+        let pos = mrope_position_tensor(&t, &h, &w, &dev).unwrap();
+        assert_eq!(pos.dims(), &[3, 5]);
+        let (cos, _sin) = rope.mrope_cos_sin(&pos).unwrap();
+        assert_eq!(cos.dims(), &[5, rope.inv_freq.dim(1).unwrap()]);
+
+        let inv: Vec<f32> = rope.inv_freq.i(0).unwrap().to_vec1().unwrap();
+        // Last image token (index 4): grid (h=1, w=1) → base 1 → h=2, w=2.
+        // Height column (index 1) must track h-position 2, not text.
+        let last: Vec<f32> = cos.i(4).unwrap().to_vec1().unwrap();
+        assert!((last[1] - (2.0 * inv[1]).cos()).abs() < 1e-5);
+    }
+
+    #[test]
+    fn get_rope_index_196_is_14x14() {
+        let mut ids = vec![1u32]; // one text token
+        ids.extend(std::iter::repeat_n(99u32, 196));
+        let (t, h, w, _delta) = get_rope_index(&ids, 99).unwrap();
+        // image base = 1. Last image token (index 196) is grid (h=13,w=13).
+        assert_eq!(*t.last().unwrap(), 1, "grid_t=1 → temporal const at base");
+        assert_eq!(h[1], 1, "first image row at base");
+        assert_eq!(w[1], 1, "first image col at base");
+        assert_eq!(h[196], 1 + 13, "last image row = base + 13");
+        assert_eq!(w[196], 1 + 13, "last image col = base + 13");
+    }
+}

crates/neuron/src/harness/tp/tp_qwen3_5.rs

@@ -526,7 +526,8 @@ impl TpQwen3_5Attention {
         &mut self,
         x: &Tensor,
         attn_mask: Option<&Tensor>,
-        offset: usize,
+        cos: &Tensor,
+        sin: &Tensor,
     ) -> candle_core::Result<Tensor> {
         let (b, l, _) = x.dims3()?;
 
@@ -559,7 +560,7 @@ impl TpQwen3_5Attention {
             .transpose(1, 2)?
             .contiguous()?;
 
-        let (q, k) = self.rotary.apply(&q, &k, offset)?;
+        let (q, k) = self.rotary.apply_cos_sin(&q, &k, cos, sin)?;
         let (k, v) = self.kv_cache.append(&k, &v)?;
         let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
         let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
@@ -807,11 +808,12 @@ impl TpQwen3_5DecoderLayer {
         &mut self,
         x: &Tensor,
         attn_mask: Option<&Tensor>,
-        offset: usize,
+        cos: &Tensor,
+        sin: &Tensor,
     ) -> candle_core::Result<Tensor> {
         let h = self.input_layernorm.forward(x)?;
         let attn_out = match &mut self.attention {
-            TpAttentionKind::Full(attn) => attn.forward(&h, attn_mask, offset)?,
+            TpAttentionKind::Full(attn) => attn.forward(&h, attn_mask, cos, sin)?,
             TpAttentionKind::Linear(net) => net.forward(&h)?,
         };
         let x = (x + attn_out)?;
@@ -834,6 +836,15 @@ pub struct TpQwen3_5Model {
     embed_tokens: Embedding,
     layers: Vec<TpQwen3_5DecoderLayer>,
     norm: Qwen3_5RmsNorm,
+    /// Replicated rotary, shared with every full-attention layer. The
+    /// model builds the per-forward cos/sin (interleaved M-RoPE for image
+    /// tokens, plain for text) once and the layers apply it. Identical on
+    /// every rank, so per-rank position ids stay consistent.
+    rotary: Arc<RotaryEmbedding>,
+    /// `offset + rope_delta` is the text-axis decode position; set from
+    /// `get_rope_index` during a vision prefill, reset in `clear_kv_cache`.
+    /// See `Qwen3_5Model::rope_delta`.
+    rope_delta: i64,
     device: Device,
     dtype: DType,
 }
@@ -900,6 +911,8 @@ impl TpQwen3_5Model {
             embed_tokens,
             layers,
             norm,
+            rotary,
+            rope_delta: 0,
             device,
             dtype,
         })
@@ -956,6 +969,8 @@ impl TpQwen3_5Model {
             embed_tokens,
             layers,
             norm,
+            rotary,
+            rope_delta: 0,
             device,
             dtype,
         })
@@ -969,6 +984,14 @@ impl TpQwen3_5Model {
         for l in &mut self.layers {
             l.clear_kv_cache();
         }
+        self.rope_delta = 0;
+    }
+
+    /// Set the decode `rope_delta` computed by `get_rope_index` during a
+    /// vision prefill, so decode after the image resumes text positions
+    /// from the image-compressed counter.
+    pub fn set_rope_delta(&mut self, delta: i64) {
+        self.rope_delta = delta;
     }
 
     fn causal_mask(&self, b: usize, tgt: usize, offset: usize) -> candle_core::Result<Tensor> {
@@ -980,64 +1003,80 @@ impl TpQwen3_5Model {
     }
 
     pub fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
-        let (b, l) = input.dims2()?;
-        let mut h = self.embed_tokens.forward(input)?;
-        let causal = if l == 1 {
-            None
-        } else {
-            Some(self.causal_mask(b, l, offset)?)
-        };
-        for layer in &mut self.layers {
-            h = layer.forward(&h, causal.as_ref(), offset)?;
-        }
-        self.norm.forward(&h)
+        self.forward_inner(input, offset, None, None, None)
     }
 
-    /// Forward with image-embedding splice (TP, replicated tower).
-    ///
-    /// Mirrors the single-GPU `Qwen3_5Model::forward_inner` splice:
-    /// embed locally, replace the rows at `image_token_id` positions
-    /// with the image patch embeddings, then run the sharded decoder
-    /// stack. The TP invariant is that every rank holds an identical
-    /// hidden state (only the attention/MLP matmuls shard, with a
-    /// trailing `AllReduce`). That holds here because every rank
-    /// encodes the *same* pixels through its *replicated* vision tower
-    /// and so produces identical `image_embeds` — no broadcast needed.
-    pub fn forward_with_vision(
+    /// Forward for a vision-prefill chunk: optional image-embedding
+    /// splice plus explicit interleaved-M-RoPE `position_ids` (the
+    /// chunk's slice of the full prompt's 3D positions). Used by
+    /// `TpQwen3_5ForCausalLM::prefill_with_images_chunked`, which
+    /// computes the positions once over the whole prompt and slices them
+    /// per chunk so every rank steps in lockstep.
+    pub fn forward_with_positions(
         &mut self,
         input: &Tensor,
         offset: usize,
-        image_embeds: &Tensor,
-        image_token_id: u32,
+        position_ids: &Tensor,
+        image_embeds: Option<&Tensor>,
+        image_token_id: Option<u32>,
+    ) -> candle_core::Result<Tensor> {
+        self.forward_inner(
+            input,
+            offset,
+            image_embeds,
+            image_token_id,
+            Some(position_ids),
+        )
+    }
+
+    /// Shared forward. Splices image embeddings at `image_token_id`
+    /// positions when present, then builds the rotary cos/sin — from the
+    /// explicit `position_ids` (interleaved M-RoPE, vision) when given,
+    /// else plain positions at `offset + rope_delta` (text / decode) —
+    /// and runs the sharded decoder stack. The TP replicated-hidden-state
+    /// invariant holds because every rank encodes the same pixels and
+    /// computes the same positions.
+    fn forward_inner(
+        &mut self,
+        input: &Tensor,
+        offset: usize,
+        image_embeds: Option<&Tensor>,
+        image_token_id: Option<u32>,
+        position_ids: Option<&Tensor>,
     ) -> candle_core::Result<Tensor> {
         let (b, l) = input.dims2()?;
         let mut h = self.embed_tokens.forward(input)?;
 
-        // Locate the image-token positions in the (pre-expanded) input
-        // ids and splice the patch rows in. Same CPU-side scan as the
-        // single-GPU path; the count must match the patch dimension or
-        // the prompt expansion is wrong.
-        let ids: Vec<u32> = input.flatten_all()?.to_vec1()?;
-        let mut positions: Vec<u32> = Vec::with_capacity(image_embeds.dim(0)?);
-        for (idx, id) in ids.iter().enumerate() {
-            if *id == image_token_id {
-                positions.push(idx as u32);
+        if let (Some(img), Some(tok_id)) = (image_embeds, image_token_id) {
+            let ids: Vec<u32> = input.flatten_all()?.to_vec1()?;
+            let mut positions: Vec<u32> = Vec::with_capacity(img.dim(0)?);
+            for (idx, id) in ids.iter().enumerate() {
+                if *id == tok_id {
+                    positions.push(idx as u32);
+                }
+            }
+            let n_img_tokens = img.dim(0)?;
+            if positions.len() != n_img_tokens {
+                candle_core::bail!(
+                    "TP forward: chunk has {} image-token positions but image_embeds carries \
+                     {} tokens — patch-count expansion / chunk slicing mismatch",
+                    positions.len(),
+                    n_img_tokens,
+                );
+            }
+            if !positions.is_empty() {
+                let img = img.to_dtype(self.dtype)?;
+                h = splice_runs(&h, &img, &positions)?;
             }
         }
-        let n_img_tokens = image_embeds.dim(0)?;
-        if positions.len() != n_img_tokens {
-            candle_core::bail!(
-                "TP forward_with_vision: prompt has {} image-token positions but \
-                 image_embeds carries {} tokens — ensure the per-image patch-count \
-                 expansion has been applied",
-                positions.len(),
-                n_img_tokens,
-            );
-        }
-        if !positions.is_empty() {
-            let img = image_embeds.to_dtype(self.dtype)?;
-            h = splice_runs(&h, &img, &positions)?;
-        }
+
+        let (cos, sin) = match position_ids {
+            Some(pos) => self.rotary.mrope_cos_sin(pos)?,
+            None => {
+                let base = (offset as i64 + self.rope_delta).max(0) as usize;
+                self.rotary.plain_cos_sin(base, l)?
+            }
+        };
 
         let causal = if l == 1 {
             None
@@ -1045,7 +1084,7 @@ impl TpQwen3_5Model {
             Some(self.causal_mask(b, l, offset)?)
         };
         for layer in &mut self.layers {
-            h = layer.forward(&h, causal.as_ref(), offset)?;
+            h = layer.forward(&h, causal.as_ref(), &cos, &sin)?;
         }
         self.norm.forward(&h)
     }
@@ -1174,21 +1213,25 @@ impl TpQwen3_5ForCausalLM {
         hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
     }
 
-    /// Forward with image-embedding splice (TP). Mirrors `forward` but
-    /// routes through `TpQwen3_5Model::forward_with_vision` so the
-    /// per-rank input embeddings get the image patches spliced in at
-    /// `image_token_id` positions before the sharded decoder stack.
-    pub fn forward_with_vision(
+    /// Forward for a vision-prefill chunk (optional image splice +
+    /// explicit interleaved-M-RoPE `position_ids`). Mirrors `forward`
+    /// but routes through `TpQwen3_5Model::forward_with_positions`.
+    pub fn forward_with_positions(
         &mut self,
         input: &Tensor,
         offset: usize,
-        image_embeds: &Tensor,
-        image_token_id: u32,
+        position_ids: &Tensor,
+        image_embeds: Option<&Tensor>,
+        image_token_id: Option<u32>,
     ) -> candle_core::Result<Tensor> {
         let (_, l) = input.dims2()?;
-        let hidden = self
-            .base
-            .forward_with_vision(input, offset, image_embeds, image_token_id)?;
+        let hidden = self.base.forward_with_positions(
+            input,
+            offset,
+            position_ids,
+            image_embeds,
+            image_token_id,
+        )?;
         hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
     }
 
@@ -1245,6 +1288,21 @@ impl TpQwen3_5ForCausalLM {
         let device = self.device().clone();
         let image_embeds = self.encode_images_concat(image_pixels)?;
 
+        // Interleaved-M-RoPE 3D position ids for the whole prompt,
+        // computed once and sliced per chunk so every rank assigns image
+        // tokens their 14×14 grid coordinates (and text after the image
+        // resumes from the compressed counter). `rope_delta` is stored on
+        // the base model for the decode that follows this prefill. Every
+        // chunk — text or image — uses the M-RoPE slice, because the image
+        // shifts the positions of the text around it.
+        let (text, height, width, delta) =
+            crate::harness::arch::qwen3_5::rope::get_rope_index(tokens, image_token_id)
+                .map_err(|e| candle_core::Error::Msg(format!("get_rope_index: {e}")))?;
+        self.base.set_rope_delta(delta);
+        let full_pos = crate::harness::arch::qwen3_5::rope::mrope_position_tensor(
+            &text, &height, &width, &device,
+        )?;
+
         let mut last_logits: Option<Tensor> = None;
         // Rows of `image_embeds` already spliced by earlier chunks. The
         // `<|image_pad|>` run is contiguous, so chunks consume embedding
@@ -1255,16 +1313,22 @@ impl TpQwen3_5ForCausalLM {
             let end = (start + chunk_size).min(tokens.len());
             let chunk = &tokens[start..end];
             let input = Tensor::new(chunk, &device)?.unsqueeze(0)?;
+            let pos_slice = full_pos.narrow(1, start, end - start)?;
             let n_here = chunk.iter().filter(|&&t| t == image_token_id).count();
             let logits = if n_here == 0 {
-                // Pure-text chunk — same forward the text prefill runs.
-                self.forward(&input, base_offset + start)?
+                self.forward_with_positions(&input, base_offset + start, &pos_slice, None, None)?
             } else {
                 // Splice the next `n_here` patch rows at this chunk's
                 // local image-pad positions.
                 let rows = image_embeds.narrow(0, img_off, n_here)?;
                 img_off += n_here;
-                self.forward_with_vision(&input, base_offset + start, &rows, image_token_id)?
+                self.forward_with_positions(
+                    &input,
+                    base_offset + start,
+                    &pos_slice,
+                    Some(&rows),
+                    Some(image_token_id),
+                )?
             };
             last_logits = Some(logits);
             start = end;