| """logits-processor: the generation last-mile, loadable through `kernels`. |
| |
| Two fused ops that every serving stack runs on the hot path: |
| |
| - ``apply_token_bitmask`` - structured / guided decoding. Set the logits of |
| disallowed tokens to -inf from a packed allow-mask (XGrammar convention: |
| bit 1 = allowed), so the next softmax cannot pick them. This is the kernel |
| behind constrained JSON and tool-calling output. |
| - ``sample`` - temperature, then the intersection of top-k / top-p (nucleus) / |
| min-p filtering, then a multinomial draw, fused into one logits -> token op |
| with a counter-based RNG that is deterministic in ``(seed, row, offset)``. |
| ``temperature <= 0`` is greedy (argmax). |
| |
| Both take ``[B, V]`` logits in float / half / bf16. Sampling parameters are |
| per-row: pass a scalar to broadcast, or a length-B tensor for a heterogeneous |
| batch (continuous batching, per-request grammars). |
| """ |
| from typing import Union |
|
|
| import torch |
|
|
| from ._ops import ops |
|
|
| __all__ = ["apply_token_bitmask", "sample"] |
|
|
| Param = Union[float, int, torch.Tensor] |
|
|
|
|
| def _row(x: Param, B: int, device, dtype) -> torch.Tensor: |
| if isinstance(x, torch.Tensor): |
| return x.to(device=device, dtype=dtype).contiguous() |
| return torch.full((B,), x, device=device, dtype=dtype) |
|
|
|
|
| def apply_token_bitmask(logits: torch.Tensor, bitmask: torch.Tensor) -> torch.Tensor: |
| """In place: set disallowed tokens' logits to -inf. |
| |
| logits: ``[B, V]`` (float/half/bf16). bitmask: ``[B, ceil(V/32)]`` int32, |
| bit 1 = token allowed. Returns ``logits`` for chaining. |
| """ |
| ops.apply_token_bitmask_inplace(logits, bitmask.to(torch.int32).contiguous()) |
| return logits |
|
|
|
|
| def sample(logits: torch.Tensor, temperature: Param = 1.0, top_k: Param = 0, |
| top_p: Param = 1.0, min_p: Param = 0.0, seed: int = 0, |
| offset: int = 0) -> torch.Tensor: |
| """Fused temperature + top-k ∩ top-p ∩ min-p + multinomial sample. |
| |
| top_k <= 0, top_p >= 1, and min_p <= 0 each disable that filter. |
| temperature <= 0 selects the argmax. Returns int64 token ids ``[B]``. |
| Same ``(seed, offset)`` and logits give the same tokens. |
| """ |
| assert logits.dim() == 2, "logits must be [B, V]" |
| B, dev = logits.size(0), logits.device |
| return ops.sample( |
| logits.contiguous(), |
| _row(temperature, B, dev, torch.float32), |
| _row(top_p, B, dev, torch.float32), |
| _row(top_k, B, dev, torch.int32), |
| _row(min_p, B, dev, torch.float32), |
| int(seed), int(offset), |
| ) |
|
|
