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', 'radix_cutlass', '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) chooses GVR when available (Blackwell + pre_idx supplied), else the CuTe DSL radix backend on Blackwell, else the radix_cutlass masked fallback. Force a specific backend with backend="radix", backend="gvr", or backend="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.

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; ignored by "radix_cutlass".

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.

type out_indices:

torch.Tensor, optional

param out_values:

Pre-allocated values buffer (same dtype as logits). 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+; 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.
"radix_cutlass" — Masked CUTLASS radix top-K (all GPUs, no: pre_idx needed).
"auto" — GVR (if pre_idx supplied) > radix (Blackwell): > radix_cutlass.

type backend:

{“radix”, “gvr”, “radix_cutlass”, “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,).

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. When workspace is None (default) buffers are allocated locally and are safe for any concurrency.

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 the required inputs (e.g. pre_idx) are missing.

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