refactor: split infer core into subpackages by concern
- 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:
@@ -0,0 +1,25 @@
|
||||
"""Execution primitives: forward passes, CUDA graphs, and sampling."""
|
||||
|
||||
from astrai.inference.runtime.executor import Executor
|
||||
from astrai.inference.runtime.graph import CudaGraphContext
|
||||
from astrai.inference.runtime.sample import (
|
||||
BaseSamplingStrategy,
|
||||
FrequencyPenaltyStrategy,
|
||||
SamplingPipeline,
|
||||
TemperatureStrategy,
|
||||
TopKStrategy,
|
||||
TopPStrategy,
|
||||
sample,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Executor",
|
||||
"CudaGraphContext",
|
||||
"BaseSamplingStrategy",
|
||||
"FrequencyPenaltyStrategy",
|
||||
"SamplingPipeline",
|
||||
"TemperatureStrategy",
|
||||
"TopKStrategy",
|
||||
"TopPStrategy",
|
||||
"sample",
|
||||
]
|
||||
@@ -0,0 +1,418 @@
|
||||
import logging
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from astrai.extension.attention_backend import (
|
||||
ATTN_BACKEND,
|
||||
CudaBackend,
|
||||
attn_backend,
|
||||
get_backend,
|
||||
)
|
||||
from astrai.inference.cache import PagePool, TaskCacheManager
|
||||
from astrai.inference.runtime.graph import CudaGraphContext
|
||||
from astrai.inference.runtime.sample import sample
|
||||
from astrai.inference.task import Task
|
||||
from astrai.inference.workspace import InferenceWorkspace
|
||||
from astrai.model.automodel import AutoModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def timed(label: str, log: Optional[logging.Logger] = None):
|
||||
"""GPU-precise timer via CUDA events; falls back to perf_counter on CPU."""
|
||||
log = log or logger
|
||||
if not log.isEnabledFor(logging.DEBUG):
|
||||
yield
|
||||
return
|
||||
use_cuda = torch.cuda.is_available()
|
||||
if use_cuda:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
else:
|
||||
tic = time.perf_counter()
|
||||
yield
|
||||
if use_cuda:
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
elapsed_ms = start.elapsed_time(end)
|
||||
else:
|
||||
elapsed_ms = (time.perf_counter() - tic) * 1000
|
||||
log.debug("%s %.2fms", label, elapsed_ms)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SamplingBatchInfo:
|
||||
"""Per-batch sampling parameters, cached across decode steps.
|
||||
|
||||
Sampling params are constant for a given ordered task set, so they are
|
||||
built once (pinned-memory async H2D) and reused until the task set
|
||||
changes. ``top_ks`` is int32 to match the native consumers.
|
||||
"""
|
||||
|
||||
temperatures: Tensor # float32 [B]
|
||||
top_ks: Tensor # int32 [B]
|
||||
top_ps: Tensor # float32 [B]
|
||||
freq_penalties: Tensor # float32 [B]
|
||||
has_freq: bool # any frequency_penalty != 0 (avoids per-step GPU .any())
|
||||
|
||||
|
||||
@dataclass
|
||||
class DecodeSteadyState:
|
||||
"""Cached decode metadata for the steady-state case.
|
||||
|
||||
When the same ordered task set decodes one token per step, sampling
|
||||
params and task signature are reused; only positions advance by 1.
|
||||
"""
|
||||
|
||||
task_sig: tuple
|
||||
positions: list[int]
|
||||
sampling_info: SamplingBatchInfo
|
||||
|
||||
|
||||
def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
|
||||
pin = str(device).startswith("cuda")
|
||||
freq_penalties = torch.tensor(
|
||||
[t.frequency_penalty for t in tasks], dtype=torch.float32, pin_memory=pin
|
||||
).to(device, non_blocking=True)
|
||||
return SamplingBatchInfo(
|
||||
temperatures=torch.tensor(
|
||||
[t.temperature for t in tasks], dtype=torch.float32, pin_memory=pin
|
||||
).to(device, non_blocking=True),
|
||||
top_ks=torch.tensor(
|
||||
[t.top_k for t in tasks], dtype=torch.int32, pin_memory=pin
|
||||
).to(device, non_blocking=True),
|
||||
top_ps=torch.tensor(
|
||||
[t.top_p for t in tasks], dtype=torch.float32, pin_memory=pin
|
||||
).to(device, non_blocking=True),
|
||||
freq_penalties=freq_penalties,
|
||||
has_freq=bool((freq_penalties != 0).any()),
|
||||
)
|
||||
|
||||
|
||||
def _warmup_cuda_graphs(
|
||||
model: AutoModel,
|
||||
pool: PagePool,
|
||||
task_cache: TaskCacheManager,
|
||||
ws: InferenceWorkspace,
|
||||
gctx: CudaGraphContext,
|
||||
max_batch_size: int,
|
||||
prompt_len: int = 1,
|
||||
device: Optional[str] = None,
|
||||
):
|
||||
dev = device or next(model.parameters()).device
|
||||
|
||||
# Prefill warmup: cuBLAS auto-tunes for the actual prompt-length tensor
|
||||
# shapes on first call (F.linear is the dominant cost). This also warms
|
||||
# up the CUDA context (driver init) and compiles the graph-capture trace
|
||||
# that follows. Custom .so kernels do NOT need this — they are pre-built.
|
||||
warmup_len = 64
|
||||
tid = "_warmup_prefill"
|
||||
if task_cache.task_alloc(tid, list(range(warmup_len))):
|
||||
with (
|
||||
torch.inference_mode(),
|
||||
timed("warmup prefill", logger),
|
||||
):
|
||||
kv = task_cache.bind([tid], ws, start_pos=0)
|
||||
ids_in = torch.arange(warmup_len, device=dev).unsqueeze(0)
|
||||
pos_in = ids_in
|
||||
model(
|
||||
ids_in,
|
||||
input_mask=pos_in.unsqueeze(-1) >= torch.arange(warmup_len, device=dev),
|
||||
kv_cache=kv,
|
||||
position_ids=pos_in,
|
||||
)
|
||||
task_cache.task_free(tid)
|
||||
|
||||
batch_sizes = [1]
|
||||
n = 2
|
||||
while n <= max_batch_size:
|
||||
batch_sizes.append(n)
|
||||
n *= 2
|
||||
if max_batch_size not in batch_sizes:
|
||||
batch_sizes.append(max_batch_size)
|
||||
|
||||
for b in batch_sizes:
|
||||
task_ids = [f"_warmup_decode_{b}_{i}" for i in range(b)]
|
||||
prompt_tokens = [list(range(prompt_len)) for _ in range(b)]
|
||||
alloc_ok = True
|
||||
for tid, pt in zip(task_ids, prompt_tokens):
|
||||
if not task_cache.task_alloc(tid, pt):
|
||||
alloc_ok = False
|
||||
break
|
||||
if not alloc_ok:
|
||||
for tid in task_ids:
|
||||
task_cache.task_free(tid)
|
||||
continue
|
||||
|
||||
with (
|
||||
torch.inference_mode(),
|
||||
attn_backend(ATTN_BACKEND.CUDA),
|
||||
timed(f"warmup decode b={b}", logger),
|
||||
):
|
||||
for step in range(2):
|
||||
seq_pos = step
|
||||
ws.position_ids[:b] = seq_pos
|
||||
for tid in task_ids:
|
||||
task_cache.task_extend(tid, seq_pos)
|
||||
kv = task_cache.bind(task_ids, ws)
|
||||
input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
|
||||
ids_buf = ws.fill_input_ids([step] * b)
|
||||
gctx.forward(
|
||||
model,
|
||||
key=(b,),
|
||||
input_ids=ids_buf.unsqueeze(1),
|
||||
input_mask=input_mask,
|
||||
kv_cache=kv,
|
||||
position_ids=ws.position_ids[:b].unsqueeze(1),
|
||||
)
|
||||
|
||||
for tid in task_ids:
|
||||
task_cache.task_free(tid)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
class Executor:
|
||||
"""Model forward passes for prefill and decode phases."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: AutoModel,
|
||||
kv_cache: PagePool,
|
||||
task_cache: TaskCacheManager,
|
||||
device: Optional[str] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
):
|
||||
self.model = model
|
||||
self.kv_cache = kv_cache
|
||||
self.task_cache = task_cache
|
||||
self.device = device or next(model.parameters()).device
|
||||
self.dtype = dtype or next(model.parameters()).dtype
|
||||
|
||||
# Per-step decode cache for the steady-state case (same ordered
|
||||
# task set decodes one token per step). Sampling params stay
|
||||
# constant; only positions advance.
|
||||
self._decode_cache: Optional[DecodeSteadyState] = None
|
||||
|
||||
# Pre-allocated fixed-shape buffers for the decode hot path
|
||||
# (input_ids, decode mask, KV bind metadata). Eagerly sized at init
|
||||
# so the workspace is CUDA-graph-capture friendly — no allocation
|
||||
# during capture.
|
||||
config = model.config
|
||||
max_q_heads = config.num_attention_heads
|
||||
head_dim = config.hidden_size // config.num_attention_heads
|
||||
self._graph_supported = CudaBackend.supports(head_dim=head_dim)
|
||||
self._workspace = InferenceWorkspace(
|
||||
max_batch_size=kv_cache.max_batch_size,
|
||||
max_seq_len=kv_cache.max_seq_len,
|
||||
max_q_heads=max_q_heads,
|
||||
head_dim=head_dim,
|
||||
device=self.device,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
|
||||
# CUDA-graph capture: one graph per (batch_size,) key.
|
||||
# Enabled at init-time via _warmup_cuda_graphs for CudaBackend
|
||||
# on supported head_dims; left disabled otherwise.
|
||||
self._graph_ctx = CudaGraphContext()
|
||||
self._try_enable_cuda_graph()
|
||||
|
||||
def _try_enable_cuda_graph(self):
|
||||
if not self._graph_supported:
|
||||
return
|
||||
|
||||
self._graph_ctx.set_enabled(True)
|
||||
_warmup_cuda_graphs(
|
||||
self.model,
|
||||
self.kv_cache,
|
||||
self.task_cache,
|
||||
self._workspace,
|
||||
self._graph_ctx,
|
||||
max_batch_size=self.kv_cache.max_batch_size,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
def _sample_logits(
|
||||
self,
|
||||
logits: Tensor,
|
||||
tasks: List[Task],
|
||||
return_logprobs: bool = False,
|
||||
info: Optional[SamplingBatchInfo] = None,
|
||||
):
|
||||
info = info or _build_sampling_batch_info(tasks, self.device)
|
||||
if info.has_freq:
|
||||
history_lists = [
|
||||
t.prompt_ids[-t.rep_window :] + t.output_ids for t in tasks
|
||||
]
|
||||
history_lens = [len(ids) for ids in history_lists]
|
||||
max_len = max(history_lens, default=0)
|
||||
padded_ids = torch.zeros(
|
||||
len(tasks), max_len, dtype=torch.long, device=self.device
|
||||
)
|
||||
padded_mask = torch.zeros(
|
||||
len(tasks), max_len, dtype=torch.bool, device=self.device
|
||||
)
|
||||
for i, ids in enumerate(history_lists):
|
||||
length = len(ids)
|
||||
padded_ids[i, :length] = torch.as_tensor(
|
||||
ids, dtype=torch.long, device=self.device
|
||||
)
|
||||
padded_mask[i, :length] = True
|
||||
else:
|
||||
padded_ids = None
|
||||
padded_mask = None
|
||||
|
||||
result = sample(
|
||||
logits,
|
||||
temperature=info.temperatures,
|
||||
top_k=info.top_ks,
|
||||
top_p=info.top_ps,
|
||||
frequency_penalty=info.freq_penalties,
|
||||
input_ids=padded_ids,
|
||||
input_mask=padded_mask,
|
||||
return_logprobs=return_logprobs,
|
||||
)
|
||||
if not return_logprobs:
|
||||
return result.tolist()
|
||||
|
||||
tokens, logprobs = result
|
||||
tokens_list = tokens.tolist()
|
||||
logprobs_list = logprobs.tolist()
|
||||
for task, logprob in zip(tasks, logprobs_list):
|
||||
task.output_logprobs.append(float(logprob))
|
||||
return list(zip(tokens_list, logprobs_list))
|
||||
|
||||
def execute_prefill(
|
||||
self,
|
||||
tasks: List[Task],
|
||||
prompt_len: int,
|
||||
start_pos: int = 0,
|
||||
return_logprobs: bool = False,
|
||||
):
|
||||
if start_pos >= prompt_len:
|
||||
return []
|
||||
|
||||
tasks = sorted(tasks, key=lambda t: t.task_id)
|
||||
batch_sz = len(tasks)
|
||||
|
||||
input_ids = torch.tensor(
|
||||
[t.prompt_ids[start_pos:prompt_len] for t in tasks],
|
||||
dtype=torch.long,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
task_ids = [t.task_id for t in tasks]
|
||||
position_ids = (
|
||||
torch.arange(start_pos, prompt_len, dtype=torch.long, device=self.device)
|
||||
.unsqueeze(0)
|
||||
.expand(batch_sz, -1)
|
||||
)
|
||||
input_mask = position_ids.unsqueeze(-1) >= torch.arange(
|
||||
prompt_len, device=self.device
|
||||
)
|
||||
|
||||
with (
|
||||
torch.inference_mode(),
|
||||
timed(f"execute_prefill b={batch_sz} prompt_len={prompt_len}", logger),
|
||||
):
|
||||
outputs = self.model(
|
||||
input_ids,
|
||||
input_mask=input_mask,
|
||||
position_ids=position_ids,
|
||||
kv_cache=self.task_cache.bind(
|
||||
task_ids,
|
||||
self._workspace,
|
||||
start_pos=start_pos,
|
||||
),
|
||||
)
|
||||
logits = outputs["logits"][:, -1, :]
|
||||
|
||||
return tasks, self._sample_logits(logits, tasks, return_logprobs)
|
||||
|
||||
def execute_decode(
|
||||
self, tasks: List[Task], return_logprobs: bool = False
|
||||
) -> List[int]:
|
||||
"""Decode next token for each task.
|
||||
|
||||
Args:
|
||||
return_logprobs: When ``True``, also record (and return)
|
||||
the log-probability of each sampled token under the
|
||||
post-strategy sampling distribution. The logprob is
|
||||
appended to ``task.output_logprobs`` and the return
|
||||
list becomes ``List[Tuple[int, float]]``.
|
||||
|
||||
Returns:
|
||||
``List[int]`` of sampled token IDs, or
|
||||
``List[Tuple[int, float]]`` of ``(token_id, logprob)`` when
|
||||
``return_logprobs`` is ``True``.
|
||||
"""
|
||||
if not tasks:
|
||||
return []
|
||||
|
||||
b = len(tasks)
|
||||
ws = self._workspace
|
||||
|
||||
# ---- pre-replay: update input buffers in-place ----
|
||||
|
||||
input_ids = ws.fill_input_ids(
|
||||
[t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks]
|
||||
)
|
||||
|
||||
task_ids = [t.task_id for t in tasks]
|
||||
cur_positions = [t.next_pos for t in tasks]
|
||||
|
||||
kv_cache = self.task_cache.bind(task_ids, ws)
|
||||
|
||||
if self.task_cache.bind_was_steady and self._decode_cache is not None:
|
||||
info = self._decode_cache.sampling_info
|
||||
ws.position_ids[:b] += 1
|
||||
else:
|
||||
info = _build_sampling_batch_info(tasks, self.device)
|
||||
ws.position_ids[:b].copy_(
|
||||
torch.tensor(cur_positions, dtype=torch.long, device=self.device)
|
||||
)
|
||||
self._decode_cache = DecodeSteadyState(tuple(task_ids), cur_positions, info)
|
||||
|
||||
total_len = max(cur_positions) + 1
|
||||
input_mask = ws.decode_mask(ws.position_ids[:b], total_len)
|
||||
|
||||
# ---- forward (graph replay or live run + capture) ----
|
||||
|
||||
use_graph = (
|
||||
self._graph_ctx.enabled
|
||||
and self._graph_supported
|
||||
and get_backend().supports_graph()
|
||||
)
|
||||
key = (b,)
|
||||
if use_graph:
|
||||
input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
|
||||
|
||||
with (
|
||||
torch.inference_mode(),
|
||||
timed(f"execute_decode forward b={b}", logger),
|
||||
):
|
||||
if use_graph:
|
||||
outputs = self._graph_ctx.forward(
|
||||
self.model,
|
||||
key=key,
|
||||
input_ids=input_ids.unsqueeze(1),
|
||||
input_mask=input_mask,
|
||||
kv_cache=kv_cache,
|
||||
position_ids=ws.position_ids[:b].unsqueeze(1),
|
||||
)
|
||||
else:
|
||||
outputs = self.model(
|
||||
input_ids.unsqueeze(1),
|
||||
input_mask=input_mask,
|
||||
kv_cache=kv_cache,
|
||||
position_ids=ws.position_ids[:b].unsqueeze(1),
|
||||
)
|
||||
logits = outputs["logits"][:, -1, :]
|
||||
|
||||
return self._sample_logits(logits, tasks, return_logprobs, info=info)
|
||||
@@ -0,0 +1,103 @@
|
||||
"""CUDA-graph capture for the decode model-forward step.
|
||||
|
||||
Mirrors SGLang's cuda-graph manager: one graph per batch size. The graph
|
||||
pair. The graph captures ``model.forward()`` with workspace-backed inputs
|
||||
(all at fixed addresses). Before each replay the caller updates the input
|
||||
buffer content in-place so the graph sees fresh data at the same tensor
|
||||
addresses.
|
||||
|
||||
Only the model forward is captured — sampling runs outside the graph
|
||||
(via ``torch.multinomial`` which consumes a mutable RNG state).
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class CudaGraphContext:
|
||||
"""CUDA-graph capture/replay for decode steps.
|
||||
|
||||
Parameters:
|
||||
enabled: When ``False``, ``forward()`` always runs the live model
|
||||
forward without capture/replay (graphs are cleared). Toggle at
|
||||
runtime via the ``set_enabled()`` method.
|
||||
|
||||
Usage::
|
||||
|
||||
gctx = CudaGraphContext()
|
||||
with torch.inference_mode():
|
||||
outputs = gctx.forward(
|
||||
model,
|
||||
key=(batch_size,),
|
||||
input_ids=workspace.input_ids[:b].unsqueeze(1),
|
||||
input_mask=input_mask,
|
||||
kv_cache=kv_cache,
|
||||
position_ids=workspace.position_ids[:b].unsqueeze(1),
|
||||
)
|
||||
|
||||
The first call at a given key runs *without* capture (warmup). The
|
||||
second call captures the graph. Subsequent calls replay the captured
|
||||
graph. A ``torch.cuda.synchronize()`` before capture drains in-flight
|
||||
work so the graph trace is clean.
|
||||
"""
|
||||
|
||||
def __init__(self, enabled: bool = False):
|
||||
self._enabled = enabled
|
||||
self._graphs: dict[tuple, torch.cuda.CUDAGraph] = {}
|
||||
self._outputs: dict[tuple, dict[str, Tensor]] = {}
|
||||
self._warmed: set[tuple] = set()
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._enabled
|
||||
|
||||
def set_enabled(self, flag: bool):
|
||||
"""Enable or disable CUDA-graph capture at runtime.
|
||||
|
||||
Disabling clears all captured graphs (frees GPU memory) and warmup
|
||||
state. Re-enabling after disable starts fresh — graphs are
|
||||
re-captured on the next warmup cycle.
|
||||
"""
|
||||
if flag == self._enabled:
|
||||
return
|
||||
self._enabled = flag
|
||||
if not flag:
|
||||
self._graphs.clear()
|
||||
self._outputs.clear()
|
||||
self._warmed.clear()
|
||||
|
||||
def forward(self, model, *, key, **kwargs) -> dict[str, Tensor]:
|
||||
"""Run ``model(**kwargs)`` via graph replay or live forward.
|
||||
|
||||
Args:
|
||||
model: callable, e.g. ``self.model.forward``.
|
||||
key: ``(batch_size,)`` — the dispatch key (one graph per batch size).
|
||||
**kwargs: arguments forwarded to ``model``. All tensor arguments
|
||||
must reside at stable addresses (workspace buffers).
|
||||
|
||||
Returns:
|
||||
The dict produced by ``model(**kwargs)``, e.g.
|
||||
``{"logits": ..., "h0": ...}``.
|
||||
"""
|
||||
if not self._enabled:
|
||||
self._outputs[key] = model(**kwargs)
|
||||
return self._outputs[key]
|
||||
|
||||
if key in self._graphs:
|
||||
self._graphs[key].replay()
|
||||
elif key in self._warmed:
|
||||
cap_output = model(**kwargs)
|
||||
torch.cuda.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
self._outputs[key] = model(**kwargs)
|
||||
self._graphs[key] = graph
|
||||
self._warmed.discard(key)
|
||||
return cap_output
|
||||
else:
|
||||
self._warmed.add(key)
|
||||
self._outputs[key] = model(**kwargs)
|
||||
return self._outputs[key]
|
||||
|
||||
def has_graph(self, key: tuple) -> bool:
|
||||
return key in self._graphs
|
||||
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user