logits-processor

The generation last-mile as loadable kernels: guided-decoding masking and fused sampling, the two hot-path ops between a model's logits and the next token that the Kernel Hub did not carry.

from kernels import get_kernel
lp = get_kernel("phanerozoic/logits-processor")

# 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 + multinomial draw
tokens = lp.sample(logits, temperature=0.7, top_k=50, top_p=0.9, min_p=0.0, seed=step)

Ops

`apply_token_bitmask(logits, bitmask) -> logits`

Sets the logits of disallowed tokens to -inf, in place, from a packed allow-mask. Bit 1 means the token is allowed (XGrammar convention). logits is [B, V] in float/half/bf16; bitmask is [B, ceil(V/32)] int32. This is the op behind constrained JSON and tool-calling output — apply it, then sample.

`sample(logits, temperature=1.0, top_k=0, top_p=1.0, min_p=0.0, seed=0, offset=0) -> tokens`

One logits -> token op: per-row temperature scaling, then the intersection of top-k, top-p (nucleus), and min-p filtering, then a multinomial draw.

top_k <= 0, top_p >= 1, min_p <= 0 each disable that filter.
temperature <= 0 is greedy (argmax).
Every sampling parameter is per-row: pass a scalar to broadcast, or a length-B tensor for a heterogeneous batch — continuous batching with a different grammar, temperature, or nucleus per request.
The draw is a counter-based RNG keyed by (seed, row, offset): deterministic and replayable. Advance offset (or seed) per decoding step.
Returns int64 token ids [B].

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 / bf16, including a vocabulary size that is 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 in their intersection, checked against an independent reference over many draws;
the empirical sampling distribution matches the filtered, renormalized softmax reference to within 0.002 (0.0003 at a 32k vocabulary);
draws are deterministic in (seed, offset) and change when offset advances;
per-row temperature / top-k / top-p / min-p are honored independently across a batch, and a bitmask followed by sample never draws a disallowed token.

How it works

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 a 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 a Gumbel-max over the kept set. Greedy is a bare argmax and unfiltered sampling is a single Gumbel-max pass, both skipping the threshold step. There is no sort and no per-row serial scan, and the only scratch is one [B, V] buffer.

Performance

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

Architectures

Built for sm80 / sm86 / sm89 / sm90 / sm100 / sm120; certified through get_kernel on all of these except sm100, for which no B200 was available to run the check.