| """spec-decode-ops: fused sampling and speculative-decoding verification. |
| |
| sample() executes the transformers logits pipeline (repetition penalty -> |
| temperature -> top-k -> top-p -> min-p -> categorical draw) in a fused form: |
| one optional segmented sort plus one kernel, sampling by the Gumbel-argmax |
| identity. Sort-free fast paths cover greedy decoding and unfiltered sampling. |
| |
| verify() implements canonical speculative-decoding rejection sampling |
| (Leviathan et al.; Chen et al.). Draft token x_i is accepted with probability |
| min(1, p_target(x_i) / p_draft(x_i)); the first rejection is replaced by a |
| sample from the renormalized residual max(p_target - p_draft, 0); full |
| acceptance appends a bonus token from the target's final position. The |
| emitted token stream is distributed exactly as the target distribution. |
| |
| Randomness is counter-based Philox parameterized by (seed, offset): results |
| are bitwise reproducible for fixed seed, offset, shape, and architecture. |
| """ |
| from typing import Optional, Tuple |
|
|
| import torch |
|
|
| from ._ops import ops |
|
|
|
|
| def sample( |
| logits: torch.Tensor, |
| temperature: float = 1.0, |
| top_k: int = 0, |
| top_p: float = 1.0, |
| min_p: float = 0.0, |
| repetition_penalty: float = 1.0, |
| prev_tokens: Optional[torch.Tensor] = None, |
| seed: int = 0, |
| offset: int = 0, |
| ) -> torch.Tensor: |
| """Sample one token per row of logits. |
| |
| Args: |
| logits: [M, V] f32/bf16/f16 CUDA tensor. |
| temperature: 0 selects greedy argmax; otherwise logits are divided by it. |
| top_k: keep the k highest-probability tokens (0 disables). |
| top_p: nucleus filtering; keep the smallest prefix of the sorted |
| distribution with cumulative probability >= top_p (1.0 disables). |
| min_p: keep tokens with p >= min_p * p_max (0 disables). |
| repetition_penalty: applied to prev_tokens before temperature |
| (l/r for l > 0 else l*r), matching transformers. |
| prev_tokens: [M, P] int64, -1 padded; required if penalty != 1. |
| seed, offset: Philox counter parameters. |
| Returns: |
| [M] int64 token ids. |
| """ |
| logits = logits.contiguous() |
| tokens = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device) |
| ops.fused_sample(tokens, logits, prev_tokens, float(temperature), int(top_k), |
| float(top_p), float(min_p), float(repetition_penalty), |
| int(seed), int(offset)) |
| return tokens |
|
|
|
|
| def filter_logits( |
| logits: torch.Tensor, |
| temperature: float = 1.0, |
| top_k: int = 0, |
| top_p: float = 1.0, |
| min_p: float = 0.0, |
| repetition_penalty: float = 1.0, |
| prev_tokens: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| """Return the filtered logits (temperature-scaled, masked to -inf), the |
| same tensor the transformers processor chain would hand to the sampler. |
| Useful for composing filtered distributions with verify().""" |
| logits = logits.contiguous() |
| out = torch.empty(logits.shape, dtype=torch.float32, device=logits.device) |
| ops.fused_filter(out, logits, prev_tokens, float(temperature), int(top_k), |
| float(top_p), float(min_p), float(repetition_penalty)) |
| return out |
|
|
|
|
| def verify( |
| target_logits: torch.Tensor, |
| draft_logits: torch.Tensor, |
| draft_tokens: torch.Tensor, |
| temperature: float = 1.0, |
| seed: int = 0, |
| offset: int = 0, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| """Speculative-decoding verification. |
| |
| Args: |
| target_logits: [B, k+1, V] target-model logits over the draft |
| positions plus one (the bonus position). |
| draft_logits: [B, k, V] draft-model logits at the same positions. |
| draft_tokens: [B, k] int64 tokens the draft model sampled. |
| temperature: sampling temperature applied to both distributions; |
| 0 selects greedy verification (accept while target argmax equals |
| the draft token). |
| Returns: |
| (accept_len [B] int64, out_tokens [B, k+1] int64). For sequence b, |
| out_tokens[b, :accept_len[b]] are accepted draft tokens and |
| out_tokens[b, accept_len[b]] is the correction or bonus token, so |
| accept_len[b] + 1 tokens are always emitted. |
| """ |
| B, kp1, V = target_logits.shape |
| k = kp1 - 1 |
| accept_len = torch.empty(B, dtype=torch.int64, device=target_logits.device) |
| out_tokens = torch.empty(B, k + 1, dtype=torch.int64, device=target_logits.device) |
| ops.spec_verify(accept_len, out_tokens, target_logits.contiguous(), |
| draft_logits.contiguous(), draft_tokens.contiguous(), |
| float(temperature), int(seed), int(offset)) |
| return accept_len, out_tokens |
|
|
|
|
| __all__ = ["sample", "filter_logits", "verify"] |
|
|
