fix: chunk QSA indexer scoring to avoid OOM on long packed sequences

src/mcore_bridge/model/modules/qsa_indexer.py

@@ -4,6 +4,11 @@
 from megatron.core.extensions.transformer_engine import TELinear
 from torch import nn
 
+# Byte budget for the transient score tile in select_token_indices_thd: the (token, block)
+# scoring is chunked over queries to avoid OOM from the full [T, n_heads, NB] fp32 tensor.
+# Tests monkeypatch this to compare one chunk (un-chunked reference) vs many chunks.
+_QSA_INDEX_SCORE_CHUNK_BYTES = 1024 * 1024 * 1024
+
 
 class Qwen4ExpTextRMSNorm(nn.Module):
 
@@ -356,20 +361,32 @@ def apply_rope(t, cos_, sin_):
         first_pack = cu_seqlens[block_doc].long() + block_in_doc_idx * R  # [NB]
         block_keys = apply_rope(pooled, cos[first_pack], sin[first_pack])  # [NB, d]
 
-        # ---- score every (token, block) pair ----
-        scores = torch.einsum('thd,kd->thk', q.float(), block_keys.float())
-        scores = torch.relu(scores).sum(dim=1) / math.sqrt(self.index_head_dim)  # [T, NB]
-
-        # ---- restrict to same-document, causally-before blocks ----
+        # ---- score every (token, block) pair, chunked over queries ----
+        # The score -> relu -> head-sum -> mask -> top-k pipeline is row-wise independent, so it is
+        # processed in query chunks and only each row's top-k survives. This keeps peak memory at
+        # O(chunk * NB) instead of materializing the full [T, n_heads, NB] fp32 score tensor (and the
+        # [T, NB] masked copy), which is O(seq^2 / compress_ratio) and OOMs at long packed sequences.
+        # Numerically identical to the un-chunked form: every row is computed the same way.
         q_nblocks = (pos_in_doc + 1) // R  # [T]
-        valid = (block_doc[None, :] == token_doc[:, None]) & \
-            (block_in_doc_idx[None, :] < q_nblocks[:, None])  # [T, NB]
-        scores = scores.masked_fill(~valid, float('-inf'))
+        k = min(self.block_topk, NB)
+        scale = math.sqrt(self.index_head_dim)
+        chunk = max(1, min(T, _QSA_INDEX_SCORE_CHUNK_BYTES // max(1, self.index_n_heads * NB * 4)))
+        top_blocks = torch.empty(T, k, dtype=torch.long, device=device)
+        keep = torch.empty(T, k, dtype=torch.bool, device=device)
+        for start in range(0, T, chunk):
+            end = min(start + chunk, T)
+            sc = torch.einsum('thd,kd->thk', q[start:end].float(), block_keys.float())
+            sc = torch.relu(sc).sum(dim=1) / scale  # [c, NB]
+            td = token_doc[start:end]
+            valid_c = (block_doc[None, :] == td[:, None]) & \
+                (block_in_doc_idx[None, :] < q_nblocks[start:end, None])  # [c, NB]
+            sc = sc.masked_fill(~valid_c, float('-inf'))
+            tb = sc.topk(k, dim=-1).indices  # [c, k] into [0, NB)
+            top_blocks[start:end] = tb
+            keep[start:end] = valid_c.gather(1, tb)
+        del sc, valid_c, tb
 
         # ---- top-k blocks -> token indices ----
-        k = min(self.block_topk, NB)
-        top_blocks = scores.topk(k, dim=-1).indices  # [T, k] into [0, NB)
-        keep = valid.gather(1, top_blocks)  # [T, k]
         arange_r = torch.arange(R, device=device)
         base = cu_seqlens[block_doc[top_blocks]].long() + block_in_doc_idx[top_blocks] * R  # [T, k]
         top_idx = (base.unsqueeze(-1) + arange_r).flatten(-2)  # [T, k*R]

tests/test_qsa_indexer.py

@@ -1,9 +1,92 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
+import pytest
 import torch
 
 from mcore_bridge.model.modules.qsa_indexer import _materialize_rope, _rotate_half
 
 
+def _packed_inputs(idx, doc_lens, hidden_size, d, device):
+    """Build a packed (thd) multi-doc input for select_token_indices_thd."""
+    T = sum(doc_lens)
+    cu = [0]
+    for L in doc_lens:
+        cu.append(cu[-1] + L)
+    cu_seqlens = torch.tensor(cu, dtype=torch.long, device=device)
+    hidden_tok = torch.randn(T, hidden_size, device=device)
+    freqs = torch.randn(T, 1, 1, d, device=device)
+    return hidden_tok, freqs, cu_seqlens
+
+
+def _make_idx(compress_ratio, budget, device):
+    from test_qwen4_exp_units import _make_config
+
+    from mcore_bridge.model.modules.qsa_indexer import QSAIndexer
+    cfg = _make_config(compress_ratio=compress_ratio, budget=budget)
+    idx = QSAIndexer(cfg).to(device)
+    with torch.no_grad():
+        idx.index_qk_proj.weight.normal_(0, 0.02)
+        idx.q_layernorm.weight.normal_(0, 0.02)
+        idx.k_layernorm.weight.normal_(0, 0.02)
+    return idx, cfg
+
+
+def test_qsa_thd_chunked_matches_single_chunk_bitwise(monkeypatch):
+    """Chunked (token, block) scoring must be bitwise identical to the un-chunked reference.
+
+    The score -> relu -> head-sum -> mask -> top-k pipeline is row-wise independent, so forcing a
+    single chunk (chunk >= T) reproduces the original full-einsum behaviour exactly; a small chunk
+    budget exercises the multi-chunk path. Both must return identical indices.
+    """
+    import mcore_bridge.model.modules.qsa_indexer as qi
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    torch.manual_seed(0)
+    idx, cfg = _make_idx(compress_ratio=4, budget=32, device=device)
+    d = cfg.indexer_head_dim
+    # each doc longer than block_topk blocks so the selection is genuinely sparse
+    per = 4 * (idx.block_topk + 4)
+    hidden_tok, freqs, cu_seqlens = _packed_inputs(idx, [per, per + 4, per + 8], cfg.hidden_size, d, device)
+
+    monkeypatch.setattr(qi, '_QSA_INDEX_SCORE_CHUNK_BYTES', 1 << 62)  # single chunk = un-chunked ref
+    one = idx.select_token_indices_thd(hidden_tok, freqs, cu_seqlens, force_materialize=True)
+    monkeypatch.setattr(qi, '_QSA_INDEX_SCORE_CHUNK_BYTES', 1024)  # tiny budget = many chunks
+    many = idx.select_token_indices_thd(hidden_tok, freqs, cu_seqlens, force_materialize=True)
+    assert torch.equal(one, many), 'chunked selection diverged from the un-chunked reference'
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='peak-memory regression needs a GPU')
+def test_qsa_thd_peak_memory_seq_shape_insensitive():
+    """Peak selection memory must not track the packed sequence shape at equal total tokens.
+
+    Pattern A = few long docs (large per-doc NB), pattern B = many short docs; both at the same
+    total token count. Before the chunked scoring the peak tracked [T, n_heads, NB] (i.e. the
+    s^2/compress_ratio term), so B (larger total NB) peaked far above A. With chunking the peak is
+    O(chunk * NB) and the two patterns must be close.
+    """
+    import mcore_bridge.model.modules.qsa_indexer as qi
+    torch.manual_seed(0)
+    idx, cfg = _make_idx(compress_ratio=4, budget=32, device='cuda')
+    d = cfg.indexer_head_dim
+    per = 4 * (idx.block_topk + 8)
+    n = 8
+    total = per * n
+    cases = {
+        'few_long': [total // 2, total // 2],  # 2 long docs
+        'many_short': [per] * n,  # n short docs, same total tokens
+    }
+    peaks = {}
+    for name, doc_lens in cases.items():
+        hidden_tok, freqs, cu_seqlens = _packed_inputs(idx, doc_lens, cfg.hidden_size, d, 'cuda')
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        idx.select_token_indices_thd(hidden_tok, freqs, cu_seqlens, force_materialize=True)
+        torch.cuda.synchronize()
+        peaks[name] = torch.cuda.max_memory_allocated()
+        del hidden_tok, freqs, cu_seqlens
+    ratio = peaks['many_short'] / max(1, peaks['few_long'])
+    assert ratio < 1.3, (f'peak memory tracks sequence shape (many_short/few_long={ratio:.2f}); peaks={peaks}; '
+                         'the (token, block) scoring is not chunked')
+
+
 def test_materialize_rope_preserves_mrope_batch_dimension():
     """MRoPE positions differ per sample, so ``freq_b`` must stay its own dim.