Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        https://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "{}"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright {yyyy} {name of copyright owner}
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       https://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,21 @@
+Jev-Spatial — Apache License, Version 2.0
+
+Based on Molmo2-ER and Molmo2 by the Allen Institute for AI and contributors.
+Model: https://huggingface.co/allenai/Molmo2-ER
+Source: https://github.com/allenai/molmo2
+Pinned revision: dab22564403d2607855bb1fffb0721285b445081
+Upstream license: Apache License, Version 2.0 (retained in LICENSE).
+
+The five native Python modules under spatial_jev/molmo2 are unmodified copies
+of the model/processor implementation at the pinned revision. Jev-Spatial adds
+a shared classification head, hierarchical inference, training and evaluation
+code. The released weights merge the trained LoRA into Molmo2-ER and omit the
+unused autoregressive vocabulary head. Configuration metadata is adapted to
+load the implementation from the installed source package.
+
+Jev / System One by TypeSafe is the project's conceptual inspiration.
+This is an independent project, not an official TypeSafe implementation or
+a reproduction of proprietary RLCD. Full acknowledgements are in README.md.
+
+Training and benchmark media are not included. Their original attribution
+and terms remain with their respective authors and distributors.

README.md

@@ -1,13 +1,32 @@
 ---
-title: Mjev Spatial
-emoji: 🔥
-colorFrom: indigo
-colorTo: purple
+title: MJev-Spatial
+emoji: 🧭
+colorFrom: yellow
+colorTo: blue
 sdk: gradio
 sdk_version: 6.30.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Spatial VLM that decides in one forward pass
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MJev-Spatial
+
+Demo of **MJev-Spatial: System One Spatial Intelligence via Native Attention Matching**
+([arXiv 2610.22276](https://arxiv.org/abs/2610.22276) · [GitHub](https://github.com/Fr0zenCrane/MJev-Spatial)).
+
+MJev-Spatial is a 4B multimodal decision model built on Qwen3.5-4B. It turns a spatial
+question into a finite candidate-selection problem and scores those candidates in a
+**single forward pass** with no generated answer tokens, reusing the language backbone's
+final full-attention layer's native query/key projections as the readout.
+
+The demo exposes the three answer spaces from the paper:
+
+- **Choice** — pick one of the supplied options (e.g. relative-direction questions).
+- **Scalar** — estimate a nonnegative measurement in meters via an interval hierarchy.
+- **Pointing** — point at a requested target in the image; the predicted point is
+  drawn on the image (the model makes three 3×3 region decisions, refining a crop).
+
+Example image: `examples/images/scene.png` from the official
+[MJev-Spatial repository](https://github.com/Fr0zenCrane/MJev-Spatial) (Apache-2.0).

app.py

@@ -0,0 +1,331 @@
+"""MJev-Spatial — one-forward-pass spatial decision demo.
+
+Model: https://huggingface.co/Fr0zenCr4nE/MJev-Spatial-Qwen-3.5-4B
+Paper: MJev-Spatial: System One Spatial Intelligence via Native Attention Matching
+(arXiv 2610.22276). Inference follows the official scripts/predict.py path:
+the model scores a finite candidate set with the language backbone's native
+final-attention Q/K projections and returns the argmax decision — no text is generated.
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: E402  -- MUST precede torch / CUDA-touching imports
+import json  # noqa: E402
+import sys  # noqa: E402
+import tempfile  # noqa: E402
+import time  # noqa: E402
+from pathlib import Path  # noqa: E402
+
+import gradio as gr  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+from PIL import Image, ImageDraw  # noqa: E402
+
+ROOT = Path(__file__).resolve().parent
+sys.path.insert(0, str(ROOT / "src"))
+
+# Vendored inference runtime (Apache-2.0) from github.com/Fr0zenCrane/MJev-Spatial
+from spatial_jev.kev import KevBuilder, build_model, make_processor, predict_batch  # noqa: E402
+from spatial_jev.numeric_tasks import ANGLE_CODEBOOK, COUNT_CODEBOOK  # noqa: E402
+
+CHECKPOINT = "Fr0zencr4nE/MJev-Spatial-Qwen-3.5-4B"
+MAX_SIDE = 512  # keep encode time and sequence length comfortable on ZeroGPU
+
+# ---------------------------------------------------------------------------
+# Model loading (module scope, eager .to("cuda") — intercepted by the spaces hijack)
+# ---------------------------------------------------------------------------
+
+ckpt_dir = CHECKPOINT  # from_pretrained accepts repo ids directly
+config = json.loads(Path(hf_hub_file(ckpt_dir, "config.json")).read_text()) if False else None
+
+
+def hf_hub_file(repo_id, filename):
+    from huggingface_hub import hf_hub_download
+    return hf_hub_download(repo_id, filename)
+
+
+config = json.loads(hf_hub_file(ckpt_dir, "config.json").read_text())
+config.update(
+    preserve_option_order=True,
+    max_image_tokens=1024,
+    max_sequence_length=max(8192, int(config.get("max_sequence_length", 4096))),
+    inference_language_dtype="bf16",
+)
+
+print("[mjev] loading processor…", flush=True)
+processor = make_processor(config)
+print("[mjev] loading model…", flush=True)
+builder = KevBuilder(processor, config)
+MODEL = build_model(config, ckpt_dir, trainable=False)
+MODEL = MODEL.to("cuda").eval()
+print("[mjev] model ready", flush=True)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def load_image(image, max_side=MAX_SIDE):
+    """PIL RGB image scaled so neon-style huge uploads don't blow the token budget."""
+    if image is None:
+        raise gr.Error("Please provide an image.")
+    img = image if isinstance(image, Image.Image) else Image.open(image)
+    img = img.convert("RGB")
+    w, h = img.size
+    scale = max_side / max(w, h)
+    if scale < 1:
+        img = img.resize((max(1, round(w * scale)), max(1, round(h * scale))), Image.Resampling.BILINEAR)
+    return img
+
+
+def parse_options(text):
+    """One option per line; 'id | text' optional — falls back to the text as id."""
+    options = []
+    for line in (text or "").strip().splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        if "|" in line:
+            oid, otext = line.split("|", 1)
+            oid, otext = oid.strip(), otext.strip() or oid.strip()
+        else:
+            oid = otext = line
+        options.append({"id": oid, "text": otext})
+    if len(options) < 2:
+        raise gr.Error("Choice tasks need at least 2 options (one per line).")
+    if len(options) > config["max_choices"]:
+        raise gr.Error(f"Too many options — this checkpoint supports at most {config['max_choices']}.")
+    return options
+
+
+def draw_point(image, xy):
+    """Overlay the predicted normalized point + crosshair on a copy of the image."""
+    img = image.copy()
+    w, h = img.size
+    x, y = xy[0] * w, xy[1] * h
+    draw = ImageDraw.Draw(img)
+    r = max(6, int(min(w, h) * 0.02))
+    for color, width in ((255, 255, 255, 255), (255, 70, 0, 255)):
+        draw.ellipse([x - r, y - r, x + r, y + r], outline=color, width=width)
+        draw.line([x - 1.6 * r, y, x + 1.6 * r, y], fill=color, width=width)
+        draw.line([x, y - 1.6 * r, x, y + 1.6 * r], fill=color, width=width)
+    return img
+
+
+def overlay_probability_grid(image, result):
+    """Pointing: warm heatmap of the final-stage 9-way scores over the selected crop."""
+    if "logits" not in result or result.get("pointing_method") != "visual_tokens":
+        return None
+    probs = np.asarray(result.get("probabilities", []), dtype=np.float64)
+    if probs.size == 0:
+        return None
+    height, width = result.get("grid_hw", [1, 1])
+    if probs.shape != (height, width):
+        return None
+    heat = np.zeros((height, width, 4), dtype=np.uint8)
+    heat[..., 0] = 255
+    heat[..., 3] = (probs / probs.max() * 255).astype(np.uint8) if probs.max() > 0 else 0
+    heat_img = Image.fromarray(heat, "RGBA").resize(image.size, Image.Resampling.BILINEAR)
+    out = image.copy()
+    out.paste(heat_img, (0, 0), heat_img)
+    return out
+
+
+def describe_path(result, kind, option_map=None):
+    """Human-readable trace of the decision path."""
+    lines = []
+    logits = result.get("logits") or []
+    path = result.get("path") or []
+    names = {
+        "top-center": "top-center", "top-left": "top-left", "top-right": "top-right",
+        "center": "center", "middle-left": "middle-left", "middle-right": "middle-right",
+        "bottom-center": "bottom-center", "bottom-left": "bottom-left",
+        "bottom-right": "bottom-right",
+    }
+    from spatial_jev.hierarchy import POINT_NAMES, decode_point, decode_scalar  # local import
+    if kind == "choice":
+        if option_map:
+            lines.append(f"selected option: **{option_map[path[0]]}**" if path else "")
+        elif result.get("prediction") is not None:
+            lines.append(f"selected option: **{result['prediction']}**")
+    elif kind == "scalar":
+        if path:
+            edges = config["scalar_codebook"]["edges_m"]
+            if path[0] == 0:
+                lines.append("coarse stage: exactly zero")
+            elif path[0] == len(edges):
+                lines.append(f"coarse stage: greater than {edges[-1]:g} m (overflow)")
+            else:
+                from spatial_jev.hierarchy import scalar_interval, scalar_fine_bins
+                lo, hi = scalar_interval(path[0], config["scalar_codebook"])
+                lines.append(f"coarse stage: interval {lo:.6g} – {hi:.6g} m")
+                if len(path) > 1:
+                    fine = scalar_fine_bins(path[0], config["scalar_codebook"])
+                    stepw = (hi - lo) / fine
+                    lines.append(f"fine stage: refined bin {path[1] + 1}/{fine} "
+                                 f"≈ {lo + path[1] * stepw:.6g} – {lo + (path[1] + 1) * stepw:.6g} m")
+            lines.append(f"decoded estimate: **{result['prediction']:.6g} m**")
+    elif kind == "point":
+        steps = ["whole image", *(f"zoom #{i}" for i in range(1, len(path)))]
+        for stage, (choice, scores) in enumerate(zip(path, logits)):
+            label = names.get(POINT_NAMES[choice], POINT_NAMES[choice])
+            best = scores[choice] if stage < len(scores) and choice < len(scores) else float("nan")
+            lines.append(f"stage {stage + 1} ({steps[stage] if stage < len(steps) else ''}): "
+                         f"region **{label}** (score {best:.3f})")
+        if kind == "point" and result.get("prediction"):
+            x, y = result["prediction"]
+            lines.append(f"decoded point: **({x:.3f}, {y:.3f})** normalized "
+                         f"(origin top-left, x→right, y→down)")
+    return "\n".join(f"- {l}" for l in lines if l)
+
+
+# ---------------------------------------------------------------------------
+# Inference endpoint
+# ---------------------------------------------------------------------------
+
+@spaces.GPU(duration=75)
+def run_spatial(image, task, question, options_text, quantity, unit, progress=gr.Progress()):
+    """Answer a spatial question about an image in one forward pass per decision stage.
+
+    Args:
+        image: input image (PIL).
+        task: one of "choice" | "scalar" | "point".
+        question: the spatial question about the image.
+        options_text: for choice tasks, one option per line ("id | text" accepted).
+        quantity: for scalar tasks, what is being estimated (e.g. "distance").
+        unit: scalar unit — "m", "count" or "deg".
+        progress: Gradio progress bar.
+    Returns:
+        (answer markdown, decision-path markdown, annotated image, raw JSON result)
+    """
+    t0 = time.perf_counter()
+    img = load_image(image)
+    if not (question or "").strip():
+        raise gr.Error("Please provide a question.")
+
+    question = question.strip()
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "input.png"
+        img.save(path)
+        if task == "choice":
+            space = {"kind": "choice", "options": parse_options(options_text)}
+        elif task == "scalar":
+            space = {"kind": "scalar", "quantity": (quantity or "value").strip() or "value",
+                     "unit": unit}
+        else:
+            space = {"kind": "point", "coordinate_system": "normalized_xy", "num_points": 1}
+        request = {"id": "demo-request", "input": {
+            "media": [{"kind": "image", "uri": str(path)}],
+            "question": question,
+            "answer_space": space,
+        }}
+        result = predict_batch(MODEL, builder, [request], torch.device("cuda"))[0]
+
+    prediction = result["prediction"]
+    if task == "point":
+        point = prediction[0] if isinstance(prediction, list) else prediction
+        overlay = draw_point(img, point)
+    else:
+        overlay = img.copy()
+
+    if task == "choice":
+        answer = str(prediction)
+    elif task == "scalar":
+        label = {"m": "m", "count": "", "deg": "°"}[unit]
+        answer = f"{prediction:g}{label}" if unit != "count" else f"{int(prediction)} {quantity or 'objects'}"
+    else:
+        answer = f"point (x, y) = ({point[0]:.3f}, {point[1]:.3f})"
+
+    trace = describe_path(result, task)
+    seconds = time.perf_counter() - t0
+    summary = f"## {answer}\n\n*(one forward pass per decision stage, {seconds:.1f}s)*"
+    raw = json.dumps({"prediction": result.get("prediction"), "path": result.get("path"),
+                      "logits": result.get("logits")}, ensure_ascii=False)
+    return summary, trace, overlay, raw
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="MJev-Spatial") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            "# 🧭 MJev-Spatial\n"
+            "Spatial intelligence via **native attention matching** — a 4B VLM that "
+            "answers spatial questions as a single candidate-scoring decision, with no "
+            "generated answer tokens. [paper](https://arxiv.org/abs/2610.22276) · "
+            "[code](https://github.com/Fr0zenCrane/MJev-Spatial) · "
+            "[checkpoint](https://huggingface.co/Fr0zencr4nE/MJev-Spatial-Qwen-3.5-4B)"
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_in = gr.Image(type="pil", label="Image", height=300)
+                task = gr.Radio(
+                    ["choice", "scalar", "point"], value="choice", label="Answer space",
+                    info="choice: pick an option · scalar: estimate a value · point: locate a target",
+                )
+                question = gr.Textbox(
+                    label="Question", placeholder="e.g. Where is the red square relative to the blue circle?",
+                    value="Where is the red square relative to the blue circle?",
+                )
+                options_text = gr.Textbox(
+                    label="Options (choice only) — one per line, 'id | text' optional",
+                    value="left | left\nright | right", lines=3,
+                )
+                quantity = gr.Textbox(
+                    label="Quantity (scalar only)", value="distance",
+                    info="What is being estimated, e.g. distance, width, height.",
+                )
+                unit = gr.Radio(["m", "count", "deg"], value="m", label="Unit (scalar only)")
+                run_btn = gr.Button("Decide", variant="primary")
+
+            with gr.Column(scale=1):
+                answer_md = gr.Markdown(label="Answer")
+                annotated = gr.Image(label="Predicted point (pointing tasks)", height=300)
+                trace_md = gr.Markdown("Decision path will appear here.")
+                with gr.Accordion("Raw result (JSON)", open=False):
+                    raw_json = gr.JSON(label="scores / path")
+
+        gr.Examples(
+            examples=[
+                ["examples/images/scene.png", "choice",
+                 "Where is the red square relative to the blue circle?",
+                 "left | left\nright | right", "distance", "m"],
+                ["examples/images/scene.png", "scalar",
+                 "What is the approximate distance between the two objects?",
+                 "", "distance", "m"],
+                ["examples/images/scene.png", "point",
+                 "Point to the blue circle.", "", "distance", "m"],
+            ],
+            inputs=[image_in, task, question, options_text, quantity, unit],
+            outputs=[answer_md, trace_md, annotated, raw_json],
+            fn=run_spatial,
+            label="Official examples (scene.png from the MJev-Spatial repo)",
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+        gr.Markdown(
+            "Pointing decodes three 3×3 region decisions into a normalized point; the "
+            "score of each stage's picked region is shown in the decision path. "
+            "Scalar questions are answered through an interval hierarchy over meters "
+            "(or an integer count / degree codebook)."
+        )
+
+    run_btn.click(
+        run_spatial,
+        inputs=[image_in, task, question, options_text, quantity, unit],
+        outputs=[answer_md, trace_md, annotated, raw_json],
+        api_name="decide",
+    )
+
+demo.launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,6 @@
+transformers==5.17.0
+accelerate
+numpy
+pillow
+safetensors
+sentencepiece

src/spatial_jev/__init__.py

@@ -0,0 +1,12 @@
+"""Jev-Spatial: a shared classifier for structured spatial decisions."""
+
+__version__ = "0.1.0"
+
+__all__ = ["JevSpatial"]
+
+
+def __getattr__(name):
+    if name == "JevSpatial":
+        from .inference import JevSpatial
+        return JevSpatial
+    raise AttributeError(name)

src/spatial_jev/hierarchy.py

@@ -0,0 +1,162 @@
+"""Deterministic scalar/point classification codebooks and visual ROI geometry."""
+
+import bisect
+import math
+
+import numpy as np
+
+# User-specified order: 正上、左上、右上、正中、左中、右中、正下、左下、右下.
+POINT_NAMES = ("top-center", "top-left", "top-right", "center", "middle-left",
+               "middle-right", "bottom-center", "bottom-left", "bottom-right")
+POINT_RC = ((0, 1), (0, 0), (0, 2), (1, 1), (1, 0), (1, 2), (2, 1), (2, 0), (2, 2))
+ROOT_BOX = (0.0, 0.0, 1.0, 1.0)
+
+
+def child_box(box, choice):
+    if not 0 <= choice < 9:
+        raise ValueError("Invalid nine-way choice")
+    x0, y0, x1, y1 = box
+    row, col = POINT_RC[choice]
+    w, h = (x1 - x0) / 3, (y1 - y0) / 3
+    return (x0 + col * w, y0 + row * h, x0 + (col + 1) * w, y0 + (row + 1) * h)
+
+
+def point_label(xy, box):
+    x0, y0, x1, y1 = box
+    x, y = xy
+    if not all(math.isfinite(v) for v in xy) or not (x0-1e-9 <= x <= x1+1e-9
+                                                       and y0-1e-9 <= y <= y1+1e-9):
+        raise ValueError("Point outside parent")
+    col = min(2, max(0, math.floor((x - x0) / (x1 - x0) * 3 + 1e-12)))
+    row = min(2, max(0, math.floor((y - y0) / (y1 - y0) * 3 + 1e-12)))
+    return POINT_RC.index((row, col))
+
+
+def point_path(xy, depth=3):
+    box, path = ROOT_BOX, []
+    for _ in range(depth):
+        choice = point_label(xy, box)
+        path.append(choice)
+        box = child_box(box, choice)
+    return path
+
+
+def decode_point(path):
+    box = ROOT_BOX
+    for choice in path:
+        box = child_box(box, choice)
+    return [(box[0] + box[2]) / 2, (box[1] + box[3]) / 2]
+
+
+def scalar_fine_bins(root_choice, codebook):
+    """Refinement cardinality is fixed by the selected root, never by the target."""
+    if not 1 <= root_choice < len(codebook["edges_m"]):
+        raise ValueError("Zero/overflow has no refinement candidates")
+    by_root = codebook.get("fine_bins_by_root")
+    if by_root is not None:
+        if len(by_root) != len(codebook["edges_m"]) - 1:
+            raise ValueError("Refinement counts must cover every positive root")
+        fine = by_root[root_choice - 1]
+    else:
+        fine = codebook["fine_bins"]
+    if type(fine) is not int or not 1 <= fine <= 32:
+        raise ValueError("Refinement counts must be integers in [1, 32]")
+    return fine
+
+
+def scalar_path(value, codebook):
+    edges = codebook["edges_m"]
+    if not math.isfinite(value) or value < 0:
+        raise ValueError("Invalid nonnegative scalar")
+    if value == 0:
+        return [0]
+    if value > edges[-1]:
+        return [len(edges)]  # Explicit overflow, never silently clipped.
+    bucket = min(bisect.bisect_right(edges, value) - 1, len(edges) - 2)
+    fine = scalar_fine_bins(bucket + 1, codebook)
+    lo, hi = edges[bucket:bucket+2]
+    local = min(fine - 1, max(0, int((value - lo) / (hi - lo) * fine)))
+    return [bucket + 1, local]
+
+
+def scalar_interval(root_choice, codebook):
+    edges = codebook["edges_m"]
+    if not 1 <= root_choice < len(edges):
+        raise ValueError("Zero/overflow has no refinement interval")
+    return edges[root_choice - 1], edges[root_choice]
+
+
+def decode_scalar(path, codebook):
+    if path[0] == 0:
+        return 0.0
+    if path[0] == len(codebook["edges_m"]):
+        return codebook["overflow_representative_m"]
+    fine = scalar_fine_bins(path[0], codebook)
+    if len(path) != 2 or not 0 <= path[1] < fine:
+        raise ValueError("Missing/invalid scalar refinement")
+    lo, hi = scalar_interval(path[0], codebook)
+    return lo + (path[1] + .5) * (hi - lo) / fine
+
+
+def decode_scalar_expectation(root_choice, fine_logits, codebook):
+    """Conditional mean of fine-bin centers, T=1, within one argmax coarse root.
+
+    No averaging across coarse roots and no target-dependent decisions. Reserved
+    zero/overflow roots retain their existing exact representatives.
+    """
+    if root_choice in (0, len(codebook["edges_m"])):
+        if len(fine_logits):
+            raise ValueError("Zero/overflow must not have a second-stage distribution")
+        return decode_scalar([root_choice], codebook)
+    fine = scalar_fine_bins(root_choice, codebook)
+    logits = np.asarray(fine_logits, dtype=np.float64)
+    if logits.shape != (fine,) or not np.isfinite(logits).all():
+        raise ValueError("Expectation requires exactly the finite valid fine logits")
+    probabilities = np.exp(logits - logits.max())
+    probabilities /= probabilities.sum()
+    centers = np.asarray([decode_scalar([root_choice, i], codebook) for i in range(fine)])
+    return float(probabilities @ centers)
+
+
+def visual_token_boxes(metadata, image_grids, image_num_crops, patches_per_crop):
+    """Map pooled tokens to normalized original-image support boxes, including padding.
+
+    Uses native subpatch_mapping (stretched stitched crop grid), not token-index slicing.
+    Each image's global crop and high-resolution crop mapping have their own geometry.
+    """
+    pooling = np.asarray(metadata["token_pooling"])
+    offset, token_offset, result = 0, 0, []
+    side = math.isqrt(patches_per_crop)
+    if side * side != patches_per_crop:
+        raise ValueError("Expected square ViT crop patch grid")
+    for image_i, (grid, crops) in enumerate(zip(image_grids, image_num_crops)):
+        mapping = np.asarray(metadata["subpatch_mapping"][image_i])
+        h, w = mapping.shape
+        patch_boxes = {}
+        for row in range(side):
+            for col in range(side):
+                patch_boxes[offset + row * side + col] = (
+                    col / side, row / side, (col + 1) / side, (row + 1) / side)
+        for row in range(h):
+            for col in range(w):
+                patch_boxes[int(mapping[row, col])] = (col / w, row / h,
+                                                       (col + 1) / w, (row + 1) / h)
+        count = int(grid[0] * grid[1] + grid[2] * grid[3])
+        for indices in pooling[token_offset:token_offset + count]:
+            boxes = [patch_boxes[int(i)] for i in indices if i >= 0]
+            if not boxes:
+                raise ValueError("Pooled visual token has no valid support")
+            result.append((image_i, min(b[0] for b in boxes), min(b[1] for b in boxes),
+                           max(b[2] for b in boxes), max(b[3] for b in boxes)))
+        token_offset += count
+        offset += int(crops) * patches_per_crop
+    if token_offset != len(pooling):
+        raise ValueError("Unmapped visual tokens")
+    return np.array(result, dtype=np.float32)
+
+
+def roi_visible(boxes, roi, image_index=0):
+    boxes = np.asarray(boxes)
+    x0, y0, x1, y1 = roi
+    return ((boxes[:, 0] == image_index) & (boxes[:, 1] < x1) & (boxes[:, 3] > x0)
+            & (boxes[:, 2] < y1) & (boxes[:, 4] > y0))

src/spatial_jev/kev.py

@@ -0,0 +1,749 @@
+"""Multimodal KEV decisions using the existing spatial schemas and controllers.
+
+Pointer formula and delimiters follow jaredpalmer/kev@f1535963 (Apache-2.0).
+Each original QA is a private causal branch, including its successive crops.
+"""
+
+import hashlib
+import json
+import math
+import random
+import re
+from contextlib import nullcontext
+from functools import lru_cache
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from .hierarchy import ROOT_BOX, child_box, decode_point, decode_scalar_expectation, point_path, scalar_fine_bins
+from .numeric_tasks import ANGLE_CODEBOOK, COUNT_CODEBOOK, codebook_for, decode_numeric, numeric_done, numeric_path
+from .qwen35 import LORA_TARGETS, configure_full_language, sample_seed, train_mode
+from .qwen35 import make_processor as make_processor  # noqa: PLC0414 -- explicit public re-export
+from .roi_replay import (
+    EMPTY_ROI_IMAGE,
+    KV_REPLAY,
+    REPLAY_VARIANTS,
+    expand_replay_inputs,
+    kv_replay_hidden,
+)
+from .unified import UnifiedBuilder
+from .visual_token_pointing import (
+    ORIGINAL_IMAGE,
+    VISUAL_TOKENS,
+    decode_visual_points,
+    image_token_candidates,
+    target_cells,
+    validate_pointing_config,
+)
+
+UPSTREAM_REVISION = "f1535963cea021439370c23127bc970b6788e730"
+SPECIAL = ("<|fim_prefix|>", "<|fim_middle|>", "<|box_start|>", "<|box_end|>", "<|fim_suffix|>")
+IMAGE = "<|vision_start|><|image_pad|><|vision_end|>"
+CONTROL = re.compile(r"<\|([A-Za-z0-9_]+)\|>")
+
+
+def escape(text):
+    return CONTROL.sub(r"<¦\1¦>", text)
+
+
+@lru_cache(maxsize=100000)
+def _file_hash(path, size, mtime_ns):
+    return hashlib.sha256(Path(path).read_bytes()).hexdigest()
+
+
+def media_key(media):
+    # Content identities, not image-pad token IDs or mutable paths, identify images.
+    path = Path(media["uri"])
+    stat = path.stat()
+    return _file_hash(str(path), stat.st_size, stat.st_mtime_ns)
+
+
+def state_key(inp):
+    return tuple(media_key(m) for m in inp["media"])
+
+
+def to_device(value, device):
+    if isinstance(value, torch.Tensor):
+        return value.to(device, non_blocking=True)
+    if isinstance(value, dict):
+        return {k: to_device(v, device) for k, v in value.items()}
+    return value
+
+
+class PointerHead(nn.Module):
+    """KEV's affine query/key projections and scaled dot product, in FP32."""
+
+    def __init__(self, hidden_size, pointer_dim=256):
+        super().__init__()
+        self.q = nn.Linear(hidden_size, pointer_dim)
+        self.k = nn.Linear(hidden_size, pointer_dim)
+        self.scale = 1 / math.sqrt(pointer_dim)
+
+    def forward(self, decisions, options, owner, slot, counts):
+        with torch.autocast(device_type=decisions.device.type, enabled=False):
+            scores = (self.k(options.float()) * self.q(decisions.float())[owner]).sum(-1)
+            scores = scores * self.scale
+            logits = scores.new_full((len(decisions), int(counts.max())), -torch.inf)
+            return logits.index_put((owner, slot), scores)
+
+
+class KevBuilder(UnifiedBuilder):
+    def __init__(self, processor, config):
+        validate_pointing_config(config)
+        if (config.get("numeric_codebooks") is not None
+                and config["numeric_codebooks"] != {"count": COUNT_CODEBOOK, "deg": ANGLE_CODEBOOK}):
+            raise ValueError("Checkpoint numeric codebooks differ from this implementation")
+        self.processor, self.config = processor, config
+        self.codebook, self.variant = config["scalar_codebook"], config.get("point_variant", "crop_refill")
+        if self.codebook.get("fine_bins_by_root") is not None:
+            for root in range(1, len(self.codebook["edges_m"])):
+                if scalar_fine_bins(root, self.codebook) > config["max_choices"]:
+                    raise ValueError("Scalar refinement exceeds model candidate budget")
+        self.visual_token_pointing = self.variant == VISUAL_TOKENS
+        self.markers = [processor.tokenizer.convert_tokens_to_ids(t) for t in SPECIAL]
+        if len(set(self.markers)) != 5:
+            raise ValueError("KEV requires five distinct existing delimiter tokens")
+
+    def stage(self, inp, seed, box=ROOT_BOX, level=0, scalar_root=None):
+        task = inp["answer_space"]["kind"]
+        if task == "point" and self.visual_token_pointing:
+            instruction = f"Request: {inp['question']}\nLocate the requested target in the image."
+            # No artificial region options: the candidates are image-token positions.
+            return SPECIAL[1] + escape(instruction) + SPECIAL[4], 0, []
+        mapping = []
+        if task == "choice":
+            options = list(inp["answer_space"]["options"])
+            if not (self.config.get("preserve_option_order", False)
+                    or inp["answer_space"].get("preserve_option_order", False)):
+                random.Random(seed).shuffle(options)
+            mapping = [o["id"] for o in options]
+            instruction = f"Question: {inp['question']}\nSelect exactly one option."
+            options = [o["text"] for o in options]
+        else:
+            # Use the existing authoritative spatial prompts, including interval semantics.
+            prompt = (
+                self.point_prompt(inp, box, level)
+                if task == "point"
+                else self.scalar_prompt(inp, scalar_root)[0]
+            )
+            instruction, tail = prompt.rsplit("\nOptions:\n", 1)
+            option_text, suffix = tail.rsplit("\nSelect exactly one option.", 1)
+            if suffix:
+                raise ValueError("Unexpected spatial prompt suffix")
+            options = [line.split(". ", 1)[1] for line in option_text.splitlines()]
+            instruction += "\nSelect exactly one option."
+        text = SPECIAL[1] + escape(instruction)
+        text += "".join(SPECIAL[2] + escape(o) + SPECIAL[3] for o in options)
+        return text + SPECIAL[4], len(options), mapping
+
+    def initial_spec(self, inp, identity):
+        if self.variant in (VISUAL_TOKENS, ORIGINAL_IMAGE, *REPLAY_VARIANTS) and inp["answer_space"]["kind"] == "point":
+            if len(inp["media"]) != 1 or inp["media"][0].get("kind", "image") != "image":
+                raise ValueError("Original-image / visual-token pointing requires exactly one still image")
+            if inp["answer_space"].get("coordinate_system", "normalized_xy") != "normalized_xy":
+                raise ValueError("Original-image / visual-token pointing requires normalized_xy coordinates")
+        text, count, mapping = self.stage(inp, sample_seed(identity, self.config["seed"]))
+        keys = state_key(inp)
+        return {
+            "text": SPECIAL[0] + IMAGE * len(keys) + text,
+            "images": self.load_images(inp),
+            "image_keys": list(keys),
+            "state_key": keys,
+            "sample_id": identity,
+            "counts": [count],
+            "mapping": mapping,
+            "box": ROOT_BOX,
+            "path": [],
+            "input": inp,
+        }
+
+    def advance(self, spec, choice):
+        inp, path = spec["input"], spec["path"]
+        task = inp["answer_space"]["kind"]
+        if task == "point" and self.visual_token_pointing:
+            raise ValueError("Visual-token pointing does not advance or crop")
+        extra = ""
+        if task == "point":
+            spec["box"] = child_box(spec["box"], choice)
+            if self.variant in REPLAY_VARIANTS:
+                spec.setdefault('roi_boxes', []).append(spec['box'])
+                extra = EMPTY_ROI_IMAGE
+            elif self.variant != ORIGINAL_IMAGE:
+                spec["images"].append(self.crop(spec["images"][0], spec["box"]))
+                spec["image_keys"].append(spec["state_key"][0] + ":crop:" + repr(spec["box"]))
+                extra = IMAGE
+        text, count, _ = self.stage(
+            inp,
+            self.config["seed"],
+            spec["box"],
+            len(path),
+            scalar_root=path[0] if task == "scalar" else None,
+        )
+        spec["text"] += f"Selected option: {choice + 1}." + extra + text
+        spec["counts"].append(count)
+
+    def training_spec(self, row):
+        spec = self.initial_spec(row["input"], row["sample_id"])
+        kind = row["input"]["answer_space"]["kind"]
+        target = row["target"]
+        if kind == "point" and self.visual_token_pointing:
+            # Labels stay out of text/images; cell indices depend on the processor's grid.
+            spec.update(target_points=target["points"], weight=row.get("_weight", 1.0))
+            return spec
+        path = (
+            [spec["mapping"].index(target["option_id"])]
+            if kind == "choice"
+            else numeric_path(target["value"], codebook_for(row["input"], self.codebook))
+            if kind == "scalar"
+            else point_path(target["points"][0])
+        )
+        for choice in path[:-1]:
+            spec["path"].append(choice)
+            self.advance(spec, choice)
+        spec.update(targets=path, weight=row.get("_weight", 1.0))
+        return spec
+
+    def encode_specs(self, specs, training=True):
+        encoded = dict(
+            self.processor(
+                text=[s["text"] for s in specs],
+                images=[im for s in specs for im in s["images"]],
+                padding=True,
+                add_special_tokens=False,
+                return_tensors="pt",
+                device="cpu",
+            )
+        )
+        encoded.pop("token_type_ids", None)
+        if self.variant in REPLAY_VARIANTS:
+            encoded = expand_replay_inputs(self, specs, encoded)
+        lengths = encoded["attention_mask"].sum(-1).tolist()
+        if max(lengths) > self.config["max_sequence_length"]:
+            raise ValueError(f"KEV sequence budget exceeded without truncation: {lengths}")
+        if "mm_token_type_ids" not in encoded:
+            encoded["mm_token_type_ids"] = (
+                encoded["input_ids"]
+                == self.processor.tokenizer.convert_tokens_to_ids("<|image_pad|>")
+            ).long()
+        positions, option_positions, owner, slot, counts, targets, weights = (
+            [],
+            [],
+            [],
+            [],
+            [],
+            [],
+            [],
+        )
+        prefix_lengths = []
+        image_cursor = 0
+        visual_owners, visual_slots, visual_mass = [], [], []
+        has_visual = False
+        for row, spec in enumerate(specs):
+            ids = encoded["input_ids"][row, : lengths[row]].tolist()
+            prefix_lengths.append(ids.index(self.markers[1]))
+            ends = [i for i, t in enumerate(ids) if t == self.markers[4]]
+            opt_ends = [i for i, t in enumerate(ids) if t == self.markers[3]]
+            first_grid = encoded["image_grid_thw"][image_cursor]
+            image_cursor += len(spec["images"])
+            if self.visual_token_pointing and spec["input"]["answer_space"]["kind"] == "point":
+                if len(ends) != 1 or opt_ends or len(spec["images"]) != 1:
+                    raise ValueError("Visual pointing must have one query, one image and no text candidates")
+                candidates, shape = image_token_candidates(
+                    ids, self.processor.tokenizer.convert_tokens_to_ids("<|image_pad|>"),
+                    first_grid.tolist(), self.processor.image_processor.merge_size,
+                )
+                if candidates[-1] >= prefix_lengths[-1] or prefix_lengths[-1] >= ends[0]:
+                    raise ValueError("Visual keys must precede the private question and query")
+                decision = len(positions)
+                count = len(candidates)
+                positions.append((row, ends[0]))
+                option_positions.extend((row, p) for p in candidates)
+                owner.extend([decision] * count)
+                slot.extend(range(count))
+                counts.append(count)
+                spec.update(counts=[count], visual_grid=shape)
+                has_visual = True
+                if training:
+                    cells = target_cells(spec["target_points"], shape)
+                    targets.append(cells[0])
+                    weights.append(spec.get("weight", 1.0))
+                    visual_owners.extend([decision] * len(cells))
+                    visual_slots.extend(cells)
+                    visual_mass.extend([1 / len(cells)] * len(cells))
+                continue
+            if len(ends) != len(spec["counts"]) or len(opt_ends) != sum(spec["counts"]):
+                raise ValueError(f"Decision/candidate markers mismatch: {spec['sample_id']}")
+            cursor = 0
+            for stage, (pos, count) in enumerate(zip(ends, spec["counts"])):
+                options = opt_ends[cursor : cursor + count]
+                cursor += count
+                if not all((ends[stage - 1] if stage else -1) < p < pos for p in options):
+                    raise ValueError("Candidate belongs to a different decision stage")
+                if not training and stage != len(ends) - 1:
+                    continue
+                owner.extend([len(positions)] * count)
+                slot.extend(range(count))
+                positions.append((row, pos))
+                option_positions.extend((row, p) for p in options)
+                counts.append(count)
+                if training:
+                    targets.append(spec["targets"][stage])
+                    weights.append(spec.get("weight", 1.0) / len(ends))
+        # One frozen vision forward per unique image/crop in this batch. Keep full grids
+        # for each row's native M-RoPE, and an explicit map back to each image occurrence.
+        keys = [key for s in specs for key in s["image_keys"]]
+        chunks = encoded["pixel_values"].split(encoded["image_grid_thw"].prod(-1).tolist())
+        unique, grids, pixels, image_owner = {}, [], [], []
+        for occurrence, (key, grid, pixel) in enumerate(
+            zip(keys, encoded["image_grid_thw"], chunks, strict=True)
+        ):
+            signature = (key, tuple(grid.tolist()))
+            if not self.config.get("deduplicate_vision", True):
+                signature = (*signature, occurrence)
+            if signature not in unique:
+                unique[signature] = len(grids)
+                grids.append(grid)
+                pixels.append(pixel)
+            image_owner.append(unique[signature])
+        encoded["pixel_values"] = torch.cat(pixels)
+        encoded["unique_grid_thw"] = torch.stack(grids)
+        encoded["image_owner"] = torch.tensor(image_owner)
+        packed = {
+            "inputs": encoded,
+            "positions": torch.tensor(positions),
+            "option_positions": torch.tensor(option_positions),
+            "owner": torch.tensor(owner),
+            "slot": torch.tensor(slot),
+            "counts": torch.tensor(counts),
+            "targets": torch.tensor(targets),
+            "weights": torch.tensor(weights),
+            "prefix_lengths": torch.tensor(prefix_lengths),
+            "lengths": torch.tensor(lengths),
+            "state_keys": [s["state_key"] for s in specs],
+            "image_keys": list(unique),
+            "qa_count": sum(s.get("weight", 1.0) for s in specs),
+            "tokens": sum(lengths),
+            "padded_tokens": encoded["input_ids"].numel(),
+        }
+        if self.variant == KV_REPLAY and 'roi_kv_sources' in encoded:
+            packed['roi_kv_copy'] = True
+        if has_visual:
+            packed["visual_pointing"] = True
+            if training:
+                packed.update(visual_target_owner=torch.tensor(visual_owners, dtype=torch.long),
+                              visual_target_slot=torch.tensor(visual_slots, dtype=torch.long),
+                              visual_target_mass=torch.tensor(visual_mass, dtype=torch.float32))
+        return packed
+
+    def collate(self, rows):
+        return self.encode_specs([self.training_spec(row) for row in rows])
+
+
+class KevSpatialModel(nn.Module):
+    def __init__(self, backbone, pointer):
+        super().__init__()
+        self.backbone, self.pointer = backbone, pointer
+        self.language_fp32 = False
+
+    @property
+    def native(self):
+        return (
+            self.backbone.get_base_model()
+            if hasattr(self.backbone, "get_base_model")
+            else self.backbone
+        )
+
+    def train(self, mode=True):
+        super().train(mode)
+        self.native.visual.eval()
+        return self
+
+    def language_context(self, device_type):
+        return (
+            torch.autocast(device_type=device_type, enabled=False)
+            if self.language_fp32
+            else nullcontext()
+        )
+
+    def embeddings(self, inputs, features=None):
+        native = self.native
+        hidden = native.get_input_embeddings()(inputs["input_ids"])
+        if features is None:
+            with torch.no_grad():
+                features = native.get_image_features(
+                    inputs["pixel_values"], inputs["unique_grid_thw"], return_dict=True
+                ).pooler_output
+        features = torch.cat([features[i] for i in inputs["image_owner"].tolist()]).to(hidden.dtype)
+        if 'roi_feature_sources' in inputs:
+            features = features.index_select(0, inputs['roi_feature_sources'])
+        mask, _ = native.get_placeholder_mask(inputs["input_ids"], hidden, image_features=features)
+        hidden = hidden.masked_scatter(mask, features)
+        positions, _ = native.get_rope_index(
+            inputs.get('roi_base_ids', inputs["input_ids"]),
+            inputs.get('roi_base_types', inputs["mm_token_type_ids"]),
+            image_grid_thw=inputs["image_grid_thw"],
+            attention_mask=inputs.get('roi_base_mask', inputs["attention_mask"]),
+        )
+        if 'roi_position_sources' in inputs:
+            positions = positions.gather(2, inputs['roi_position_sources'][None].expand(3, -1, -1))
+        return hidden, positions
+
+    def readout(self, hidden, packed, branch_only=False):
+        if self.pointer is None:
+            raise ValueError("Native QK requires multimodal positions; use model.forward")
+        p, o = packed["positions"], packed["option_positions"]
+        if branch_only:
+            p = p.clone()
+            o = o.clone()
+            p[:, 1] -= packed["prefix_lengths"][p[:, 0]]
+            o[:, 1] -= packed["prefix_lengths"][o[:, 0]]
+        return self.pointer(
+            hidden[p[:, 0], p[:, 1]],
+            hidden[o[:, 0], o[:, 1]],
+            packed["owner"],
+            packed["slot"],
+            packed["counts"],
+        )
+
+    def forward(self, packed, shared_prefix=False):
+        if packed.get("visual_pointing", False) and shared_prefix:
+            raise ValueError("Visual-token keys are in the prefix; use the full causal-row forward")
+        if self.pointer is None and shared_prefix:
+            raise ValueError("Native QK currently requires independent causal rows")
+        hidden, positions = self.embeddings(packed["inputs"])
+        with self.language_context(hidden.device.type):
+            if self.pointer is None and getattr(self.native.language_model.config, "kev_readout_layer_indices", None):
+                from .kev_multilayer_qk import multilayer_forward
+
+                return multilayer_forward(self.native.language_model, hidden, positions,
+                                          packed["inputs"]["attention_mask"], packed)
+            if packed.get('roi_kv_copy', False):
+                hidden = kv_replay_hidden(self.native.language_model, hidden, positions,
+                                         packed['inputs']['attention_mask'], packed['inputs']['roi_kv_sources'])
+            elif shared_prefix:
+                from .kev_prefix import shared_hidden
+
+                hidden = shared_hidden(self.native.language_model, hidden, positions, packed)
+            else:
+                hidden = self.native.language_model(
+                    inputs_embeds=hidden,
+                    position_ids=positions,
+                    attention_mask=packed["inputs"]["attention_mask"],
+                    use_cache=False,
+                ).last_hidden_state
+            if self.pointer is None:
+                lm = self.native.language_model
+                rope_positions = positions[1:] if positions.shape[0] == 4 else positions
+                return lm.native_qk(hidden, packed, lm.rotary_emb(hidden, rope_positions))
+        return self.readout(hidden, packed, branch_only=shared_prefix)
+
+
+def build_model(config, checkpoint=None, trainable=True):
+    from transformers import Qwen3_5ForConditionalGeneration, Qwen3_5TextConfig, Qwen3_5TextModel
+
+    validate_pointing_config(config)
+    mode = train_mode(config)
+    kind = config.get("readout_kind", "added_pointer")
+    if kind not in (
+        "added_pointer",
+        "native_last_qk_mean",
+        "native_merged_qk256",
+        "native_last_qk_concat",
+    ):
+        raise ValueError(f"Unsupported KEV readout: {kind}")
+    native_qk = kind != "added_pointer"
+    fusion_layers = config.get("readout_layer_indices")
+    if fusion_layers is not None and (kind != "native_last_qk_mean"
+            or config.get("layer_fusion") != "equal_logits" or config.get("gradient_checkpointing", False)
+            or config.get("point_variant", "crop_refill") != "crop_refill"):
+        raise ValueError("Multilayer ablation fixes mean heads, equal logits, old pointing and no gradient checkpointing")
+    if checkpoint:
+        saved = json.loads((Path(checkpoint) / "config.json").read_text())
+        if train_mode(saved) != mode:
+            raise ValueError("Cannot load a KEV checkpoint with a different training mode")
+        if saved.get("readout_kind", "added_pointer") != kind:
+            raise ValueError("Cannot load a KEV checkpoint with a different readout")
+        if saved.get("readout_layer_indices") != fusion_layers or saved.get("layer_fusion") != config.get("layer_fusion"):
+            raise ValueError("Cannot load a checkpoint with a different multilayer readout")
+        if saved.get("point_variant", "crop_refill") != config.get("point_variant", "crop_refill"):
+            raise ValueError("Cannot load a checkpoint with a different pointing method")
+    if native_qk and (mode != "full_language" or config.get("shared_prefix", False)):
+        raise ValueError("Native QK requires full_language and independent causal rows")
+    full = Qwen3_5ForConditionalGeneration.from_pretrained(
+        config["base_model"],
+        local_files_only=config.get("local_files_only", False),
+        dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+    )
+    backbone, hidden_size = full.model, full.config.text_config.hidden_size
+    del full
+    backbone.requires_grad_(False)
+    if mode == "full_language":
+        if checkpoint:
+            directory = Path(checkpoint) / "language_model"
+            text_config = Qwen3_5TextConfig.from_pretrained(
+                directory, local_files_only=config.get("local_files_only", False)
+            )
+            text_class = Qwen3_5TextModel
+            if native_qk:
+                if getattr(text_config, "kev_readout_kind", None) != kind:
+                    raise ValueError("Language checkpoint readout marker disagrees with config")
+                if (getattr(text_config, "kev_readout_layer_indices", None) != fusion_layers
+                        or getattr(text_config, "kev_layer_fusion", None) != config.get("layer_fusion")):
+                    raise ValueError("Language checkpoint multilayer markers disagree with config")
+                if getattr(text_config, "kev_readout_fp32", False) != config.get(
+                    "native_readout_fp32", False
+                ):
+                    raise ValueError("Language checkpoint readout precision disagrees with config")
+                expected_layer = config.get(
+                    "readout_layer_index", text_config.num_hidden_layers - 1
+                )
+                if (
+                    getattr(
+                        text_config, "kev_readout_layer_index", text_config.num_hidden_layers - 1
+                    )
+                    != expected_layer
+                ):
+                    raise ValueError("Language checkpoint readout layer disagrees with config")
+                from .kev_native_qk import NativeQKTextModel
+
+                text_class = NativeQKTextModel
+            backbone.language_model, loading = text_class.from_pretrained(
+                directory,
+                config=text_config,
+                local_files_only=config.get("local_files_only", False),
+                dtype=torch.float32,
+                attn_implementation="sdpa",
+                output_loading_info=True,
+            )
+            if any(
+                loading.get(k)
+                for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
+            ):
+                raise ValueError(f"Incomplete language checkpoint load: {loading}")
+            backbone.config.text_config = backbone.language_model.config
+        elif native_qk:
+            from .kev_native_qk import prune_native_qk
+
+            prune_native_qk(backbone.language_model, kind, config.get("readout_layer_index"))
+        configure_full_language(backbone, trainable=trainable)
+        if native_qk:
+            if fusion_layers is not None:
+                from .kev_multilayer_qk import configure_multilayer
+
+                configure_multilayer(backbone.language_model, fusion_layers)
+            backbone.language_model.native_qk.readout_fp32 = config.get(
+                "native_readout_fp32", False
+            )
+            backbone.language_model.config.kev_readout_fp32 = config.get(
+                "native_readout_fp32", False
+            )
+    else:
+        from peft import LoraConfig, PeftModel, get_peft_model
+
+        if checkpoint:
+            backbone = PeftModel.from_pretrained(
+                backbone, str(Path(checkpoint) / "adapter"), is_trainable=trainable
+            )
+        else:
+            backbone = get_peft_model(
+                backbone,
+                LoraConfig(
+                    r=config["lora_rank"],
+                    lora_alpha=config["lora_alpha"],
+                    lora_dropout=config["lora_dropout"],
+                    target_modules=LORA_TARGETS,
+                    bias="none",
+                ),
+            )
+    from .training_alignment import seed_head
+
+    pointer = (
+        None
+        if native_qk
+        else seed_head(config, lambda: PointerHead(hidden_size, config.get("pointer_dim", 256)))
+    )
+    if checkpoint and pointer is not None:
+        from safetensors.torch import load_file
+
+        pointer.load_state_dict(load_file(str(Path(checkpoint) / "pointer.safetensors")))
+    if trainable and config.get("gradient_checkpointing", False):
+        native = backbone.get_base_model() if mode == "lora" else backbone
+        native.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    model = KevSpatialModel(backbone, pointer)
+    inference_dtype = config.get("inference_language_dtype", "bf16")
+    if inference_dtype not in ("bf16", "fp32"):
+        raise ValueError("inference_language_dtype must be bf16 or fp32")
+    if inference_dtype == "fp32":
+        if trainable:
+            raise ValueError("FP32 language compute is an inference verification option")
+        model.native.language_model.float()
+        model.language_fp32 = True
+    if any(p.requires_grad for p in model.native.visual.parameters()):
+        raise ValueError("Vision and merger must remain frozen")
+    if trainable and mode == "full_language":
+        if not all(
+            p.requires_grad and p.dtype == torch.float32
+            for p in model.native.language_model.parameters()
+        ):
+            raise ValueError("Full language training requires every language parameter in FP32")
+        if any("lora_" in n for n, _ in model.named_parameters()):
+            raise ValueError("Full language training must not contain LoRA adapters")
+    if not trainable:
+        model.requires_grad_(False).eval()
+    return model
+
+
+def save_checkpoint(model, processor, config, path, step):
+    from safetensors.torch import save_file
+
+    path = Path(path)
+    path.mkdir(parents=True, exist_ok=False)
+    mode = train_mode(config)
+    if mode == "full_language":
+        model.native.language_model.save_pretrained(
+            path / "language_model", safe_serialization=True, max_shard_size="4GB"
+        )
+    else:
+        model.backbone.save_pretrained(path / "adapter", safe_serialization=True)
+    if model.pointer is not None:
+        save_file(
+            {k: v.detach().cpu().contiguous() for k, v in model.pointer.state_dict().items()},
+            str(path / "pointer.safetensors"),
+        )
+    else:
+        metadata_path = path / "language_model/config.json"
+        metadata = json.loads(metadata_path.read_text())
+        metadata["architectures"] = ["NativeQKTextModel"]
+        metadata_path.write_text(json.dumps(metadata, indent=2) + "\n")
+    processor.save_pretrained(path / "processor")
+    (path / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    (path / "checkpoint.json").write_text(
+        json.dumps(
+            {
+                "step": step,
+                "format": "kev-spatial/v1",
+                "train_mode": mode,
+                "readout_kind": config.get("readout_kind", "added_pointer"),
+                "point_variant": config.get("point_variant", "crop_refill"),
+                "readout_layer_index": getattr(
+                    model.native.language_model.config, "kev_readout_layer_index", None
+                ),
+                "readout_layer_indices": config.get("readout_layer_indices"),
+                "layer_fusion": config.get("layer_fusion"),
+                "checkpoint_policy": config.get("checkpoint_policy", "full"),
+                "upstream_revision": UPSTREAM_REVISION,
+                "base_revision": config["base_revision"],
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+
+
+def validate_checkpoint_parameters(model, config, checkpoint):
+    """Compare saved tensors shard by shard without loading a second full model."""
+    from safetensors import safe_open
+    from safetensors.torch import load_file
+
+    checkpoint = Path(checkpoint)
+    if train_mode(config) == "full_language":
+        loaded = model.native.language_model.state_dict()
+        seen = set()
+        for shard in sorted((checkpoint / "language_model").glob("*.safetensors")):
+            with safe_open(shard, framework="pt", device="cpu") as stream:
+                for name in stream.keys():  # noqa: SIM118 -- safetensors reader is not a mapping
+                    if name in seen or name not in loaded:
+                        raise ValueError(f"Unexpected checkpoint parameter: {name}")
+                    tensor = stream.get_tensor(name)
+                    if not torch.equal(tensor, loaded[name].detach().cpu()):
+                        raise ValueError(f"Language checkpoint round trip differs: {name}")
+                    seen.add(name)
+        if seen != loaded.keys():
+            raise ValueError("Language checkpoint is missing state tensors")
+    else:
+        from peft import get_peft_model_state_dict
+
+        saved = load_file(str(checkpoint / "adapter/adapter_model.safetensors"))
+        loaded = get_peft_model_state_dict(model.backbone)
+        if saved.keys() != loaded.keys() or any(
+            not torch.equal(saved[k], loaded[k].detach().cpu()) for k in saved
+        ):
+            raise ValueError("Adapter checkpoint round trip is not exact")
+    if model.pointer is None:
+        if (checkpoint / "pointer.safetensors").exists():
+            raise ValueError("Native QK checkpoint must not contain an added pointer")
+        return
+    head = load_file(str(checkpoint / "pointer.safetensors"))
+    loaded_head = model.pointer.state_dict()
+    if head.keys() != loaded_head.keys() or any(
+        not torch.equal(v, loaded_head[k].detach().cpu()) for k, v in head.items()
+    ):
+        raise ValueError("Pointer checkpoint round trip is not exact")
+
+
+@torch.inference_mode()
+def predict_batch(model, builder, requests, device, runtime=None, point_decoding=None, scalar_decoding="argmax"):
+    """Free rollout: only requests (no targets), with a private crop history per QA."""
+    visual = getattr(builder, "visual_token_pointing", False)
+    if scalar_decoding not in ("argmax", "fine_expectation"):
+        raise ValueError("Unknown scalar decoding mode")
+    if point_decoding is not None and (not visual or point_decoding not in ("argmax", "components")):
+        raise ValueError("point_decoding overrides require visual-token pointing")
+    if visual and runtime is not None:
+        raise ValueError("Visual-token pointing currently uses the uncached causal-row forward")
+    states, results = [], [None] * len(requests)
+    for index, request in enumerate(requests):
+        spec = builder.initial_spec(request["input"], request["id"])
+        spec.update(index=index, logits=[])
+        states.append(spec)
+    while states:
+        packed = to_device(builder.encode_specs(states, training=False), device)
+        logits = model(packed) if runtime is None else runtime(packed)
+        choices = logits.argmax(-1).tolist()
+        following = []
+        for state, choice, scores in zip(states, choices, logits.float().cpu()):
+            path, inp = state["path"], state["input"]
+            if visual and inp["answer_space"]["kind"] == "point":
+                values = scores[:state["counts"][-1]].tolist()
+                results[state["index"]] = {
+                    **decode_visual_points(values, state["visual_grid"],
+                        decoder=point_decoding or builder.config.get("point_decoder", "argmax"),
+                        relative_threshold=builder.config.get("point_component_threshold", 0.1)),
+                    "path": [], "logits": [values],
+                }
+                continue
+            path.append(choice)
+            state["logits"].append(scores[: state["counts"][-1]].tolist())
+            kind = inp["answer_space"]["kind"]
+            done = (
+                kind == "choice"
+                or (kind == "point" and len(path) == 3)
+                or (
+                    kind == "scalar"
+                    and numeric_done(path, codebook_for(inp, builder.codebook))
+                )
+            )
+            if done:
+                prediction = (
+                    state["mapping"][choice]
+                    if kind == "choice"
+                    else decode_point(path)
+                    if kind == "point"
+                    else decode_numeric(path, codebook_for(inp, builder.codebook))
+                )
+                if (kind == "scalar" and inp["answer_space"].get("unit", "m") == "m"
+                        and scalar_decoding == "fine_expectation"):
+                    prediction = decode_scalar_expectation(
+                        path[0], state["logits"][-1] if len(path) == 2 else [], builder.codebook)
+                results[state["index"]] = {
+                    "prediction": prediction,
+                    "path": path,
+                    "logits": state["logits"],
+                }
+            else:
+                builder.advance(state, choice)
+                following.append(state)
+        states = following
+    return results

src/spatial_jev/kev_native_qk.py

@@ -0,0 +1,155 @@
+"""Spatial decisions from the pretrained final attention's native Q/K scores.
+
+Only existing parameters survive: the first N-1 layers and the last input norm,
+Q-only rows, K projection and Q/K norms. No final model norm or learned pointer.
+"""
+
+from contextlib import nullcontext
+
+import torch
+from torch import nn
+from transformers import Qwen3_5TextModel
+from transformers.models.qwen3_5.modeling_qwen3_5 import rotate_half
+
+NATIVE_QK = "native_last_qk_mean"
+MERGED_QK = "native_merged_qk256"
+CONCAT_QK = "native_last_qk_concat"
+READOUT_KINDS = (NATIVE_QK, MERGED_QK, CONCAT_QK)
+
+
+class NativeQKReadout(nn.Module):
+    def __init__(self, layer, merge_heads=False, concat_heads=False):
+        super().__init__()
+        if merge_heads and concat_heads:
+            raise ValueError("Weight merging and feature concatenation are distinct readouts")
+        self.concat_heads = concat_heads
+        if layer.block_type != "full_attention":
+            raise ValueError("Native QK readout requires a final full-attention layer")
+        attention = layer.self_attn
+        if attention.q_proj.bias is not None or attention.k_proj.bias is not None:
+            raise ValueError("This readout expects the Qwen3.5 bias-free projections")
+        self.head_dim = attention.head_dim
+        self.groups = attention.num_key_value_groups
+        self.heads = attention.q_proj.out_features // (2 * self.head_dim)
+        self.kv_heads = self.heads // self.groups
+        self.scale = attention.scaling
+        self.readout_fp32 = False
+        self.input_layernorm = layer.input_layernorm
+        self.k_proj = attention.k_proj
+        self.q_norm, self.k_norm = attention.q_norm, attention.k_norm
+        # Q and output gate are interleaved PER HEAD, not two global halves.
+        weight = attention.q_proj.weight
+        selected = weight.view(self.heads, 2 * self.head_dim, -1)[:, : self.head_dim]
+        if merge_heads:
+            # Average PRETRAINED WEIGHTS in FP32, before applying the shared
+            # Q/K Norms. Averaging the four GQA keys equals averaging their
+            # sixteen aligned copies. No learned fusion or random initialization.
+            selected = selected.float().mean(dim=0)
+            key_weight = attention.k_proj.weight
+            merged_key = key_weight.view(self.kv_heads, self.head_dim, -1).float().mean(dim=0)
+            self.k_proj = nn.Linear(weight.shape[1], self.head_dim, bias=False, device="meta")
+            self.k_proj.weight = nn.Parameter(
+                merged_key.detach().clone(), requires_grad=key_weight.requires_grad
+            )
+            self.heads = self.kv_heads = self.groups = 1
+        self.q_proj = nn.Linear(
+            weight.shape[1], self.heads * self.head_dim, bias=False, device="meta"
+        )
+        self.q_proj.weight = nn.Parameter(
+            selected.reshape(self.heads * self.head_dim, -1).detach().clone(),
+            requires_grad=weight.requires_grad,
+        )
+
+    def readout_parameters(self):
+        """Merged projections and their norms; input_layernorm stays in backbone."""
+        for module in (self.q_proj, self.k_proj, self.q_norm, self.k_norm):
+            yield from module.parameters()
+
+    @staticmethod
+    def rotate(states, indices, cos, sin):
+        # Same partial RoPE as native apply_rotary_pos_emb; gather original
+        # multimodal positions rather than renumbering selected candidates.
+        cos = cos[indices[:, 0], indices[:, 1]][:, None]
+        sin = sin[indices[:, 0], indices[:, 1]][:, None]
+        width = cos.shape[-1]
+        rotated, unchanged = states[..., :width], states[..., width:]
+        return torch.cat((rotated * cos + rotate_half(rotated) * sin, unchanged), dim=-1)
+
+    def forward(self, hidden, packed, position_embeddings):
+        context = (
+            torch.autocast(device_type=hidden.device.type, enabled=False)
+            if self.readout_fp32
+            else nullcontext()
+        )
+        with context:
+            return self._scores(hidden, packed, position_embeddings)
+
+    def _scores(self, hidden, packed, position_embeddings):
+        p, o = packed["positions"], packed["option_positions"]
+        decisions, options = hidden[p[:, 0], p[:, 1]], hidden[o[:, 0], o[:, 1]]
+        if self.readout_fp32:
+            decisions, options = decisions.float(), options.float()
+            position_embeddings = tuple(p.float() for p in position_embeddings)
+        decisions = self.input_layernorm(decisions)
+        options = self.input_layernorm(options)
+        queries = self.q_norm(self.q_proj(decisions).view(-1, self.heads, self.head_dim))
+        keys = self.k_norm(self.k_proj(options).view(-1, self.kv_heads, self.head_dim))
+        queries = self.rotate(queries, p, *position_embeddings)
+        keys = self.rotate(keys, o, *position_embeddings)
+        with torch.autocast(device_type=hidden.device.type, enabled=False):
+            keys = keys.float().repeat_interleave(self.groups, dim=1)
+            if self.concat_heads:
+                # Retain all per-head Norm/RoPE operations, then concatenate
+                # channels on BOTH sides. Repeated GQA keys stay weight-tied.
+                # For H heads this is sqrt(H) times the mean-head logit,
+                # not cross-head mixing. No additional norm or learned scale.
+                query_features = queries.float().flatten(start_dim=1)
+                key_features = keys.flatten(start_dim=1)
+                scores = (query_features[packed["owner"]] * key_features).sum(-1)
+                scores = scores * (self.heads * self.head_dim) ** -0.5
+            else:
+                scores = (queries.float()[packed["owner"]] * keys).sum(-1) * self.scale
+                # Average logits, NOT per-head probabilities. Candidate masking
+                # is independent for every QA/stage; padded slots carry no mass.
+                scores = scores.mean(dim=-1)
+            logits = scores.new_full((len(p), int(packed["counts"].max())), -torch.inf)
+            return logits.index_put((packed["owner"], packed["slot"]), scores)
+
+
+def prune_native_qk(language_model, kind=NATIVE_QK, layer_index=None):
+    if kind not in READOUT_KINDS:
+        raise ValueError(f"Unsupported native readout: {kind}")
+    if hasattr(language_model, "native_qk"):
+        raise ValueError("Native QK pruning was already applied")
+    layer_index = len(language_model.layers) - 1 if layer_index is None else layer_index
+    if not isinstance(layer_index, int) or not 0 <= layer_index < len(language_model.layers):
+        raise ValueError(f"Invalid readout layer index: {layer_index}")
+    last = language_model.layers[layer_index]
+    language_model.native_qk = NativeQKReadout(
+        last, merge_heads=kind == MERGED_QK, concat_heads=kind == CONCAT_QK
+    )
+    language_model.native_qk.readout_fp32 = getattr(
+        language_model.config, "kev_readout_fp32", False
+    )
+    language_model.layers = nn.ModuleList(list(language_model.layers[:layer_index]))
+    # Execute only the original layers BEFORE the selected full attention.
+    # Later layers cannot affect this readout and must not enter the optimizer.
+    language_model.norm = nn.Identity()
+    # Keep the original layer configuration for checkpoint reconstruction.
+    language_model.config.kev_readout_kind = kind
+    language_model.config.kev_readout_layer_index = layer_index
+    return language_model
+
+
+class NativeQKTextModel(Qwen3_5TextModel):
+    """Explicit loader for the pruned checkpoint; never load it as an ordinary LM."""
+
+    def __init__(self, config):
+        if getattr(config, "kev_readout_kind", None) not in READOUT_KINDS:
+            raise ValueError("Missing native QK checkpoint architecture marker")
+        super().__init__(config)
+        prune_native_qk(
+            self,
+            config.kev_readout_kind,
+            getattr(config, "kev_readout_layer_index", config.num_hidden_layers - 1),
+        )

src/spatial_jev/kev_prefix.py

@@ -0,0 +1,325 @@
+"""Functional hybrid prefixes with native multimodal embeddings and 3-axis RoPE.
+
+Prefix and the layer-wise training schedule are adapted from
+jaredpalmer/kev@f1535963cea021439370c23127bc970b6788e730/kev/shared_prefix.py
+(Apache-2.0; see third_party/kev/). Branches never mutate a saved prefix.
+"""
+
+import os
+from collections import OrderedDict
+
+import torch
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+class Prefix:
+    record_past = True
+
+    def __init__(self):
+        self.keys = self.values = self.conv = None
+        self.recurrent_states = {}
+
+    @property
+    def layers(self):
+        return self
+
+    def __getitem__(self, layer_idx):
+        return self
+
+    def has_previous_state(self, layer_idx=None, state_idx=0):
+        return self.conv is not None
+
+    def update_conv_state(self, conv_input, layer_idx, conv_kernel_size, **kwargs):
+        if self.conv is not None:
+            return torch.cat([self.conv, conv_input], -1)
+        self.conv = F.pad(conv_input, (max(0, conv_kernel_size - conv_input.shape[-1]), 0))[
+            ..., -conv_kernel_size:
+        ]
+        return conv_input
+
+    def update_recurrent_state(self, recurrent_state, layer_idx, **kwargs):
+        self.recurrent_states.setdefault(0, recurrent_state)
+        return recurrent_state
+
+    def update(self, keys, values, layer_idx, *args, **kwargs):
+        if self.keys is None:
+            self.keys, self.values = keys, values
+            return keys, values
+        return torch.cat([self.keys, keys], -2), torch.cat([self.values, values], -2)
+
+    def branch(self, owner):
+        other = Prefix()
+        pick = lambda t: None if t is None else t.index_select(0, owner)
+        other.keys, other.values, other.conv = pick(self.keys), pick(self.values), pick(self.conv)
+        other.recurrent_states = {k: pick(v) for k, v in self.recurrent_states.items()}
+        return other
+
+    def nbytes(self):
+        tensors = [self.keys, self.values, self.conv, *self.recurrent_states.values()]
+        return sum(t.numel() * t.element_size() for t in tensors if t is not None)
+
+    def compact(self):
+        # Cached slices must own their storage so byte accounting does not hide a
+        # complete prefill tensor behind a small convolution window or KV view.
+        self.keys = None if self.keys is None else self.keys.clone().contiguous()
+        self.values = None if self.values is None else self.values.clone().contiguous()
+        self.conv = None if self.conv is None else self.conv.clone().contiguous()
+        self.recurrent_states = {
+            k: v.clone().contiguous() for k, v in self.recurrent_states.items()
+        }
+        return self
+
+
+def attention_mask(allow, hidden, implementation):
+    allow = allow[:, None]
+    if implementation == "sdpa":
+        return allow
+    return hidden.new_zeros(allow.shape).masked_fill(~allow, torch.finfo(hidden.dtype).min)
+
+
+def split_rows(hidden, positions, packed):
+    prefixes = packed["prefix_lengths"].tolist()
+    lengths = packed["lengths"].tolist()
+    unique, first, owners = {}, [], []
+    for i, key in enumerate(packed["state_keys"]):
+        if key not in unique:
+            unique[key] = len(first)
+            first.append(i)
+        reference = first[unique[key]]
+        n = prefixes[i]
+        if n != prefixes[reference] or not torch.equal(
+            packed["inputs"]["input_ids"][i, :n], packed["inputs"]["input_ids"][reference, :n]
+        ):
+            raise ValueError("Same state identity compiled to different prefix tokens")
+        if not torch.equal(positions[:, i, :n], positions[:, reference, :n]):
+            raise ValueError("Same state identity compiled to different multimodal positions")
+        owners.append(unique[key])
+    ns, nb = max(prefixes), max(n - p for n, p in zip(lengths, prefixes))
+    s = hidden.new_zeros((len(first), ns, hidden.shape[-1]))
+    b = hidden.new_zeros((len(lengths), nb, hidden.shape[-1]))
+    sp = positions.new_zeros((3, len(first), ns))
+    bp = positions.new_zeros((3, len(lengths), nb))
+    sm = torch.zeros(s.shape[:2], dtype=torch.bool, device=hidden.device)
+    bm = torch.zeros(b.shape[:2], dtype=torch.bool, device=hidden.device)
+    for j, i in enumerate(first):
+        n = prefixes[i]
+        s[j, -n:], sp[:, j, -n:], sm[j, -n:] = hidden[i, :n], positions[:, i, :n], True
+    for i, (p, n) in enumerate(zip(prefixes, lengths)):
+        b[i, : n - p], bp[:, i, : n - p], bm[i, : n - p] = (
+            hidden[i, p:n],
+            positions[:, i, p:n],
+            True,
+        )
+    return s, b, sp, bp, sm, bm, torch.tensor(owners, device=hidden.device)
+
+
+def masks(lm, s, b, sm, bm, owner):
+    ns, nb, device = s.shape[1], b.shape[1], s.device
+    causal = lambda n: torch.ones(n, n, dtype=torch.bool, device=device).tril()
+    eye = lambda n: torch.eye(n, dtype=torch.bool, device=device)
+    impl = lm.config._attn_implementation
+    s_allow = (causal(ns)[None] & sm[:, None, :]) | eye(ns)[None]
+    b_allow = torch.cat(
+        [
+            sm[owner][:, None, :].expand(-1, nb, -1),
+            (causal(nb)[None] & bm[:, None, :]) | eye(nb)[None],
+        ],
+        -1,
+    )
+    ms = {
+        "full_attention": None
+        if impl == "sdpa" and bool(sm.all())
+        else attention_mask(s_allow, s, impl),
+        "linear_attention": sm.to(s.dtype),
+    }
+    mb = {"full_attention": attention_mask(b_allow, b, impl), "linear_attention": bm.to(b.dtype)}
+    return ms, mb
+
+
+def shared_hidden(lm, hidden, positions, packed):
+    s, b, sp, bp, sm, bm, owner = split_rows(hidden, positions, packed)
+    ms, mb = masks(lm, s, b, sm, bm, owner)
+    rs, rb = lm.rotary_emb(s, sp), lm.rotary_emb(b, bp)
+    for layer in lm.layers:
+        # Capture the current layer, never a late-bound loop variable in recomputation.
+        def step(s, b, layer=layer):
+            prefix = Prefix()
+            s = layer.forward(
+                s,
+                position_embeddings=rs,
+                attention_mask=ms[layer.block_type],
+                past_key_values=prefix,
+            )
+            b = layer.forward(
+                b,
+                position_embeddings=rb,
+                attention_mask=mb[layer.block_type],
+                past_key_values=prefix.branch(owner),
+            )
+            return s * sm[..., None], b
+
+        if lm.gradient_checkpointing and lm.training and torch.is_grad_enabled():
+            s, b = checkpoint(step, s, b, use_reentrant=False)
+        else:
+            s, b = step(s, b)
+    return lm.norm(b)
+
+
+class ByteLRU:
+    def __init__(self, max_bytes):
+        self.max_bytes, self.bytes = max_bytes, 0
+        self.entries = OrderedDict()
+        self.hits = self.misses = 0
+
+    def get(self, key):
+        if key not in self.entries:
+            self.misses += 1
+            return None
+        self.hits += 1
+        self.entries.move_to_end(key)
+        return self.entries[key][0]
+
+    def put(self, key, value, size):
+        if size > self.max_bytes:
+            return
+        if key in self.entries:
+            self.bytes -= self.entries.pop(key)[1]
+        while self.entries and self.bytes + size > self.max_bytes:
+            self.bytes -= self.entries.popitem(last=False)[1][1]
+        self.entries[key] = (value, size)
+        self.bytes += size
+
+    def clear(self):
+        self.entries.clear()
+        self.bytes = 0
+
+
+class KevRuntime:
+    """An immutable-model session with byte-bounded image and hybrid prefix caches.
+
+    Every round recomputes only the private question history from the original state.
+    All branches of one state are batched; attention/conv/recurrent prefixes are read-only.
+    """
+
+    def __init__(self, model, config, prefix_bytes=8 * 1024**3, image_bytes=2 * 1024**3,
+                 enable_native_qk=False):
+        from .roi_replay import REPLAY_VARIANTS
+        if config.get('point_variant') in REPLAY_VARIANTS:
+            raise ValueError('ROI replay uses uncached full rows, not the persistent hybrid-prefix runtime')
+        if config.get("point_variant") == "visual_tokens":
+            raise ValueError("Visual-token pointing requires prefix keys; cached branch-only readout is unsupported")
+        if not enable_native_qk and config.get("readout_kind") in (
+            "native_last_qk_mean",
+            "native_merged_qk256",
+            "native_last_qk_concat",
+        ):
+            raise ValueError("Native QK currently uses the uncached causal-row forward")
+        if enable_native_qk and (config.get("readout_kind") != "native_last_qk_mean"
+                or config.get("readout_layer_indices") is not None
+                or config.get("gradient_checkpointing", False)):
+            raise ValueError("Experimental native cache requires single-layer mean and no checkpointing")
+        import hashlib
+        import json
+
+        self.model = model
+        self.signature = hashlib.sha256(json.dumps(config, sort_keys=True).encode()).hexdigest()
+        self.versions = tuple((p._version, p.dtype, p.device) for p in model.parameters())
+        self.prefix_cache, self.image_cache = ByteLRU(prefix_bytes), ByteLRU(image_bytes)
+
+    def clear(self):
+        self.prefix_cache.clear()
+        self.image_cache.clear()
+
+    @torch.inference_mode()
+    def __call__(self, packed):
+        if packed.get("visual_pointing", False):
+            raise ValueError("Visual-token pointing requires the uncached causal-row forward")
+        if self.model.training or self.versions != tuple(
+            (p._version, p.dtype, p.device) for p in self.model.parameters()
+        ):
+            raise ValueError("Create a fresh runtime after changing model weights or training mode")
+        model, inputs = self.model, packed["inputs"]
+        device_type = inputs["input_ids"].device.type
+        precision = (
+            torch.is_autocast_enabled(device_type),
+            str(torch.get_autocast_dtype(device_type)),
+            torch.get_float32_matmul_precision(),
+            os.environ.get("TRITON_F32_DEFAULT", "tf32"),
+        )
+        scope = (self.signature, precision)
+        features = [self.image_cache.get((scope, k)) for k in packed["image_keys"]]
+        missing = [i for i, f in enumerate(features) if f is None]
+        if missing:
+            chunks = inputs["pixel_values"].split(inputs["unique_grid_thw"].prod(-1).tolist())
+            fs = model.native.get_image_features(
+                torch.cat([chunks[i] for i in missing]),
+                inputs["unique_grid_thw"][missing],
+                return_dict=True,
+            ).pooler_output
+            for i, f in zip(missing, fs):
+                f = f.clone().contiguous()
+                features[i] = f
+                self.image_cache.put(
+                    (scope, packed["image_keys"][i]), f, f.numel() * f.element_size()
+                )
+        hidden, positions = model.embeddings(inputs, features=features)
+        with model.language_context(hidden.device.type):
+            return self.language_forward(packed, hidden, positions, scope)
+
+    def language_forward(self, packed, hidden, positions, scope):
+        model, inputs = self.model, packed["inputs"]
+        s, b, sp, bp, sm, bm, owner = split_rows(hidden, positions, packed)
+        lm = model.native.language_model
+        result = torch.zeros_like(b)
+        for state in range(len(s)):
+            rows = (owner == state).nonzero().flatten()
+            first = int(rows[0])
+            ns = int(packed["prefix_lengths"][first])
+            key = (
+                scope,
+                packed["state_keys"][first],
+                str(hidden.dtype),
+                tuple(inputs["input_ids"][first, :ns].tolist()),
+                tuple(positions[:, first, :ns].flatten().tolist()),
+            )
+            entry = self.prefix_cache.get(key)
+            if entry is None:
+                hs = s[state : state + 1, -ns:]
+                rs = lm.rotary_emb(hs, sp[:, state : state + 1, -ns:])
+                entry = []
+                for layer in lm.layers:
+                    prefix = Prefix()
+                    hs = layer.forward(hs, position_embeddings=rs, past_key_values=prefix)
+                    entry.append(prefix.compact())
+                self.prefix_cache.put(key, entry, sum(p.nbytes() for p in entry))
+            # Different states need not inherit the longest unrelated question's
+            # padding. Their final readouts only index these real branch lengths.
+            branch_length = int(bm[rows].sum(-1).max())
+            hb, mask = b[rows, :branch_length], bm[rows, :branch_length]
+            rb = lm.rotary_emb(hb, bp[:, rows, :branch_length])
+            owners = torch.zeros(len(rows), dtype=torch.long, device=hb.device)
+            _, mb = masks(
+                lm, s[state : state + 1, -ns:], hb, sm[state : state + 1, -ns:], mask, owners
+            )
+            for layer, prefix in zip(lm.layers, entry):
+                hb = layer.forward(
+                    hb,
+                    position_embeddings=rb,
+                    attention_mask=mb[layer.block_type],
+                    past_key_values=prefix.branch(owners),
+                )
+            result[rows, :branch_length] = lm.norm(hb)
+        if model.pointer is None:
+            # The cached image prefix is unchanged. Candidates and decisions are
+            # private text suffix tokens; retain their ORIGINAL M-RoPE coordinates
+            # while translating tensor indices into the returned suffix buffer.
+            selected = dict(packed)
+            for key in ("positions", "option_positions"):
+                indices = packed[key].clone()
+                indices[:, 1] -= packed["prefix_lengths"][indices[:, 0]]
+                if (indices[:, 1] < 0).any():
+                    raise ValueError("Native suffix cache cannot read keys inside the prefix")
+                selected[key] = indices
+            return lm.native_qk(result, selected, lm.rotary_emb(result, bp))
+        return model.readout(result, packed, branch_only=True)

src/spatial_jev/numeric_tasks.py

@@ -0,0 +1,71 @@
+"""Input-selected numeric codebooks; never interpret counts or degrees as meters."""
+
+import math
+
+COUNT_CODEBOOK = {"version": "integer-32x32-v1", "unit": "count", "fine_bins": 32,
+                  "maximum": 1023}
+ANGLE_CODEBOOK = {"version": "degrees-30x32-v1", "unit": "deg", "fine_bins": 32,
+                  "edges_m": list(range(0, 361, 30)), "overflow_representative_m": 360.0}
+# edges_m is the historical controller's internal key, not the physical unit.
+# Public answer spaces, prompts, validators and decoded units remain explicit.
+
+
+def codebook_for(inp, metric):
+    unit = inp["answer_space"].get("unit", "m")
+    if unit == "m":
+        return metric
+    if unit == "count":
+        return COUNT_CODEBOOK
+    if unit == "deg":
+        return ANGLE_CODEBOOK
+    raise ValueError(f"Unsupported numeric unit: {unit}")
+
+
+def numeric_path(value, codebook):
+    from .hierarchy import scalar_path
+    if codebook.get("unit", "m") != "count":
+        if codebook.get("unit") == "deg" and not 0 <= value <= 360:
+            raise ValueError("Angle outside [0, 360]")
+        return scalar_path(value, codebook)
+    if isinstance(value, bool) or not math.isfinite(value) or int(value) != value or not 0 <= value <= 1023:
+        raise ValueError("Count outside integer codebook [0, 1023]")
+    return [int(value) // 32, int(value) % 32]
+
+
+def numeric_done(path, codebook):
+    return len(path) == 2 or (codebook.get("unit", "m") != "count" and path[0] in (0, len(codebook["edges_m"])))
+
+
+def decode_numeric(path, codebook):
+    from .hierarchy import decode_scalar
+    if codebook.get("unit", "m") != "count":
+        return decode_scalar(path, codebook)
+    if len(path) != 2 or any(not 0 <= n < 32 for n in path):
+        raise ValueError("Invalid integer decision path")
+    return 32 * path[0] + path[1]
+
+
+def numeric_prompt(inp, root, codebook):
+    """Non-metric prompts. Metric prompts intentionally keep their old exact bytes."""
+    if codebook["unit"] == "count":
+        options = ([f"count {32*i} through {32*i+31} (inclusive)" for i in range(32)]
+                   if root is None else [f"count exactly {32*root+i}" for i in range(32)])
+        instruction = "Answer the counting question. Select the integer count or inclusive count range."
+    else:
+        from .hierarchy import scalar_interval
+        edges = codebook["edges_m"]
+        if root is None:
+            options = ["exactly zero degrees"]
+            options += [f"{lo:g} to {hi:g} degrees" for lo, hi in zip(edges, edges[1:])]
+            options[1] = "greater than zero and less than 30 degrees"
+            options += ["greater than 360 degrees (out of supported range)"]
+        else:
+            lo, hi = scalar_interval(root, codebook)
+            step = (hi - lo) / codebook["fine_bins"]
+            options = [f"{lo+i*step:.6g} to {lo+(i+1)*step:.6g} degrees" for i in range(32)]
+        instruction = ("Estimate the angle in degrees, preserving the reference frame and rotation "
+                       "direction specified in the question. Ranges include the lower endpoint and "
+                       "exclude the upper, except that 360 is included. Exact zero has its own class.")
+    return (f"Question: {inp['question']}\n{instruction}\nOptions:\n"
+            + "\n".join(f"{i+1}. {s}" for i, s in enumerate(options))
+            + "\nSelect exactly one option."), len(options)

src/spatial_jev/point_probability_map.py

@@ -0,0 +1,125 @@
+"""Nine-way conditional trees and connected-component multi-point readout."""
+
+import numpy as np
+
+from .hierarchy import POINT_RC
+
+
+def softmax9(logits):
+    values = np.asarray(logits, dtype=np.float64)
+    if values.shape != (9,) or not np.isfinite(values).all():
+        raise ValueError("Expected nine finite logits")
+    probabilities = np.exp(values - values.max())
+    return probabilities / probabilities.sum()
+
+
+def grid_location(path):
+    row = col = 0
+    for choice in path:
+        dr, dc = POINT_RC[choice]
+        row, col = row * 3 + dr, col * 3 + dc
+    return row, col
+
+
+def branch_spec(builder, parent, choice):
+    # Copy mutable histories, but share immutable PIL images and input metadata.
+    child = dict(parent)
+    for key in ("images", "image_keys", "counts", "mapping", "path"):
+        child[key] = list(parent[key])
+    if 'roi_boxes' in parent:
+        child['roi_boxes'] = list(parent['roi_boxes'])
+    child["path"].append(choice)
+    builder.advance(child, choice)
+    return child
+
+
+def probability_map(builder, request, evaluate_nodes, mode="full", root_relative_cutoff=0.01):
+    """Three levels, 729 leaves; labels/masks/counts are not accepted as arguments.
+
+    root_pruned uses only the first decision to omit low-mass root subtrees.
+    Unvisited leaves remain explicitly unknown, with their missing mass reported.
+    No probability renormalization is applied after pruning.
+    """
+    if mode not in ("full", "root_pruned"):
+        raise ValueError("Unknown tree policy")
+    if not 0 <= root_relative_cutoff <= 1:
+        raise ValueError("Invalid root cutoff")
+    if request["input"]["answer_space"]["kind"] != "point":
+        raise ValueError("Point input required")
+    grid = np.zeros((27, 27), dtype=np.float64)
+    known = np.zeros((27, 27), dtype=bool)
+    nodes, pruned = [], []
+    frontier = [(builder.initial_spec(request["input"], request["id"]), 1.0)]
+    for depth in range(3):
+        logits = evaluate_nodes([state for state, _ in frontier])
+        if len(logits) != len(frontier):
+            raise ValueError("Node evaluator returned a different batch size")
+        following = []
+        for (state, mass), scores in zip(frontier, logits):
+            probabilities = softmax9(scores)
+            path = list(state["path"])
+            nodes.append({"path": path, "mass": mass, "logits": list(map(float, scores)),
+                          "probabilities": probabilities.tolist()})
+            cutoff = probabilities.max() * root_relative_cutoff
+            for choice, probability in enumerate(probabilities):
+                child_path = [*path, choice]
+                child_mass = mass * float(probability)
+                if depth == 2:
+                    row, col = grid_location(child_path)
+                    if known[row, col]:
+                        raise ValueError("Duplicate leaf location")
+                    grid[row, col], known[row, col] = child_mass, True
+                elif depth == 0 and mode == "root_pruned" and probability < cutoff:
+                    pruned.append({"path": child_path, "mass": child_mass})
+                else:
+                    following.append((branch_spec(builder, state, choice), child_mass))
+        frontier = following
+    retained, omitted = float(grid.sum()), sum(r["mass"] for r in pruned)
+    if not np.isclose(retained + omitted, 1, atol=1e-12, rtol=0):
+        raise ValueError("Probability mass was lost or duplicated")
+    return {"probabilities": grid, "known": known, "nodes": nodes,
+            "evaluated_nodes": len(nodes), "evaluated_leaves": int(known.sum()),
+            "retained_mass": retained, "pruned_mass": omitted, "pruned_branches": pruned}
+
+
+def component_points(probabilities, known=None, relative_threshold=0.1, connectivity=8):
+    """Pick each superlevel component's maximum cell, with deterministic tie breaks."""
+    grid = np.asarray(probabilities, dtype=np.float64)
+    if grid.ndim != 2 or not np.isfinite(grid).all() or (grid < 0).any():
+        raise ValueError("Invalid spatial probability grid")
+    if not 0 < relative_threshold <= 1 or connectivity not in (4, 8):
+        raise ValueError("Invalid connected-component settings")
+    known = np.ones_like(grid, dtype=bool) if known is None else np.asarray(known, dtype=bool)
+    if known.shape != grid.shape:
+        raise ValueError("Known-cell mask shape mismatch")
+    peak = float(grid[known].max()) if known.any() else 0.0
+    threshold = peak * relative_threshold
+    mask = known & (grid >= threshold) & (grid > 0)
+    visited = np.zeros_like(mask)
+    neighbors = [(-1, 0), (1, 0), (0, -1), (0, 1)]
+    if connectivity == 8:
+        neighbors += [(-1, -1), (-1, 1), (1, -1), (1, 1)]
+    height, width = grid.shape
+    components, points = [], []
+    for row, col in zip(*np.where(mask)):
+        if visited[row, col]:
+            continue
+        visited[row, col] = True
+        stack, cells = [(int(row), int(col))], []
+        while stack:
+            y, x = stack.pop()
+            cells.append((y, x))
+            for dy, dx in neighbors:
+                yy, xx = y + dy, x + dx
+                if 0 <= yy < height and 0 <= xx < width and mask[yy, xx] and not visited[yy, xx]:
+                    visited[yy, xx] = True
+                    stack.append((yy, xx))
+        # Prefer the highest probability, then topmost/leftmost on exact ties.
+        y, x = min(cells, key=lambda cell: (-grid[cell], cell[0], cell[1]))
+        point = [(x + 0.5) / width, (y + 0.5) / height]
+        points.append(point)
+        components.append({"cells": [list(cell) for cell in cells], "peak_cell": [y, x],
+                           "peak_probability": float(grid[y, x]),
+                           "mass": sum(float(grid[cell]) for cell in cells)})
+    return {"points": points, "components": components, "threshold": threshold,
+            "relative_threshold": relative_threshold, "connectivity": connectivity}

src/spatial_jev/qwen35.py

@@ -0,0 +1,352 @@
+"""Qwen3.5 multimodal backbone with the existing spatial decision contract.
+
+Each QA is one causal row, with crops revealed at their actual decision stage.
+This baseline deliberately uses complete-prefix rollout before cache optimization.
+"""
+
+import hashlib
+import json
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from .hierarchy import ROOT_BOX, child_box, decode_point, decode_scalar, point_path, scalar_path
+from .unified import SharedClassifier, UnifiedBuilder
+
+LORA_TARGETS = (
+    r"language_model\.layers\.\d+\.(?:"
+    r"self_attn\.(?:q_proj|k_proj|v_proj|o_proj)|"
+    r"linear_attn\.(?:in_proj_qkv|in_proj_z|in_proj_a|in_proj_b|out_proj)|"
+    r"mlp\.(?:gate_proj|up_proj|down_proj))"
+)
+
+
+def sample_seed(sample_id, seed):
+    return seed ^ int(hashlib.sha256(sample_id.encode()).hexdigest()[:8], 16)
+
+
+def train_mode(config):
+    mode = config.get("train_mode", "lora")
+    if mode not in ("lora", "full_language"):
+        raise ValueError(f"Unsupported train mode: {mode}")
+    return mode
+
+
+def configure_full_language(backbone, trainable=True):
+    """Keep the original vision/merger frozen; FP32 master weights for full updates."""
+    backbone.requires_grad_(False)
+    backbone.language_model.float().requires_grad_(trainable)
+
+
+def make_processor(config):
+    from transformers import AutoProcessor
+
+    processor = AutoProcessor.from_pretrained(
+        config["base_model"], local_files_only=config.get("local_files_only", False)
+    )
+    processor.tokenizer.padding_side = "right"
+    processor.image_processor.size = {
+        "shortest_edge": config["min_image_tokens"] * 32 * 32,
+        "longest_edge": config["max_image_tokens"] * 32 * 32,
+    }
+    return processor
+
+
+class Qwen35SpatialModel(nn.Module):
+    def __init__(self, backbone, classifier):
+        super().__init__()
+        self.backbone, self.classifier = backbone, classifier
+
+    def forward(self, inputs, positions, counts):
+        out = self.backbone(**inputs, use_cache=False)
+        hidden = out.last_hidden_state[positions[:, 0], positions[:, 1]]
+        return self.classifier(hidden, counts)
+
+
+def build_model(config, checkpoint=None, trainable=True):
+    from transformers import Qwen3_5ForConditionalGeneration, Qwen3_5TextConfig, Qwen3_5TextModel
+
+    mode = train_mode(config)
+    full = Qwen3_5ForConditionalGeneration.from_pretrained(
+        config["base_model"],
+        local_files_only=True,
+        dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+    )
+    backbone, hidden = full.model, full.config.text_config.hidden_size
+    del full
+    backbone.requires_grad_(False)
+    if mode == "full_language":
+        if checkpoint:
+            text_directory = Path(checkpoint) / "language_model"
+            text_config = Qwen3_5TextConfig.from_pretrained(text_directory, local_files_only=True)
+            backbone.language_model = Qwen3_5TextModel.from_pretrained(
+                text_directory,
+                config=text_config,
+                local_files_only=True,
+                dtype=torch.float32,
+                attn_implementation="sdpa",
+            )
+            backbone.config.text_config = backbone.language_model.config
+        configure_full_language(backbone, trainable)
+    else:
+        from peft import LoraConfig, PeftModel, get_peft_model
+
+        if checkpoint:
+            backbone = PeftModel.from_pretrained(
+                backbone, str(Path(checkpoint) / "adapter"), is_trainable=trainable
+            )
+        else:
+            backbone = get_peft_model(
+                backbone,
+                LoraConfig(
+                    r=config["lora_rank"],
+                    lora_alpha=config["lora_alpha"],
+                    lora_dropout=config["lora_dropout"],
+                    target_modules=LORA_TARGETS,
+                    bias="none",
+                ),
+            )
+    from .training_alignment import seed_head
+
+    classifier = seed_head(config, lambda: SharedClassifier(hidden, config["max_choices"]))
+    if checkpoint:
+        from safetensors.torch import load_file
+
+        classifier.load_state_dict(load_file(str(Path(checkpoint) / "classifier.safetensors")))
+    if trainable and config.get("gradient_checkpointing", False):
+        native = backbone.get_base_model() if mode == "lora" else backbone
+        native.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    model = Qwen35SpatialModel(backbone, classifier)
+    if any(
+        p.requires_grad
+        for n, p in backbone.named_parameters()
+        if n.startswith("visual.") or ".visual." in n
+    ):
+        raise ValueError("Vision encoder and merger must remain frozen")
+    if trainable and not any(
+        "linear_attn" in n and p.requires_grad for n, p in backbone.named_parameters()
+    ):
+        raise ValueError("DeltaNet must have trainable parameters")
+    if not trainable:
+        model.requires_grad_(False).eval()
+    return model
+
+
+def move_inputs(inputs, device):
+    return {
+        key: value.to(
+            device,
+            dtype=torch.bfloat16 if value.is_floating_point() else value.dtype,
+            non_blocking=True,
+        )
+        for key, value in inputs.items()
+    }
+
+
+class Qwen35Builder(UnifiedBuilder):
+    def __init__(self, processor, config):
+        self.processor, self.config = processor, config
+        self.codebook, self.variant = config["scalar_codebook"], "crop_refill"
+        rendered = self.render("Example question", 0)
+        suffix = rendered[rendered.rindex("<|im_start|>assistant") :]
+        self.header = processor.tokenizer.encode(suffix, add_special_tokens=False)
+
+    def render(self, text, image_count):
+        content = [{"type": "image"} for _ in range(image_count)]
+        content.append({"type": "text", "text": text})
+        return self.processor.apply_chat_template(
+            [{"role": "user", "content": content}],
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
+
+    def training_spec(self, row, seed):
+        inp, task = row["input"], row["input"]["answer_space"]["kind"]
+        text, count, mapping = self.initial(inp, seed)
+        images = self.load_images(inp)
+        original, box = images[0], ROOT_BOX
+        if task == "choice":
+            path = [mapping.index(row["target"]["option_id"])]
+        elif task == "scalar":
+            path = scalar_path(row["target"]["value"], self.codebook)
+        else:
+            path = point_path(row["target"]["points"][0])
+        counts = [count]
+        for stage in range(1, len(path)):
+            if task == "point":
+                box = child_box(box, path[stage - 1])
+                question, count = self.point_prompt(inp, box, stage), 9
+                extra = [self.crop(original, box)]
+            else:
+                question, count = self.scalar_prompt(inp, path[0])
+                extra = []
+            text += f"{path[stage - 1] + 1}<|im_end|>\n" + self.render(question, len(extra))
+            images.extend(extra)
+            counts.append(count)
+        return {
+            "text": text,
+            "images": images,
+            "counts": counts,
+            "targets": path,
+            "sample_id": row["sample_id"],
+            "weight": row.get("_weight", 1.0),
+        }
+
+    def encode_specs(self, specs, training=True):
+        encoded = dict(
+            self.processor(
+                text=[s["text"] for s in specs],
+                images=[im for s in specs for im in s["images"]],
+                padding=True,
+                return_tensors="pt",
+                device="cpu",
+            )
+        )
+        encoded.pop("token_type_ids", None)
+        lengths = encoded["attention_mask"].sum(-1).tolist()
+        if max(lengths) > self.config["max_sequence_length"]:
+            raise ValueError(f"Sequence budget exceeded without truncation: {lengths}")
+        positions, counts, targets, weights = [], [], [], []
+        if training:
+            for index, spec in enumerate(specs):
+                ids = encoded["input_ids"][index, : lengths[index]].tolist()
+                ends = [
+                    i + len(self.header) - 1
+                    for i in range(len(ids) - len(self.header) + 1)
+                    if ids[i : i + len(self.header)] == self.header
+                ]
+                if len(ends) != len(spec["targets"]):
+                    raise ValueError(f"Decision positions mismatch: {spec['sample_id']}")
+                positions.extend((index, p) for p in ends)
+                counts.extend(spec["counts"])
+                targets.extend(spec["targets"])
+                weights.extend([spec["weight"] / len(ends)] * len(ends))
+        else:
+            positions = [(i, n - 1) for i, n in enumerate(lengths)]
+            counts = [s["count"] for s in specs]
+        return {
+            "inputs": encoded,
+            "positions": torch.tensor(positions),
+            "counts": torch.tensor(counts),
+            "targets": torch.tensor(targets),
+            "weights": torch.tensor(weights, dtype=torch.float32),
+            "qa_count": sum(s.get("weight", 1.0) for s in specs),
+            "tokens": sum(lengths),
+            "padded_tokens": encoded["input_ids"].numel(),
+        }
+
+    def collate(self, rows):
+        return self.encode_specs(
+            [
+                self.training_spec(row, sample_seed(row["sample_id"], self.config["seed"]))
+                for row in rows
+            ]
+        )
+
+
+@torch.inference_mode()
+def predict_batch(model, builder, requests, device):
+    """Free rollout: this function accepts inputs and IDs, never target labels."""
+    states, results = [], [None] * len(requests)
+    for index, request in enumerate(requests):
+        inp = request["input"]
+        text, count, mapping = builder.initial(
+            inp, sample_seed(request["id"], builder.config["seed"])
+        )
+        states.append(
+            {
+                "index": index,
+                "input": inp,
+                "text": text,
+                "count": count,
+                "mapping": mapping,
+                "images": builder.load_images(inp),
+                "path": [],
+                "logits": [],
+                "box": ROOT_BOX,
+            }
+        )
+    while states:
+        packed = builder.encode_specs(states, training=False)
+        logits = model(
+            move_inputs(packed["inputs"], device),
+            packed["positions"].to(device),
+            packed["counts"].to(device),
+        )
+        choices = logits.argmax(-1).tolist()
+        following = []
+        for state, choice, scores in zip(states, choices, logits.float().cpu()):
+            path, inp = state["path"], state["input"]
+            path.append(choice)
+            state["logits"].append(scores[: state["count"]].tolist())
+            task = inp["answer_space"]["kind"]
+            done = (
+                task == "choice"
+                or (task == "point" and len(path) == 3)
+                or (
+                    task == "scalar"
+                    and (len(path) == 2 or choice in (0, len(builder.codebook["edges_m"])))
+                )
+            )
+            if done:
+                prediction = (
+                    state["mapping"][choice]
+                    if task == "choice"
+                    else decode_point(path)
+                    if task == "point"
+                    else decode_scalar(path, builder.codebook)
+                )
+                results[state["index"]] = {
+                    "prediction": prediction,
+                    "path": path,
+                    "logits": state["logits"],
+                }
+                continue
+            if task == "point":
+                state["box"] = child_box(state["box"], choice)
+                question, count = builder.point_prompt(inp, state["box"], len(path)), 9
+                extra = [builder.crop(state["images"][0], state["box"])]
+            else:
+                question, count = builder.scalar_prompt(inp, choice)
+                extra = []
+            state["text"] += f"{choice + 1}<|im_end|>\n" + builder.render(question, len(extra))
+            state["images"].extend(extra)
+            state["count"] = count
+            following.append(state)
+        states = following
+    return results
+
+
+def save_checkpoint(model, processor, config, path, step):
+    from safetensors.torch import save_file
+
+    path = Path(path)
+    path.mkdir(parents=True, exist_ok=False)
+    mode = train_mode(config)
+    if mode == "full_language":
+        model.backbone.language_model.save_pretrained(
+            path / "language_model", safe_serialization=True, max_shard_size="4GB"
+        )
+    else:
+        model.backbone.save_pretrained(path / "adapter", safe_serialization=True)
+    save_file(
+        {k: v.detach().cpu().contiguous() for k, v in model.classifier.state_dict().items()},
+        str(path / "classifier.safetensors"),
+    )
+    processor.save_pretrained(path / "processor")
+    (path / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    (path / "checkpoint.json").write_text(
+        json.dumps(
+            {
+                "step": step,
+                "format": f"qwen35-{mode}-shared-classifier",
+                "train_mode": mode,
+                "base_revision": config["base_revision"],
+            },
+            indent=2,
+        )
+        + "\n"
+    )

src/spatial_jev/roi_replay.py

@@ -0,0 +1,154 @@
+"""Local original-image token replay, optionally transplanting attention K/V.
+
+No pixels are cropped/resized after the original image. Copy merged vision
+tokens whose support intersects the current rectangle, in original row-major
+order, retaining original M-RoPE coordinates. The KV variant also gathers the
+original post-RoPE K and V at every intact full-attention layer. DeltaNet, Q,
+gates, residuals and MLPs are computed normally; this is a representation ablation,
+not a claim that all computations at replayed token slots have been eliminated.
+"""
+
+import math
+
+import torch
+
+VIT_REPLAY = 'roi_vit_replay'
+KV_REPLAY = 'roi_attention_kv_replay'
+REPLAY_VARIANTS = (VIT_REPLAY, KV_REPLAY)
+EMPTY_ROI_IMAGE = '<|vision_start|><|vision_end|>'
+
+
+def roi_cells(box, shape):
+    """Positive-area overlap with a normalized rectangle, no target-dependent zoom."""
+    h, w = map(int, shape)
+    x0, y0, x1, y1 = box
+    if h < 1 or w < 1 or not (0 <= x0 < x1 <= 1 and 0 <= y0 < y1 <= 1):
+        raise ValueError('Invalid original-image ROI/grid')
+    # Avoid including a neighboring cell from floating-point noise at exact edges.
+    left, top = math.floor(x0*w+1e-9), math.floor(y0*h+1e-9)
+    right, bottom = math.ceil(x1*w-1e-9), math.ceil(y1*h-1e-9)
+    return [y*w+x for y in range(top, bottom) for x in range(left, right)]
+
+
+def expand_replay_inputs(builder, specs, encoded):
+    """Insert virtual local-image slots after processing only the real original(s).
+
+Base input tensors are retained solely for authoritative native M-RoPE creation;
+gather maps align expanded slots, original visual features and original positions.
+The ordinary causal mask still uses SEQUENCE order, not repeated spatial IDs.
+"""
+    if not any(s.get('roi_boxes') for s in specs):
+        return encoded
+    tokenizer = builder.processor.tokenizer
+    image_id, start_id, end_id = [tokenizer.convert_tokens_to_ids(s) for s in
+                                ('<|image_pad|>', '<|vision_start|>', '<|vision_end|>')]
+    merge = builder.processor.image_processor.merge_size
+    original_ids, original_mask = encoded['input_ids'], encoded['attention_mask']
+    original_types = encoded.get('mm_token_type_ids', (original_ids == image_id).long())
+    new_ids, rope_sources, kv_sources, feature_sources = [], [], [], []
+    grid_cursor = feature_cursor = 0
+    for row, spec in enumerate(specs):
+        n = int(original_mask[row].sum())
+        ids = original_ids[row, :n].tolist()
+        grids = encoded['image_grid_thw'][grid_cursor:grid_cursor+len(spec['images'])]
+        n_features = sum(int(g.prod())//merge**2 for g in grids)
+        image_positions = [i for i, token in enumerate(ids) if token == image_id]
+        if len(image_positions) != n_features:
+            raise ValueError('Original visual token/grid count mismatch')
+        original_feature = {pos: feature_cursor+i for i, pos in enumerate(image_positions)}
+        boxes = spec.get('roi_boxes', [])
+        selected = []
+        if boxes:
+            if len(grids) != 1 or int(grids[0, 0]) != 1:
+                raise ValueError('ROI replay requires one still original image')
+            shape = (int(grids[0, 1])//merge, int(grids[0, 2])//merge)
+            selected = [roi_cells(box, shape) for box in boxes]
+        tokens, positions, aliases = [], [], []
+        slot = 0
+        for i, token in enumerate(ids):
+            if token == end_id and i and ids[i-1] == start_id:
+                if slot >= len(selected):
+                    raise ValueError('Unregistered local-image slot')
+                for cell in selected[slot]:
+                    source = image_positions[cell]
+                    tokens.append(image_id)
+                    positions.append(source)
+                    # All originals precede every insertion, hence their sequence
+                    # indices remain identical in the expanded row.
+                    if source >= len(tokens)-1:
+                        raise ValueError('Replay may only reference earlier original tokens')
+                    aliases.append(source)
+                    feature_sources.append(feature_cursor+cell)
+                slot += 1
+            tokens.append(token)
+            positions.append(i)
+            aliases.append(len(tokens)-1)
+            if token == image_id:
+                feature_sources.append(original_feature[i])
+        if slot != len(boxes):
+            raise ValueError('Missing local-image marker')
+        new_ids.append(tokens)
+        rope_sources.append(positions)
+        kv_sources.append(aliases)
+        grid_cursor += len(grids)
+        feature_cursor += n_features
+    length = max(map(len, new_ids))
+    if length > builder.config['max_sequence_length']:
+        raise ValueError('ROI replay exceeds sequence budget; refusing truncation')
+    batch = len(specs)
+    ids = original_ids.new_full((batch, length), tokenizer.pad_token_id)
+    mask = original_mask.new_zeros((batch, length))
+    rope = original_ids.new_zeros((batch, length))
+    aliases = torch.arange(length, dtype=original_ids.dtype).expand(batch, -1).clone()
+    for row, tokens in enumerate(new_ids):
+        n = len(tokens)
+        ids[row, :n] = torch.tensor(tokens)
+        mask[row, :n] = 1
+        rope[row, :n] = torch.tensor(rope_sources[row])
+        aliases[row, :n] = torch.tensor(kv_sources[row])
+    encoded.update(input_ids=ids, attention_mask=mask, mm_token_type_ids=(ids == image_id).long(),
+                   roi_base_ids=original_ids, roi_base_mask=original_mask, roi_base_types=original_types,
+                   roi_position_sources=rope, roi_kv_sources=aliases,
+                   roi_feature_sources=torch.tensor(feature_sources, dtype=torch.long))
+    return encoded
+
+
+class CopyAttentionKV:
+    """Stateless, differentiable aliasing at the native post-RoPE cache boundary."""
+
+    def __init__(self, sources):
+        self.sources = sources
+        self.calls = 0
+
+    def update(self, keys, values, layer_idx, *args, **kwargs):
+        if keys.shape[:3] != values.shape[:3] or keys.shape[0] != self.sources.shape[0] or keys.shape[2] != self.sources.shape[1]:
+            raise ValueError('KV replay shape mismatch')
+        self.calls += 1
+        def gather(tensor):
+            indices = self.sources[:, None, :, None].expand(-1, tensor.shape[1], -1, tensor.shape[-1])
+            return tensor.gather(2, indices)
+        return gather(keys), gather(values)
+
+
+def kv_replay_hidden(lm, hidden, positions, attention_mask, sources):
+    """Scoped native-attention hooks; never modify recurrent/conv states or weights."""
+    if lm.gradient_checkpointing:
+        raise ValueError('Scoped KV replay requires gradient_checkpointing=False')
+    cache, handles, count = CopyAttentionKV(sources), [], 0
+    def inject(module, args, kwargs):
+        if kwargs.get('past_key_values') is not None:
+            raise ValueError('ROI replay must not mix with a persistent prefix cache')
+        return args, {**kwargs, 'past_key_values': cache}
+    try:
+        for layer in lm.layers:
+            if layer.block_type == 'full_attention':
+                count += 1
+                handles.append(layer.self_attn.register_forward_pre_hook(inject, with_kwargs=True))
+        result = lm(inputs_embeds=hidden, position_ids=positions,
+                    attention_mask=attention_mask, use_cache=False).last_hidden_state
+        if cache.calls != count or not count:
+            raise ValueError('Every intact full-attention layer must copy its own K/V')
+        return result
+    finally:
+        for handle in handles:
+            handle.remove()

src/spatial_jev/schema.py

@@ -0,0 +1,138 @@
+"""Typed pilot records, strict target parsing and answer-free prompt construction."""
+
+import ast
+import math
+import re
+
+SCHEMA_VERSION = "spatial-decision/v2"
+EXTENDED_SCHEMA_VERSION = "spatial-decision/v3"
+TASKS = ("choice", "scalar", "point")
+UNIT_FACTORS = {"m": 1.0, "meter": 1.0, "meters": 1.0,
+                "cm": 0.01, "centimeter": 0.01, "centimeters": 0.01,
+                "mm": 0.001, "millimeter": 0.001, "millimeters": 0.001,
+                "inch": 0.0254, "inches": 0.0254, "ft": 0.3048,
+                "foot": 0.3048, "feet": 0.3048}
+_NUMBER = r"[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?"
+
+
+def parse_measurement(answer, question=""):
+    """Accept a scalar, optionally with an explicit length unit; never extract CoT numbers."""
+    match = re.fullmatch(rf"\s*({_NUMBER})\s*([a-zA-Z]+)?\s*\.?\s*", answer)
+    if not match:
+        raise ValueError("not_a_scalar_answer")
+    unit = match.group(2)
+    if unit is None:
+        units = re.findall(r"\b(" + "|".join(UNIT_FACTORS) + r")\b", question.lower())
+        factors = {UNIT_FACTORS[u] for u in units}
+        if len(factors) != 1:
+            raise ValueError("missing_or_ambiguous_unit")
+        unit = units[-1]
+    unit = unit.lower()
+    if unit not in UNIT_FACTORS:
+        raise ValueError("unsupported_unit")
+    value = float(match.group(1)) * UNIT_FACTORS[unit]
+    if not math.isfinite(value) or value < 0:
+        raise ValueError("invalid_measurement")
+    return value, unit
+
+
+def parse_single_point(answer):
+    try:
+        value = ast.literal_eval(answer.strip())
+    except (ValueError, SyntaxError) as exc:
+        raise ValueError("not_a_point_literal") from exc
+    if not isinstance(value, (list, tuple)) or len(value) != 1:
+        raise ValueError("not_a_single_point")
+    point = value[0]
+    if not isinstance(point, (list, tuple)) or len(point) != 2:
+        raise ValueError("not_xy")
+    if any(isinstance(x, bool) or not isinstance(x, (int, float)) for x in point):
+        raise ValueError("non_numeric_point")
+    if any(not math.isfinite(x) or not 0 <= x <= 1 for x in point):
+        raise ValueError("point_out_of_range")
+    return [float(x) for x in point]
+
+
+def clean_question(question):
+    question = question.replace("<image>", "").strip()
+    question = re.split(r"\s*Your answer should be formatted", question, flags=re.I)[0]
+    question = re.sub(r"\s*Answer the question using a single word or phrase\.?", "",
+                      question, flags=re.I)
+    return question.strip()
+
+
+def canonical_metric_question(question):
+    question = clean_question(question)
+    # Do not rewrite numeric lengths in a premise without also converting their values.
+    if re.search(rf"{_NUMBER}\s*(?:cm|mm|meters?|centimeters?|inches|feet)\b",
+                 question, re.I):
+        raise ValueError("numeric_unit_in_question_premise")
+    return re.sub(r"\b(?:centimeters?|cm|millimeters?|mm|inches|feet|foot|ft)\b",
+                  "meters", question, flags=re.I)
+
+
+def validate_record(record, max_choices=32):
+    if record.get("schema_version") not in (SCHEMA_VERSION, EXTENDED_SCHEMA_VERSION):
+        raise ValueError("unsupported_schema")
+    inp = record["input"]
+    if not inp["question"].strip() or not inp["media"]:
+        raise ValueError("empty_input")
+    if any(m["kind"] != "image" for m in inp["media"]):
+        raise ValueError("pilot_requires_images")
+    space, target = inp["answer_space"], record["target"]
+    kind = space["kind"]
+    if kind == "choice":
+        options = space["options"]
+        ids = [o["id"] for o in options]
+        if not 2 <= len(options) <= max_choices or len(set(ids)) != len(ids):
+            raise ValueError("invalid_options")
+        if target["option_id"] not in ids:
+            raise ValueError("target_missing_from_options")
+    elif kind == "scalar":
+        allowed = ("m", "count", "deg") if record["schema_version"] == EXTENDED_SCHEMA_VERSION else ("m",)
+        if space["unit"] not in allowed or target["unit"] != space["unit"]:
+            raise ValueError("pilot_scalar_requires_meters")
+        if isinstance(target["value"], bool) or not math.isfinite(target["value"]) or target["value"] < 0:
+            raise ValueError("invalid_scalar")
+        if space["unit"] == "count":
+            if space.get("quantity") != "count" or int(target["value"]) != target["value"] or target["value"] > 1023:
+                raise ValueError("invalid_integer_count")
+        if space["unit"] == "deg":
+            if target["value"] > 360 or not space.get("angle_semantics"):
+                raise ValueError("invalid_angle_semantics_or_range")
+    elif kind == "point":
+        if len(inp["media"]) != 1:
+            raise ValueError("pilot_point_requires_one_image")
+        if space["coordinate_system"] != "normalized_xy" or space["num_points"] != 1:
+            raise ValueError("unsupported_point_space")
+        parse_single_point(repr(target["points"]))
+    else:
+        raise ValueError("unsupported_task")
+
+
+def compile_prompt(inp, permutation=None):
+    """Only input is accepted; targets and provenance cannot enter the prompt."""
+    space = inp["answer_space"]
+    kind = space["kind"]
+    text = f"Question: {inp['question']}\n"
+    option_ids = []
+    if kind == "choice":
+        options = space["options"]
+        order = list(range(len(options))) if permutation is None else list(permutation)
+        if sorted(order) != list(range(len(options))):
+            raise ValueError("invalid_permutation")
+        text += "Options:\n"
+        for i, index in enumerate(order):
+            option = options[index]
+            text += f"{i + 1}. {option['text']}\n"
+            option_ids.append(option["id"])
+        text += "Select exactly one option."
+    elif kind == "scalar":
+        unit = {"m": "meters", "count": "count units (an integer)", "deg": "degrees"}[space["unit"]]
+        text += f"Estimate {space['quantity']}. Return a nonnegative value in {unit}."
+    elif kind == "point":
+        text += ("Return one valid point (x, y) in the original image, normalized to [0, 1]. "
+                 "The origin is top-left; x increases rightward and y increases downward.")
+    else:
+        raise ValueError("unsupported_task")
+    return text, option_ids

src/spatial_jev/training_alignment.py

@@ -0,0 +1,56 @@
+"""Shared controls for the three full-language architecture experiments."""
+
+import hashlib
+import json
+import math
+
+import torch
+
+
+def old_head_factors(config, step, total_steps):
+    """Exactly the old runner's zero-based, pre-optimizer-step LR multipliers."""
+    decay = 0.1 + 0.9 * (1 + math.cos(math.pi * step / total_steps)) / 2
+    warm = min(1, (step + 1) / config["warmup_steps"])
+    backbone_warm = min(1, max(0, step + 1 - config["head_warmup_steps"]) / config["warmup_steps"])
+    return backbone_warm * decay, warm * decay
+
+
+def seed_head(config, factory):
+    if "head_seed" not in config:
+        return factory()
+    with torch.random.fork_rng(devices=[]):
+        torch.random.default_generator.manual_seed(config["head_seed"])
+        return factory()
+
+
+def state_hash(module):
+    return hashlib.sha256(
+        b"".join(
+            value.detach().float().cpu().contiguous().numpy().tobytes()
+            for value in module.state_dict().values()
+        )
+    ).hexdigest()
+
+
+def canonical_rows(rows, path, expected_sha256):
+    """Restore the ORIGINAL flat order before shuffling, including padding row 0."""
+    if hashlib.sha256(path.read_bytes()).hexdigest() != expected_sha256:
+        raise ValueError("Canonical training file hash changed")
+    by_id = {row["sample_id"]: row for row in rows}
+    identities = []
+    with path.open() as stream:
+        for line in stream:
+            source = json.loads(line)
+            identity = source["sample_id"]
+            identities.append(identity)
+            if identity not in by_id or any(
+                source.get(key) != by_id[identity].get(key) for key in ("input", "target")
+            ):
+                raise ValueError(f"Canonical QA contents differ: {identity}")
+    if (
+        len(by_id) != len(rows)
+        or len(set(identities)) != len(identities)
+        or set(identities) != by_id.keys()
+    ):
+        raise ValueError("Canonical QA identities do not match the training dataset")
+    return [by_id[identity] for identity in identities]

src/spatial_jev/unified.py

@@ -0,0 +1,357 @@
+"""One shared classifier with scalar interval decisions and three-step pointing."""
+
+import json
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+from torch import nn
+
+from .hierarchy import (POINT_NAMES, ROOT_BOX, child_box, decode_point, decode_scalar,
+                        point_path, roi_visible, scalar_fine_bins, scalar_interval, scalar_path,
+                        visual_token_boxes)
+from .schema import compile_prompt
+from .numeric_tasks import codebook_for, numeric_prompt
+
+
+def causal_visual_mask(length, visual_types, stages, query_start=0):
+    """Causal text plus bidirectional vision ONLY within the same reveal stage.
+
+    Native Molmo2's all-image bidirectionality would leak future GT-selected crops.
+    Old prefix queries are not recomputed during incremental inference.
+    """
+    keys = torch.arange(length)
+    queries = keys[query_start:]
+    allowed = keys[None] <= queries[:, None]
+    visual_types = torch.as_tensor(visual_types, dtype=torch.bool)
+    stages = torch.as_tensor(stages, dtype=torch.long)
+    allowed |= (visual_types[query_start:, None] & visual_types[None]
+                & (stages[query_start:, None] == stages[None]))
+    return torch.zeros((1, 1, len(queries), length)).masked_fill(~allowed[None, None], -torch.inf)
+
+
+def apply_roi_mask(mask, visual_positions, boxes, roi, start=0, end=None):
+    keep = roi_visible(boxes, roi)
+    if not keep.any():
+        raise ValueError("ROI has no visual support")
+    blocked = torch.as_tensor(np.asarray(visual_positions)[~keep], dtype=torch.long)
+    mask[:, :, start:end, blocked] = -torch.inf
+    return int(keep.sum())
+
+
+def move_inputs(inputs, device):
+    return {key: value.to(device=device, dtype=torch.bfloat16)
+            if value.is_floating_point() else value.to(device) for key, value in inputs.items()}
+
+
+def collate_training(examples, pad_id):
+    """Pad independent examples without cross-example attention or image-pooling offsets."""
+    length = max(p["inputs"]["input_ids"].shape[1] for p in examples)
+    batch = len(examples)
+    ids = torch.full((batch, length), pad_id, dtype=torch.long)
+    mask = torch.full((batch, 1, length, length), -torch.inf)
+    positions, counts, targets, weights = [], [], [], []
+    for index, packed in enumerate(examples):
+        n = packed["inputs"]["input_ids"].shape[1]
+        ids[index, :n] = packed["inputs"]["input_ids"][0]
+        mask[index, :, :n, :n] = packed["inputs"]["attention_mask"][0]
+        padding = torch.arange(n, length)
+        mask[index, 0, padding, padding] = 0
+        positions.extend((index, int(p)) for p in packed["positions"])
+        counts.extend(packed["counts"].tolist())
+        targets.extend(packed["targets"].tolist())
+        weights.extend([1 / len(packed["targets"])] * len(packed["targets"]))
+    inputs = {"input_ids": ids, "attention_mask": mask}
+    for key in ("pixel_values", "image_token_pooling", "image_grids", "image_num_crops"):
+        # Native pooling indices are local to each image; the model adds offsets itself.
+        inputs[key] = torch.cat([p["inputs"][key] for p in examples], dim=0)
+    return {"inputs": inputs, "positions": torch.tensor(positions), "counts": torch.tensor(counts),
+            "targets": torch.tensor(targets), "loss_weights": torch.tensor(weights) / batch}
+
+
+class SharedClassifier(nn.Module):
+    def __init__(self, hidden_size, max_choices=32):
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, max_choices)
+
+    def forward(self, hidden, counts):
+        with torch.autocast(device_type=hidden.device.type, enabled=False):
+            logits = self.linear(self.norm(hidden.float()))
+            invalid = torch.arange(logits.shape[-1], device=logits.device)[None] >= counts[:, None]
+            return logits.masked_fill(invalid, -torch.inf)
+
+
+class UnifiedSpatialModel(nn.Module):
+    def __init__(self, backbone, classifier):
+        super().__init__()
+        self.backbone, self.classifier = backbone, classifier
+
+    def forward(self, inputs, positions, counts):
+        out = self.backbone(**inputs, use_cache=False)
+        hidden = (out.last_hidden_state[positions[:, 0], positions[:, 1]] if positions.ndim == 2
+                  else out.last_hidden_state[0, positions])
+        return self.classifier(hidden, counts)
+
+
+def build_unified_model(config, checkpoint=None, trainable=True):
+    from peft import LoraConfig, PeftModel, get_peft_model
+    from transformers import AutoModelForImageTextToText
+
+    full = AutoModelForImageTextToText.from_pretrained(
+        config["base_model"], trust_remote_code=True, local_files_only=True,
+        dtype=torch.bfloat16, attn_implementation="sdpa")
+    backbone, hidden = full.model, full.config.text_config.hidden_size
+    del full
+    if checkpoint:
+        backbone = PeftModel.from_pretrained(backbone, str(Path(checkpoint) / "adapter"),
+                                             is_trainable=trainable)
+    else:
+        backbone = get_peft_model(backbone, LoraConfig(
+            r=config["lora_rank"], lora_alpha=config["lora_alpha"],
+            lora_dropout=config["lora_dropout"], target_modules=config["lora_targets"], bias="none"))
+    classifier = SharedClassifier(hidden, config["max_choices"])
+    if checkpoint:
+        from safetensors.torch import load_file
+        classifier.load_state_dict(load_file(str(Path(checkpoint) / "classifier.safetensors")))
+    if trainable and config.get("gradient_checkpointing", True):
+        backbone.get_base_model().gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs={"use_reentrant": False})
+    model = UnifiedSpatialModel(backbone, classifier)
+    if any(p.requires_grad for n, p in backbone.named_parameters() if "vision_backbone" in n):
+        raise ValueError("Visual encoder/projector must be frozen in this ablation")
+    if not trainable:
+        model.requires_grad_(False).eval()
+    return model
+
+
+def save_unified_checkpoint(model, processor, config, output, step, optimizer=None):
+    from safetensors.torch import save_file
+
+    output = Path(output)
+    output.mkdir(parents=True, exist_ok=False)
+    model.backbone.save_pretrained(output / "adapter", safe_serialization=True)
+    save_file({k: v.detach().cpu().contiguous() for k, v in model.classifier.state_dict().items()},
+              str(output / "classifier.safetensors"))
+    processor.save_pretrained(output / "processor")
+    (output / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    if optimizer:
+        torch.save({"step": step, "optimizer": optimizer.state_dict()}, output / "optimizer.pt")
+    (output / "checkpoint.json").write_text(json.dumps({
+        "step": step, "format": "base_plus_lora_plus_one_classifier",
+        "point_variant": config["point_variant"], "base_revision": config["base_revision"],
+        "calibrated": False}, indent=2) + "\n")
+
+
+class UnifiedBuilder:
+    def __init__(self, processor, config):
+        self.processor, self.config = processor, config
+        self.codebook = config["scalar_codebook"]
+        self.variant = config["point_variant"]
+        if self.variant not in ("single_image", "crop_refill", "roi_mask"):
+            raise ValueError("Unknown point ablation")
+        self.patch_id = processor.tokenizer.convert_tokens_to_ids("<im_patch>")
+        self.header = processor.tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
+
+    def render(self, text, image_count):
+        content = [{"type": "image"} for _ in range(image_count)]
+        content.append({"type": "text", "text": text})
+        # Render each reveal separately. Rendering all messages at once hoists future crops.
+        return self.processor.apply_chat_template([{"role": "user", "content": content}],
+                                                  tokenize=False, add_generation_prompt=True)
+
+    @staticmethod
+    def load_images(inp):
+        images = []
+        for media in inp["media"]:
+            with Image.open(media["uri"]) as image:
+                images.append(image.convert("RGB"))
+        return images
+
+    @staticmethod
+    def crop(original, box):
+        w, h = original.size
+        extent = (box[0] * w, box[1] * h, box[2] * w, box[3] * h)
+        size = (max(1, round(extent[2] - extent[0])), max(1, round(extent[3] - extent[1])))
+        return original.transform(size, Image.Transform.EXTENT, extent, Image.Resampling.BILINEAR)
+
+    def point_prompt(self, inp, box, level):
+        bounds = ", ".join(f"{x:.6f}" for x in box)
+        options = "\n".join(f"{i+1}. {name}" for i, name in enumerate(POINT_NAMES))
+        return (f"Request: {inp['question']}\nChoose one valid point satisfying the request. "
+                f"Localization step {level+1} of 3. The current rectangle in the ORIGINAL image "
+                f"is (left, top, right, bottom) = ({bounds}), normalized to [0, 1]. "
+                "The origin is top-left, x increases rightward and y downward. "
+                "Any image appended after the original is a zoomed view of its current rectangle. "
+                "Divide the CURRENT rectangle into three equal rows and three equal columns. "
+                f"Which region contains the point?\nOptions:\n{options}\nSelect exactly one option.")
+
+    def scalar_prompt(self, inp, root=None):
+        cb = codebook_for(inp, self.codebook)
+        if cb.get("unit", "m") != "m":
+            return numeric_prompt(inp, root, cb)
+        if root is None:
+            options = ["exactly zero meters"]
+            options += [f"{lo:g} to {hi:g} meters" for lo, hi in zip(cb["edges_m"], cb["edges_m"][1:])]
+            options[1] = f"greater than zero and less than {cb['edges_m'][1]:g} meters"
+            options += [f"greater than {cb['edges_m'][-1]:g} meters"]
+            instruction = "Select the range containing the nonnegative measurement."
+        else:
+            lo, hi = scalar_interval(root, cb)
+            fine = scalar_fine_bins(root, cb)
+            step = (hi - lo) / fine
+            options = [f"{lo+i*step:.6g} to {lo+(i+1)*step:.6g} meters"
+                       for i in range(fine)]
+            instruction = f"Refine the previously selected range {lo:g} to {hi:g} meters."
+        text = (f"Question: {inp['question']}\n{instruction} "
+                "Ranges include the lower endpoint and exclude the upper endpoint, "
+                "except that the maximum finite endpoint is included. Exact zero has its own class."
+                "\nOptions:\n" + "\n".join(f"{i+1}. {s}" for i, s in enumerate(options))
+                + "\nSelect exactly one option.")
+        return text, len(options)
+
+    def initial(self, inp, seed):
+        task, mapping = inp["answer_space"]["kind"], []
+        if task == "choice":
+            order = list(range(len(inp["answer_space"]["options"])))
+            if not self.config.get("preserve_option_order", False):
+                random.Random(seed).shuffle(order)
+            text, mapping = compile_prompt(inp, order)
+            count = len(mapping)
+        elif task == "scalar":
+            text, count = self.scalar_prompt(inp)
+        else:
+            text, count = self.point_prompt(inp, ROOT_BOX, 0), 9
+        return self.render(text, len(inp["media"])), count, mapping
+
+    def encode(self, text, images, metadata=False, continuation=False):
+        kwargs = {"text": [text], "return_tensors": "pt", "padding": False,
+                  "return_mm_token_type_ids": True}
+        if images:
+            kwargs.update(images=images, return_pointing_metadata=metadata)
+        encoded = self.processor(**kwargs)
+        meta = encoded.pop("metadata", None)
+        if continuation:
+            bos = self.processor.tokenizer.bos_token_id or self.processor.tokenizer.eos_token_id
+            if encoded["input_ids"][0, 0].item() != bos:
+                raise ValueError("Unexpected continuation BOS")
+            for key in ("input_ids", "attention_mask", "token_type_ids"):
+                encoded[key] = encoded[key][:, 1:]
+        return dict(encoded), meta
+
+    def geometry(self, inputs, metadata):
+        positions = (inputs["input_ids"][0] == self.patch_id).nonzero().flatten().numpy()
+        boxes = visual_token_boxes(metadata, inputs["image_grids"].tolist(),
+                                   inputs["image_num_crops"].tolist(), inputs["pixel_values"].shape[1])
+        if len(positions) != len(boxes):
+            raise ValueError("Patch positions and geometric supports differ")
+        return positions, boxes
+
+    def training_example(self, row, seed):
+        inp, task = row["input"], row["input"]["answer_space"]["kind"]
+        text, count, mapping = self.initial(inp, seed)
+        images = self.load_images(inp)
+        original = images[0]
+        if task == "choice":
+            path = [mapping.index(row["target"]["option_id"])]
+        elif task == "scalar":
+            path = scalar_path(row["target"]["value"], self.codebook)
+        else:
+            path = point_path(row["target"]["points"][0])
+        counts, boxes, box = [count], [ROOT_BOX], ROOT_BOX
+        for stage in range(1, len(path)):
+            if task == "point":
+                box = child_box(box, path[stage - 1])
+                question, count = self.point_prompt(inp, box, stage), 9
+                extra = [self.crop(original, box)] if self.variant == "crop_refill" else []
+            else:
+                question, count = self.scalar_prompt(inp, path[0])
+                extra = []
+            text += f"{path[stage-1]+1}<|im_end|>\n" + self.render(question, len(extra))
+            images.extend(extra)
+            counts.append(count)
+            boxes.append(box)
+        needs_geometry = task == "point" and self.variant == "roi_mask"
+        inputs, metadata = self.encode(text, images, metadata=needs_geometry)
+        ids = inputs["input_ids"][0].tolist()
+        positions = [i + len(self.header) - 1 for i in range(len(ids) - len(self.header) + 1)
+                     if ids[i:i+len(self.header)] == self.header]
+        if len(positions) != len(path):
+            raise ValueError("Decision marker count mismatch")
+        if len(ids) > self.config["max_sequence_length"]:
+            raise ValueError("Sequence budget exceeded")
+        types = inputs.pop("token_type_ids")[0]
+        stages = torch.zeros(len(ids), dtype=torch.long)
+        for stage in range(1, len(path)):
+            stages[positions[stage-1]+1:positions[stage]+1] = stage
+        mask = causal_visual_mask(len(ids), types, stages)
+        if needs_geometry:
+            visual_positions, supports = self.geometry(inputs, metadata)
+            for stage in range(1, len(path)):
+                apply_roi_mask(mask, visual_positions, supports, boxes[stage],
+                               positions[stage-1]+1, positions[stage]+1)
+        inputs["attention_mask"] = mask
+        return {"inputs": inputs, "positions": torch.tensor(positions),
+                "counts": torch.tensor(counts), "targets": torch.tensor(path),
+                "mapping": mapping, "boxes": boxes}
+
+
+@torch.inference_mode()
+def predict_unified(model, builder, inp, device, seed, forced_path=None):
+    """No targets are accepted. forced_path is solely for cache-equivalence integration tests."""
+    task = inp["answer_space"]["kind"]
+    images = builder.load_images(inp)
+    original = images[0]
+    text, count, mapping = builder.initial(inp, seed)
+    need_geometry = task == "point" and builder.variant == "roi_mask"
+    encoded, meta = builder.encode(text, images, metadata=need_geometry)
+    geom = builder.geometry(encoded, meta) if need_geometry else None
+    all_types, all_stages = [], []
+    path, logits_trace, support_counts = [], [], []
+    box, cache, previous_length, total_input_tokens = ROOT_BOX, None, 0, 0
+    for stage in range(3 if task == "point" else 2 if task == "scalar" else 1):
+        types = encoded.pop("token_type_ids")[0].tolist()
+        all_types.extend(types)
+        all_stages.extend([stage] * len(types))
+        total_input_tokens += encoded["input_ids"].shape[1]
+        if total_input_tokens > builder.config["max_sequence_length"]:
+            raise ValueError("Rollout sequence budget exceeded")
+        mask = causal_visual_mask(total_input_tokens, all_types, all_stages, previous_length)
+        if need_geometry and stage > 0:
+            support_counts.append(apply_roi_mask(mask, *geom, box))
+        encoded["attention_mask"] = mask
+        inputs = move_inputs(encoded, device)
+        inputs["cache_position"] = torch.arange(previous_length, total_input_tokens, device=device)
+        output = model.backbone(**inputs, past_key_values=cache, use_cache=True)
+        logits = model.classifier(output.last_hidden_state[:, -1],
+                                  torch.tensor([count], device=device))
+        logits_trace.append(logits[0, :count].float().cpu().tolist())
+        choice = int(logits.argmax(-1).item()) if forced_path is None else forced_path[stage]
+        path.append(choice)
+        cache, previous_length = output.past_key_values, total_input_tokens
+        if task == "choice":
+            break
+        if task == "scalar":
+            if stage == 1 or choice in (0, len(builder.codebook["edges_m"])):
+                break
+            question, count = builder.scalar_prompt(inp, choice)
+            extra = []
+        else:
+            box = child_box(box, choice)
+            if stage == 2:
+                break
+            question, count = builder.point_prompt(inp, box, stage + 1), 9
+            extra = [builder.crop(original, box)] if builder.variant == "crop_refill" else []
+        text = f"{choice+1}<|im_end|>\n" + builder.render(question, len(extra))
+        encoded, _ = builder.encode(text, extra, continuation=True)
+    if task == "choice":
+        prediction = mapping[path[0]]
+    elif task == "scalar":
+        prediction = decode_scalar(path, builder.codebook)
+    else:
+        prediction = decode_point(path)
+    return {"prediction": prediction, "path": path, "logits": logits_trace,
+            "mapping": mapping, "input_tokens": total_input_tokens,
+            "roi_support_counts": support_counts}

src/spatial_jev/visual_token_pointing.py

@@ -0,0 +1,118 @@
+"""Direct question-Q / image-token-K pointing on Qwen's merged visual grid.
+
+The candidates are LLM image-token positions, not raw ViT keys (a different
+projection space). Geometry follows image_grid_thw after spatial merging, in
+row-major order, with no crop, label-dependent input, or learned coordinate head.
+"""
+
+import math
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+from .point_probability_map import component_points
+from .roi_replay import REPLAY_VARIANTS
+
+VISUAL_TOKENS = "visual_tokens"
+ORIGINAL_IMAGE = "original_image_nine_way"
+
+
+def validate_pointing_config(config):
+    variant = config.get("point_variant", "crop_refill")
+    if variant not in ("crop_refill", VISUAL_TOKENS, ORIGINAL_IMAGE, *REPLAY_VARIANTS):
+        raise ValueError(f"Unsupported KEV pointing method: {variant}")
+    if variant == ORIGINAL_IMAGE and config.get("point_depth", 3) != 3:
+        raise ValueError("Original-image nine-way pointing requires three decisions")
+    if variant in REPLAY_VARIANTS and (
+            config.get('point_depth', 3) != 3 or config.get('readout_kind') != 'native_last_qk_mean'
+                or config.get('readout_layer_indices') is not None or config.get('readout_layer_index', 31) != 31
+            or config.get('shared_prefix', False) or config.get('gradient_checkpointing', False)):
+        raise ValueError('ROI replay fixes L32 mean, three decisions and no persistent cache/checkpointing')
+    if variant != VISUAL_TOKENS:
+        return
+    if config.get("point_depth", 1) != 1:
+        raise ValueError("Visual-token pointing is a single decision, not a crop tree")
+    if config.get("shared_prefix", False):
+        raise ValueError("Visual-token keys require the full causal row; shared_prefix is unsupported")
+    if config.get("point_loss", "uniform_cell_ce") != "uniform_cell_ce":
+        raise ValueError("Visual-token pointing requires uniform_cell_ce")
+    if config.get("point_decoder", "argmax") not in ("argmax", "components"):
+        raise ValueError("Unknown visual-token point decoder")
+    if not 0 < config.get("point_component_threshold", 0.1) <= 1:
+        raise ValueError("Invalid visual-token component threshold")
+
+
+def image_token_candidates(ids, image_token_id, grid_thw, merge_size):
+    """One original image only; reject video, missing tokens and ambiguous grids."""
+    t, h, w = [int(v) for v in grid_thw]
+    if merge_size < 1 or t != 1 or min(h, w) < 1 or h % merge_size or w % merge_size:
+        raise ValueError("Expected one spatially mergeable still-image grid")
+    shape = (h // merge_size, w // merge_size)
+    positions = [i for i, token in enumerate(ids) if token == image_token_id]
+    if len(positions) != math.prod(shape) or positions != list(range(positions[0], positions[-1] + 1)):
+        raise ValueError("Image-token positions do not match the merged visual grid")
+    return positions, shape
+
+
+def target_cells(points, shape):
+    """Map normalized training labels to occupied cells, without modifying inputs."""
+    if not isinstance(points, (list, tuple)) or not points:
+        raise ValueError("Point supervision must contain at least one normalized point")
+    height, width = shape
+    if min(height, width) < 1:
+        raise ValueError("Empty visual grid")
+    cells = set()
+    for point in points:
+        if (not isinstance(point, (list, tuple)) or len(point) != 2
+                or any(isinstance(v, bool) or not isinstance(v, (int, float))
+                       or not math.isfinite(v) or not 0 <= v <= 1 for v in point)):
+            raise ValueError("Point supervision must use finite normalized coordinates")
+        x, y = point
+        col, row = min(int(x * width), width - 1), min(int(y * height), height - 1)
+        cells.add(row * width + col)
+    return sorted(cells)
+
+
+def decision_losses(logits, packed):
+    """Ordinary CE unchanged; multiple occupied visual cells share one QA's mass.
+
+    Sparse targets avoid 0 * -inf on padded candidate slots. Coordinates in the
+    same visual cell are intentionally indistinguishable at this resolution.
+    """
+    targets = packed["targets"].to(logits.device)
+    if "visual_target_owner" not in packed:
+        return F.cross_entropy(logits, targets, reduction="none")
+    log_probs = F.log_softmax(logits.float(), dim=-1)
+    losses = -log_probs.gather(1, targets[:, None]).squeeze(1)
+    owners = packed["visual_target_owner"].to(logits.device)
+    slots = packed["visual_target_slot"].to(logits.device)
+    mass = packed["visual_target_mass"].to(logits.device)
+    visual_losses = torch.zeros_like(losses).index_add(0, owners, -log_probs[owners, slots] * mass)
+    selected = owners.unique()
+    return losses.index_copy(0, selected, visual_losses[selected])
+
+
+def decode_visual_points(logits, shape, decoder="argmax", relative_threshold=0.1):
+    """Decode only model scores and known geometry; never receive GT/count/masks."""
+    values = np.asarray(logits, dtype=np.float64)
+    height, width = shape
+    if min(height, width) < 1 or values.shape != (height * width,) or not np.isfinite(values).all():
+        raise ValueError("Expected one finite score per merged visual token")
+    probabilities = np.exp(values - values.max())
+    probabilities /= probabilities.sum()
+    grid = probabilities.reshape(height, width)
+    if decoder == "argmax":
+        index = int(values.argmax())
+        y, x = divmod(index, width)
+        points = [[(x + 0.5) / width, (y + 0.5) / height]]
+    elif decoder == "components":
+        components = component_points(grid, relative_threshold=relative_threshold, connectivity=8)["components"]
+        components.sort(key=lambda c: (-c["peak_probability"], *c["peak_cell"]))
+        points = [[(c["peak_cell"][1] + 0.5) / width, (c["peak_cell"][0] + 0.5) / height]
+                  for c in components]
+    else:
+        raise ValueError(f"Unknown visual-token point decoder: {decoder}")
+    return {"prediction": points[0], "points": points, "grid_hw": list(shape),
+            "probabilities": grid.tolist(), "pointing_method": VISUAL_TOKENS,
+            "point_decoder": decoder, "point_component_threshold": relative_threshold if decoder == "components" else None}