flashinfer.top_k_varlen

flashinfer.top_k_varlen(logits: Tensor, seq_lens: Tensor, top_k: int, *, row_starts: Tensor | None = None, max_window_len: int | None = None, absolute_indices: bool = False, page_table: Tensor | None = None, page_size: int | None = None, compress_ratio: int | None = None, next_n: int | None = None, pre_idx: Tensor | None = None, return_values: bool = False, out_indices: Tensor | None = None, out_values: Tensor | None = None, backend: Literal['radix', 'gvr', 'gvr_2', 'radix_cutlass', 'radix_filter', 'auto'] = 'auto', load_balance: bool = True, workspace: dict | None = None) → Tuple[Tensor, Tensor | None]

Top-k selection per row over rows of different valid lengths.

logits holds one row per query token and one score per key column; num_cols is the width of that static buffer, the same for every row. In decode mode (the default), each row is a request’s newest token (one row per request, or next_n rows with speculative tokens) and its valid columns are the request’s cached keys, which start at column 0: with one row per request and no compression they are [0, seq_lens[i]), and the seq_lens entry below gives their count under next_n > 1 or compress_ratio > 1. In windowed (prefill) mode (row_starts given), each row is one prompt token of a packed chunk and its valid columns are its request’s window [row_starts[r], row_starts[r] + seq_lens[r]). The result is the top_k indices of every row’s valid columns, -1 padded where a row has fewer valid columns than top_k, and with return_values=True also the selected scores. In decode mode, an index is the logits column, which is also the key’s position in its request’s cache, or the key’s physical KV slot when page_table is given. In windowed (prefill) mode, an index is by default the position inside the window: 0 is row_starts[r], the request’s first key, so a window position addresses the request’s cache exactly like a decode index. With absolute_indices=True an index is instead the logits column row_starts[r] + position, which addresses the chunk’s packed key sequence directly.

Backend selection

backend="auto" (default) runs the highest-ranked backend among those whose admission checks pass for the call. The checks cover dtype, top_k, layout and alignment, the mode arguments, the GPU and the installed CuTe DSL. The ranking is a static table over dtype, top_k, num_cols, request count and whether a valid hint (pre_idx) is given; it never reads the contents of seq_lens, which would cost a device synchronization. Every backend returns the exact top-k set of each row for finite and infinite scores (ties at the K-th value may be broken differently), so auto only ever affects speed. The table is maintained from measurements, revised between releases, and not part of the contract. Force a backend with backend="radix", "gvr", "gvr_2", "radix_filter" or "radix_cutlass". top_k_varlen.is_backend_supported(name, compute_capability) tells whether a backend exists for a GPU.

Windowed (prefill) mode

row_starts switches a call into windowed (prefill) mode. max_window_len and absolute_indices apply only there (a decode call refuses max_window_len and absolute_indices=True with ValueError); next_n, compress_ratio and pre_idx belong to decode mode (a windowed call refuses any value but None, next_n=1 and compress_ratio=1 included, with ValueError).

The mode exists for the DSA indexer’s prefill layout, where two things move a row’s keys away from column 0: several requests share one packed key axis, and chunked prefill or prefix caching puts a request’s earlier keys on that axis without rows. In decode every row is one request whose keys start at column 0. In prefill the indexer scores one chunk of the batch in a single grouped GEMM: the keys of every request, including the prefix cached by earlier chunks, occupy consecutive column ranges of one packed key axis, while only the chunk’s new tokens have rows. A row’s valid keys therefore start at its request’s column offset, row_starts[r], and span the cached prefix plus the causal part of the chunk, seq_lens[r] columns; cells outside that window belong to other requests or to later tokens and are never read. Toy chunk, request A (2 new tokens) and request B (1 cached token b0, 3 new tokens), top_k=2:

           col: 0   1  | 2   3   4   5       row_starts  seq_lens
                a0  a1 | b0  b1  b2  b3
row 0  a0     [ *   .  | .   .   .   .  ]       0           1
row 1  a1     [ *   *  | .   .   .   .  ]       0           2
row 2  b1     [ .   .  | *   *   .   .  ]       2           2
row 3  b2     [ .   .  | *   *   *   .  ]       2           3
row 4  b3     [ .   .  | *   *   *   *  ]       2           4

* marks a key inside the row’s window, . a cell outside it (garbage from the GEMM, never read as a candidate). row_starts says where the row’s request begins; seq_lens becomes the window length (cached prefix + position + 1; the cached b0 is a column, not a row). Row 4 selecting b0 and b2 as its two top-k indices (top_k=2 here) is returned as the positions 0 and 2 inside B (the default: a paged consumer maps them through B’s page table) or as the columns 2 and 4 of logits (absolute_indices=True: the ragged consumer reads the packed KV buffer by column). max_window_len (here 4, the longest window) is only a bound the backend specializes on before the launch and defaults to num_cols; while every window stays within the bound, it never changes which keys win.

Rules of the mode: a window shorter than top_k returns all its positions; a window end past num_cols is clamped to num_cols. Under a captured graph the contents of row_starts and seq_lens may change between replays while every window stays within max_window_len. The backends that serve the mode, with their limits and CUDA-graph warm-up rule, are listed under backend.

param logits:

2-D CUDA tensor of shape (num_rows, num_cols) in float32, bfloat16 or float16; the dtypes, layouts and alignments each backend accepts are listed under backend.

type logits:

torch.Tensor

param seq_lens:

1-D int32 tensor with one entry per request, shape (num_rows // next_n,): each request’s KV-cache length. Units: in decode mode, KV tokens, before compression (the kernel divides by compress_ratio); in windowed (prefill) mode, columns of logits, after compression. Base case, next_n == 1 and compress_ratio == 1: request i is row i and ranks its first seq_lens[i] columns; columns at or beyond that count are not candidates. With speculative decoding, next_n > 1: request i owns rows i * next_n to i * next_n + next_n - 1, and row t of the request (t from 0 to next_n - 1) ranks its first seq_lens[i] - next_n + t + 1 columns, so the last draft token sees the whole cache and each earlier one sees one token less. With KV compression, compress_ratio > 1: every column stands for compress_ratio KV tokens and the token count is divided by it (floor). The general per-row count is therefore max(0, (seq_lens[i] - next_n + t + 1) // compress_ratio). seq_lens stays in tokens even with compress_ratio > 1: the kernel converts tokens to columns itself, so a value erroneously given in columns would be converted twice and rank too few keys. A row with nothing to rank (its count is zero, as for an empty padding slot) comes back as -1 in every slot. A row whose length is larger than the buffer is wide (count above num_cols) is ranked over the num_cols keys the buffer holds; nothing is read beyond it. In windowed (prefill) mode there is one entry per row and seq_lens[r] is the window length of row r in columns, after compression (see row_starts).

type seq_lens:

torch.Tensor

param top_k:

Number of top elements per row; the accepted values are backend-specific (see backend).

type top_k:

int

param row_starts:

Windowed (prefill) mode only; passing it switches the mode on (see “Windowed (prefill) mode” above). 1-D int32 tensor, contiguous, one entry per row: the column where row r’s window begins, so row r ranks logits[r, row_starts[r] : row_starts[r] + seq_lens[r]]. Starts must lie within the row; they are device data and are not validated on the host. Default None.

type row_starts:

torch.Tensor, optional

param max_window_len:

Windowed (prefill) mode only: an upper bound on every window length seq_lens[r], in columns; defaults to num_cols. The bound is consumed before the launch and fixed for a captured graph, so it must hold at every replay. Pass the batch’s longest window (seq_lens.max()) when it is smaller than num_cols; the tighter the bound, the more the backend can specialize. A positive int; a value above num_cols counts as num_cols. Refused with ValueError without row_starts.

type max_window_len:

int, optional

param absolute_indices:

Windowed (prefill) mode only: write each hit as its column of logits (window-local index + row_starts[r]) instead of the window-local position; -1 padding is unchanged. In the DSA prefill layout the column axis is the chunk’s packed key sequence, so this is the flattened-KV position a ragged sparse-attention kernel consumes directly; a paged consumer keeps the default and maps the local position through its page table. True without row_starts raises ValueError. Default False.

type absolute_indices:

bool, optional

param page_table:

Decode mode only. The request’s page map: an int32 tensor, contiguous, with one row per request and at least ceil(num_cols / page_size) columns (one entry per page of the request’s cache, so every column of logits maps through it); columns and page_size are in the units of logits. When given, each selected column c of request q is written as the physical KV slot page_table[q, c // page_size] * page_size + c % page_size instead of the column; -1 padding is unchanged. The backends that serve it, and their limits, are listed under backend. Default None.

type page_table:

torch.Tensor, optional

param page_size:

Columns per page_table entry, a power of two up to 2**30. Any value but None without page_table raises ValueError. Default None, meaning 1.

type page_size:

int, optional

param compress_ratio:

Decode mode only: KV-index compression factor (1 for DSv3.2, 4 for DSv4), every logit column standing for compress_ratio consecutive KV tokens. It converts the token lengths in seq_lens into per-row column counts, max(0, (seq_lens[i] - next_n + t + 1) // compress_ratio) (see seq_lens); only the kernel can do that when one request length serves next_n rows. The accepted values are backend-specific (see backend). Any value but None with row_starts raises ValueError: windowed (prefill) mode has one length per row and takes it in columns already, so there is nothing to convert, and a compressed layer passes its windows in compressed columns. Default None, meaning 1.

type compress_ratio:

int, optional

param next_n:

Decode mode only: rows per request, the query tokens a request contributes to the current step (1 without speculative decoding); row r belongs to request r // next_n and the per-row column counts are given under seq_lens. Any value but None with row_starts raises ValueError: windowed (prefill) mode has one length per row, so a request’s extra tokens are plain rows with their own windows. Default None, meaning 1.

type next_n:

int, optional

param pre_idx:

Decode mode only: a hint of each request’s top-k indices, a warm start for the backends that take one (how each backend consumes or ignores it is under backend), as a contiguous, 16-byte-aligned int32 tensor of shape [seq_lens.shape[0], top_k], one row per request. A hint that breaks this is discarded with a RuntimeWarning and the call runs hint-free (an explicit backend that requires one is refused). Any value but None with row_starts raises ValueError. Default None.

type pre_idx:

torch.Tensor, optional

param return_values:

When True, also return the selected logit values. Padded slots hold a backend-specific value (see backend); test indices >= 0 rather than the value. Default False.

type return_values:

bool, optional

param out_indices:

Pre-allocated buffer for the indices: int32, contiguous, 16-byte aligned, of shape (num_rows, top_k) or flat with num_rows * top_k elements. The result is written into it and returned as a (num_rows, top_k) view of it. Default None: a new tensor is allocated.

type out_indices:

torch.Tensor, optional

param out_values:

Pre-allocated buffer for the values, in the dtype of logits and under the layout rules of out_indices; used only with return_values=True. Default None.

type out_values:

torch.Tensor, optional

param backend:

Backend to use. Default "auto".

"radix" — CuTe DSL radix top-K.

SM100, SM103, SM107, SM110, SM120 and SM121 with an installed CuTe DSL. fp32, bf16 or fp16; contiguous logits; no top_k limit; any compress_ratio; pre_idx accepted and ignored; padded value slots are left unwritten. Decode mode only.

"gvr" — GVR (Guess-Verify-Refine) kernel.

SM100, SM103 or SM107 with an installed CuTe DSL. fp32, bf16 or fp16; num_cols a multiple of 16 // itemsize (8 for fp16/bf16, 4 for fp32); top_k in {512, 1024, 2048}; compress_ratio 1 or 4. Requires a well-formed pre_idx; its contents affect only the speed, never the result. load_balance=True balances rows of unequal length and caps the batch at 1024 requests; False has no cap and is the better choice when all rows have similar lengths; both are CUDA-graph safe. Padded value slots hold -FLT_MAX converted to dtype. Decode mode only.

"gvr_2" — Self-sampling GVR V2 kernel.

SM100, SM103 or SM107 with an installed CuTe DSL. fp32 only; unit inner stride, row stride a multiple of 4 elements (for a single row, the width), 16-byte-aligned base, non-overlapping rows (a row stride wider than num_cols is fine); top_k in {512, 1024, 2048}; compress_ratio 1 or 4; pre_idx optional, affecting only the speed; load_balance ignored; padded value slots hold finfo(float32).min; NaN ordering is implementation-specific. Decode mode, and the only backend of windowed (prefill) mode. CUDA graphs: a geometry must be compiled before it is captured. Warm up each decode geometry (row count, num_cols and row stride, top_k, next_n, compress_ratio, hinted or not) with one eager call on the stream that will capture, and each windowed geometry (row count, top_k, absolute_indices and the bound: max_window_len, else num_cols) with an eager call or flashinfer.topk_varlen.warmup_prefill on that stream; flashinfer.topk_varlen.prefill_ready tells whether a windowed geometry would launch without compiling. The warm-up also creates the backend’s workspace slab (21 MB per device and CUDA stream) that the graphs captured on that stream share, so those graphs must not be replayed concurrently with each other; graphs captured on different streams never share one. A capture whose geometry or slab is missing raises RuntimeError. Slabs live for the process unless released with flashinfer.topk_varlen.release_gvr2_resources(device), which synchronizes the device, drops them and returns the bytes freed; graphs captured against them must then be re-captured after a fresh warm-up. A caller-provided "gvr2_workspace" (see workspace) replaces the slab. Paged output (page_table) is decode mode only, hint-free (no pre_idx) and indices only (no return_values); page_size is part of its warm-up key, so warm up each paged geometry with one eager paged call on the capturing stream.

"radix_cutlass" — Masked CUTLASS radix top-K.

Every GPU FlashInfer supports (SM75 and up). fp32, bf16 or fp16; contiguous logits; no top_k limit; any compress_ratio; pre_idx accepted and ignored; padded value slots hold 0. Decode mode only.

"radix_filter" — Filtered-radix top-K.

SM100, SM103 or SM107 with nvidia-cutlass-dsl >= 4.8. fp32, bf16 or fp16; any layout (a non-unit inner stride or a base not 32-byte aligned costs a copy); 1 <= top_k <= 16384; compress_ratio == 1 only; pre_idx accepted and ignored; nondeterministic ties; padded value slots hold -inf. Decode mode only.

"auto" — Automatic selection.

Picks among the entries above as described under “Backend selection”.

type backend:

{“radix”, “gvr”, “gvr_2”, “radix_cutlass”, “radix_filter”, “auto”}, optional

param load_balance:

"gvr" only: chooses between its two kernel paths (see backend). Every other backend ignores it. Default True.

type load_balance:

bool, optional

param workspace:

Caller-provided scratch buffers, keyed by backend (the "gvr_2" slab’s lifetime is under backend):

• "gvr" with load_balance=True: "gvr_order_row" (int32[M], M the smallest power of 2 in [64, 1024] that is >= seq_lens.shape[0]) and "gvr_counters" (int32[2]); without them the path allocates per call, which is safe for any concurrency.

• "gvr_2": "gvr2_workspace", a tensor of at least flashinfer.topk_varlen.kernels.gvr2_topk_host.workspace_bytes() = 20,973,568 bytes, zero-initialized before first use and 16-byte aligned (a smaller or misaligned tensor raises RuntimeError); it replaces the backend’s cached per-stream slab, for decode and windowed calls alike.

A workspace dict must not be shared by launches that can overlap (concurrent streams, or graphs replayed concurrently).

type workspace:

dict, optional

returns:

(indices, values) – Always a 2-tuple. indices is int32[num_rows, top_k]. values holds the selected logits (same dtype as logits) when return_values=True, otherwise None.

rtype:

Tuple[torch.Tensor, Optional[torch.Tensor]]

raises BackendSupportedError:

If backend="auto" finds no backend whose checks pass for the call (which backends serve each mode, and on which GPUs, is under backend), or an explicit backend is not registered for the device. Admission runs before the argument checks below, so such a call raises this error even when a ValueError condition also holds. skip_check=True bypasses admission, not the argument checks, and keeps two refusals that prevent a silently wrong launch: row_starts with a backend that does not serve windowed (prefill) mode, and an explicit "gvr" without a hint.

raises ValueError:

If logits / seq_lens / next_n violate the shape, dtype or grouping contract; if an argument is used outside its mode (next_n, compress_ratio or pre_idx with row_starts; max_window_len or absolute_indices=True without it); if row_starts is not a contiguous 1-D int32 tensor with one entry per row, max_window_len not a positive int or absolute_indices not a bool; if page_table comes with row_starts, pre_idx or return_values, is not a contiguous 2-D int32 tensor with one row per request covering the buffer width, or page_size is given without page_table or is not a power of two up to 2**30; if an out_indices / out_values buffer is outside the contract above; or if an explicit backend’s checker rejects the problem (reported as “Problem size is not supported”).

raises RuntimeError:

Under CUDA-graph capture, when a backend’s launcher for the geometry is not compiled yet or the capturing stream has no workspace slab (the warm-up and slab rules are in that backend’s entry under backend); or when a caller-passed workspace buffer fails its size or alignment check.

Warns:

RuntimeWarning – pre_idx was malformed and has been discarded; the call ran hint-free.

Examples

Decode mode (the default): one row per request, seq_lens in KV tokens, indices are logits columns.

>>> import torch, flashinfer
>>> _ = torch.manual_seed(42)
>>> B, N, top_k = 32, 8192, 1024
>>> logits = torch.randn(B, N, dtype=torch.float32, device="cuda")
>>> seq_lens = torch.randint(top_k + 1, N, (B,), dtype=torch.int32, device="cuda")
>>> indices, _ = flashinfer.top_k_varlen(logits, seq_lens, top_k)
>>> for i in range(B):  # every hit is >= the K-th largest valid score
...     s = seq_lens[i].item()
...     kth = torch.topk(logits[i, :s], top_k).values[-1]
...     assert (logits[i, :s][indices[i].long()] < kth).sum() == 0
>>> # the next step may pass the previous result as a hint
>>> logits_next = torch.randn(B, N, dtype=torch.float32, device="cuda")
>>> indices_next, _ = flashinfer.top_k_varlen(logits_next, seq_lens + 1, top_k, pre_idx=indices)

Decode mode with speculative tokens and KV compression: next_n rows per request (row r belongs to request r // next_n), one column per compress_ratio KV tokens; seq_lens stays in tokens and each row’s column count follows from it (see seq_lens).

>>> next_n, compress_ratio = 2, 4
>>> logits_s = torch.randn(B * next_n, N, dtype=torch.float32, device="cuda")
>>> tokens = torch.randint((top_k + 1) * compress_ratio, N * compress_ratio, (B,), dtype=torch.int32, device="cuda")
>>> indices_s, _ = flashinfer.top_k_varlen(logits_s, tokens, top_k, next_n=next_n, compress_ratio=compress_ratio)
>>> for r in range(B * next_n):
...     i, t = divmod(r, next_n)
...     c = (tokens[i].item() - next_n + t + 1) // compress_ratio  # row r's valid columns
...     kth = torch.topk(logits_s[r, :c], top_k).values[-1]
...     assert (logits_s[r, :c][indices_s[r].long()] < kth).sum() == 0

Windowed (prefill) mode: one row per query token of a packed chunk, each ranking its own window [row_starts[r], row_starts[r] + seq_lens[r]) given in columns; next_n, compress_ratio and pre_idx are not passed (the backends that serve the mode are listed under backend).

>>> # two requests of L keys each on one packed key axis of 2 * L columns;
>>> # a request's j-th token (0-based) ranks the request's first j + 1 keys
>>> R, L, top_k = 2, 2048, 512
>>> logits_p = torch.randn(R * L, R * L, dtype=torch.float32, device="cuda")
>>> row_starts = torch.tensor([q * L for q in range(R) for _ in range(L)], dtype=torch.int32, device="cuda")
>>> seq_lens_p = torch.tensor([j + 1 for _ in range(R) for j in range(L)], dtype=torch.int32, device="cuda")
>>> local, _ = flashinfer.top_k_varlen(logits_p, seq_lens_p, top_k, row_starts=row_starts, max_window_len=L)
>>> columns, _ = flashinfer.top_k_varlen(logits_p, seq_lens_p, top_k, row_starts=row_starts, max_window_len=L, absolute_indices=True)
>>> # the two frames address the same hits: window position + row_starts[r]
>>> # is the logits column (compared through the selected values: a row's
>>> # slot order is not fixed, and a tie at the K-th value may go either way)
>>> hit = local >= 0
>>> assert torch.equal(hit, columns >= 0)
>>> v_pos = logits_p.gather(1, torch.where(hit, local + row_starts[:, None], 0).long())
>>> v_col = logits_p.gather(1, torch.where(hit, columns, 0).long())
>>> assert torch.equal(v_pos.sort(dim=1).values, v_col.sort(dim=1).values)
>>> assert local[0].max() == 0 and (local[0] >= 0).sum() == 1  # a 1-key window: one hit, -1 padded
>>> r = R * L - 1  # the last token of request 1 ranks the request's whole window, columns [L, 2L)
>>> w = logits_p[r, L : 2 * L]
>>> assert (w[local[r].long()] < torch.topk(w, top_k).values[-1]).sum() == 0

See also

flashinfer.top_k

General-purpose radix/clusters top-K (uniform lengths).