Card: standardized form with hero

15c529d verified 16 days ago

4.82 kB

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