Upload joey-1

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,129 @@
+---
+license: mit
+library_name: transformers
+base_model: Qwen/Qwen3.5-2B
+pipeline_tag: text-classification
+language:
+- en
+tags:
+- decision-model
+- classification
+- calibration
+- routing
+- guardrails
+- vision
+- multimodal
+- baseten
+---
+
+# joey-1
+
+joey-1 is a 2B parameter **decision model** post-trained from [Qwen3.5-2B](https://huggingface.co/Qwen/Qwen3.5-2B) by Base Labs at Baseten.
+Provide the model a state (text, optionally with images) and a set of questions and it will return a calibrated probability for every option,
+answering all questions from **one shared-prefix prefill**. On tested benchmarks:
+
+- **63.3 on the public split of the Jev Decision Index**, ahead of every listed decision model of 2.5B parameters or fewer.
+- **62.9 on a replica of the Artificial Analysis Decision Index**.
+
+For tasks that need more accuracy, see the 9B sibling [elephant-1](https://huggingface.co/baselabs/elephant-1) and the 27B [mammoth-1](https://huggingface.co/baselabs/mammoth-1).
+
+## 🗒️ Model Details
+
+| Model | Parameters | Description |
+|---|---|---|
+| [Qwen3.5-2B](https://huggingface.co/Qwen/Qwen3.5-2B) | 2B | General-purpose base model |
+| **[joey-1](https://huggingface.co/baselabs/joey-1)** | 2B | Post-trained for single-pass, calibrated decisions |
+| [elephant-1](https://huggingface.co/baselabs/elephant-1) | 9B | Larger sibling |
+| [mammoth-1](https://huggingface.co/baselabs/mammoth-1) | 27B | Largest sibling, highest accuracy |
+
+- **Inputs**: text and images
+- **Training**: full fine-tune of the language model on text and image decisions; the vision encoder is frozen
+- **Attention**: block-bidirectional mask on full-attention layers.
+- **Max sequence length**: 16,384 tokens
+
+## 🏃 How to use
+
+### Self-hosted
+
+`setup.sh` installs the pinned serving runtime (OneJev + `serve.patch`). It needs one CUDA GPU and the `pytorch/pytorch:2.8.0-cuda12.8-cudnn9-devel` image.
+
+```bash
+hf download baselabs/joey-1 --local-dir joey-1 && bash joey-1/setup.sh
+```
+
+```python
+from qev.engine import DecisionEngine
+from qev.schema import SystemOneRequest
+
+engine = DecisionEngine("baselabs/joey-1", device="cuda:0")
+
+questions = {
+    "refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"},
+    "team": {
+        "type": "choice",
+        "instructions": "Which team should handle this?",
+        "criteria": {"billing": "Charges, refunds, invoices",
+                     "technical": "App or site faults",
+                     "fraud": "Suspected unauthorised use"},
+    },
+    "urgency": {
+        "type": "score",
+        "instructions": "How urgent is this?",
+        "criteria": ["Can wait", "Today", "Blocking the customer now"],
+    },
+}
+resp, _ = engine.decide(SystemOneRequest(
+    state="I was charged twice this month, please refund one of them.", questions=questions))
+print(resp.answers)
+```
+
+| `type` | `criteria` | answer fields |
+|---|---|---|
+| `noul`: yes or no | optional | `noul`: P(yes) |
+| `choice`: one of named options | `{name: description}` | `choice`, `probabilities` |
+| `score`: ordered levels | list of levels, lowest first | `score`, `probabilities` |
+
+Requests longer than the trained length limits are refused, not truncated.
+
+## 📊 Performance
+
+### Jev Decision Index, public split (our runs)
+
+| Model | Public score |
+|---|---:|
+| elephant-1 | 72.2 |
+| **joey-1** | **63.3** |
+
+| Area | joey-1 |
+|---|---:|
+| Knowledge & reasoning | 50.0 |
+| Language | 62.8 |
+| Retrieval & classification | 71.2 |
+| Tools & automation | 85.8 |
+| Arts & human taste | 54.8 |
+
+### Artificial Analysis Decision Index (our replica)
+
+| Model | Score |
+|---|---:|
+| elephant-1 | 65.2 |
+| **joey-1** | **62.9** |
+
+## 🔧 Training
+
+- **Data**: 6.21M decision examples (Mix 8), the same mix as elephant-1.
+- **Objective**: cross-entropy plus a Brier calibration term (λ = 1), with a trained "I don't know" option.
+- **Recipe**: 1 epoch (43,390 steps), peak LR 1e-5, AdamW (no weight decay), 3% warm-up, cosine decay, bf16.
+- **Hardware**: 2 nodes × 8 H100 (16 GPUs).
+
+## Citation
+
+```bibtex
+@article{baseten2026mammoth,
+  author  = {Base Labs, Baseten},
+  title   = {mammoth-1 and elephant-1: Frontier Decision Models},
+  journal = {Base Labs},
+  year    = {2026},
+  note    = {https://labs.baseten.co/articles/mammoth-1-and-elephant-1-frontier-decision-models},
+}
+```

chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is true %}
+        {{- '<think>\n' }}
+    {%- else %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,115 @@
+{
+  "architectures": [
+    "Qwen3_5ForConditionalGeneration"
+  ],
+  "dtype": "float32",
+  "image_token_id": 248056,
+  "model_type": "qwen3_5",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_output_gate": true,
+    "bos_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 6144,
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 16,
+    "linear_value_head_dim": 128,
+    "mamba_ssm_dtype": "float32",
+    "max_position_embeddings": 262144,
+    "mlp_only_layers": [],
+    "model_type": "qwen3_5_text",
+    "mtp_num_hidden_layers": 1,
+    "mtp_use_dedicated_embeddings": false,
+    "num_attention_heads": 8,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 2,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        11,
+        11,
+        10
+      ],
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 10000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "vocab_size": 248320
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.19.0",
+  "use_cache": false,
+  "video_token_id": 248057,
+  "vision_config": {
+    "deepstack_visual_indexes": [],
+    "depth": 24,
+    "dtype": "bfloat16",
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1024,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "model_type": "qwen3_5_vision",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 2048,
+    "patch_size": 16,
+    "rope_parameters": {
+      "rope_theta": 10000.0,
+      "rope_type": "axial"
+    },
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 248054,
+  "vision_start_token_id": 248053,
+  "modeling_overrides": {
+    "attention": {
+      "full_attention": "block_bidirectional",
+      "linear_attention": "causal"
+    },
+    "readout": "slot_lm",
+    "head_dim": 256,
+    "n_slots": 256
+  }
+}

constraints.txt

@@ -0,0 +1,65 @@
+# Pinned package versions of the training and eval stack, minus torch's own CUDA libraries, which follow the image.
+accelerate==1.15.0
+annotated-doc==0.0.5
+annotated-types==0.8.0
+anyio==4.15.1
+astunparse==1.6.3
+brotli==1.2.0
+causal-conv1d==1.7.0
+click==8.5.0
+cmake==4.0.0
+dnspython==2.7.0
+einops==0.8.2
+expecttest==0.3.0
+fastapi==0.143.1
+fla-core==0.5.2
+flash-linear-attention==0.5.2
+fsspec==2025.3.2
+googleapis-common-protos==1.75.5
+h11==0.16.0
+hf-xet==1.7.0
+httpcore2==2.13.1
+httpx2==2.13.1
+huggingface_hub==2.2.0
+hypothesis==6.131.7
+idna==3.20
+lintrunner==0.12.7
+markdown-it-py==4.2.0
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.4.2
+ninja==1.11.1.4
+opentelemetry-api==1.45.1
+opentelemetry-exporter-http-transport==0.66b1
+opentelemetry-exporter-otlp-common==0.66b1
+opentelemetry-exporter-otlp-proto-common==1.45.1
+opentelemetry-exporter-otlp-proto-http==1.45.1
+opentelemetry-proto==1.45.1
+opentelemetry-sdk==1.45.1
+opentelemetry-semantic-conventions==0.66b1
+optree==0.15.0
+peft==0.21.2
+pillow==11.0.0
+protobuf==7.36.2
+pyarrow==26.0.0
+pydantic==2.14.1
+pydantic_core==2.50.1
+python-etcd==0.4.5
+regex==2026.9.29
+rich==15.0.0
+safetensors==0.8.0
+shellingham==1.5.4
+six==1.17.0
+sortedcontainers==2.4.0
+starlette==1.8.0
+sympy==1.13.3
+tokenizers==0.23.3
+transformers==5.19.0
+typer==0.27.3
+types-dataclasses==0.6.6
+typing-inspection==0.4.4
+typing_extensions==4.16.0
+uvicorn==0.54.0
+wandb==0.30.0
+xxhash==4.0.1
+zstandard==0.23.0

engine.py

@@ -0,0 +1,55 @@
+"""Decision Index engine for this checkpoint, on OneJev's qev runtime @81ce62f + serve.patch (setup.sh installs both).
+
+The checkpoint's readout.json (attention, readout) and optional calibration.json are read at load, so the model is
+served exactly as it was trained: one shared-prefix prefill per request and one batched forward over the question
+suffixes. No prompt tuning, truncation or option pruning; requests over the trained length limits are refused.
+
+    python -m decision_index pipeline --engine engine:DecisionModel --option model=<repo or path> --out runs/<name>
+"""
+from decision_index.engines import Engine, Unsupported
+from qev.engine import DecisionEngine
+from qev.schema import SystemOneRequest
+
+FIELDS = {"choice": {"type", "choice", "probabilities"}, "noul": {"type", "noul"},
+          "score": {"type", "score", "probabilities"}}
+
+
+class DecisionModel(Engine):
+    name = "decision-model"
+    latency = "In-process qev decide() wall time including prompt rendering; excludes model loading."
+
+    def __init__(self, model, device="cuda:0", **options):
+        super().__init__(model=model, device=device, **options)
+        self.engine = DecisionEngine(model, device=device)
+        info = self.engine.info
+        self.provenance = {"kind": "full fine-tune", "runtime": "OneJev qev @81ce62f + serve.patch", "model": model,
+                           "fork_mode": info.fork_mode, "attention": info.extra["attention"],
+                           "readout": info.extra["readout"], "calibration": info.extra["calibration"]}
+
+    def __call__(self, state, questions):
+        for q in questions.values():
+            if q.get("type") not in FIELDS:
+                raise Unsupported(f"question type {q.get('type')}")
+        request = SystemOneRequest(state=state, questions=questions)
+        try:
+            resp, _ = self.engine.decide(request)
+        except ValueError as e:
+            if "limit is" in str(e):
+                raise Unsupported(str(e)) from e
+            raise
+        except RuntimeError as e:
+            # The batched prefix fork overflows 32-bit indexing on the largest multi-question requests; the
+            # sequential branches are OneJev's exact fallback for the same computation.
+            if "32BitIndexMath" not in str(e):
+                raise
+            mode, self.engine.info.fork_mode = self.engine.info.fork_mode, "sequential"
+            try:
+                resp, _ = self.engine.decide(request)
+            finally:
+                self.engine.info.fork_mode = mode
+        answers = {k: a.model_dump(include=FIELDS[a.type]) for k, a in resp.answers.items()}
+        return {"model": self.options["model"], "answers": answers}, None
+
+    def synchronize(self):
+        import torch
+        torch.cuda.synchronize(self.engine.device)

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 248044,
+  "transformers_version": "5.19.0",
+  "use_cache": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c503aded23bae311dff80a63326503fb720572c178f07a2142332d55f5f3afc
+size 5443677680

processor_config.json

@@ -0,0 +1,61 @@
+{
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_processor_type": "Qwen2VLImageProcessor",
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "merge_size": 2,
+    "patch_size": 16,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "size": {
+      "longest_edge": 16777216,
+      "shortest_edge": 65536
+    },
+    "temporal_patch_size": 2
+  },
+  "processor_class": "Qwen3VLProcessor",
+  "video_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "fps": 2,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "max_frames": 768,
+    "max_video_tokens": 768,
+    "merge_size": 2,
+    "min_frames": 4,
+    "patch_size": 16,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "size": {
+      "longest_edge": 25165824,
+      "shortest_edge": 4096
+    },
+    "temporal_patch_size": 2,
+    "video_processor_type": "Qwen3VLVideoProcessor"
+  }
+}

readout.json

@@ -0,0 +1,5 @@
+{
+  "readout": "slot_lm",
+  "head_dim": 256,
+  "n_slots": 256
+}

serve.patch

@@ -0,0 +1,785 @@
+diff --git a/qev/engine.py b/qev/engine.py
+index 2e4c7b0..cb79e4f 100644
+--- a/qev/engine.py
++++ b/qev/engine.py
+@@ -10,6 +10,7 @@ tolerance of 0.05 and only guards against structural failures (an unknown cache
+ """
+ from __future__ import annotations
+ 
++import contextlib
+ import copy
+ import hashlib
+ import logging
+@@ -19,6 +20,7 @@ from pathlib import Path
+ from typing import Any
+ 
+ from .answers import make_choice, make_noul, make_score
++from . import readout as ro
+ from .calibrate import Calibration, softmax
+ from .prompt import (DEFAULT_STYLE, LETTERS, MAX_SLOTS, PROMPT_VERSION, SLOT_LABELS, PromptStyle, RenderedQuestion,
+                      build_messages, render_question)
+@@ -102,11 +104,20 @@ class DecisionEngine:
+         self.model_type = config.model_type
+         self.template_kwargs = {"enable_thinking": False} if "qwen" in self.model_type else {}
+         self.slot_ids = self._verify_slots()
++        rcfg = ro.ReadoutConfig.load(model_path) if local else ro.ReadoutConfig()
++        self.readout = ro.Readout.attach(self.model, rcfg, slot_ids=self.slot_ids,
++                                         head_path=Path(model_path) / ro.HEAD_FILE if rcfg.head else None,
++                                         device=self.device)
++        if rcfg.head:
++            self.max_slots = min(self.max_slots, rcfg.n_slots)
++        if calibration is None and local and (Path(model_path) / "calibration.json").exists():
++            self.calibration = Calibration.load(Path(model_path) / "calibration.json")
+         self.pad_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
+         self.info = EngineInfo(model_path=model_path, model_type=self.model_type, dtype=dtype, device=str(self.device),
+                                fork_mode="sequential", fork_max_abs_diff=None,
+                                transformers_version=transformers.__version__, torch_version=torch.__version__,
+-                               extra={"head_dtype": head_dtype})
++                               extra={"head_dtype": head_dtype, "attention": rcfg.attention, "readout": rcfg.readout,
++                                      "calibration": self.calibration.source})
+         if fork_mode == "auto":
+             self._select_fork_mode()
+         else:
+@@ -190,6 +201,44 @@ class DecisionEngine:
+         with torch.no_grad():
+             torch.cuda.empty_cache() if self.device.type == "cuda" else None
+ 
++    def _split(self, state: Any, rendered: list[RenderedQuestion]):
++        """Prefix ids, suffix ids, and per branch the option token spans relative to its suffix (None for slot_lm).
++        The layout is the one train.common.batch_scores computes, so serving matches training: under noncausal
++        attention the cache boundary must be the prefix block boundary."""
++        prefix_ids, suffixes = self._prefix_and_suffix_ids(state, rendered)
++        cfg = self.readout.config
++        if not cfg.needs_layout:
++            return prefix_ids, suffixes, [None] * len(rendered)
++        P = len(prefix_ids)
++        spans = []
++        for r, suf in zip(rendered, suffixes):
++            prompt = self._render_prompt(state, r.suffix)
++            start = prompt.rindex(r.suffix)
++            ids, prefix, sp = ro.token_layout(self.tokenizer, prompt, start,
++                                              [(start + a, start + b) for a, b in r.option_spans or []])
++            if ids != prefix_ids + suf:
++                raise ValueError("layout tokenization differs from the branch tokenization")
++            if cfg.noncausal and prefix != P:
++                raise ValueError(f"prefix block ends at token {prefix} but the shared cache at {P}")
++            if any(a < P for a, _ in sp):
++                raise ValueError("an option span starts inside the shared prefix")
++            spans.append([(a - P, b - P) for a, b in sp])
++        return prefix_ids, suffixes, spans
++
++    def _blocks(self, prefixes: list[int], lengths: list[int], width: int):
++        if not self.readout.config.noncausal:
++            return contextlib.nullcontext()
++        return ro.blocks(ro.block_ids(prefixes, lengths, width, self.device))
++
++    def _scores(self, o, hidden, row: int, logit_pos: int, end: int, n_slots: int, spans) -> tuple[list[float], float]:
++        """Scores of one branch from a forward's output: the slot letters' LM logits, or the readout head over the
++        captured final hidden states (`end` and `spans` index the suffix tokens of this row)."""
++        if self.readout.head is None:
++            return self._readout(o.logits[row, logit_pos], n_slots)
++        h = hidden[row: row + 1]
++        scores = self.readout.head(h[:, end], ro.pool_options(h, [spans], n_slots))[0]
++        return scores.tolist(), 1.0
++
+     def _prefix_and_suffix_ids(self, state: Any, rendered: list[RenderedQuestion]) -> tuple[list[int], list[list[int]]]:
+         """Tokenize every branch fully, then find the longest common token prefix that ends at or
+         before the end of the state block, so the shared prefix is exact by construction."""
+@@ -218,21 +267,25 @@ class DecisionEngine:
+     def _run_sequential(self, state: Any, rendered: list[RenderedQuestion]) -> list[BranchResult]:
+         import torch
+ 
+-        prefix_ids, suffixes = self._prefix_and_suffix_ids(state, rendered)
++        prefix_ids, suffixes, spans = self._split(state, rendered)
+         dev = self.device
++        P = len(prefix_ids)
+         with torch.inference_mode():
+-            out = self.model(input_ids=torch.tensor([prefix_ids], device=dev),
+-                             attention_mask=torch.ones((1, len(prefix_ids)), dtype=torch.long, device=dev),
+-                             use_cache=True, return_dict=True, logits_to_keep=1)
++            with self._blocks([P], [P], P):
++                out = self.model(input_ids=torch.tensor([prefix_ids], device=dev),
++                                 attention_mask=torch.ones((1, P), dtype=torch.long, device=dev),
++                                 use_cache=True, return_dict=True, logits_to_keep=1)
+             base = out.past_key_values
+             del out
+             results = []
+-            for r, suf in zip(rendered, suffixes):
++            for r, suf, sp in zip(rendered, suffixes, spans):
+                 cache = copy.deepcopy(base)
+-                o = self.model(input_ids=torch.tensor([suf], device=dev),
+-                               attention_mask=torch.ones((1, len(prefix_ids) + len(suf)), dtype=torch.long, device=dev),
+-                               past_key_values=cache, use_cache=True, return_dict=True, logits_to_keep=1)
+-                logits, mass = self._readout(o.logits[0, -1], r.n_slots)
++                with self._blocks([P], [P + len(suf)], P + len(suf)):
++                    o = self.model(input_ids=torch.tensor([suf], device=dev),
++                                   attention_mask=torch.ones((1, P + len(suf)), dtype=torch.long, device=dev),
++                                   past_key_values=cache, use_cache=True, return_dict=True, logits_to_keep=1)
++                hidden = self.readout.capture.pop() if self.readout.capture else None
++                logits, mass = self._scores(o, hidden, 0, -1, len(suf) - 1, r.n_slots, sp)
+                 results.append(BranchResult(r.qid, r.kind, r.labels, r.descriptions, logits, mass, len(suf)))
+                 del o, cache
+             del base
+@@ -266,22 +319,23 @@ class DecisionEngine:
+         caller falls back to sequential branches."""
+         import torch
+ 
+-        prefix_ids, suffixes = self._prefix_and_suffix_ids(state, rendered)
++        prefix_ids, suffixes, spans = self._split(state, rendered)
+         dev = self.device
+         P = len(prefix_ids)
+         chunk = max(1, min(len(rendered), self.fork_token_budget // max(P, 1)))
+         results: list[BranchResult | None] = [None] * len(rendered)
+         with torch.inference_mode():
+-            out = self.model(input_ids=torch.tensor([prefix_ids], device=dev),
+-                             attention_mask=torch.ones((1, P), dtype=torch.long, device=dev),
+-                             use_cache=True, return_dict=True, logits_to_keep=1)
++            with self._blocks([P], [P], P):
++                out = self.model(input_ids=torch.tensor([prefix_ids], device=dev),
++                                 attention_mask=torch.ones((1, P), dtype=torch.long, device=dev),
++                                 use_cache=True, return_dict=True, logits_to_keep=1)
+             base = out.past_key_values
+             del out
+             start = 0
+             while start < len(rendered):
+                 idx = list(range(start, min(start + chunk, len(rendered))))
+                 try:
+-                    self._run_chunk(base, P, idx, rendered, suffixes, results)
++                    self._run_chunk(base, P, idx, rendered, suffixes, results, spans)
+                 except torch.OutOfMemoryError:
+                     torch.cuda.empty_cache()
+                     if chunk == 1:
+@@ -294,7 +348,7 @@ class DecisionEngine:
+             del base
+         return results
+ 
+-    def _run_chunk(self, base, P: int, idx: list[int], rendered, suffixes, results) -> None:
++    def _run_chunk(self, base, P: int, idx: list[int], rendered, suffixes, results, spans=None) -> None:
+         import torch
+ 
+         dev = self.device
+@@ -311,13 +365,16 @@ class DecisionEngine:
+             ends.append(len(s) - 1)
+         keep = sorted(set(ends))
+         cache = self._fork(base, len(idx))
++        spans = spans or [None] * len(rendered)
+         try:
+-            o = self.model(input_ids=ids, attention_mask=mask, position_ids=pos, past_key_values=cache,
+-                           use_cache=True, return_dict=True,
+-                           logits_to_keep=torch.tensor(keep, dtype=torch.long, device=dev))
++            with self._blocks([P] * len(idx), [P + len(s) for s in sufs], P + width):
++                o = self.model(input_ids=ids, attention_mask=mask, position_ids=pos, past_key_values=cache,
++                               use_cache=True, return_dict=True,
++                               logits_to_keep=torch.tensor(keep, dtype=torch.long, device=dev))
++            hidden = self.readout.capture.pop() if self.readout.capture else None
+             for row, i in enumerate(idx):
+                 r = rendered[i]
+-                logits, mass = self._readout(o.logits[row, keep.index(ends[row])], r.n_slots)
++                logits, mass = self._scores(o, hidden, row, keep.index(ends[row]), ends[row], r.n_slots, spans[i])
+                 results[i] = BranchResult(r.qid, r.kind, r.labels, r.descriptions, logits, mass, len(sufs[row]))
+             del o
+         finally:
+diff --git a/qev/prompt.py b/qev/prompt.py
+index 8eaadd3..de49cca 100644
+--- a/qev/prompt.py
++++ b/qev/prompt.py
+@@ -74,13 +74,14 @@ class RenderedQuestion:
+     labels: list[str]
+     descriptions: list[Any]
+     suffix: str
++    option_spans: list[tuple[int, int]] | None = None
+     n_slots: int = field(init=False)
+ 
+     def __post_init__(self) -> None:
+         self.n_slots = len(self.labels)
+ 
+ 
+-def _options_block(names: Sequence[str], descriptions: Sequence[Any], show_names: bool = True) -> str:
++def _options_lines(names: Sequence[str], descriptions: Sequence[Any], show_names: bool = True) -> list[str]:
+     lines = []
+     for letter, name, desc in zip(SLOT_LABELS, names, descriptions):
+         text = render_entry(desc)
+@@ -90,7 +91,11 @@ def _options_block(names: Sequence[str], descriptions: Sequence[Any], show_names
+             lines.append(f"{letter}. {name}: {text}")
+         else:
+             lines.append(f"{letter}. {name}")
+-    return "\n".join(lines)
++    return lines
++
++
++def _options_block(names: Sequence[str], descriptions: Sequence[Any], show_names: bool = True) -> str:
++    return "\n".join(_options_lines(names, descriptions, show_names))
+ 
+ 
+ def render_question(qid: str, q: Question, order: Sequence[int] | None = None, style: PromptStyle = DEFAULT_STYLE) -> RenderedQuestion:
+@@ -127,6 +132,7 @@ def render_question(qid: str, q: Question, order: Sequence[int] | None = None, s
+         descs = [descs[i] for i in order]
+ 
+     letters = list(SLOT_LABELS[: len(labels)])
++    spans = None
+     if style.layout == "json":
+         payload = {"criterion": instructions,
+                    "options": [{"letter": l, "description": (f"{n}: {render_entry(d)}" if (style.show_names and n is not None and render_entry(d)) else (render_entry(d) or n))}
+@@ -134,12 +140,17 @@ def render_question(qid: str, q: Question, order: Sequence[int] | None = None, s
+         suffix = json.dumps(payload, ensure_ascii=False)[1:]
+     else:
+         header = style.question_header if kind != "score" else style.score_header
++        head = f"{header} {instructions}\n\nOptions:\n"
++        lines = _options_lines(names, descs, style.show_names)
++        spans, pos = [], len(head)
++        for line in lines:
++            spans.append((pos, pos + len(line)))
++            pos += len(line) + 1
+         suffix = (
+-            f"{header} {instructions}\n\n"
+-            f"Options:\n{_options_block(names, descs, style.show_names)}\n\n"
++            head + "\n".join(lines) + "\n\n"
+             + (style.answer_line if len(letters) <= len(LETTERS) else style.answer_line_many).format(letters=", ".join(letters))
+         )
+-    return RenderedQuestion(qid=qid, kind=kind, labels=labels, descriptions=descs, suffix=suffix)
++    return RenderedQuestion(qid=qid, kind=kind, labels=labels, descriptions=descs, suffix=suffix, option_spans=spans)
+ 
+ 
+ def state_block(state: Any, style: PromptStyle = DEFAULT_STYLE) -> str:
+diff --git a/qev/readout.py b/qev/readout.py
+new file mode 100644
+index 0000000..02707ca
+@@ -0,0 +1,269 @@
++"""Readout variants, shared by training (train.common / train.sft) and serving (qev.engine).
++
++attention:
++  causal     the base model's mask.
++  noncausal  in the full-attention layers every token attends bidirectionally within its block and causally to
++             earlier blocks. A rendered prompt has two blocks: the shared prefix (system turn and state block) and
++             the question suffix (question, options, answer line, generation prompt). The prefix never sees a
++             suffix, so its cache is still shared exactly by all questions of a request, and one forward of a
++             suffix scores all of its options. The GatedDeltaNet (linear-attention) layers are recurrent and stay
++             causal.
++
++readout:
++  slot_lm    softmax over the LM-head logits of the slot letters at the answer position.
++  head       a learned head over the final (normed) hidden states:
++                 score_i = <w_i, h_ans> + <W_q h_ans, W_k mean(h over option i's line)> / sqrt(d)
++             w has one row per slot and starts as the LM-head rows of the slot letters, W_k starts at zero, so the
++             untrained head reproduces slot_lm exactly; the pointer term scores each option by its content.
++
++Both settings are recorded next to the weights, so a checkpoint is always served the way it was trained. The
++attention mask is declared per layer type under `modeling_overrides` in `config.json`, where serving stacks
++read it; `readout.json` holds the rest. Checkpoints without the overrides block carry `attention` in `readout.json`.
++"""
++from __future__ import annotations
++
++import bisect
++import contextlib
++import contextvars
++import importlib
++import json
++import math
++from dataclasses import asdict, dataclass
++from pathlib import Path
++from typing import Sequence
++
++import torch
++
++from .prompt import MAX_SLOTS
++
++CONFIG_FILE = "readout.json"
++MODEL_CONFIG_FILE = "config.json"
++OVERRIDES_KEY = "modeling_overrides"
++FULL_ATTENTION_MASK = {"causal": "causal", "noncausal": "block_bidirectional"}
++HEAD_FILE = "readout_head.pt"
++ATTENTION = ("causal", "noncausal")
++READOUT = ("slot_lm", "head")
++
++
++@dataclass
++class ReadoutConfig:
++    attention: str = "causal"
++    readout: str = "slot_lm"
++    head_dim: int = 256
++    n_slots: int = MAX_SLOTS
++
++    def __post_init__(self) -> None:
++        if self.attention not in ATTENTION:
++            raise ValueError(f"attention must be one of {ATTENTION}, got {self.attention!r}")
++        if self.readout not in READOUT:
++            raise ValueError(f"readout must be one of {READOUT}, got {self.readout!r}")
++
++    @property
++    def noncausal(self) -> bool:
++        return self.attention == "noncausal"
++
++    @property
++    def head(self) -> bool:
++        return self.readout == "head"
++
++    @property
++    def needs_layout(self) -> bool:
++        return self.noncausal or self.head
++
++    @classmethod
++    def load(cls, model_dir: str | Path) -> "ReadoutConfig":
++        path, model_path = Path(model_dir) / CONFIG_FILE, Path(model_dir) / MODEL_CONFIG_FILE
++        fields = json.loads(path.read_text()) if path.exists() else {}
++        overrides = json.loads(model_path.read_text()).get(OVERRIDES_KEY) if model_path.exists() else None
++        if overrides:
++            mask = overrides["attention"]["full_attention"]
++            attention = {v: k for k, v in FULL_ATTENTION_MASK.items()}.get(mask)
++            if attention is None:
++                raise ValueError(f"unsupported full_attention mask {mask!r} in {model_path}")
++            if fields.get("attention", attention) != attention:
++                raise ValueError(f"{path} says attention {fields['attention']!r} but {model_path} says {mask!r}")
++            fields["attention"] = attention
++        return cls(**fields)
++
++    def save(self, model_dir: str | Path) -> None:
++        fields = asdict(self)
++        model_path = Path(model_dir) / MODEL_CONFIG_FILE
++        if model_path.exists():  # LoRA adapters have no config.json; their readout.json keeps the attention
++            attention = fields.pop("attention")
++            config = json.loads(model_path.read_text())
++            config[OVERRIDES_KEY] = {"attention": {"full_attention": FULL_ATTENTION_MASK[attention],
++                                                   "linear_attention": "causal"}, **fields}
++            model_path.unlink()  # may be a symlink to the source checkpoint (model_v5/mkvariant.py)
++            model_path.write_text(json.dumps(config, indent=2) + "\n")
++        (Path(model_dir) / CONFIG_FILE).write_text(json.dumps(fields, indent=2))
++
++
++_BLOCKS: contextvars.ContextVar = contextvars.ContextVar("qev_readout_blocks", default=None)
++
++
++def install_noncausal(model) -> None:
++    """Route block ids (set with `blocks`) into the full-attention mask of the model's modeling module. Outside a
++    `blocks` context the model is unchanged. Needs a mask-based attention implementation (sdpa, eager, flex)."""
++    impl = getattr(model.config.get_text_config(), "_attn_implementation", None) or model.config._attn_implementation
++    if impl not in ("sdpa", "eager", "flex_attention"):
++        raise ValueError(f"noncausal attention needs sdpa, eager or flex_attention, not {impl}")
++    mod = importlib.import_module(type(model).__module__)
++    if getattr(mod, "_qev_noncausal", False):
++        return
++    if not hasattr(mod, "create_causal_mask"):
++        raise ValueError(f"{mod.__name__} builds no causal mask to extend")
++    original = mod.create_causal_mask
++
++    def create_causal_mask(*args, **kwargs):
++        ids = _BLOCKS.get()
++        if ids is not None and kwargs.get("block_sequence_ids") is None:
++            kwargs["block_sequence_ids"] = ids
++        return original(*args, **kwargs)
++
++    mod.create_causal_mask = create_causal_mask
++    mod._qev_noncausal = True
++
++
++@contextlib.contextmanager
++def blocks(ids):
++    """Block ids [B, kv_length] over absolute positions (cached prefix included); None leaves the mask causal."""
++    token = _BLOCKS.set(ids)
++    try:
++        yield
++    finally:
++        _BLOCKS.reset(token)
++
++
++def block_ids(prefix_lens: Sequence[int], lengths: Sequence[int], width: int, device) -> torch.Tensor:
++    """0 on the prefix, 1 on the suffix, -1 on right padding (causal, and masked as a key anyway)."""
++    pos = torch.arange(width, device=device)[None]
++    p = torch.tensor(list(prefix_lens), device=device)[:, None]
++    n = torch.tensor(list(lengths), device=device)[:, None]
++    return torch.where(pos < p, 0, torch.where(pos < n, 1, -1))
++
++
++def token_layout(tokenizer, prompt: str, suffix_start: int, option_spans: Sequence[tuple[int, int]] | None):
++    """Token ids of `prompt`, the prefix length (tokens that end at or before the suffix) and each option line's
++    token span [start, end), from the tokenizer's character offsets. Option spans are characters of the prompt."""
++    enc = tokenizer(prompt, add_special_tokens=False, return_offsets_mapping=True)
++    ids, offs = enc["input_ids"], enc["offset_mapping"]
++    ends = [e for _, e in offs]
++    starts = [s for s, _ in offs]
++    prefix = bisect.bisect_right(ends, suffix_start)
++    spans = []
++    for cs, ce in option_spans or []:
++        a = bisect.bisect_right(ends, cs)
++        b = bisect.bisect_left(starts, ce)
++        if b <= a:
++            raise ValueError(f"option span {cs}:{ce} covers no token")
++        spans.append((a, b))
++    return ids, prefix, spans
++
++
++class ReadoutHead(torch.nn.Module):
++    def __init__(self, hidden: int, n_slots: int = MAX_SLOTS, head_dim: int = 256):
++        super().__init__()
++        self.slot = torch.nn.Parameter(torch.zeros(n_slots, hidden))
++        self.q = torch.nn.Linear(hidden, head_dim, bias=False)
++        self.k = torch.nn.Linear(hidden, head_dim, bias=False)
++        torch.nn.init.normal_(self.q.weight, std=hidden ** -0.5)
++        torch.nn.init.zeros_(self.k.weight)
++
++    @classmethod
++    def from_lm_head(cls, lm_head_weight, slot_ids: Sequence[int], head_dim: int) -> "ReadoutHead":
++        rows = lm_head_weight.detach()[list(slot_ids)].float()
++        head = cls(rows.shape[1], rows.shape[0], head_dim)
++        with torch.no_grad():
++            head.slot.copy_(rows.to(head.slot.device))
++        return head
++
++    def forward(self, h_ans, h_opt):
++        """h_ans [B, H], h_opt [B, K, H] -> scores [B, K] (fp32)."""
++        h_ans, h_opt = h_ans.float(), h_opt.float()
++        k = h_opt.shape[1]
++        slot = h_ans @ self.slot[:k].T
++        ptr = torch.einsum("bd,bkd->bk", self.q(h_ans), self.k(h_opt)) / math.sqrt(self.q.out_features)
++        return slot + ptr
++
++
++def pool_options(hidden, spans: Sequence[Sequence[tuple[int, int]]], kmax: int):
++    """Mean of hidden [B, T, H] over each option's token span -> [B, kmax, H] fp32 (zeros past a row's options)."""
++    bsz, _, dim = hidden.shape
++    dev = hidden.device
++    rows, toks, segs = [], [], []
++    for b, sp in enumerate(spans):
++        for k, (s, e) in enumerate(sp):
++            n = e - s
++            rows.append(torch.full((n,), b, dtype=torch.long))
++            toks.append(torch.arange(s, e, dtype=torch.long))
++            segs.append(torch.full((n,), b * kmax + k, dtype=torch.long))
++    rows, toks, segs = (torch.cat(x).to(dev) for x in (rows, toks, segs))
++    flat = hidden[rows, toks].float()
++    total = torch.zeros(bsz * kmax, dim, dtype=torch.float32, device=dev).index_add(0, segs, flat)
++    count = torch.zeros(bsz * kmax, dtype=torch.float32, device=dev).index_add(
++        0, segs, torch.ones_like(segs, dtype=torch.float32))
++    return (total / count.clamp(min=1)[:, None]).view(bsz, kmax, dim)
++
++
++def final_norm(model):
++    """The final norm of a (possibly multimodal) Qwen-style decoder: its output is the readout hidden state."""
++    inner = getattr(model, "model", model)
++    inner = getattr(inner, "language_model", inner)
++    norm = getattr(inner, "norm", None)
++    if norm is None:
++        raise ValueError(f"cannot find the final norm of {type(model).__name__}")
++    return norm
++
++
++class HiddenCapture:
++    """Keeps the final-norm output of the latest forward (with its graph, so the head trains the backbone)."""
++
++    def __init__(self, model):
++        self.value = None
++        self.handle = final_norm(model).register_forward_hook(self._hook)
++
++    def _hook(self, module, args, output):
++        self.value = output
++
++    def pop(self):
++        out, self.value = self.value, None
++        if out is None:
++            raise RuntimeError("no hidden state captured; the forward did not reach the final norm")
++        return out
++
++
++@dataclass
++class Readout:
++    """A model's readout at runtime: its config, and for `head` the head module and the hidden-state hook."""
++
++    config: ReadoutConfig
++    head: ReadoutHead | None = None
++    capture: HiddenCapture | None = None
++
++    @classmethod
++    def attach(cls, model, config: ReadoutConfig, slot_ids: Sequence[int] | None = None,
++               head_path: str | Path | None = None, device=None) -> "Readout":
++        """Install the variant on `model`. A head is loaded from `head_path` when given, else initialised from the
++        model's LM-head rows of `slot_ids` (call before the model is sharded)."""
++        if config.noncausal:
++            install_noncausal(model)
++        if not config.head:
++            return cls(config)
++        if head_path is not None:
++            state = torch.load(head_path, map_location="cpu", weights_only=True)
++            head = ReadoutHead(state["slot"].shape[1], state["slot"].shape[0], state["q.weight"].shape[0])
++            head.load_state_dict(state)
++        else:
++            head = ReadoutHead.from_lm_head(model.get_output_embeddings().weight, slot_ids[: config.n_slots],
++                                            config.head_dim)
++        head = head.to(device or next(model.parameters()).device)
++        return cls(config, head, HiddenCapture(model))
++
++    def save(self, out_dir: str | Path) -> None:
++        self.config.save(out_dir)
++        if self.head is not None:
++            torch.save({k: v.detach().cpu() for k, v in self.head.state_dict().items()}, Path(out_dir) / HEAD_FILE)
++
++    @property
++    def max_slots(self) -> int:
++        return self.config.n_slots
+diff --git a/train/common.py b/train/common.py
+index 331206f..896e6d3 100644
+--- a/train/common.py
++++ b/train/common.py
+@@ -85,6 +85,8 @@ class Example:
+     target: list[float]
+     source: str
+     task: str
++    suffix_start: int = -1
++    option_spans: list[tuple[int, int]] | None = None
+ 
+ 
+ PAD_FAMILIES = ("next_action",)
+@@ -185,6 +187,8 @@ def render_example(row: dict, tmpl, idx: int, seed: int | None = None, augment:
+         compact = augment and rng.random() < 0.5
+     state = serialize_state(row.get("state"), compact)
+     prompt = tmpl(state, rendered.suffix)
++    suffix_start = prompt.rindex(rendered.suffix)
++    spans = [(suffix_start + a, suffix_start + b) for a, b in rendered.option_spans or []]
+     return Example(
+         idx=idx,
+         prompt=prompt,
+@@ -193,6 +197,8 @@ def render_example(row: dict, tmpl, idx: int, seed: int | None = None, augment:
+         target=target_vector(row["target"], rendered.labels),
+         source=str(row.get("source", "unknown")),
+         task=str(row.get("task", "unknown")),
++        suffix_start=suffix_start,
++        option_spans=spans or None,
+     )
+ 
+ 
+@@ -244,6 +250,7 @@ class Bundle:
+     template_kwargs: dict
+     processor: Any = None
+     counter: Any = None
++    readout: Any = None
+ 
+     @property
+     def multimodal(self) -> bool:
+@@ -403,6 +410,65 @@ def readout_logits(model, input_ids, attention_mask, ends):
+     return out.logits[rows, pos]
+ 
+ 
++def batch_scores(bundle: Bundle, model, examples: Sequence[Example], rows: Sequence[dict], device):
++    """Option scores [B, Kmax] (-inf past each row's options) and their mask for a right-padded batch, under the
++    bundle's readout (qev.readout): causal or noncausal attention, slot_lm or head scoring. The same function scores
++    training steps, in-training evals and train.eval_mm, and qev.engine reproduces it for serving."""
++    import torch
++
++    from qev import readout as ro
++
++    rd = bundle.readout
++    cfg = rd.config if rd is not None else ro.ReadoutConfig()
++    ks = [len(ex.labels) for ex in examples]
++    if bundle.multimodal:
++        from train import mm
++
++        inputs, ends = mm.mm_inputs(bundle.processor, [ex.prompt for ex in examples],
++                                    [r.get("media") or [] for r in rows], device)
++        shifts = [bundle.counter.media_extra(r.get("media") or []) for r in rows]
++    else:
++        input_ids, attention_mask, ends = pad_batch([bundle.encode(ex.prompt) for ex in examples], bundle.tokenizer.pad_token_id, device)
++        inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
++        shifts = [0] * len(examples)
++    ctx = ro.blocks(None)
++    spans = None
++    if cfg.needs_layout:
++        prefixes, spans = [], []
++        for ex, shift, end in zip(examples, shifts, ends):
++            ids, prefix, sp = ro.token_layout(bundle.tokenizer, ex.prompt, ex.suffix_start, ex.option_spans)
++            if len(ids) + shift != end + 1:
++                raise ValueError(f"layout of row {ex.idx} has {len(ids) + shift} tokens, the batch has {end + 1}")
++            if cfg.head and len(sp) != len(ex.labels):
++                raise ValueError(f"row {ex.idx}: {len(sp)} option spans for {len(ex.labels)} options")
++            prefixes.append(prefix + shift)
++            spans.append([(a + shift, b + shift) for a, b in sp])
++        if cfg.noncausal:
++            width = inputs["input_ids"].shape[1]
++            ctx = ro.blocks(ro.block_ids(prefixes, [e + 1 for e in ends], width, device))
++    with ctx:
++        if cfg.head:
++            out = model(**inputs, use_cache=False, return_dict=True, logits_to_keep=1)
++            hidden = rd.capture.pop()
++            kmax = max(ks)
++            rows_ix = torch.arange(len(examples), device=hidden.device)
++            h_ans = hidden[rows_ix, torch.tensor(ends, device=hidden.device)]
++            scores = rd.head(h_ans, ro.pool_options(hidden, spans, kmax))
++            # every FSDP unit (the LM head too) must take part in the backward
++            scores = scores + 0.0 * out.logits.float().sum()
++            mask = torch.zeros((len(ks), kmax), dtype=torch.bool, device=scores.device)
++            for i, k in enumerate(ks):
++                mask[i, :k] = True
++            return scores.masked_fill(~mask, float("-inf")), mask
++        if bundle.multimodal:
++            from train import mm
++
++            vocab_logits = mm.mm_readout_logits(model, inputs, ends)
++        else:
++            vocab_logits = readout_logits(model, inputs["input_ids"], inputs["attention_mask"], ends)
++    return slot_view(vocab_logits, bundle.slot_ids, ks)
++
++
+ def slot_view(vocab_logits, slot_ids, ks):
+     """Gather the first K slot logits per row into a padded [B, Kmax] tensor plus a mask."""
+     import torch
+@@ -557,11 +623,7 @@ def collect_slot_logits(bundle: Bundle, rows: Sequence[dict], max_len: int = 409
+             readout_logits(bundle.model, ids, torch.ones_like(ids), [0])
+         for bi, batch in enumerate(batches):
+             exs = [examples[i] for i in batch]
+-            id_lists = [bundle.encode(ex.prompt) for ex in exs]
+-            input_ids, attention_mask, ends = pad_batch(id_lists, pad_id, dev)
+-            ks = [len(ex.labels) for ex in exs]
+-            vocab_logits = readout_logits(bundle.model, input_ids, attention_mask, ends)
+-            sel, _ = slot_view(vocab_logits, bundle.slot_ids, ks)
++            sel, _ = batch_scores(bundle, bundle.model, exs, [rows[ex.idx] for ex in exs], dev)
+             sel = sel.float().cpu().tolist()
+             for r, ex in enumerate(exs):
+                 out.append(Scored(kind=ex.kind, k=len(ex.labels), logits=sel[r][: len(ex.labels)],
+diff --git a/train/eval_mm.py b/train/eval_mm.py
+index e87e72d..03a6022 100644
+--- a/train/eval_mm.py
++++ b/train/eval_mm.py
+@@ -48,6 +48,13 @@ def main() -> int:
+         mm.set_root(args.media_root)
+     bundle = common.load_bundle(args.model, device=args.device, dtype="bfloat16", multimodal=args.multimodal,
+                                 chat_template=args.chat_template)
++    from qev.readout import HEAD_FILE, Readout, ReadoutConfig
++
++    rcfg = ReadoutConfig.load(args.model)
++    head_path = Path(args.model) / HEAD_FILE
++    bundle.readout = Readout.attach(bundle.model, rcfg, slot_ids=bundle.slot_ids,
++                                    head_path=head_path if rcfg.head else None)
++    print(f"readout: attention {rcfg.attention}, scores {rcfg.readout}", flush=True)
+     results = {}
+     for spec in args.evals:
+         name, _, path = spec.partition("=")
+diff --git a/train/sft.py b/train/sft.py
+index 4f6f7bd..ae428d9 100644
+--- a/train/sft.py
++++ b/train/sft.py
+@@ -108,6 +108,14 @@ def build_parser() -> argparse.ArgumentParser:
+                         "train one step on each listed micro-batch (one GPU), peak memory reset per step, no final eval")
+     p.add_argument("--dump-first-grad", default=None,
+                    help="test hook: after the first step's all-reduce, rank 0 saves the adapter gradients here")
++    p.add_argument("--attention", choices=["causal", "noncausal"], default="causal",
++                   help="noncausal: bidirectional within the prefix and within the question suffix in the full-attention "
++                        "layers (qev.readout)")
++    p.add_argument("--readout", choices=["slot_lm", "head"], default="slot_lm",
++                   help="head: learned slot + pointer head over the final hidden states instead of LM-head letter logits")
++    p.add_argument("--head-dim", type=int, default=256, help="pointer projection width of the readout head")
++    p.add_argument("--head-lr", type=float, default=1e-4, help="peak learning rate of the pointer projections "
++                                                              "(the head's slot rows use --lr)")
+     p.add_argument("--disable-fla-kernels", action="store_true",
+                    help="force the reference PyTorch gated delta rule instead of the fla kernel "
+                         "(needed on H100/H200 with Triton in [3.4.0, 3.7.1); correct but much slower)")
+@@ -326,10 +334,7 @@ def collect_mm(bundle, rows, max_len: int, batch_tokens: int, max_rows: int, dev
+             common.readout_logits(bundle.model, ids, torch.ones_like(ids), [0])
+         for b in batches:
+             exs = [examples[usable[j]] for j in b]
+-            inputs, ends = mm.mm_inputs(bundle.processor, [ex.prompt for ex in exs],
+-                                        [rows[ex.idx].get("media") or [] for ex in exs], device)
+-            vocab_logits = mm.mm_readout_logits(bundle.model, inputs, ends)
+-            sel, _ = common.slot_view(vocab_logits, bundle.slot_ids, [len(ex.labels) for ex in exs])
++            sel, _ = common.batch_scores(bundle, bundle.model, exs, [rows[ex.idx] for ex in exs], device)
+             sel = sel.float().cpu().tolist()
+             for r, ex in enumerate(exs):
+                 out.append(common.Scored(kind=ex.kind, k=len(ex.labels), logits=sel[r][: len(ex.labels)],
+@@ -339,6 +344,13 @@ def collect_mm(bundle, rows, max_len: int, batch_tokens: int, max_rows: int, dev
+ 
+ 
+ def save_weights(model, bundle, out_dir: Path, full: bool, sharded: bool = False, is_main: bool = True) -> None:
++    _save_weights(model, bundle, out_dir, full, sharded, is_main)
++    if is_main and bundle.readout is not None:
++        out_dir.mkdir(parents=True, exist_ok=True)
++        bundle.readout.save(out_dir)
++
++
++def _save_weights(model, bundle, out_dir: Path, full: bool, sharded: bool = False, is_main: bool = True) -> None:
+     """LoRA: the adapter. Full fine-tuning: the whole model in bf16 (from the fp32 master weights, copied to CPU),
+     with tokenizer and processor, loadable like the base checkpoint. Sharded (--fsdp): every rank must call it,
+     the weights are gathered to rank 0's CPU and written there as HF safetensors shards."""
+@@ -560,6 +572,8 @@ def main(argv=None) -> int:
+         raise SystemExit("--fsdp is implemented without --share-prefix")
+     if args.disable_fla_kernels:
+         common.block_fla_kernels()
++    if (args.attention != "causal" or args.readout != "slot_lm") and args.share_prefix:
++        raise SystemExit("--share-prefix supports only causal slot_lm")
+ 
+     import torch
+     from peft import LoraConfig, get_peft_model
+@@ -600,6 +614,16 @@ def main(argv=None) -> int:
+     if args.plan_only:
+         return plan_only(bundle, rows, args, world)
+     log.info("model_type=%s template_kwargs=%s", bundle.model_type, bundle.template_kwargs)
++    from qev.readout import Readout, ReadoutConfig
++
++    bundle.readout = Readout.attach(bundle.model, ReadoutConfig(attention=args.attention, readout=args.readout,
++                                                                head_dim=args.head_dim),
++                                    slot_ids=bundle.slot_ids, device=torch.device(args.device))
++    head = bundle.readout.head
++    if head is not None and world > 1:
++        for prm in head.parameters():
++            dist.broadcast(prm.data, 0)
++    log.info("readout: attention %s, scores %s", args.attention, args.readout)
+ 
+     if args.full_finetune:
+         model = bundle.model
+@@ -696,6 +720,13 @@ def main(argv=None) -> int:
+     opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95), eps=1e-8,
+                             fused=bool(args.full_finetune))
+     sched = get_cosine_schedule_with_warmup(opt, num_warmup_steps=warmup, num_training_steps=max(total_steps, 1))
++    head_opt = head_sched = None
++    if head is not None:
++        head_opt = torch.optim.AdamW([{"params": [head.slot], "lr": args.lr},
++                                      {"params": [head.q.weight, head.k.weight], "lr": args.head_lr}],
++                                     weight_decay=0.0, betas=(0.9, 0.95), eps=1e-8)
++        head_sched = get_cosine_schedule_with_warmup(head_opt, num_warmup_steps=warmup,
++                                                     num_training_steps=max(total_steps, 1))
+ 
+     logf = open(log_path, "a", encoding="utf-8") if is_main else open(os.devnull, "w")
+     step = 0
+@@ -729,25 +760,16 @@ def main(argv=None) -> int:
+                 examples.append(common.render_example(rows[i], bundle.template_for(rows[i]), i, seed=seed,
+                                                       augment=args.augment, option_pool=OPTION_POOL,
+                                                       pad_prob=args.pad_options))
+-            if bundle.multimodal:
+-                from train import mm
+-
+-                inputs, ends = mm.mm_inputs(bundle.processor, [ex.prompt for ex in examples],
+-                                            [rows[i].get("media") or [] for i in idx], device)
+-                with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_autocast(args)):
+-                    vocab_logits = mm.mm_readout_logits(model, inputs, ends)
+-            else:
+-                id_lists = [bundle.encode(ex.prompt) for ex in examples]
+-                input_ids, attention_mask, ends = common.pad_batch(id_lists, pad_id, device)
+-                with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_autocast(args)):
+-                    vocab_logits = common.readout_logits(model, input_ids, attention_mask, ends)
++            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_autocast(args)):
++                sel, mask = common.batch_scores(bundle, model, examples, [rows[i] for i in idx], device)
+             ks = [len(ex.labels) for ex in examples]
+             targets = [ex.target for ex in examples]
+         kmax = max(ks)
+         tgt = torch.zeros((len(ks), kmax), dtype=torch.float32, device=device)
+         for r, t in enumerate(targets):
+             tgt[r, : len(t)] = torch.tensor(t, dtype=torch.float32, device=device)
+-        sel, mask = common.slot_view(vocab_logits, bundle.slot_ids, ks)
++        if kind == "group":
++            sel, mask = common.slot_view(vocab_logits, bundle.slot_ids, ks)
+         return sel, mask, tgt
+ 
+     if eval_sets and args.eval_at_start:
+@@ -802,6 +824,17 @@ def main(argv=None) -> int:
+             gnorm = torch.nn.utils.clip_grad_norm_(params, args.max_grad_norm)
+             if hasattr(gnorm, "full_tensor"):
+                 gnorm = gnorm.full_tensor()
++            if head is not None:
++                for prm in head.parameters():
++                    if prm.grad is None:
++                        prm.grad = torch.zeros_like(prm)
++                    if world > 1:
++                        dist.all_reduce(prm.grad)
++                        prm.grad /= grad_scale
++                torch.nn.utils.clip_grad_norm_(list(head.parameters()), args.max_grad_norm)
++                head_opt.step()
++                head_sched.step()
++                head_opt.zero_grad(set_to_none=True)
+             try:
+                 opt.step()
+             except torch.OutOfMemoryError:

setup.sh

@@ -0,0 +1,16 @@
+# Usage: bash setup.sh [WORKDIR=.]   (image: pytorch/pytorch:2.8.0-cuda12.8-cudnn9-devel, one CUDA GPU)
+# Installs the serving runtime (OneJev @81ce62f + serve.patch) and the Decision Index kit, pinned as in our runs.
+set -e
+B=$(cd "$(dirname "$0")" && pwd); W=${1:-.}
+ONEJEV=81ce62f1597c91e46767d4d02d6ac2e18534fe94; KIT=f5a6e4fd85fed05e77878dfc4966d0568bfb45ec
+command -v git >/dev/null || { apt-get update -qq; apt-get install -y -qq git >/dev/null; }
+cd "$W"
+[ -d OneJev ] || git clone -q https://github.com/OmniJev/OneJev.git
+git -C OneJev checkout -q -f $ONEJEV && git -C OneJev clean -fdq && git -C OneJev apply "$B/serve.patch"
+[ -d decision-index ] || git clone -q https://github.com/apolinario/decision-index.git
+git -C decision-index checkout -q -f $KIT
+# Keep the image's CUDA build of torch: unpinned, pip resolves torch>=2.6 to the newest wheel (another CUDA).
+python -c "import torch, torchvision; print(f'torch=={torch.__version__}\ntorchvision=={torchvision.__version__}')" > image-torch.txt
+pip install -q -c "$B/constraints.txt" -c image-torch.txt -e "./OneJev[train]" huggingface_hub hf_xet brotli flash-linear-attention httpx
+pip install -q -c "$B/constraints.txt" -c image-torch.txt --no-build-isolation causal-conv1d
+pip install -q --no-deps -e ./decision-index

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5cd9732badce41de57e6efce8302930ded1c1188c5f81feb2bd6c24c4a1941f
+size 19989339

tokenizer_config.json

@@ -0,0 +1,49 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "local_files_only": true,
+  "merges_file": null,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "processor_class": "Qwen3VLProcessor",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen3_5Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}