serve_decide.py: answer read-out that works with speculative decoding (+ fallback), system2_only; serve.sh: --max-logprobs 256, MTP=1 option

README.md

@@ -172,7 +172,8 @@ Lower `threshold` or `think_budget` to trade accuracy for speed; `thinking: "off
 
 MMLU-Pro and BBH (except six questions) ran with speculative decoding, which does not change the output distribution.
 Six BBH questions with 18 options hit a vLLM error in the speculative-decoding read-out path and were answered by the same
-server without speculative decoding.
+server without speculative decoding; the current `serve_decide.py` reads answers through a path that works with
+speculative decoding.
 
 ## Images
 
@@ -289,13 +290,13 @@ requests.post("http://localhost:8000/v1/chat/completions", json={
 * Adaptive thinking adds nothing when it does not think, and about 1 second per 250 thinking tokens when it does
   (see [Latency](#latency)).
 
-**Faster thinking with speculative decoding.** Adding
-`--speculative-config '{"model": "google/gemma-4-26B-A4B-it-assistant", "num_speculative_tokens": 4}'` to `serve.sh`
-(Google's 0.9 GB draft model for this backbone) speeds up System 2 about 1.8–1.9× on one B200: 247 → 438 tokens per
-second for a single request, 7,072 → 13,505 tokens per second at 128 concurrent requests (mean acceptance length 3.5–3.7
-of 4). System 1 decisions are unchanged, but System 1 throughput at high concurrency drops (257 → 140 decisions per second
-with 64 clients), and in this vLLM build the System 2 answer read-out fails for questions with more than 16 options. Use it
-when you mostly think, and serve System 1-heavy traffic without it.
+**Faster thinking with speculative decoding.** `MTP=1 bash GEV-26B-Decide/serve.sh` adds Google's 0.9 GB draft model for
+this backbone (`--speculative-config '{"model": "google/gemma-4-26B-A4B-it-assistant", "num_speculative_tokens": 4}'`).
+On one B200 it speeds up System 2 about 1.8–1.9×: 247 → 438 tokens per second for a single request, 7,072 → 13,505 tokens
+per second at 128 concurrent requests (mean acceptance length 3.5–3.7 of 4). System 1 decisions are unchanged, but System 1
+throughput at high concurrency drops (257 → 140 decisions per second with 64 clients); use it when you mostly think.
+`serve_decide.py` reads answers with the token restriction that works under speculative decoding, which needs
+`--max-logprobs 256` (set in `serve.sh`).
 
 If you call `/v1/completions` for System 1 yourself, pass `top_k: 0` and `top_p: 1.0`: the model's generation config sets
 `top_k=64` and `top_p=0.95`, which vLLM applies as request defaults and which would truncate the returned probabilities.

serve.sh

@@ -1,16 +1,22 @@
 #!/usr/bin/env bash
 # Serve autotrust/GEV-26B-Decide with vLLM: one engine, both systems, text and images.
 #   System 2: the unmodified gemma-4-26B-A4B-it (served name autotrust/GEV-26B-Decide) on the OpenAI endpoints
-#   System 1: POST /v1/decide (2-256 options, optional adaptive thinking), or the LoRA module "jev-decision" directly
+#   System 1: POST /v1/decide (2-256 options, adaptive thinking), or the LoRA module "jev-decision" directly
 # serve_decide.py is the standard vLLM OpenAI server (same flags) with /v1/decide added.
 # Requirements: a vLLM build with Gemma-4 support plus patches/vllm-gemma4-lm-head-lora.patch (LoRA on Gemma-4's tied
 # lm_head, vocabulary 262,144), tested with a vLLM development build from September 2026.
 # MAX_MODEL_LEN: up to 262144 (the backbone's native context).
+# MTP=1: speculative decoding with google/gemma-4-26B-A4B-it-assistant (System 2 about 1.8-1.9x faster; lower System 1
+#        throughput at high concurrency).
 set -e
 MODEL_DIR=${MODEL_DIR:-GEV-26B-Decide}
 [ -d "$MODEL_DIR" ] || hf download autotrust/GEV-26B-Decide --local-dir "$MODEL_DIR"
+EXTRA=()
+if [ "${MTP:-0}" = "1" ]; then
+  EXTRA=(--speculative-config '{"model": "google/gemma-4-26B-A4B-it-assistant", "num_speculative_tokens": 4}')
+fi
 exec python3 "$MODEL_DIR/serve_decide.py" --model "$MODEL_DIR" --served-model-name autotrust/GEV-26B-Decide \
   --enable-lora --max-lora-rank 32 --lora-modules jev-decision="$MODEL_DIR/adapter_vllm" \
-  --logprobs-mode processed_logprobs --max-model-len ${MAX_MODEL_LEN:-65536} --enable-prefix-caching \
+  --logprobs-mode processed_logprobs --max-logprobs 256 --max-model-len ${MAX_MODEL_LEN:-65536} --enable-prefix-caching \
   --max-num-seqs 256 --trust-request-chat-template --limit-mm-per-prompt '{"image": 8}' \
-  --port ${PORT:-8000}
+  "${EXTRA[@]}" --port ${PORT:-8000}

serve_decide.py

@@ -86,6 +86,7 @@ class DecideRequest(BaseModel):
     chat_template_kwargs: dict | None = None
     return_reasoning: bool = False
     debug: bool = False
+    system2_only: bool = False   # skip System 1: think and return System 2's answer distribution (for callers that already have p1)
 
 
 # ------------------------------------------------------------------ setup
@@ -137,6 +138,7 @@ def setup(args):
            "{%- endif -%}{%- endfor -%}{%- endif -%}{%- endfor -%}")
     names = args.served_model_name
     S.update(tok=tok, has_effort="reasoning_effort" in (tok.chat_template or ""), think_cache={},
+             max_logprobs=int(getattr(args, "max_logprobs", 20) or 20),
              head=head, temps=temps, labels=[l for l, _ in dec], label_ids=[t for _, t in dec], ans_labels=ans, prof=prof,
              think_pre=pre, think_post=post, raw=raw, end_think_id=tok.convert_tokens_to_ids(prof["end_think"]),
              model=(names[0] if isinstance(names, list) else names) or args.model, model_type=mtype)
@@ -194,11 +196,32 @@ class Ctx:
             self.completion_tokens += out.usage.completion_tokens or 0
 
 
+def _readout_allowed(ids):
+    """Read-out path: restrict the next token to `ids` and return all of them as top logprobs when the server allows that
+    many (--max-logprobs); otherwise ask for `ids` explicitly with logprob_token_ids (chunks of 128). The first path also
+    works with speculative decoding; logprob_token_ids does not in this vLLM build. Both give the same distribution."""
+    return ids if len(ids) <= S["max_logprobs"] else None
+
+
 async def _logprobs(ctx, content, has_img, ids, model, allowed=None):
+    """One-token read-out with a fallback to the other read-out path if the first one fails."""
+    try:
+        return await _logprobs_once(ctx, content, has_img, ids, model, allowed)
+    except Upstream:
+        raise
+    except Exception:
+        other = None if allowed is not None else (ids if len(ids) <= S["max_logprobs"] else None)
+        if other is None and allowed is None:
+            raise
+        return await _logprobs_once(ctx, content, has_img, ids, model, other)
+
+
+async def _logprobs_once(ctx, content, has_img, ids, model, allowed=None):
     """One-token read-out: {token_id: logprob} for `ids` after the content (text or chat parts)."""
     lp = {}
-    for i in range(0, len(ids), CHUNK):
-        chunk = ids[i: i + CHUNK]
+    step = CHUNK if allowed is None else len(ids)
+    for i in range(0, len(ids), step):
+        chunk = ids[i: i + step]
         kw = dict(logprob_token_ids=chunk) if allowed is None else dict(allowed_token_ids=allowed)
         if has_img:
             r = ChatCompletionRequest(model=model, messages=[{"role": "user", "content": content}], chat_template=S["raw"],
@@ -228,7 +251,7 @@ async def s1_pass(ctx, kind, parts, has_img, question, opts):
         ids = head["verbalizer_ids"][lo:hi]; bias = head["bias"][lo:hi]; lines = opts
     content = ([{"type": "text", "text": f"{prof['prefix']}[kind] {kind}\n[state] "}] + parts +
                [{"type": "text", "text": f"\n[question] {question}\n[options]\n" + "\n".join(lines) + "\n[decision]:"}])
-    lp = await _logprobs(ctx, content, has_img, ids, LORA, allowed=ids if len(ids) <= 16 else None)
+    lp = await _logprobs(ctx, content, has_img, ids, LORA, allowed=_readout_allowed(ids))
     return _softmax([(max(lp.get(t, -1e9), -1e9) + b) / temps[kind] for t, b in zip(ids, bias)])
 
 
@@ -302,7 +325,7 @@ async def s2_dist(ctx, kind, parts, has_img, question, opts, budget, want_text,
     ntok = out.usage.completion_tokens if out.usage else None
     read = user + [{"type": "text", "text": thought + S["prof"]["end_think"] + S["prof"]["after_think"] + "Answer: ("}]
     ids = [t for _, t in labs]
-    lp = await _logprobs(ctx, read, True, ids, S["model"], allowed=ids if len(ids) <= 16 else None)
+    lp = await _logprobs(ctx, read, True, ids, S["model"], allowed=_readout_allowed(ids))
     p = _softmax([max(lp.get(t, -1e9), -1e9) for t in ids])
     if kind == "noul":
         p = [p[1], p[0]]  # (A) yes / (B) no -> [P(false), P(true)]
@@ -356,6 +379,15 @@ async def decide(req: DecideRequest, raw: Request):
     parts, has_img = _parts(req.state)
     ctx = Ctx(raw)
     try:
+        if req.system2_only:
+            if req.kind == "score":
+                return _err("system2_only is supported for noul and choice")
+            p2, think = await s2_dist(ctx, req.kind, parts, has_img, req.question, opts, req.think_budget, req.return_reasoning, tmpl)
+            k = max(range(len(p2)), key=p2.__getitem__)
+            return {"kind": req.kind, "options": opts, "probabilities": p2, "choice_index": k, "choice": opts[k], "system": 2,
+                    "model": S["model"], "thinking": {"used": True, "budget": req.think_budget, **think},
+                    "usage": {"prompt_tokens": ctx.prompt_tokens, "completion_tokens": ctx.completion_tokens},
+                    "elapsed_seconds": time.time() - t0, "num_model_requests": ctx.requests}
         p1 = await s1_dist(ctx, req.kind, parts, has_img, req.question, opts, strategy)
         probs, think = p1, None
         if thinking == "on" or (thinking == "auto" and max(p1) < tau):