diff --git a/python/src/coreai_models/primitives/ios/rope.py b/python/src/coreai_models/primitives/ios/rope.py index a2bf20bd..e6c086f4 100644 --- a/python/src/coreai_models/primitives/ios/rope.py +++ b/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 diff --git a/python/tests/test_model_units/test_primitives/test_ios/test_mrope_host.py b/python/tests/test_model_units/test_primitives/test_ios/test_mrope_host.py new file mode 100644 index 00000000..8e95cde9 --- /dev/null +++ b/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)