Top-k and Sampling Operators¶

本页暂无中文版。以下为英文原文。

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

The top-k selector returns the indices of the largest scores in each row. The masks set every logit a filter drops to -inf, so a softmax over the result renormalizes over what is kept. The samplers draw from a distribution, and take the random seed and offset as tensors so a draw is reproducible.

Top-k selection¶

tileops.attention.TopKSelectFwdOp ¶

The topk highest-scoring key positions of each query row's own window.

Row (b, s, g) selects from index_score[b, s, starts[b, s]:ends[b, s], g].

Two deviations from torch.topk, which returns its indices sorted:

The indices come back in no particular order along the topk axis. Two calls on one input select the same positions and may place them in different slots.
A window holding fewer than topk positions fills the rest with seq_len_kv, one past the last key, which selects nothing.

init ¶

__init__(
    topk,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from each call.

Parameters:

topk (int) –

Manifest params.topk, int.
target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    index_score,
    starts,
    ends,
)

Select each query row's topk highest-scoring keys inside its window.

Parameters:

index_score (Tensor) –

Scores [batch, seq_len, seq_len_kv, kv_group], float32.
starts (Tensor) –

First key of each row's window [batch, seq_len], int32.
ends (Tensor) –

One past the last key of each row's window [batch, seq_len], int32.

Returns:

Tensor –

Selected key positions [batch, seq_len, kv_group, topk], int32.

Logits masks¶

tileops.sampling.TopKMaskFwdOp ¶

Top-k logit filter: keeps the k[b] largest logits of row b, the rest -inf.

Every logit equal to the k[b]-th largest is kept, so a tie at the threshold keeps more than k[b], and a row with k[b] >= V is returned unchanged, as FlashInfer's top_k_mask_logits does. The output has the shape and dtype of logits, and each kept logit is passed through bit for bit.

NaN ranks above every number, as torch.sort places it, and is masked wherever the row is filtered, as logits >= kth is false at a NaN; a row with k[b] >= V keeps its NaNs.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    logits,
    k,
)

Mask each row of logits to its k largest.

Parameters:

logits (Tensor) –

[B, V] logits, float16, bfloat16 or float32.
k (Tensor) –

[B] int32 number of logits each row keeps, at least 1.

Returns:

Tensor –

[B, V] logits of logits' dtype, -inf where masked.

tileops.sampling.MinPMaskFwdOp ¶

Min-p logit filter: masks every token less probable than min_p[b] times row b's top.

A token's probability is below min_p[b] * max(prob[b]) exactly when its logit is below max_logit + log(min_p[b]); those logits become -inf and the rest pass through bit for bit. 0 <= min_p <= 1 is the caller's obligation. The output has the shape and dtype of logits.

min_p = 0 masks nothing, the value vLLM gives a request that disables min-p, and min_p = 1 keeps only the logits equal to the row max. A row holding a NaN has a NaN threshold, so every comparison against it is false and the row passes through whole, as torch.amax and < leave it.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    logits,
    min_p,
)

Mask the tokens of each row of logits below its min_p threshold.

Parameters:

logits (Tensor) –

[B, V] logits, float16, bfloat16 or float32.
min_p (Tensor) –

[B] float32 fraction of the row's top probability a token must reach.

Returns:

Tensor –

[B, V] logits of logits' dtype, -inf where masked.

tileops.sampling.TopPMaskFwdOp ¶

Top-p logit filter: keeps the most probable tokens of row b until they hold p[b].

With prob = softmax(logits[b]), a token survives while the total probability of the tokens strictly more probable than it is below p[b], so tokens tied at the boundary are kept together, as FlashInfer's top_p_renorm_probs keeps them. Masked logits become -inf; the rest pass through bit for bit. 0 <= p <= 1 is the caller's obligation. The output has the shape and dtype of logits.

p = 0 masks the row whole. p = 1 keeps every token whose probability survives the row's float32 sum, and masks the ones the sum rounds away with the zero-probability ones. A row whose maximum is not finite, a NaN or an infinity among its logits, has an all-NaN softmax, so nothing in it is masked and it passes through whole.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    logits,
    p,
)

Mask each row of logits to its top-p nucleus.

Parameters:

logits (Tensor) –

[B, V] logits, float16, bfloat16 or float32.
p (Tensor) –

[B] float32 probability mass each row keeps.

Returns:

Tensor –

[B, V] logits of logits' dtype, -inf where masked.

tileops.sampling.TopKTopPMaskFwdOp ¶

Top-k then top-p logit filter: TopPMaskFwdOp(TopKMaskFwdOp(logits, k), p).

The top-k filter of TopKMaskFwdOp runs first; top-p's probabilities are then the softmax over the tokens it kept, with ties at either boundary kept together. 0 < p < 1 is the caller's obligation. The output has the shape and dtype of logits, and each kept logit is passed through bit for bit, except a NaN, which comes back as a quiet NaN rather than its own payload.

NaN ranks above every number, as torch.sort places it, so it is masked wherever k filters the row; a row with k[b] >= V keeps its NaNs, and a row whose largest surviving logit is NaN or +inf has no finite probability for top-p to compare, so top-p removes nothing from it.

The top-p bound is settled from float32 weights summed in bin order, where the reference sums them in sorted order, so the two can disagree only on a token whose exclusive probability sits within a few rounding steps of p.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    logits,
    k,
    p,
)

Mask each row of logits to its top k, then to the top-p nucleus of those.

Parameters:

logits (Tensor) –

[B, V] logits, float16, bfloat16 or float32.
k (Tensor) –

[B] int32 number of logits each row keeps, at least 1.
p (Tensor) –

[B] float32 probability mass each row keeps among its top k.

Returns:

Tensor –

[B, V] logits of logits' dtype, -inf where masked.

Sampling¶

tileops.sampling.SamplingFromProbsFwdOp ¶

Categorical draw: one index per row, with probability proportional to probs[b].

The row need not sum to 1: the draw is scaled by the row's total. A zero-weight token is never drawn. Each row finite, non-negative and with a positive total is the caller's obligation. The draw is a function of the Philox state (seed, offset): the same pair and inputs give the same samples. Which stream a pair selects belongs to the implementation, so two implementations need not draw the same samples from one pair.

The in-tree kernel takes the stream workloads/sampling.py states: the row's uniform is the first output word of Philox4x32-10 keyed by (seed, offset) and counted by the row index. The counter names no launch fact, so the same pair and inputs give the same samples whatever the batch and whatever grid the call takes.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    probs,
    seed,
    offset,
)

Draw one token per row of probs.

Parameters:

probs (Tensor) –

[B, V] float32 unnormalized probabilities.
seed (Tensor) –

[1] int64 Philox seed.
offset (Tensor) –

[1] int64 Philox offset.

Returns:

Tensor –

[B] int32 drawn indices.

tileops.sampling.ChainSpeculativeSamplingFwdOp ¶

Chain speculative sampling: verifies each request's N draft tokens in order.

Draft token i is accepted with probability min(1, target / draft) at its id, as FlashInfer's chain_speculative_sampling tests u * draft < target. At the first rejection a token is drawn from normalize(max(0, target - draft)) of that position and the rest of the row is -1; when every draft is accepted a bonus token is drawn from target row N. num_accepted is the accepted prefix length, the bonus excluded. Every probability row finite, non-negative and normalized, and each draft token of positive draft probability, is the caller's obligation.

The draws are a function of the Philox state (seed, offset): the same pair and inputs give the same result. Which stream a pair selects belongs to the implementation, so two implementations need not draw the same tokens from one pair.

The in-tree kernel takes the stream workloads/sampling.py states: uniform j of a row is the first output word of Philox4x32-10 keyed by (seed, offset) and counted by (j, row), draws 0 to N - 1 testing the drafts and draw N placing the token. The counter names no launch fact and the acceptance test is the reference's own float32 comparison, so num_accepted and the accepted prefix are the reference's exactly, whatever the batch and whatever grid the call takes.

init ¶

__init__(
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from each call.

Parameters:

target (Target, default: None ) –

Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward ¶

forward(
    draft_probs,
    draft_token_ids,
    target_probs,
    seed,
    offset,
)

Verify the draft chains and draw the token after each accepted prefix.

Parameters:

draft_probs (Tensor) –

[B, N, V] float32 draft model probabilities.
draft_token_ids (Tensor) –

[B, N] int32 draft tokens, each in [0, V).
target_probs (Tensor) –

[B, N + 1, V] float32 target model probabilities.
seed (Tensor) –

[1] int64 Philox seed.
offset (Tensor) –

[1] int64 Philox offset.

Returns:

Tensor –

output_token_ids, [B, N + 1] int32: the accepted drafts, the drawn token,
Tensor –

then -1; and num_accepted, [B] int32.

Top-k and Sampling Operators¶

Top-k selection¶

tileops.attention.TopKSelectFwdOp ¶

__init__ ¶

forward ¶

Logits masks¶

tileops.sampling.TopKMaskFwdOp ¶

__init__ ¶

forward ¶

tileops.sampling.MinPMaskFwdOp ¶

__init__ ¶

forward ¶

tileops.sampling.TopPMaskFwdOp ¶

__init__ ¶

forward ¶

tileops.sampling.TopKTopPMaskFwdOp ¶

__init__ ¶

forward ¶

Sampling¶

tileops.sampling.SamplingFromProbsFwdOp ¶

__init__ ¶

forward ¶

tileops.sampling.ChainSpeculativeSamplingFwdOp ¶

__init__ ¶

forward ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶

init ¶