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 43538ddca37..51fd005f104 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -505,6 +505,24 @@ applies at the next start. Network, key, MCP and program keys are not editable t `STRATA_SPEC_GUMBEL=1` raise the accepted-draft rate on some machines; `STRATA_SPEC_PROB=1` is speculative rejection sampling. `STRATA_SPEC_DEPTH=1` prints acceptance by draft depth and `STRATA_SPEC_MIN_TEMP=0.9` limits sampled drafting to hotter requests. Try them on your own card before relying on them. +- **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 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 + a 400. A token whose probability could not be read gets no entry, and the reply goes on. On a server started with + `--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/server.py b/serve/server.py index 76b5a2c7971..7a268a46a6b 100644 --- a/serve/server.py +++ b/serve/server.py @@ -80,7 +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 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 @@ -1456,6 +1456,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 @@ -1474,35 +1483,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 @@ -1779,6 +1809,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) @@ -1833,7 +1869,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) @@ -1961,7 +1998,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) @@ -2195,6 +2233,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 = 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 \ f"GEN {int(max_new)}{self.sampling_keys(sampling or {})}" @@ -2222,7 +2266,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 @@ -2249,7 +2296,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(): @@ -2292,9 +2350,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 @@ -3720,6 +3781,40 @@ 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] + + +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 + + 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.""" @@ -5814,6 +5909,141 @@ def work(c=c, box=box): PP_DONE_TAIL = 64 + +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. + + 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`: 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. + + 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): + 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.""" + 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: + # 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. + # 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 = [] + 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, t, lp) + out.append(e) + return out + + 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 cand: + x["logprob"] -= lse + 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): """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 @@ -5847,6 +6077,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), 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 + 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) @@ -5863,9 +6100,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 @@ -5905,11 +6142,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 "") @@ -5936,11 +6176,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 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 = [] @@ -5956,12 +6265,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() @@ -7146,6 +7466,25 @@ 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 + # 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 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 + # 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/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 diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py new file mode 100644 index 00000000000..0cd2bab8730 --- /dev/null +++ b/serve/test_logprobs.py @@ -0,0 +1,561 @@ +"""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 json +import math +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, 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") + + +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 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, 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.""" + + 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}}} + + +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) + + 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]) + + 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.""" + + 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_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 + 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_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 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) + 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_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 + # (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_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")) + 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_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 + # 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 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] + + 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/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() diff --git a/src/core/verify.cpp b/src/core/verify.cpp index 3d6f48b2e13..5d98fbd8ec3 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_); } @@ -908,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); @@ -4147,6 +4177,56 @@ 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) { + // Unreachable from the decoder: its GPU gate selects CPU scoring for this case. + 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 ((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 0b153ee16aa..3a5c6814e9e 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -151,6 +151,59 @@ 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 语义一致。 + // server 只在能从 tokenizer 读到词表大小时把越界 id 拒为 400；引擎自己的 vocab() 是权威， + // 越界的 id 不出现在行上，也就不会越界读 logits + for (const int32_t id : allowed) { + 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"); + 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"); +} + +#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); @@ -12338,6 +12391,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 @@ -12377,6 +12435,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; @@ -14770,6 +14840,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]); @@ -15187,9 +15258,31 @@ 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; if (sfx_on) sfx.append(outv[(size_t) i]);