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 hinted fp32 on datacentre Blackwell-class GPUs, 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. Logits at or beyond seq_lens[i] are excluded from the search. 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" and "gvr_2" backends; 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 (auto picks a hint-free backend; an explicit "gvr" / "gvr_2" request is refused).

type pre_idx:

torch.Tensor, optional

param compress_ratio:

KV-index compression factor (1 for DSv3.2, 4 for DSv4). 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); requires pre_idx (hints steer sampling, never exactness) and fp32 logits; 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) geometry with one eager call before capture (an uncompiled launcher 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 hinted fp32; radix_filter for hint-free fp32 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 per-device cached 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, shared by every launch on that device: 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 two such launches in flight at once — two eager streams, or two CUDA graphs replayed concurrently — race on it. Every concurrently active stream, and every CUDA graph that may replay concurrently with another launch, must pass its own "gvr2_workspace"; stream-ordered use needs nothing.

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" / "gvr_2" 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

See also

flashinfer.top_k

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