From b3a6edf67c9144b8639fbb19803928117cc33644 Mon Sep 17 00:00:00 2001 From: taking-lying-flat <1615405@qq.com> Date: Tue, 1 Sep 2026 08:58:05 +0800 Subject: [PATCH] [bugfix] preserve QSA MRoPE batch dimension --- src/mcore_bridge/model/modules/qsa_indexer.py | 33 ++++++++++++------ tests/test_qsa_indexer.py | 34 +++++++++++++++++++ 2 files changed, 56 insertions(+), 11 deletions(-) create mode 100644 tests/test_qsa_indexer.py diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py index 5a2c706c..a08b9a8d 100644 --- a/src/mcore_bridge/model/modules/qsa_indexer.py +++ b/src/mcore_bridge/model/modules/qsa_indexer.py @@ -27,6 +27,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, tp_group=None): @@ -67,10 +77,12 @@ def select_mask(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch hidden_states: ``[s, b, h]`` (mcore layout), pre-attention input -- the same tensor the reference indexer consumes. ``s`` is the local sequence length when sequence parallelism is enabled. - 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 @@ -111,11 +123,10 @@ 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_): @@ -123,8 +134,8 @@ def apply_rope(t, cos_, sin_): 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 @@ -133,7 +144,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()) diff --git a/tests/test_qsa_indexer.py b/tests/test_qsa_indexer.py new file mode 100644 index 00000000..c2f2b897 --- /dev/null +++ b/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])