diff --git a/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py b/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py index 6231c0b1..b1ef7899 100644 --- a/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py +++ b/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py @@ -1,42 +1,12 @@ -# VENDORED FILE -- DO NOT EDIT LOCALLY. -# -# refer : radixark/miles, miles_plugins/models/qwen3_8_next/ops/kernel/qsa_block_sparse_attn.py -# -# CONTRACT THIS KERNEL RELIES ON (verified against our indexer, see qsa_kernels.py): -# The selection rows must be whole, contiguous blocks of `block_size` tokens. The -# kernel tests membership at block granularity and restores exactness with the lo/hi -# range test, so a hand-built row like [0, 5] with block_size=4 is NOT computed as -# "attend {0,5}" -- it attends the whole 0..5 range (measured: 16.62 vs 50.0 on a -# probe where only the selected keys should contribute). Our QSAIndexer satisfies -# this by construction (`top_blocks * R + arange(R)`); there is no runtime check, so -# any new caller must uphold it by construction too. +# Vendored from radixark/miles: miles/kernels/attention/qsa/block_sparse_attention.py (24bd5fb716). +# Local change only: int64 tile addressing in the kernels (256K bitmaps exceed 2 GiB). """Tensor-core QSA sparse attention for training: forward + backward. Same semantics as ``qsa_sparse_attn.py`` -- each query attends exactly the tokens in -its selection row -- but reached a different way, because the gather-per-query form -cannot use ``tl.dot``: with a distinct key set per query it has to materialise a -``[BQ, BK, D]`` tile and reduce it with ALU math. Measured at the production shape -(T=25k, 12 q-heads, D=256, budget 2048) that costs 2.5 s forward and 13.9 s -forward+backward per layer, which is ~330x slower than a DENSE causal flash kernel over -the same tensors (9.4 ms / 39.2 ms) even though dense does 6x more FLOPs. QSA at 25k -tokens only removes ~6x of the work, so paying 300x for the privilege is a large loss: -on the 16-node agentic run those three QSA layers per pipeline stage were ~49 s of a -~52 s micro-batch. - -So this kernel walks key tiles with ``tl.dot`` and masks each (query, key) pair to the -query's own selection. Tiles nobody in the query tile selected are skipped entirely via -a CSR-style per-query-tile list, so it beats a dense sweep by the coverage ratio rather -than merely matching it. The result is exact, not an approximation: - -- selection membership is tested at BLOCK granularity (the indexer picks blocks of - ``BLK`` consecutive tokens), which is a superset of the row's token list, because - expanding a block drops only the tokens past the query's own position; -- the dropped ones come back out via ``lo``/``hi``, the inclusive key range the caller - already computes (sequence start .. query position). - -bf16 inputs feed the dots with fp32 accumulation, the flash-attention convention, which -is also what the sglang kernel this mirrors does -- the old kernel's fp32 ALU path was -the odd one out. +its selection row -- but it walks key tiles with ``tl.dot`` and masks each (query, key) +pair to the query's selection, skipping tiles no query in the tile selected. The result +is exact: block-granular membership is a superset of the row, and ``lo``/``hi`` (the +inclusive key range, sequence start .. query position) removes the extra tokens. """ import torch @@ -78,7 +48,6 @@ def _qsa_bs_fwd_kernel( BK: tl.constexpr, BLK: tl.constexpr, ): - # The selection bitmap exceeds 2 GiB at long context lengths. pid_t = tl.program_id(0).to(tl.int64) pid_h = tl.program_id(1) kv_head = pid_h // GROUP @@ -101,20 +70,15 @@ def _qsa_bs_fwd_kernel( l_i = tl.zeros((BQ, ), tl.float32) acc = tl.zeros((BQ, D), tl.float32) - # Only the key tiles this query tile actually selected, so the kernel beats a dense - # sweep by the coverage ratio instead of merely matching it. The list is the union - # over the tile's queries; per-query exactness comes from the mask below. + # the tile list is the union over the tile's queries; per-query exactness comes from the mask n_tiles = tl.load(KCNT + pid_t) for i in range(0, n_tiles): kt = tl.load(KLIST + pid_t * stride_kl + i).to(tl.int64) offs_k = kt * BK + tl.arange(0, BK) k_in = offs_k < T - # per-sequence block grid: after the packed-indexer fix a sequence's blocks start - # at its own first token, which is not a multiple of BLK in a packed batch. - # Looked up per TOKEN, not per 4-token group: a sequence whose start is not a - # multiple of BLK has its per-sequence blocks straddling the global groups, so a - # group-granular lookup would assign some tokens to the wrong block. + # per-token block lookup: a packed sequence's blocks start at its own first token, + # which need not be a multiple of BLK blk = blk_base[:, None] + (offs_k[None, :] - tok_base[:, None]) // BLK sel = tl.load( SEL + offs_q[:, None] * stride_st + blk, @@ -222,11 +186,7 @@ def _qsa_bs_dq_kernel( offs_k = kt * BK + tl.arange(0, BK) k_in = offs_k < T - # per-sequence block grid: after the packed-indexer fix a sequence's blocks start - # at its own first token, which is not a multiple of BLK in a packed batch. - # Looked up per TOKEN, not per 4-token group: a sequence whose start is not a - # multiple of BLK has its per-sequence blocks straddling the global groups, so a - # group-granular lookup would assign some tokens to the wrong block. + # per-token block lookup, as in the forward kernel blk = blk_base[:, None] + (offs_k[None, :] - tok_base[:, None]) // BLK sel = tl.load( SEL + offs_q[:, None] * stride_st + blk, @@ -292,12 +252,10 @@ def _qsa_bs_dkdv_kernel( BK: tl.constexpr, BLK: tl.constexpr, ): - """dK/dV keyed on the KEY tile, so nothing needs atomics. + """dK/dV keyed on the key tile, so nothing needs atomics. - Launched per KV head, not per query head: with GQA the ``GROUP`` query heads sharing a - KV head all contribute to the same dK/dV, so they have to be summed here. Writing them - from separate programs would have them overwrite each other (the gather kernel got away - with it only because it used atomic_add). + Launched per KV head: the ``GROUP`` query heads sharing a KV head are summed here, + since separate programs would overwrite each other's dK/dV. """ pid_k = tl.program_id(0).to(tl.int64) kv_head = tl.program_id(1) @@ -382,71 +340,70 @@ def selection_to_block_bitmap(indices: Tensor, num_tokens: int, block_size: int) return flags -def build_tile_index(sel: Tensor, bq: int, bk: int, block_size: int) -> tuple[Tensor, Tensor]: +def _tile_mask(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int) -> Tensor: + """``[NQT, NKT]`` bool: some query in query tile ``i`` selected a block with a token in key tile ``j``. + + Key tiles are absolute token ranges ``[j * bk, (j + 1) * bk)``, as the kernels walk them. + ``block_first`` / ``block_last`` ``[NB]`` are each block's inclusive token span; packed + block ids drift from ``token // block_size`` once a sequence length is off the grid. + """ + T, nb = sel.shape + nqt = -(-T // bq) + nkt = -(-T // bk) + pad_q = nqt * bq - T + if pad_q: + sel = torch.nn.functional.pad(sel, (0, 0, 0, pad_q)) + picked = sel.view(nqt, bq, nb).amax(dim=1).to(torch.int32) + tile = torch.zeros(nqt, nkt, dtype=torch.int32, device=sel.device) + # block_size <= bk, so a block touches at most the tiles of its first and last token + tile.index_add_(1, block_first.long() // bk, picked) + tile.index_add_(1, block_last.long() // bk, picked) + return tile > 0 + + +def _compact_tile_mask(mat: Tensor) -> tuple[Tensor, Tensor]: + """``[R, C]`` bool -> (ascending column ids per row ``[R, maxc]`` int32, counts ``[R]`` int32).""" + cnt = mat.sum(dim=1).to(torch.int32) + maxc = max(int(cnt.max().item()), 1) + order = torch.argsort((~mat).to(torch.int8), dim=1, stable=True) + return order[:, :maxc].contiguous().to(torch.int32), cnt + + +def build_tile_index(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int) -> tuple[Tensor, Tensor]: """``[T, NB]`` block flags -> (``klist`` [NQT, maxc] int32, ``kcnt`` [NQT] int32). ``klist[i]`` lists, ascending, the key tiles that at least one query in query-tile ``i`` selected. Per-query exactness still comes from the in-kernel mask; this only decides which tiles are worth visiting. """ - T, nb = sel.shape - bpt = bk // block_size - nqt = -(-T // bq) - nkt = -(-nb // bpt) - pad_b = nkt * bpt - nb - pad_q = nqt * bq - T - if pad_b or pad_q: - sel = torch.nn.functional.pad(sel, (0, pad_b, 0, pad_q)) - tile = sel.view(nqt, bq, nkt, bpt).amax(dim=3).amax(dim=1) > 0 # [nqt, nkt] - kcnt = tile.sum(dim=1).to(torch.int32) - maxc = max(int(kcnt.max().item()), 1) - order = torch.argsort((~tile).to(torch.int8), dim=1, stable=True) - klist = order[:, :maxc].contiguous().to(torch.int32) - return klist, kcnt + return _compact_tile_mask(_tile_mask(sel, block_first, block_last, bq, bk)) -def build_tile_index_pair(sel: Tensor, bq: int, bk: int, block_size: int): +def build_tile_index_pair(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int): """Both directions of the tile map: (klist, kcnt) per query tile, (qlist, qcnt) per key tile. The transposed half is what lets dK/dV be keyed on the key tile and so avoid atomics. """ - T, nb = sel.shape - bpt = bk // block_size - nqt = -(-T // bq) - nkt = -(-nb // bpt) - pad_b = nkt * bpt - nb - pad_q = nqt * bq - T - padded = sel - if pad_b or pad_q: - padded = torch.nn.functional.pad(sel, (0, pad_b, 0, pad_q)) - tile = padded.view(nqt, bq, nkt, bpt).amax(dim=3).amax(dim=1) > 0 - - def compact(mat): - cnt = mat.sum(dim=1).to(torch.int32) - maxc = max(int(cnt.max().item()), 1) - order = torch.argsort((~mat).to(torch.int8), dim=1, stable=True) - return order[:, :maxc].contiguous().to(torch.int32), cnt - - klist, kcnt = compact(tile) - qlist, qcnt = compact(tile.t().contiguous()) + tile = _tile_mask(sel, block_first, block_last, bq, bk) + klist, kcnt = _compact_tile_mask(tile) + qlist, qcnt = _compact_tile_mask(tile.t().contiguous()) return klist, kcnt, qlist, qcnt class _QSABlockSparseAttn(torch.autograd.Function): @staticmethod - def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, scale, block_size): + def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, block_first, block_last, scale, block_size): T, Hq, D = q.shape S, Hkv, _ = k.shape assert Hq % Hkv == 0 group = Hq // Hkv qc, kc, vc = q.contiguous(), k.contiguous(), v.contiguous() selc = sel.contiguous() - BQ_, BK_ = 64, 64 - klist, kcnt = build_tile_index(selc, BQ_, BK_, block_size) + BQ, BK = 64, 64 + klist, kcnt = build_tile_index(selc, block_first, block_last, BQ, BK) o = torch.empty(T, Hq, D, device=q.device, dtype=torch.float32) lse = torch.empty(Hq, T, device=q.device, dtype=torch.float32) - BQ, BK = 64, 64 grid = (triton.cdiv(T, BQ), Hq) _qsa_bs_fwd_kernel[grid]( qc, @@ -482,7 +439,7 @@ def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, scale, block_size): num_warps=8, num_stages=2, ) - ctx.save_for_backward(qc, kc, vc, selc, lo, hi, blk_base, tok_base, klist, kcnt, o, lse) + ctx.save_for_backward(qc, kc, vc, selc, lo, hi, blk_base, tok_base, block_first, block_last, o, lse) ctx.scale = scale ctx.group = group ctx.block_size = block_size @@ -490,7 +447,7 @@ def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, scale, block_size): @staticmethod def backward(ctx, grad_out): - qc, kc, vc, selc, lo, hi, blk_base, tok_base, _klist_fwd, _kcnt_fwd, o, lse = ctx.saved_tensors + qc, kc, vc, selc, lo, hi, blk_base, tok_base, block_first, block_last, o, lse = ctx.saved_tensors T, Hq, D = qc.shape do = grad_out.contiguous().to(qc.dtype) # delta once, in torch: both backward kernels need it and neither should redo it @@ -500,7 +457,7 @@ def backward(ctx, grad_out): dv = torch.zeros(vc.shape, device=vc.device, dtype=torch.float32) BQ, BK = 64, 32 - klist, kcnt, qlist, qcnt = build_tile_index_pair(selc, BQ, BK, ctx.block_size) + klist, kcnt, qlist, qcnt = build_tile_index_pair(selc, block_first, block_last, BQ, BK) common = ( qc.stride(0), qc.stride(1), @@ -572,10 +529,10 @@ def backward(ctx, grad_out): num_warps=8, num_stages=1, ) - return dq.to(qc.dtype), dk.to(kc.dtype), dv.to(vc.dtype), None, None, None, None, None, None, None + return dq.to(qc.dtype), dk.to(kc.dtype), dv.to(vc.dtype), None, None, None, None, None, None, None, None, None -def qsa_block_sparse_attention_triton( +def qsa_block_sparse_attention( q: Tensor, k: Tensor, v: Tensor, @@ -584,15 +541,19 @@ def qsa_block_sparse_attention_triton( hi: Tensor, blk_base: Tensor, tok_base: Tensor, + block_first: Tensor, + block_last: Tensor, scale: float, block_size: int = 4, ) -> Tensor: """``q`` [T, Hq, D], ``k``/``v`` [S, Hkv, D], ``sel_blocks`` [T, NB] uint8. ``lo``/``hi`` are the inclusive key range per query; ``blk_base``/``tok_base`` place - the query's sequence in the packed block grid (both zero for a single sequence). + the query's sequence in the packed block grid (both zero for a single sequence); + ``block_first``/``block_last`` [NB] are each block's inclusive token span. """ - return _QSABlockSparseAttn.apply(q, k, v, sel_blocks, lo, hi, blk_base, tok_base, scale, block_size) + return _QSABlockSparseAttn.apply(q, k, v, sel_blocks, lo, hi, blk_base, tok_base, block_first, block_last, scale, + block_size) def qsa_sparse_attention_from_indices(q: Tensor, @@ -609,4 +570,6 @@ def qsa_sparse_attention_from_indices(q: Tensor, lo = torch.where(valid, indices, torch.full_like(indices, big)).min(dim=1).values.to(torch.int32) hi = torch.where(valid, indices, torch.full_like(indices, -1)).max(dim=1).values.to(torch.int32) zeros = torch.zeros(T, dtype=torch.int32, device=q.device) - return qsa_block_sparse_attention_triton(q, k, v, sel, lo, hi, zeros, zeros, scale, block_size) + block_first = torch.arange(sel.shape[1], dtype=torch.int32, device=q.device) * block_size + block_last = (block_first + block_size).clamp_max(T) - 1 + return qsa_block_sparse_attention(q, k, v, sel, lo, hi, zeros, zeros, block_first, block_last, scale, block_size) diff --git a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py index fa57d751..543ca59a 100644 --- a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py +++ b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py @@ -1,9 +1,14 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """QSA sparse attention wrappers around the vendored tensor-core triton kernel. -The kernel itself lives in ``qsa_block_sparse_attn.py``, adapted from -miles PR #2777 (commit 0f5dff4). This file owns only the glue mcore needs: -sbhd<->thd flattening, context parallelism, and the ``core_attention`` shim. +The kernel itself lives in ``qsa_block_sparse_attn.py`` (vendored from miles). This +file owns the glue: the packed block form, sbhd<->thd flattening, context +parallelism, and the ``core_attention`` shim. + +The indexer selects whole blocks on each document's own grid (``doc_start + m * R``), +which need not sit on the global ``token // R`` grid. The packed path therefore builds +``blk_base``/``tok_base``/``block_first``/``block_last`` from ``cu_seqlens`` in +``_PackedBlockLayout``; ``qsa_sparse_attention_from_indices`` is global-grid only. WHY A KERNEL AT ALL (the decision this file encodes) QSA needs a per-query key set, which as an attention mask is TE's @@ -28,7 +33,7 @@ from mcore_bridge.utils import get_env_args, get_logger -from .qsa_block_sparse_attn import qsa_sparse_attention_from_indices +from .qsa_block_sparse_attn import qsa_block_sparse_attention logger = get_logger() @@ -109,16 +114,93 @@ def _cp_gathered_to_logical_order(seq_len: int, cp_size: int, device) -> torch.T return torch.cat([torch.arange(c * chunk, (c + 1) * chunk, device=device) for c in order]) -def qsa_sparse_attention_thd(q, k, v, indices, scale, block_size): - """``q`` [T, Hq, D], ``k``/``v`` [S, Hkv, D], ``indices`` [T, K] (-1 pad). +def _align_cu_seqlens(cu, total: int): + """Return ``cu`` ending at ``total``, accepting a missing leading 0 or a batch + offset; ``None`` when it cannot describe ``total`` tokens.""" + cu = cu.reshape(-1).to(dtype=torch.long) + if cu.numel() > 0 and int(cu[0]) != 0: + if int(cu[-1]) == total: + cu = torch.cat([cu.new_zeros(1), cu]) + elif int(cu[-1]) - int(cu[0]) == total: + cu = cu - cu[0] + if not (cu.numel() > 1 and int(cu[-1]) == total): + return None + return cu + + +def _packed_cu_seqlens(packed_seq_params, total: int): + """Resolve the ``cu_seqlens`` matching ``total`` packed tokens; padded first (it is + the pack space CP gathers into).""" + if packed_seq_params is None: + raise ValueError('QSA packed attention needs packed_seq_params.cu_seqlens_q to place every ' + f'sequence on its own block grid, but none was passed for {total} tokens.') + for name in ('cu_seqlens_q_padded', 'cu_seqlens_q'): + cu = getattr(packed_seq_params, name, None) + if cu is None: + continue + aligned = _align_cu_seqlens(cu, total) + if aligned is not None: + return aligned + raise ValueError(f'QSA cannot align cu_seqlens with the {total} tokens it is about to attend: ' + f'cu_seqlens_q={getattr(packed_seq_params, "cu_seqlens_q", None)}, ' + f'cu_seqlens_q_padded={getattr(packed_seq_params, "cu_seqlens_q_padded", None)}.') + + +class _PackedBlockLayout: + """Per-sequence block grid: block ``m`` covers ``[start + m*R, start + (m+1)*R)`` + and ids form one global numbering of size ``sum(ceil(len_d / R))``; the exact + per-block token spans feed the kernel's tile map. + """ - Token-space indices: works for packed (thd) inputs directly, and for any - pre-flattened token dimension. + def __init__(self, cu_seqlens, compress_ratio: int, device): + cu = cu_seqlens.to(device=device, dtype=torch.long) + lengths = cu[1:] - cu[:-1] + if lengths.numel() == 0 or bool((lengths < 0).any()): + raise ValueError(f'QSA block layout needs non-negative sequence lengths, got cu={cu.tolist()}') + blocks_per_seq = -(-lengths // compress_ratio) + self.num_blocks = int(blocks_per_seq.sum().item()) + if self.num_blocks == 0: + raise ValueError('QSA block layout needs at least one non-empty sequence.') + seq_block_start = torch.cat([blocks_per_seq.new_zeros(1), blocks_per_seq.cumsum(0)[:-1]]) + seq_token_start = cu[:-1] + # per token: its sequence start and that sequence's first block id + seg = torch.repeat_interleave(torch.arange(lengths.numel(), device=device), lengths) + self.token_start = seq_token_start.index_select(0, seg) + self.token_block_start = seq_block_start.index_select(0, seg) + # per block: exact inclusive token span + block_seq = torch.repeat_interleave(torch.arange(lengths.numel(), device=device), blocks_per_seq) + block_local = torch.arange(self.num_blocks, device=device) - seq_block_start.index_select(0, block_seq) + self.block_first_token = seq_token_start.index_select(0, block_seq) + block_local * compress_ratio + self.block_last_token = torch.minimum(self.block_first_token + compress_ratio, cu[1:].index_select( + 0, block_seq)) - 1 + + +def _block_form(indices, layout: _PackedBlockLayout, block_size: int): + """Rows -> ``(sel, lo, hi, blk_base, tok_base, block_first, block_last)``: a + token's block id is ``blk_base + (token - tok_base) // R`` and ``lo``/``hi`` is + ``[sequence start, query position]`` (block membership supplies the holes).""" + device = indices.device + total = indices.shape[0] + tok_base = layout.token_start.to(torch.int32) + blk_base = layout.token_block_start.to(torch.int32) + sel = torch.zeros(total, layout.num_blocks, dtype=torch.uint8, device=device) + valid = indices >= 0 + if bool(valid.any()): + rows = torch.arange(total, device=device).unsqueeze(1).expand_as(indices) + blk = blk_base.unsqueeze(1).long() + (indices - tok_base.unsqueeze(1)).div(block_size, rounding_mode='floor') + sel[rows[valid], blk[valid].long().clamp_(0, layout.num_blocks - 1)] = 1 + lo = tok_base + hi = torch.arange(total, dtype=torch.int32, device=device) + return (sel, lo, hi, blk_base, tok_base, layout.block_first_token.to(torch.int32), + layout.block_last_token.to(torch.int32)) + + +def qsa_sparse_attention_thd(q, k, v, indices, cu_seqlens, scale, block_size): + """``q`` [T, Hq, D], ``k``/``v`` [S, Hkv, D], ``indices`` [T, K] (-1 pad). - ``block_size`` must be the indexer's ``compress_ratio``. The kernel tests - selection membership per block of that many consecutive tokens, so passing a - different value silently changes which keys are attended -- see the contract - note in qsa_block_sparse_attn.py. + ``cu_seqlens`` gives every sequence's start (``_PackedBlockLayout``). ``block_size`` + must be the indexer's ``compress_ratio``: a different value silently changes which + keys are attended. """ if not HAVE_TRITON or q.device.type == 'cpu': raise RuntimeError('QSA sparse attention requires triton and accelerator tensors ' @@ -132,41 +214,42 @@ def qsa_sparse_attention_thd(q, k, v, indices, scale, block_size): # len(q). Callers must equalise first -- _forward_cp does this by scattering # the local query shard into a full-length buffer. raise ValueError(f'QSA sparse attention needs len(q) == len(k), got {q.shape[0]} vs {k.shape[0]}.') - return qsa_sparse_attention_from_indices(q, k, v, indices.contiguous(), scale, block_size) + cu = _align_cu_seqlens(cu_seqlens, q.shape[0]) if cu_seqlens is not None else None + if cu is None: + raise ValueError(f'QSA packed attention needs cu_seqlens whose last entry is {q.shape[0]}, ' + f'got {None if cu_seqlens is None else cu_seqlens.tolist()}.') + layout = _PackedBlockLayout(cu, block_size, q.device) + sel, lo, hi, blk_base, tok_base, block_first, block_last = _block_form(indices.contiguous(), layout, block_size) + return qsa_block_sparse_attention(q, k, v, sel, lo, hi, blk_base, tok_base, block_first, block_last, scale, + block_size) -def qsa_sparse_attention(q, k, v, indices, scale, block_size): +def qsa_sparse_attention(q, k, v, indices, scale, block_size, cu_seqlens=None): """QSA sparse attention over mcore layouts. sbhd: ``q`` [s, b, Hq, D], ``k``/``v`` [s, b, Hkv, D], ``indices`` - [b, s, K] in per-sample sequence space. + [b, s, K] in per-sample sequence space (each sample is one sequence). thd: ``q`` [T, Hq, D], ``k``/``v`` [T, Hkv, D], ``indices`` [T, K] in - pack space (document boundaries already clamped torch-side). + pack space, with ``cu_seqlens`` giving the document boundaries. Raises rather than returning ``None``: the caller only reaches here for packing or CP>1, and for those a dense fallback would silently diverge from sparse inference. """ if q.dim() == 3: - return qsa_sparse_attention_thd(q, k, v, indices, scale, block_size) + return qsa_sparse_attention_thd(q, k, v, indices, cu_seqlens, scale, block_size) if q.dim() != 4: raise ValueError(f'qsa_sparse_attention expected 3D (thd) or 4D (sbhd) q, got {tuple(q.shape)}') s, b, hq, d = q.shape sk = k.shape[0] - # Flatten to token space, batch-major: token t = r * sk + p. The per-sample - # sequence-space indices are offset by the batch start; -1 padding stays -1. - # The offset is a whole multiple of sk, so block alignment survives it only - # when sk % block_size == 0; guard rather than corrupt the selection. - if sk % block_size: - raise ValueError(f'sbhd QSA needs the kv sequence length ({sk}) to be a multiple of ' - f'block_size ({block_size}); otherwise flattening to token space shifts ' - 'each sample off the block grid the kernel indexes by.') q_f = q.permute(1, 0, 2, 3).reshape(b * s, hq, d) k_f = k.permute(1, 0, 2, 3).reshape(b * sk, *k.shape[2:]) v_f = v.permute(1, 0, 2, 3).reshape(b * sk, *v.shape[2:]) + # batch-major flattening: token t = r * sk + p, one sequence per sample off = torch.arange(b, device=q.device).view(b, 1, 1) * sk idx_f = torch.where(indices >= 0, indices + off, indices.new_full((), -1)).reshape(b * s, -1) - out_f = qsa_sparse_attention_thd(q_f, k_f, v_f, idx_f, scale, block_size) + cu = torch.arange(b + 1, device=q.device, dtype=torch.int64) * sk + out_f = qsa_sparse_attention_thd(q_f, k_f, v_f, idx_f, cu, scale, block_size) return out_f.view(b, s, hq, d).permute(1, 0, 2, 3) @@ -215,10 +298,13 @@ def forward(self, if attention_mask is not None and attention_mask.dtype in (torch.int32, torch.int64): scale = self.softmax_scale if self.softmax_scale is not None else query.shape[-1]**-0.5 cp_size = self.config.context_parallel_size + # Under CP the local query is only a shard, so _forward_cp resolves its own cu. + cu = _packed_cu_seqlens(packed_seq_params, query.shape[0]) \ + if query.dim() == 3 and cp_size == 1 else None if cp_size > 1: - out = self._forward_cp(query, key, value, attention_mask, scale, packed_seq_params) + out = self._forward_cp(query, key, value, attention_mask, scale, packed_seq_params, cu) else: - out = qsa_sparse_attention(query, key, value, attention_mask, scale, self.block_size) + out = qsa_sparse_attention(query, key, value, attention_mask, scale, self.block_size, cu_seqlens=cu) # mcore's sbhd core attention returns heads flattened ([s, b, h*d]); # the kernel yields [s, b, h, d]. thd stays [t, h, d] -- mcore # reshapes it itself right after (see attention.py thd branch). @@ -235,7 +321,7 @@ def forward(self, packed_seq_params=packed_seq_params, **kwargs) - def _forward_cp(self, query, key, value, indices, scale, packed_seq_params): + def _forward_cp(self, query, key, value, indices, scale, packed_seq_params, cu_q): from megatron.core import mpu from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region cp_size = mpu.get_context_parallel_world_size() @@ -245,9 +331,9 @@ def _forward_cp(self, query, key, value, indices, scale, packed_seq_params): if thd: # the gathered k/v live in the padded pack space, so the query # positions must come from the padded cu as well - cu_q = ( - packed_seq_params.cu_seqlens_q_padded - if packed_seq_params.cu_seqlens_q_padded is not None else packed_seq_params.cu_seqlens_q) + if cu_q is None: + # the gathered length is the local shard times the group size + cu_q = _packed_cu_seqlens(packed_seq_params, key.shape[0] * cp_size) q_pos = _cp_query_global_positions_thd(cu_q, cp_size, cp_rank, device) # rank-major gathered positions -> permutation back to global packed # order (same construction as mcore DSA's packed kv reorder) @@ -270,6 +356,9 @@ def _gather_full(t): key_full = _gather_full(key) value_full = _gather_full(value) + if thd and int(cu_q[-1]) != key_full.shape[0]: + raise ValueError(f'QSA thd CP: cu_seqlens ends at {int(cu_q[-1])} but the gathered keys have ' + f'{key_full.shape[0]} tokens; packed_seq_params do not match this shard.') # The kernel derives the key bound from the query count (T, Hq, D = q.shape, # then `offs_k < T`), so it structurally requires len(q) == len(k). Under CP # the queries are a 1/cp_size shard while k/v are now full length, so scatter @@ -278,8 +367,11 @@ def _gather_full(t): # skips, so they cost tile launches but produce nothing. if thd: local_idx = indices[q_pos] - out_full = qsa_sparse_attention(*self._scatter_q_to_full(query, key_full, value_full, local_idx, q_pos), - scale, self.block_size) + out_full = qsa_sparse_attention( + *self._scatter_q_to_full(query, key_full, value_full, local_idx, q_pos), + scale, + self.block_size, + cu_seqlens=cu_q) return out_full.index_select(0, q_pos) # sbhd: token-space kernel on the batch-major flattening (t = r*sk + p) local_idx = indices[:, q_pos] @@ -291,7 +383,10 @@ def _gather_full(t): q_f = query.permute(1, 0, 2, 3).reshape(sq * b, query.shape[2], query.shape[3]) # batch-major token ids of this rank's rows: sample r contributes q_pos + r*sk rows = (q_pos[None, :] + torch.arange(b, device=device).view(b, 1) * sk).reshape(-1) - out_f = qsa_sparse_attention(*self._scatter_q_to_full(q_f, k_f, v_f, idx_f, rows), scale, self.block_size) + # each sample is one sequence of length sk in the flattened buffer + cu = torch.arange(b + 1, device=device, dtype=torch.int64) * sk + out_f = qsa_sparse_attention( + *self._scatter_q_to_full(q_f, k_f, v_f, idx_f, rows), scale, self.block_size, cu_seqlens=cu) out_f = out_f.index_select(0, rows) return out_f.view(b, sq, query.shape[2], query.shape[3]).permute(1, 0, 2, 3) diff --git a/tests/test_qsa_sparse_kernels.py b/tests/test_qsa_sparse_kernels.py index f0513013..9bf89a47 100644 --- a/tests/test_qsa_sparse_kernels.py +++ b/tests/test_qsa_sparse_kernels.py @@ -1,19 +1,22 @@ """QSA sparse attention numerical and long-context address regressions.""" -import importlib.util import os import pytest import torch -from pathlib import Path pytest.importorskip('triton') -_spec = importlib.util.spec_from_file_location( - 'qsa_sparse_under_test', - Path(__file__).parents[1] / 'src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py') -kernels = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(kernels) + +from mcore_bridge.model.modules.kernels import qsa_kernels as qk # noqa: E402 +from mcore_bridge.model.modules.kernels.qsa_block_sparse_attn import qsa_block_sparse_attention # noqa: E402 + pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason='QSA kernels require CUDA') +def _block_spans(num_blocks, block_size, total, device='cuda'): + first = torch.arange(num_blocks, dtype=torch.int32, device=device) * block_size + last = (first + block_size).clamp_max(total) - 1 + return first, last + + @pytest.mark.parametrize('packed', [False, True]) def test_qsa_forward_backward_match_dense_attention(packed): torch.manual_seed(123) @@ -28,7 +31,8 @@ def test_qsa_forward_backward_match_dense_attention(packed): lo = torch.where(pos >= 19, 19, 0).int() if packed else torch.zeros_like(pos, dtype=torch.int32) hi = pos.int() zero = torch.zeros_like(lo) - actual = kernels.qsa_block_sparse_attention_triton(q, k, v, sel, lo, hi, zero, zero, d**-0.5, 4) + first, last = _block_spans(sel.shape[1], 4, t) + actual = qsa_block_sparse_attention(q, k, v, sel, lo, hi, zero, zero, first, last, d**-0.5, 4) allowed = sel[:, pos // 4].bool() & (pos[None, :] >= lo[:, None]) & (pos[None, :] <= hi[:, None]) scores = torch.einsum('thd,shd->hts', q.float(), k.expand(-1, h, -1).float()) * d**-0.5 probs = scores.masked_fill(~allowed[None], float('-inf')).softmax(-1) @@ -41,6 +45,134 @@ def test_qsa_forward_backward_match_dense_attention(packed): torch.testing.assert_close(a, e, rtol=0.04, atol=0.015) +def _packed_indexer_rows(lengths, ratio, topk=2): + """Indexer-shaped rows on packed documents: whole document-relative blocks with + holes (0, 2, 4, ...) plus the query's causal tail.""" + cu = torch.tensor([0, *lengths]).cumsum(0) + total = int(cu[-1]) + starts = cu[:-1] + seg = torch.zeros(total, dtype=torch.long) + seg[starts[1:]] = 1 + seg = seg.cumsum(0) + positions = torch.arange(total) - starts[seg] + rows = [] + for t in range(total): + doc_start = int(starts[seg[t]]) + p = int(positions[t]) + visible = (p + 1) // ratio + blocks = list(range(0, visible, 2))[:topk] + toks = [doc_start + b * ratio + j for b in blocks for j in range(ratio)] + qn = (p + 1) // ratio + toks += list(range(doc_start + qn * ratio, min(doc_start + qn * ratio + ratio, t + 1))) + rows.append(sorted({x for x in toks if doc_start <= x <= t})) + idx = torch.full((total, max(len(r) for r in rows)), -1, dtype=torch.int64, device='cuda') + for t, row in enumerate(rows): + idx[t, :len(row)] = torch.tensor(row, device='cuda') + return idx, cu.int().cuda() + + +def _dense_reference(q, k, v, idx, scale): + total, hq, _ = q.shape + rep = hq // k.shape[1] + allowed = torch.zeros(total, total, dtype=torch.bool, device=q.device) + for t, row in enumerate(idx): + allowed[t, row[row >= 0]] = True + kk = k.repeat_interleave(rep, dim=1).float() + vv = v.repeat_interleave(rep, dim=1).float() + scores = torch.einsum('thd,shd->hts', q.float(), kk) * scale + probs = scores.masked_fill(~allowed[None], float('-inf')).softmax(-1) + return torch.einsum('hts,shd->thd', probs, vv).to(q.dtype) + + +@pytest.mark.parametrize('lengths,ratio', [ + ([3, 20], 4), + ([61, 128], 4), + ([1, 3, 20, 37, 9], 8), + ([5, 61, 6], 3), + ([5, 61, 6], 5), +]) +def test_qsa_off_grid_packed_matches_dense(lengths, ratio): + """Off-grid packed documents must attend exactly their selection (fwd + bwd).""" + torch.manual_seed(7) + idx, cu = _packed_indexer_rows(lengths, ratio) + total, d, hq = idx.shape[0], 32, 2 + hkv = 1 + scale = d**-0.5 + q0 = (torch.randn(total, hq, d, device='cuda', dtype=torch.bfloat16) * 0.2) + k0 = (torch.randn(total, hkv, d, device='cuda', dtype=torch.bfloat16) * 0.2) + v0 = (torch.randn(total, hkv, d, device='cuda', dtype=torch.bfloat16) * 0.2) + go = torch.randn(total, hq, d, device='cuda', dtype=torch.bfloat16) + + q, k, v = q0.clone().requires_grad_(), k0.clone().requires_grad_(), v0.clone().requires_grad_() + out = qk.qsa_sparse_attention_thd(q, k, v, idx, cu, scale, ratio) + grads = torch.autograd.grad(out, (q, k, v), go) + + q, k, v = q0.clone().requires_grad_(), k0.clone().requires_grad_(), v0.clone().requires_grad_() + ref = _dense_reference(q, k, v, idx, scale) + ref_grads = torch.autograd.grad(ref.float(), (q, k, v), go.float()) + + torch.testing.assert_close(out.float(), ref.float(), rtol=0.03, atol=0.01) + for a, e in zip(grads, ref_grads): + torch.testing.assert_close(a.float(), e.float(), rtol=0.04, atol=0.02) + + +def test_qsa_sbhd_flattening_keeps_each_sample_on_its_own_grid(): + """sbhd flattening treats each sample as its own sequence.""" + torch.manual_seed(11) + s, b, d, hq, hkv, ratio = 20, 2, 32, 2, 1, 4 + # per-sample rows on the sample's own grid, with holes + idx = torch.full((b, s, 16), -1, dtype=torch.int64, device='cuda') + for sample in range(b): + for p in range(s): + visible = (p + 1) // ratio + blocks = list(range(0, visible, 2))[:2] + toks = [blk * ratio + j for blk in blocks for j in range(ratio)] + qn = (p + 1) // ratio + toks += list(range(qn * ratio, min(qn * ratio + ratio, p + 1))) + row = sorted({x for x in toks if x <= p}) + idx[sample, p, :len(row)] = torch.tensor(row, device='cuda') + + q = torch.randn(s, b, hq, d, device='cuda', dtype=torch.bfloat16) + k = torch.randn(s, b, hkv, d, device='cuda', dtype=torch.bfloat16) + v = torch.randn(s, b, hkv, d, device='cuda', dtype=torch.bfloat16) + out = qk.qsa_sparse_attention(q, k, v, idx, d**-0.5, ratio) + + idx_f = torch.where(idx >= 0, idx + torch.arange(b, device='cuda').view(b, 1, 1) * s, idx.new_full( + (), -1)).reshape(b * s, -1) + q_f = q.permute(1, 0, 2, 3).reshape(b * s, hq, d) + k_f = k.permute(1, 0, 2, 3).reshape(b * s, hkv, d) + v_f = v.permute(1, 0, 2, 3).reshape(b * s, hkv, d) + ref = _dense_reference(q_f, k_f, v_f, idx_f, d**-0.5) + out_f = out.permute(1, 0, 2, 3).reshape(b * s, hq, d) + torch.testing.assert_close(out_f.float(), ref.float(), rtol=0.03, atol=0.01) + + +def test_qsa_core_attention_shim_resolves_packed_cu(): + """forward must resolve cu_seqlens from packed_seq_params for packed inputs.""" + from mcore_bridge.model.modules.kernels import QSASparseCoreAttention + + class _Config: + + context_parallel_size = 1 + indexer_compress_ratio = 4 + + class _PackedSeqParams: + + cu_seqlens_q = torch.tensor([0, 3, 23], dtype=torch.int32) + cu_seqlens_q_padded = None + + layer = QSASparseCoreAttention(core_attention=None, config=_Config()) + idx, _ = _packed_indexer_rows([3, 20], 4) + total, d, hq, hkv = idx.shape[0], 32, 2, 1 + torch.manual_seed(3) + q = torch.randn(total, hq, d, device='cuda', dtype=torch.bfloat16) + k = torch.randn(total, hkv, d, device='cuda', dtype=torch.bfloat16) + v = torch.randn(total, hkv, d, device='cuda', dtype=torch.bfloat16) + out = layer(q, k, v, idx, packed_seq_params=_PackedSeqParams()) + ref = _dense_reference(q, k, v, idx, d**-0.5) + torch.testing.assert_close(out.float(), ref.float(), rtol=0.03, atol=0.01) + + @pytest.mark.skipif( os.environ.get('MCORE_BRIDGE_TEST_LONG_QSA') != '1', reason='Set MCORE_BRIDGE_TEST_LONG_QSA=1 for the 16 GiB bitmap regression') @@ -58,7 +190,8 @@ def test_qsa_256k_bitmap_forward_backward_cross_int32_boundary(): sel[:, 0] = 1 lo = torch.zeros(t, device='cuda', dtype=torch.int32) hi = torch.full_like(lo, 3) - out = kernels.qsa_block_sparse_attention_triton(q, k, v, sel, lo, hi, lo, lo, d**-0.5, 4) + first, last = _block_spans(sel.shape[1], 4, t) + out = qsa_block_sparse_attention(q, k, v, sel, lo, hi, lo, lo, first, last, d**-0.5, 4) torch.cuda.synchronize() torch.testing.assert_close(out, torch.ones_like(out), rtol=0, atol=0) out.sum().backward() diff --git a/tests/test_qsa_tile_index.py b/tests/test_qsa_tile_index.py new file mode 100644 index 00000000..f76ae7af --- /dev/null +++ b/tests/test_qsa_tile_index.py @@ -0,0 +1,141 @@ +"""CPU-only: a packed selection must attend exactly its keys, and the tile map must +list every attended key's tile.""" +import pytest +import random + +from mcore_bridge.model.modules.kernels import qsa_kernels as qk +from mcore_bridge.model.modules.kernels.qsa_block_sparse_attn import build_tile_index, build_tile_index_pair + +torch = pytest.importorskip('torch') +pytest.importorskip('triton') + +BQ = 64 +_rng = random.Random(0) +PACKS = [ + ([61, 128], 4), + ([61, 61, 61, 200], 4), + ([64, 128], 4), + ([1, 1, 1, 1, 1, 200], 4), + ([3, 20], 4), + ([1, 3, 20, 37, 9], 4), + *[([_rng.randint(1, 300) for _ in range(_rng.randint(2, 6))], _rng.choice([3, 4, 5, 8])) for _ in range(6)], +] + + +def _selection_rows(lens, ratio, seed=0): + """Indexer-shaped rows: whole document-relative blocks (with holes) + causal tail.""" + cu = torch.tensor([0, *lens]).cumsum(0) + total = int(cu[-1]) + starts = cu[:-1] + seg = torch.zeros(total, dtype=torch.long) + seg[starts[1:]] = 1 + seg = seg.cumsum(0) + positions = torch.arange(total) - starts[seg] + layout = qk._PackedBlockLayout(cu, ratio, 'cpu') + g = torch.Generator().manual_seed(seed) + + rows = [] + for t in range(total): + seq = int(seg[t]) + p = int(positions[t]) + visible = (p + 1) // ratio + blocks = torch.randperm(visible, generator=g)[:min(2, visible)].sort().values if visible else \ + torch.zeros(0, dtype=torch.long) + full = (blocks[:, None] * ratio + torch.arange(ratio)[None, :]).reshape(-1) + int(starts[seq]) + qn = (p + 1) // ratio + tail = torch.arange(int(starts[seq]) + qn * ratio, min(int(starts[seq]) + qn * ratio + ratio, t + 1)) + row = torch.cat([full, tail]).unique() + rows.append(row) + idx = torch.full((total, max(r.numel() for r in rows)), -1, dtype=torch.int64) + for t, row in enumerate(rows): + idx[t, :row.numel()] = row + return layout, idx + + +def _admitted(layout, idx, block_size, sel, lo, hi): + """Kernel membership rule: selected block column and inside [lo, hi].""" + total = idx.shape[0] + keys = torch.arange(total) + tok_base = layout.token_start + blk_base = layout.token_block_start + blk = blk_base.unsqueeze(1) + (keys[None, :] - tok_base.unsqueeze(1)) // block_size + inside = (blk >= 0) & (blk < sel.shape[1]) + gathered = torch.gather(sel.long(), 1, blk.clamp(0, sel.shape[1] - 1)) + return (gathered != 0) & inside & (keys[None, :] >= lo[:, None]) & (keys[None, :] <= hi[:, None]) + + +def _tile_sets(lst, cnt): + return [set(lst[i, :int(cnt[i])].tolist()) for i in range(cnt.numel())] + + +@pytest.mark.parametrize('lens,ratio', PACKS, ids=lambda x: str(x)) +def test_packed_off_grid_attends_exactly_the_selection(lens, ratio): + layout, idx = _selection_rows(lens, ratio) + sel, lo, hi, blk_base, tok_base, block_first, block_last = qk._block_form(idx, layout, ratio) + + exact = torch.zeros(idx.shape[0], idx.shape[0], dtype=torch.bool) + rows = torch.arange(idx.shape[0]).unsqueeze(1).expand_as(idx) + valid = idx >= 0 + exact[rows[valid], idx[valid]] = True + # rows may list the current block whole; the kernel clips at [lo, hi] + keys = torch.arange(idx.shape[0]) + exact &= (keys[None, :] >= lo[:, None]) & (keys[None, :] <= hi[:, None]) + + admitted = _admitted(layout, idx, ratio, sel, lo, hi) + extra = int((admitted & ~exact).sum()) + missed = int((exact & ~admitted).sum()) + assert extra == 0, f'{lens} R={ratio}: {extra} keys attended outside the selection' + assert missed == 0, f'{lens} R={ratio}: {missed} selected keys dropped' + + +@pytest.mark.parametrize('lens,ratio', PACKS, ids=lambda x: str(x)) +def test_packed_off_grid_tile_lists_cover_every_attended_key(lens, ratio): + layout, idx = _selection_rows(lens, ratio) + sel, lo, hi, blk_base, tok_base, block_first, block_last = qk._block_form(idx, layout, ratio) + admitted = _admitted(layout, idx, ratio, sel, lo, hi) + total = idx.shape[0] + + for bk, pair in ((64, False), (32, True)): + if pair: + klist, kcnt, qlist, qcnt = build_tile_index_pair(sel, block_first, block_last, BQ, bk) + else: + klist, kcnt = build_tile_index(sel, block_first, block_last, BQ, bk) + qlist, qcnt = build_tile_index_pair(sel, block_first, block_last, BQ, bk)[2:] + k_tiles = _tile_sets(klist, kcnt) + q_tiles = _tile_sets(qlist, qcnt) + missing = 0 + for t in range(total): + for key in admitted[t].nonzero().flatten().tolist(): + if key // bk not in k_tiles[t // BQ]: + missing += 1 + if t // BQ not in q_tiles[key // bk]: + missing += 1 + assert missing == 0, f'{lens} R={ratio} BK={bk}: tile lists miss {missing} (query, key) pairs' + + +class _PackedSeqParams: + + def __init__(self, cu, padded=None): + self.cu_seqlens_q = cu + self.cu_seqlens_q_padded = padded + + +@pytest.mark.parametrize( + 'cu,padded,total,expected', + [ + ([0, 3, 23], None, 23, [0, 3, 23]), # plain padding_free packing + ([3, 23], None, 23, [0, 3, 23]), # missing leading 0 + ([5, 8, 28], None, 23, [0, 3, 23]), # per-batch offset + ([0, 3, 23], [0, 4, 28], 28, [0, 4, 28]), # CP: padded space wins + ([0, 3, 23], None, 28, None), # nothing describes 28 tokens -> loud + ]) +def test_packed_cu_seqlens_alignment(cu, padded, total, expected): + psp = _PackedSeqParams(None if cu is None else torch.tensor(cu, dtype=torch.int32), + None if padded is None else torch.tensor(padded, dtype=torch.int32)) + if expected is None: + with pytest.raises(ValueError): + qk._packed_cu_seqlens(psp, total) + return + aligned = qk._packed_cu_seqlens(psp, total) + assert aligned.tolist() == expected + assert aligned.dtype == torch.long