| """Letter readout of google/gemma-4-31B-it served by vLLM in-process, with the same prompts and options as jevlab.reader. |
| |
| Same token ids as Reader: the rows are built by jevlab.reader.Layout (decider or header layout, labels A..J, single-token |
| wide labels above 10 options), and the read policy (read_twice, twice_max_tokens, reread_below, reverse_below) follows |
| Reader.__call__ step for step. tools/vllm_check.py checks both on CPU against Reader with a mocked vLLM. |
| |
| What differs is the readout. Each row is one vLLM request with max_tokens=1 and `logprob_token_ids` set to the labels of |
| its options (stock vLLM SamplingParams, at most 128 ids per request, so wider questions are read in chunks of 128 with |
| the same prompt). vLLM returns log-probabilities of the logits after Gemma's final soft cap (30 * tanh(x / 30)), not the |
| raw logits Reader reads, so Reader's T(n) = 10.124 - 1.633 ln(n) does not carry over. The probabilities are a softmax of |
| the label log-probabilities at a temperature per question type: |
| |
| choice: T = temperature_choice + temperature_choice_log_slope * ln(n options) |
| noul: T = temperature_noul, and noul_bias is added to the "yes" score after dividing by T |
| |
| The defaults (choice 4.263, noul 7.609, no slope, no bias) are wcole3/jev-brain's published fit for this model in the |
| header layout on vLLM (MIT, data/gemma-4-31b/header/calibration.json, fitted outside the suite). They are a starting |
| point; refit them on our own off-suite data. temperature_scheme="options" applies Reader's T(n) instead (for checks). |
| |
| Why in-process (vllm.LLM with prompt token ids) and not an OpenAI-compatible server: the kit runs one request at a time |
| in one process, as the board measures it; the engine loads like Reader (`--engine jevlab.vllm_reader:VllmReader`), with |
| no second process, port or HTTP client to manage on a shared machine, and token ids go to vLLM unchanged (checked on |
| every output). vLLM's own engine core still runs in a child process; latency is the wall time of the generate calls. |
| |
| All rows of one read pass of a request go to vLLM in one generate call. With prefix caching, the rows share the KV blocks |
| of their common prefix (chat head, question list, state); the re-read and reversed passes hit the same cache. Rows above |
| the token limit (max_model_len - 1) are refused as Unsupported, never truncated. max_model_len defaults to 131,072 (the |
| board's maintainers lift context caps to that or more for their private runs; the model's window is 262,144): with FP8 |
| weights a request of that length needs about 23 GiB of KV cache, which the default gpu_memory_utilization=0.70 holds. |
| |
| python -m decision_index run --engine jevlab.vllm_reader:VllmReader --option layout=header \ |
| --rows sample-2000.jsonl.gz --out runs/vllm-header-sample2000 |
| """ |
| import math |
| import os |
|
|
| from decision_index.engines.base import Engine, Unsupported |
|
|
| from jevlab.reader import LAYOUTS, Layout, _flag, render_question |
|
|
| REPO, REVISION = "google/gemma-4-31B-it", "842da3794eaa0b77d5f08bae87a17459d91ff475" |
| MAX_LABEL_IDS = 128 |
| TEMPERATURE_SCHEMES = ("type", "options") |
| |
| HF_ONLY_OPTIONS = ("device", "chunk_tokens", "cache_prefix", "cudnn_sdp", "pack_questions", "pack_tokens", "attn_block") |
| LENGTH_MARKERS = ("maximum model length", "maximum context length", "longer than the maximum") |
|
|
|
|
| def softmax(scores): |
| top = max(scores) |
| weights = [math.exp(x - top) for x in scores] |
| total = sum(weights) |
| return [w / total for w in weights] |
|
|
|
|
| class VllmBackend: |
| """Log-probabilities of given token ids at the position after each prompt, from one vLLM decoding step. |
| `llm` and `sampling_params` can be stand-ins with the same interface (tools/vllm_check.py).""" |
|
|
| def __init__(self, llm=None, sampling_params=None, **llm_kwargs): |
| if llm is None: |
| from vllm import LLM |
|
|
| llm = LLM(**llm_kwargs) |
| if sampling_params is None: |
| from vllm import SamplingParams as sampling_params |
| self.llm, self.SamplingParams = llm, sampling_params |
|
|
| def read(self, prompts, label_ids): |
| """-> (one list of log-probabilities per prompt, in the order of its label_ids; prompt tokens served from the |
| prefix cache per prompt). Label ids go in chunks of 128; each chunk is sent with the prompts that read one.""" |
| found = [{} for _ in prompts] |
| cached = [None] * len(prompts) |
| width = max(map(len, label_ids), default=0) |
| for start in range(0, width, MAX_LABEL_IDS): |
| rows = [k for k, ids in enumerate(label_ids) if len(ids) > start] |
| params = [self.SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=label_ids[k][start:start + MAX_LABEL_IDS], detokenize=False) |
| for k in rows] |
| outs = self.llm.generate([{"prompt_token_ids": prompts[k]} for k in rows], params, use_tqdm=False) |
| if len(outs) != len(rows): |
| raise RuntimeError(f"vLLM returned {len(outs)} outputs for {len(rows)} prompts") |
| for k, o in zip(rows, outs): |
| got = o.prompt_token_ids if isinstance(o.prompt_token_ids, list) else list(o.prompt_token_ids) |
| if got != prompts[k]: |
| raise RuntimeError("vLLM read other token ids than the ones sent (order or added tokens)") |
| step = o.outputs[0].logprobs[0] |
| for t in label_ids[k][start:start + MAX_LABEL_IDS]: |
| found[k][t] = float(step[t].logprob) |
| if cached[k] is None: |
| cached[k] = int(o.num_cached_tokens or 0) |
| return [[found[k][t] for t in ids] for k, ids in enumerate(label_ids)], cached |
|
|
| def close(self): |
| engine = getattr(self.llm, "llm_engine", None) |
| core = getattr(engine, "engine_core", None) |
| if core is not None and hasattr(core, "shutdown"): |
| core.shutdown() |
|
|
|
|
| class VllmReader(Engine): |
| name = "jevlab-vllm-reader" |
| latency = ("In-process request wall time including prompt construction: one vLLM generate call per read pass of a request " |
| "(all its rows together, shared prefix from vLLM's prefix cache); excludes model loading.") |
|
|
| def __init__(self, model=REPO, revision=REVISION, quantization="fp8_per_block", dtype="bfloat16", max_model_len=131072, max_tokens=None, |
| gpu_memory_utilization=0.70, kv_cache_dtype="auto", attention_backend="TRITON_ATTN", prefix_caching=True, |
| prefix_cache_retention_interval=32, max_num_batched_tokens=8192, max_num_seqs=256, enforce_eager=False, language_model_only=True, |
| deep_gemm=False, seed=0, llm_options=None, layout="decider", read_twice=False, twice_max_tokens=0, reread_below=0.0, reverse_below=0.0, |
| temperature_scheme="type", temperature_choice=4.263, temperature_choice_log_slope=0.0, temperature_noul=7.609, noul_bias=0.0, |
| temperature=10.124, temperature_log_slope=-1.633, backend=None, tokenizer=None, system=None, **options): |
| ignored = {k: options.pop(k) for k in HF_ONLY_OPTIONS if k in options} |
| if options: |
| raise ValueError(f"unknown VllmReader options {sorted(options)}") |
| super().__init__(**ignored) |
| if layout not in LAYOUTS: |
| raise ValueError(f"layout must be one of {LAYOUTS}, not {layout!r}") |
| if temperature_scheme not in TEMPERATURE_SCHEMES: |
| raise ValueError(f"temperature_scheme must be one of {TEMPERATURE_SCHEMES}, not {temperature_scheme!r}") |
| if isinstance(quantization, str) and quantization.strip().lower() in ("", "none", "null", "bf16", "bfloat16"): |
| quantization = None |
| if isinstance(attention_backend, str) and attention_backend.strip().lower() in ("", "auto", "none", "null"): |
| attention_backend = None |
| self.max_model_len = int(max_model_len) |
| self.limit = int(max_tokens) if max_tokens else self.max_model_len - 1 |
| if not 0 < self.limit < self.max_model_len: |
| raise ValueError(f"max_tokens must be below max_model_len ({self.max_model_len}): one position is the answer slot") |
| self.header = layout == "header" |
| self.read_twice = _flag(read_twice) |
| self.twice_max_tokens, self.reread_below = int(twice_max_tokens), float(reread_below) |
| self.reverse_below = float(reverse_below) |
| self.scheme = temperature_scheme |
| self.t_choice, self.t_choice_slope = float(temperature_choice), float(temperature_choice_log_slope) |
| self.t_noul, self.noul_bias = float(temperature_noul), float(noul_bias) |
| self.temperature, self.temperature_log_slope = float(temperature), float(temperature_log_slope) |
| self.model_id = model |
| if tokenizer is None: |
| from transformers import AutoTokenizer |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model, revision=revision) |
| self.tok = tokenizer |
| self.system = system or None |
| self.layout = Layout(tokenizer, system=self.system) |
| self.llm_kwargs = dict(model=model, revision=revision, tokenizer_revision=revision, dtype=dtype, quantization=quantization, |
| max_model_len=self.max_model_len, gpu_memory_utilization=float(gpu_memory_utilization), kv_cache_dtype=kv_cache_dtype, |
| enable_prefix_caching=_flag(prefix_caching), prefix_cache_retention_interval=prefix_cache_retention_interval, |
| max_num_batched_tokens=int(max_num_batched_tokens), max_num_seqs=int(max_num_seqs), enforce_eager=_flag(enforce_eager), |
| language_model_only=_flag(language_model_only), attention_backend=attention_backend, seed=int(seed), |
| logprobs_mode="raw_logprobs", generation_config="vllm", **dict(llm_options or {})) |
| if backend is None and not _flag(deep_gemm): |
| |
| os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") |
| self.backend = backend if backend is not None else VllmBackend(**self.llm_kwargs) |
| temps = ({"choice": f"{self.t_choice} + {self.t_choice_slope} * ln(n options)", "noul": self.t_noul, "noul_bias": self.noul_bias} |
| if self.scheme == "type" else {"all": f"max({self.temperature} + {self.temperature_log_slope} * ln(n options), 0.05)"}) |
| self.provenance = { |
| "kind": "inference technique", |
| "repo": model, |
| "revision": revision, |
| "serving": "vLLM offline LLM, in-process, prompt token ids", |
| "vllm_options": {k: v for k, v in self.llm_kwargs.items() if k not in ("model", "revision", "tokenizer_revision")}, |
| "VLLM_USE_DEEP_GEMM": os.environ.get("VLLM_USE_DEEP_GEMM"), |
| "max_row_tokens": self.limit, |
| "read_twice": self.read_twice, |
| "twice_max_tokens": self.twice_max_tokens, |
| "reread_below": self.reread_below, |
| "reverse_below": self.reverse_below, |
| "layout": layout, |
| **({"system": self.system} if self.system else {}), |
| "temperature_scheme": self.scheme, |
| "temperatures": temps, |
| "ignored_options": sorted(ignored), |
| "policy": ("The state is preceded by the list of every question of the request ('Questions that will be asked about the context:'), " |
| "the header layout of wcole3/jev-brain. " if self.header else "") |
| + (f"A fixed system turn comes first: {self.system!r}. " if self.system else "") |
| + "One row per question in the model's chat template with thinking off, ending at 'Answer: ('; vLLM reads one step and returns " |
| "the log-probabilities (after the final soft cap) of the option labels (logprob_token_ids), which go through a softmax at a " |
| + ("temperature per question type. " if self.scheme == "type" else "temperature set by the number of options. ") |
| + "Options keep the request's order. Rows above the token limit are refused, never cut." |
| + (" The user turn (state, question, options) is written twice, joined by 'Read the same input again before answering:'" |
| + (f", in requests of at most {self.twice_max_tokens} tokens read once" if self.twice_max_tokens > 0 else "") + "." if self.read_twice else "") |
| + (f" Questions of the other requests whose most likely option has probability below {self.reread_below} are read again, written twice " |
| "the same way, after the first read." if self.reread_below > 0 else "") |
| + (f" A question whose most likely option then has probability below {self.reverse_below} is read once more with its options in " |
| "reverse order, and the two distributions are averaged." if self.reverse_below > 0 else ""), |
| } |
|
|
| def runtime(self): |
| info = {} |
| try: |
| import torch |
| import vllm |
|
|
| info.update(vllm=vllm.__version__, torch=torch.__version__, cuda=torch.version.cuda) |
| from vllm.platforms import current_platform |
|
|
| info["gpu"] = current_platform.get_device_name(0) |
| except Exception as exc: |
| info.setdefault("note", f"runtime details unavailable: {type(exc).__name__}") |
| try: |
| import transformers |
|
|
| info["transformers"] = transformers.__version__ |
| except ImportError: |
| pass |
| return info |
|
|
| def close(self): |
| if hasattr(self.backend, "close"): |
| self.backend.close() |
|
|
| def _read(self, ctx, pieces, ns): |
| """-> [(label log-probabilities, prompt tokens served from the prefix cache)] for the rows ctx + piece.""" |
| try: |
| scores, cached = self.backend.read([ctx + piece for piece in pieces], [self.layout.label_ids[:n] for n in ns]) |
| except ValueError as exc: |
| if any(m in str(exc) for m in LENGTH_MARKERS): |
| raise Unsupported(f"maximum context length: {exc}") from exc |
| raise |
| return list(zip(scores, cached)) |
|
|
| def _probs(self, scores, kind, reverse=False): |
| """Softmax over the shown options; with reverse the options are in reverse order (yes first for noul).""" |
| n = len(scores) |
| if self.scheme == "options": |
| t, bias = max(self.temperature + self.temperature_log_slope * math.log(n), 0.05), 0.0 |
| elif kind == "noul": |
| t, bias = self.t_noul, self.noul_bias |
| else: |
| t, bias = max(self.t_choice + self.t_choice_slope * math.log(n), 0.05), 0.0 |
| x = [s / t for s in scores] |
| if bias: |
| x[0 if reverse else 1] += bias |
| return softmax(x) |
|
|
| def _check(self, ctx, piece): |
| if len(ctx) + len(piece) > self.limit: |
| raise Unsupported(f"maximum context length: a row of {len(ctx) + len(piece)} tokens exceeds the declared limit of {self.limit}") |
|
|
| def __call__(self, state, questions): |
| |
| rendered = [(key, q["type"], *render_question(q)) for key, q in questions.items()] |
| texts = [text for _, _, text, _, _ in rendered] if self.header else None |
| ctx = self.layout.context(state, texts) |
| again = self.layout.again(state, texts) if self.read_twice or self.reread_below > 0 else None |
| rows = [(key, kind, self.layout.question(text, options), self.layout.question(text, options, again) if again else None, keys) |
| for key, kind, text, options, keys in rendered] |
| |
| fits = [again is not None and len(ctx) + len(twice) <= self.limit for _, _, _, twice, _ in rows] |
| once_tokens = len(ctx) + sum(len(once) for _, _, once, _, _ in rows) |
| whole = self.read_twice and all(fits) and (self.twice_max_tokens <= 0 or once_tokens <= self.twice_max_tokens) |
| first = [twice if whole else once for _, _, once, twice, _ in rows] |
| for piece in first: |
| self._check(ctx, piece) |
| ns = [len(keys) for *_, keys in rows] |
| kinds = [kind for _, kind, *_ in rows] |
| read = self._read(ctx, first, ns) |
| pieces = list(first) |
| reread = self.reread_below > 0 and not whole |
| unsure = [i for i, (scores, _) in enumerate(read) if reread and fits[i] and max(self._probs(scores, kinds[i])) < self.reread_below] |
| firsts = {} |
| if unsure: |
| for i, again_read in zip(unsure, self._read(ctx, [rows[i][3] for i in unsure], [ns[i] for i in unsure])): |
| firsts[i] = read[i][0] |
| read[i] = again_read |
| pieces[i] = rows[i][3] |
| probs = [self._probs(scores, kind) for (scores, _), kind in zip(read, kinds)] |
| flipped = {} |
| if self.reverse_below > 0: |
| low = [i for i, pr in enumerate(probs) if max(pr) < self.reverse_below] |
| if low: |
| twice = [whole or i in firsts for i in low] |
| rev = [self.layout.question(rendered[i][2], rendered[i][3][::-1], again if t else None) for i, t in zip(low, twice)] |
| for piece in rev: |
| self._check(ctx, piece) |
| for i, (scores, _) in zip(low, self._read(ctx, rev, [ns[i] for i in low])): |
| back = self._probs(scores, kinds[i], reverse=True)[::-1] |
| flipped[i] = scores |
| probs[i] = [(a + b) / 2 for a, b in zip(probs[i], back)] |
| answers, raw = {}, {} |
| for i, ((key, kind, _, _, keys), (scores, cached)) in enumerate(zip(rows, read)): |
| raw[key] = {"label_logprobs": [round(x, 6) for x in scores], "prompt_tokens": len(ctx) + len(pieces[i]), "cached_prefix_tokens": cached, |
| "read_twice": whole or i in firsts} |
| if i in firsts: |
| raw[key]["first_label_logprobs"] = [round(x, 6) for x in firsts[i]] |
| if i in flipped: |
| raw[key]["reversed_label_logprobs"] = [round(x, 6) for x in flipped[i]] |
| probs_i = probs[i] |
| if kind == "noul": |
| answers[key] = {"type": "noul", "noul": probs_i[1]} |
| else: |
| answers[key] = {"type": "choice", "choice": keys[max(range(len(keys)), key=probs_i.__getitem__)], "probabilities": dict(zip(keys, probs_i))} |
| return {"model": self.model_id, "answers": answers, "usage": {"input_tokens": sum(v["prompt_tokens"] for v in raw.values())}}, raw |
|
|
