From 104becfd7cbc342b1cf93c500d42fee3bd182daf Mon Sep 17 00:00:00 2001 From: xiezhipeng-git Date: Thu, 15 Oct 2026 18:09:08 +0800 Subject: [PATCH 01/14] OpenAI logprobs / top_logprobs / allowed_token_ids for the CUDA/HIP engine Adds per-request `logprobs` / `top_logprobs` to `/v1/chat/completions` (streamed or not) on the CUDA/HIP engine, plus `allowed_token_ids` with vLLM's semantics. Upstream already has this on SYCL (835d30b1) and Metal (31f26683); CUDA/HIP is the backend still missing it, and this fills it in the same shape: `logprobs=K` and `allowed=a,b,c` on the engine's GEN line, and an `LP [id:logprob ...]` line after each `T `, written in the same flush -- the same protocol the SYCL port speaks, so one server-side reader serves both. Engine half (src/program/generate.cpp): - two request keys in the `--serve` key=value line, skipped as unknown by an engine that predates them; - `emit_lp_line()` takes the logprob from the verify window's head logits before sampling, penalties and temperature, a max, a logsumexp and, for `top_logprobs`, a partial_sort over the vocabulary -- one printf per token; - the pipeline path asks `V1(A)` (the verifier of window A, whose tokens that loop prints) and the serial path `ver`, matching each token's own window; - off by default: a request without `logprobs` produces today's output, byte for byte, and the per-token cost is one integer compare. Server half (serve/server.py): - `top_logprobs` / `allowed_ids` translate the request; `sampling_keys` appends `logprobs=K` and `allowed=` only when asked; - `generate` reads each token's `LP` line before yielding the token (want_lp off = the path is unchanged); a held line is consumed before the next read, so an engine that emits no LP lines loses no lines, and a DONE read while looking for one still ends the request; - `LogprobEntries` decodes the token bytes incrementally (a character split across tokens lands in the one that completes it), so the entries, joined, spell the reply; the stop token is not listed; - `allowed_token_ids` keeps only the named ids and renormalises over them, so a caller asking about specific candidates gets a distribution over exactly those; - `structured_chunks` re-aligns the entries when `response_format: json_schema` rewrites the content: if a run of tokens spells what is returned, or spells the model's own JSON when that is the same JSON value as the canonical form, the content is returned as the model wrote it so the tokens still spell it; otherwise there are no logprobs -- never a set that does not match. Tests (serve/test_logprobs.py, no GPU): 17 items in python -m unittest serve.test_logprobs -- the request keys, the LP read (including the drain), the reply shape byte for byte, streaming, thinking's two parts each spelled by its own tokens, allowed_token_ids renormalised to 1, the schema cases, and the span helper. serve's own suites (server, structured, responses, logit_bias, security) pass unchanged: 392 tests. Off by default, additive only: a request that does not ask for logprobs gets the same bytes as before. --- serve/server.py | 248 +++++++++++++++++++++++++++++++++-- serve/test_logprobs.py | 270 +++++++++++++++++++++++++++++++++++++++ src/program/generate.cpp | 60 +++++++++ 3 files changed, 570 insertions(+), 8 deletions(-) create mode 100644 serve/test_logprobs.py diff --git a/serve/server.py b/serve/server.py index 5a119ec7cfc..1ef83d69b5d 100644 --- a/serve/server.py +++ b/serve/server.py @@ -73,6 +73,7 @@ from serve import logit_bias as bias_api # noqa: E402 from serve import lifeline # noqa: E402 from serve.structured import StructuredOutputError, prepare_format, validated_json # noqa: E402 +from serve.structured import _extract_json # noqa: E402 from serve import responses as responses_api # noqa: E402 from serve.responses import ResponsesError, error_body as responses_error_body # noqa: E402 from serve import images as images_api # noqa: E402 @@ -1394,6 +1395,15 @@ def sampling_keys(sampling: dict) -> str: pn = prefix.get("tokens") if isinstance(pn, int) and not isinstance(pn, bool) and pn > 0: keys += f" pin={pn}" + # logprobs: the engine writes an `LP` line after each `T` line (an engine without them ignores + # the key, so an older engine answers exactly as it did). Absent = no LP lines at all. + k = top_logprobs(sampling) + if k is not None: + keys += f" logprobs={k}" + ids = allowed_ids(sampling) + if ids: + # only these ids are printed; the engine's comma-separated list, renormalised below + keys += " allowed=" + ",".join(str(i) for i in ids) return keys + StrataEngine.projection_key(sampling) + bias_api.engine_key(sampling.get("logit_bias")) @staticmethod @@ -2133,6 +2143,12 @@ def generate(self, ids, max_new, sampling, cancel, embeddings=None): return self.progress, self.progress_ms, self.reused = None, 0, 0 self.prefill_tok_s_mean = None + # logprobs: the engine writes "LP [id:logprob ...]" right after each "T " (the same + # flush), read here before the token is yielded, so whoever takes the token finds its line in + # lp_lines. Not asked for = the deque stays empty and nothing on this path changes. + want_lp = top_logprobs(sampling or {}) is not None + self.lp_lines = collections.deque() + carry = NO_LINE # a line read for a token's LP line, held until it is used # an image request takes the same sampling keys as text (#75: it used to decode greedily whatever was asked) head = f"GENI {int(max_new)}{self.sampling_keys(sampling or {})} {embeddings}" if embeddings else \ f"GEN {int(max_new)}{self.sampling_keys(sampling or {})}" @@ -2160,7 +2176,10 @@ def generate(self, ids, max_new, sampling, cancel, embeddings=None): raise self._silent(f"the engine said nothing for {time.monotonic() - heard:.0f} s during " "the request") try: - line = self.lines.get(timeout=min(wait, 0.5)) + if carry is not NO_LINE: # the line read after a token for its LP line + line, carry = carry, NO_LINE + else: + line = self.lines.get(timeout=min(wait, 0.5)) except queue.Empty: if cancel.is_set(): return @@ -2187,7 +2206,18 @@ def generate(self, ids, max_new, sampling, cancel, embeddings=None): allow = silence if cancel.is_set(): return - yield int(line[2:]) + t = int(line[2:]) + if want_lp: + # the engine writes the token and its `LP` line in one flush, so the line is + # here now; an engine without logprobs leaves the next line a normal one + try: + carry = self.lines.get(timeout=0.5) + except queue.Empty: + carry = NO_LINE + if isinstance(carry, str) and carry.startswith("LP "): + self.lp_lines.append((t, carry)) + carry = NO_LINE + yield t elif line.startswith("PP "): f = line.split() if len(f) >= 3 and f[1].isdigit() and f[2].isdigit(): @@ -2230,9 +2260,12 @@ def generate(self, ids, max_new, sampling, cancel, embeddings=None): while True: left = allow - (time.monotonic() - heard) if allow > 0 else None try: - if left is not None and left <= 0: + if carry is not NO_LINE: # read already, after the last token + line, carry = carry, NO_LINE + elif left is not None and left <= 0: raise queue.Empty - line = self.lines.get(timeout=left) + else: + line = self.lines.get(timeout=left) except queue.Empty: raise self._silent("the engine did not finish the request after it was stopped (STOP) " f"within {allow:.0f} s") from None @@ -3243,6 +3276,31 @@ def push(self, t: int) -> str: OPENAI_MAX_STOP = 4 # OpenAI's limit on `stop` +def top_logprobs(req: dict): + """The K for the engine's logprobs=K, or None when the request did not ask. + + OpenAI: "logprobs": true with "top_logprobs" 0..20; the legacy completions form, + "logprobs": N, is taken as N.""" + lp = req.get("logprobs") + if lp is True: + k = req.get("top_logprobs") + return max(0, min(20, k)) if isinstance(k, int) and not isinstance(k, bool) else 0 + if isinstance(lp, int) and not isinstance(lp, bool) and lp >= 0: + return min(20, lp) + return None + + +def allowed_ids(req: dict) -> list: + """`allowed_token_ids`: the ids a caller named, or [] (vLLM's semantics, see LogprobEntries).""" + ids = req.get("allowed_token_ids") + if not isinstance(ids, (list, tuple)): + return [] + return [x for x in ids if isinstance(x, int) and not isinstance(x, bool) and x >= 0] + + +NO_LINE = object() # StrataEngine.generate: no line held after a token + + def stop_strings(req: dict) -> list[str]: """The request's stop strings: OpenAI's `stop` (a string or a list of up to 4) or Anthropic's `stop_sequences` (a list). Empty strings are left out. A bad value is a ValueError, so a 400 before anything is generated.""" @@ -5253,6 +5311,90 @@ def work(c=c, box=box): PP_DONE_TAIL = 64 + +class LogprobEntries: + """The engine's LP lines as OpenAI's logprobs.content entries. + + Each token's "token" is the text it completes (an incremental UTF-8 decode, as the reply's + text is made), so a character split across tokens lands in the token that finishes it and + the tokens, joined, spell the reply. "bytes" are the token's own bytes. The stop token + ends the reply and is not listed (OpenAI does not list it either). Tokens that make no + text of the reply are left out, because spell() gives each piece of text the tokens that + spell it. + + `allowed_token_ids` (vLLM's semantics): only those ids are kept, and the logprobs are + renormalised over them, so a caller asking about specific candidates gets a distribution + over exactly those instead of a top-K list it has to hope contains them. + """ + + def __init__(self, svc, allowed=()): + self.tok, self.stop_ids = svc.tok, svc.stop_ids + self.allowed = [i for i in allowed] + self.inc = codecs.getincrementaldecoder("utf-8")(errors="replace") + self.pending = [] # entries read but not yet given to a piece of text + + def spell(self, eng, text: str) -> list: + """The entries whose tokens spell `text`, the next piece of the reply.""" + self.pending += self.take(eng) + at = [0] + for e in self.pending: + at.append(at[-1] + len(e["token"])) + joined = "".join(e["token"] for e in self.pending) + i = joined.find(text) + while i >= 0 and text: + if i in at: + a = at.index(i) # the first token there + end = i + len(text) + b = next((j for j in range(a + 1, len(at)) if at[j] == end), None) + if b is not None: + out, self.pending = self.pending[a:b], self.pending[b:] + return out + i = joined.find(text, i + 1) + out, self.pending = self.pending, [] + return out + + def raw(self, t): + if hasattr(self.tok, "token_bytes"): + return self.tok.token_bytes(t) + return self.tok.decode([t]).encode("utf-8") + + def take(self, eng) -> list: + lines, out = getattr(eng, "lp_lines", None), [] + while lines: + t, line = lines.popleft() + if t in self.stop_ids: + continue + f = line.split() + if len(f) < 2 or f[1] == "nan": # no head logits for that window row: no distribution + continue + b = self.raw(t) + e = {"token": self.inc.decode(b), "logprob": float(f[1]), "bytes": list(b), + "top_logprobs": []} + alts = f[2:] + if self.allowed: + alts = [x for x in alts if int(x.partition(":")[0]) in self.allowed] + for x in alts: + i, _, v = x.partition(":") + alt = self.raw(int(i)) + e["top_logprobs"].append({"token": alt.decode("utf-8", "replace"), + "logprob": float(v), "bytes": list(alt)}) + if self.allowed: + e = self._renormalise(e) + out.append(e) + return out + + def _renormalise(self, e): + """vLLM: keep only the named ids and normalise over them, so they sum to 1.""" + if not e["top_logprobs"]: + return None + mx = max(x["logprob"] for x in e["top_logprobs"]) + se = sum(math.exp(x["logprob"] - mx) for x in e["top_logprobs"]) + lse = mx + math.log(se) + for x in e["top_logprobs"]: + x["logprob"] -= lse + e["logprob"] = e["top_logprobs"][0]["logprob"] if e["top_logprobs"] else e["logprob"] + return e + def prompt_progress(svc: Service): """The last PP line as llama.cpp's `prompt_progress`, so a client written for llama.cpp reads both servers: the same four fields, `time_ms` elapsed since prompt reading started, and `processed` counting from token 0 so that a @@ -5286,6 +5428,13 @@ def chunk(delta, finish=None): "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]} yield chunk({"role": "assistant", "content": ""}) + lp = LogprobEntries(svc, allowed_ids(req)) if top_logprobs(req) is not None else None + + def with_lp(c, key): + got = lp.spell(svc.engine, c["choices"][0]["delta"][key]) if lp is not None else None + if got: + c["choices"][0]["logprobs"] = {key: got} + return c calls = 0 streamed = {} # tool call id -> index, for calls sent piece by piece finished = set() # ... and the ones whose final tool_call came (#211) @@ -5302,9 +5451,9 @@ def chunk(delta, finish=None): elif kind == "event": ev: Event = x if ev.kind == "reasoning" and ev.text: - yield chunk({"reasoning_content": ev.text}) + yield with_lp(chunk({"reasoning_content": ev.text}), "reasoning_content") elif ev.kind == "content" and ev.text: - yield chunk({"content": ev.text}) + yield with_lp(chunk({"content": ev.text}), "content") elif ev.kind == "tool_start": streamed[ev.call.id] = calls calls += 1 @@ -5344,11 +5493,14 @@ def _is_json(text: str) -> bool: def openai_collect(chunks) -> dict: content, reasoning, by_index, last, mcp = [], [], {}, None, [] + lps = {} # "content" / "reasoning_content" -> logprob entries for c in chunks: if c is None: # a heartbeat continue if c.get("strata_mcp"): mcp.append(c["strata_mcp"]) + for k, v in (c["choices"][0].get("logprobs") or {}).items(): + lps.setdefault(k, []).extend(v) d = c["choices"][0]["delta"] content.append(d.get("content") or "") reasoning.append(d.get("reasoning_content") or "") @@ -5373,11 +5525,80 @@ def openai_collect(chunks) -> dict: out = {"id": last["id"], "object": "chat.completion", "created": last["created"], "model": last["model"], "choices": [{"index": 0, "message": msg, "finish_reason": last["choices"][0]["finish_reason"]}], "usage": last["usage"]} + if lps: + out["choices"][0]["logprobs"] = {"content": lps.get("content", []), "refusal": None, **( + {"reasoning_content": lps["reasoning_content"]} if lps.get("reasoning_content") else {})} if last.get("timings"): out["timings"] = last["timings"] return out +def logprob_span(text: str, part: str, entries: list): + """The entries whose tokens spell `part`, a piece of `text`, or None when the tokens do not + spell `text` or no token boundary falls at both ends of `part`.""" + if "".join(e["token"] for e in entries) != text: + return None + start = text.find(part) + if start < 0: + return None + at, first = 0, None + for i, e in enumerate(entries): + if at == start and first is None: + first = i + if at == start + len(part) and first is not None: + return entries[first:i] + at += len(e["token"]) + return entries[first:] if first is not None and at == start + len(part) else None + + +def _same_json(a: str, b: str) -> bool: + """Are two strings the same JSON value (whitespace and key order aside)? No -> False.""" + try: + return json.loads(a) == json.loads(b) + except Exception: + return False + + +def _structured_logprobs(choice: dict, content: str): + """`content`'s logprobs, and the content the tokens actually spell. + + Returns `(content, logprobs)`: on failure `(the content as validated, None)`. + + `structured_chunks` returns the canonical JSON the validator produced, which the model never + generated, so its tokens do not spell it. Two honest outcomes: + + 1. a run of tokens spells what is being returned -> use it as it is; + 2. the tokens spell the model's own JSON, which is the same JSON value (the fenced-with-spaces + case, the common one) -> return the model's form, so `''.join(token) == content` again. + + Neither holds (a rearranged key order, or a reply that is not that JSON at all) and there are no + logprobs: a caller that joins the tokens must not silently get something that does not match. + """ + lps = choice.get("logprobs") + if not lps or not lps.get("content"): + return content, None + raw = choice["message"]["content"] or "" + entries = lps["content"] + if raw == content: # not rewritten: the entries are already its + return content, _with_reasoning(lps, entries) + span = logprob_span(raw, content, entries) + if span is not None: + return content, _with_reasoning(lps, span) + inner = _extract_json(raw) if _extract_json else "" + if inner and _same_json(inner, content): + span = logprob_span(raw, inner, entries) + if span is not None: + return inner, _with_reasoning(lps, span) # the model's own form: the tokens spell it + return content, None + + +def _with_reasoning(lps: dict, entries: list) -> dict: + out = {"content": entries} + if lps.get("reasoning_content"): + out["reasoning_content"] = lps["reasoning_content"] + return out + + def structured_chunks(chunks, validator): """Buffer structured streams so a client never receives unvalidated content.""" buffered = [] @@ -5393,12 +5614,23 @@ def structured_chunks(chunks, validator): result = openai_collect(buffered) choice = result["choices"][0] content = validated_json(choice["message"]["content"], validator, choice["finish_reason"]) + # logprobs: the entries spell what the model wrote, which this function replaces with the + # canonical JSON. `_structured_logprobs` gives back the form the tokens spell, or None. + content, logprobs = _structured_logprobs(choice, content) yield buffered[0] delta = {"content": content} if choice["message"].get("reasoning_content"): delta["reasoning_content"] = choice["message"]["reasoning_content"] - yield {**buffered[0], "choices": [{"index": 0, "delta": delta, "finish_reason": None}]} - yield buffered[-1] + piece = {"index": 0, "delta": delta, "finish_reason": None} + if logprobs: + piece["logprobs"] = logprobs + yield {**buffered[0], "choices": [piece]} + end = buffered[-1] + if end["choices"][0].get("logprobs"): + # those entries went out in the piece above (or were dropped with it): the closing + # chunk must not carry a second copy + end = {**end, "choices": [{k: v for k, v in end["choices"][0].items() if k != "logprobs"}]} + yield end finally: chunks.close() diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py new file mode 100644 index 00000000000..ef801a4d8d1 --- /dev/null +++ b/serve/test_logprobs.py @@ -0,0 +1,270 @@ +"""OpenAI logprobs on /v1/chat/completions, from the engine's LP lines: python -m unittest serve.test_logprobs. + +The engine (src/program/generate.cpp) writes an `LP [id:logprob ...]` line right after +each `T ` when the request carries `logprobs=K`; this is the server half that turns those +lines into OpenAI's logprobs, and the parts that are easy to get wrong: + + - `allowed_token_ids`: only the named ids, renormalised over them (vLLM's semantics). + - `response_format: json_schema` rewrites the content, so the tokens no longer spell what is + returned; the entries are re-aligned or dropped, never returned wrong. +""" +from __future__ import annotations + +import collections +import json +import queue +import threading +import unittest +import urllib.request +from pathlib import Path + +from serve.frontend import ChatTemplate +from serve.server import (ByteTokenizer, MockEngine, Service, StrataEngine, logprob_span, + serve, top_logprobs, allowed_ids) + +ROOT = Path(__file__).resolve().parent.parent +TEMPLATE = ChatTemplate(ROOT / "serve" / "chat_template.jinja") + + +class Tokenizer(ByteTokenizer): + def token_bytes(self, token): + return self.SPECIALS[token - 256].encode() if token >= 256 else bytes([token]) + + +class ThinkTokenizer(Tokenizer): + # two extra special tokens that decode as the thinking marks; ALWAYS so the byte-level + # encode() maps them wherever they appear ("" would loop forever in startswith, so they + # are real strings, like the engine's own think tokens are) + SPECIALS = ByteTokenizer.SPECIALS + ["", ""] + ALWAYS = ("", "") + + +class LpEngine(MockEngine): + """MockEngine with the LP line of each token on lp_lines before the token is yielded, as the + engine leaves it; token i has the log-probability -(i + 1) / 100 and four named alternatives.""" + + def __init__(self, tok, script, lines=True): + super().__init__(tok, script, max_context=16384) + self.with_lines, self.sampling = lines, None + + def generate(self, ids, max_new, sampling, cancel, embeddings=None): + self.sampling, self.lp_lines = dict(sampling or {}), collections.deque() + for i, t in enumerate(super().generate(ids, max_new, sampling, cancel, embeddings)): + if self.with_lines: + self.lp_lines.append(( + t, f"LP {-(i + 1) / 100:.6f} 65:-0.500000 72:-1.250000 " + f"88:-3.000000 91:-4.200000\n")) + yield t + + +SCHEMA = {"type": "json_schema", "json_schema": {"name": "a", "strict": True, "schema": { + "type": "object", "properties": {"rent": {"type": "string"}, "quote": {"type": "string"}}, + "required": ["rent", "quote"], "additionalProperties": False}}} + + +class Keys(unittest.TestCase): + """What the request asks for, and what reaches the engine line.""" + + def test_what_the_request_asks(self): + self.assertEqual(top_logprobs({"logprobs": True}), 0) + self.assertEqual(top_logprobs({"logprobs": True, "top_logprobs": 3}), 3) + self.assertEqual(top_logprobs({"logprobs": True, "top_logprobs": 50}), 20) + self.assertEqual(top_logprobs({"logprobs": 2}), 2) # the legacy completions form + for req in ({}, {"logprobs": False}, {"logprobs": None}, {"top_logprobs": 3}): + self.assertIsNone(top_logprobs(req)) + + def test_allowed_ids(self): + self.assertEqual(allowed_ids({"allowed_token_ids": [1, 2, 3]}), [1, 2, 3]) + for req in ({}, {"allowed_token_ids": "1,2"}, {"allowed_token_ids": [True]}, {"allowed_token_ids": [-1]}): + self.assertEqual(allowed_ids(req), []) + + def test_the_engine_key_only_when_asked(self): + self.assertIn(" logprobs=0", StrataEngine.sampling_keys({"logprobs": True})) + self.assertIn(" logprobs=5", StrataEngine.sampling_keys({"logprobs": True, "top_logprobs": 5})) + self.assertNotIn("logprobs", StrataEngine.sampling_keys({})) + self.assertNotIn("logprobs", StrataEngine.sampling_keys({"logprobs": False, "temperature": 0})) + + def test_allowed_reaches_the_engine_line(self): + keys = StrataEngine.sampling_keys({"logprobs": True, "top_logprobs": 2, + "allowed_token_ids": [7, 8, 9]}) + self.assertIn(" logprobs=2 ", keys) + self.assertIn(" allowed=7,8,9", keys) + + +class Generate(unittest.TestCase): + """The token's LP line is read next to it, and only when it was asked for.""" + + def engine(self, *lines): + e = StrataEngine.__new__(StrataEngine) + e.proc = __import__("types").SimpleNamespace( + stdin=__import__("io").StringIO(), poll=lambda: None) + e.lines, e.can_stop, e.max_context, e.silence_s = queue.Queue(), False, 262144, 0 + for line in lines: + e.lines.put(line) + return e + + def test_a_token_is_yielded_with_its_lp_line_read(self): + e = self.engine("PP 3 3", "T 72", "LP -0.100000 72:-0.1 65:-2.5", "T 105", + "LP -0.200000 105:-0.2 65:-3", "T 257", "LP -0.010000", + "DONE 3 3 0 3 stop 0 0 0") + seen = [] + for t in e.generate([1, 2, 3], 8, {"logprobs": True, "top_logprobs": 2}, threading.Event()): + if t is not None: + seen.append((t, [x[0] for x in e.lp_lines])) + self.assertEqual(seen, [(72, [72]), (105, [72, 105]), (257, [72, 105, 257])]) + self.assertIn(" logprobs=2 ", e.proc.stdin.getvalue()) + self.assertEqual(e.lp_lines[1][1], "LP -0.200000 105:-0.2 65:-3") + + def test_an_engine_without_lp_lines(self): + e = self.engine("T 72", "T 105", "DONE 3 3 0 2 stop 0 0 0") + toks = [t for t in e.generate([1], 8, {"logprobs": True}, threading.Event()) if t is not None] + self.assertEqual(toks, [72, 105]) + self.assertEqual(list(e.lp_lines), []) + self.assertEqual(e.last["finish"], "stop") + + def test_not_asked_reads_as_before(self): + e = self.engine("T 72", "LP -0.1", "T 105", "DONE 3 3 0 2 stop 0 0 0") + toks = [t for t in e.generate([1], 8, {}, threading.Event()) if t is not None] + self.assertEqual(toks, [72, 105]) + self.assertEqual(list(e.lp_lines), []) + + def test_a_consumer_that_stops_early_does_not_lose_the_done_read_after_the_token(self): + # the DONE line was read while looking for the token's LP line: the drain must take it, + # not wait for another + e = self.engine("T 72", "DONE 3 3 0 1 length 0 0 0", None) + gen = e.generate([1], 1, {"logprobs": True}, threading.Event()) + self.assertEqual(next(gen), 72) + gen.close() + self.assertEqual(e.last["finish"], "length") + + +class Api(unittest.TestCase): + def setUp(self): + self.tok = Tokenizer() + self.engine = LpEngine(self.tok, "Café costs €5.") + self.svc = Service(self.engine, self.tok, TEMPLATE) + self.httpd = serve(self.svc, port=0) + self.base = f"http://127.0.0.1:{self.httpd.server_address[1]}" + + def tearDown(self): + self.httpd.shutdown() + self.httpd.server_close() + + def post(self, body, stream=False): + body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 200, + "chat_template_kwargs": {"enable_thinking": False}, **body} + if stream: + body["stream"] = True + req = urllib.request.Request(self.base + "/v1/chat/completions", json.dumps(body).encode(), + {"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=30) as r: + raw = r.read().decode() + if not stream: + return json.loads(raw) + return [json.loads(x[6:]) for x in raw.splitlines() if x.startswith("data: {")] + + def test_the_tokens_spell_the_content(self): + out = self.post({"logprobs": True, "top_logprobs": 1}) + choice = out["choices"][0] + content = choice["message"]["content"] + lps = choice["logprobs"]["content"] + self.assertEqual(content, "Café costs €5.") + self.assertEqual("".join(e["token"] for e in lps), content) + self.assertEqual(len(lps), len(content.encode())) # one per byte token; no stop token + self.assertEqual(b"".join(bytes(e["bytes"]) for e in lps), content.encode()) + e_acute = [e for e in lps if len(e["bytes"]) == 1 and e["bytes"][0] >= 0x80] + self.assertEqual([e["token"] for e in e_acute[:2]], ["", "é"]) # split across two tokens + self.assertEqual(lps[0]["logprob"], -0.01) + self.assertEqual(lps[0]["top_logprobs"][0]["token"], "A") + self.assertEqual(self.engine.sampling.get("logprobs"), True) + + def test_streamed(self): + chunks = self.post({"logprobs": True}, stream=True) + content = "".join(c["choices"][0]["delta"].get("content") or "" for c in chunks if c.get("choices")) + toks = [e["token"] for c in chunks if c.get("choices") + for e in (c["choices"][0].get("logprobs") or {}).get("content", [])] + self.assertEqual(content, "Café costs €5.") + self.assertEqual("".join(toks), content) + + def test_with_thinking_each_part_is_spelled_by_its_tokens(self): + tok = ThinkTokenizer() + svc = Service(LpEngine(tok, "Two times three.\n\nSix."), tok, + TEMPLATE) + httpd = serve(svc, port=0) + try: + self.base = f"http://127.0.0.1:{httpd.server_address[1]}" + for stream in (False, True): + out = self.post({"logprobs": True, "chat_template_kwargs": {"enable_thinking": True}}, + stream) + chunks = out if stream else [{"choices": [{"delta": out["choices"][0]["message"], + "logprobs": out["choices"][0]["logprobs"]}]}] + for key in ("content", "reasoning_content"): + text = "".join(c["choices"][0]["delta"].get(key) or "" for c in chunks if c.get("choices")) + toks = "".join(e["token"] for c in chunks if c.get("choices") + for e in (c["choices"][0].get("logprobs") or {}).get(key, [])) + self.assertTrue(text) + self.assertEqual(toks, text, (key, stream)) + self.assertNotIn("think>", toks) + finally: + httpd.shutdown() + httpd.server_close() + + def test_not_asked_no_logprobs(self): + out = self.post({}) + self.assertNotIn("logprobs", out["choices"][0]) + + def test_an_engine_without_lp_lines_gives_none(self): + self.engine.with_lines = False + out = self.post({"logprobs": True}) + self.assertEqual(out["choices"][0]["message"]["content"], "Café costs €5.") + self.assertNotIn("logprobs", out["choices"][0]) + + def test_allowed_token_ids_are_renormalised_over_the_named_ids(self): + out = self.post({"logprobs": True, "allowed_token_ids": [65, 72]}) + lps = out["choices"][0]["logprobs"]["content"] + self.assertTrue(lps) + for e in lps: + ids = [t["token"] for t in e["top_logprobs"]] + self.assertEqual(ids, ["A", "H"]) # only the named ones + import math + total = sum(math.exp(t["logprob"]) for t in lps[0]["top_logprobs"]) + self.assertAlmostEqual(total, 1.0, places=6) # renormalised: they sum to 1 + + def test_json_schema_answer_as_written_with_its_logprobs(self): + # MockEngine appends the stop token to the script it is built with, and reports + # finish_reason=length when it runs out -- so the engine is rebuilt with this answer + # rather than re-pointed at a raw token list (which would drop the stop token). + written = '```json\n{"rent": "$1,250", "quote": "x"}\n```' + self.engine = self.svc.engine = LpEngine(self.tok, written) + out = self.post({"logprobs": True, "temperature": 0, "response_format": SCHEMA}) + choice = out["choices"][0] + content = choice["message"]["content"] + self.assertEqual(content, '{"rent": "$1,250", "quote": "x"}') # the JSON as written + self.assertEqual("".join(e["token"] for e in choice["logprobs"]["content"]), content) + streamed = self.post({"logprobs": True, "response_format": SCHEMA}, stream=True) + toks = [e["token"] for c in streamed if c.get("choices") + for e in (c["choices"][0].get("logprobs") or {}).get("content", [])] + self.assertEqual("".join(toks), content) + plain = self.post({"temperature": 0, "response_format": SCHEMA}) # not asked: as before + self.assertEqual(plain["choices"][0]["message"]["content"], + '{"rent":"$1,250","quote":"x"}') + self.assertNotIn("logprobs", plain["choices"][0]) + + +class Span(unittest.TestCase): + def entries(self, *toks): + return [{"token": t} for t in toks] + + def test_span(self): + es = self.entries("```", "json", "\n", "{", "}", "\n", "```") + self.assertEqual([e["token"] for e in logprob_span("```json\n{}\n```", "{}", es)], ["{", "}"]) + self.assertEqual(len(logprob_span("{}", "{}", self.entries("{", "}"))), 2) + + def test_no_span(self): + self.assertIsNone(logprob_span("x{}", "{}", self.entries("x{", "}"))) # no boundary at the start + self.assertIsNone(logprob_span("{}x", "{}", self.entries("{", "}x"))) # none at the end + self.assertIsNone(logprob_span("{}", "{}", self.entries("{"))) # do not spell it + + +if __name__ == "__main__": + unittest.main() diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 01b99f06cf2..01ed09de45b 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -144,6 +144,46 @@ namespace strata::metal_engine { void set_model_shard(const char* shard1); } / #include #include +// logprobs=K / allowed=...: one `LP` line for the token just emitted, from `ver`'s head +// logits for the window row it came from (before sampling, penalties and temperature). +// Format matches the SYCL port's LP line so one server-side reader serves both. +void emit_lp_line(const strata::core::Verifier& ver, int t, int32_t token, int k, + const std::vector& allowed) { + if (k < 0) return; + static std::vector row; + row.resize((size_t) ver.vocab()); + if (!ver.copy_logits(t, row.data())) { + std::printf("LP nan\n"); + return; + } + float mx = -INFINITY; + for (const float v : row) mx = std::max(mx, v); + double se = 0.0; + for (const float v : row) se += std::exp((double) v - (double) mx); + const double lse = (double) mx + std::log(se); + std::printf("LP %.6f", (double) row[(size_t) token] - lse); + if (!allowed.empty()) { + // 只要调用方点名的那几个 id:server 侧再按它们归一化 ⇒ 与 allowed_token_ids 语义一致 + for (const int32_t id : allowed) { + if (id == token) continue; + std::printf(" %d:%.6f", (int) id, (double) row[(size_t) id] - lse); + } + std::printf("\n"); + return; + } + if (k > 0) { + static std::vector ord; + ord.resize(row.size()); + for (size_t v = 0; v < row.size(); ++v) ord[v] = (int32_t) v; + const int kk = std::min(k, (int) row.size()); + std::partial_sort(ord.begin(), ord.begin() + kk, ord.end(), + [&](int32_t a, int32_t b) { return row[(size_t) a] > row[(size_t) b]; }); + for (int j = 0; j < kk; ++j) + std::printf(" %d:%.6f", ord[(size_t) j], (double) row[(size_t) ord[(size_t) j]] - lse); + } + std::printf("\n"); +} + namespace { // Windows' WDDM driver model: native Windows, or WSL2 (its GPU goes through /dev/dxg to the Windows driver). There, // pinning a large arena into two CUDA contexts leaves WDDM refusing every later allocation (the 5080 + 3090 rig); @@ -11948,6 +11988,11 @@ int main(int argc, char** argv) { std::string req_bias_text; bool req_bias_present = false; int req_cvec = 1; // cvec=0|1: a loaded control vector for this request (on when absent) + // logprobs=K (0..20): after each "T id" an "LP logprob id:logprob ..." line. + // -1 (absent) = no LP lines, so a server that never asks sees today's output. + int req_logprobs = -1; + // allowed=a,b,c: with logprobs, print only these token ids (renormalized server-side). + std::vector req_allowed; // ckpt=0: a one-shot call whose turn no later request extends. No checkpoint at its last turn boundary // (so no split there) nor every --prompt-cache-every tokens, and its session is neither continued nor // parked after it. It still resumes from a checkpoint it matches, and still saves the system-prompt root @@ -11987,6 +12032,18 @@ int main(int argc, char** argv) { else if (key == "temperature") req_temperature = fv; else if (key == "top_p") req_top_p = fv; else if (key == "top_k") req_top_k = std::atoi(tok.c_str() + eq + 1); + else if (key == "logprobs") req_logprobs = std::atoi(tok.c_str() + eq + 1); + else if (key == "allowed") { + const std::string csv = tok.substr(eq + 1); + size_t p = 0; + while (p < csv.size()) { + const size_t c = csv.find(",", p); + const std::string one = csv.substr(p, c == std::string::npos ? std::string::npos : c - p); + if (!one.empty()) req_allowed.push_back((int32_t) std::atoi(one.c_str())); + if (c == std::string::npos) break; + p = c + 1; + } + } else if (key == "min_p") req_min_p = fv; else if (key == "penalty_last_n") req_penalty_last_n = std::atoi(tok.c_str() + eq + 1); else if (key == "penalty_repeat") req_penalty_repeat = fv; @@ -12034,6 +12091,7 @@ int main(int argc, char** argv) { } const int64_t n = (int64_t) ids.size(); req_imgs.clear(); + req_allowed.clear(); if (geni && !o.vision) { std::printf("ERR this engine was started without --vision\n"); continue; } if (geni || !mrope_identity) { // positions for every cell this request can reach; the identity again for a text request @@ -14369,6 +14427,7 @@ int main(int argc, char** argv) { bool eos = false; for (int i = 0; !nf && i <= a && produced_n < max_new && !eos; ++i) { std::printf("T %d\n", (int) outp[(size_t) i]); + emit_lp_line(V1(A), i, outp[(size_t) i], req_logprobs, req_allowed); strata::core::progress_beat(); ++produced_n; if (o.suffix_draft > 0) sfx.append(outp[(size_t) i]); @@ -14782,6 +14841,7 @@ int main(int argc, char** argv) { bool eos = false; for (int i = merged ? a : 0; i <= a && produced_n < max_new && !eos; ++i) { // (merged: row T-1 only) std::printf("T %d\n", (int) outv[(size_t) i]); + emit_lp_line(ver, i, outv[(size_t) i], req_logprobs, req_allowed); strata::core::progress_beat(); ++produced_n; if (sfx_on) sfx.append(outv[(size_t) i]); From a4349ec95e10da985d62a62968ed71cad2211470 Mon Sep 17 00:00:00 2001 From: Leonardo Murillo Date: Thu, 22 Oct 2026 11:45:47 -0300 Subject: [PATCH 02/14] serve: logprobs lines belong to the request that asked for them; bounded replies; unreadable LP lines skipped - LPLines: the engine's lp_lines deque is tagged with the request's cancel Event (one per request, passed unchanged from the handler through Service.run to Engine.generate). LogprobEntries is created with the same Event and reads only lines carrying it. A reply's tail (text held by the parser or a stop string) is spelled after the request has released the queue; by then the next request may have replaced lp_lines with its own. The tail now reads its own lines and never consumes the next request's. - logprobs require max_tokens / max_completion_tokens 1..4096 (400 otherwise): entries of a non-streamed or json_schema reply are held until it ends. - an LP line that does not parse (nan, -nan(ind), a cut line) or carries a non-finite value gives no entry for that token instead of failing the request. - tests: two requests on one engine where the first one's last "." is held by its stop string and sent while the second is generating (each reply's tokens spell that reply); unreadable lines; the max_tokens bound. Follows #3200. --- serve/server.py | 81 ++++++++++++++++++++++---------- serve/test_logprobs.py | 104 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 157 insertions(+), 28 deletions(-) diff --git a/serve/server.py b/serve/server.py index 1ef83d69b5d..5a747ad6a42 100644 --- a/serve/server.py +++ b/serve/server.py @@ -2147,7 +2147,7 @@ def generate(self, ids, max_new, sampling, cancel, embeddings=None): # flush), read here before the token is yielded, so whoever takes the token finds its line in # lp_lines. Not asked for = the deque stays empty and nothing on this path changes. want_lp = top_logprobs(sampling or {}) is not None - self.lp_lines = collections.deque() + self.lp_lines = LPLines(cancel) carry = NO_LINE # a line read for a token's LP line, held until it is used # an image request takes the same sampling keys as text (#75: it used to decode greedily whatever was asked) head = f"GENI {int(max_new)}{self.sampling_keys(sampling or {})} {embeddings}" if embeddings else \ @@ -3300,6 +3300,11 @@ def allowed_ids(req: dict) -> list: NO_LINE = object() # StrataEngine.generate: no line held after a token +# A reply with logprobs keeps every entry (its token, bytes and up to 20 alternatives) until it ends when it is not +# streamed or answers a json_schema, and spell() matches text against the entries still pending: a bounded reply +# keeps that memory and work bounded. +LOGPROBS_MAX_TOKENS = 4096 + def stop_strings(req: dict) -> list[str]: """The request's stop strings: OpenAI's `stop` (a string or a list of up to 4) or Anthropic's `stop_sequences` @@ -5312,6 +5317,17 @@ def work(c=c, box=box): +class LPLines(collections.deque): + """One request's LP lines, tagged with its cancel Event (one per request, handed unchanged from the handler + through Service.run to Engine.generate). The last events of a request - the parser's and the stop strings' + tail - come after it has left the queue, when the next request may already have replaced the engine's lp_lines + with its own: LogprobEntries reads only the lines tagged with its own request's Event.""" + + def __init__(self, owner=None): + super().__init__() + self.owner = owner + + class LogprobEntries: """The engine's LP lines as OpenAI's logprobs.content entries. @@ -5327,11 +5343,13 @@ class LogprobEntries: over exactly those instead of a top-K list it has to hope contains them. """ - def __init__(self, svc, allowed=()): + def __init__(self, svc, allowed=(), owner=None): self.tok, self.stop_ids = svc.tok, svc.stop_ids self.allowed = [i for i in allowed] self.inc = codecs.getincrementaldecoder("utf-8")(errors="replace") self.pending = [] # entries read but not yet given to a piece of text + self.owner = owner # the request's cancel Event (LPLines) + self.lines = None # this request's current LPLines, kept once seen def spell(self, eng, text: str) -> list: """The entries whose tokens spell `text`, the next piece of the reply.""" @@ -5359,28 +5377,39 @@ def raw(self, t): return self.tok.decode([t]).encode("utf-8") def take(self, eng) -> list: - lines, out = getattr(eng, "lp_lines", None), [] - while lines: - t, line = lines.popleft() - if t in self.stop_ids: - continue - f = line.split() - if len(f) < 2 or f[1] == "nan": # no head logits for that window row: no distribution - continue - b = self.raw(t) - e = {"token": self.inc.decode(b), "logprob": float(f[1]), "bytes": list(b), - "top_logprobs": []} - alts = f[2:] - if self.allowed: - alts = [x for x in alts if int(x.partition(":")[0]) in self.allowed] - for x in alts: - i, _, v = x.partition(":") - alt = self.raw(int(i)) - e["top_logprobs"].append({"token": alt.decode("utf-8", "replace"), - "logprob": float(v), "bytes": list(alt)}) - if self.allowed: - e = self._renormalise(e) - out.append(e) + # the engine's current lines are this request's only when they carry its Event (a later generate pass + # of the same request makes new ones: the old ones are read to the end first); another request's never are + cur, old = getattr(eng, "lp_lines", None), None + if cur is not self.lines and getattr(cur, "owner", None) is self.owner: + old, self.lines = self.lines, cur + out = [] + for lines in (old, self.lines): + while lines: + t, line = lines.popleft() + if t in self.stop_ids: + continue + f = line.split() + if len(f) < 2: + continue + try: # "nan" (no head logits for that window row), "-nan(ind)" + lp = float(f[1]) # (MSVC) or a cut line: no entry for that token + alts = [(int(i), float(v)) for i, _, v in (x.partition(":") for x in f[2:])] + except ValueError: + continue + if not math.isfinite(lp) or not all(math.isfinite(v) for _, v in alts): + continue + b = self.raw(t) + e = {"token": self.inc.decode(b), "logprob": lp, "bytes": list(b), + "top_logprobs": []} + if self.allowed: + alts = [(i, v) for i, v in alts if i in self.allowed] + for i, v in alts: + alt = self.raw(i) + e["top_logprobs"].append({"token": alt.decode("utf-8", "replace"), + "logprob": v, "bytes": list(alt)}) + if self.allowed: + e = self._renormalise(e) + out.append(e) return out def _renormalise(self, e): @@ -5428,7 +5457,7 @@ def chunk(delta, finish=None): "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]} yield chunk({"role": "assistant", "content": ""}) - lp = LogprobEntries(svc, allowed_ids(req)) if top_logprobs(req) is not None else None + lp = LogprobEntries(svc, allowed_ids(req), cancel) if top_logprobs(req) is not None else None def with_lp(c, key): got = lp.spell(svc.engine, c["choices"][0]["delta"][key]) if lp is not None else None @@ -6809,6 +6838,8 @@ def _openai(self, req): raise ValueError("structured response_format with tools/MCP is not supported") svc.load() max_req = max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 0) # 0/-1: the rest + if top_logprobs(req) is not None and not 0 < max_req <= LOGPROBS_MAX_TOKENS: + raise ValueError(f"logprobs: max_tokens (or max_completion_tokens) 1..{LOGPROBS_MAX_TOKENS} is required") bias_api.validate_request(req, svc.engine, len(getattr(svc.tok, "tokens", ())) or None) use_mcp = req.get("strata_mcp") is True and svc.mcp is not None # the web app's opt-in (serve/mcp.py) own = {t.get("name") for t in tools or [] if isinstance(t, dict)} # #592: a second line of defence diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index ef801a4d8d1..96aa4022924 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -10,16 +10,16 @@ """ from __future__ import annotations -import collections import json import queue import threading import unittest import urllib.request from pathlib import Path +from types import SimpleNamespace from serve.frontend import ChatTemplate -from serve.server import (ByteTokenizer, MockEngine, Service, StrataEngine, logprob_span, +from serve.server import (ByteTokenizer, LogprobEntries, LPLines, MockEngine, Service, StrataEngine, logprob_span, serve, top_logprobs, allowed_ids) ROOT = Path(__file__).resolve().parent.parent @@ -48,7 +48,7 @@ def __init__(self, tok, script, lines=True): self.with_lines, self.sampling = lines, None def generate(self, ids, max_new, sampling, cancel, embeddings=None): - self.sampling, self.lp_lines = dict(sampling or {}), collections.deque() + self.sampling, self.lp_lines = dict(sampling or {}), LPLines(cancel) for i, t in enumerate(super().generate(ids, max_new, sampling, cancel, embeddings)): if self.with_lines: self.lp_lines.append(( @@ -57,6 +57,59 @@ def generate(self, ids, max_new, sampling, cancel, embeddings=None): yield t +class NextEngine(LpEngine): + """LpEngine for two requests: the second one stops after its first token (its LP line is in lp_lines) + until the first request's reply has been read.""" + + def __init__(self, tok, scripts, first_out): + super().__init__(tok, scripts) + self.first_out, self.calls = first_out, 0 + self.first_in, self.first_generated, self.second_started = (threading.Event(), threading.Event(), + threading.Event()) + + def generate(self, ids, max_new, sampling, cancel, embeddings=None): + self.calls += 1 + mine = self.calls + if mine == 1: + self.first_in.set() + try: + for i, t in enumerate(super().generate(ids, max_new, sampling, cancel, embeddings)): + if mine == 2 and i == 0: + self.second_started.set() + self.first_out.wait(10) + yield t + finally: + if mine == 1: + self.first_generated.set() + + +class HandOver: + """Service.fifo that, once the first request has generated, lets its turn go and then waits until the + next request is generating before the first one goes on.""" + + def __init__(self, engine): + self.lock, self.engine, self.waited = threading.Lock(), engine, False + + def acquire(self, *a, **k): + return self.lock.acquire(*a, **k) + + def release(self): + self.lock.release() + if self.engine.first_generated.is_set() and not self.waited: + self.waited = True + self.engine.second_started.wait(10) + + def locked(self): + return self.lock.locked() + + def __enter__(self): + self.acquire() + return self + + def __exit__(self, *exc): + self.release() + + SCHEMA = {"type": "json_schema", "json_schema": {"name": "a", "strict": True, "schema": { "type": "object", "properties": {"rent": {"type": "string"}, "quote": {"type": "string"}}, "required": ["rent", "quote"], "additionalProperties": False}}} @@ -90,6 +143,15 @@ def test_allowed_reaches_the_engine_line(self): self.assertIn(" logprobs=2 ", keys) self.assertIn(" allowed=7,8,9", keys) + def test_unreadable_lines_give_no_entry(self): + a = threading.Event() + eng = SimpleNamespace(lp_lines=LPLines(a)) + for line in ("LP nan", "LP -nan(ind)", "LP inf", "LP -0.1 65:x", "LP -0.1 65:-inf", "LP -0.1 65", + "LP -0.200000 65:-1.0"): + eng.lp_lines.append((72, line)) + got = LogprobEntries(SimpleNamespace(tok=Tokenizer(), stop_ids=set()), owner=a).take(eng) + self.assertEqual([e["logprob"] for e in got], [-0.2]) + class Generate(unittest.TestCase): """The token's LP line is read next to it, and only when it was asked for.""" @@ -178,6 +240,42 @@ def test_the_tokens_spell_the_content(self): self.assertEqual(lps[0]["top_logprobs"][0]["token"], "A") self.assertEqual(self.engine.sampling.get("logprobs"), True) + def test_a_request_reads_only_its_own_lines(self): + # A's last "." is held by its stop string and sent after A has left the queue, when B is + # already generating on the same engine: each reply's tokens must spell that reply + first_out = threading.Event() + engine = NextEngine(self.tok, ["Café costs €5.", "Bonjour!"], first_out) + svc = Service(engine, self.tok, TEMPLATE) + svc.fifo = HandOver(engine) + httpd = serve(svc, port=0) + got = {} + try: + self.base = f"http://127.0.0.1:{httpd.server_address[1]}" + a = threading.Thread(target=lambda: got.update(a=self.post({"logprobs": True, "stop": [".!"]}))) + b = threading.Thread(target=lambda: got.update(b=self.post({"logprobs": True}))) + a.start() + self.assertTrue(engine.first_in.wait(10)) + b.start() + a.join(30) + first_out.set() + b.join(30) + finally: + first_out.set() + httpd.shutdown() + httpd.server_close() + for key, text in (("a", "Café costs €5."), ("b", "Bonjour!")): + choice = got[key]["choices"][0] + self.assertEqual(choice["message"]["content"], text) + self.assertEqual("".join(e["token"] for e in choice["logprobs"]["content"]), text, key) + + def test_logprobs_need_a_bounded_reply(self): + for max_tokens in (0, -1, 5000): # 0 / -1: "the rest of the context" + with self.assertRaises(urllib.error.HTTPError) as cm: + self.post({"logprobs": True, "max_tokens": max_tokens}) + self.assertEqual(cm.exception.code, 400) + self.assertIn("logprobs", self.post({"logprobs": True, "max_tokens": 4096})["choices"][0]) + self.assertNotIn("logprobs", self.post({"max_tokens": 0})["choices"][0]) # not asked: as before + def test_streamed(self): chunks = self.post({"logprobs": True}, stream=True) content = "".join(c["choices"][0]["delta"].get("content") or "" for c in chunks if c.get("choices")) From 900a5b82867475899525ca3dff6b653bc00cf9b4 Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Thu, 22 Oct 2026 23:22:34 +0800 Subject: [PATCH 03/14] serve: allowed_token_ids keeps the token when no named id survives the filter Follows the logprobs-lines change; the one defect left open there. `_renormalise` returned None when the filter left no candidate. The entry was appended all the same, so the next `spell()` failed on `e["token"]`: an unstreamed reply came back as 400 "the request is malformed (TypeError: 'NoneType' object is not subscriptable)", a streamed one stopped after the role chunk. - the chosen token is put back among the candidates when the caller named it. The line carries its logprob on the head and never repeats it among the named ids, so `allowed_token_ids` naming only the token the model chose used to leave nothing to normalise over -- on a real engine too, not only under the mock. It is then the whole distribution, at 1 (logprob 0), which is what vLLM's mask gives. - an entry is never None and is never dropped: the reply's tokens have to keep spelling its text. A line that carries none of the named ids keeps the chosen token with its own logprob and an empty candidate list. - the entry's `logprob` is the chosen token's own value. It used to be copied from the first named alternative, which is another token's number. - `top_logprobs` is sorted most likely first, as the unfiltered top-k already is: with the chosen token put back, insertion order no longer says anything. The class docstring said "vLLM's semantics". vLLM masks generation with these ids (logits.masked_fill_(mask, -inf) before sampling) and reports the unmasked distribution; here they filter what is reported, so the chosen token can be outside them. The docstring says that now. Tests (serve/test_logprobs.py, no GPU): a line whose named ids are all filtered out, and one naming only the chosen token, at the LogprobEntries level; the same two over HTTP; a choice of ten ids with the reply "11", so only the ended reply settles which one; and a held tail of two tokens ("5." as the start of "5.!"). 26 tests in the file. Not run: a real engine, and the parallel path (its batched generate never sets lp_lines: a reading of the code, unchanged here). --- serve/server.py | 45 ++++++++++++----- serve/test_logprobs.py | 107 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 137 insertions(+), 15 deletions(-) diff --git a/serve/server.py b/serve/server.py index 5a747ad6a42..19baff3f392 100644 --- a/serve/server.py +++ b/serve/server.py @@ -5338,9 +5338,11 @@ class LogprobEntries: text of the reply are left out, because spell() gives each piece of text the tokens that spell it. - `allowed_token_ids` (vLLM's semantics): only those ids are kept, and the logprobs are - renormalised over them, so a caller asking about specific candidates gets a distribution - over exactly those instead of a top-K list it has to hope contains them. + `allowed_token_ids`: only those ids are kept, and the logprobs are renormalised over them, + so a caller asking about specific candidates gets a distribution over exactly those instead + of a top-K list it has to hope contains them. vLLM, which named the field, masks generation + with these ids instead (every other logit -inf before sampling) and reports the unmasked + distribution; here they filter what is reported, so the chosen token can be outside them. """ def __init__(self, svc, allowed=(), owner=None): @@ -5408,20 +5410,39 @@ def take(self, eng) -> list: e["top_logprobs"].append({"token": alt.decode("utf-8", "replace"), "logprob": v, "bytes": list(alt)}) if self.allowed: - e = self._renormalise(e) + e = self._renormalise(e, t, lp) out.append(e) return out - def _renormalise(self, e): - """vLLM: keep only the named ids and normalise over them, so they sum to 1.""" - if not e["top_logprobs"]: - return None - mx = max(x["logprob"] for x in e["top_logprobs"]) - se = sum(math.exp(x["logprob"] - mx) for x in e["top_logprobs"]) + def _renormalise(self, e, t, lp): + """Keep only the named ids and normalise over them, so they sum to 1. + + The line carries the chosen token's own logprob on its head and the named ids after it, and never + repeats the chosen one among them. So the chosen token is put back when the caller named it: with + `allowed_token_ids` naming only the token the model chose, that token is then the whole distribution, + at 1 (logprob 0), which is what vLLM's mask gives -- here it used to leave nothing to normalise over, + which returned None, the entry was appended all the same, and the reply failed on e["token"]. A line + that carries none of the named ids has nothing to normalise over either: the entry keeps the chosen + token and its own logprob, with no alternatives. It is never None and never dropped, because the + reply's tokens have to keep spelling its text. + """ + cand = e["top_logprobs"] + own = None + if t in self.allowed: # named, and not repeated on the line: it is a candidate + b = self.raw(t) + own = {"token": b.decode("utf-8", "replace"), "logprob": lp, "bytes": list(b)} + cand = [own] + cand + if not cand: + return e # no named id observed for this token: its own logprob, no top + mx = max(x["logprob"] for x in cand) + se = sum(math.exp(x["logprob"] - mx) for x in cand) lse = mx + math.log(se) - for x in e["top_logprobs"]: + for x in cand: x["logprob"] -= lse - e["logprob"] = e["top_logprobs"][0]["logprob"] if e["top_logprobs"] else e["logprob"] + cand.sort(key=lambda x: -x["logprob"]) # most likely first, as the unfiltered top-k already is + e["top_logprobs"] = cand + if own is not None: # the chosen token is in the distribution: report its own value + e["logprob"] = own["logprob"] # (it used to copy the first alternative's, another token's) return e def prompt_progress(svc: Service): diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index 96aa4022924..b4f329593e7 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -11,6 +11,7 @@ from __future__ import annotations import json +import math import queue import threading import unittest @@ -47,16 +48,29 @@ def __init__(self, tok, script, lines=True): super().__init__(tok, script, max_context=16384) self.with_lines, self.sampling = lines, None + def line(self, i, t): + """The LP line of token t, the i-th of the reply: the same four ids whatever was asked.""" + return (f"LP {-(i + 1) / 100:.6f} 65:-0.500000 72:-1.250000 " + f"88:-3.000000 91:-4.200000\n") + def generate(self, ids, max_new, sampling, cancel, embeddings=None): self.sampling, self.lp_lines = dict(sampling or {}), LPLines(cancel) for i, t in enumerate(super().generate(ids, max_new, sampling, cancel, embeddings)): if self.with_lines: - self.lp_lines.append(( - t, f"LP {-(i + 1) / 100:.6f} 65:-0.500000 72:-1.250000 " - f"88:-3.000000 91:-4.200000\n")) + self.lp_lines.append((t, self.line(i, t))) yield t +class NamedEngine(LpEngine): + """LpEngine that prints the request's own `allowed_token_ids`, as the engine's emit_lp_line does: + `LP : ...`, the chosen id left out (it is on the head).""" + + def line(self, i, t): + named = [int(x) for x in (self.sampling or {}).get("allowed_token_ids") or [] if int(x) != t] + alts = "".join(f" {x}:{-0.700000 - 0.010000 * j:.6f}" for j, x in enumerate(named)) + return f"LP {-(i + 1) / 100:.6f}{alts}\n" + + class NextEngine(LpEngine): """LpEngine for two requests: the second one stops after its first token (its LP line is in lp_lines) until the first request's reply has been read.""" @@ -152,6 +166,25 @@ def test_unreadable_lines_give_no_entry(self): got = LogprobEntries(SimpleNamespace(tok=Tokenizer(), stop_ids=set()), owner=a).take(eng) self.assertEqual([e["logprob"] for e in got], [-0.2]) + def test_a_filtered_line_with_nothing_left_keeps_its_token(self): + # none of the named ids on the line: the entry keeps the chosen token and its own logprob, and no + # alternatives. It used to be None, appended all the same, and the reply failed on e["token"]. + a = threading.Event() + eng = SimpleNamespace(lp_lines=LPLines(a)) + eng.lp_lines.append((72, "LP -0.300000 65:-0.500000")) + got = LogprobEntries(SimpleNamespace(tok=Tokenizer(), stop_ids=set()), [70], a).take(eng) + self.assertEqual([(e["token"], e["logprob"], e["top_logprobs"]) for e in got], [("H", -0.3, [])]) + + def test_a_named_token_the_line_does_not_repeat_is_its_own_distribution(self): + # the engine prints the chosen token's logprob on the head and the named ids after it, never repeating + # the chosen one: naming only that id leaves it alone in the distribution, at 1 (logprob 0) + a = threading.Event() + eng = SimpleNamespace(lp_lines=LPLines(a)) + eng.lp_lines.append((72, "LP -0.300000 65:-0.500000")) + got = LogprobEntries(SimpleNamespace(tok=Tokenizer(), stop_ids=set()), [72], a).take(eng) + self.assertEqual(got[0]["top_logprobs"], [{"token": "H", "logprob": 0.0, "bytes": [72]}]) + self.assertEqual(got[0]["logprob"], 0.0) + class Generate(unittest.TestCase): """The token's LP line is read next to it, and only when it was asked for.""" @@ -328,6 +361,74 @@ def test_allowed_token_ids_are_renormalised_over_the_named_ids(self): total = sum(math.exp(t["logprob"]) for t in lps[0]["top_logprobs"]) self.assertAlmostEqual(total, 1.0, places=6) # renormalised: they sum to 1 + def test_allowed_token_ids_with_nothing_to_renormalise(self): + # the mock's alternatives are 65/72/88/91: naming another id leaves every line with no survivor. Each + # token keeps its own logprob and an empty candidate list, and the reply still comes back whole (this + # used to be a 400 on the unstreamed path and a reply cut after the role chunk when streamed). + for body in ({"logprobs": True, "allowed_token_ids": [66]}, + {"logprobs": True, "allowed_token_ids": [66, 70, 71]}): + choice = self.post(body)["choices"][0] + lps = choice["logprobs"]["content"] + self.assertEqual(choice["message"]["content"], "Café costs €5.") + self.assertEqual("".join(e["token"] for e in lps), "Café costs €5.") + self.assertEqual([e["top_logprobs"] for e in lps], [[]] * len(lps)) + self.assertEqual(lps[0]["logprob"], -0.01) # the token's own, as the engine reported it + streamed = self.post(body, stream=True) + toks = [e["token"] for c in streamed if c.get("choices") + for e in (c["choices"][0].get("logprobs") or {}).get("content", [])] + self.assertEqual("".join(toks), "Café costs €5.") + + def test_a_named_token_the_model_chose_is_in_its_distribution(self): + # 46 is the reply's final ".": the engine does not repeat the chosen token among the named ids, so + # naming only it used to leave nothing to normalise over. It is the whole distribution now, at 0. + lps = self.post({"logprobs": True, "allowed_token_ids": [46]})["choices"][0]["logprobs"]["content"] + dot = [e for e in lps if e["token"] == "."] + self.assertEqual(len(dot), 1) + self.assertEqual(dot[0]["top_logprobs"], [{"token": ".", "logprob": 0.0, "bytes": [46]}]) + self.assertEqual(dot[0]["logprob"], 0.0) + self.assertEqual([e["top_logprobs"] for e in lps if e["token"] != "."], [[]] * (len(lps) - 1)) + + def test_a_choice_of_ten_that_only_the_ended_reply_settles(self): + # ten named ids, and a reply that is two of them: "1" alone is already one of the choices, so which + # one was picked is only known when the reply ends. Here "11": each token's entries cover the ten + # named ids (the chosen one included, because the engine never repeats it), and they sum to 1. + self.engine = self.svc.engine = NamedEngine(self.tok, "11") + digits = [ord(c) for c in "0123456789"] + for stream in (False, True): + out = self.post({"logprobs": True, "top_logprobs": 10, "allowed_token_ids": digits}, stream) + chunks = out if stream else [{"choices": [{"delta": out["choices"][0]["message"], + "logprobs": out["choices"][0]["logprobs"]}]}] + content = "".join(c["choices"][0]["delta"].get("content") or "" for c in chunks if c.get("choices")) + lps = [e for c in chunks if c.get("choices") + for e in (c["choices"][0].get("logprobs") or {}).get("content", [])] + self.assertEqual(content, "11") + self.assertEqual([e["token"] for e in lps], ["1", "1"]) # the first token alone is a choice too + self.assertEqual("".join(e["token"] for e in lps), content) + for e in lps: + ids = [t["token"] for t in e["top_logprobs"]] + self.assertEqual(sorted(ids), sorted("0123456789")) # the ten named ids, and only them + self.assertIn(e["token"], ids) # the chosen one is one of them + self.assertEqual(ids[0], e["token"]) # ... and the most likely, here + self.assertEqual([t["logprob"] for t in e["top_logprobs"]], + sorted((t["logprob"] for t in e["top_logprobs"]), reverse=True)) + self.assertAlmostEqual(sum(math.exp(t["logprob"]) for t in e["top_logprobs"]), 1.0, places=6) + + def test_a_held_tail_of_more_than_one_token(self): + # "5." is the start of the stop string, so the parser holds two tokens and sends them as one piece + # after the request has left the queue: both entries come from this request's own lines all the same + chunks = self.post({"logprobs": True, "stop": ["5.!"]}, stream=True) + content = [c for c in chunks if c.get("choices") and c["choices"][0]["delta"].get("content")] + self.assertEqual(content[-1]["choices"][0]["delta"]["content"], "5.") # held, then flushed whole + self.assertEqual([e["token"] for e in content[-1]["choices"][0]["logprobs"]["content"]], ["5", "."]) + text = "".join(c["choices"][0]["delta"].get("content") or "" for c in chunks if c.get("choices")) + self.assertEqual(text, "Café costs €5.") + self.assertEqual("".join(e["token"] for c in chunks if c.get("choices") + for e in (c["choices"][0].get("logprobs") or {}).get("content", [])), text) + choice = self.post({"logprobs": True, "stop": ["5.!"]})["choices"][0] # and not streamed + self.assertEqual(choice["message"]["content"], "Café costs €5.") + self.assertEqual("".join(e["token"] for e in choice["logprobs"]["content"]), "Café costs €5.") + self.assertEqual("".join(e["token"] for e in choice["logprobs"]["content"][-2:]), "5.") + def test_json_schema_answer_as_written_with_its_logprobs(self): # MockEngine appends the stop token to the script it is built with, and reports # finish_reason=length when it runs out -- so the engine is rebuilt with this answer From 597571c8f140def427316301656ad173387783fb Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Fri, 23 Oct 2026 02:36:46 +0800 Subject: [PATCH 04/14] docs: logprobs on the chat API: the keys, the cap, and what allowed_token_ids does The CUDA/HIP engine has had logprobs since #3200 and DETAILS.md said nothing about them, and the reply cap the logprobs path now needs (max_tokens 1..4096, a 400 without it) is something a client can hit without warning. One bullet in "Using it", the same voice as its neighbours: the keys and where the entries land, that the tokens spell the reply, the cap, what allowed_token_ids keeps and renormalises, and that a token whose probability could not be read gets no entry. --- docs/DETAILS.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 07c9b3e68a2..d8658125481 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -903,6 +903,16 @@ print(r.choices[0].message.content) `max_tokens` (the rest of the context) get a budget of three quarters of it. `"answer_room": false` in `strata-.json` turns it off. Qwen recommends up to 262,144 tokens of reasoning and 131,072 of answer: a client should not cap the reply far below what the context allows. +- **logprobs: how likely each token was (CUDA/HIP).** `"logprobs": true` with `"top_logprobs": 0..20` on + `/v1/chat/completions`, streamed or not, adds `choices[0].logprobs.content[]` to the reply (`token`, `logprob`, + `bytes`, `top_logprobs`); with thinking on, the thinking's entries come under `reasoning_content` instead. The + tokens, joined, spell the reply exactly (a character split across two tokens lands in the one that finishes it). + A reply that asks for logprobs needs `max_tokens` or `max_completion_tokens` of 1 to 4,096 and gets a 400 + without it, because a reply that is not streamed holds its entries until it ends. `"allowed_token_ids": [ids]` + keeps only those ids in `top_logprobs`, most likely first, renormalised to sum to 1; the chosen token is one of + them when it was named (vLLM's field, but vLLM masks generation with it, while here it filters what is reported, + so a token outside the named ids keeps its own `logprob` and an empty list). A token whose probability could not + be read gets no entry, and the reply goes on. - **A reply that ends inside its thinking (#1053, opt-in).** Some turns write a sentence of reasoning and then the end-of-turn token with no ``: the content is empty and an agent stops. `"reasoning_close_retry": true` in `strata-.json` closes the thinking once (as the thinking budget does) and continues, once per request, only for From 22fe5184dc45a12e61fd0cc812ba8fc89a866651 Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Fri, 23 Oct 2026 05:47:40 +0800 Subject: [PATCH 05/14] engine: allowed_token_ids: keep the list the request named `GEN ... logprobs=K allowed=a,b,c` parsed the ids into `req_allowed` and then cleared the vector three lines later, next to `req_imgs.clear()`, before either decode loop ran. So `emit_lp_line` always took the top-k branch and `allowed_token_ids` only ever filtered what was already in top-k. On Swift 1.5 IQ3_XXS, `allowed_token_ids: [16]` ("1") returned the chosen id twice, at 0.5 each, and ten named digits (`0`-`9`, ids 15-24) for the reply `11` covered 2 and then 8 of the ten in the two entries (2 and 5 with `top_logprobs: 5`): ids outside top-k were simply lost. Remove the clear. Measured after it, same model, same requests: `[16]` gives one entry per token at logprob 0, and the ten digits give all ten candidates in every entry, sorted, summing to 1. Not run: HIP, SYCL and Metal builds (this machine has one CUDA card). The solo and speculative decode loops are the only `emit_lp_line` call sites; the batch loop is unchanged. --- src/program/generate.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/program/generate.cpp b/src/program/generate.cpp index 01ed09de45b..73ac43ff5d7 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -12091,7 +12091,6 @@ int main(int argc, char** argv) { } const int64_t n = (int64_t) ids.size(); req_imgs.clear(); - req_allowed.clear(); if (geni && !o.vision) { std::printf("ERR this engine was started without --vision\n"); continue; } if (geni || !mrope_identity) { // positions for every cell this request can reach; the identity again for a text request From b55e9fdc43381d8d4ea7353e4fb3ca5b0723187d Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Fri, 23 Oct 2026 05:47:40 +0800 Subject: [PATCH 06/14] serve: logprobs are refused on the --parallel / --batch path The engine's batch loop prints `BT ` and never calls `emit_lp_line` (the only two call sites are the solo and speculative decode loops), so a server started with `--parallel` could only answer a logprobs request with the field missing -- which a client cannot tell from a reply that had no tokens. Refuse it with the same 400 the cap check raises, before anything is generated; a request that does not ask is untouched. Measured with `--parallel 2` on Swift 1.5 IQ3_XXS: logprobs, streamed and not, both 400; a plain request still 200 with the same content. Test: test_logprobs_are_refused_on_the_batch_path (red when the check is disabled). The docs bullet and the class docstring say it. serve: 908 tests, 0 errors, 20 skipped (`python -m unittest discover -s serve -t . -p "test_*.py"`). --- docs/DETAILS.md | 3 ++- serve/server.py | 10 ++++++++++ serve/test_logprobs.py | 12 ++++++++++++ 3 files changed, 24 insertions(+), 1 deletion(-) diff --git a/docs/DETAILS.md b/docs/DETAILS.md index d8658125481..62d83276b33 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -912,7 +912,8 @@ print(r.choices[0].message.content) keeps only those ids in `top_logprobs`, most likely first, renormalised to sum to 1; the chosen token is one of them when it was named (vLLM's field, but vLLM masks generation with it, while here it filters what is reported, so a token outside the named ids keeps its own `logprob` and an empty list). A token whose probability could not - be read gets no entry, and the reply goes on. + be read gets no entry, and the reply goes on. On a server started with `--parallel`, a request that asks for + logprobs gets a 400: the engine's batched loop writes no LP lines, so there is nothing to report. - **A reply that ends inside its thinking (#1053, opt-in).** Some turns write a sentence of reasoning and then the end-of-turn token with no ``: the content is empty and an agent stops. `"reasoning_close_retry": true` in `strata-.json` closes the thinking once (as the thinking budget does) and continues, once per request, only for diff --git a/serve/server.py b/serve/server.py index 19baff3f392..3564cf4d55f 100644 --- a/serve/server.py +++ b/serve/server.py @@ -5343,6 +5343,9 @@ class LogprobEntries: of a top-K list it has to hope contains them. vLLM, which named the field, masks generation with these ids instead (every other logit -inf before sampling) and reports the unmasked distribution; here they filter what is reported, so the chosen token can be outside them. + + Not on `--parallel`: the engine's batched loop prints `BT ` and never an LP line, + so a request that asks for logprobs on such a server is refused (400) before it starts. """ def __init__(self, svc, allowed=(), owner=None): @@ -6861,6 +6864,13 @@ def _openai(self, req): max_req = max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 0) # 0/-1: the rest if top_logprobs(req) is not None and not 0 < max_req <= LOGPROBS_MAX_TOKENS: raise ValueError(f"logprobs: max_tokens (or max_completion_tokens) 1..{LOGPROBS_MAX_TOKENS} is required") + if top_logprobs(req) is not None and getattr(svc.engine, "batch", 0): + # --parallel/--batch: the batched loop prints "BT " and never calls emit_lp_line, so no + # LP line ever reaches the server there (src/program/generate.cpp). Refused rather than answered + # with the field missing, which a client cannot tell from a reply that had no tokens. + raise ValueError("logprobs is not available when the engine runs requests in batch slots " + "(--parallel/--batch): it writes no LP lines there; start the server without " + "--parallel") bias_api.validate_request(req, svc.engine, len(getattr(svc.tok, "tokens", ())) or None) use_mcp = req.get("strata_mcp") is True and svc.mcp is not None # the web app's opt-in (serve/mcp.py) own = {t.get("name") for t in tools or [] if isinstance(t, dict)} # #592: a second line of defence diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index b4f329593e7..811b07f5fe8 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -309,6 +309,18 @@ def test_logprobs_need_a_bounded_reply(self): self.assertIn("logprobs", self.post({"logprobs": True, "max_tokens": 4096})["choices"][0]) self.assertNotIn("logprobs", self.post({"max_tokens": 0})["choices"][0]) # not asked: as before + def test_logprobs_are_refused_on_the_batch_path(self): + # --parallel/--batch: the engine's batched loop prints "BT " and never calls emit_lp_line + # (src/program/generate.cpp), so no LP line reaches the server and the field would simply be absent. + # Refused up front instead, before anything is generated; a request that does not ask is untouched. + self.engine.batch = 2 + with self.assertRaises(urllib.error.HTTPError) as cm: + self.post({"logprobs": True}) + self.assertEqual(cm.exception.code, 400) + self.assertIn("--parallel", cm.exception.read().decode()) + self.engine.batch = 0 + self.assertIn("logprobs", self.post({"logprobs": True})["choices"][0]) + def test_streamed(self): chunks = self.post({"logprobs": True}, stream=True) content = "".join(c["choices"][0]["delta"].get("content") or "" for c in chunks if c.get("choices")) From 26a2d4803a9db92a6c58de092fac8b774c515fd2 Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Fri, 23 Oct 2026 06:26:14 +0800 Subject: [PATCH 07/14] serve: an allowed_token_ids id past the vocabulary is a 400 emit_lp_line indexes its logits row with each named id as given (src/program/generate.cpp), so an allowed_token_ids entry >= the model's vocabulary reads out of bounds. The server checks the ids against the tokenizer's vocabulary before anything is sent and raises the same 400 as the other logprobs checks; a request that asks for no logprobs, or names ids inside the vocabulary, is unchanged. Test: Api.test_allowed_token_ids_outside_the_vocabulary_is_a_400 (vocabulary 300: id 300 -> 400, id 299 -> 200). This one comes from reading the code -- no crash was seen -- and the docs bullet says it. serve: 909 tests, 0 errors, 20 skipped (python -m unittest discover -s serve -t . -p "test_*.py"). --- docs/DETAILS.md | 7 ++++--- serve/server.py | 8 ++++++++ serve/test_logprobs.py | 11 +++++++++++ 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 62d83276b33..19e27194bfa 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -911,9 +911,10 @@ print(r.choices[0].message.content) without it, because a reply that is not streamed holds its entries until it ends. `"allowed_token_ids": [ids]` keeps only those ids in `top_logprobs`, most likely first, renormalised to sum to 1; the chosen token is one of them when it was named (vLLM's field, but vLLM masks generation with it, while here it filters what is reported, - so a token outside the named ids keeps its own `logprob` and an empty list). A token whose probability could not - be read gets no entry, and the reply goes on. On a server started with `--parallel`, a request that asks for - logprobs gets a 400: the engine's batched loop writes no LP lines, so there is nothing to report. + so a token outside the named ids keeps its own `logprob` and an empty list). An id past the model's vocabulary is + a 400. A token whose probability could not be read gets no entry, and the reply goes on. On a server started with + `--parallel`, a request that asks for logprobs gets a 400: the engine's batched loop writes no LP lines, so there is + nothing to report. - **A reply that ends inside its thinking (#1053, opt-in).** Some turns write a sentence of reasoning and then the end-of-turn token with no ``: the content is empty and an agent stops. `"reasoning_close_retry": true` in `strata-.json` closes the thinking once (as the thinking budget does) and continues, once per request, only for diff --git a/serve/server.py b/serve/server.py index 3564cf4d55f..15affdfa958 100644 --- a/serve/server.py +++ b/serve/server.py @@ -6871,6 +6871,14 @@ def _openai(self, req): raise ValueError("logprobs is not available when the engine runs requests in batch slots " "(--parallel/--batch): it writes no LP lines there; start the server without " "--parallel") + n_vocab = len(getattr(svc.tok, "tokens", ())) or None + if top_logprobs(req) is not None and n_vocab: + # emit_lp_line indexes the logits row with each named id as given, so an id past the vocabulary + # would read out of bounds: refused here, before anything is sent + bad = [i for i in allowed_ids(req) if i >= n_vocab] + if bad: + raise ValueError(f"allowed_token_ids: {bad[0]} is outside the vocabulary " + f"(0..{n_vocab - 1})") bias_api.validate_request(req, svc.engine, len(getattr(svc.tok, "tokens", ())) or None) use_mcp = req.get("strata_mcp") is True and svc.mcp is not None # the web app's opt-in (serve/mcp.py) own = {t.get("name") for t in tools or [] if isinstance(t, dict)} # #592: a second line of defence diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index 811b07f5fe8..6d008721d0c 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -321,6 +321,17 @@ def test_logprobs_are_refused_on_the_batch_path(self): self.engine.batch = 0 self.assertIn("logprobs", self.post({"logprobs": True})["choices"][0]) + def test_allowed_token_ids_outside_the_vocabulary_is_a_400(self): + # emit_lp_line indexes its logits row with each named id, so an id past the vocabulary would read out of + # bounds; the request is refused before anything is generated + self.tok.tokens = [str(i) for i in range(300)] # a tokenizer that knows its vocabulary + with self.assertRaises(urllib.error.HTTPError) as cm: + self.post({"logprobs": True, "top_logprobs": 1, "allowed_token_ids": [300]}) + self.assertEqual(cm.exception.code, 400) + self.assertIn("allowed_token_ids", cm.exception.read().decode()) + self.assertIn("logprobs", self.post({"logprobs": True, "top_logprobs": 1, + "allowed_token_ids": [299]})["choices"][0]) + def test_streamed(self): chunks = self.post({"logprobs": True}, stream=True) content = "".join(c["choices"][0]["delta"].get("content") or "" for c in chunks if c.get("choices")) From eca6513b01188e8dea88901c73c63449656d22c6 Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Sun, 25 Oct 2026 20:17:25 +0800 Subject: [PATCH 08/14] serve: a streamed logprobs reply takes any max_tokens The 1..4096 bound on max_tokens / max_completion_tokens was applied to every request that asked for logprobs. Its reason only holds for a reply that keeps its entries until it ends: a reply that is not streamed (openai_collect holds the whole list to build the body) or that answers a json_schema (structured_chunks buffers it to validate). A streamed reply hands each entry out with its chunk and keeps none, so it is bounded by its own length. So a streamed reply is bounded by the model's context, as on OpenAI, whose request schema has max_tokens optional and no link between it and logprobs (vLLM has none either: max_tokens defaults to 16 and logprobs only checks itself). Two behaviours change, both of them requests that could not have worked before: - `"logprobs": true, "stream": true` without max_tokens is answered instead of a 400. The OpenAI SDKs send no max_tokens by default. - `"logprobs": true, "stream": true, "max_tokens": 5000` is answered instead of a 400. A reply that is not streamed, or that answers a json_schema, still needs max_tokens of 1..4096 and gets a 400 without it; the message now says why and which case it is. Tests: serve.test_logprobs 30/30 (28 before; the two new ones fail with a 400 on the previous commit). test_a_streamed_reply_needs_no_bound covers 5000, 0 (the rest of the context) and no max_tokens at all, and checks that the entries still spell the reply; test_a_streamed_json_schema_reply_is_still_bounded keeps the cap on that path. --- docs/DETAILS.md | 6 ++++-- serve/server.py | 13 +++++++++---- serve/test_logprobs.py | 21 +++++++++++++++++++++ 3 files changed, 34 insertions(+), 6 deletions(-) diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 19e27194bfa..f8da0d4ac63 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -907,8 +907,10 @@ print(r.choices[0].message.content) `/v1/chat/completions`, streamed or not, adds `choices[0].logprobs.content[]` to the reply (`token`, `logprob`, `bytes`, `top_logprobs`); with thinking on, the thinking's entries come under `reasoning_content` instead. The tokens, joined, spell the reply exactly (a character split across two tokens lands in the one that finishes it). - A reply that asks for logprobs needs `max_tokens` or `max_completion_tokens` of 1 to 4,096 and gets a 400 - without it, because a reply that is not streamed holds its entries until it ends. `"allowed_token_ids": [ids]` + A reply that asks for logprobs and is not streamed -- or that answers a `json_schema`, which is buffered to be + validated -- needs `max_tokens` or `max_completion_tokens` of 1 to 4,096 and gets a 400 without it, because such + a reply holds every entry until it ends. A streamed one hands each entry out with its chunk and keeps none, so it + is bounded only by the context, as on OpenAI. `"allowed_token_ids": [ids]` keeps only those ids in `top_logprobs`, most likely first, renormalised to sum to 1; the chosen token is one of them when it was named (vLLM's field, but vLLM masks generation with it, while here it filters what is reported, so a token outside the named ids keeps its own `logprob` and an empty list). An id past the model's vocabulary is diff --git a/serve/server.py b/serve/server.py index 15affdfa958..2d310ae22b6 100644 --- a/serve/server.py +++ b/serve/server.py @@ -3301,8 +3301,10 @@ def allowed_ids(req: dict) -> list: NO_LINE = object() # StrataEngine.generate: no line held after a token # A reply with logprobs keeps every entry (its token, bytes and up to 20 alternatives) until it ends when it is not -# streamed or answers a json_schema, and spell() matches text against the entries still pending: a bounded reply -# keeps that memory and work bounded. +# streamed or answers a json_schema (structured_chunks buffers it to validate), and spell() matches text against the +# entries still pending: a bounded reply keeps that memory and work bounded. A streamed reply hands each entry out +# with its chunk and drops it, so it is bounded by the reply's own length and needs no such cap -- the same reason +# OpenAI itself only bounds a reply by the model's context. LOGPROBS_MAX_TOKENS = 4096 @@ -6862,8 +6864,11 @@ def _openai(self, req): raise ValueError("structured response_format with tools/MCP is not supported") svc.load() max_req = max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 0) # 0/-1: the rest - if top_logprobs(req) is not None and not 0 < max_req <= LOGPROBS_MAX_TOKENS: - raise ValueError(f"logprobs: max_tokens (or max_completion_tokens) 1..{LOGPROBS_MAX_TOKENS} is required") + kept = not req.get("stream") or validator is not None # its entries are held until the reply ends + if top_logprobs(req) is not None and kept and not 0 < max_req <= LOGPROBS_MAX_TOKENS: + raise ValueError(f"logprobs: max_tokens (or max_completion_tokens) 1..{LOGPROBS_MAX_TOKENS} is " + f"required when the reply holds its entries until it ends (not streamed, or a " + f"json_schema); a streamed reply takes any length") if top_logprobs(req) is not None and getattr(svc.engine, "batch", 0): # --parallel/--batch: the batched loop prints "BT " and never calls emit_lp_line, so no # LP line ever reaches the server there (src/program/generate.cpp). Refused rather than answered diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index 6d008721d0c..65937900605 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -302,6 +302,7 @@ def test_a_request_reads_only_its_own_lines(self): self.assertEqual("".join(e["token"] for e in choice["logprobs"]["content"]), text, key) def test_logprobs_need_a_bounded_reply(self): + # a reply that is not streamed holds its entries (token, bytes, up to 20 alternatives) until it ends for max_tokens in (0, -1, 5000): # 0 / -1: "the rest of the context" with self.assertRaises(urllib.error.HTTPError) as cm: self.post({"logprobs": True, "max_tokens": max_tokens}) @@ -309,6 +310,26 @@ def test_logprobs_need_a_bounded_reply(self): self.assertIn("logprobs", self.post({"logprobs": True, "max_tokens": 4096})["choices"][0]) self.assertNotIn("logprobs", self.post({"max_tokens": 0})["choices"][0]) # not asked: as before + def test_a_streamed_reply_needs_no_bound(self): + # A streamed reply hands each entry out with its chunk and keeps none, so the cap above has no reason to + # apply: OpenAI itself bounds a reply only by the model's context, and its SDKs send no max_tokens by + # default. 5000 and "the rest of the context" must both be answered, as they are without logprobs. + for max_tokens in (5000, 0, None): # None: the field is not sent at all + chunks = self.post({"logprobs": True, "max_tokens": max_tokens}, stream=True) + text = "".join(c["choices"][0]["delta"].get("content") or "" for c in chunks if c.get("choices")) + toks = [e["token"] for c in chunks if c.get("choices") + for e in (c["choices"][0].get("logprobs") or {}).get("content", [])] + self.assertEqual(text, "Café costs €5.", max_tokens) + self.assertEqual("".join(toks), text, max_tokens) # an entry for every token, streamed or not + + def test_a_streamed_json_schema_reply_is_still_bounded(self): + # structured_chunks buffers the whole reply to validate it, logprobs included: that reply is not + # streamed in the sense above, so the cap stays + for max_tokens in (0, 5000): + with self.assertRaises(urllib.error.HTTPError) as cm: + self.post({"logprobs": True, "max_tokens": max_tokens, "response_format": SCHEMA}, stream=True) + self.assertEqual(cm.exception.code, 400) + def test_logprobs_are_refused_on_the_batch_path(self): # --parallel/--batch: the engine's batched loop prints "BT " and never calls emit_lp_line # (src/program/generate.cpp), so no LP line reaches the server and the field would simply be absent. From 4b7339e16743745002f946fe3b5e71246ce9804a Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Sun, 25 Oct 2026 21:46:45 +0800 Subject: [PATCH 09/14] serve: a logprobs reply is bounded by max_tokens, not by a cap of 4096 The 1..4096 bound on max_tokens / max_completion_tokens was ours alone: OpenAI's request schema makes max_tokens optional and ties no limit to logprobs, and vLLM's defaults to 16 or None and its logprobs only check themselves. A reply is already bounded by max_tokens -- or by the rest of the context when it is not sent -- so the constant was a second bound on top of one that exists, and it answered 400 to requests that would have worked. What it was protecting is real, but it is the reply's own length: a reply that is not streamed, or that answers a json_schema (buffered to be validated), holds every entry (its token, bytes and up to 20 alternatives) until it ends, while a streamed reply hands each entry out with its chunk and keeps none. Both are now answered at any length; a client that wants that memory bounded streams. Three behaviours change, all of them a 400 before: - not streamed, no max_tokens (what the OpenAI SDKs send by default) is answered - not streamed with max_tokens 0 / -1 ("the rest of the context") is answered - not streamed, or a json_schema, with max_tokens above 4096 is answered Tests: serve.test_logprobs 30/30. test_logprobs_take_any_max_tokens (200 / 5000 / 0 / -1 / not sent, not streamed) and test_a_json_schema_reply_takes_any_max_tokens (200 / 0 / 5000, streamed and not) replace the two that asserted the 400. The whole serve suite: 1064 tests, 1 error, 21 skipped -- the error is serve.test_winjob.test_children_end_with_a_killed_server, which fails the same way on a clean upstream/main checkout here and touches no logprobs code. Hardware, Swift-1.5 IQ3_XXS with --spec 4 --mtp, 22 cases: max_tokens absent / 0 / 200 / 4096 / 5000, streamed and not, each 200 with every token carrying an entry; json_schema at absent / 0 / 200 / 5000, streamed and not, each 200; allowed_token_ids over ten digits and [16] unchanged; an id past the vocabulary still a 400. --- docs/DETAILS.md | 8 ++++---- serve/server.py | 16 ++++------------ serve/test_logprobs.py | 40 ++++++++++++++++++++++++---------------- 3 files changed, 32 insertions(+), 32 deletions(-) diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 3839dbd766f..2b170c149e0 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -509,10 +509,10 @@ applies at the next start. Network, key, MCP and program keys are not editable t `/v1/chat/completions`, streamed or not, adds `choices[0].logprobs.content[]` to the reply (`token`, `logprob`, `bytes`, `top_logprobs`); with thinking on, the thinking's entries come under `reasoning_content` instead. The tokens, joined, spell the reply exactly (a character split across two tokens lands in the one that finishes it). - A reply that asks for logprobs and is not streamed -- or that answers a `json_schema`, which is buffered to be - validated -- needs `max_tokens` or `max_completion_tokens` of 1 to 4,096 and gets a 400 without it, because such - a reply holds every entry until it ends. A streamed one hands each entry out with its chunk and keeps none, so it - is bounded only by the context, as on OpenAI. `"allowed_token_ids": [ids]` + A reply is bounded by `max_tokens` -- or by the rest of the context when it is not sent -- and by no cap of our + own, as on OpenAI and vLLM, neither of which ties a length limit to logprobs. A reply that is not streamed, or + that answers a `json_schema` (buffered to be validated), does hold every entry until it ends, so a long one holds + a lot of them; a streamed reply hands each entry out with its chunk and keeps none. `"allowed_token_ids": [ids]` keeps only those ids in `top_logprobs`, most likely first, renormalised to sum to 1; the chosen token is one of them when it was named (vLLM's field, but vLLM masks generation with it, while here it filters what is reported, so a token outside the named ids keeps its own `logprob` and an empty list). An id past the model's vocabulary is diff --git a/serve/server.py b/serve/server.py index d483472a0bb..b5796c1203b 100644 --- a/serve/server.py +++ b/serve/server.py @@ -3777,13 +3777,6 @@ def allowed_ids(req: dict) -> list: NO_LINE = object() # StrataEngine.generate: no line held after a token -# A reply with logprobs keeps every entry (its token, bytes and up to 20 alternatives) until it ends when it is not -# streamed or answers a json_schema (structured_chunks buffers it to validate), and spell() matches text against the -# entries still pending: a bounded reply keeps that memory and work bounded. A streamed reply hands each entry out -# with its chunk and drops it, so it is bounded by the reply's own length and needs no such cap -- the same reason -# OpenAI itself only bounds a reply by the model's context. -LOGPROBS_MAX_TOKENS = 4096 - def stop_strings(req: dict) -> list[str]: """The request's stop strings: OpenAI's `stop` (a string or a list of up to 4) or Anthropic's `stop_sequences` @@ -7432,11 +7425,10 @@ def _openai(self, req): raise ValueError("structured response_format with tools/MCP is not supported") svc.load() max_req = max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 0) # 0/-1: the rest - kept = not req.get("stream") or validator is not None # its entries are held until the reply ends - if top_logprobs(req) is not None and kept and not 0 < max_req <= LOGPROBS_MAX_TOKENS: - raise ValueError(f"logprobs: max_tokens (or max_completion_tokens) 1..{LOGPROBS_MAX_TOKENS} is " - f"required when the reply holds its entries until it ends (not streamed, or a " - f"json_schema); a streamed reply takes any length") + # No cap of our own on the reply's length when logprobs are asked for: the reply is already bounded by + # max_tokens (by the rest of the context when it is not sent), which is the only bound OpenAI and vLLM + # put on it either. A reply that is not streamed does hold every entry (its token, bytes and up to 20 + # alternatives) until it ends, so a long one holds a lot: streaming hands each entry out with its chunk. if top_logprobs(req) is not None and getattr(svc.engine, "batch", 0): # --parallel/--batch: the batched loop prints "BT " and never calls emit_lp_line, so no # LP line ever reaches the server there (src/program/generate.cpp). Refused rather than answered diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index 65937900605..ca92bcf0f9a 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -301,18 +301,19 @@ def test_a_request_reads_only_its_own_lines(self): self.assertEqual(choice["message"]["content"], text) self.assertEqual("".join(e["token"] for e in choice["logprobs"]["content"]), text, key) - def test_logprobs_need_a_bounded_reply(self): - # a reply that is not streamed holds its entries (token, bytes, up to 20 alternatives) until it ends - for max_tokens in (0, -1, 5000): # 0 / -1: "the rest of the context" - with self.assertRaises(urllib.error.HTTPError) as cm: - self.post({"logprobs": True, "max_tokens": max_tokens}) - self.assertEqual(cm.exception.code, 400) - self.assertIn("logprobs", self.post({"logprobs": True, "max_tokens": 4096})["choices"][0]) + def test_logprobs_take_any_max_tokens(self): + # No cap of our own: the reply is already bounded by max_tokens (the rest of the context when it is not + # sent), which is the only bound OpenAI and vLLM put on it too -- openai-node's max_tokens is optional + # and vLLM's defaults to 16 or None, neither tying it to logprobs. + for max_tokens in (200, 5000, 0, -1, None): # 0 / -1: "the rest of the context"; None: not sent + choice = self.post({"logprobs": True, "top_logprobs": 1, "max_tokens": max_tokens})["choices"][0] + self.assertEqual("".join(e["token"] for e in choice["logprobs"]["content"]), + "Café costs €5.", max_tokens) self.assertNotIn("logprobs", self.post({"max_tokens": 0})["choices"][0]) # not asked: as before def test_a_streamed_reply_needs_no_bound(self): - # A streamed reply hands each entry out with its chunk and keeps none, so the cap above has no reason to - # apply: OpenAI itself bounds a reply only by the model's context, and its SDKs send no max_tokens by + # A streamed reply hands each entry out with its chunk and keeps none, so it is bounded by the reply's own + # length. OpenAI itself bounds a reply only by the model's context, and its SDKs send no max_tokens by # default. 5000 and "the rest of the context" must both be answered, as they are without logprobs. for max_tokens in (5000, 0, None): # None: the field is not sent at all chunks = self.post({"logprobs": True, "max_tokens": max_tokens}, stream=True) @@ -322,13 +323,20 @@ def test_a_streamed_reply_needs_no_bound(self): self.assertEqual(text, "Café costs €5.", max_tokens) self.assertEqual("".join(toks), text, max_tokens) # an entry for every token, streamed or not - def test_a_streamed_json_schema_reply_is_still_bounded(self): - # structured_chunks buffers the whole reply to validate it, logprobs included: that reply is not - # streamed in the sense above, so the cap stays - for max_tokens in (0, 5000): - with self.assertRaises(urllib.error.HTTPError) as cm: - self.post({"logprobs": True, "max_tokens": max_tokens, "response_format": SCHEMA}, stream=True) - self.assertEqual(cm.exception.code, 400) + def test_a_json_schema_reply_takes_any_max_tokens(self): + # structured_chunks buffers the reply to validate it, logprobs included: it holds the entries until the + # end, as a reply that is not streamed does. That memory is bounded by max_tokens, not by a cap of ours. + written = '```json\n{"rent": "$1,250", "quote": "x"}\n```' + self.engine = self.svc.engine = LpEngine(self.tok, written) + content = '{"rent": "$1,250", "quote": "x"}' + for max_tokens in (200, 0, 5000): + choice = self.post({"logprobs": True, "max_tokens": max_tokens, "response_format": SCHEMA})["choices"][0] + self.assertEqual(choice["message"]["content"], content, max_tokens) + self.assertEqual("".join(e["token"] for e in choice["logprobs"]["content"]), content, max_tokens) + chunks = self.post({"logprobs": True, "max_tokens": max_tokens, "response_format": SCHEMA}, stream=True) + toks = [e["token"] for c in chunks if c.get("choices") + for e in (c["choices"][0].get("logprobs") or {}).get("content", [])] + self.assertEqual("".join(toks), content, (max_tokens, "streamed")) def test_logprobs_are_refused_on_the_batch_path(self): # --parallel/--batch: the engine's batched loop prints "BT " and never calls emit_lp_line From 1878675f00e0ac9bc380ab9ba8f05db700ded926 Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Sun, 25 Oct 2026 22:59:29 +0800 Subject: [PATCH 10/14] serve: logprobs with "parallel" for a max_tokens 1 reply, and on a replica server without slots A batch slot prints "BT " and its decode loop never calls emit_lp_line, so logprobs was refused outright there. An admission does write one: a max_tokens 1 request is asked of the engine as BGEN and read as "GEN 1" (generate.cpp:12090), the single-token path that prints "T " and the LP line after it in one flush. _control() now reads that line (want_lp) and keeps it on the request's own LPLines, held on its cancel Event -- the engine's lp_lines is one attribute and several slots read control lines at once, so it cannot be shared between them. A longer reply is still a 400: past that first token its entries would cover a prefix only, which a client cannot tell from the model's own answer. This is the line gt50 drew on #1916 ("LP lines for the BT slot phase can be a later step"), and the read itself follows the server-only fix karls0r posted there (+57/-8, engine unchanged). The other half is a refusal that should not have been there at all. ReplicaEngine.batch totals how many requests its engines run at once and counts a replica without slots as one (serve/replicas.py), and the check asked it "does this engine use batch slots". A server started with --replicas 2 and no --parallel therefore answered 400 to every logprobs request, while reading every reply by the GEN path, which does write the lines. batch_slots_on() asks each replica's own figure instead. Tests: serve.test_logprobs 32/32 -- the two new ones in Replicas are no slots behind two replicas, and slots on one of them; serve.test_logprobs_parallel 6/6 -- a decision, allowed_token_ids on the admission, four decisions at once each with its own, not asked, a longer reply a 400, and an engine that writes no LP lines at all. Hardware, Swift-1.5 IQ3_XXS, --spec 4 --mtp, --parallel 2: max_tokens 1 + top_logprobs 5 -> 200, one entry spelling the reply (A -0.062427, then B -2.848, " A" -7.415); allowed_token_ids over 11 ids -> 200, the named ids as the candidates, sum(exp) 1.0, a chosen token outside them keeping its own logprob; max_tokens 5, max_tokens absent and max_completion_tokens 0 -> 400 naming max_tokens 1; max_completion_tokens 1 -> 200 with its entry; streamed -> one chunk carrying the entry for "Odd"; four decisions at once on two slots -> 200 each, 1.6 s for all four, every entry matching its reply; a plain request and a 64-token reply unchanged. The whole serve suite: 1073 tests, 1 failure, 1 error, 21 skipped. Both are known on a clean upstream/main here and touch no logprobs code: test_responses.test_server_close_ends_the_sampler (1 != 2, order-sensitive -- it passes when its file runs alone, which it did in the 83-test run just before this one) and test_winjob's running() on a pid this environment cannot read. --- docs/DETAILS.md | 5 +- serve/server.py | 101 +++++++++++++++++++-------- serve/test_logprobs.py | 24 ++++++- serve/test_logprobs_parallel.py | 117 ++++++++++++++++++++++++++++++++ 4 files changed, 213 insertions(+), 34 deletions(-) create mode 100644 serve/test_logprobs_parallel.py diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 2b170c149e0..fb796ec4df8 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -517,8 +517,9 @@ applies at the next start. Network, key, MCP and program keys are not editable t them when it was named (vLLM's field, but vLLM masks generation with it, while here it filters what is reported, so a token outside the named ids keeps its own `logprob` and an empty list). An id past the model's vocabulary is a 400. A token whose probability could not be read gets no entry, and the reply goes on. On a server started with - `--parallel`, a request that asks for logprobs gets a 400: the engine's batched loop writes no LP lines, so there is - nothing to report. + `--parallel`, only a `max_tokens` 1 request gets them: it is the one a batch slot decodes through the single-token + path, which writes an LP line. A longer reply there is a 400, since the slot's `BT` lines carry none -- entries for + a prefix only, which a client cannot tell from the model's own answer. - **Non-finite logits end the request and drop the cached state (on by default).** If a verify window's logits are NaN, the window emits no tokens, the request ends, and the next request drops the live session, checkpoints and parked conversations and reads its prompt from token 0. The server retries a request that had sent nothing once, transparently. diff --git a/serve/server.py b/serve/server.py index dd622e15a7c..658ce827b2f 100644 --- a/serve/server.py +++ b/serve/server.py @@ -1484,35 +1484,56 @@ def _send(self, text: str): except (OSError, AttributeError): # the pipe is gone (or a restart's close() took the process): raise EngineDied(f"the engine stopped unexpectedly (exit code {self.exit_code()})") from None - def _control(self, cancel, on_token, stop_when=None, on_quiet=None): + def _control(self, cancel, on_token, stop_when=None, on_quiet=None, want_lp=False): """Reads the control lines of the request on them (GEN / BGEN), yielding None heartbeats. Calls on_token(id) for each `T`; returns ("done", None) at DONE, or ("badm", continues) at BADM (after DONE). `stop_when()` true sends STOP once (the request is then read to its DONE). `on_quiet()` runs whenever the engine has been - quiet for half a second: between two PP lines of a read whose chunks take longer than that.""" + quiet for half a second: between two PP lines of a read whose chunks take longer than that. + + `want_lp`: a batch admission's token is the only one this reader ever sees with an `LP` line after it (a + slot's `BT` lines carry none), so a request asking for logprobs there is a `max_tokens` 1 one: read the + line here, for the request itself (its own lines: the engine's is shared by the slots' readers).""" stopped = False beat = time.monotonic() + held = NO_LINE # a line read while looking for a token's LP line while True: - try: - # A client can leave before the first PP/token. Poll its Event as often as _take_control does; - # the ten-second HTTP heartbeat is independent of how promptly STOP reaches the engine. - line = self.lines.get(timeout=0.5) - except queue.Empty: - if cancel.is_set() and not stopped: - self._send("STOP") - stopped = True - if on_quiet is not None: - on_quiet() - if time.monotonic() - beat >= 10.0: - beat = time.monotonic() - yield None - continue + if held is not NO_LINE: + line, held = held, NO_LINE + else: + try: + # A client can leave before the first PP/token. Poll its Event as often as _take_control does; + # the ten-second HTTP heartbeat is independent of how promptly STOP reaches the engine. + line = self.lines.get(timeout=0.5) + except queue.Empty: + if cancel.is_set() and not stopped: + self._send("STOP") + stopped = True + if on_quiet is not None: + on_quiet() + if time.monotonic() - beat >= 10.0: + beat = time.monotonic() + yield None + continue beat = time.monotonic() if line is None: raise EngineDied(f"the engine stopped unexpectedly (exit code {self.exit_code()})") if not line.startswith(("PP ", "INFO")): btrace("ctl<", self._ctl_mode, line.strip()[:60]) if line.startswith("T "): - on_token(int(line[2:])) + t = int(line[2:]) + if want_lp: + # the engine writes the token and its LP line in one flush: the next line is that one, or a + # line this reader must still see (held for the next turn, as generate() does) + try: + held = self.lines.get(timeout=0.5) + except queue.Empty: + held = NO_LINE + if isinstance(held, str) and held.startswith("LP "): + own = getattr(cancel, "_strata_lp_lines", None) + if own is not None: + own.append((t, held)) + held = NO_LINE + on_token(t) if not stopped and (cancel.is_set() or (stop_when is not None and stop_when())): self._send("STOP") stopped = True @@ -1789,6 +1810,12 @@ def generate_batched(self, ids, max_new, sampling, cancel, embeddings=None): self.progress, self.progress_ms, self.reused = None, 0, 0 born = self.gen # the engine process this request is sent to keys = self.sampling_keys(sampling or {}) + # logprobs: only an admission's token can have an LP line here (the slot's "BT" lines carry none), so the + # server refuses these requests unless max_tokens is 1. The lines go to the request itself, through its + # cancel Event: several requests read control lines at once, and the engine's lp_lines is one attribute. + want_lp = top_logprobs(sampling or {}) is not None + if want_lp: + cancel._strata_lp_lines = LPLines(cancel) out: list[int] = [] pending: list[int] = [] prompt, left = list(ids), int(max_new) @@ -1843,7 +1870,8 @@ def give_way(): # at a PP line, and between two far apart (the engine takes it self.slot_claims.add(reserved) if reserved is not None: self._send(f"BYIELD {reserved}") - for x in self._control(cancel, pending.append, stop_when=others, on_quiet=give_way): + for x in self._control(cancel, pending.append, stop_when=others, on_quiet=give_way, + want_lp=want_lp): while pending: t = pending.pop(0) out.append(t) @@ -1971,7 +1999,8 @@ def note_progress(): if self.progress and live["state"] == "reading": live["prompt_read"] = self.progress[0] try: - for x in self._control(cancel, pending.append, on_quiet=lambda: (ask(), note_progress())): + for x in self._control(cancel, pending.append, on_quiet=lambda: (ask(), note_progress()), + want_lp=want_lp): while pending: t = pending.pop(0) out.append(t) @@ -3775,6 +3804,15 @@ def allowed_ids(req: dict) -> list: return [x for x in ids if isinstance(x, int) and not isinstance(x, bool) and x >= 0] +def batch_slots_on(engine) -> bool: + """Does this engine run requests in batch slots ("parallel", whose `BT` lines carry no LP lines)? + + Asked of each engine behind a ReplicaEngine: its own `batch` is how many requests it can run at once, so it + counts a replica without slots as one (serve/replicas.py) -- a server with `--replicas 2` and no `--parallel` + has none, and reads every reply by the GEN path, which does write them.""" + return any(int(getattr(e, "batch", 0) or 0) for e in (getattr(engine, "engines", None) or [engine])) + + NO_LINE = object() # StrataEngine.generate: no line held after a token @@ -5900,8 +5938,8 @@ class LogprobEntries: with these ids instead (every other logit -inf before sampling) and reports the unmasked distribution; here they filter what is reported, so the chosen token can be outside them. - Not on `--parallel`: the engine's batched loop prints `BT ` and never an LP line, - so a request that asks for logprobs on such a server is refused (400) before it starts. + On `--parallel`: only a `max_tokens` 1 request gets entries (an admission, decoded through the single-token path + that writes an LP line); the slot's `BT` lines carry none, so a longer reply is refused (400) before it starts. """ def __init__(self, svc, allowed=(), owner=None): @@ -5939,8 +5977,11 @@ def raw(self, t): def take(self, eng) -> list: # the engine's current lines are this request's only when they carry its Event (a later generate pass - # of the same request makes new ones: the old ones are read to the end first); another request's never are - cur, old = getattr(eng, "lp_lines", None), None + # of the same request makes new ones: the old ones are read to the end first); another request's never are. + # A batch admission keeps its own lines on the request's cancel Event: the engine's is one attribute, and + # the next request to read control lines may already have replaced it. + own = getattr(self.owner, "_strata_lp_lines", None) + cur, old = (own if own is not None else getattr(eng, "lp_lines", None)), None if cur is not self.lines and getattr(cur, "owner", None) is self.owner: old, self.lines = self.lines, cur out = [] @@ -7430,13 +7471,13 @@ def _openai(self, req): # max_tokens (by the rest of the context when it is not sent), which is the only bound OpenAI and vLLM # put on it either. A reply that is not streamed does hold every entry (its token, bytes and up to 20 # alternatives) until it ends, so a long one holds a lot: streaming hands each entry out with its chunk. - if top_logprobs(req) is not None and getattr(svc.engine, "batch", 0): - # --parallel/--batch: the batched loop prints "BT " and never calls emit_lp_line, so no - # LP line ever reaches the server there (src/program/generate.cpp). Refused rather than answered - # with the field missing, which a client cannot tell from a reply that had no tokens. - raise ValueError("logprobs is not available when the engine runs requests in batch slots " - "(--parallel/--batch): it writes no LP lines there; start the server without " - "--parallel") + if top_logprobs(req) is not None and batch_slots_on(svc.engine) and max_req != 1: + # --parallel/--batch: a slot's "BT " lines carry no LP line, and the batched loop never + # calls emit_lp_line (src/program/generate.cpp); an admission -- a max_tokens 1 request -- is read + # as GEN 1, which does write them. Past that one token a reply would hold entries for a prefix + # only, which a client cannot tell from the model's own answer, so the longer one is refused. + raise ValueError('logprobs is available with "parallel" (batch slots) only for max_tokens 1; ' + 'run the server without "--parallel" for longer replies') n_vocab = len(getattr(svc.tok, "tokens", ())) or None if top_logprobs(req) is not None and n_vocab: # emit_lp_line indexes the logits row with each named id as given, so an id past the vocabulary diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index ca92bcf0f9a..a21237db782 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -20,8 +20,8 @@ from types import SimpleNamespace from serve.frontend import ChatTemplate -from serve.server import (ByteTokenizer, LogprobEntries, LPLines, MockEngine, Service, StrataEngine, logprob_span, - serve, top_logprobs, allowed_ids) +from serve.server import (ByteTokenizer, LogprobEntries, LPLines, MockEngine, Service, StrataEngine, batch_slots_on, + logprob_span, serve, top_logprobs, allowed_ids) ROOT = Path(__file__).resolve().parent.parent TEMPLATE = ChatTemplate(ROOT / "serve" / "chat_template.jinja") @@ -502,6 +502,26 @@ def test_json_schema_answer_as_written_with_its_logprobs(self): self.assertNotIn("logprobs", plain["choices"][0]) +class Replicas(unittest.TestCase): + """`--replicas`: a ReplicaEngine's own `batch` is how many requests its engines run at once and counts a replica + without slots as one, so asking it "does this engine use batch slots" refused logprobs on a server that has + none -- one that reads every reply by the GEN path, which does write LP lines.""" + + @staticmethod + def replica(*batches): + from serve.replicas import ReplicaEngine + return ReplicaEngine([SimpleNamespace(gen=1, batch=b) for b in batches]) + + def test_no_slots_behind_replicas(self): + self.assertFalse(batch_slots_on(self.replica(0, 0))) + self.assertFalse(batch_slots_on(SimpleNamespace(batch=0))) # no replicas: the engine's own figure + + def test_slots_behind_replicas(self): + self.assertTrue(batch_slots_on(self.replica(2, 3))) # every one of them + self.assertTrue(batch_slots_on(self.replica(0, 2))) # one is enough + self.assertTrue(batch_slots_on(SimpleNamespace(batch=4))) + + class Span(unittest.TestCase): def entries(self, *toks): return [{"token": t} for t in toks] diff --git a/serve/test_logprobs_parallel.py b/serve/test_logprobs_parallel.py new file mode 100644 index 00000000000..cc494aa3e43 --- /dev/null +++ b/serve/test_logprobs_parallel.py @@ -0,0 +1,117 @@ +"""Logprobs with "parallel" (batch slots): python -m unittest serve.test_logprobs_parallel. + +The fake engine of serve/test_parallel.py behind the real StrataEngine and Service, writing an LP line after each +T line when the GEN / BGEN line asks for logprobs. A typed decision (max_tokens 1) is read on the control lines +(an admission: BGEN, one T, DONE, BADM 0) and gets its logprobs; longer replies are refused with a 400, since the +slot's BT lines carry none. +""" +import json +import threading +import unittest +import urllib.error +import urllib.request +from unittest import mock + +import serve.test_parallel as tp + +# LP after each T when asked: the token's log-probability is -(prompt length) / 1000 (so two requests can be told +# apart) and "A" (65) is its one alternative at -1.5 -- or the named ids themselves, at -1, when the line carries +# `allowed=` (the engine's own branch: it prints those and no top-K) +_T_LINE = ' print(f"T {t}", flush=True)\n' +assert tp.FAKE_BATCH.count(_T_LINE) == 1 +FAKE_LP = tp.FAKE_BATCH.replace(_T_LINE, _T_LINE + ( + ' if " logprobs=" in line:\n' + ' al = [x[8:] for x in line.split() if x.startswith("allowed=")]\n' + ' cs = (" " + " ".join(i + ":-1.000000" for i in al[0].split(","))) if al else " 65:-1.500000"\n' + ' print(f"LP {-len(ids) / 1000:.6f}" + cs, flush=True)\n')) + + +class LogprobsWithBatchSlots(unittest.TestCase): + start, tearDown, get = tp.ParallelService.start, tp.ParallelService.tearDown, tp.ParallelService.get + + def start_lp(self, fake=FAKE_LP, slots=2): + with mock.patch.object(tp, "FAKE_BATCH", fake): + self.start(slots) + + def chat(self, text, max_tokens=1, logprobs=True, allowed=None): + body = {"messages": [{"role": "user", "content": text}], "max_tokens": max_tokens, + "reasoning_effort": "none"} + if logprobs: + body |= {"logprobs": True, "top_logprobs": 1} + if allowed: + body["allowed_token_ids"] = allowed + req = urllib.request.Request(self.base + "/v1/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=60) as r: + return json.loads(r.read().decode()) + + @staticmethod + def entries(reply): + return ((reply["choices"][0].get("logprobs") or {}).get("content")) or [] + + def test_a_decision_gets_its_logprobs(self): + self.start_lp() + r = self.chat("Is it A or B?") + self.assertEqual(r["choices"][0]["message"]["content"], "o") + e = self.entries(r) + self.assertEqual(len(e), 1) + self.assertEqual(e[0]["token"], "o") + self.assertLess(e[0]["logprob"], 0) + self.assertEqual(e[0]["top_logprobs"], [{"token": "A", "logprob": -1.5, "bytes": [65]}]) + self.assertTrue(any(x.startswith("BGEN") for x in self.log.read_text().splitlines())) # an admission + + def test_the_named_ids_reach_the_slot_and_are_what_is_reported(self): + # allowed_token_ids travels in the same BGEN keys as logprobs, so the engine's own branch answers with + # exactly those ids: the admission path carries the whole feature, not just the top-K one + self.start_lp() + e = self.entries(self.chat("Is it A or B?", allowed=[65])) + self.assertEqual(len(e), 1) + self.assertEqual([x["token"] for x in e[0]["top_logprobs"]], ["A"]) + self.assertEqual(e[0]["top_logprobs"][0]["logprob"], 0) # the whole distribution: just that id + + def test_two_at_once_each_get_their_own(self): + self.start_lp() + short, long_ = "A or B?", "A or B? " + "x" * 300 + alone = {t: self.entries(self.chat(t))[0]["logprob"] for t in (short, long_)} + self.assertNotEqual(alone[short], alone[long_]) + for _ in range(5): + got, errs = {}, [] + + def run(t): + try: + got[t] = self.entries(self.chat(t))[0]["logprob"] + except Exception as exc: # noqa: BLE001 - reported below + errs.append(repr(exc)) + th = [threading.Thread(target=run, args=(t,)) for t in (short, long_, short, long_)] + for x in th: + x.start() + for x in th: + x.join(60) + self.assertEqual(errs, []) + self.assertEqual(got, alone) + + def test_not_asked_nothing_changes(self): + self.start_lp() + r = self.chat("Is it A or B?", logprobs=False) + self.assertEqual(r["choices"][0]["message"]["content"], "o") + self.assertIsNone(r["choices"][0].get("logprobs")) + + def test_longer_replies_are_refused_with_batch_slots(self): + self.start_lp() + with self.assertRaises(urllib.error.HTTPError) as cm: + self.chat("Is it A or B?", max_tokens=8) + self.assertEqual(cm.exception.code, 400) + self.assertIn("max_tokens 1", cm.exception.read().decode()) + # the server is still serving, and a reply that does not ask is untouched + self.assertEqual(self.chat("still serving?", max_tokens=16, logprobs=False)["choices"][0]["message"]["content"], + "ok, done.") + + def test_an_engine_without_lp_lines(self): + self.start_lp(fake=tp.FAKE_BATCH) # an older engine: the key is ignored, no LP line at all + r = self.chat("Is it A or B?") + self.assertEqual(r["choices"][0]["message"]["content"], "o") + self.assertEqual(self.entries(r), []) + + +if __name__ == "__main__": + unittest.main() From 447e49a2c49bc877656a7a572e950c468d603596 Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Mon, 26 Oct 2026 00:49:17 +0800 Subject: [PATCH 11/14] engine: an allowed_token_ids id past the engine's own vocabulary is left off the LP line 26a2d480 refuses an allowed_token_ids id past the vocabulary on the server, but only when the tokenizer reports its size; with no size to check against, the request reaches the engine and emit_lp_line indexes its logits row with the id as given, reading out of bounds. The engine's own vocab() is authoritative, so the named ids are checked against it where they are printed: an id outside 0..vocab()-1 does not appear on the line, and a reply whose named ids are all out of range keeps the chosen token's own logprob with no alternatives -- the shape the server's "no named id observed" case already produces and documents (_renormalise). The cost is one integer compare per named id, a list the server bounds by the vocabulary; no other path moves. This one is from reading the code -- no out-of-bounds read was observed (the server's tokenizer always reports a size today). Incremental CUDA build on this machine (RTX 4090 config, build-lp) compiles and links: [157/158] Linking CXX executable strata.exe. --- src/program/generate.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/program/generate.cpp b/src/program/generate.cpp index b30118dccd8..d93387b0fc2 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -170,9 +170,11 @@ void emit_lp_line(const strata::core::Verifier& ver, int t, int32_t token, int k const double lse = (double) mx + std::log(se); std::printf("LP %.6f", (double) row[(size_t) token] - lse); if (!allowed.empty()) { - // 只要调用方点名的那几个 id:server 侧再按它们归一化 ⇒ 与 allowed_token_ids 语义一致 + // 只要调用方点名的那几个 id:server 侧再按它们归一化 ⇒ 与 allowed_token_ids 语义一致。 + // server 只在能从 tokenizer 读到词表大小时把越界 id 拒为 400;引擎自己的 vocab() 是权威, + // 越界的 id 不出现在行上,也就不会越界读 logits for (const int32_t id : allowed) { - if (id == token) continue; + if (id == token || id < 0 || (int64_t) id >= ver.vocab()) continue; std::printf(" %d:%.6f", (int) id, (double) row[(size_t) id] - lse); } std::printf("\n"); From a82a9b4d469dd1e8fdfe986402e86407edc966f9 Mon Sep 17 00:00:00 2001 From: xiezhipeng-git <651326514@qq.com> Date: Mon, 26 Oct 2026 00:49:28 +0800 Subject: [PATCH 12/14] serve: structured's JSON extractor is a public name (extract_json) server.py's json_schema logprobs alignment reads the model's own JSON form through it, which imported a private name across module lines -- a rename in structured.py would break the server for no reason the compiler can name. The function keeps its behaviour, gains the docstring a public name owes, and the always-true "if _extract_json" at the call site (the import above it already fails when the name is missing) goes with it. No reply changes shape. python -m unittest serve.test_structured serve.test_logprobs serve.test_logprobs_parallel -- 48 tests, OK (structured 10, logprobs 32, parallel 6). --- serve/server.py | 5 ++--- serve/structured.py | 7 +++++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/serve/server.py b/serve/server.py index 658ce827b2f..7a268a46a6b 100644 --- a/serve/server.py +++ b/serve/server.py @@ -80,8 +80,7 @@ from serve import launch_identity # noqa: E402 from serve import logit_bias as bias_api # noqa: E402 from serve import lifeline # noqa: E402 -from serve.structured import StructuredOutputError, prepare_format, validated_json # noqa: E402 -from serve.structured import _extract_json # noqa: E402 +from serve.structured import StructuredOutputError, prepare_format, validated_json, extract_json # noqa: E402 from serve import responses as responses_api # noqa: E402 from serve.responses import ResponsesError, error_body as responses_error_body # noqa: E402 from serve import images as images_api # noqa: E402 @@ -6236,7 +6235,7 @@ def _structured_logprobs(choice: dict, content: str): span = logprob_span(raw, content, entries) if span is not None: return content, _with_reasoning(lps, span) - inner = _extract_json(raw) if _extract_json else "" + inner = extract_json(raw) if inner and _same_json(inner, content): span = logprob_span(raw, inner, entries) if span is not None: diff --git a/serve/structured.py b/serve/structured.py index 9bed53484d7..0c5b6ddb71f 100644 --- a/serve/structured.py +++ b/serve/structured.py @@ -152,7 +152,10 @@ def no_remote(uri): return messages, validator -def _extract_json(text: str) -> str: +def extract_json(text: str) -> str: + """The first JSON object in `text`: a ```-fenced block's inside when there is one, else the text + from its first `{`. Empty when there is nothing to take. The server's json_schema logprobs + alignment reads the model's own JSON form with it.""" s = (text or "").strip() if not s: return "" @@ -204,7 +207,7 @@ def constant(value): if finish != "stop": raise StructuredOutputError(f"structured output was incomplete (finish_reason={finish}); increase the output budget") try: - value = json.loads(_extract_json(text), object_pairs_hook=pairs, parse_constant=constant) + value = json.loads(extract_json(text), object_pairs_hook=pairs, parse_constant=constant) canonical = json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")) except (ValueError, TypeError) as exc: raise StructuredOutputError(f"model did not return valid JSON: {exc}") from exc From a2ea8f95e3dfc556a93ee5f862603c09ce718e46 Mon Sep 17 00:00:00 2001 From: envy-ai <177786885+envy-ai@users.noreply.github.com> Date: Sun, 25 Oct 2026 15:28:21 -0400 Subject: [PATCH 13/14] [Engine] score generated token logprobs on GPU --- CMakeLists.txt | 11 ++ docs/DETAILS.md | 3 + include/strata/core/verify.hpp | 12 ++ include/strata/kernels/logprobs.hpp | 29 +++++ serve/test_logprobs.py | 20 +++ src/core/verify.cpp | 64 ++++++++++ src/kernels/cuda/logprobs.cu | 188 ++++++++++++++++++++++++++++ src/kernels/logprobs_parity.cpp | 174 +++++++++++++++++++++++++ src/program/generate.cpp | 32 +++++ 9 files changed, 533 insertions(+) create mode 100644 include/strata/kernels/logprobs.hpp create mode 100644 src/kernels/cuda/logprobs.cu create mode 100644 src/kernels/logprobs_parity.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index c4e6e8a185c..e681e588152 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -527,6 +527,10 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) endif() # plan v0.3 P6: llama.cpp's block layouts and codebook grids for the i-quant formats (MIT, third_party/ggml) target_include_directories(strata_kernels PRIVATE third_party/ggml) + # gfx906 uses 64-lane waves; this reduction's fixed 32-lane groups are for CUDA and wave32 HIP. + if(NOT STRATA_HIP_GFX906) + target_sources(strata_kernels PRIVATE src/kernels/cuda/logprobs.cu) + endif() # Match the pinned llama.cpp CUDA arithmetic contract for this adapter only. if(STRATA_ENABLE_CUDA) set_source_files_properties(src/kernels/cuda/native_mmvq.cu src/kernels/cuda/native_bf16.cu @@ -1261,6 +1265,13 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) target_link_libraries(sampler_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) add_test(NAME sampler_parity COMMAND sampler_parity --selftest) set_tests_properties(sampler_parity PROPERTIES ENVIRONMENT "STRATA_OLD_SAMPLER=0;STRATA_SAMPLER_ONE_BLOCK=0") + if(NOT STRATA_HIP_GFX906) + add_executable(logprobs_parity src/kernels/logprobs_parity.cpp) + target_include_directories(logprobs_parity PRIVATE ${_strata_gpu_include_directories}) + target_link_libraries(logprobs_parity PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + add_test(NAME logprobs_parity COMMAND logprobs_parity) + set_tests_properties(logprobs_parity PROPERTIES SKIP_RETURN_CODE 77) + endif() # the GPU verifier of probabilistic draft acceptance against the host reference (STRATA_SPEC_PROB) add_executable(spec_verify_parity src/kernels/spec_verify_parity.cpp) target_include_directories(spec_verify_parity PRIVATE ${_strata_gpu_include_directories}) diff --git a/docs/DETAILS.md b/docs/DETAILS.md index fb796ec4df8..51fd005f104 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -520,6 +520,9 @@ applies at the next start. Network, key, MCP and program keys are not editable t `--parallel`, only a `max_tokens` 1 request gets them: it is the one a batch slot decodes through the single-token path, which writes an LP line. A longer reply there is a 400, since the slot's `BT` lines carry none -- entries for a prefix only, which a client cannot tell from the model's own answer. + In serial CUDA and wave32 HIP decode, the engine scores accepted tokens and top IDs on the GPU and transfers + those scores. Equal-logit top IDs are ordered by token ID; score rounding can differ from the CPU path. Named + `allowed_token_ids` and pipelined windows keep the full-logit CPU scoring path. - **Non-finite logits end the request and drop the cached state (on by default).** If a verify window's logits are NaN, the window emits no tokens, the request ends, and the next request drops the live session, checkpoints and parked conversations and reads its prompt from token 0. The server retries a request that had sent nothing once, transparently. diff --git a/include/strata/core/verify.hpp b/include/strata/core/verify.hpp index dc6613da748..8c8cca283f5 100644 --- a/include/strata/core/verify.hpp +++ b/include/strata/core/verify.hpp @@ -27,6 +27,7 @@ #include "strata/core/layer.hpp" #include "strata/core/session.hpp" #include "strata/kernels/sampler.hpp" +#include "strata/kernels/logprobs.hpp" #include @@ -105,6 +106,12 @@ class Verifier { bool run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err); /// Diagnostics: row `t` of the last window's head logits (n_vocab floats) to the host. Valid after run(). bool copy_logits(int t, float* host) const; +#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906) + /// Raw head-softmax scores for accepted tokens in the first `rows` of the last window. + /// Returns one chosen token plus `top_k` alternatives per row, before the next window overwrites logits. + bool logprobs(int rows, const int32_t* chosen, int top_k, strata::kernels::LogprobItem* out, + std::string& err); +#endif int64_t vocab() const { return next_ ? next_->vocab() : n_vocab_; } /// #879 GUARD: the first row of the last finished window (run, run_slot_rows, batch_poll, pl_finish) whose head /// logits held a NaN or an infinity, else -1; nonfinite_pos() is that row's position. Every window graph @@ -645,6 +652,11 @@ class Verifier { float* hit_xs_ = nullptr; void* hit_scratch_ = nullptr; float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr; +#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906) + void* logprobs_scratch_ = nullptr; + int32_t* logprobs_chosen_ = nullptr; + strata::kernels::LogprobItem* logprobs_out_ = nullptr; +#endif uint16_t* sh_bf16_ = nullptr; uint8_t* arg_scratch_ = nullptr; ///< argmax_rows' partials and counters int32_t* one_ = nullptr; ///< device {1}: the n_keep of a one-token window, which commits itself diff --git a/include/strata/kernels/logprobs.hpp b/include/strata/kernels/logprobs.hpp new file mode 100644 index 00000000000..f8a9cbdb1fc --- /dev/null +++ b/include/strata/kernels/logprobs.hpp @@ -0,0 +1,29 @@ +#pragma once + +#include +#include + +namespace strata::kernels { + +// One chosen token followed by K highest-scoring alternatives for each row. +// Scores are natural log probabilities under the unmodified, full-vocabulary +// softmax of the head logits. The chosen token need not be in the top K. +struct LogprobItem { + int32_t id; + float logprob; +}; + +constexpr int kLogprobsMaxTop = 20; + +// Device scratch for logprobs_extract. The head vocabulary must contain at most +// 262144 entries, as in the split sampler. Returns zero for unsupported sizes. +size_t logprobs_scratch_bytes(int rows, int vocab, int top_k); + +// All pointers, including chosen and out, are device pointers. Scratch is at +// least logprobs_scratch_bytes(rows, vocab, top_k) bytes. The caller preserves +// logits until the stream has finished, and transfers only rows*(K+1) items. +// Returns false without launching for invalid dimensions. +bool logprobs_extract(const float* logits, int rows, int vocab, const int32_t* chosen, + int top_k, void* scratch, LogprobItem* out, void* stream); + +} // namespace strata::kernels diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index a21237db782..0cd2bab8730 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -273,6 +273,26 @@ def test_the_tokens_spell_the_content(self): self.assertEqual(lps[0]["top_logprobs"][0]["token"], "A") self.assertEqual(self.engine.sampling.get("logprobs"), True) + def test_engine_wire_keeps_chosen_score_separate_from_raw_top_k(self): + # The LP header is the sampled token's full-vocabulary score. Following pairs are raw top-K IDs; + # the chosen token can be outside that list, and the server must not replace its score with one of them. + class TwoAlternativeEngine(LpEngine): + def line(self, i, t): + return "LP -2.500000 65:-0.500000 72:-1.250000\n" + + self.engine = self.svc.engine = TwoAlternativeEngine(self.tok, "Z") + for stream in (False, True): + out = self.post({"logprobs": True, "top_logprobs": 2}, stream) + if stream: + entries = [entry for chunk in out if chunk.get("choices") + for entry in chunk["choices"][0].get("logprobs", {}).get("content", [])] + else: + entries = out["choices"][0]["logprobs"]["content"] + self.assertEqual([entry["token"] for entry in entries], ["Z"]) + self.assertEqual(entries[0]["logprob"], -2.5) + self.assertEqual([alt["token"] for alt in entries[0]["top_logprobs"]], ["A", "H"]) + self.assertEqual([alt["logprob"] for alt in entries[0]["top_logprobs"]], [-0.5, -1.25]) + def test_a_request_reads_only_its_own_lines(self): # A's last "." is held by its stop string and sent after A has left the queue, when B is # already generating on the same engine: each reply's tokens must spell that reply diff --git a/src/core/verify.cpp b/src/core/verify.cpp index 3d6f48b2e13..ce37ed60925 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -75,6 +75,7 @@ #include #include #include +#include #if defined(STRATA_GPU_METAL) // The Metal backend's decode window (metal/src/engine/window.cpp): this file's init/run/commit hand over to it. @@ -635,6 +636,11 @@ Verifier::~Verifier() { if (tap_scratch_) cudaFree(tap_scratch_); if (h_tap_) cudaFreeHost(h_tap_); if (logit_bias_device_) cudaFree(logit_bias_device_); +#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906) + if (logprobs_scratch_) cudaFree(logprobs_scratch_); + if (logprobs_chosen_) cudaFree(logprobs_chosen_); + if (logprobs_out_) cudaFree(logprobs_out_); +#endif if (cs_ && cs_ != ext_stream_) cudaStreamDestroy(cs_); // set_stream: the stage's stream, shared, not ours if (sh_cs_) cudaStreamDestroy(sh_cs_); if (copy_) { cudaStreamSynchronize(copy_); cudaStreamDestroy(copy_); } @@ -4147,6 +4153,64 @@ bool Verifier::copy_logits(int t, float* host) const { cudaMemcpyDeviceToHost) == cudaSuccess; } +#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906) +bool Verifier::logprobs(int rows, const int32_t* chosen, int top_k, + strata::kernels::LogprobItem* out, std::string& err) { + if (next_ != nullptr) return next_->logprobs(rows, chosen, top_k, out, err); + if (head_logits_ == nullptr || chosen == nullptr || out == nullptr || rows < 1 || + rows > last_t_ || rows > max_t_ || top_k < 0 || top_k > strata::kernels::kLogprobsMaxTop || + top_k > n_vocab_ || n_vocab_ > std::numeric_limits::max()) { + err = "verify: invalid logprobs rows, top_k, or head state"; + return false; + } + for (int t = 0; t < rows; ++t) + if (chosen[t] < 0 || chosen[t] >= n_vocab_ || nonfinite_in(t)) { + err = "verify: logprobs selected an invalid token or a non-finite head row"; + return false; + } + const size_t scratch_bytes = strata::kernels::logprobs_scratch_bytes( + max_t_, (int) n_vocab_, std::min(strata::kernels::kLogprobsMaxTop, (int) n_vocab_)); + if (scratch_bytes == 0) { + err = "verify: logprobs vocabulary is outside the supported range"; + return false; + } + const OnDevice on(device_); + auto failed = [&](const char* step, cudaError_t code) { + err = std::string("verify: logprobs ") + step + ": " + cudaGetErrorString(code); + return false; + }; + cudaError_t rc = cudaSuccess; + if (!logprobs_scratch_ && (rc = cudaMalloc(&logprobs_scratch_, scratch_bytes)) != cudaSuccess) + return failed("scratch allocation", rc); + if (!logprobs_chosen_ && (rc = cudaMalloc((void**) &logprobs_chosen_, (size_t) max_t_ * sizeof(int32_t))) != cudaSuccess) + return failed("chosen-token allocation", rc); + if (!logprobs_out_ && + (rc = cudaMalloc((void**) &logprobs_out_, (size_t) max_t_ * + (strata::kernels::kLogprobsMaxTop + 1) * + sizeof(strata::kernels::LogprobItem))) != cudaSuccess) + return failed("result allocation", rc); + if ((rc = cudaMemcpyAsync(logprobs_chosen_, chosen, (size_t) rows * sizeof(int32_t), + cudaMemcpyHostToDevice, cs_)) != cudaSuccess) + return failed("chosen-token upload", rc); + if (!strata::kernels::logprobs_extract(head_logits_, rows, (int) n_vocab_, logprobs_chosen_, top_k, + logprobs_scratch_, logprobs_out_, cs_)) + return failed("kernel launch", cudaGetLastError()); + if ((rc = cudaMemcpyAsync(out, logprobs_out_, (size_t) rows * (top_k + 1) * + sizeof(strata::kernels::LogprobItem), cudaMemcpyDeviceToHost, cs_)) != cudaSuccess) + return failed("result download", rc); + if ((rc = cudaStreamSynchronize(cs_)) != cudaSuccess) return failed("completion", rc); + for (int t = 0; t < rows; ++t) + for (int j = 0; j <= top_k; ++j) { + const auto& item = out[(size_t) t * (top_k + 1) + j]; + if (item.id < 0 || item.id >= n_vocab_ || !std::isfinite(item.logprob)) { + err = "verify: logprobs returned a non-finite or invalid result"; + return false; + } + } + return true; +} +#endif + // ================================ PIPELINED WINDOWS (see verify.hpp) ================================ // diff --git a/src/kernels/cuda/logprobs.cu b/src/kernels/cuda/logprobs.cu new file mode 100644 index 00000000000..89e0bc87c30 --- /dev/null +++ b/src/kernels/cuda/logprobs.cu @@ -0,0 +1,188 @@ +#include "strata/kernels/logprobs.hpp" + +#include + +#include + +namespace strata::kernels { +namespace { + +constexpr int kThreads = 128; +constexpr int kPerThread = 32; +constexpr int kSpan = kThreads * kPerThread; +constexpr int kMaxBlocks = 64; + +__device__ __forceinline__ void take_best(float& score, int& id, float other, int other_id) { + if (other > score || (other == score && other_id < id)) { + score = other; + id = other_id; + } +} + +__device__ __forceinline__ void warp_best(float& score, int& id) { + for (int step = 16; step > 0; step >>= 1) { + const float v = __shfl_down_sync(0xffffffffu, score, step); + const int i = __shfl_down_sync(0xffffffffu, id, step); + take_best(score, id, v, i); + } + score = __shfl_sync(0xffffffffu, score, 0); + id = __shfl_sync(0xffffffffu, id, 0); +} + +__device__ __forceinline__ float warp_max(float x) { + for (int step = 16; step > 0; step >>= 1) x = fmaxf(x, __shfl_down_sync(0xffffffffu, x, step)); + return __shfl_sync(0xffffffffu, x, 0); +} + +__device__ __forceinline__ double warp_sum(double x) { + for (int step = 16; step > 0; step >>= 1) x += __shfl_down_sync(0xffffffffu, x, step); + return __shfl_sync(0xffffffffu, x, 0); +} + +// One block owns 4096 logits. It reads them once into registers; the stable +// local log-sum-exp and top-K selection share the same loads. The 61 blocks of +// a 248320-token vocabulary run across the GPU rather than serially on one SM. +__global__ void logprobs_part(const float* __restrict__ logits, int vocab, int blocks, int k, + float2* __restrict__ sums, int2* __restrict__ candidates) { + const int row = blockIdx.y, block = blockIdx.x; + const int tid = threadIdx.x, base = block * kSpan + tid; + float values[kPerThread]; + float local_max = -INFINITY; +#pragma unroll + for (int j = 0; j < kPerThread; ++j) { + const int id = base + kThreads * j; + const float v = id < vocab ? logits[(size_t) row * vocab + id] : -INFINITY; + values[j] = v; + local_max = fmaxf(local_max, v); + } + __shared__ float warp_mx[4]; + __shared__ double warp_sm[4]; + __shared__ float block_mx; + __shared__ int chosen_id; + __shared__ float warp_score[4]; + __shared__ int warp_id[4]; + const int warp = tid >> 5, lane = tid & 31; + const float mx = warp_max(local_max); + if (lane == 0) warp_mx[warp] = mx; + __syncthreads(); + if (tid == 0) { + float m = warp_mx[0]; + for (int w = 1; w < 4; ++w) m = fmaxf(m, warp_mx[w]); + block_mx = m; + } + __syncthreads(); + double local_sum = 0.0; +#pragma unroll + for (int j = 0; j < kPerThread; ++j) + if (values[j] != -INFINITY) local_sum += exp((double) values[j] - (double) block_mx); + const double sm = warp_sum(local_sum); + if (lane == 0) warp_sm[warp] = sm; + __syncthreads(); + if (tid == 0) { + sums[(size_t) row * blocks + block] = make_float2(block_mx, + (float) (warp_sm[0] + warp_sm[1] + warp_sm[2] + warp_sm[3])); + } + __syncthreads(); + + for (int rank = 0; rank < k; ++rank) { + float best = -INFINITY; + int id = vocab; +#pragma unroll + for (int j = 0; j < kPerThread; ++j) { + const int candidate_id = base + kThreads * j; + if (candidate_id < vocab) take_best(best, id, values[j], candidate_id); + } + warp_best(best, id); + if (lane == 0) { warp_score[warp] = best; warp_id[warp] = id; } + __syncthreads(); + if (tid == 0) { + float b = warp_score[0]; + int i = warp_id[0]; + for (int w = 1; w < 4; ++w) take_best(b, i, warp_score[w], warp_id[w]); + chosen_id = i; + candidates[((size_t) row * blocks + block) * k + rank] = make_int2(i, __float_as_int(b)); + } + __syncthreads(); + if (chosen_id >= base && chosen_id < base + kThreads * kPerThread && + (chosen_id - base) % kThreads == 0) + values[(chosen_id - base) / kThreads] = -INFINITY; + __syncthreads(); + } +} + +// One warp combines all block sums and merges their sorted top-K lists. Only +// K+1 pairs per row leave the GPU, even for a 248320-token vocabulary. +__global__ void logprobs_finish(const float* __restrict__ logits, int vocab, int blocks, int k, + const int32_t* __restrict__ chosen, const float2* __restrict__ sums, + const int2* __restrict__ candidates, LogprobItem* __restrict__ out) { + const int row = blockIdx.x, lane = threadIdx.x; + float local_max = -INFINITY; + for (int b = lane; b < blocks; b += 32) + local_max = fmaxf(local_max, sums[(size_t) row * blocks + b].x); + const float mx = warp_max(local_max); + double local_sum = 0.0; + for (int b = lane; b < blocks; b += 32) { + const float2 part = sums[(size_t) row * blocks + b]; + if (part.y > 0.0f) local_sum += (double) part.y * exp((double) part.x - (double) mx); + } + const double log_sum = log(warp_sum(local_sum)); + const int picked = chosen[row]; + if (lane == 0) out[(size_t) row * (k + 1)] = {picked, (float) (((double) logits[(size_t) row * vocab + picked] - (double) mx) - log_sum)}; + + // Each lane owns up to two sorted block lists. The first K elements of a + // union must occur in the first K elements of one of its constituent lists. + int cursor[2] = {0, 0}; + float head_score[2] = {-INFINITY, -INFINITY}; + int head_id[2] = {vocab, vocab}; + for (int j = 0; j < 2; ++j) { + const int b = lane + 32 * j; + if (b < blocks && k > 0) { + const int2 p = candidates[((size_t) row * blocks + b) * k]; + head_id[j] = p.x; + head_score[j] = __int_as_float(p.y); + } + } + for (int rank = 0; rank < k; ++rank) { + float score = head_score[0]; + int id = head_id[0]; + take_best(score, id, head_score[1], head_id[1]); + warp_best(score, id); + if (lane == 0) out[(size_t) row * (k + 1) + rank + 1] = {id, (float) (((double) score - (double) mx) - log_sum)}; + for (int j = 0; j < 2; ++j) { + const int b = lane + 32 * j; + if (b >= blocks || head_id[j] != id) continue; + if (++cursor[j] < k) { + const int2 p = candidates[((size_t) row * blocks + b) * k + cursor[j]]; + head_id[j] = p.x; + head_score[j] = __int_as_float(p.y); + } else { + head_id[j] = vocab; + head_score[j] = -INFINITY; + } + } + } +} + +} // namespace + +size_t logprobs_scratch_bytes(int rows, int vocab, int top_k) { + if (rows < 1 || vocab < 1 || vocab > kSpan * kMaxBlocks || top_k < 0 || + top_k > kLogprobsMaxTop || top_k > vocab) return 0; + const size_t blocks = ((size_t) vocab + kSpan - 1) / kSpan; + return (size_t) rows * blocks * (sizeof(float2) + (size_t) top_k * sizeof(int2)); +} + +bool logprobs_extract(const float* logits, int rows, int vocab, const int32_t* chosen, + int top_k, void* scratch, LogprobItem* out, void* stream) { + if (!logits || !chosen || !scratch || !out || !logprobs_scratch_bytes(rows, vocab, top_k)) return false; + const int blocks = (vocab + kSpan - 1) / kSpan; + auto* sums = static_cast(scratch); + auto* candidates = reinterpret_cast(sums + (size_t) rows * blocks); + logprobs_part<<>>( + logits, vocab, blocks, top_k, sums, candidates); + logprobs_finish<<>>( + logits, vocab, blocks, top_k, chosen, sums, candidates, out); + return cudaPeekAtLastError() == cudaSuccess; +} + +} // namespace strata::kernels diff --git a/src/kernels/logprobs_parity.cpp b/src/kernels/logprobs_parity.cpp new file mode 100644 index 00000000000..cb0d0f4f482 --- /dev/null +++ b/src/kernels/logprobs_parity.cpp @@ -0,0 +1,174 @@ +#include "strata/kernels/logprobs.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using strata::kernels::LogprobItem; + +namespace { + +void cuda_ok(cudaError_t rc, const char* what) { + if (rc != cudaSuccess) { + std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(rc)); + std::exit(1); + } +} + +struct Case { + int vocab; + int rows; + int top_k; + const char* label; +}; + +float value_for(int row, int id, int vocab) { + if (row == 1) return 1.0e20f; // large common offset catches LSE cancellation + if (row == 2) { + if (id == 0 || id == vocab / 2 || id == vocab - 1) return 1.0e20f; + if (id % 5 == 0) return -1.0e20f; + return (float) ((id * 13) % 17) - 8.0f; + } + // Deterministic ties across and within 4096-logit partitions, plus a wide score range. + if (id == 1 || id == 4095 || id == 4096) return 100.0f; + return (float) ((id * 37 + row * 11) % 101) / 10.0f; +} + +struct Expected { + int id; + float logprob; +}; + +std::vector reference(const std::vector& logits, int vocab, int row, + int chosen, int top_k) { + const float* values = logits.data() + (size_t) row * vocab; + const double mx = (double) *std::max_element(values, values + vocab); + double sum = 0.0; + for (int i = 0; i < vocab; ++i) sum += std::exp((double) values[i] - mx); + const double log_sum = std::log(sum); + std::vector order(vocab); + for (int i = 0; i < vocab; ++i) order[i] = i; + std::partial_sort(order.begin(), order.begin() + top_k, order.end(), [&](int a, int b) { + if (values[a] != values[b]) return values[a] > values[b]; + return a < b; + }); + std::vector result; + result.reserve((size_t) top_k + 1); + result.push_back({chosen, (float) (((double) values[chosen] - mx) - log_sum)}); + for (int i = 0; i < top_k; ++i) + result.push_back({order[i], (float) (((double) values[order[i]] - mx) - log_sum)}); + return result; +} + +bool run_case(const Case& c) { + const size_t count = (size_t) c.rows * c.vocab; + std::vector h_logits(count); + std::vector h_chosen(c.rows); + for (int r = 0; r < c.rows; ++r) { + for (int v = 0; v < c.vocab; ++v) h_logits[(size_t) r * c.vocab + v] = value_for(r, v, c.vocab); + h_chosen[r] = r == 1 ? c.vocab - 1 : (r * 31 + 7) % c.vocab; + } + + const size_t scratch_bytes = strata::kernels::logprobs_scratch_bytes(c.rows, c.vocab, c.top_k); + if (!scratch_bytes) { + std::fprintf(stderr, "%s: valid shape returned zero scratch size\n", c.label); + return false; + } + float* d_logits = nullptr; + int32_t* d_chosen = nullptr; + void* d_scratch = nullptr; + LogprobItem* d_out = nullptr; + const size_t out_count = (size_t) c.rows * (c.top_k + 1); + cuda_ok(cudaMalloc(&d_logits, count * sizeof(float)), "cudaMalloc logits"); + cuda_ok(cudaMalloc(&d_chosen, c.rows * sizeof(int32_t)), "cudaMalloc chosen"); + cuda_ok(cudaMalloc(&d_scratch, scratch_bytes), "cudaMalloc scratch"); + cuda_ok(cudaMalloc(&d_out, out_count * sizeof(LogprobItem)), "cudaMalloc output"); + cuda_ok(cudaMemcpy(d_logits, h_logits.data(), count * sizeof(float), cudaMemcpyHostToDevice), "copy logits"); + cuda_ok(cudaMemcpy(d_chosen, h_chosen.data(), c.rows * sizeof(int32_t), cudaMemcpyHostToDevice), "copy chosen"); + + bool ok = strata::kernels::logprobs_extract(d_logits, c.rows, c.vocab, d_chosen, c.top_k, + d_scratch, d_out, nullptr); + if (!ok) { + std::fprintf(stderr, "%s: extraction refused a valid input\n", c.label); + ok = false; + } else { + std::vector got(out_count); + cuda_ok(cudaMemcpy(got.data(), d_out, out_count * sizeof(LogprobItem), cudaMemcpyDeviceToHost), "copy output"); + for (int r = 0; r < c.rows && ok; ++r) { + const auto want = reference(h_logits, c.vocab, r, h_chosen[r], c.top_k); + for (int j = 0; j <= c.top_k; ++j) { + const auto& actual = got[(size_t) r * (c.top_k + 1) + j]; + const auto& expected = want[(size_t) j]; + if (actual.id != expected.id || !std::isfinite(actual.logprob) || + std::abs(actual.logprob - expected.logprob) > 3.0e-5f) { + std::fprintf(stderr, "%s row=%d item=%d got=(%d, %.9g) expected=(%d, %.9g)\n", + c.label, r, j, actual.id, actual.logprob, expected.id, expected.logprob); + ok = false; + break; + } + } + } + if (ok && c.vocab > 100000) { + const float expected = -std::log((float) c.vocab); + for (int j = 0; j <= c.top_k; ++j) { + const auto& item = got[(size_t) (c.top_k + 1) + j]; + if (std::abs(item.logprob - expected) > 3.0e-5f) { + std::fprintf(stderr, "%s all-equal analytic score: got %.9g expected %.9g\n", + c.label, item.logprob, expected); + ok = false; + break; + } + } + } + } + + cuda_ok(cudaFree(d_out), "free output"); + cuda_ok(cudaFree(d_scratch), "free scratch"); + cuda_ok(cudaFree(d_chosen), "free chosen"); + cuda_ok(cudaFree(d_logits), "free logits"); + return ok; +} + +bool invalid_shapes_refused() { + return strata::kernels::logprobs_scratch_bytes(0, 1, 0) == 0 && + strata::kernels::logprobs_scratch_bytes(1, 0, 0) == 0 && + strata::kernels::logprobs_scratch_bytes(1, 262145, 0) == 0 && + strata::kernels::logprobs_scratch_bytes(1, 16, -1) == 0 && + strata::kernels::logprobs_scratch_bytes(1, 16, 21) == 0 && + strata::kernels::logprobs_scratch_bytes(1, 4, 5) == 0; +} + +} // namespace + +int main() { + int device_count = 0; + const cudaError_t device_status = cudaGetDeviceCount(&device_count); + if (device_status != cudaSuccess || device_count == 0) { + std::fprintf(stderr, "logprobs parity requires a CUDA device\n"); + return 77; + } + const Case cases[] = { + {1, 2, 0, "single-token vocabulary K=0"}, + {7, 2, 7, "K equals vocabulary"}, + {17, 3, 0, "K=0"}, + {4097, 2, 1, "partial-block K=1"}, + {8193, 4, 20, "multiple-block rows (accepted MTP rows)"}, + {248320, 4, 20, "full 61-block vocabulary"}, + }; + for (const auto& c : cases) { + if (!run_case(c)) return 1; + } + if (!invalid_shapes_refused()) { + std::fprintf(stderr, "invalid dimensions were accepted\n"); + return 1; + } + std::puts("logprobs parity passed"); + return 0; +} diff --git a/src/program/generate.cpp b/src/program/generate.cpp index d93387b0fc2..3a5c6814e9e 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -193,6 +193,17 @@ void emit_lp_line(const strata::core::Verifier& ver, int t, int32_t token, int k std::printf("\n"); } +#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906) +// The ordinary decode path extracts all accepted rows in one GPU call. Keep +// the LP wire and chosen-token-plus-top-K shape; GPU ties sort by token ID. +void emit_lp_items_line(const strata::kernels::LogprobItem* row, int k) { + std::printf("LP %.6f", (double) row[0].logprob); + for (int j = 1; j <= k; ++j) + std::printf(" %d:%.6f", (int) row[j].id, (double) row[j].logprob); + std::printf("\n"); +} +#endif + namespace { // Windows' WDDM driver model: native Windows, or WSL2 (its GPU goes through /dev/dxg to the Windows driver). There, // pinning a large arena into two CUDA contexts leaves WDDM refusing every later allocation (the 5080 + 3090 rig); @@ -15247,9 +15258,30 @@ int main(int argc, char** argv) { draft_accepted += a; } first_window = false; +#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906) + // Named allowed IDs use the existing CPU path, which emits every named + // ID rather than top K. Unsupported vocabulary/K sizes use it too. + const bool lp_gpu = req_logprobs >= 0 && req_allowed.empty() && + ver.vocab() <= std::numeric_limits::max() && + strata::kernels::logprobs_scratch_bytes(1, (int) ver.vocab(), req_logprobs) != 0; + std::vector lp_items; + if (lp_gpu) { + lp_items.resize((size_t) (a + 1) * (req_logprobs + 1)); + if (!ver.logprobs(a + 1, outv.data(), req_logprobs, lp_items.data(), err)) { + if (adapt_thr.joinable()) adapt_thr.join(); + std::printf("ERR %s\n", err.c_str()); + return 1; + } + } +#endif bool eos = false; for (int i = merged ? a : 0; i <= a && produced_n < max_new && !eos; ++i) { // (merged: row T-1 only) std::printf("T %d\n", (int) outv[(size_t) i]); +#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906) + if (lp_gpu) + emit_lp_items_line(lp_items.data() + (size_t) i * (req_logprobs + 1), req_logprobs); + else +#endif emit_lp_line(ver, i, outv[(size_t) i], req_logprobs, req_allowed); strata::core::progress_beat(); ++produced_n; From 7265fc75dfbbf714cc03d740e95b41f9ea0586fb Mon Sep 17 00:00:00 2001 From: envy-ai <177786885+envy-ai@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:36:31 -0400 Subject: [PATCH 14/14] [Engine] allocate verifier logprobs buffers at startup --- src/core/verify.cpp | 34 +++++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/src/core/verify.cpp b/src/core/verify.cpp index ce37ed60925..5d98fbd8ec3 100644 --- a/src/core/verify.cpp +++ b/src/core/verify.cpp @@ -914,6 +914,30 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState& return false; } cudaMemset(arena_, 0, count.used); +#if !defined(STRATA_GPU_METAL) && !defined(STRATA_HIP_GFX906) + // Keep logprobs buffers separate from the arena, but fail allocation at startup. + // Unsupported vocabularies allocate nothing and retain the decoder's CPU path. + const size_t scratch_bytes = n_vocab_ <= std::numeric_limits::max() + ? strata::kernels::logprobs_scratch_bytes( + max_t_, (int) n_vocab_, std::min(strata::kernels::kLogprobsMaxTop, (int) n_vocab_)) + : 0; + if (scratch_bytes != 0) { + auto failed = [&](const char* step, cudaError_t code) { + err = std::string("verify: logprobs ") + step + ": " + cudaGetErrorString(code); + return false; + }; + cudaError_t rc = cudaSuccess; + if (!logprobs_scratch_ && (rc = cudaMalloc(&logprobs_scratch_, scratch_bytes)) != cudaSuccess) + return failed("scratch allocation", rc); + if (!logprobs_chosen_ && (rc = cudaMalloc((void**) &logprobs_chosen_, (size_t) max_t_ * sizeof(int32_t))) != cudaSuccess) + return failed("chosen-token allocation", rc); + if (!logprobs_out_ && + (rc = cudaMalloc((void**) &logprobs_out_, (size_t) max_t_ * + (strata::kernels::kLogprobsMaxTop + 1) * + sizeof(strata::kernels::LogprobItem))) != cudaSuccess) + return failed("result allocation", rc); + } +#endif if (hflag_out_d_ != nullptr && hcount_d_ == nullptr) { // the overlapped split's hand-off counter if (cudaMalloc((void**) &hcount_d_, 64) != cudaSuccess) { err = "verify: hand-off counter allocation failed"; return false; } cudaMemset(hcount_d_, 0, 64); @@ -4171,6 +4195,7 @@ bool Verifier::logprobs(int rows, const int32_t* chosen, int top_k, const size_t scratch_bytes = strata::kernels::logprobs_scratch_bytes( max_t_, (int) n_vocab_, std::min(strata::kernels::kLogprobsMaxTop, (int) n_vocab_)); if (scratch_bytes == 0) { + // Unreachable from the decoder: its GPU gate selects CPU scoring for this case. err = "verify: logprobs vocabulary is outside the supported range"; return false; } @@ -4180,15 +4205,6 @@ bool Verifier::logprobs(int rows, const int32_t* chosen, int top_k, return false; }; cudaError_t rc = cudaSuccess; - if (!logprobs_scratch_ && (rc = cudaMalloc(&logprobs_scratch_, scratch_bytes)) != cudaSuccess) - return failed("scratch allocation", rc); - if (!logprobs_chosen_ && (rc = cudaMalloc((void**) &logprobs_chosen_, (size_t) max_t_ * sizeof(int32_t))) != cudaSuccess) - return failed("chosen-token allocation", rc); - if (!logprobs_out_ && - (rc = cudaMalloc((void**) &logprobs_out_, (size_t) max_t_ * - (strata::kernels::kLogprobsMaxTop + 1) * - sizeof(strata::kernels::LogprobItem))) != cudaSuccess) - return failed("result allocation", rc); if ((rc = cudaMemcpyAsync(logprobs_chosen_, chosen, (size_t) rows * sizeof(int32_t), cudaMemcpyHostToDevice, cs_)) != cudaSuccess) return failed("chosen-token upload", rc);