flashinfer.top_k_page_table_transform

flashinfer.top_k_page_table_transform(input: Tensor, src_page_table: Tensor, lengths: Tensor, k: int, row_to_batch: Tensor | None = None, deterministic: bool = False, tie_break: int = TopKTieBreak.NONE, dsa_graph_safe: bool = False, row_starts: Tensor | None = None, page_table_row_starts: Tensor | None = None, *, page_size: int = 1, out: Tensor | None = None, out_raw_indices: Tensor | None = None) → Tensor

Fused Top-K selection + Page Table Transform for sparse attention.

This function performs top-k selection on input scores and translates the selected indices through a page table in a single fused kernel. Each page-table entry represents page_size consecutive score positions. For each selected local index idx in row i:

physical_page = src_page_table[
    batch_idx, page_table_row_start[i] + idx // page_size
]
output[i, j] = physical_page * page_size + idx % page_size

where batch_idx is determined by row_to_batch[i] if provided, otherwise i. topk_indices are relative to row_starts[i].

Parameters:

  • input (torch.Tensor) – Input scores tensor of shape (num_rows, max_len). Supported dtypes: float32, float16, bfloat16.

  • src_page_table (torch.Tensor) – Source page table of shape (batch_size, max_page_table_length) with dtype int32. Entries used by selected indices must be nonnegative, and each resulting physical_page * page_size + offset must fit in signed int32. These value constraints are not checked at runtime.

  • lengths (torch.Tensor) – Actual KV lengths per row of shape (num_rows,) with dtype int32.

  • k (int) – Number of top elements to select from each row.

  • row_to_batch (Optional[torch.Tensor], optional) – Mapping from row index to batch index of shape (num_rows,) with dtype int32. If None, uses 1:1 mapping (row_idx == batch_idx). Default is None.

  • deterministic (bool, optional) – If True, uses deterministic mode. Default is False (non-deterministic, which is faster).

  • tie_break (int, optional) –

    Tie-breaking mode for equal values at the selection boundary. Supported modes are (or use TopKTieBreak enum values):

    • 0: no explicit index tie-break

    • 1: prefer smaller indices

    • 2: prefer larger indices

    Default is 0. Tie-breaking controls which boundary elements are selected; it does not imply deterministic output ordering. Set deterministic=True when repeatable output ordering is also required.

  • dsa_graph_safe (bool, optional) – If True, force FilteredTopK path and graph-safe vectorization (VEC_SIZE=1). Default is False.

  • row_starts (Optional[torch.Tensor], optional) – Per-row start indices of shape (num_rows,) with dtype int32. Top-k is computed over [row_starts[i], row_starts[i] + lengths[i]) for row i. Default is None (equivalent to all zeros).

  • page_table_row_starts (Optional[torch.Tensor], optional) – Per-row page-table start indices of shape (num_rows,) with dtype int32, measured in page-table entries. If None, defaults to row_starts, so score and page-table windows share the same start. When page_size > 1 and row_starts is provided, this argument must also be provided because the two starts use different units.

  • page_size (int, optional) – Number of score positions represented by each page-table entry. Must be a positive power of two no greater than 2**30. Setting this to 1 preserves the one-entry-per-score behavior. Default is 1.

  • out (Optional[torch.Tensor], optional) – Optional contiguous int32 output buffer of shape (num_rows, k). Supplying this buffer avoids an allocation and is CUDA-graph friendly.

  • out_raw_indices (Optional[torch.Tensor], optional) – Optional contiguous int32 output buffer of shape (num_rows, k). Receives selected indices relative to each score window before page-table translation. Padding positions are set to -1 and remain positionally aligned with out. Must not overlap out.

Returns:

output – Physical indices of shape (num_rows, k) with dtype int32. This is the same tensor as out when one is supplied. Positions beyond actual length are set to -1.

Return type:

torch.Tensor

Note

  • This is specifically designed for sparse attention’s second stage.

  • input may have padding between rows, but its last dimension must be contiguous.

  • If lengths[i] <= k, raw indices are 0..lengths[i]-1 and remaining positions are set to -1.

Examples

>>> import torch
>>> import flashinfer
>>> num_rows = 8
>>> max_len = 4096
>>> k = 256
>>> scores = torch.randn(num_rows, max_len, device="cuda", dtype=torch.float16)
>>> src_page_table = torch.randint(0, 1000, (num_rows, max_len), device="cuda", dtype=torch.int32)
>>> lengths = torch.full((num_rows,), max_len, device="cuda", dtype=torch.int32)
>>> output = flashinfer.top_k_page_table_transform(scores, src_page_table, lengths, k)
>>> output.shape
torch.Size([8, 256])