perf(qsa): score prefill rows in tiles of 16 query rows per program

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(

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",
 ]

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,