diff --git a/python/freetoken/core.py b/python/freetoken/core.py index fd691d180..1487dc218 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -24,6 +24,8 @@ class SamplingParams: # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one # appears in the decoded output; the matched substring (and anything after) is trimmed. stop_strs: list[str] = field(default_factory=list) + presence_penalty: float = 0.0 + frequency_penalty: float = 0.0 @property def is_greedy(self) -> bool: @@ -64,6 +66,7 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False + output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: assert self.input_ids.is_cpu diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 3bdae7695..c6fb2f40e 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, List import torch @@ -16,6 +16,7 @@ class BatchSamplingArgs: top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None greedy_mask: torch.Tensor | None = None + penalties: list[tuple[int, torch.Tensor, float, float]] = field(default_factory=list) def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: @@ -59,8 +60,20 @@ class Sampler: def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] + penalties = [] + for row, req in enumerate(batch.reqs): + p = req.sampling_params + if not (p.presence_penalty or p.frequency_penalty) or not req.can_decode: + continue + if req.output_token_counts is None: + req.output_token_counts = torch.zeros( + self.vocab_size, dtype=torch.int32, device=self.device + ) + penalties.append( + (row, req.output_token_counts, p.presence_penalty, p.frequency_penalty) + ) if all(is_greedy): - return BatchSamplingArgs(temperatures=None) + return BatchSamplingArgs(temperatures=None, penalties=penalties) MIN_P = MIN_T = 1e-6 # Greedy outputs are selected explicitly in sample(); use neutral sampling @@ -83,17 +96,30 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: greedy_mask = ( make_device_tensor(is_greedy, torch.bool, self.device) if any(is_greedy) else None ) - return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask) + return BatchSamplingArgs( + temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, penalties=penalties + ) @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + if args.penalties: + logits = logits.float().clone() + for row, counts, presence, frequency in args.penalties: + logits[row] -= frequency * counts + presence * (counts > 0) if args.temperatures is None: # greedy sampling - return torch.argmax(logits, dim=-1) - tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) - if args.greedy_mask is not None: - # Mixed batches still run probability sampling for all rows, but - # greedy rows must follow argmax's deterministic tie-breaking. - greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) - tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + tokens = torch.argmax(logits, dim=-1) + else: + tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) + if args.greedy_mask is not None: + # Mixed batches still run probability sampling for all rows, but + # greedy rows must follow argmax's deterministic tie-breaking. + greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) + tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + # Update on the sampling stream: overlapped scheduling can prepare the next + # batch before the previous token reaches Req.input_ids on the CPU. + for row, counts, _, _ in args.penalties: + counts.scatter_add_( + 0, tokens[row : row + 1].long(), counts.new_ones(1) + ) return tokens diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py index ffd717280..d53398c18 100644 --- a/python/freetoken/server/api_models.py +++ b/python/freetoken/server/api_models.py @@ -74,8 +74,8 @@ class ChatCompletionRequest(BaseModel): stream: bool = False stream_options: StreamOptions | None = None stop: str | list[str] | None = None - presence_penalty: float = 0.0 - frequency_penalty: float = 0.0 + presence_penalty: float = Field(default=0.0, allow_inf_nan=False) + frequency_penalty: float = Field(default=0.0, allow_inf_nan=False) chat_template_kwargs: dict[str, Any] = Field(default_factory=dict) reasoning_effort: str | None = None # DeepSeek-wire thinking toggle ({"type": "enabled"|"disabled"}). Any so a @@ -112,8 +112,8 @@ class CompletionRequest(BaseModel): stream: bool = False stream_options: StreamOptions | None = None stop: str | list[str] | None = None - presence_penalty: float = 0.0 - frequency_penalty: float = 0.0 + presence_penalty: float = Field(default=0.0, allow_inf_nan=False) + frequency_penalty: float = Field(default=0.0, allow_inf_nan=False) ignore_eos: bool = False logprobs: int | None = None echo: bool = False diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 91056e8d1..58292c47b 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -167,6 +167,8 @@ def resolve_sampling( model_sampling: dict[str, Any], stop: str | list[str] | None = None, default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, ) -> SamplingParams: """Map a protocol's sampling fields onto the engine's neutral SamplingParams, filling unspecified fields from the checkpoint's recommended defaults.""" @@ -188,6 +190,8 @@ def pick(value, key, framework): top_k=pick(top_k, "top_k", -1), top_p=pick(top_p, "top_p", 1.0), stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, ) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index dc2f73a97..99a428354 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -79,6 +79,8 @@ def chat_request_to_genspec( model_sampling=model_sampling, stop=req.stop, default_max_tokens=default_max_tokens, + presence_penalty=req.presence_penalty, + frequency_penalty=req.frequency_penalty, ), chat_template_kwargs=ctk, template_tools=_tools_for_template(req), @@ -550,6 +552,8 @@ def _resolve_sampling( model_sampling=model_sampling, stop=req.stop, default_max_tokens=default_max_tokens, + presence_penalty=req.presence_penalty, + frequency_penalty=req.frequency_penalty, )