fix(sampling): apply presence and frequency penalties

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

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

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

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,
     )
 
 

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,
     )