feat(server): OpenAI-compatible logprobs for chat and legacy completions

python/freetoken/core.py

@@ -25,6 +25,10 @@ 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)
+    # Sampled-token logprobs (OpenAI `logprobs`/`top_logprobs`): when on, the sampler
+    # reports the chosen token's raw (pre-temperature) logprob and top-k alternatives.
+    logprobs: bool = False
+    top_logprobs: int = 0
 
     @property
     def is_greedy(self) -> bool:

python/freetoken/engine/engine.py

@@ -288,6 +288,11 @@ class ForwardOutput(NamedTuple):
     next_tokens_gpu: torch.Tensor
     next_tokens_cpu: torch.Tensor
     copy_done_event: torch.cuda.Event
+    # Sampled-token logprobs (None unless some request in the batch asked): CPU
+    # copies covered by copy_done_event, padded to the batch max top_logprobs.
+    chosen_logprobs_cpu: torch.Tensor | None = None
+    top_ids_cpu: torch.Tensor | None = None
+    top_logprobs_cpu: torch.Tensor | None = None
 
 
 class Engine:
@@ -930,9 +935,17 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
         batch_logits = logits[: batch.size]
         next_tokens_gpu = self.sampler.sample(batch_logits, args).to(torch.int32)
         next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True)
+        logprobs_out = self.sampler.compute_logprobs(batch_logits, next_tokens_gpu, args)
         copy_done_event = torch.cuda.Event()
         copy_done_event.record(self.stream)
-        return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)
+        if logprobs_out is None:
+            return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)
+        chosen_logprobs, top_ids, top_logprobs = logprobs_out
+        return ForwardOutput(
+            next_tokens_gpu, next_tokens_cpu, copy_done_event,
+            chosen_logprobs_cpu=chosen_logprobs, top_ids_cpu=top_ids,
+            top_logprobs_cpu=top_logprobs,
+        )
 
     @torch.inference_mode()
     def _warmup_prefill(self) -> None:

python/freetoken/engine/sample.py

@@ -15,6 +15,8 @@ class BatchSamplingArgs:
     temperatures: torch.Tensor | None
     top_k: torch.Tensor | None = None
     top_p: torch.Tensor | None = None
+    logprob_rows: torch.Tensor | None = None
+    max_top_logprobs: int = 0
 
 
 def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
@@ -57,8 +59,21 @@ class Sampler:
 
     def prepare(self, batch: Batch) -> BatchSamplingArgs:
         params = [r.sampling_params for r in batch.reqs]
+        want_logprobs = [p.logprobs for p in params]
+        logprob_rows = (
+            make_device_tensor(want_logprobs, torch.bool, self.device)
+            if any(want_logprobs)
+            else None
+        )
         if all(p.is_greedy for p in params):
-            return BatchSamplingArgs(temperatures=None)
+            max_top_logprobs = max((p.top_logprobs for p in params if p.logprobs), default=0)
+            if max_top_logprobs > self.vocab_size:
+                max_top_logprobs = self.vocab_size
+            return BatchSamplingArgs(
+                temperatures=None,
+                logprob_rows=logprob_rows,
+                max_top_logprobs=max_top_logprobs,
+            )
 
         MIN_P = MIN_T = 1e-6
         ts = [max(0.0 if p.is_greedy else p.temperature, MIN_T) for p in params]
@@ -70,11 +85,76 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             top_k = make_device_tensor(top_ks, torch.int32, self.device)
         if any(p < 1.0 for p in top_ps):
             top_p = make_device_tensor(top_ps, torch.float32, self.device)
-        return BatchSamplingArgs(temperatures, top_k=top_k, top_p=top_p)
+        max_top_logprobs = max((p.top_logprobs for p in params if p.logprobs), default=0)
+        if max_top_logprobs > self.vocab_size:
+            max_top_logprobs = self.vocab_size
+        return BatchSamplingArgs(
+            temperatures,
+            top_k=top_k,
+            top_p=top_p,
+            logprob_rows=logprob_rows,
+            max_top_logprobs=max_top_logprobs,
+        )
 
     @nvtx_annotate("Sampler")
     def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
         with torch.cuda.nvtx.range("Sampler"):
             if args.temperatures is None:  # greedy sampling
                 return torch.argmax(logits, dim=-1)
             return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p)
+
+    def compute_logprobs(
+        self,
+        logits: torch.Tensor,
+        sampled_tokens: torch.Tensor,
+        args: BatchSamplingArgs,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
+        if args.logprob_rows is None:
+            return None
+
+        requested_rows = torch.nonzero(args.logprob_rows, as_tuple=False).flatten()
+        if requested_rows.numel() == 0:
+            return None
+
+        request_logits = logits.index_select(0, requested_rows).float()
+        # Reported values are raw model logprobs (pre-temperature log_softmax over logits).
+        request_logprobs = torch.log_softmax(request_logits, dim=-1)
+
+        request_tokens = sampled_tokens.to(dtype=torch.long, device=logits.device).index_select(
+            0, requested_rows
+        )
+        request_row_idx = torch.arange(requested_rows.numel(), device=logits.device)
+        request_chosen_logprobs = request_logprobs[request_row_idx, request_tokens]
+
+        chosen_logprobs = torch.full(
+            (logits.shape[0],), float("nan"), dtype=torch.float32, device=logits.device
+        )
+        chosen_logprobs.index_copy_(0, requested_rows, request_chosen_logprobs)
+
+        if args.max_top_logprobs > 0:
+            request_top_logprobs, request_top_ids = torch.topk(
+                request_logprobs, k=args.max_top_logprobs, dim=-1
+            )
+            top_ids = torch.full(
+                (logits.shape[0], args.max_top_logprobs),
+                -1,
+                dtype=torch.int32,
+                device=logits.device,
+            )
+            top_logprobs = torch.full(
+                (logits.shape[0], args.max_top_logprobs),
+                float("-inf"),
+                dtype=torch.float32,
+                device=logits.device,
+            )
+            top_ids[requested_rows] = request_top_ids.to(torch.int32)
+            top_logprobs[requested_rows] = request_top_logprobs
+        else:
+            top_ids = torch.empty((logits.shape[0], 0), dtype=torch.int32, device=logits.device)
+            top_logprobs = torch.empty((logits.shape[0], 0), dtype=torch.float32, device=logits.device)
+
+        return (
+            chosen_logprobs.to("cpu", non_blocking=True),
+            top_ids.to("cpu", non_blocking=True),
+            top_logprobs.to("cpu", non_blocking=True),
+        )

python/freetoken/message/frontend.py

@@ -56,6 +56,9 @@ class UserReply(BaseFrontendMsg):
     finish_reason: str | None = None
     # The stop string that ended generation (Anthropic reports it as stop_reason='stop_sequence').
     matched_stop: str | None = None
+    # Neutral sampled-token logprobs entry for this token (see
+    # tokenizer.detokenize.build_logprobs_entry); None when the request did not ask.
+    logprobs: dict | None = None
 
 
 @dataclass

python/freetoken/message/tokenizer.py

@@ -36,6 +36,11 @@ class DetokenizeMsg(BaseTokenizerMsg):
     # The request's stop strings (None when it has none), so the detokenizer can hold back
     # a trailing partial-stop prefix instead of streaming it and then needing to retract.
     stop_strs: list[str] | None = None
+    # Sampled-token logprobs (None unless the request asked): the chosen token's raw
+    # logprob and the top alternatives, already cut to this request's top_logprobs.
+    chosen_logprob: float | None = None
+    top_ids: list[int] | None = None
+    top_logprobs: list[float] | None = None
     # KV page-pool usage snapshot at this step (not-evictable used/total), passed
     # through to the frontend for the shell status bar. 0/0 for owned-KV models.
     kv_used_pages: int = 0

python/freetoken/scheduler/scheduler.py

@@ -303,8 +303,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
         if last_data is None:
             return
 
-        batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1]
-        copy_done.synchronize()
+        batch, outputs = last_data[0].batch, last_data[1]
+        next_tokens_cpu = outputs.next_tokens_cpu
+        outputs.copy_done_event.synchronize()
         reply: List[DetokenizeMsg] = []
         new_finished_reqs: Set[Req] = set()
         with self.cache_manager.lazy_free_region():
@@ -337,6 +338,19 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                 next_token = next_tokens_cpu[i]
                 req.append_host(next_token.unsqueeze(0))
                 next_token = int(next_token.item())
+
+                row_chosen_logprob: float | None = None
+                row_top_ids: list[int] | None = None
+                row_top_logprobs: list[float] | None = None
+                if req.sampling_params.logprobs and outputs.chosen_logprobs_cpu is not None:
+                    row_chosen_logprob = float(outputs.chosen_logprobs_cpu[i].item())
+                    requested_top = req.sampling_params.top_logprobs
+                    if requested_top > 0 and outputs.top_ids_cpu is not None:
+                        row_top_ids = [int(t) for t in outputs.top_ids_cpu[i, :requested_top].tolist()]
+                        row_top_logprobs = outputs.top_logprobs_cpu[i, :requested_top].tolist()
+                    else:
+                        row_top_ids = []
+                        row_top_logprobs = []
                 # EOS / stop-string -> "stop", output budget exhausted -> "length";
                 # EOS and stop strings win over length.
                 hit_length = not req.can_decode
@@ -368,6 +382,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                         finish_reason=finish_reason,
                         matched_stop=matched_stop,
                         stop_strs=req.sampling_params.stop_strs or None,
+                        chosen_logprob=row_chosen_logprob,
+                        top_ids=row_top_ids,
+                        top_logprobs=row_top_logprobs,
                     )
                 )
 

python/freetoken/server/api_models.py

@@ -76,6 +76,10 @@ class ChatCompletionRequest(BaseModel):
     stop: str | list[str] | None = None
     presence_penalty: float = 0.0
     frequency_penalty: float = 0.0
+    # Sampled-token logprobs (OpenAI chat semantics): top_logprobs (0..20) requires
+    # logprobs=true; entries cover generated tokens only (no prompt logprobs).
+    logprobs: bool = False
+    top_logprobs: int | None = None
     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

python/freetoken/server/generation.py

@@ -70,6 +70,10 @@ class ReasoningDelta:
 @dataclass
 class ContentDelta:
     text: str
+    # Neutral sampled-token logprob entries riding this delta (see UserReply.logprobs);
+    # None when the request did not ask. Parser buffering can attach several entries
+    # to one delta.
+    logprobs: list[dict] | None = None
 
 
 @dataclass
@@ -126,6 +130,9 @@ class GenResult:
     completion_tokens: int
     matched_stop: str | None = None
     cached_tokens: int = 0
+    # Neutral sampled-token logprob entries, one per sampled token (empty when the
+    # request did not ask).
+    logprobs: list[dict] = field(default_factory=list)
 
 
 @dataclass
@@ -559,10 +566,17 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte
     completion_tokens = 0
     cached_tokens = 0
     pending = ""
+    pending_logprobs: list[dict] = []
     parse_tools = spec.parse_tools
     reasoning_parser = _make_reasoning_parser(spec, state)
     specials = _leaked_special_tokens(state)
 
+    def _content_delta(text: str) -> ContentDelta:
+        nonlocal pending_logprobs
+        logprobs = pending_logprobs or None
+        pending_logprobs = []
+        return ContentDelta(text, logprobs=logprobs)
+
     tool_parser: FunctionCallParser | None = None
     if parse_tools:
         try:
@@ -616,7 +630,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
                     out.append(done)
                 stripped = strip_special_tokens(payload, specials)
                 if stripped and not (stripped.strip() == "" and suppress_ws):
-                    out.append(ContentDelta(stripped))
+                    out.append(_content_delta(stripped))
                     if stripped.strip():
                         suppress_ws = False
                 continue
@@ -653,6 +667,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         prompt_tokens += ack.prompt_tokens_delta
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
+        if getattr(ack, "logprobs", None) is not None:
+            pending_logprobs.append(ack.logprobs)
         content_delta = ack.incremental_output
         if reasoning_parser is not None and content_delta:
             reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta)
@@ -667,7 +683,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             elif parse_tools:
                 pending += content_delta
             else:
-                yield ContentDelta(strip_special_tokens(content_delta, specials))
+                yield _content_delta(strip_special_tokens(content_delta, specials))
         if ack.finished:
             engine_finish_reason = getattr(ack, "finish_reason", None)
             engine_matched_stop = getattr(ack, "matched_stop", None)
@@ -688,7 +704,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             elif parse_tools:
                 pending += flush_content
             else:
-                yield ContentDelta(strip_special_tokens(flush_content, specials))
+                yield _content_delta(strip_special_tokens(flush_content, specials))
 
     # Engine reason ("stop"/"length"); a tool call overrides it, but a truncation (length) wins.
     finish_reason = engine_finish_reason or "stop"
@@ -716,7 +732,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         if residual:
             stripped = strip_special_tokens(residual, specials)
             if stripped and not (stripped.strip() == "" and suppress_ws):
-                yield ContentDelta(stripped)
+                yield _content_delta(stripped)
         if calls_emitted and finish_reason != "length":
             finish_reason = "tool_calls"
     else:
@@ -725,13 +741,14 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
             normal_text, tool_calls = parsed
             normal_text = strip_special_tokens(normal_text, specials)
             if normal_text:
-                yield ContentDelta(normal_text)
+                yield _content_delta(normal_text)
             yield ToolCallsDelta(tool_calls)
             if finish_reason != "length":
                 finish_reason = "tool_calls"
         elif parse_tools and pending:
-            yield ContentDelta(strip_special_tokens(pending, specials))
+            yield _content_delta(strip_special_tokens(pending, specials))
 
+    # Entries without a content delta are intentionally dropped in streaming mode.
     yield GenDone(
         finish_reason, prompt_tokens, completion_tokens,
         matched_stop=engine_matched_stop, cached_tokens=cached_tokens,
@@ -742,6 +759,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
     """Protocol-neutral non-streaming generation: accumulate, split reasoning, parse
     tool calls, strip special tokens. The adapters format the GenResult into their wire."""
     full_content = ""
+    logprob_entries: list[dict] = []
     prompt_tokens = 0
     completion_tokens = 0
     cached_tokens = 0
@@ -754,6 +772,8 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
         full_content += ack.incremental_output
+        if getattr(ack, "logprobs", None) is not None:
+            logprob_entries.append(ack.logprobs)
         if ack.finished:
             engine_finish_reason = getattr(ack, "finish_reason", None)
             engine_matched_stop = getattr(ack, "matched_stop", None)
@@ -779,4 +799,5 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult:
         completion_tokens=completion_tokens,
         matched_stop=engine_matched_stop,
         cached_tokens=cached_tokens,
+        logprobs=logprob_entries,
     )