---
library_name: kernels
license: apache-2.0
---

# 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 sorted softmax is carved by top-k, top-p, and min-p in turn, then four thousand real draws accumulate under the kept bars](https://huggingface.co/kernels/phanerozoic/logits-processor/resolve/main/media/hero.gif)

*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

```python
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.

## References

Gumbel-max sampling; radix select; the XGrammar token-bitmask convention;
the transformers logits-processor chain as the semantic reference.

## License

Apache-2.0.