diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py index f739daa..85595b2 100644 --- a/src/mcore_bridge/model/modules/qsa_indexer.py +++ b/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] diff --git a/tests/test_qsa_indexer.py b/tests/test_qsa_indexer.py index 22f9b3c..419aa50 100644 --- a/tests/test_qsa_indexer.py +++ b/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.