From 58a252ed3f2411f4d2a8bf58e3be308657ade17f Mon Sep 17 00:00:00 2001 From: mac Date: Wed, 26 Aug 2026 23:23:29 +0300 Subject: [PATCH 1/2] engine: compute sampled-token logprobs on request Adds the engine half of OpenAI logprobs support. SamplingParams gains logprobs/top_logprobs; the sampler computes log_softmax over the PRE-temperature logits (raw model distribution, so temperature/top-k/top-p do not change reported values), gathers the chosen token and batch-max top-k, and ships CPU copies covered by the existing copy_done_event. The scheduler attaches per-request values (cut to each request's own top_logprobs) to DetokenizeMsg; the detokenizer builds a neutral entry (token text via single-id decode + UTF-8 bytes, so clients can reassemble partial-UTF-8 pieces) onto UserReply.logprobs. Zero cost when off: no row asked -> no mask tensor, no log_softmax, ForwardOutput carries None. Message fields default to None, so old and new peers interoperate. Stop-string trimming can hide final visible text; entries still cover every sampled token. --- python/freetoken/core.py | 4 + python/freetoken/engine/engine.py | 15 +- python/freetoken/engine/sample.py | 84 +++++++++++- python/freetoken/message/frontend.py | 3 + python/freetoken/message/tokenizer.py | 5 + python/freetoken/scheduler/scheduler.py | 21 ++- python/freetoken/tokenizer/detokenize.py | 29 ++++ python/freetoken/tokenizer/server.py | 15 +- tests/engine/test_sample_logprobs.py | 167 +++++++++++++++++++++++ tests/tokenizer/test_logprobs_entry.py | 59 ++++++++ 10 files changed, 396 insertions(+), 6 deletions(-) create mode 100644 tests/engine/test_sample_logprobs.py create mode 100644 tests/tokenizer/test_logprobs_entry.py diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539cb..b001fc3e6 100644 --- a/python/freetoken/core.py +++ b/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: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2d..15121ba60 100644 --- a/python/freetoken/engine/engine.py +++ b/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: diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 01d14b1aa..3d5c27e15 100644 --- a/python/freetoken/engine/sample.py +++ b/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,7 +85,16 @@ 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: @@ -78,3 +102,59 @@ def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: 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), + ) diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index 24725567b..3da582cf6 100644 --- a/python/freetoken/message/frontend.py +++ b/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 diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 33b75c785..f658c9cc4 100644 --- a/python/freetoken/message/tokenizer.py +++ b/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 diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 48923e3b0..a5a3c2fe1 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/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, ) ) diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1a6..1505a56c5 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -145,3 +145,32 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: del self.decode_map[msg.uid] return incremental_strs + + +def build_logprobs_entry( + tokenizer: PreTrainedTokenizerBase, + token_id: int, + chosen_logprob: float, + top_ids: list[int] | None, + top_logprobs: list[float] | None, +) -> dict: + """Neutral sampled-token logprobs entry (UserReply.logprobs). Values are the raw + pre-temperature logprobs the sampler computed; token text comes from a single-id + decode, with UTF-8 bytes alongside so clients can reassemble partial-UTF-8 pieces.""" + + def _fields(tid: int) -> tuple[str, list[int]]: + text = tokenizer.decode([tid]) + return text, list(text.encode("utf-8")) + + token, token_bytes = _fields(token_id) + top = [] + for tid, logprob in zip(top_ids or [], top_logprobs or []): + text, data = _fields(int(tid)) + top.append({"token_id": int(tid), "token": text, "bytes": data, "logprob": float(logprob)}) + return { + "token_id": token_id, + "token": token, + "bytes": token_bytes, + "logprob": float(chosen_logprob), + "top": top, + } diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 530e862d0..d4f6a03ac 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -144,7 +144,7 @@ def tokenize_worker( tokenizer = load_tokenizer(tokenizer_path) logger = init_logger(__name__, f"tokenizer_{tokenizer_id}") - from .detokenize import DetokenizeManager + from .detokenize import DetokenizeManager, build_logprobs_entry from .tokenize import TokenizeManager tokenize_manager = TokenizeManager(tokenizer) @@ -223,6 +223,19 @@ def tokenize_worker( swa_used_tokens=msg.swa_used_tokens, swa_total_tokens=msg.swa_total_tokens, gpu_mem_bytes=msg.gpu_mem_bytes, + # Stop-string trimming can hide final visible text; keep one logprob + # entry per sampled token. + logprobs=( + build_logprobs_entry( + detokenize_manager.tokenizer, + msg.next_token, + msg.chosen_logprob, + msg.top_ids, + msg.top_logprobs, + ) + if msg.chosen_logprob is not None + else None + ), ) for msg, reply in zip(detokenize_msg, replies, strict=True) ] diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py new file mode 100644 index 000000000..aabbdc5fe --- /dev/null +++ b/tests/engine/test_sample_logprobs.py @@ -0,0 +1,167 @@ +from types import SimpleNamespace + +import torch + +from freetoken.engine.sample import BatchSamplingArgs, Sampler + + +def test_compute_logprobs_matches_raw_log_softmax_and_sorted_top() -> None: + sampler = Sampler(torch.device("cpu"), vocab_size=4) + + logits = torch.tensor( + [ + [2.0, 0.0, 1.0, -1.0], + [0.0, -1.0, 1.0, 3.0], + [1.0, 2.0, 3.0, 4.0], + ], + dtype=torch.float32, + ) + sampled_tokens = torch.tensor([2, 3, 0], dtype=torch.long) + args = BatchSamplingArgs( + temperatures=torch.tensor([1.0, 1.0, 1.0], dtype=torch.float32), + logprob_rows=torch.tensor([True, False, True], dtype=torch.bool), + max_top_logprobs=3, + ) + + result = sampler.compute_logprobs(logits, sampled_tokens, args) + + assert result is not None + chosen_logprobs, top_ids, top_logprobs = result + expected = torch.log_softmax(logits.float(), dim=-1) + + assert torch.isclose(chosen_logprobs[0], expected[0, sampled_tokens[0]]) + assert torch.isnan(chosen_logprobs[1]) + assert torch.isclose(chosen_logprobs[2], expected[2, sampled_tokens[2]]) + + expected_top0 = torch.topk(expected[0], k=3) + expected_top2 = torch.topk(expected[2], k=3) + assert torch.equal(top_ids[0], expected_top0.indices) + assert torch.allclose(top_logprobs[0], expected_top0.values) + assert torch.equal(top_ids[2], expected_top2.indices) + assert torch.allclose(top_logprobs[2], expected_top2.values) + + assert torch.equal(top_ids[1], torch.full((3,), -1, dtype=torch.int32)) + assert torch.isneginf(top_logprobs[1]).all() + assert top_ids.shape == (3, 3) + assert top_logprobs.shape == (3, 3) + + +def test_compute_logprobs_ignores_temperatures() -> None: + sampler = Sampler(torch.device("cpu"), vocab_size=4) + + logits = torch.tensor( + [ + [0.5, 1.0, 2.0, 3.0], + [4.0, 3.0, 2.0, 1.0], + ], + dtype=torch.float32, + ) + sampled_tokens = torch.tensor([1, 2], dtype=torch.long) + logprob_rows = torch.tensor([True, True], dtype=torch.bool) + + cold = BatchSamplingArgs( + temperatures=torch.full((2,), 0.3, dtype=torch.float32), + logprob_rows=logprob_rows, + max_top_logprobs=2, + ) + hot = BatchSamplingArgs( + temperatures=torch.full((2,), 2.5, dtype=torch.float32), + logprob_rows=logprob_rows, + max_top_logprobs=2, + ) + + cold_out = sampler.compute_logprobs(logits, sampled_tokens, cold) + hot_out = sampler.compute_logprobs(logits, sampled_tokens, hot) + + assert cold_out is not None and hot_out is not None + cold_chosen, cold_top_ids, cold_top_logprobs = cold_out + hot_chosen, hot_top_ids, hot_top_logprobs = hot_out + + assert torch.allclose(cold_chosen, hot_chosen) + assert torch.equal(cold_top_ids, hot_top_ids) + assert torch.allclose(cold_top_logprobs, hot_top_logprobs) + + +def test_compute_logprobs_returns_none_without_requested_rows() -> None: + sampler = Sampler(torch.device("cpu"), vocab_size=4) + logits = torch.zeros((2, 4), dtype=torch.float32) + sampled_tokens = torch.tensor([0, 1], dtype=torch.long) + args = BatchSamplingArgs( + temperatures=torch.ones(2), + logprob_rows=torch.zeros(2, dtype=torch.bool), + max_top_logprobs=2, + ) + + assert sampler.compute_logprobs(logits, sampled_tokens, args) is None + + +def test_compute_logprobs_handles_zero_top_logprobs() -> None: + sampler = Sampler(torch.device("cpu"), vocab_size=4) + logits = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0], + [4.0, 3.0, 2.0, 1.0], + [0.1, 0.2, 0.3, 0.4], + ], + dtype=torch.float32, + ) + sampled_tokens = torch.tensor([3, 0, 1], dtype=torch.long) + args = BatchSamplingArgs( + temperatures=torch.full((3,), 1.0), + logprob_rows=torch.tensor([True, False, True], dtype=torch.bool), + max_top_logprobs=0, + ) + + result = sampler.compute_logprobs(logits, sampled_tokens, args) + + assert result is not None + chosen_logprobs, top_ids, top_logprobs = result + expected = torch.log_softmax(logits, dim=-1) + assert torch.isclose(chosen_logprobs[0], expected[0, 3]) + assert torch.isnan(chosen_logprobs[1]) + assert torch.isclose(chosen_logprobs[2], expected[2, 1]) + assert top_ids.shape == (3, 0) + assert top_logprobs.shape == (3, 0) + + +def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatch) -> None: + # prepare() pins host tensors for the H2D copy; pinning needs a CUDA context, + # so stub the transfer helper to keep this test host-agnostic. + import freetoken.engine.sample as sample_mod + + monkeypatch.setattr( + sample_mod, "make_device_tensor", + lambda data, dtype, device: torch.tensor(data, dtype=dtype), + ) + sampler = Sampler(torch.device("cpu"), vocab_size=5) + + batch = SimpleNamespace( + reqs=[ + SimpleNamespace(sampling_params=SimpleNamespace(logprobs=True, top_logprobs=17, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(logprobs=False, top_logprobs=0, is_greedy=True)), + ] + ) + args = sampler.prepare(batch) + assert args.max_top_logprobs == sampler.vocab_size + + logits = torch.tensor( + [ + [1.0, 0.0, -1.0, 0.5, 2.0], + [2.0, 1.0, 0.0, -1.0, -2.0], + ], + dtype=torch.float32, + ) + sampled_tokens = torch.tensor([4, 0], dtype=torch.long) + + result = sampler.compute_logprobs(logits, sampled_tokens, args) + assert result is not None + chosen_logprobs, top_ids, top_logprobs = result + + assert chosen_logprobs.shape == (2,) + assert torch.isclose(chosen_logprobs[0], torch.log_softmax(logits[0], dim=-1)[4]) + assert torch.isnan(chosen_logprobs[1]) + + assert top_ids.shape == (2, sampler.vocab_size) + assert top_logprobs.shape == (2, sampler.vocab_size) + assert torch.equal(top_ids[1], torch.full((sampler.vocab_size,), -1, dtype=torch.int32)) + assert torch.isneginf(top_logprobs[1]).all() diff --git a/tests/tokenizer/test_logprobs_entry.py b/tests/tokenizer/test_logprobs_entry.py new file mode 100644 index 000000000..89cf35cb5 --- /dev/null +++ b/tests/tokenizer/test_logprobs_entry.py @@ -0,0 +1,59 @@ +from freetoken.tokenizer.detokenize import build_logprobs_entry + + +class _FakeTokenizer: + def __init__(self, token_table: dict[int, str]) -> None: + self.token_table = token_table + + def decode(self, token_ids: list[int]) -> str: + return "".join(self.token_table[token_id] for token_id in token_ids) + + +def test_build_logprobs_entry_matches_contract() -> None: + tokenizer = _FakeTokenizer({1: "hel", 2: "lo", 3: "é"}) + entry = build_logprobs_entry( + tokenizer, + token_id=3, + chosen_logprob=-0.123, + top_ids=[3, 2, 1], + top_logprobs=[-0.1, -0.9, -1.5], + ) + + assert entry["token_id"] == 3 + assert entry["token"] == "é" + assert entry["bytes"] == list("é".encode("utf-8")) + assert entry["logprob"] == -0.123 + + top = entry["top"] + assert len(top) == 3 + assert top[0] == { + "token_id": 3, + "token": "é", + "bytes": list("é".encode("utf-8")), + "logprob": -0.1, + } + assert top[1] == { + "token_id": 2, + "token": "lo", + "bytes": list("lo".encode("utf-8")), + "logprob": -0.9, + } + assert top[2] == { + "token_id": 1, + "token": "hel", + "bytes": list("hel".encode("utf-8")), + "logprob": -1.5, + } + + +def test_build_logprobs_entry_empty_top() -> None: + tokenizer = _FakeTokenizer({7: "x", 8: "y"}) + entry = build_logprobs_entry( + tokenizer, + token_id=7, + chosen_logprob=-2.0, + top_ids=[], + top_logprobs=[], + ) + + assert entry["top"] == [] From ea7ae6d6d718ecd2c648474145aab8cf72f5148a Mon Sep 17 00:00:00 2001 From: mac Date: Wed, 26 Aug 2026 23:25:35 +0300 Subject: [PATCH 2/2] server: OpenAI logprobs for chat and legacy completions API half of logprobs support. Chat: logprobs + top_logprobs (0..20) yield choice.logprobs.content entries {token, logprob, bytes, top_logprobs[]}, streaming and non-streaming. Completions: the legacy integer field (0..5) yields {tokens, token_logprobs, top_logprobs, text_offset} with absolute offsets across the stream; echo+logprobs stays rejected (prompt logprobs need prefill logits and are out of scope here). The protocol-neutral event layer carries entries on ContentDelta (a list -- parser buffering can release several tokens' text in one delta); entries always accumulate on GenResult for the non-streaming path, and reasoning/tool-call buffering carries pending entries onto the next content delta. Formatting lives in server/logprobs.py; wire compatibility follows the OpenAI shapes. --- python/freetoken/server/api_models.py | 4 + python/freetoken/server/generation.py | 33 +++- python/freetoken/server/logprobs.py | 60 ++++++ python/freetoken/server/openai_api.py | 86 +++++--- tests/server/test_logprobs_api.py | 274 ++++++++++++++++++++++++++ 5 files changed, 423 insertions(+), 34 deletions(-) create mode 100644 python/freetoken/server/logprobs.py create mode 100644 tests/server/test_logprobs_api.py diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py index ffd717280..7bd43e642 100644 --- a/python/freetoken/server/api_models.py +++ b/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 diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908a..e4627a06b 100644 --- a/python/freetoken/server/generation.py +++ b/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, ) diff --git a/python/freetoken/server/logprobs.py b/python/freetoken/server/logprobs.py new file mode 100644 index 000000000..861856954 --- /dev/null +++ b/python/freetoken/server/logprobs.py @@ -0,0 +1,60 @@ +"""OpenAI logprobs formatting: neutral engine entries -> wire shapes.""" + +from __future__ import annotations + +from typing import Any + + +def chat_logprobs_error(req: Any) -> str | None: + if req.top_logprobs is not None and not req.logprobs: + return "top_logprobs requires logprobs=true" + if req.top_logprobs is not None and not 0 <= req.top_logprobs <= 20: + return "top_logprobs must be between 0 and 20" + return None + + +def completion_logprobs_error(req: Any) -> str | None: + if req.logprobs is not None and not 0 <= req.logprobs <= 5: + return "logprobs must be between 0 and 5" + if req.echo and req.logprobs is not None: + return "echo with logprobs is not supported" + return None + + +def chat_content_entry(entry: dict) -> dict: + return { + "token": entry["token"], + "logprob": entry["logprob"], + "bytes": entry["bytes"], + "top_logprobs": [ + { + "token": candidate["token"], + "logprob": candidate["logprob"], + "bytes": candidate["bytes"], + } + for candidate in entry["top"] + ], + } + + +def completions_logprobs(entries: list[dict], start_offset: int = 0) -> dict: + tokens: list[str] = [] + token_logprobs: list[float] = [] + top_logprobs: list[dict[str, float]] = [] + text_offset: list[int] = [] + offset = start_offset + + for entry in entries: + token = entry["token"] + tokens.append(token) + token_logprobs.append(entry["logprob"]) + top_logprobs.append({candidate["token"]: candidate["logprob"] for candidate in entry["top"]}) + text_offset.append(offset) + offset += len(token) + + return { + "tokens": tokens, + "token_logprobs": token_logprobs, + "top_logprobs": top_logprobs, + "text_offset": text_offset, + } diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd263..82c32bb78 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -22,6 +22,12 @@ ) from .function_call_parser import ToolCallItem from .request_logger import log_request +from .logprobs import ( + chat_content_entry, + chat_logprobs_error, + completion_logprobs_error, + completions_logprobs, +) from .generation import ( ContentDelta, GenDone, @@ -66,17 +72,20 @@ def chat_request_to_genspec( thinking_type = _thinking_type(req) if req.reasoning_effort or thinking_type: ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type) + sampling_params = resolve_sampling( + temperature=req.temperature, + top_k=req.top_k, + top_p=req.top_p, + max_tokens=req.max_tokens, + ignore_eos=req.ignore_eos, + model_sampling=model_sampling, + stop=req.stop, + ) + sampling_params.logprobs = req.logprobs + sampling_params.top_logprobs = req.top_logprobs or 0 return GenSpec( messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages]), - sampling_params=resolve_sampling( - temperature=req.temperature, - top_k=req.top_k, - top_p=req.top_p, - max_tokens=req.max_tokens, - ignore_eos=req.ignore_eos, - model_sampling=model_sampling, - stop=req.stop, - ), + sampling_params=sampling_params, chat_template_kwargs=ctk, template_tools=_tools_for_template(req), parser_tools=(_all_tool_dicts(req.tools) if _should_parse_tools(req) else None), @@ -154,6 +163,9 @@ async def handle_chat_completion( return create_error_response("function_call is not supported; use tools/tool_choice instead") if req.logit_bias is not None: return create_error_response("logit_bias is not supported") + logprobs_error = chat_logprobs_error(req) + if logprobs_error is not None: + return create_error_response(logprobs_error, param="top_logprobs") if _response_format_unsupported(req.response_format): return create_error_response( "response_format json_object/json_schema is not supported (no constrained decoding)", @@ -208,18 +220,20 @@ async def handle_chat_completion( if result.tool_calls: message["tool_calls"] = _tool_calls_to_openai(result.tool_calls) + choice: dict[str, Any] = { + "index": 0, + "message": message, + "finish_reason": result.finish_reason, + } + if req.logprobs: + choice["logprobs"] = {"content": [chat_content_entry(e) for e in result.logprobs]} + return { "id": f"chatcmpl-{uid}", "object": "chat.completion", "created": int(time.time()), "model": req.model, - "choices": [ - { - "index": 0, - "message": message, - "finish_reason": result.finish_reason, - } - ], + "choices": [choice], "usage": _usage( result.prompt_tokens, result.completion_tokens, @@ -272,13 +286,10 @@ async def stream_chat_completion_chunks( ) ) elif isinstance(ev, ContentDelta): - yield _sse( - _chat_chunk( - req, - uid, - [{"delta": {"content": ev.text}, "index": 0, "finish_reason": None}], - ) - ) + choice: dict[str, Any] = {"delta": {"content": ev.text}, "index": 0, "finish_reason": None} + if ev.logprobs: + choice["logprobs"] = {"content": [chat_content_entry(e) for e in ev.logprobs]} + yield _sse(_chat_chunk(req, uid, [choice])) elif isinstance(ev, ToolCallStart): open_tool = { "index": tool_calls_sent, @@ -412,6 +423,7 @@ async def handle_completion( uid = state.new_user() await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=_resolve_sampling(req, model_sampling))) text = "" + entries: list[dict] = [] finish_reason = "stop" async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): @@ -420,10 +432,17 @@ async def handle_completion( completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens text += ack.incremental_output + if req.logprobs is not None and ack.logprobs is not None: + entries.append(ack.logprobs) if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break - choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None}) + choices.append({ + "index": index, + "text": text, + "finish_reason": finish_reason, + "logprobs": completions_logprobs(entries) if req.logprobs is not None else None, + }) return { "id": f"cmpl-{uuid.uuid4().hex}", @@ -437,6 +456,7 @@ async def handle_completion( async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) -> AsyncIterator[bytes]: prompt_tokens = 0 + text_offset = 0 completion_tokens = 0 cached_tokens = 0 finish_reason = "stop" @@ -460,11 +480,16 @@ async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) "text": ack.incremental_output, "index": 0, "finish_reason": None, - "logprobs": None, + "logprobs": ( + completions_logprobs([ack.logprobs], text_offset) + if req.logprobs is not None and ack.logprobs is not None + else None + ), } ], } ) + text_offset += len(ack.incremental_output) if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break @@ -518,7 +543,7 @@ def _resolve_sampling( req: ChatCompletionRequest | CompletionRequest, model_sampling: dict[str, Any], ) -> SamplingParams: - return resolve_sampling( + sampling_params = resolve_sampling( temperature=req.temperature, top_k=req.top_k, top_p=req.top_p, @@ -527,6 +552,10 @@ def _resolve_sampling( model_sampling=model_sampling, stop=req.stop, ) + if isinstance(req, CompletionRequest) and req.logprobs is not None: + sampling_params.logprobs = True + sampling_params.top_logprobs = req.logprobs + return sampling_params def _tools_for_template(req: ChatCompletionRequest) -> list[dict[str, Any]] | None: @@ -627,8 +656,9 @@ def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool def _completion_unsupported_reason(req: CompletionRequest) -> str | None: if _is_token_prompt(req.prompt): return "OpenAI token-id prompt inputs are not supported; pass text prompt strings instead" - if req.logprobs is not None: - return "logprobs is not supported" + logprobs_error = completion_logprobs_error(req) + if logprobs_error is not None: + return logprobs_error if req.echo: return "echo is not supported" if req.suffix is not None: diff --git a/tests/server/test_logprobs_api.py b/tests/server/test_logprobs_api.py new file mode 100644 index 000000000..3b4fd5e32 --- /dev/null +++ b/tests/server/test_logprobs_api.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace + +from freetoken.message import TokenizeMsg, UserReply +from freetoken.server.openai_api import ( + ChatCompletionRequest, + CompletionRequest, + chat_request_to_genspec, + handle_chat_completion, + handle_completion, + stream_chat_completion_chunks, + stream_completion_chunks, +) + + +def run(coro): + return asyncio.run(coro) + + +class FakeState: + def __init__(self, replies: list[UserReply], reasoning_parser: str | None = None) -> None: + self.config = SimpleNamespace( + model_path="/models/unit-model", + served_model_name="unit-model", + tool_call_parser="llama3", + reasoning_parser=reasoning_parser, + ) + self.replies = replies + self.sent: TokenizeMsg | None = None + + def new_user(self) -> int: + return 42 + + async def send_one(self, msg): + self.sent = msg + + async def wait_for_ack(self, uid: int): + assert uid == 42 + for reply in self.replies: + yield reply + + +def reply(text: str, *, finished: bool = False, logprobs: dict | None = None) -> UserReply: + return UserReply( + uid=42, + incremental_output=text, + finished=finished, + prompt_tokens_delta=3 if not text else 0, + completion_tokens_delta=1 if text else 0, + cached_tokens=0, + logprobs=logprobs, + ) + + +def entry(token_id: int, token: str, logprob: float) -> dict: + return { + "token_id": token_id, + "token": token, + "bytes": list(token.encode("utf-8")), + "logprob": logprob, + "top": [ + { + "token_id": token_id, + "token": token, + "bytes": list(token.encode("utf-8")), + "logprob": logprob, + }, + { + "token_id": token_id + 1, + "token": "x", + "bytes": [120], + "logprob": logprob - 1, + }, + ], + } + + +def chat_request(**kwargs) -> ChatCompletionRequest: + payload = { + "model": "client-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 8, + } + payload.update(kwargs) + return ChatCompletionRequest(**payload) + + +def parse_sse(chunks: list[bytes]) -> list[dict | str]: + events: list[dict | str] = [] + for chunk in chunks: + for line in chunk.decode().splitlines(): + if line.startswith("data: "): + data = line.removeprefix("data: ") + events.append(data if data == "[DONE]" else json.loads(data)) + return events + + +def test_chat_non_stream_logprobs() -> None: + first = entry(1, "Hello", -0.1) + second = entry(2, "!", -0.2) + result = run( + handle_chat_completion( + chat_request(logprobs=True, top_logprobs=2), + None, + FakeState([reply("Hello", logprobs=first), reply("!", finished=True, logprobs=second)]), + {}, + ) + ) + + content = result["choices"][0]["logprobs"]["content"] + assert content == [ + { + "token": "Hello", + "logprob": -0.1, + "bytes": [72, 101, 108, 108, 111], + "top_logprobs": [ + {"token": "Hello", "logprob": -0.1, "bytes": [72, 101, 108, 108, 111]}, + {"token": "x", "logprob": -1.1, "bytes": [120]}, + ], + }, + { + "token": "!", + "logprob": -0.2, + "bytes": [33], + "top_logprobs": [ + {"token": "!", "logprob": -0.2, "bytes": [33]}, + {"token": "x", "logprob": -1.2, "bytes": [120]}, + ], + }, + ] + + without_logprobs = run( + handle_chat_completion( + chat_request(), + None, + FakeState([reply("Hello", finished=True, logprobs=first)]), + {}, + ) + ) + assert without_logprobs["choices"][0].get("logprobs") is None + + +def test_chat_stream_logprobs_follow_content_deltas() -> None: + first = entry(1, "Hello", -0.1) + state = FakeState( + [ + reply("Hello", logprobs=first), + reply(" world", finished=True), + ] + ) + req = chat_request(stream=True, logprobs=True, top_logprobs=2) + spec = chat_request_to_genspec(req, {}) + events = parse_sse(run(_collect(stream_chat_completion_chunks(42, req, state, spec)))) + + content_choices = [ + event["choices"][0] + for event in events + if isinstance(event, dict) + and event["choices"] + and event["choices"][0]["delta"].get("content") + ] + assert content_choices[0]["logprobs"]["content"][0]["token"] == "Hello" + assert "logprobs" not in content_choices[1] + + +def test_completion_logprobs_non_stream_and_stream() -> None: + first = entry(1, "Hi", -0.1) + second = entry(2, "!", -0.2) + req = CompletionRequest(model="client-model", prompt="hello", logprobs=2, max_tokens=8) + + result = run( + handle_completion( + req, + None, + FakeState([reply("Hi", logprobs=first), reply("!", finished=True, logprobs=second)]), + {}, + ) + ) + assert result["choices"][0]["logprobs"] == { + "tokens": ["Hi", "!"], + "token_logprobs": [-0.1, -0.2], + "top_logprobs": [ + {"Hi": -0.1, "x": -1.1}, + {"!": -0.2, "x": -1.2}, + ], + "text_offset": [0, 2], + } + + events = parse_sse( + run( + _collect( + stream_completion_chunks( + 42, + CompletionRequest( + model="client-model", + prompt="hello", + logprobs=2, + max_tokens=8, + stream=True, + ), + FakeState([reply("Hi", finished=True, logprobs=first)]), + ) + ) + ) + ) + chunk = next(event for event in events if isinstance(event, dict) and event["choices"][0]["text"]) + assert chunk["choices"][0]["logprobs"] == { + "tokens": ["Hi"], + "token_logprobs": [-0.1], + "top_logprobs": [{"Hi": -0.1, "x": -1.1}], + "text_offset": [0], + } + + +def test_logprobs_validation_errors() -> None: + chat_top_out_of_range = run( + handle_chat_completion(chat_request(logprobs=True, top_logprobs=25), None, FakeState([]), {}) + ) + assert chat_top_out_of_range.status_code == 400 + + chat_missing_flag = run( + handle_chat_completion(chat_request(top_logprobs=1), None, FakeState([]), {}) + ) + assert chat_missing_flag.status_code == 400 + + completion_out_of_range = run( + handle_completion( + CompletionRequest(model="client-model", prompt="hello", logprobs=7), + None, + FakeState([]), + {}, + ) + ) + assert completion_out_of_range.status_code == 400 + + completion_echo = run( + handle_completion( + CompletionRequest(model="client-model", prompt="hello", echo=True, logprobs=1), + None, + FakeState([]), + {}, + ) + ) + assert completion_echo.status_code == 400 + + +def test_reasoning_logprob_is_carried_to_next_content_delta() -> None: + reasoning_entry = entry(1, "thought", -0.1) + state = FakeState( + [ + reply("thought", logprobs=reasoning_entry), + reply("answer", finished=True), + ], + reasoning_parser="qwen3", + ) + req = chat_request(stream=True, logprobs=True, top_logprobs=2) + spec = chat_request_to_genspec(req, {}) + events = parse_sse(run(_collect(stream_chat_completion_chunks(42, req, state, spec)))) + + content_choice = next( + event["choices"][0] + for event in events + if isinstance(event, dict) + and event["choices"] + and event["choices"][0]["delta"].get("content") == "answer" + ) + assert content_choice["logprobs"]["content"][0]["token"] == "thought" + + +async def _collect(iterator): + return [chunk async for chunk in iterator]