File size: 4,704 Bytes
6aac1a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """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"]
|