logits-mask-xpu

Guided-decoding masking and fused sampling for Intel GPUs, loadable through kernels: the two hot-path operations between a model's logits and the next token. The reference baselines are torch.argmax, matched bit-exactly, and the reference softmax, matched empirically to 8.34e-03 over 20,000 draws.

Constrained decoding and sampling run every token of every request, and on Intel hardware the usual implementation hits a wall that has nothing to do with speed: torch.topk raises at a 128k vocabulary on this device, exceeding its work-group size limit. This kernel finds top-k and top-p thresholds by histogram refinement instead of sorting, so a large-vocabulary sampled decode works at all, and the grammar mask costs one bit per token rather than a float tensor.

On an Intel Iris Xe at a 128,256-token vocabulary: greedy at 0.177 ms against eager's 0.252, and the full top-k plus top-p pipeline at 2.19 ms with no eager baseline to compare against, because torch.topk raises on this device at this vocabulary.

Usage

from kernels import get_kernel

lm = get_kernel("phanerozoic/logits-mask-xpu", version=1, trust_remote_code=True)

bits = lm.bitmask_from_allowed(allowed_bool)   # [B, V] bool -> [B, ceil(V/32)] int32
lm.apply_allowed_mask(logits, bits)            # in place, disallowed -> -inf

tokens = lm.sample(logits, temperature=0.8, top_k=40, top_p=0.95)

temperature <= 0 selects greedy argmax. top_k <= 0 disables the top-k stage and top_p >= 1 disables the nucleus stage. Pass uniform [B] for deterministic draws. version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

top-k and top-p thresholds are found by successive histogram refinement rather than by sorting. A 256-bin histogram over the value range locates the bin holding the cutoff, and two further refinements narrow it to about 2^-24 of the original range. That is a fixed six passes over the vocabulary regardless of its size, with no sort and no scratch buffer.

The inverse-CDF draw is cooperative: one contiguous chunk per work-item, an exclusive scan over chunk masses locates the chunk holding the target, and only that chunk is walked. top_k == 1 is dispatched to argmax directly rather than relying on threshold precision to separate the maximum from its nearest neighbour. apply_allowed_mask takes the allow-set as a packed bitmask, one bit per token, so a constrained-decoding grammar's per-step mask costs V / 32 int32 words rather than a float tensor.

Measured

On Intel Iris Xe (Gen12LP, 96 EUs), torch 2.13.0+xpu, 30 iterations after 5 warmup, batch 8, vocabulary 128256:

torch.topk raises on this device at a 128k vocabulary, exceeding its work-group size limit, so there is no eager baseline for the sampled path to compare against: the kernel is the only route. The table is reproduced by benchmarks/bench.py.

Validation

Ten checks on Intel Iris Xe, torch 2.13.0+xpu. Masking is bit-exact: allowed logits are unchanged and every disallowed logit is -inf. Greedy sampling is bit-exact against torch.argmax, masked greedy is verified to land inside the allow-set, and top_k=1 reduces to argmax exactly. Over 40 trials with varying draws, top_k=5 never selects outside the true top-5. Distribution correctness is checked empirically: 20,000 draws over an 8-token vocabulary track the reference softmax to 8.34e-03 maximum frequency deviation; top_p=1 retains the full support, and top_p=0.5 on a peaked distribution collapses to the mode in all 200 draws.

Requirements and limits

• float32 logits, int32 bitmask.
• Intel GPU with a working Level Zero driver and a torch XPU build.
• The threshold is resolved to about 2^-24 of the value range; exact tie-breaking at that boundary is not guaranteed, which is why top_k == 1 dispatches to argmax.
• The ops ship fake implementations, so they trace under torch.compile without breaking the graph.

References

Holtzman et al., "The Curious Case of Neural Text Degeneration" (nucleus sampling, 2020); the XGrammar token-bitmask convention; radix/histogram selection.

License

Apache-2.0.