vLLM backend: raise on a label missing from top_logprobs instead of filling -30

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]:

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 "<i>", 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]])