| """Decision Index engine, vLLM path, one instance: direct slot pass for every question, learned router, stock-model rationale only for the escalated questions, adapted slot pass. |
| |
| Same decisions as decision_engine_v6r:RoutedSlotEngine (Hugging Face) but served by vLLM (CUDA graphs on, prefix cache on). vLLM does not return hidden states, so the model runner's |
| `sample(hidden_states, input_batch, ...)` is wrapped (single-process engine core, VLLM_ENABLE_V1_MULTIPROCESSING=0): for every request the final-norm hidden states at the wanted slot |
| positions are kept, and the answer-slot head (frozen LM-head rows of the answer tokens + trained residual, head.pt) and the router head (CKPT/router-direct-brief.npz) run on them in torch. |
| The rationale is generated by the stock model (adapter off, greedy, at most `brief_cap` tokens) and cleaned like the training rationales (dense/contexts.py). |
| |
| VLLM_ENABLE_V1_MULTIPROCESSING=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 /mnt/datasets/venv-vllm/bin/python -m decision_index run --engine decision_engine_vr:VllmRoutedEngine \ |
| --option model=Qwen/Qwen3.5-9B --option run=RUN --option ckpt=final [--option mode=routed|direct|brief] [--option tau=0.05] --suite-dir ... --rows ... --out ... |
| Options: model, run, ckpt, mode (routed | direct | brief), tau (override; inf = never escalate), brief_cap (64), max_tokens (prompt limit, 32768), gpu_util (0.88), max_loras lora (converted adapter dir; made next to the checkpoint when absent). |
| Per-option (parallel slot) questions are always answered by the direct pass (v6b recipe). Temperature: none. |
| """ |
| import json |
| import os |
| import shutil |
| import sys |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, HERE) |
| os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| import answer_slot as AS |
| from common import QTYPES |
| try: |
| from decision_index.engines.base import Engine, Unsupported |
| except ImportError: |
| from engine_base import Engine, Unsupported |
| from dense.contexts import Contexts |
|
|
|
|
| def _lg(p): |
| p = np.clip(p, 1e-4, 1 - 1e-4) |
| return np.log(p / (1 - p)) |
|
|
|
|
| class VllmRoutedEngine(Engine): |
| name = "dense-slot-routed-vllm" |
| latency = ("In-process request wall time (vLLM, one instance, CUDA graphs): the direct slot pass of all questions of the request in one batch, the router head, and for the escalated " |
| "questions a batched stock-model rationale (adapter off, at most 64 tokens) followed by the adapted slot pass. Excludes model loading.") |
|
|
| def __init__(self, model="Qwen/Qwen3.5-9B", run=None, ckpt="final", mode="routed", tau=None, brief_cap=64, max_tokens=32768, gpu_util=0.88, lora=None, **options): |
| super().__init__(**options) |
| if not run: |
| raise ValueError("VllmRoutedEngine needs a trained run (option run=...)") |
| from transformers import AutoTokenizer |
| from vllm import LLM, SamplingParams |
| from vllm.lora.request import LoRARequest |
| self.mode = str(mode) |
| if self.mode not in ("routed", "direct", "brief"): |
| raise ValueError("mode must be routed, direct or brief") |
| ck = ckpt if os.path.isdir(ckpt) else os.path.join(run, ckpt) |
| self.tok = AutoTokenizer.from_pretrained(model) |
| if self.tok.pad_token_id is None: |
| self.tok.pad_token_id = self.tok.eos_token_id |
| self.codec = AS.SlotCodec(self.tok, state_max=None, option_max=None, template_kwargs={"enable_thinking": False}, mask_id=self.tok.pad_token_id or 0, strip_label_keys=True, autoregressive=True) |
| self.brief_cap = int(brief_cap) |
| self.ctx = Contexts(self.tok, self.codec, brief_cap=self.brief_cap) |
| self.eos = self.tok.convert_tokens_to_ids("<|im_end|>") |
| self.model_id = model |
| lora_dir = lora or os.path.join(ck, "lora-vllm") |
| if not os.path.isfile(os.path.join(lora_dir, "adapter_model.safetensors")): |
| self._convert_lora(os.path.join(ck, "lora"), lora_dir) |
| self.limit = int(max_tokens) |
| self.llm = LLM(model, dtype="bfloat16", max_model_len=self.limit + 256, max_num_batched_tokens=self.limit + 256, gpu_memory_utilization=float(gpu_util), enable_prefix_caching=True, |
| enable_lora=True, max_lora_rank=32, max_loras=1) |
| self.eng = self.llm.llm_engine |
| core = self.eng.engine_core |
| core = core.engine_core if hasattr(core, "engine_core") else core |
| self.runner = core.model_executor.driver_worker.worker.model_runner |
| self.lora = LoRARequest("slot", 1, lora_dir) |
| self.sp_slot = SamplingParams(temperature=0, max_tokens=1) |
| self.sp_slot_nc = SamplingParams(temperature=0, max_tokens=1, skip_reading_prefix_cache=True) |
| self.sp_gen = SamplingParams(temperature=0, max_tokens=self.brief_cap, stop_token_ids=[self.eos]) |
| self.cap, self.n_req = {}, 0 |
| orig = self.runner.sample |
| cap = self.cap |
|
|
| def wrapped(hidden_states, input_batch, *a, **k): |
| for b, rid in enumerate(input_batch.req_ids): |
| w = cap.get(rid.rsplit("-", 1)[0]) |
| if w is None: |
| continue |
| start, comp, n = int(input_batch.query_start_loc_np[b]), int(input_batch.num_computed_tokens_np[b]), int(input_batch.num_scheduled_tokens[b]) |
| for p in w["pos"]: |
| if comp <= p < comp + n and p not in w["got"]: |
| w["got"][p] = hidden_states[start + p - comp].detach().clone() |
| return orig(hidden_states, input_batch, *a, **k) |
| self.runner.sample = wrapped |
| |
| m = self.runner.get_model() if hasattr(self.runner, "get_model") else self.runner.model |
| w = next(p for n, p in m.named_parameters() if n.endswith("lm_head.weight")) if any(n.endswith("lm_head.weight") for n, _ in m.named_parameters()) else m.get_input_embeddings().weight |
| self.cand_w = w[torch.tensor(self.codec.cand, device=w.device)].detach().clone() |
| self.softcap = getattr(m.config.get_text_config(), "final_logit_softcapping", None) if hasattr(m, "config") else None |
| self.head = AS.SlotResidual(self.cand_w.size(1), len(self.codec.cand)).to("cuda").eval() |
| self.head.load_state_dict(torch.load(os.path.join(ck, "head.pt"), map_location="cuda")) |
| |
| blob = np.load(os.path.join(ck, "router-direct-brief.npz")) |
| self.knots = blob["knots"] |
| self.tau = float(blob["tau"]) if tau in (None, "") else float(tau) |
| self.W = {p: {k: blob[f"{p}_{k}"] for k in ("mh", "sh", "mc", "sc", "w")} for p in ("w", "f")} |
| self.provenance = {"kind": "dense-answer-slot", "repo": model, "run": run, "checkpoint": ck, "mode": self.mode, "tau": self.tau, "layout": AS.LAYOUT, "engine": "vLLM, hidden states from a wrapped model-runner sample()", |
| "context_limit_tokens": self.limit, "temperature": "none", |
| "policy": "Direct answer-slot pass for every question; a learned head escalates a question to a stock-model rationale of at most %d tokens before the slot is read. Prompts over the limit are refused as unsupported, never truncated." % self.brief_cap} |
|
|
| @staticmethod |
| def _convert_lora(src, dst): |
| from safetensors.torch import load_file, save_file |
| os.makedirs(dst, exist_ok=True) |
| sd = load_file(os.path.join(src, "adapter_model.safetensors")) |
| new = {k.replace("base_model.model.model.layers.", "base_model.model.model.language_model.layers.", 1): v for k, v in sd.items()} |
| assert len(new) == len(sd) and any("language_model" in k for k in new) |
| save_file(new, os.path.join(dst, "adapter_model.safetensors"), metadata={"format": "pt"}) |
| shutil.copy(os.path.join(src, "adapter_config.json"), os.path.join(dst, "adapter_config.json")) |
|
|
| |
| def _item(self, state, q): |
| if q["type"] not in ("choice", "noul"): |
| raise Unsupported("Unsupported question type " + str(q["type"])) |
| keys = list(q["criteria"]) if q["type"] == "choice" else ["false", "true"] |
| if not keys: |
| raise Unsupported("question without options") |
| crit = q["criteria"] if q["type"] == "choice" else (q.get("criteria") or {}) |
| k = len(keys) |
| text = "" if state in (None, "", {}, []) else (state if isinstance(state, str) else json.dumps(state, ensure_ascii=False, separators=(",", ":"))) |
| return keys, {"qtype": QTYPES[q["type"]], "target": [1.0 / k] * k, "label": 0, "state": text, "source": "index", "qdef": {"type": q["type"], "instructions": q.get("instructions", ""), "criteria": crit}} |
|
|
| |
| def _add(self, ids, sp, lora=None): |
| self.n_req += 1 |
| rid = f"r{self.n_req}" |
| self.eng.add_request(rid, {"prompt_token_ids": ids}, sp, lora_request=lora) |
| return rid |
|
|
| def _drain(self, rids): |
| wait, done = set(rids), {} |
| while len(done) < len(wait): |
| for o in self.eng.step(): |
| if o.request_id in wait and o.finished: |
| done[o.request_id] = o |
| return done |
|
|
| def _states(self, jobs): |
| """jobs: list of (ids, slot positions) -> list of (S, H) bf16 hidden states at the slot positions (adapter on). A slot position that sits in the prefix cache is |
| not recomputed, so such a job is run again with the prefix-cache read skipped.""" |
| res, todo = [None] * len(jobs), list(range(len(jobs))) |
| for attempt in (0, 1): |
| rids = {} |
| for i in todo: |
| ids, pos = jobs[i] |
| self.n_req += 1 |
| rid = f"s{self.n_req}" |
| self.cap[rid] = {"pos": set(pos), "got": {}} |
| self.eng.add_request(rid, {"prompt_token_ids": ids}, self.sp_slot_nc if attempt else self.sp_slot, lora_request=self.lora) |
| rids[rid] = i |
| self._drain(rids) |
| todo = [] |
| for rid, i in rids.items(): |
| got, pos = self.cap.pop(rid)["got"], jobs[i][1] |
| if all(p in got for p in pos): |
| res[i] = torch.stack([got[p] for p in pos]) |
| else: |
| todo.append(i) |
| if not todo: |
| return res |
| raise RuntimeError("slot hidden state not captured") |
|
|
| @torch.inference_mode() |
| def _readout(self, hs, r): |
| """Option logits (K,) and pooled slot state (H,) from slot hidden states hs (S, H), as DensePipeline.readout.""" |
| dev = "cuda" |
| t = lambda x: torch.tensor([x], device=dev) |
| k = len(r["target"]) |
| with torch.amp.autocast("cuda", dtype=torch.bfloat16): |
| h = hs[None] |
| lin = F.linear(h, self.cand_w).float() |
| if self.softcap: |
| lin = self.softcap * torch.tanh(lin / self.softcap) |
| z = lin + self.head(h).float() |
| zs = z.gather(1, t(r["slot_of_opt"]).unsqueeze(-1).expand(-1, -1, z.size(-1))) |
| pos = zs.gather(2, t(r["pos_idx"]).unsqueeze(-1)).squeeze(-1) |
| neg = zs.gather(2, t(r["neg_idx"]).clamp(min=0).unsqueeze(-1)).squeeze(-1) |
| logits = (pos - neg * (t(r["neg_idx"]) >= 0)).masked_fill(~torch.ones(1, k, dtype=torch.bool, device=dev), -1e4)[0] |
| return logits.float(), hs.float().mean(0) |
|
|
| def _score(self, pooled, conf, k): |
| x = _lg(np.array([conf], dtype=np.float64)) |
| C = np.concatenate([x[:, None], np.maximum(0, x[:, None] - self.knots[None]), np.log([[float(k)]]), np.log([[float(k)]]) * x[:, None]], 1) |
| h = pooled.double().cpu().numpy()[None] |
| f = lambda M: float(np.clip((np.concatenate([(h - M["mh"]) / M["sh"], (C - M["mc"]) / M["sc"], np.ones((1, 1))], 1) @ M["w"])[0], 0, 1)) |
| return f(self.W["w"]) * f(self.W["f"]) |
|
|
| @staticmethod |
| def _slot_ids(r): |
| """Prompt ids and slot positions; a single slot is the last position, so everything after it is dropped (causal).""" |
| sp = list(r["slot_pos"]) |
| return (r["ids"][:sp[0] + 1] if len(sp) == 1 else r["ids"]), sp |
|
|
| |
| def __call__(self, state, questions): |
| qids = list(questions) |
| items = {qid: self._item(state, questions[qid]) for qid in qids} |
| direct = {} |
| for qid in qids: |
| keys, item = items[qid] |
| r = self.codec.render_ar(item) |
| if len(r["ids"]) + self.brief_cap > self.limit: |
| raise Unsupported(f"prompt of {len(r['ids'])} tokens exceeds the {self.limit}-token limit") |
| direct[qid] = r |
| jobs = [self._slot_ids(direct[qid]) for qid in qids] |
| hs = self._states(jobs) |
| answers, raw, esc, probs_of = {}, {}, [], {} |
| for qid, h in zip(qids, hs): |
| keys, item = items[qid] |
| r = direct[qid] |
| logits, pooled = self._readout(h, r) |
| p = torch.softmax(logits, -1).cpu().numpy().astype(np.float64) |
| p = p / p.sum() |
| q = questions[qid] |
| score = None |
| routable = q["type"] == "choice" and AS.slot_mode(r["qtype"], len(keys)) != "parallel" |
| go = False |
| if routable and self.mode == "brief": |
| go = True |
| elif routable and self.mode == "routed": |
| score = self._score(pooled, float(p.max()), len(keys)) |
| go = score >= self.tau |
| if go: |
| esc.append(qid) |
| raw[qid] = {"routed": "brief", "router_score": None if score is None else round(score, 5), "direct_conf": round(float(p.max()), 4)} |
| continue |
| probs_of[qid] = dict(zip(keys, (float(x) for x in p))) |
| raw[qid] = {"routed": "direct", "router_score": None if score is None else round(score, 5), "prompt_tokens": len(r["ids"]), "direct_conf": round(float(p.max()), 4)} |
| if esc: |
| gen_prompts = [self.ctx.gen_prompt(items[qid][1]) for qid in esc] |
| for qid, gp in zip(esc, gen_prompts): |
| if len(gp) + self.brief_cap > self.limit: |
| raise Unsupported(f"prompt of {len(gp)} tokens exceeds the {self.limit}-token limit") |
| rids = {self._add(gp, self.sp_gen): qid for gp, qid in zip(gen_prompts, esc)} |
| outs = self._drain(rids) |
| rendered = {} |
| for rid, qid in rids.items(): |
| g = list(outs[rid].outputs[0].token_ids) |
| if self.eos in g: |
| g = g[:g.index(self.eos)] |
| rat = self.ctx.clean_brief(g) |
| keys, item = items[qid] |
| r = self.ctx.render({**item, "rationale": {"brief": {"tokens": rat, "n": len(rat)}}}, "brief") |
| if len(r["ids"]) > self.limit: |
| raise Unsupported(f"prompt of {len(r['ids'])} tokens exceeds the {self.limit}-token limit") |
| rendered[qid] = (r, rat) |
| order = list(rendered) |
| hs2 = self._states([self._slot_ids(rendered[qid][0]) for qid in order]) |
| for qid, h in zip(order, hs2): |
| r, rat = rendered[qid] |
| keys = items[qid][0] |
| logits, _ = self._readout(h, r) |
| p = torch.softmax(logits, -1).cpu().numpy().astype(np.float64) |
| p = p / p.sum() |
| probs_of[qid] = dict(zip(keys, (float(x) for x in p))) |
| raw[qid].update({"prompt_tokens": len(r["ids"]), "rationale_tokens": len(rat), "rationale": self.tok.decode(rat)}) |
| for qid in qids: |
| probs = probs_of[qid] |
| answers[qid] = {"type": "choice", "choice": max(probs, key=probs.get), "probabilities": probs} if questions[qid]["type"] == "choice" else {"type": "noul", "noul": probs["true"]} |
| return {"model": self.model_id, "answers": answers, "usage": {"input_tokens": sum(v["prompt_tokens"] for v in raw.values())}}, raw |
|
|
