[bugfix] preserve QSA MRoPE batch dimension

src/mcore_bridge/model/modules/qsa_indexer.py

@@ -26,6 +26,16 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor:
     return torch.cat((-x2, x1), dim=-1)
 
 
+def _materialize_rope(freqs: torch.Tensor, seq_len: int, dtype: torch.dtype, mscale: float):
+    """Materialize batch-aware RoPE cos/sin tensors in QSA layout."""
+    # [s, freq_b, 1, d] -> [freq_b, s, d]. Keeping freq_b separate is
+    # required for MRoPE, whose positions can differ between samples.
+    f = freqs[:seq_len].squeeze(2).permute(1, 0, 2)
+    cos = (torch.cos(f) * mscale).to(dtype)
+    sin = (torch.sin(f) * mscale).to(dtype)
+    return cos, sin
+
+
 class QSAIndexer(nn.Module):
     # refer: transformers Qwen4ExpTextQSAIndexer
     def __init__(self, config):
@@ -64,10 +74,12 @@ def select_mask(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch
         Args:
             hidden_states: ``[s, b, h]`` (mcore layout), pre-attention input --
                 the same tensor the reference indexer consumes.
-            freqs: mcore rotary frequencies ``[s, 1, 1, rot_dim]``. mcore stores
-                angles rather than cos/sin, so they are materialized here the way
-                ``_patch_apply_rotary_pos_emb`` does (``cos(freqs) * mscale``),
-                keeping the indexer's RoPE identical to the attention's.
+            freqs: mcore rotary frequencies ``[s, freq_b, 1, rot_dim]``, where
+                ``freq_b`` is 1 for ordinary RoPE or ``b`` for batch-dependent
+                MRoPE. mcore stores angles rather than cos/sin, so they are
+                materialized here the way ``_patch_apply_rotary_pos_emb`` does
+                (``cos(freqs) * mscale``), keeping the indexer's RoPE identical
+                to the attention's.
 
         Returns:
             ``[b, 1, s, s]`` bool mask where True marks a *masked-out* key (TE's
@@ -100,20 +112,19 @@ def select_mask(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch
         q = self.q_layernorm(q)
 
         # ---- materialize cos/sin from mcore freqs ----
-        # freqs: [s, 1, 1, rot_dim] -> [s, rot_dim]; mscale mirrors the attention path.
+        # freqs: [s, freq_b, 1, rot_dim] -> [freq_b, s, rot_dim].
+        # mscale mirrors the attention path.
         mscale = getattr(self.config, 'attention_scaling', 1.0) or 1.0
-        f = freqs.reshape(freqs.shape[0], -1)[:s]
-        cos = (torch.cos(f) * mscale).to(q.dtype)
-        sin = (torch.sin(f) * mscale).to(q.dtype)
+        cos, sin = _materialize_rope(freqs, s, q.dtype, mscale)
         rot = cos.shape[-1]
 
         def apply_rope(t, cos_, sin_):
             t_rope, t_pass = t[..., :rot], t[..., rot:]
             t_rope = (t_rope * cos_) + (_rotate_half(t_rope) * sin_)
             return torch.cat((t_rope, t_pass), dim=-1)
 
-        # queries rotate at their own position: cos [s, rot] -> [1, s, 1, rot]
-        q = apply_rope(q, cos[None, :, None, :], sin[None, :, None, :])
+        # queries rotate at their own positions; freq_b broadcasts when it is 1.
+        q = apply_rope(q, cos[:, :, None, :], sin[:, :, None, :])
 
         # ---- pool every block once (shared across queries) ----
         usable = max_blocks * R
@@ -122,7 +133,7 @@ def apply_rope(t, cos_, sin_):
         pooled = self.k_layernorm(pooled)
         # blocks rotate at their first token's position
         starts = torch.arange(max_blocks, device=device) * R
-        block_keys = apply_rope(pooled, cos[starts][None], sin[starts][None])  # [b, nb, d]
+        block_keys = apply_rope(pooled, cos[:, starts, :], sin[:, starts, :])  # [b, nb, d]
 
         # ---- score all (query, block) pairs ----
         scores = torch.einsum('bqhd,bkd->bqhk', q.float(), block_keys.float())

tests/test_qsa_indexer.py

@@ -0,0 +1,34 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib.util
+import torch
+from pathlib import Path
+
+
+def _load_materialize_rope():
+    module_path = Path(__file__).parents[1] / 'src/mcore_bridge/model/modules/qsa_indexer.py'
+    spec = importlib.util.spec_from_file_location('qsa_indexer', module_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module._materialize_rope
+
+
+def test_materialize_rope_preserves_mrope_batch_dimension():
+    materialize_rope = _load_materialize_rope()
+    seq_len, batch_size, rope_dim = 16, 2, 64
+    freqs = torch.randn(seq_len, batch_size, 1, rope_dim)
+
+    cos, sin = materialize_rope(freqs, seq_len, torch.float32, 1.0)
+    expected = freqs.squeeze(2).permute(1, 0, 2)
+
+    assert cos.shape == (batch_size, seq_len, rope_dim)
+    torch.testing.assert_close(cos, expected.cos())
+    torch.testing.assert_close(sin, expected.sin())
+
+    changed_freqs = freqs.clone()
+    changed_freqs[:, 1].add_(0.5)
+    changed_cos, changed_sin = materialize_rope(changed_freqs, seq_len, torch.float32, 1.0)
+
+    torch.testing.assert_close(changed_cos[0], cos[0])
+    torch.testing.assert_close(changed_sin[0], sin[0])
+    assert not torch.equal(changed_cos[1], cos[1])
+    assert not torch.equal(changed_sin[1], sin[1])