Precompute M-RoPE cos/sin on the host (workaround for #280)

python/src/coreai_models/primitives/ios/rope.py

@@ -4,6 +4,7 @@
 # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
 import os
+from collections.abc import Sequence
 
 import torch
 
@@ -23,6 +24,15 @@ def rotate_half(x: torch.Tensor) -> torch.Tensor:
 def rope_gather_cached_cos_sin(
     position_ids: torch.Tensor, cos_cached: torch.Tensor, sin_cached: torch.Tensor
 ) -> list[torch.Tensor]:
+    """Gather cached cos/sin by position for in-graph RoPE.
+
+    The exported composite survives exactly one lookup per graph, and only when
+    position_ids is the graph's own position input; a second lookup, or a derived
+    operand, SIGSEGVs delegate specialization. Multi-position (M-RoPE) models
+    (Qwen2/3-VL) need a lookup per row, so precompute cos/sin on the host
+    (gather_cos_sin_host / mrope_cos_sin) and feed them as rope_cos/rope_sin inputs.
+    Runtime composite limitation; see #280.
+    """
     position_ids = position_ids.to(torch.int32)
     rope_cos = cos_cached[position_ids]
     rope_sin = sin_cached[position_ids]
@@ -45,10 +55,8 @@ def apply_rope(
     rope_sin: torch.Tensor,
     head_axis: int = 1,
 ) -> torch.Tensor:
-    # `head_axis` is where the (broadcast) head dimension lives in `x`. The
-    # default of 1 matches the `(B, n_heads, S, head_dim)` head-format layout.
-    # Callers that keep Q/K in `(B, S, n_heads, head_dim)` pass `head_axis=2` so
-    # cos/sin broadcast over heads without an extra transpose into head format.
+    # head_axis is where the broadcast head dim lives: 1 for (B, heads, S, dim),
+    # 2 for (B, S, heads, dim) -- lets cos/sin broadcast without a transpose.
     rope_cos = rope_cos.unsqueeze(head_axis)
     rope_sin = rope_sin.unsqueeze(head_axis)
 
@@ -121,5 +129,53 @@ def _compute_sin_and_cos(self, dtype: torch.dtype = torch.float32) -> None:
             self.sin_cached = torch.nn.Buffer(torch.sin(emb).to(dtype=dtype), persistent=False)
 
     def gather_cos_sin(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        """Gather the cached cos/sin values using the position_ids"""
+        """Gather the cached cos/sin values using the position_ids.
+
+        This goes through the ``coreai::rope_gather_cached_cos_sin`` custom op, which
+        lowers to the in-graph gather composite. That composite tolerates only one
+        lookup per graph, fed straight from the graph's position input -- see the
+        warning on :func:`rope_gather_cached_cos_sin`. For multi-position (M-RoPE)
+        models, or any case needing more than one lookup, do the lookup on the host via
+        :meth:`gather_cos_sin_host` / :meth:`mrope_cos_sin` and pass cos/sin in as graph
+        inputs.
+        """
         return rope_gather_cached_cos_sin(position_ids, self.cos_cached, self.sin_cached)
+
+    def gather_cos_sin_host(
+        self, position_ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Host-side cos/sin gather (plain indexing), bypassing the in-graph composite.
+
+        Precompute rope_cos/rope_sin outside the graph and feed them as inputs -- the
+        supported pattern for M-RoPE models (a lookup per time/height/width row), where
+        the in-graph composite survives only one lookup (see rope_gather_cached_cos_sin).
+        Shaped like gather_cos_sin.
+        """
+        position_ids = position_ids.to(torch.int32)
+        return self.cos_cached[position_ids], self.sin_cached[position_ids]
+
+    def mrope_cos_sin(
+        self,
+        rows: Sequence[torch.Tensor],
+        section_masks: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Combine M-RoPE position rows into one cos/sin on the host.
+
+        ``rows`` is one position_ids per section (time/height/width); ``section_masks``
+        is ``(len(rows), head_dim)``, row ``s`` being 1 on the head-dim entries that take
+        their position from ``rows[s]``. The masks let one path serve both contiguous
+        (Qwen2-VL) and interleaved (Qwen3.5) layouts, and collapse to ordinary 1-D RoPE
+        for text-only prompts. Host-only (via gather_cos_sin_host), shaped like gather_cos_sin.
+        """
+        if len(rows) != section_masks.shape[0]:
+            msg = f"expected {section_masks.shape[0]} position rows, got {len(rows)}"
+            raise ValueError(msg)
+        cos_total: torch.Tensor | None = None
+        sin_total: torch.Tensor | None = None
+        for row, mask in zip(rows, section_masks, strict=True):
+            row_cos, row_sin = self.gather_cos_sin_host(row)
+            mask = mask.to(row_cos.dtype)
+            cos_total = row_cos * mask if cos_total is None else cos_total + row_cos * mask
+            sin_total = row_sin * mask if sin_total is None else sin_total + row_sin * mask
+        assert cos_total is not None and sin_total is not None
+        return cos_total, sin_total

python/tests/test_model_units/test_primitives/test_ios/test_mrope_host.py

@@ -0,0 +1,102 @@
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+"""Host-side RoPE gather for multi-position (M-RoPE) models.
+
+Covers the issue #280 workaround: the in-graph ``rope_cached_cos_sin_gather``
+composite tolerates only one lookup per graph, so M-RoPE models precompute cos/sin on
+the host and feed them in as graph inputs.
+"""
+
+import pytest
+import torch
+
+from coreai_models.primitives.ios.rope import RoPECache
+
+
+def _interleaved_masks(head_dim: int, n_sections: int = 3) -> torch.Tensor:
+    """Interleaved M-RoPE section masks (Qwen3.5 layout): freq f -> row f % n."""
+    frequency = torch.arange(head_dim) % (head_dim // 2)
+    return torch.stack([(frequency % n_sections == s).float() for s in range(n_sections)])
+
+
+def _contiguous_masks(head_dim: int, n_sections: int = 3) -> torch.Tensor:
+    """Contiguous M-RoPE section masks (Qwen2-VL ``mrope_section`` layout)."""
+    band = head_dim // n_sections
+    masks = torch.zeros(n_sections, head_dim)
+    for s in range(n_sections):
+        hi = head_dim if s == n_sections - 1 else (s + 1) * band
+        masks[s, s * band : hi] = 1.0
+    return masks
+
+
+def test_gather_cos_sin_host_matches_composite_op():
+    """Host gather must return the same values as the custom-op gather path."""
+    rope = RoPECache(head_dim=32, max_cache_size=64)
+    pos = torch.tensor([[3, 7, 1]], dtype=torch.int32)
+
+    host_cos, host_sin = rope.gather_cos_sin_host(pos)
+    op_cos, op_sin = rope.gather_cos_sin(pos)
+
+    torch.testing.assert_close(host_cos, op_cos)
+    torch.testing.assert_close(host_sin, op_sin)
+    # And it is a plain index into the cache.
+    torch.testing.assert_close(host_cos, rope.cos_cached[pos])
+    torch.testing.assert_close(host_sin, rope.sin_cached[pos])
+
+
+def test_gather_cos_sin_host_shape():
+    rope = RoPECache(head_dim=32, max_cache_size=64)
+    pos = torch.arange(8, dtype=torch.int32).unsqueeze(0)
+    cos, sin = rope.gather_cos_sin_host(pos)
+    assert cos.shape == (1, 8, 32)
+    assert sin.shape == (1, 8, 32)
+
+
+@pytest.mark.parametrize("masks_fn", [_interleaved_masks, _contiguous_masks])
+def test_mrope_identical_rows_collapse_to_1d(masks_fn):
+    """Partitioning masks + identical rows == the ordinary single-row embedding."""
+    head_dim = 36
+    rope = RoPECache(head_dim=head_dim, max_cache_size=64)
+    masks = masks_fn(head_dim)
+    # Masks must partition the head dim for the collapse to hold.
+    torch.testing.assert_close(masks.sum(0), torch.ones(head_dim))
+
+    pos = torch.tensor([[0, 3, 7]], dtype=torch.int32)
+    cos_m, sin_m = rope.mrope_cos_sin([pos, pos, pos], masks)
+    cos_1, sin_1 = rope.gather_cos_sin_host(pos)
+
+    torch.testing.assert_close(cos_m, cos_1)
+    torch.testing.assert_close(sin_m, sin_1)
+
+
+def test_mrope_selects_the_right_row_per_dim():
+    """Each head-dim entry must take its value from the section its mask marks."""
+    head_dim = 36
+    rope = RoPECache(head_dim=head_dim, max_cache_size=64)
+    masks = _interleaved_masks(head_dim)
+
+    rows = [
+        torch.tensor([[2, 5]], dtype=torch.int32),
+        torch.tensor([[7, 1]], dtype=torch.int32),
+        torch.tensor([[3, 9]], dtype=torch.int32),
+    ]
+    cos, sin = rope.mrope_cos_sin(rows, masks)
+    per_row = [rope.gather_cos_sin_host(r) for r in rows]
+
+    for d in range(head_dim):
+        # Exactly one section owns this dim.
+        assert masks[:, d].sum().item() == 1.0
+        s = int(masks[:, d].argmax())
+        torch.testing.assert_close(cos[..., d], per_row[s][0][..., d])
+        torch.testing.assert_close(sin[..., d], per_row[s][1][..., d])
+
+
+def test_mrope_row_count_mismatch_raises():
+    rope = RoPECache(head_dim=36, max_cache_size=64)
+    masks = _interleaved_masks(36)  # 3 sections
+    pos = torch.tensor([[0, 1]], dtype=torch.int32)
+    with pytest.raises(ValueError, match="position rows"):
+        rope.mrope_cos_sin([pos, pos], masks)