logits-processor

Guided-decoding masking and fused sampling, loadable through kernels: the two hot-path operations between a model's logits and the next token. The reference baselines are the transformers processor chain (kept sets exact on non-degenerate inputs) and a sort-based sampler (distributions matched to 1e-3).

Every decoded token ends with the same two steps: zero out whatever a grammar disallows, then filter and draw. Done eagerly that is a chain of full-vocabulary kernels, sorts included, repeated every token of every request. This kernel does the whole thing in one pass with no sort: the grammar mask applies from a packed bitmask at one bit per token, the intersection of top-k, top-p, and min-p resolves as a single keep-threshold by radix-select, and the token draws by Gumbel-max, with per-row parameters so a continuous batch can carry a different grammar and temperature per request.

The distribution carved live: top-k, then the top-p waterline, then the min-p floor, leaving the keep set; 4,000 kernel draws land with zero outside the keep set, total variation 0.021 against the renormalized softmax, and the full pipeline runs 1.13 ms at a 128,256-token vocabulary; bitmasked tokens are never drawn.

Usage

from kernels import get_kernel
lp = get_kernel("phanerozoic/logits-processor", version=1, trust_remote_code=True)

# 1) structured / guided decoding: -inf the tokens a grammar disallows
lp.apply_token_bitmask(logits, bitmask)          # logits [B, V], bitmask [B, ceil(V/32)] int32

# 2) fused temperature + top-k, top-p, min-p intersection + multinomial draw
tokens = lp.sample(logits, temperature=0.7, top_k=50, top_p=0.9, min_p=0.0, seed=step)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Symbol	Purpose
`apply_token_bitmask(logits, bitmask)`	set disallowed logits to `-inf` in place; bit 1 = allowed (XGrammar convention)
`sample(logits, temperature, top_k, top_p, min_p, seed, offset)`	one `logits -> token` op: scaling, filter intersection, draw

top_k <= 0, top_p >= 1, min_p <= 0 each disable that filter; temperature <= 0 is greedy.
Every sampling parameter is per-row: scalar to broadcast, or a length-B tensor for a heterogeneous batch.
The draw is counter-based, keyed by (seed, row, offset): deterministic and replayable; advance offset per decoding step.

Method

The three filters all keep the largest-probability tokens above a cutoff, so their intersection is a single keep-threshold and no ordering of the vocabulary is needed. One threadblock handles each row: a parallel reduction takes the max logit (and the argmax, for greedy); a second pass forms the softmax; the keep-threshold is found by radix-select over the softmax values (a per-warp-privatized histogram, four 8-bit passes resolving the exact boundary); and the token is drawn by Gumbel-max over the kept set. Greedy and unfiltered sampling skip the threshold step. There is no sort and no per-row serial scan.

Measured

On H200 against a sort-based reference (segmented radix sort plus per-row cutoff and inverse-CDF draw), batch 256, 128k vocabulary: greedy 49x, unfiltered sampling 59x, top-p 9.9x, top-k 2.4x; top-p at batch 1024 and a 32k vocabulary 5.9x. Every path is faster with tokens distributed identically, matching the sort-based sampler to 1e-3.

Correctness

Verified through get_kernel on A100 (sm80), A10G (sm86), L4 (sm89), H200 (sm90), and RTX PRO 6000 (sm120):

the bitmask sets exactly the disallowed logits to -inf in float, half, and bf16, including vocabularies not a multiple of 32;
greedy and top_k=1 return the argmax; the sampled token always lies in the exact top-k / top-p / min-p keep-set and their intersection;
the empirical sampling distribution matches the filtered, renormalized softmax to within 0.002 (0.0003 at a 32k vocabulary);
draws are deterministic in (seed, offset); per-row parameters are honored independently across a batch; a bitmask followed by sample never draws a disallowed token.

Requirements and limits

Built for sm80 / sm86 / sm89 / sm90 / sm100 / sm120; certified through get_kernel on all except sm100.
Scratch is one [B, V] buffer; logits in f32, half, or bf16.