flashinfer.top_k_varlen¶

flashinfer.top_k_varlen(logits: Tensor, seq_lens: Tensor, top_k: int, pre_idx: Tensor | None = None, compress_ratio: int = 1, next_n: int = 1, 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 over batched decode-step logits.

Selects the top-top_k elements from each row of logits, respecting per-request KV-cache lengths given by seq_lens.

Backend selection¶

backend="auto" (default) ranks the backends that can run the call by shape and dtype (see _top_k_varlen_heuristic): gvr_2 for fp32 on datacentre Blackwell-class GPUs, with or without a hint (a missing hint costs ~10% on most shapes), radix_filter in the measured large-N regions, gvr for large hinted half-precision batches of mid-length rows, otherwise the CuTe DSL radix backend on Blackwell, with the radix_cutlass masked fallback in its fp32 big corner and on every other GPU. Force a specific backend with backend="radix", "gvr", "gvr_2", "radix_filter" or "radix_cutlass".

param logits:

2-D float tensor of shape (num_rows, max_seq_len). Supported dtypes: float32, bfloat16, float16. For the "gvr" backend the row width max_seq_len must be a multiple of 16 // itemsize (8 for fp16/bf16, 4 for fp32) so each row is 16-byte aligned for GVR’s 128-bit vectorized loads; a ValueError is raised otherwise. The "radix_cutlass" backend has no such constraint.

type logits:

torch.Tensor

param seq_lens:

1-D int32 tensor of shape (num_rows // next_n,) with the effective KV-cache length per request, in KV-token units regardless of compress_ratio (the raw sequence length, not a length already divided by the compression factor). Row t of request i (t in [0, next_n)) ranks its first max(0, (seq_lens[i] - next_n + t + 1) // compress_ratio) logit columns (padded or evicted slots can make the numerator negative; such rows have no valid columns and are emitted as all -1); logits at or beyond that count are excluded from the search. Passing lengths already expressed in compressed / logit-column units together with compress_ratio > 1 would divide the search range twice. A row whose length exceeds the logits width (the dynamic length has outgrown the static buffer, e.g. under CUDA-graph replay) is clamped to the width by every backend: only the scores present in the buffer are ranked.

type seq_lens:

torch.Tensor

param top_k:

Number of top elements per row. GVR backend supports {512, 1024, 2048}; radix backend has no restriction.

type top_k:

int

param pre_idx:

int32[num_rows // next_n, top_k] — top-K KV-cache indices selected by this same layer at the previous token’s decode step. GVR exploits the strong correlation between a layer’s attention pattern at step t and step t+1; the kernel internally applies a +1 offset (DSv3.2) so the previous step’s indices land correctly in the current step’s grown KV-cache space. pre_idx[:, 0] must be the argmax index. Required by the "gvr" backend; optional for "gvr_2" (the kernels use it only as a sampling anchor, so a hint-free call is exact and only somewhat slower); ignored by the radix backends. Must be a contiguous, 16-byte-aligned int32 CUDA tensor on logits.device of shape [seq_lens.shape[0], top_k]: a hint that violates this is discarded with a RuntimeWarning and the call runs hint-free (gvr_2 with its hint-free engines; an explicit "gvr" request is refused).

type pre_idx:

torch.Tensor, optional

param compress_ratio:

KV-index compression factor (1 for DSv3.2, 4 for DSv4): every logit column stands for compress_ratio consecutive KV tokens. seq_lens stay in KV-token units; the kernels derive each row’s valid column count as max(0, (seq_len - next_n + t + 1) // compress_ratio) (see seq_lens). Default 1.

type compress_ratio:

int, optional

param next_n:

Speculative-decode temporal stride. Default 1.

type next_n:

int, optional

param return_values:

When True also return the selected logit values. Default False.

type return_values:

bool, optional

param out_indices:

Pre-allocated int32[num_rows, top_k] output buffer: contiguous, 16-byte aligned, on logits.device, of that shape or a 1-D buffer of exactly num_rows * top_k elements (viewed in place and returned as its [num_rows, top_k] view over the same storage, not as the caller’s object); a wider, taller or otherwise shaped buffer raises ValueError.

type out_indices:

torch.Tensor, optional

param out_values:

Pre-allocated values buffer (same dtype as logits, same layout rules as out_indices). Only used when return_values=True.

type out_values:

torch.Tensor, optional

param backend:

Backend to use. Default "auto".

"radix" — CuTe DSL single-pass multi-CTA radix top-K: (Blackwell sm_100+ incl. Rubin sm_107; native varlen, no pre_idx, no logit masking).
"gvr" — GVR kernel (Blackwell sm_100+ only; requires: pre_idx). load_balance selects the LB vs single-CTA path.
"gvr_2" — self-sampling GVR V2 (TRT-LLM PR #17821 port):: sample-calibrated threshold ladders, exact tie-interchangeable top-K, one launch per batch. Datacentre Blackwell-class only (sm_100/103, or Rubin sm_107); fp32 logits; pre_idx optional (hints steer sampling, never exactness — without one the kernels run hint-free on the identity sample); top_k in {512, 1024, 2048}. load_balance is ignored (the kernel families load-balance internally). CUDA graphs: warm up each (num_rows, N, top_k, next_n, compress_ratio, hinted-or-not) geometry with one eager call, ON THE STREAM THAT WILL CAPTURE, before capture: the eager call compiles the launcher (hinted and hint-free are distinct compiled engines, so hint mode IS a warm-up dimension) and creates that stream’s default workspace slab (one per device and stream, so graphs captured on different streams never share scratch); either missing raises loudly under capture; replays may change seq_lens CONTENTS freely in either direction. All finite values, +inf and -inf are tie-aware exact (TRT-LLM #18501/#18625 ported); NaN ordering is implementation-specific.
"radix_cutlass" — Masked CUTLASS radix top-K (all GPUs, no: pre_idx needed).
"radix_filter" — DKG filtered-radix top-K (coarse histogram →: filter → on-chip refine): hint-free like "radix", fp32/fp16/bf16, top_k in [1, 16384], compress_ratio == 1 only, row-relative indices only (no fused page-table transform, nondeterministic ties). Datacentre Blackwell-class (sm_100/103/107); requires nvidia-cutlass-dsl >= 4.8. pre_idx is accepted and ignored (the kernel takes no hint), as for "radix" and "radix_cutlass".
"auto" — shape/dtype-aware selection tracking the: measured per-config winner: gvr_2 for fp32, hinted or hint-free (except hint-free top_k >= 2048 at N <= 4096 for a single row, and N <= top_k where every row is the identity answer: radix_filter wins both); radix_filter for the remaining hint-free fp32 cells from 32K columns up (all N once batch >= 256) and for the bf16/fp16 mid/large regions; gvr for fp32 when batch * max_seq_len >= 2^22 and for half precision only when batch >= 256 with 32K-512K columns; radix otherwise, with radix_cutlass preferred in its fp32 big-batch/long-row corner. Batches whose rows are mostly far shorter than max_seq_len are a known blind spot (auto cannot read seq_lens without a sync); prefer backend="radix" there.

type backend:

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

param load_balance:

Selects the GVR kernel path (ignored by the radix backend). Default True.

True (default) — two-kernel LB path (GvrTopKLBPrepareKernel +: GvrTopKLBKernel): a prepare kernel classifies requests into long/short buckets, then the main kernel splits each long row across a CTA cluster and packs short rows. Best for the ragged decode batches GVR targets.
False — single-kernel path (GvrTopKKernel): one CTA per row,: no prepare step. Faster when the batch has no length variance (all rows short, or all long).

Both settings are CUDA-graph safe (no host branch on device data).

type load_balance:

bool, optional

param workspace:

Reusable workspace buffers for the GVR load_balance=True path. When provided, the LB prepare and decode kernels read/write these tensors instead of allocating fresh ones per call — useful in decode loops where the same batch size is reused at every step.

Required keys (both int32 on the same device as logits):

"gvr_order_row": shape (M,) where M is the smallest power of 2 in [64, 1024] that is >= seq_lens.shape[0].
"gvr_counters": shape (2,).

For the "gvr_2" backend the optional key "gvr2_workspace" (CUDA tensor of at least flashinfer.topk_varlen.kernels.gvr2_topk_host.workspace_bytes() = 20,973,568 bytes, zero-initialized before first use, 16-byte aligned) overrides the cached default slab.

Warning

Do not share the same workspace dict across concurrent CUDA streams — each stream must have its own workspace to avoid races on the device tensors. For the "gvr" backend, workspace=None (default) allocates per call and is safe for any concurrency. For "gvr_2", workspace=None resolves to one slab per device and CUDA stream: it holds the cross-CTA counters, offsets and candidate buffer of the multi-CTA streaming path (selected by the kernel’s route() from row count, envelope and top_k; small batches with long rows), so launches that can overlap must not share it — and with per-stream slabs, two eager streams or two CUDA graphs captured on different streams and replayed concurrently do not. The slab is created by the first eager launch on a stream; a graph captured on a stream that has never run an eager gvr_2 launch (and whose plan needs the slab) raises under capture instead of allocating from the graph pool. Graphs captured on the SAME stream share that stream’s slab and must not be replayed concurrently with each other; pass a private "gvr2_workspace" for that pattern. Slabs are keyed by the raw CUDA stream handle: there is one 21 MB slab per distinct handle that has run an eager slab-using launch. PyTorch hands out torch.cuda.Stream() objects from a fixed pool per device and priority (two Stream objects may share one slab when the pool wraps — they also share the raw stream, so their launches are ordered), but stream handles created outside that pool (external streams, other priorities) each add a slab, so the total is not bounded by the pool size. The 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 eager warm-up.

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 the requested backend is not supported on the current device, or an explicit "gvr" request has no usable pre_idx.

raises ValueError:

If logits / seq_lens / next_n violate the shape, dtype or grouping contract, an out_indices / out_values buffer is outside the contract above, or an explicit backend’s checker rejects the problem (reported as “Problem size is not supported”).

Warns:

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

Examples

>>> import torch, flashinfer
>>> torch.manual_seed(42)
>>> B, N_max, top_k = 32, 8192, 1024
>>>
>>> # Step t: no prior indices; use radix to get the first top-K.
>>> # Each request has a different KV-cache length in [top_k+1, N_max-1].
>>> logits = torch.randn(B, N_max, dtype=torch.bfloat16, device="cuda")
>>> seq_lens_t = torch.randint(top_k + 1, N_max, (B,), dtype=torch.int32, device="cuda")
>>> indices_t, _ = flashinfer.top_k_varlen(logits, seq_lens_t, top_k, backend="radix_cutlass")
>>> # Reference check: every selected value must be >= the K-th largest.
>>> for i in range(B):
...     s = seq_lens_t[i].item()
...     kth = torch.topk(logits[i, :s].float(), top_k).values[-1]
...     assert (logits[i, :s].float()[indices_t[i].long()] < kth - 1e-5).sum() == 0
>>>
>>> # Step t+1: one new token appended per request; seq_lens grows by 1.
>>> logits_t1 = torch.randn(B, N_max, dtype=torch.bfloat16, device="cuda")
>>> seq_lens_t1 = seq_lens_t + 1
>>> # Pass indices_t as pre_idx; GVR uses it to warm-start the threshold search.
>>> indices_t1, _ = flashinfer.top_k_varlen(logits_t1, seq_lens_t1, top_k, pre_idx=indices_t)
>>> for i in range(B):
...     s = seq_lens_t1[i].item()
...     kth = torch.topk(logits_t1[i, :s].float(), top_k).values[-1]
...     assert (logits_t1[i, :s].float()[indices_t1[i].long()] < kth - 1e-5).sum() == 0