5080deploy: OpenAI compat: stream/non-stream parser parity, chat logprobs under the parsers, logprobs/usage fixes, cache report + decode interleave

python/freetoken/core.py

@@ -54,6 +54,9 @@ class SamplingParams:
     # reports the chosen token's raw (pre-temperature) logprob and top-k alternatives.
     logprobs: bool = False
     top_logprobs: int = 0
+    # Set by the tokenizer worker when the rendered prompt leaves a <think> block open, so
+    # the detokenizer counts an unclosed generation as reasoning tokens.
+    starts_in_reasoning: bool = False
 
     @property
     def is_greedy(self) -> bool:

python/freetoken/engine/sample.py

@@ -57,7 +57,10 @@ class BatchSamplingArgs:
     greedy_mask: torch.Tensor | None = None
     penalties: list[RowPenalty] = field(default_factory=list)
     plan: LogitsPlan | None = None
-    logprob_rows: torch.Tensor | None = None
+    # Rows that asked for logprobs, as an int64 device index built on the host (no
+    # nonzero sync). None with logprob_all_rows set when every row asks: no gather then.
+    logprob_row_idx: torch.Tensor | None = None
+    logprob_all_rows: bool = False
     max_top_logprobs: int = 0
 
 
@@ -244,9 +247,12 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             )
         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)
+        logprob_all_rows = bool(want_logprobs) and all(want_logprobs)
+        logprob_row_idx = (
+            make_device_tensor(
+                [i for i, want in enumerate(want_logprobs) if want], torch.int64, self.device
+            )
+            if any(want_logprobs) and not logprob_all_rows
             else None
         )
         max_top_logprobs = min(
@@ -257,7 +263,8 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
                 temperatures=None,
                 penalties=penalties,
                 plan=plan,
-                logprob_rows=logprob_rows,
+                logprob_row_idx=logprob_row_idx,
+                logprob_all_rows=logprob_all_rows,
                 max_top_logprobs=max_top_logprobs,
             )
 
@@ -289,7 +296,8 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs:
             greedy_mask=greedy_mask,
             penalties=penalties,
             plan=plan,
-            logprob_rows=logprob_rows,
+            logprob_row_idx=logprob_row_idx,
+            logprob_all_rows=logprob_all_rows,
             max_top_logprobs=max_top_logprobs,
         )
 
@@ -324,46 +332,50 @@ def compute_logprobs(
         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:
+        rows = args.logprob_row_idx
+        if not args.logprob_all_rows and (rows is None or rows.numel() == 0):
             return None
 
-        request_logits = logits.index_select(0, requested_rows).float()
+        request_logits = (logits if rows is None else logits.index_select(0, 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_tokens = sampled_tokens.to(dtype=torch.long, device=logits.device)
+        if rows is not None:
+            request_tokens = request_tokens.index_select(0, rows)
+        request_row_idx = torch.arange(request_logprobs.shape[0], 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 rows is None:
+            chosen_logprobs = request_chosen_logprobs
+        else:
+            chosen_logprobs = torch.full(
+                (logits.shape[0],), float("nan"), dtype=torch.float32, device=logits.device
+            )
+            chosen_logprobs.index_copy_(0, 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
+            if rows is None:
+                top_ids = request_top_ids.to(torch.int32)
+                top_logprobs = request_top_logprobs
+            else:
+                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[rows] = request_top_ids.to(torch.int32)
+                top_logprobs[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)

python/freetoken/message/frontend.py

@@ -62,9 +62,13 @@ 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
+    # Neutral sampled-token logprobs entries released with this reply (see
+    # tokenizer.detokenize.build_logprobs_entry), normally this token's own; None when
+    # there are none. Each entry's "text_offset" is where its token starts in the generated
+    # text. With stop strings an entry waits for its token's text, and entries
+    # of tokens a matched stop string trimmed arrive marked "stop_trimmed"; the token it
+    # ends inside carries "stop_visible", the length of its text that was released.
+    logprobs: list[dict] | None = None
 
 
 @dataclass

python/freetoken/message/tokenizer.py

@@ -40,6 +40,8 @@ class DetokenizeMsg(BaseTokenizerMsg):
     keep_stop_str: bool = False
     # Decode this request with skip_special_tokens (SamplingParams.skip_special_tokens).
     skip_special_tokens: bool = False
+    # The prompt left a reasoning block open (SamplingParams.starts_in_reasoning).
+    starts_in_reasoning: 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

python/freetoken/scheduler/scheduler.py

@@ -414,6 +414,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                         stop_strs=req.sampling_params.stop_strs or None,
                         keep_stop_str=req.sampling_params.include_stop_str_in_output,
                         skip_special_tokens=req.sampling_params.skip_special_tokens,
+                        starts_in_reasoning=req.sampling_params.starts_in_reasoning,
                         prefill_ms=(
                             0.0 if started_at is None
                             else (time.monotonic() - started_at) * 1000.0

python/freetoken/server/api_models.py

@@ -77,9 +77,9 @@ class ChatCompletionRequest(BaseModel):
     presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False)
     frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0, allow_inf_nan=False)
     # 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=true; entries cover message-content tokens only (no prompt, reasoning or
+    # tool-call tokens). 400 for parser families whose hidden tokens the segmenter
+    # cannot tell apart from content (see _chat_logprobs_conflict).
     logprobs: bool = False
     top_logprobs: int | None = None
     chat_template_kwargs: dict[str, Any] = Field(default_factory=dict)