ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: split infer core into subpackages by concern

Browse Source 
- Eliminate core/ directory into cache/, runtime/, network/ subpackages plus flat modules
- Split cache.py (647 lines) into cache/{buffer,strategy,pool}.py by layer
- Add explicit ContiguousStrategy, make AllocationStrategy a real ABC
- Move TaskCacheState to cache/strategy.py, drop string forward references
- Rename api/ to network/, server.py to app.py
- Move sample.py into runtime/ alongside executor and graph
- Simplify TaskCacheManager.__init__ to single pool param
- Expose pool.strategy and pool.req_pool as public properties
- Fix KVCache import in attention_backend.py (TYPE_CHECKING guard)
- Fix steady-state decode reading uninitialized position_ids on first step
This commit is contained in:
ViperEkura
2026-08-08 23:43:05 +08:00
parent 3fa7e66676
commit 0c1b7664c1
37 changed files with 920 additions and 800 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/runtime/sample.py
+386
View File
@@ -0,0 +1,386 @@
"""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 bool((temperature == 0).all())
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)
tokens = torch.multinomial(
torch.softmax(transformed, dim=-1), num_samples=1
).squeeze(-1)
if not return_logprobs:
return tokens
log_probs = torch.log_softmax(transformed.float(), dim=-1)
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]``.
"""
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,
)