"""Composable sampling strategies for logit transformation. Implements the Strategy pattern: each sampling technique (temperature, top-k, top-p, frequency penalty) is a pluggable strategy that can be composed into a pipeline. All strategies accept both scalar and per-sample tensor parameters, so a single pipeline works for any batch size. """ from abc import ABC, abstractmethod from typing import List, Optional, Union import torch from torch import Tensor class BaseSamplingStrategy(ABC): """Abstract base for a logit transformation strategy.""" @abstractmethod def apply( self, logits: Tensor, filter_value: float = -float("inf"), input_ids: Optional[Tensor] = None, input_mask: Optional[Tensor] = None, ) -> Tensor: """Applies the strategy to logits. Args: logits: Raw logits tensor (batch, vocab_size). filter_value: Value assigned to filtered-out positions. input_ids: Previously generated token IDs ``[batch, seq_len]``, padded with 0. Used by frequency penalty. input_mask: Boolean mask ``[batch, seq_len]``, True for real tokens, False for padding. Used to exclude padding from penalty computation. Returns: Transformed logits tensor. """ raise NotImplementedError class TemperatureStrategy(BaseSamplingStrategy): """Divides logits by temperature to control randomness. Args: temperature: Scalar or ``[batch]`` tensor. """ def __init__(self, temperature: Union[float, Tensor] = 1.0): self.temperature = temperature def apply( self, logits: Tensor, filter_value: float = -float("inf"), input_ids: Optional[Tensor] = None, input_mask: Optional[Tensor] = None, ) -> Tensor: t = self.temperature if isinstance(t, Tensor): t = t.to(logits.device, non_blocking=True).view(-1, 1) t = torch.clamp(t, min=1e-8) if (t != 1.0).any(): logits = logits / t elif t != 1.0: logits = logits / max(t, 1e-8) return logits class TopKStrategy(BaseSamplingStrategy): """Keeps only the top-k logits, setting the rest to filter_value. Args: top_k: Scalar or ``[batch]`` tensor (0 disables). """ def __init__(self, top_k: Union[int, Tensor] = 0): self.top_k = top_k def apply( self, logits: Tensor, filter_value: float = -float("inf"), input_ids: Optional[Tensor] = None, input_mask: Optional[Tensor] = None, ) -> Tensor: tk = self.top_k if isinstance(tk, Tensor): tk = tk.to(logits.device, non_blocking=True).long().clamp(min=0) max_k = int(tk.max().item()) if max_k <= 0: return logits max_k = min(max_k, logits.size(-1)) values, _ = torch.topk(logits, max_k, dim=-1) per_row_k = tk.clamp(max=max_k) thresholds = torch.full_like(logits[..., -1:], -float("inf")) positive = per_row_k > 0 if positive.any(): row_idx = torch.arange(logits.size(0), device=logits.device)[positive] thresholds[positive] = values[ row_idx, per_row_k[positive] - 1 ].unsqueeze(-1) logits[logits < thresholds] = filter_value return logits if tk > 0: k = min(tk, logits.size(-1)) thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:] logits[logits < thresholds] = filter_value return logits class TopPStrategy(BaseSamplingStrategy): """Nucleus (top-p) filtering: keeps the smallest set of tokens whose cumulative probability exceeds top_p. Args: top_p: Scalar or ``[batch]`` tensor (1.0 disables). """ def __init__(self, top_p: Union[float, Tensor] = 1.0): self.top_p = top_p def _apply( self, logits: Tensor, top_p: Union[float, Tensor], filter_value: float ) -> Tensor: sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) remove = cum_probs > top_p remove[..., 1:] = remove[..., :-1].clone() remove[..., 0] = False mask = torch.zeros_like(logits, dtype=torch.bool) mask.scatter_(1, sorted_indices, remove) logits[mask] = filter_value return logits def apply( self, logits: Tensor, filter_value: float = -float("inf"), input_ids: Optional[Tensor] = None, input_mask: Optional[Tensor] = None, ) -> Tensor: tp = self.top_p if isinstance(tp, Tensor): tp = tp.to(logits.device, non_blocking=True) if (tp < 1.0).any(): logits = self._apply(logits, tp.view(-1, 1), filter_value) elif tp < 1.0: logits = self._apply(logits, tp, filter_value) return logits class FrequencyPenaltyStrategy(BaseSamplingStrategy): """Penalizes tokens based on how many times they appeared in history. Subtracts ``penalty * count(token)`` from each token's logit, where ``count(token)`` is the number of occurrences in the generation history (prompt + output). A penalty of ``0.0`` disables the strategy. Unlike repetition penalty (which only checks *presence*), frequency penalty scales linearly with occurrence count: the first use is penalized once, the third use three times. This allows natural repetition of common words while suppressing degenerate loops. Reference: OpenAI API ``frequency_penalty`` parameter. Args: penalty: Scalar or ``[batch]`` tensor (0.0 disables, range -2.0~2.0). """ def __init__(self, penalty: Union[float, Tensor] = 0.0): self.penalty = penalty def apply( self, logits: Tensor, filter_value: float = -float("inf"), input_ids: Optional[Tensor] = None, input_mask: Optional[Tensor] = None, ) -> Tensor: if input_ids is None: return logits p = self.penalty if isinstance(p, Tensor): p = p.to(logits.device, non_blocking=True).view(-1, 1) if (p == 0.0).all(): return logits elif p == 0.0: return logits input_ids = input_ids.to(logits.device, non_blocking=True) if input_mask is not None: input_mask = input_mask.to(logits.device, non_blocking=True) masked_ids = input_ids.clone() masked_ids[~input_mask] = -1 else: masked_ids = input_ids batch_sz, seq_len = masked_ids.shape vocab_size = logits.size(-1) if isinstance(p, Tensor): penalty_per_row = p.expand(batch_sz, 1) else: penalty_per_row = torch.full( (batch_sz, 1), float(p), device=logits.device, dtype=logits.dtype ) counts = torch.zeros( batch_sz, vocab_size, device=logits.device, dtype=logits.dtype ) valid_mask = masked_ids >= 0 if valid_mask.any(): valid_ids = masked_ids[valid_mask] row_indices = ( torch.arange(batch_sz, device=logits.device) .unsqueeze(1) .expand_as(masked_ids)[valid_mask] ) counts.index_put_( (row_indices, valid_ids), torch.ones_like(valid_ids, dtype=logits.dtype), accumulate=True, ) return logits - penalty_per_row * counts class SamplingPipeline(BaseSamplingStrategy): """Composes multiple sampling strategies into a single transformation. Strategies are applied sequentially in the order they are provided, matching the original temperature -> top-k -> top-p ordering. Usage:: pipeline = SamplingPipeline([ TemperatureStrategy(0.8), TopKStrategy(50), TopPStrategy(0.95), ]) logits = pipeline.apply(logits) token = pipeline.sample(logits) # softmax + multinomial """ def __init__(self, strategies: List[BaseSamplingStrategy]): self.strategies = strategies def apply( self, logits: Tensor, filter_value: float = -float("inf"), input_ids: Optional[Tensor] = None, input_mask: Optional[Tensor] = None, ) -> Tensor: for strategy in self.strategies: logits = strategy.apply(logits, filter_value, input_ids, input_mask) return logits @staticmethod def _is_greedy(temperature: Union[float, Tensor]) -> bool: if isinstance(temperature, Tensor): return temperature.numel() == 1 and temperature.item() == 0 return temperature == 0 @torch.inference_mode() def sample( self, logits: Tensor, filter_value: float = -float("inf"), input_ids: Optional[Tensor] = None, input_mask: Optional[Tensor] = None, return_logprobs: bool = False, ): """Apply strategies then sample (softmax + multinomial). Short-circuits to ``argmax`` when temperature is exactly 0 (deterministic / greedy decode). Args: logits: Raw logits ``[batch, vocab_size]``. input_ids: Previously generated token IDs ``[batch, seq_len]``. input_mask: Boolean mask for ``input_ids`` padding. return_logprobs: If ``True``, return ``(tokens, logprobs)`` where ``logprobs[i]`` is the log-probability of ``tokens[i]`` under the (post-strategy) sampling distribution. Returns: Sampled token IDs ``[batch]``, or — when ``return_logprobs`` is ``True`` — a ``(token_ids, chosen_logprobs)`` tuple. """ if self._is_greedy_pipeline(): tokens = logits.argmax(dim=-1) if not return_logprobs: return tokens log_probs = torch.log_softmax(logits.float(), dim=-1) chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1) return tokens, chosen transformed = self.apply(logits, filter_value, input_ids, input_mask) log_probs = torch.log_softmax(transformed.float(), dim=-1) tokens = torch.multinomial( torch.softmax(transformed, dim=-1), num_samples=1 ).squeeze(-1) if not return_logprobs: return tokens chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1) return tokens, chosen def _is_greedy_pipeline(self) -> bool: """True if the first strategy is greedy temperature (temp=0).""" if not self.strategies: return False first = self.strategies[0] return isinstance(first, TemperatureStrategy) and self._is_greedy( first.temperature ) @torch.inference_mode() def sample( logits: Tensor, temperature: Union[float, Tensor] = 1.0, top_k: Union[int, Tensor] = 0, top_p: Union[float, Tensor] = 1.0, frequency_penalty: Union[float, Tensor] = 0.0, input_ids: Optional[Tensor] = None, input_mask: Optional[Tensor] = None, filter_value: float = -float("inf"), return_logprobs: bool = False, ): """Apply sampling strategies then sample (softmax + multinomial). Shortcut for ``SamplingPipeline(...).sample(logits, return_logprobs=)``. When **temperature** is exactly 0 (scalar or single-element tensor) the function short-circuits to ``argmax`` for deterministic decode. When **frequency_penalty** is 0 (the common decode case), the entire frequency penalty computation — including the O(batch * vocab) count tensor allocation — is skipped. Args: logits: Raw logits ``[batch, vocab_size]``. frequency_penalty: Penalty per occurrence for repeated tokens (0.0 disables, range -2.0~2.0). input_ids: Previously generated token IDs ``[batch, seq_len]``. input_mask: Boolean mask for ``input_ids`` padding. return_logprobs: If ``True``, also return the log-probability of each sampled token under the (post-strategy) sampling distribution — useful for RL rollout (PPO/GRPO importance ratios). Returns: Sampled token IDs ``[batch]``, or — when ``return_logprobs`` is ``True`` — a ``(token_ids, chosen_logprobs)`` tuple where ``chosen_logprobs`` has shape ``[batch]``. """ greedy = ( ( isinstance(temperature, Tensor) and temperature.numel() == 1 and temperature.item() == 0 ) if isinstance(temperature, Tensor) else temperature == 0 ) if greedy: tokens = logits.argmax(dim=-1) if not return_logprobs: return tokens log_probs = torch.log_softmax(logits.float(), dim=-1) chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1) return tokens, chosen has_freq = ( (isinstance(frequency_penalty, Tensor) and (frequency_penalty != 0).any()) if isinstance(frequency_penalty, Tensor) else frequency_penalty != 0 ) strategies: List[BaseSamplingStrategy] = [ TemperatureStrategy(temperature), TopKStrategy(top_k), TopPStrategy(top_p), ] if has_freq: strategies.append(FrequencyPenaltyStrategy(frequency_penalty)) return SamplingPipeline(strategies).sample( logits, filter_value=filter_value, input_ids=input_ids, input_mask=input_mask, return_logprobs=return_logprobs, )