From 543cf0cea930809fb8b2fa26607f422884706156 Mon Sep 17 00:00:00 2001 From: tak Date: Wed, 30 Sep 2026 12:05:04 +0900 Subject: [PATCH] vLLM backend: raise on a label missing from top_logprobs instead of filling -30 vLLM's default --logprobs-mode, raw_logprobs, computes logprobs before allowed_token_ids masks the vocabulary (vllm/v1/sample/sampler.py), so the top K returned for a K-label request can hold other tokens in place of labels. The backend filled every missing label with -30.0, which turns into a confident wrong answer with no symptom. Raise LabelTokenError naming the missing ids and the tokens that were returned, and document --logprobs-mode processed_logprobs in the generate-server example. --- anyjev/backends/vllm.py | 23 +++++++++--- tests/test_truncate_and_vllm.py | 66 +++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 6 deletions(-) diff --git a/anyjev/backends/vllm.py b/anyjev/backends/vllm.py index 03a82a5..3c69076 100644 --- a/anyjev/backends/vllm.py +++ b/anyjev/backends/vllm.py @@ -4,11 +4,14 @@ cover every level. **raw / L0 / L1** -- a generate server. One request per prompt with `max_tokens=1`, -`allowed_token_ids` restricted to the label tokens and `logprobs=K`. vLLM reports -logprobs after its logit processors, so the K entries are exactly the labels, -normalized over them. Prefix caching makes the K permutations of one state cheap. - - vllm serve Qwen/Qwen3-8B --enable-prefix-caching +`allowed_token_ids` restricted to the label tokens and `logprobs=K`. Start it with +`--logprobs-mode processed_logprobs`: vLLM's default, `raw_logprobs`, computes logprobs +before `allowed_token_ids` masks the vocabulary, so the top K can hold other tokens in +place of labels. With processed logprobs the K entries are exactly the labels. A label +missing from them raises `LabelTokenError`. Prefix caching makes the K permutations of +one state cheap. + + vllm serve Qwen/Qwen3-8B --enable-prefix-caching --logprobs-mode processed_logprobs Decider(VLLMBackend("http://localhost:8000", "Qwen/Qwen3-8B")) **L2** -- an embed server whose pooler is told to return the last position's hidden @@ -36,6 +39,8 @@ import numpy as np +from anyjev.readout import LabelTokenError + class VLLMBackend: def __init__(self, base_url: str, model: str, tokenizer_name: str | None = None, @@ -73,7 +78,13 @@ def _one(self, prompt: str, ids: Sequence[int]) -> np.ndarray: if key in top: by_id[tid] = float(top[key]) break - return np.array([by_id.get(tid, -30.0) for tid in ids], dtype=np.float64) + missing = [tid for tid in ids if tid not in by_id] + if missing: + raise LabelTokenError( + f"label token ids {missing} are not in the server's top_logprobs {sorted(top)}; " + "start vLLM with --logprobs-mode processed_logprobs so that logprobs are taken " + "after allowed_token_ids") + return np.array([by_id[tid] for tid in ids], dtype=np.float64) def next_token_logprobs(self, prompts: Sequence[str], token_ids: Sequence[Sequence[int]]) -> List[np.ndarray]: diff --git a/tests/test_truncate_and_vllm.py b/tests/test_truncate_and_vllm.py index 1892fc1..91e6eb4 100644 --- a/tests/test_truncate_and_vllm.py +++ b/tests/test_truncate_and_vllm.py @@ -133,3 +133,69 @@ def test_the_config_comes_from_the_source_not_the_served_alias(stub, monkeypatch be = _backend(stub, monkeypatch) assert be.name == "served-alias" assert be.source == "some/real-model" + + +class _Completions(BaseHTTPRequestHandler): + """Answers /v1/completions with a fixed `top_logprobs` entry, keyed by token string.""" + + top: dict = {} + + def do_POST(self): # noqa: N802 - http.server's name + self.rfile.read(int(self.headers["Content-Length"])) + payload = json.dumps({"choices": [{"logprobs": {"top_logprobs": [self.top]}}]}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *a): + pass + + +class _IdTok: + """Renders token id i as "", both ways the backend looks a label up.""" + + @staticmethod + def decode(ids): + return f"<{ids[0]}>" + + @staticmethod + def convert_ids_to_tokens(i): + return f"<{i}>" + + +@pytest.fixture() +def completions(monkeypatch): + def serve(top): + handler = type("H", (_Completions,), {"top": top}) + srv = HTTPServer(("127.0.0.1", 0), handler) + threading.Thread(target=srv.serve_forever, daemon=True).start() + servers.append(srv) + be = _backend(f"http://127.0.0.1:{srv.server_port}", monkeypatch) + be.tokenizer = _IdTok() + return be + + servers = [] + yield serve + for srv in servers: + srv.shutdown() + + +def test_next_token_logprobs_reads_each_label_by_id(completions): + pytest.importorskip("transformers") + be = completions({"<2>": -2.5, "<1>": -0.1}) + (row,) = be.next_token_logprobs(["p"], [[1, 2]]) + assert np.allclose(row, [-0.1, -2.5]) + + +def test_a_label_missing_from_top_logprobs_is_an_error(completions): + """vLLM's default `--logprobs-mode raw_logprobs` computes logprobs before `allowed_token_ids` + masks the vocabulary, so the top-K can hold a non-label token in place of a label. A label + filled in with a constant would be a confident wrong answer with no symptom.""" + pytest.importorskip("transformers") + from anyjev.readout import LabelTokenError + + be = completions({"<1>": -0.1, "<9>": -1.0}) # label 2 fell out, token 9 took its place + with pytest.raises(LabelTokenError, match=r"\[2\].*processed_logprobs"): + be.next_token_logprobs(["p"], [[1, 2]])