5080deploy: merge upstream #224, OpenAI-compatible logprobs

python/freetoken/core.py

@@ -48,6 +48,10 @@ class SamplingParams:
     # Decode with skip_special_tokens. Off by default here: the reasoning and tool parsers
     # consume the <think>/tool tags from the decoded text.
     skip_special_tokens: bool = False
+    # 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

@@ -391,6 +391,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:
@@ -1187,14 +1192,22 @@ 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)
         if self.moe_offload_cache is not None and self.moe_offload_cache.collect_stats and batch.is_decode:
             self._moe_stats_step += 1
             if self._moe_stats_step >= MOE_STATS_INTERVAL:
                 self._moe_stats_step = 0
                 self._emit_moe_stats()
-        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,
+        )
 
     def _warmup_prefill_lens(self) -> list[int]:
         """Prefill lengths that cross every size bucket the in-repo Triton prefill kernels specialize on."""

python/freetoken/engine/sample.py

@@ -49,6 +49,8 @@ class BatchSamplingArgs:
     top_p: torch.Tensor | None = None
     greedy_mask: torch.Tensor | None = None
     plan: LogitsPlan | 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:
@@ -232,8 +234,22 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
         params = [r.sampling_params for r in batch.reqs]
         is_greedy = [p.is_greedy for p in params]
         plan = self._plan(batch, params)
+        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
+        )
+        max_top_logprobs = min(
+            max((p.top_logprobs for p in params if p.logprobs), default=0), self.vocab_size
+        )
         if all(is_greedy):
-            return BatchSamplingArgs(temperatures=None, plan=plan)
+            return BatchSamplingArgs(
+                temperatures=None,
+                plan=plan,
+                logprob_rows=logprob_rows,
+                max_top_logprobs=max_top_logprobs,
+            )
 
         MIN_P = MIN_T = 1e-6
         # Greedy outputs are selected explicitly in sample(); use neutral sampling
@@ -257,7 +273,13 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             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, plan=plan
+            temperatures,
+            top_k=top_k,
+            top_p=top_p,
+            greedy_mask=greedy_mask,
+            plan=plan,
+            logprob_rows=logprob_rows,
+            max_top_logprobs=max_top_logprobs,
         )
 
     @nvtx_annotate("Sampler")
@@ -274,3 +296,59 @@ def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
                 greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype)
                 tokens = torch.where(args.greedy_mask, greedy_tokens, tokens)
             return tokens
+
+    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

@@ -62,6 +62,9 @@ class UserReply(BaseFrontendMsg):
     # On the finished reply: tokens generated up to and including the reasoning end tag
     # (0 when the model emitted none) -> usage.completion_tokens_details.reasoning_tokens.
     reasoning_tokens: int = 0
+    # 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

@@ -40,6 +40,11 @@ class DetokenizeMsg(BaseTokenizerMsg):
     keep_stop_str: bool = False
     # Decode this request with skip_special_tokens (SamplingParams.skip_special_tokens).
     skip_special_tokens: bool = False
+    # 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

@@ -319,8 +319,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():
@@ -353,6 +354,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
@@ -394,6 +408,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                             0.0 if started_at is None
                             else (time.monotonic() - started_at) * 1000.0
                         ),
+                        chosen_logprob=row_chosen_logprob,
+                        top_ids=row_top_ids,
+                        top_logprobs=row_top_logprobs,
                     )
                 )
 

python/freetoken/server/api_models.py

@@ -76,6 +76,12 @@ 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).
+    # Fail-closed (400) when a reasoning parser or tool parsing is active: entries
+    # must align 1:1 with visible content tokens (see _chat_logprobs_conflict).
+    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

@@ -73,6 +73,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
@@ -153,6 +157,9 @@ class GenResult:
     cached_tokens: int = 0
     reasoning_tokens: int = 0
     timings: GenTimings = field(default_factory=GenTimings)
+    # Neutral sampled-token logprob entries, one per sampled token (empty when the
+    # request did not ask).
+    logprobs: list[dict] = field(default_factory=list)
 
 
 @dataclass
@@ -765,10 +772,17 @@ async def _generate_events_impl(
     cached_tokens = 0
     reasoning_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:
@@ -779,6 +793,14 @@ async def _generate_events_impl(
             tool_parser = candidate
     frag_stable = tool_parser.args_fragments_prefix_stable() if tool_parser else True
 
+    # Logprob entries are collected only on the passthrough path: with a reasoning
+    # parser or tool parsing active, text can be hidden, held back, or reclassified,
+    # so an entry could come to describe a token that never reaches visible content.
+    # The API layer fail-closes those request combinations; this guard is the
+    # structural half of the contract — a hidden-token entry must never ride a
+    # later visible delta, whatever the caller.
+    collect_logprobs = reasoning_parser is None and tool_parser is None and not parse_tools
+
     # Streaming tool-call assembly: detectors emit fragments (name first, then
     # argument diffs); they accumulate here and the call is emitted complete when
     # the next call starts, trailing text arrives, or the stream ends.
@@ -822,7 +844,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
@@ -861,6 +883,8 @@ def _route_tool_text(piece: str) -> list[GenEvent]:
         completion_tokens += ack.completion_tokens_delta
         cached_tokens += ack.cached_tokens
         reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0))
+        if collect_logprobs and 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)
@@ -875,7 +899,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)
@@ -896,7 +920,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"
@@ -924,7 +948,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:
@@ -933,13 +957,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,
@@ -952,6 +977,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer
     """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
@@ -967,6 +993,8 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer
         cached_tokens += ack.cached_tokens
         reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0))
         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)
@@ -994,4 +1022,5 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any, timer: _Timer
         cached_tokens=cached_tokens,
         reasoning_tokens=reasoning_tokens,
         timings=timer.finish(),
+        logprobs=logprob_entries,
     )