From aab6bdd83e6705577d5c3e3e7260c0f461d7b8f5 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 21 Oct 2026 12:12:49 -0400 Subject: [PATCH] perf(qsa): score prefill rows in tiles of 16 query rows per program The prefill indexer score ran one query row per program and padded the 4 index heads to a 16-wide dot, so every row re-read all visible keys and 3/4 of the tensor-core work was padding. qsa_mqa_paged_rows gives each program 16 rows of one request (a host-built tile map cut at the cu_seqlens request boundaries) and does one [BLOCK_N, 16] dot per head, summing the heads in the one-row kernel's order, so every score is bitwise the old one. Decode (graph-captured) keeps the one-row kernel. Microbench, one 4096-row chunk at 249K context: 18.5 -> 3.2 ms per layer, no spills (134 regs). Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846 --- python/freetoken/attention/qsa_sparse.py | 53 ++++- .../freetoken/kernel/triton/qsa/__init__.py | 4 +- python/freetoken/kernel/triton/qsa/score.py | 221 ++++++++++++++++++ tests/kernels/test_qsa_score_rows.py | 136 +++++++++++ 4 files changed, 402 insertions(+), 12 deletions(-) create mode 100644 tests/kernels/test_qsa_score_rows.py diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py index 7d20092efc..9d2a3bc8b0 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -50,6 +50,7 @@ # Block-score transient budget (vLLM's number): the fp32 [rows, n_blocks] logits tile is # 256 KB per row at a 1M-token context, so a long prefill must be scored in row chunks. _LOGITS_WORKSPACE_BYTES = 128 << 20 +_SCORE_ROWS = 16 # prefill query rows per score program (qsa_mqa_paged_rows) TORCH_TOPK_ENV = "FREETOKEN_QSA_TORCH_TOPK" @@ -96,6 +97,8 @@ class QSASparseMetadata(BaseAttnMetadata): # fetch query so no mid-forward eviction can take them) and the physical out_loc. write_pages: torch.Tensor | None = None out_loc_gpu: torch.Tensor | None = None + # Prefill scoring: row-tile maps per (start, end) row chunk, built by the first QSA layer. + score_tiles: dict | None = None # fmt: on def get_last_indices(self, bs: int) -> torch.Tensor: @@ -464,6 +467,8 @@ def _select(self, index, md: QSASparseMetadata, slot: int) -> torch.Tensor: expand_qsa_block_indices, qsa_index_norm_rope, qsa_mqa_paged, + qsa_mqa_paged_rows, + qsa_score_tile_map, ) rows = index.q.shape[0] @@ -488,22 +493,48 @@ def _select(self, index, md: QSASparseMetadata, slot: int) -> torch.Tensor: columns = md.block_table.shape[1] * self.cmp_page_size indices = self._scratch("indices", rows, self.select_width, dtype=torch.int32) rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1)) + row_tiles = not md.is_decode and self.index_heads == 4 for start in range(0, rows, rows_per_chunk): end = min(start + rows_per_chunk, rows) chunk = slice(start, end) logits = self._scratch("logits", end - start, columns, dtype=torch.float32) visible = self._scratch("visible", end - start, dtype=torch.int32) - qsa_mqa_paged( - q_index[chunk], - cmp_pages, - md.block_table, - md.token_to_req[chunk], - positions[chunk], - md.seq_lens, - self.ratio, - logits, - visible, - ) + if row_tiles: + # Prefill: several query rows of one request per program (decode, which a + # graph may capture, keeps the one-row kernel). Same scores bit for bit. + if md.score_tiles is None: + md.score_tiles = {} + tile_map = md.score_tiles.get((start, end)) + if tile_map is None: + tile_map = qsa_score_tile_map( + md.qo_indptr_cpu, start, end, _SCORE_ROWS, self.device + ) + md.score_tiles[(start, end)] = tile_map + qsa_mqa_paged_rows( + q_index[chunk], + cmp_pages, + md.block_table, + tile_map, + md.token_to_req[chunk], + positions[chunk], + md.seq_lens, + self.ratio, + logits, + visible, + rows_per_tile=_SCORE_ROWS, + ) + else: + qsa_mqa_paged( + q_index[chunk], + cmp_pages, + md.block_table, + md.token_to_req[chunk], + positions[chunk], + md.seq_lens, + self.ratio, + logits, + visible, + ) blocks = self._scratch("blocks", end - start, self.block_topk, dtype=torch.int32) self._top_blocks(logits, visible, blocks) expand_qsa_block_indices( diff --git a/python/freetoken/kernel/triton/qsa/__init__.py b/python/freetoken/kernel/triton/qsa/__init__.py index 91753b7922..b04d27f050 100644 --- a/python/freetoken/kernel/triton/qsa/__init__.py +++ b/python/freetoken/kernel/triton/qsa/__init__.py @@ -3,7 +3,7 @@ from .attend import qsa_sparse_paged_attention from .compress import qsa_compress_groups, qsa_index_norm_rope, qsa_store_rows from .expand import expand_qsa_block_indices -from .score import qsa_mqa_paged +from .score import qsa_mqa_paged, qsa_mqa_paged_rows, qsa_score_tile_map from .topk import qsa_block_topk, qsa_block_topk_scratch_width __all__ = [ @@ -13,6 +13,8 @@ "qsa_compress_groups", "qsa_index_norm_rope", "qsa_mqa_paged", + "qsa_mqa_paged_rows", + "qsa_score_tile_map", "qsa_sparse_paged_attention", "qsa_store_rows", ] diff --git a/python/freetoken/kernel/triton/qsa/score.py b/python/freetoken/kernel/triton/qsa/score.py index 7b65da8d6f..24526d0e86 100644 --- a/python/freetoken/kernel/triton/qsa/score.py +++ b/python/freetoken/kernel/triton/qsa/score.py @@ -115,6 +115,227 @@ def _qsa_mqa_paged_kernel( ) +@triton.jit +def _qsa_mqa_paged_rows_kernel( + q_ptr, + k_cache_ptr, + page_table_ptr, + tile_map_ptr, + token_to_req_ptr, + query_positions_ptr, + sequence_lengths_ptr, + visible_blocks_ptr, + logits_ptr, + stride_q_row, + stride_q_head, + stride_q_dim, + stride_cache_block, + stride_cache_token, + stride_cache_dim, + stride_table_req, + stride_table_page, + stride_logits_row, + num_tiles, + num_columns, + num_pages, + num_requests, + score_divisor, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, + ROWS: tl.constexpr, + TILES_PER_PROG: tl.constexpr, + STAGES: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +) -> None: + """Prefill scoring, ROWS query rows x 4 heads per program. + + The rows of a tile all belong to one request (``tile_map`` holds each tile's first row + and row count, cut at request boundaries), so one key tile serves every row. Each head + is its own ``[BLOCK_N, ROWS]`` dot and the head sum keeps the one-row kernel's order + ``(h0 + h1) + (h2 + h3)``: every score is bitwise the one-row kernel's. + """ + tile = tl.program_id(0) + row0 = tl.load(tile_map_ptr + tile) + nrows = tl.load(tile_map_ptr + num_tiles + tile) + dims = tl.arange(0, BLOCK_D) + lanes = tl.arange(0, ROWS) + row_live = lanes < nrows + rows = row0 + lanes + request = tl.load(token_to_req_ptr + row0) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + query_position = tl.load(query_positions_ptr + rows, mask=row_live, other=0) + sequence_length = tl.load( + sequence_lengths_ptr + safe_request, + mask=(request >= 0) & (request < num_requests), + other=0, + ) + visible = tl.minimum( + (query_position + 1) // COMPRESS_RATIO, + sequence_length // COMPRESS_RATIO, + ) + visible = tl.where(row_live, visible, 0) + if tl.program_id(1) == 0: + tl.store(visible_blocks_ptr + rows, visible, mask=row_live) + visible_max = tl.max(visible, axis=0) + tile_start = tl.program_id(1) * TILES_PER_PROG + # Top-k is bounded by visible_blocks, so columns beyond it need no value. + if tile_start * BLOCK_N >= visible_max: + return + tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(visible_max, BLOCK_N)) + tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N)) + + q_rows = q_ptr + rows[None, :].to(tl.int64) * stride_q_row + dims[:, None] * stride_q_dim + q_mask = row_live[None, :] & (dims[:, None] < HEAD_DIM) + query0 = tl.load(q_rows + 0 * stride_q_head, mask=q_mask, other=0.0) + query1 = tl.load(q_rows + 1 * stride_q_head, mask=q_mask, other=0.0) + query2 = tl.load(q_rows + 2 * stride_q_head, mask=q_mask, other=0.0) + query3 = tl.load(q_rows + 3 * stride_q_head, mask=q_mask, other=0.0) + column_offsets = tl.arange(0, BLOCK_N) + out_rows = logits_ptr + rows[None, :].to(tl.int64) * stride_logits_row + for col_tile in tl.range(tile_start, tile_end, num_stages=STAGES): + columns = col_tile * BLOCK_N + column_offsets + live = columns < visible_max + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1) + page_offset = columns % PAGE_SIZE + physical_page = tl.load( + page_table_ptr + + safe_request * stride_table_req + + logical_page * stride_table_page, + mask=live, + other=-1, + ) + page_valid = live & (physical_page >= 0) & (physical_page < num_pages) + # physical_page * block stride can overflow int32 for large caches. + safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr + + safe_physical_page[:, None] * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :] * stride_cache_dim, + mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + s0 = tl.maximum(tl.dot(keys, query0, out_dtype=tl.float32), 0.0) + s1 = tl.maximum(tl.dot(keys, query1, out_dtype=tl.float32), 0.0) + s2 = tl.maximum(tl.dot(keys, query2, out_dtype=tl.float32), 0.0) + s3 = tl.maximum(tl.dot(keys, query3, out_dtype=tl.float32), 0.0) + score = ((s0 + s1) + (s2 + s3)) / score_divisor + tl.store( + out_rows + columns[:, None], + tl.where(page_valid[:, None], score, -float("inf")), + mask=(columns[:, None] < visible[None, :]) & (columns[:, None] < num_columns), + ) + + +def qsa_score_tile_map( + qo_indptr: torch.Tensor, start: int, end: int, rows_per_tile: int, device: torch.device +) -> torch.Tensor: + """``[2, tiles]`` int32 (first row, row count) of rows ``[start, end)``, relative to + ``start``, in tiles of at most ``rows_per_tile`` rows that never cross a request boundary + of the host-side ``qo_indptr`` (cu_seqlens).""" + + bounds = [int(x) for x in qo_indptr.tolist()] + firsts: list[int] = [] + counts: list[int] = [] + for lo, hi in zip(bounds, bounds[1:]): + lo, hi = max(lo, start), min(hi, end) + for row in range(lo, hi, rows_per_tile): + firsts.append(row - start) + counts.append(min(rows_per_tile, hi - row)) + host = torch.tensor([firsts, counts], dtype=torch.int32) + if device.type == "cuda": + host = host.pin_memory() + return host.to(device, non_blocking=True) + + +SCORE_ROWS = 16 + + +def qsa_mqa_paged_rows( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + tile_map: torch.Tensor, + token_to_req: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + compress_ratio: int, + logits: torch.Tensor, + visible_blocks: torch.Tensor, + score_scale: float | None = None, + rows_per_tile: int = SCORE_ROWS, +) -> tuple[torch.Tensor, torch.Tensor]: + """Prefill form of :func:`qsa_mqa_paged`: the same scores bit for bit, with + ``rows_per_tile`` query rows of one request per program instead of one row padded to 16 + heads. ``tile_map`` comes from :func:`qsa_score_tile_map` over the same rows and must use + the same ``rows_per_tile``. Needs exactly 4 index heads.""" + + if q.ndim != 3 or q.shape[1] != 4 or q.shape[2] <= 0: + raise ValueError("QSA row-tiled scoring needs a [rows, 4, head_dim] query") + if k_cache.ndim != 4 or k_cache.shape[2] != 1 or k_cache.shape[3] != q.shape[2]: + raise ValueError("QSA cache must be [pages, page_size, 1, head_dim] matching the query") + if token_to_req.shape != (q.shape[0],) or query_positions.shape != (q.shape[0],): + raise ValueError("QSA request mapping and positions must match query rows") + if sequence_lengths.shape != (page_table.shape[0],): + raise ValueError("QSA sequence lengths must match page-table requests") + if q.dtype not in (torch.bfloat16, torch.float16) or k_cache.dtype is not q.dtype: + raise ValueError( + f"QSA scoring is 16-bit only, got query={q.dtype} keys={k_cache.dtype}" + ) + if tile_map.ndim != 2 or tile_map.shape[0] != 2 or tile_map.dtype != torch.int32: + raise ValueError("QSA tile map must be [2, tiles] int32") + if rows_per_tile < 16 or rows_per_tile & (rows_per_tile - 1): + raise ValueError("rows_per_tile must be a power of two >= 16 (the dot's N)") + score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale + columns = logits.shape[1] + tiles = tile_map.shape[1] + if not q.shape[0] or not columns or not tiles: + return logits, visible_blocks + BLOCK_N = 64 + BLOCK_D = max(16, triton.next_power_of_2(q.shape[2])) + tiles_per_program = 8 + _qsa_mqa_paged_rows_kernel[(tiles, triton.cdiv(columns, BLOCK_N * tiles_per_program))]( + q, + k_cache, + page_table, + tile_map, + token_to_req, + query_positions, + sequence_lengths, + visible_blocks, + logits, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(3), + page_table.stride(0), + page_table.stride(1), + logits.stride(0), + tiles, + columns, + k_cache.shape[0], + page_table.shape[0], + float(score_divisor), + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=page_table.shape[1], + HEAD_DIM=q.shape[2], + BLOCK_N=BLOCK_N, + BLOCK_D=BLOCK_D, + ROWS=rows_per_tile, + TILES_PER_PROG=tiles_per_program, + STAGES=2, + COMPRESS_RATIO=compress_ratio, + num_warps=4, + ) + return logits, visible_blocks + + def qsa_mqa_paged( q: torch.Tensor, k_cache: torch.Tensor, diff --git a/tests/kernels/test_qsa_score_rows.py b/tests/kernels/test_qsa_score_rows.py new file mode 100644 index 0000000000..fec0898679 --- /dev/null +++ b/tests/kernels/test_qsa_score_rows.py @@ -0,0 +1,136 @@ +"""The prefill row-tiled QSA score kernel (``qsa_mqa_paged_rows``) against the one-row kernel. + +Both must give the same visible counts, bitwise-equal logits on every visible column, and so the +same top-k block selection, for ragged multi-request prefill chunks whose request starts are not +aligned to the row tile, and for a row slice in the middle of a forward (the ``_select`` chunking). +""" + +from __future__ import annotations + +import pytest +import torch + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +CMP_PAGE = 16 +RATIO = 4 +HEADS = 4 +DIM = 128 +TOPK = 512 + + +def _case(lengths, rows_per_req, seed, with_invalid=False): + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(seed) + max_blocks = max(lengths) // RATIO + width = -(-max_blocks // CMP_PAGE) + bs = len(lengths) + total = bs * width + 3 + block_table = torch.randperm(total, device=device, generator=gen)[: bs * width] + block_table = block_table.reshape(bs, width).to(torch.int32) + if with_invalid: + block_table[0, 1] = -1 + block_table[-1, 0] = total + 5 + cache = torch.randn(total, CMP_PAGE, 1, DIM, device=device, dtype=torch.bfloat16, generator=gen) + rows = sum(rows_per_req) + q = torch.randn(rows, HEADS, DIM, device=device, dtype=torch.bfloat16, generator=gen) + token_to_req = torch.repeat_interleave( + torch.arange(bs, dtype=torch.int32), torch.tensor(rows_per_req) + ).to(device) + positions = torch.cat( + [ + torch.arange(n - r, n, dtype=torch.int32, device=device) + for n, r in zip(lengths, rows_per_req) + ] + ) + seq_lens = torch.tensor(lengths, dtype=torch.int32, device=device) + qo_indptr = torch.tensor([0] + list(rows_per_req), dtype=torch.int32).cumsum(0).to(torch.int32) + return q, cache, block_table, token_to_req, positions, seq_lens, qo_indptr, width * CMP_PAGE + + +def _run_both(case, start, end, rows_per_tile): + from freetoken.kernel.triton.qsa import ( + qsa_block_topk, + qsa_mqa_paged, + qsa_mqa_paged_rows, + qsa_score_tile_map, + ) + + q, cache, table, t2r, pos, seq, qo, columns = case + sl = slice(start, end) + n = end - start + out = [] + for rows_kernel in (False, True): + logits = torch.full((n, columns), float("nan"), dtype=torch.float32, device=q.device) + visible = torch.full((n,), -7, dtype=torch.int32, device=q.device) + if rows_kernel: + tiles = qsa_score_tile_map(qo, start, end, rows_per_tile, q.device) + qsa_mqa_paged_rows( + q[sl], cache, table, tiles, t2r[sl], pos[sl], seq, RATIO, logits, visible, + rows_per_tile=rows_per_tile, + ) + else: + qsa_mqa_paged(q[sl], cache, table, t2r[sl], pos[sl], seq, RATIO, logits, visible) + blocks = torch.empty((n, TOPK), dtype=torch.int32, device=q.device) + qsa_block_topk(logits.clone(), visible, blocks) + out.append((logits, visible, blocks)) + return out + + +def _assert_same(old, new): + (lo, vo, bo), (ln, vn, bn) = old, new + assert torch.equal(vo, vn) + col = torch.arange(lo.shape[1], device=lo.device) + live = col[None, :] < vo[:, None].to(torch.int64) + assert live.any() + a, b = lo[live], ln[live] + assert not torch.isnan(a).any() and not torch.isnan(b).any() + # bitwise, including the sign of zero and -inf on invalid pages + assert torch.equal(a.view(torch.int32), b.view(torch.int32)), ( + f"{(a.view(torch.int32) != b.view(torch.int32)).sum().item()} of {a.numel()} scores differ" + ) + assert torch.equal(bo, bn) + + +@requires_cuda +@pytest.mark.parametrize("rows_per_tile", [16, 32]) +@pytest.mark.parametrize( + "lengths, rows_per_req", + [ + ([9000], [517]), + ([20000, 3001], [37, 300]), + ([700, 12345, 64, 5000], [5, 210, 33, 17]), + ([4096], [4096]), + ], +) +def test_rows_kernel_is_bitwise_the_one_row_kernel(lengths, rows_per_req, rows_per_tile): + case = _case(lengths, rows_per_req, seed=sum(lengths) + rows_per_tile) + rows = sum(rows_per_req) + old, new = _run_both(case, 0, rows, rows_per_tile) + _assert_same(old, new) + + +@requires_cuda +@pytest.mark.parametrize("start, end", [(100, 612), (0, 512), (517, 818)]) +def test_rows_kernel_on_a_row_slice(start, end): + case = _case([30000, 2000, 9000], [600, 101, 117], seed=start + end, with_invalid=True) + old, new = _run_both(case, start, end, 16) + _assert_same(old, new) + + +@requires_cuda +def test_rows_kernel_at_long_context(): + """A late 249K-context chunk: ~62K visible blocks per row, 512-row slice.""" + case = _case([249000], [512], seed=3) + old, new = _run_both(case, 0, 512, 16) + _assert_same(old, new) + + +def test_tile_map_never_crosses_a_request(): + from freetoken.kernel.triton.qsa import qsa_score_tile_map + + qo = torch.tensor([0, 5, 40, 41, 100], dtype=torch.int32) + tiles = qsa_score_tile_map(qo, 3, 90, 16, torch.device("cpu")) + firsts, counts = tiles.tolist() + assert firsts == [0, 2, 18, 34, 37, 38, 54, 70, 86] + assert counts == [2, 16, 16, 3, 1, 16, 16, 16, 1]