Decision-2.0-Eos-0.8B Core ML: shared prefix + chunked packed questions, fp16 package, runtime, parity reports

.gitattributes

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

Decision2EosPacked.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

Decision2EosPacked.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

Decision2EosPacked.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "4DDB3568-C852-4A8A-8E49-47C9F5FCB612": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "7323AD02-87D1-4987-95F3-AC1D30199D66": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "4DDB3568-C852-4A8A-8E49-47C9F5FCB612"
+}

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://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
+
+       http://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.

README.md

@@ -0,0 +1,101 @@
+---
+license: apache-2.0
+base_model: vllm-sr/Decision-2.0-Eos-0.8B
+tags:
+- coreml
+- decision-making
+- on-device
+- system-one
+---
+
+# Decision-2.0-Eos-0.8B · Core ML
+
+[vllm-sr/Decision-2.0-Eos-0.8B](https://huggingface.co/vllm-sr/Decision-2.0-Eos-0.8B) (vLLM Semantic Router, Apache-2.0)
+converted to Core ML for on-device use on Apple silicon. Decision 2.0 answers Choice, Yes/No and Score questions
+about one input, with a probability for every option and no text generation.
+
+This port answers **every question of a request in one model call**. The text all the questions share (the context) is
+read once, and each question then continues from it as if it were alone. Eos uses a Qwen3.5 hybrid backbone (Gated
+DeltaNet + attention), so "continues from it" covers the DeltaNet recurrent state and conv history as well as the
+attention keys. See [How it works](#how-it-works).
+
+| | |
+|---|---|
+| Source | `vllm-sr/Decision-2.0-Eos-0.8B` @ `3594047d` (Qwen3.5-0.8B-Base backbone + shared decision head) |
+| Package | `Decision2EosPacked.mlpackage`, fp16, 1.4 GB, five functions sharing one copy of the weights |
+| Functions | `S128_C256_N32`, `S256_C512_N64`, `S512_C768_N96`, `S512_C1024_N128`, `S1024_C1024_N128` (shared-prefix tokens / packed question tokens / option slots) |
+| Requires | macOS 15+ / iOS 18+; run on the GPU (`CPU_AND_GPU`) |
+
+Kai 0.6B (Qwen3 backbone) is at [FluidInference/decision-2.0-kai-coreml](https://huggingface.co/FluidInference/decision-2.0-kai-coreml),
+with the same Python runtime.
+
+## Quickstart
+
+```bash
+pip install coremltools tokenizers numpy
+hf download FluidInference/decision-2.0-eos-coreml --local-dir decision-2.0-eos-coreml
+cd decision-2.0-eos-coreml && python example.py
+```
+
+`decision2_coreml.py` mirrors the upstream `system_one()` output (`choice`/`noul`/`score`, `probabilities`,
+`confidence`, `legend`). It picks the smallest function that fits a request. A request that fits none is split into
+chunks, each repeating the shared prefix, on whichever function needs the fewest milliseconds.
+
+## Fidelity
+
+Against the upstream runtime (fp32, Transformers 5.18, MPS) on the **typed-decisions TEST** split
+(`LocalLLaMA/typed-decisions`, 400 requests, 2,000 decisions):
+
+| | Upstream | Core ML fp16 |
+|---|---:|---:|
+| Choice accuracy (600) | 0.430 | 0.433 |
+| Yes/No accuracy (600) | 0.598 | 0.598 |
+| Score accuracy (800) | 0.379 | 0.376 |
+| Answers that differ from upstream | — | 4 / 2,000 (all near-ties: upstream top-2 margin ≤ 0.003) |
+| Max probability difference | — | 0.0077 |
+
+Tokenization (via `tokenizers`) matches the upstream encoder on all 2,000 questions. The fp32 PyTorch version of the
+graph (`conversion/eos_graph.py`, `EosChunked`) matches upstream to 2e-6. The accuracies are only for the comparison
+with upstream: they are not the model card's benchmark protocol.
+
+## Speed (MacBook Pro M5 Pro, 24 GB, macOS 27)
+
+| Request | Upstream PyTorch, MPS | Core ML (GPU) |
+|---|---:|---:|
+| Card quickstart, 3 questions | 306 ms | **34 ms** |
+| typed-decisions request, 5 questions (median) | 1,595 ms | **110 ms** |
+| One email, 64 questions | 15,647 ms | **874 ms** (13 calls) |
+
+Times are full requests (tokenize, pack, predict, answer) and exclude model loading. Per call: S128_C256 34 ms,
+S256_C512 64 ms, S512_C768 109 ms, S512_C1024 131 ms, S1024_C1024 176 ms. On a Mac the upstream runtime has none of its
+GPU kernels (flash-linear-attention, causal-conv1d), so the speedup is specific to this machine. The Neural Engine was
+not used: on Kev-0.8B, which has the same backbone, it was about 24× slower than the GPU.
+
+## How it works
+
+One call takes the shared prefix (right-padded to S) and the questions' suffixes packed end to end (C tokens):
+
+- **Attention layers:** one mask. Prefix tokens are causal; a suffix token sees the real prefix tokens and earlier
+  tokens of its own question.
+- **Gated DeltaNet layers:** the prefix runs the chunked delta rule (chunk 64; padded positions get β = g = 0, exact
+  no-ops), which yields its final recurrent state S₀. The packed questions also run in chunks of 64. A token whose
+  question began in an earlier chunk continues from the carried state; a token whose question began in this chunk
+  starts from S₀ (`cont`). Decays and the WY inverse are masked to the question within each chunk (`seg_chunks`), and
+  only the question of a chunk's last token carries its state forward (`last_seg`). A question's first three
+  convolution lags read the prefix's last three inputs (`lag_keep`, `lag_tail`).
+- **Readout:** Decision 2.0's head (bilinear + MLP over LayerNorm'd hidden states) at each option's last token and the
+  question's final token (`cand_idx`, `query_idx`).
+
+So every question sees exactly its own upstream row (`Context … Question … Options … Decision:`), and the delta-rule
+cost stays linear in packed tokens. Full input list: `coreml_config.json`.
+
+## Conversion
+
+`conversion/`: `eos_graph.py` (Decision 2.0 graph on top of `kev_stages.py` / `qwen35_export.py`, the Qwen3.5 export
+from [Kev-0.8B Core ML](https://huggingface.co/FluidInference/kev-0.8b-coreml)), `convert_eos.py S cP N` (torch 2.7.0,
+coremltools 9.0), `combine.py`, parity and timing scripts.
+
+## License
+
+Apache-2.0, as the source model. `LICENSE`, `tokenizer*.json` and `decision_config.json` are copied unchanged from the
+source repository.

conversion/cases.py

@@ -0,0 +1,19 @@
+LEVELS5 = ["None", "Low", "Medium", "High", "Critical"]
+CASES = [
+ ("The order arrived damaged yesterday. The customer has a receipt and asks for a replacement today.",
+  {"route": {"type":"choice","instructions":"Which team should handle this request?",
+     "criteria":{"returns":"Refunds, replacements and damaged deliveries","billing":"Payments, invoices and charges","technical":"Product setup and faults"}},
+   "receipt": {"type":"noul","instructions":"Does the customer have a receipt?"},
+   "urgency": {"type":"score","instructions":"How urgent is this request?","criteria":["Routine","Soon","Today"]}}),
+ ("Hi team, our production database went down at 3am and customers cannot log in. We need someone on this immediately. Invoice #4521 is also overdue but that can wait.",
+  {"severity": {"type":"score","instructions":"How severe is the incident?","criteria":LEVELS5},
+   "outage": {"type":"noul","instructions":"Is there an active outage?"},
+   "billing": {"type":"noul","instructions":"Is a billing topic mentioned?"},
+   "dept": {"type":"choice","instructions":"Which department?","criteria":{"sre":"Site reliability and on-call","finance":"Accounts receivable","sales":"New business","hr":"People operations"}},
+   "sentiment": {"type":"choice","instructions":"Customer sentiment?","criteria":{"pos":"Positive","neu":"Neutral","neg":"Negative"}}}),
+ ({"user":"What's the capital of Australia?","history":[]},
+  {"route": {"type":"choice","instructions":"Which model should answer?",
+     "criteria":{"small":"A small fast model for simple factual lookups","large":"A large reasoning model for complex multi-step problems","code":"A coding model"}},
+   "jailbreak": {"type":"noul","instructions":"Is this a jailbreak attempt?"},
+   "difficulty": {"type":"score","instructions":"How hard is this query?","criteria":LEVELS5}}),
+]

conversion/check_release.py

@@ -0,0 +1,40 @@
+import json, sys, time
+import numpy as np, pandas as pd
+sys.path.insert(0, sys.argv[3])
+from decision2_coreml import Decision2CoreML, encode as r_encode, question_to_row as r_row
+from transformers import AutoTokenizer
+import pack, q64ctx
+ref_root, refpath, rel = sys.argv[1], sys.argv[2], sys.argv[3]
+tok = AutoTokenizer.from_pretrained(ref_root)
+m = Decision2CoreML(rel)
+d = pd.read_parquet("data/typed/all/test-00000-of-00001.parquet")
+ref = {j["id"]: j["answers"] for j in map(json.loads, open(refpath))}
+def lab(a): return ("true" if a["noul"] >= .5 else "false") if a["type"] == "noul" else max(a["probabilities"], key=a["probabilities"].get)
+from collections import defaultdict
+acc = defaultdict(lambda: [0, 0, 0]); margins = []
+tokdiff = flips = n = 0; worst = 0; ts = []
+for r in d.itertuples():
+    st, qs = json.loads(r.state), json.loads(r.questions)
+    for qid, q in qs.items():  # tokenization vs upstream encoder
+        up = pack.rows(tok, st, {qid: q})[0][2]; mine = r_encode(r_row(st, q), m.tokenizer)
+        tokdiff += up["ids"] != mine["ids"] or up["candidate_positions"] != mine["candidate_positions"]
+    t = time.time(); a = m.system_one(state=st, questions=qs)["answers"]; ts.append((time.time() - t) * 1000)
+    for q in qs:
+        rq, aq = ref[r.id][q], a[q]; n += 1; flips += lab(rq) != lab(aq)
+        pr = rq.get("probabilities", {"true": rq.get("noul")}); pa = aq.get("probabilities", {"true": aq.get("noul")})
+        g = json.loads(r.gold)[q]["label"]; s_ = acc[rq["type"]]; s_[0] += 1; s_[1] += lab(rq) == g; s_[2] += lab(aq) == g
+        if lab(rq) != lab(aq):
+            ps = sorted(pr.values(), reverse=True) if len(pr) > 1 else sorted([pr["true"], 1 - pr["true"]], reverse=True)
+            margins.append(round(ps[0] - ps[1], 4))
+        worst = max(worst, max(abs(pr[k] - pa[k]) for k in pr))
+        for f in ("confidence", "score"):
+            if f in rq: assert abs(rq[f] - aq[f]) < 0.05, (r.id, q, f)
+        assert set(rq) == set(aq), (set(rq) ^ set(aq))
+print("flip margins (upstream top-2)", margins)
+for k, (c, h1, h2) in sorted(acc.items()): print(f"  {k:6s} n={c} acc upstream {h1/c:.3f} coreml {h2/c:.3f}")
+print(f"typed: token mismatches {tokdiff}, {n} decisions, flips {flips}, max|dp| {worst:.4f}, request p50 {np.median(ts):.1f} ms")
+t = time.time(); a = m.system_one(state=q64ctx.state, questions=q64ctx.Q); m.system_one(state=q64ctx.state, questions=q64ctx.Q)
+ts = []
+for _ in range(5):
+    t = time.time(); m.system_one(state=q64ctx.state, questions=q64ctx.Q); ts.append((time.time() - t) * 1000)
+print(f"64 questions: {len(m._calls([r_encode(r_row(q64ctx.state, q), m.tokenizer) for q in q64ctx.Q.values()]))} calls, p50 {np.median(ts):.0f} ms end to end")

conversion/combine.py

@@ -0,0 +1,11 @@
+"""Merge single-bucket packages into one multifunction package (weights stored once; macOS 15+)."""
+import sys
+import coremltools as ct
+out, specs = sys.argv[1], sys.argv[2:]  # name=path ...
+desc = ct.utils.MultiFunctionDescriptor()
+for s in specs:
+    name, path = s.split("=")
+    desc.add_function(path, src_function_name="main", target_function_name=name)
+desc.default_function_name = specs[1].split("=")[0] if len(specs) > 1 else specs[0].split("=")[0]
+ct.utils.save_multifunction(desc, out)
+print("saved", out)

conversion/convert_eos.py

@@ -0,0 +1,37 @@
+import sys, time
+import numpy as np, torch, coremltools as ct
+from eos_graph import load_chunked, load_eos, load_rows
+# packed: convert_eos.py S P N   |  rows: convert_eos.py S BxQ N
+S, N = int(sys.argv[1]), int(sys.argv[3])
+if "x" in sys.argv[2]:
+    B, Q = map(int, sys.argv[2].split("x")); P = B * Q
+    g = load_rows("eos", S, B, Q); tag = f"S{S}_R{B}x{Q}_N{N}"
+elif sys.argv[2].startswith("c"):
+    P = int(sys.argv[2][1:]); g = load_chunked("eos", S, P); tag = f"S{S}_C{P}_N{N}"
+else:
+    P = int(sys.argv[2]); g = load_eos("eos", S, P); tag = f"S{S}_P{P}_N{N}"
+chunked = sys.argv[2].startswith("c")
+out = sys.argv[4] if len(sys.argv) > 4 else f"coreml/eos_{tag}.mlpackage"
+T, R = S + P, g.cfg.rotary_dim
+ex = (torch.zeros(1, T, dtype=torch.int32), torch.zeros(T, R), torch.zeros(T, R), torch.ones(S), torch.zeros(3, S),
+      torch.eye(P), torch.zeros(3, P), torch.zeros(3, P, 3))
+if chunked:
+    M = P // 64
+    ex += (torch.eye(64).expand(M, 64, 64).contiguous(), torch.zeros(P), torch.zeros(M, 64))
+ex += (torch.zeros(N, dtype=torch.int32), torch.zeros(N, dtype=torch.int32))
+t = time.time()
+with torch.no_grad():
+    tr = torch.jit.trace(g, ex, check_trace=False)
+f16 = np.float16
+m = ct.convert(tr, convert_to="mlprogram", minimum_deployment_target=ct.target.macOS15,
+    inputs=[ct.TensorType(name="input_ids", shape=(1, T), dtype=np.int32),
+            ct.TensorType(name="cos", shape=(T, R), dtype=f16), ct.TensorType(name="sin", shape=(T, R), dtype=f16),
+            ct.TensorType(name="valid", shape=(S,), dtype=f16), ct.TensorType(name="tail_onehot", shape=(3, S), dtype=f16),
+            ct.TensorType(name="segment", shape=(P, P), dtype=f16), ct.TensorType(name="lag_keep", shape=(3, P), dtype=f16),
+            ct.TensorType(name="lag_tail", shape=(3, P, 3), dtype=f16)]
+    + ([ct.TensorType(name="seg_chunks", shape=(P // 64, 64, 64), dtype=f16), ct.TensorType(name="cont", shape=(P,), dtype=f16),
+        ct.TensorType(name="last_seg", shape=(P // 64, 64), dtype=f16)] if chunked else [])
+    + [ct.TensorType(name="cand_idx", shape=(N,), dtype=np.int32), ct.TensorType(name="query_idx", shape=(N,), dtype=np.int32)],
+    outputs=[ct.TensorType(name="logits", dtype=np.float32)], compute_precision=ct.precision.FLOAT16)
+m.short_description = "Decision-2.0-Eos-0.8B (vllm-sr): shared prefix + packed questions, one call"
+m.save(out); print("saved", out, f"{time.time()-t:.0f}s")

conversion/eos_graph.py

@@ -0,0 +1,351 @@
+"""Decision-2.0-Eos-0.8B (Qwen3.5 hybrid) as one fused Core ML graph: shared prefix + packed questions.
+
+Built on the Kev-0.8B FusedPass (kev_stages.py): the shared prefix is the "state" (S tokens, right-padded; padded
+positions are exact no-ops in the delta rule and masked in attention), the questions' suffixes are packed in P
+tokens. Each suffix restarts from the prefix's recurrent state and conv tail (segment mask), and attends to the
+prefix and itself. The readout is Decision 2.0's shared bilinear + MLP head at each option's last token and the
+question's final token (gather indices into the packed region).
+"""
+
+import json
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_file
+from torch import nn
+
+import kev_stages
+from kev_stages import MASK_VALUE, FusedPass
+from qwen35_export import TextConfig, rope_cos_sin
+
+LAGS = 3
+
+
+class EosFused(FusedPass):
+    def __init__(self, cfg: TextConfig, state_len: int, packed_len: int, lane: int | None = None, chunk: int = 64):
+        # lane = None: questions may span the whole packed region (inverse over all P)
+        super().__init__(cfg, state_len, packed_len, batch=1, max_options=2, chunk_size=chunk,
+                         lane=lane or packed_len)
+        h = cfg.hidden_size
+        del self.q, self.k, self.pointer_scale, self.inverse_temperature
+        self.embed_tokens = nn.Embedding(1, h)  # resized on load
+        self.candidate_norm = nn.LayerNorm(h)
+        self.query_norm = nn.LayerNorm(h)
+        self.key = nn.Linear(h, 256, bias=False)
+        self.query = nn.Linear(h, 256, bias=False)
+        self.candidate_mlp = nn.Linear(h, 256)
+        self.query_mlp = nn.Linear(h, 256, bias=False)
+        self.scalar = nn.Linear(256, 1, bias=False)
+
+    def forward(self, input_ids, cos, sin, valid, tail_onehot, segment, lag_keep, lag_tail, cand_idx, query_idx):
+        """input_ids [1, S+P]; cos/sin [S+P, R]; valid [S]; tail_onehot [3, S]; segment [P, P]; lag_keep [3, P];
+        lag_tail [3, P, 3]; cand_idx/query_idx [N] (indices into the packed region) -> logits [N]."""
+        S, P = self.S, self.P
+        bias = torch.cat([
+            torch.cat([self.state_causal, self.state_block], dim=1),
+            torch.cat([((1.0 - valid) * MASK_VALUE)[None, :].expand(P, S), (1.0 - segment) * MASK_VALUE], dim=1),
+        ], dim=0)
+        x = self.embed_tokens(input_ids)
+        for layer in self.decoder.layers:
+            h = layer.input_layernorm(x)
+            if layer.is_linear:
+                out = self._delta(layer.linear_attn, h, valid, tail_onehot, segment, lag_keep, lag_tail)
+            else:
+                out = self._attention(layer.self_attn, h, cos, sin, bias)
+            x = x + out
+            x = x + layer.mlp(layer.post_attention_layernorm(x))
+        xq = self.norm(x[0, S:])
+        cand = self.candidate_norm(xq[cand_idx].float())
+        qry = self.query_norm(xq[query_idx].float())
+        bilinear = (self.key(cand) * self.query(qry)).sum(-1) / 16.0
+        nonlinear = self.scalar(F.gelu(self.candidate_mlp(cand) + self.query_mlp(qry))).squeeze(-1)
+        return bilinear + nonlinear
+
+
+def load_eos(root: str, state_len: int, packed_len: int, lane: int | None = None) -> EosFused:
+    cfg = TextConfig(json.load(open(f"{root}/backbone/config.json")))
+    state = {k: v.float() for k, v in load_file(f"{root}/backbone/model.safetensors").items()}
+    model = EosFused(cfg, state_len, packed_len, lane)
+    model.decoder.load_merged(state)
+    model.norm.load_state_dict({"weight": state["norm.weight"]})
+    model.embed_tokens = nn.Embedding.from_pretrained(state["embed_tokens.weight"], freeze=True)
+    head = {k: v.float() for k, v in load_file(f"{root}/decision_head.safetensors").items()}
+    missing, unexpected = model.load_state_dict(head, strict=False)
+    assert not unexpected and not [m for m in missing if m.split(".")[0] in
+                                   ("candidate_norm", "query_norm", "key", "query", "candidate_mlp",
+                                    "query_mlp", "scalar")], (missing, unexpected)
+    return model.eval()
+
+
+def prefix_len(encoded):
+    seqs = [e["ids"] for e in encoded]
+    limit = min(min(e["candidate_positions"][0] for e in encoded), min(len(s) for s in seqs) - 1)
+    for i in range(limit):
+        if any(s[i] != seqs[0][i] for s in seqs):
+            return i
+    return limit
+
+
+def eos_inputs(cfg: TextConfig, encoded, S: int, P: int, pad_id: int, lane: int | None = None,
+               row_len: int | None = None):
+    """Fused inputs for a group of encoded rows (vendored `encode` output) sharing their common token prefix.
+    With `lane`, a question that would cross a lane boundary starts at the next lane; with `row_len`, question j
+    starts at j * row_len (EosRows)."""
+    pre = prefix_len(encoded)
+    if not 1 <= pre <= S:
+        raise ValueError(f"prefix {pre} tokens does not fit S{S}")
+    ids = np.full(S + P, pad_id, dtype=np.int64)
+    pos = np.zeros(S + P, dtype=np.int64)
+    ids[:pre] = encoded[0]["ids"][:pre]
+    pos[:pre] = np.arange(pre)
+    pos[pre:S] = pre - 1
+    valid = np.zeros(S, dtype=np.float32)
+    valid[:pre] = 1
+    tail = np.zeros((LAGS, S), dtype=np.float32)
+    for i in range(LAGS):  # tail row i = pre-conv input at position pre - LAGS + i (zero before the start)
+        p = pre - LAGS + i
+        if p >= 0:
+            tail[i, p] = 1
+    segment = np.eye(P, dtype=np.float32)
+    keep = np.zeros((LAGS, P), dtype=np.float32)
+    lag_tail = np.zeros((LAGS, P, LAGS), dtype=np.float32)
+    pos[S:] = pre
+    cand, qry, owner = [], [], []
+    start = 0
+    for j, e in enumerate(encoded):
+        suffix = e["ids"][pre:]
+        n = len(suffix)
+        if row_len:
+            start = j * row_len
+            if n > row_len:
+                raise ValueError(f"question of {n} tokens > row {row_len}")
+        elif lane and start // lane != (start + n - 1) // lane:
+            start = (start // lane + 1) * lane
+        if start + n > P:
+            raise ValueError(f"questions exceed P{P}")
+        ids[S + start : S + start + n] = suffix
+        pos[S + start : S + start + n] = np.arange(pre, pre + n)
+        segment[start : start + n, start : start + n] = np.tril(np.ones((n, n), dtype=np.float32))
+        for p in range(n):
+            for s in range(1, LAGS + 1):
+                if p >= s:
+                    keep[s - 1, start + p] = 1
+                else:
+                    lag_tail[s - 1, start + p, LAGS + p - s] = 1
+        for c in e["candidate_positions"]:
+            cand.append(start + c - pre)
+            qry.append(start + e["query_position"] - pre)
+            owner.append(j)
+        start += n
+    cos, sin = rope_cos_sin(cfg, torch.from_numpy(pos)[None].expand(3, -1))
+    return {
+        "input_ids": ids[None].astype(np.int32), "cos": cos.numpy(), "sin": sin.numpy(), "valid": valid,
+        "tail_onehot": tail, "segment": segment, "lag_keep": keep, "lag_tail": lag_tail,
+    }, cand, qry, owner, start
+
+
+def with_slots(x, cand, qry, N):
+    if len(cand) > N:
+        raise ValueError(f"{len(cand)} options > N{N}")
+    pad = N - len(cand)
+    return {**x, "cand_idx": np.array(cand + [0] * pad, dtype=np.int32),
+            "query_idx": np.array(qry + [0] * pad, dtype=np.int32)}
+
+
+class EosRows(EosFused):
+    """Questions in B rows of Q tokens (one question per row, right-padded): the question delta rule runs chunked
+    (chunk 64) from the prefix's final state, batched over rows, instead of one P x P segment inverse."""
+
+    def __init__(self, cfg: TextConfig, state_len: int, rows: int, row_len: int, chunk: int = 64):
+        assert state_len % chunk == 0 and row_len % chunk == 0
+        super().__init__(cfg, state_len, rows * row_len, lane=row_len, chunk=chunk)
+        self.B, self.Q = rows, row_len
+
+    def _delta(self, net, x, valid, tail_onehot, segment, lag_keep, lag_tail):
+        cfg, S, P, B, Q = self.cfg, self.S, self.P, self.B, self.Q
+        T = S + P
+        H, Dk, Dv, kh = cfg.lin_v_heads, cfg.lin_k_dim, cfg.lin_v_dim, cfg.lin_k_heads
+        rep = H // kh
+        pre = net.in_proj_qkv(x)[0]
+        pre_s, pre_p = pre[:S], pre[S:]
+        tail = torch.matmul(tail_onehot, pre_s)
+        weight = net.conv1d.weight[:, 0]
+        lags = weight.shape[1] - 1
+        conv_s, conv_p = pre_s * weight[:, lags], pre_p * weight[:, lags]
+        for s in range(1, lags + 1):
+            conv_s = conv_s + F.pad(pre_s[: S - s], (0, 0, s, 0)) * weight[:, lags - s]
+            shifted = F.pad(pre_p[: P - s], (0, 0, s, 0)) * lag_keep[s - 1][:, None] + torch.matmul(lag_tail[s - 1], tail)
+            conv_p = conv_p + shifted * weight[:, lags - s]
+        qkv = F.silu(torch.cat([conv_s, conv_p], dim=0))
+        q = qkv[:, : net.key_dim].reshape(T, kh, Dk)
+        k = qkv[:, net.key_dim : 2 * net.key_dim].reshape(T, kh, Dk)
+        v = qkv[:, 2 * net.key_dim :].reshape(T, H, Dv).transpose(0, 1)
+        z = net.in_proj_z(x).reshape(1, T, H, Dv)
+        beta = torch.sigmoid(net.in_proj_b(x))[0].transpose(0, 1)
+        g = (-torch.exp(net.A_log) * F.softplus(net.in_proj_a(x) + net.dt_bias))[0].transpose(0, 1)
+        q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6)
+        k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6)
+        q = (q[:, :, None].expand(T, kh, rep, Dk).reshape(T, H, Dk) * (Dk**-0.5)).transpose(0, 1)
+        k = k[:, :, None].expand(T, kh, rep, Dk).reshape(T, H, Dk).transpose(0, 1)
+
+        C, N = net.chunk, net.num_chunks
+        new_v, k_cumdecay, q_dec, k_dec, attn_intra, chunk_decay = kev_stages._delta_chunk_terms(
+            net, q[:, :S].reshape(H, N, C, Dk), k[:, :S].reshape(H, N, C, Dk), v[:, :S].reshape(H, N, C, Dv),
+            (beta[:, :S] * valid).reshape(H, N, C, 1), (g[:, :S] * valid).reshape(H, N, C))
+        state = torch.zeros(H, Dk, Dv, dtype=x.dtype)
+        outs = []
+        for i in range(N):
+            v_new = new_v[:, i] - torch.matmul(k_cumdecay[:, i], state)
+            outs.append(torch.matmul(q_dec[:, i], state) + torch.matmul(attn_intra[:, i], v_new))
+            state = state * chunk_decay[:, i] + torch.matmul(k_dec[:, i].transpose(-1, -2), v_new)
+
+        G, M = H * B, Q // C
+        new_v, k_cumdecay, q_dec, k_dec, attn_intra, chunk_decay = kev_stages._delta_chunk_terms(
+            net, q[:, S:].reshape(G, M, C, Dk), k[:, S:].reshape(G, M, C, Dk), v[:, S:].reshape(G, M, C, Dv),
+            beta[:, S:].reshape(G, M, C, 1), g[:, S:].reshape(G, M, C))
+        st = state[:, None].expand(H, B, Dk, Dv).reshape(G, Dk, Dv)
+        rows = []
+        for i in range(M):
+            v_new = new_v[:, i] - torch.matmul(k_cumdecay[:, i], st)
+            rows.append(torch.matmul(q_dec[:, i], st) + torch.matmul(attn_intra[:, i], v_new))
+            if i + 1 < M:
+                st = st * chunk_decay[:, i] + torch.matmul(k_dec[:, i].transpose(-1, -2), v_new)
+        outs.append(torch.stack(rows, dim=1).reshape(H, P, Dv))
+
+        core = torch.cat(outs, dim=1).transpose(0, 1)[None]
+        core = net.norm(core, z)
+        return net.out_proj(core.reshape(1, T, H * Dv))
+
+
+def load_rows(root: str, state_len: int, rows: int, row_len: int) -> EosRows:
+    cfg = TextConfig(json.load(open(f"{root}/backbone/config.json")))
+    state = {k: v.float() for k, v in load_file(f"{root}/backbone/model.safetensors").items()}
+    model = EosRows(cfg, state_len, rows, row_len)
+    model.decoder.load_merged(state)
+    model.norm.load_state_dict({"weight": state["norm.weight"]})
+    model.embed_tokens = nn.Embedding.from_pretrained(state["embed_tokens.weight"], freeze=True)
+    head = {k: v.float() for k, v in load_file(f"{root}/decision_head.safetensors").items()}
+    model.load_state_dict(head, strict=False)
+    return model.eval()
+
+
+class EosChunked(EosFused):
+    """Packed questions with a chunked delta rule (chunk C) instead of one P x P segment inverse.
+
+    Per chunk, a token starts either from the carried recurrent state (its question began in an earlier chunk,
+    `cont` = 1) or from the prefix's final state S0 (its question began in this chunk). `seg_chunks` [M, C, C] is the
+    segment mask restricted to each chunk (j <= t, same question); `last_seg` [M, C] marks the tokens in the question
+    of the chunk's last token, the only one whose state is carried to the next chunk."""
+
+    def __init__(self, cfg: TextConfig, state_len: int, packed_len: int, chunk: int = 64):
+        assert state_len % chunk == 0 and packed_len % chunk == 0
+        super().__init__(cfg, state_len, packed_len, lane=packed_len, chunk=chunk)
+        self.C, self.M = chunk, packed_len // chunk
+        self.register_buffer("eye_c", torch.eye(chunk), persistent=False)
+
+    def forward(self, input_ids, cos, sin, valid, tail_onehot, segment, lag_keep, lag_tail, seg_chunks, cont,
+                last_seg, cand_idx, query_idx):
+        self._extra = (seg_chunks, cont, last_seg)
+        return super().forward(input_ids, cos, sin, valid, tail_onehot, segment, lag_keep, lag_tail, cand_idx, query_idx)
+
+    def _delta(self, net, x, valid, tail_onehot, segment, lag_keep, lag_tail):
+        seg_chunks, cont, last_seg = self._extra
+        cfg, S, P, C, M = self.cfg, self.S, self.P, self.C, self.M
+        T = S + P
+        H, Dk, Dv, kh = cfg.lin_v_heads, cfg.lin_k_dim, cfg.lin_v_dim, cfg.lin_k_heads
+        rep = H // kh
+        pre = net.in_proj_qkv(x)[0]
+        pre_s, pre_p = pre[:S], pre[S:]
+        tail = torch.matmul(tail_onehot, pre_s)
+        weight = net.conv1d.weight[:, 0]
+        lags = weight.shape[1] - 1
+        conv_s, conv_p = pre_s * weight[:, lags], pre_p * weight[:, lags]
+        for s in range(1, lags + 1):
+            conv_s = conv_s + F.pad(pre_s[: S - s], (0, 0, s, 0)) * weight[:, lags - s]
+            shifted = F.pad(pre_p[: P - s], (0, 0, s, 0)) * lag_keep[s - 1][:, None] + torch.matmul(lag_tail[s - 1], tail)
+            conv_p = conv_p + shifted * weight[:, lags - s]
+        qkv = F.silu(torch.cat([conv_s, conv_p], dim=0))
+        q = qkv[:, : net.key_dim].reshape(T, kh, Dk)
+        k = qkv[:, net.key_dim : 2 * net.key_dim].reshape(T, kh, Dk)
+        v = qkv[:, 2 * net.key_dim :].reshape(T, H, Dv).transpose(0, 1)
+        z = net.in_proj_z(x).reshape(1, T, H, Dv)
+        beta = torch.sigmoid(net.in_proj_b(x))[0].transpose(0, 1)
+        g = (-torch.exp(net.A_log) * F.softplus(net.in_proj_a(x) + net.dt_bias))[0].transpose(0, 1)
+        q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6)
+        k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6)
+        q = (q[:, :, None].expand(T, kh, rep, Dk).reshape(T, H, Dk) * (Dk**-0.5)).transpose(0, 1)
+        k = k[:, :, None].expand(T, kh, rep, Dk).reshape(T, H, Dk).transpose(0, 1)
+
+        N = net.num_chunks
+        new_v, k_cumdecay, q_dec, k_dec, attn_intra, chunk_decay = kev_stages._delta_chunk_terms(
+            net, q[:, :S].reshape(H, N, net.chunk, Dk), k[:, :S].reshape(H, N, net.chunk, Dk),
+            v[:, :S].reshape(H, N, net.chunk, Dv), (beta[:, :S] * valid).reshape(H, N, net.chunk, 1),
+            (g[:, :S] * valid).reshape(H, N, net.chunk))
+        state = torch.zeros(H, Dk, Dv, dtype=x.dtype)
+        outs = []
+        for i in range(N):
+            v_new = new_v[:, i] - torch.matmul(k_cumdecay[:, i], state)
+            outs.append(torch.matmul(q_dec[:, i], state) + torch.matmul(attn_intra[:, i], v_new))
+            state = state * chunk_decay[:, i] + torch.matmul(k_dec[:, i].transpose(-1, -2), v_new)
+        s0 = state
+
+        # questions: packed, chunked, segment-aware
+        qp = q[:, S:].reshape(H, M, C, Dk)
+        kp = k[:, S:].reshape(H, M, C, Dk)
+        vp = v[:, S:].reshape(H, M, C, Dv)
+        bp = beta[:, S:].reshape(H, M, C, 1)
+        gp = g[:, S:].reshape(H, M, C, 1)
+        cum = torch.matmul(seg_chunks, gp).squeeze(-1)  # [H, M, C]: decay since max(chunk start, question start)
+        pair_decay = torch.exp((cum[..., :, None] - cum[..., None, :]) * seg_chunks + (seg_chunks - 1.0) * 1e4)
+        k_beta = kp * bp
+        kkt = torch.matmul(k_beta, kp.transpose(-1, -2)) * pair_decay
+        attn = torch.matmul(qp, kp.transpose(-1, -2)) * pair_decay
+        inv = kev_stages._unit_lower_inverse((self.eye_c + kkt * (seg_chunks - self.eye_c)).reshape(H * M, C, C), C)
+        inv = inv.reshape(H, M, C, C)
+        decay = torch.exp(cum)[..., None]  # [H, M, C, 1]
+        new_v = torch.matmul(inv, vp * bp)
+        k_cd = torch.matmul(inv, k_beta * decay)
+        q_dec = qp * decay
+        ls = last_seg[None, :, :]
+        to_last = (torch.exp((cum[..., -1:] - cum) * ls) * ls)[..., None]  # [H, M, C, 1]
+        k_last = kp * to_last
+        last_decay = torch.exp(cum[..., -1:])[..., None]  # [H, M, 1, 1]
+        cont_c = cont.reshape(M, C, 1)
+        carried = s0
+        for i in range(M):
+            c = cont_c[i]
+            v_new = new_v[:, i] - c * torch.matmul(k_cd[:, i], carried) - (1.0 - c) * torch.matmul(k_cd[:, i], s0)
+            outs.append(c * torch.matmul(q_dec[:, i], carried) + (1.0 - c) * torch.matmul(q_dec[:, i], s0)
+                        + torch.matmul(attn[:, i], v_new))
+            if i + 1 < M:
+                cl = cont_c[i, -1:]
+                init = cl * carried + (1.0 - cl) * s0
+                carried = init * last_decay[:, i] + torch.matmul(k_last[:, i].transpose(-1, -2), v_new)
+
+        core = torch.cat(outs, dim=1).transpose(0, 1)[None]
+        core = net.norm(core, z)
+        return net.out_proj(core.reshape(1, T, H * Dv))
+
+
+def load_chunked(root: str, state_len: int, packed_len: int) -> EosChunked:
+    cfg = TextConfig(json.load(open(f"{root}/backbone/config.json")))
+    state = {k: v.float() for k, v in load_file(f"{root}/backbone/model.safetensors").items()}
+    model = EosChunked(cfg, state_len, packed_len)
+    model.decoder.load_merged(state)
+    model.norm.load_state_dict({"weight": state["norm.weight"]})
+    model.embed_tokens = nn.Embedding.from_pretrained(state["embed_tokens.weight"], freeze=True)
+    model.load_state_dict({k: v.float() for k, v in load_file(f"{root}/decision_head.safetensors").items()},
+                          strict=False)
+    return model.eval()
+
+
+def chunk_inputs(segment: np.ndarray, C: int = 64):
+    """seg_chunks [M, C, C], cont [P], last_seg [M, C] from the packed segment mask [P, P]."""
+    P = segment.shape[0]
+    M = P // C
+    seg_chunks = np.stack([segment[i * C : (i + 1) * C, i * C : (i + 1) * C] for i in range(M)])
+    starts = np.argmax(segment > 0, axis=1)  # first position of each token's question (segment row's first 1)
+    cont = (starts < (np.arange(P) // C) * C).astype(np.float32)
+    last_seg = np.stack([segment[(i + 1) * C - 1, i * C : (i + 1) * C] for i in range(M)])
+    return seg_chunks.astype(np.float32), cont, last_seg.astype(np.float32)

conversion/kev_stages.py

@@ -0,0 +1,661 @@
+"""Two-stage Kev export: read the state once, then answer every question from its cache.
+
+Kev's server does this with a KV + recurrent-state cache; a row-per-question export re-reads the state for every
+question. Stage A (`StatePass`) runs the state tokens and returns, per layer, what a continuation needs:
+
+- full-attention layers: keys (after norm and RoPE) and values of the state positions;
+- Gated DeltaNet layers: the recurrent state after the last real token, and the last `conv_kernel - 1` pre-convolution
+  inputs (the causal conv's history).
+
+The state is right-padded to a bucket. Padded positions get beta = 0 and g = 0 in the delta rule, which makes them
+exact no-ops on the recurrent state (the WY inverse row of a beta = 0 position is the identity row), so the state
+comes out as of the last real token. Attention keys of padded positions are masked in stage B.
+
+Stage B (`QuestionPass`) runs B question branches at once, each right-padded to Q tokens and positioned right after
+the state, on top of that cache, and applies Kev's pointer head at <decide> and each </opt>.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from qwen35_export import DecoderChunk, RMSNorm, TextConfig, rotate_half
+
+MASK_VALUE = -1e4
+MASKED_INVERSE = True
+
+
+def _rope(x, cos, sin, r):
+    return torch.cat([x[..., :r] * cos + rotate_half(x[..., :r]) * sin, x[..., r:]], dim=-1)
+
+
+def _delta_chunk_terms(net, q, k, v, beta, g):
+    """Chunked delta-rule terms for tensors shaped [G, N, C, D] (beta [G, N, C, 1], g [G, N, C])."""
+    g_col = g.unsqueeze(-1)
+    cum = torch.matmul(net.prefix_sum_t, g_col).squeeze(-1)
+    to_end = torch.matmul(net.suffix_sum_t, g_col).squeeze(-1)
+    lead = g.shape[:2]
+    pair = torch.matmul(net.pair_sum_t, g_col).reshape(*lead, net.chunk, net.chunk)
+    pair_decay = torch.exp(pair) * net.lower_incl
+    k_beta, v_beta = k * beta, v * beta
+    kkt = torch.matmul(k_beta, k.transpose(-1, -2)) * pair_decay
+    attn_intra = torch.matmul(q, k.transpose(-1, -2)) * pair_decay
+    inv = _unit_lower_inverse(net.eye + kkt * net.strict_lower, net.chunk) if MASKED_INVERSE else \
+        net.unit_lower_inverse(net.eye + kkt * net.strict_lower)
+    new_v = torch.matmul(inv, v_beta)
+    k_cumdecay = torch.matmul(inv, k_beta * torch.exp(cum)[..., None])
+    q_dec = q * torch.exp(cum)[..., None]
+    k_dec = k * torch.exp(to_end)[..., None]
+    chunk_decay = torch.exp(cum[..., -1:])[..., None]
+    return new_v, k_cumdecay, q_dec, k_dec, attn_intra, chunk_decay
+
+
+def _delta_inputs(net, x, conv_input_prefix=None):
+    """Projections of a DeltaNet layer for x [B, L, D]. Returns q, k, v (heads expanded), z, beta, g, and the
+    pre-convolution qkv [B, L, conv]."""
+    cfg = net.cfg
+    B, L = x.shape[0], x.shape[1]
+    H, Dk, Dv, kh = cfg.lin_v_heads, cfg.lin_k_dim, cfg.lin_v_dim, cfg.lin_k_heads
+    rep = H // kh
+    pre = net.in_proj_qkv(x)  # [B, L, conv]
+    seq = pre.transpose(1, 2)  # [B, conv, L]
+    if conv_input_prefix is None:
+        seq = F.pad(seq, (cfg.conv_kernel - 1, 0))
+    else:
+        seq = torch.cat([conv_input_prefix, seq], dim=-1)
+    qkv = F.silu(net.conv1d(seq)).transpose(1, 2)  # [B, L, conv]
+    q = qkv[..., : net.key_dim].reshape(B, L, kh, Dk)
+    k = qkv[..., net.key_dim : 2 * net.key_dim].reshape(B, L, kh, Dk)
+    v = qkv[..., 2 * net.key_dim :].reshape(B, L, H, Dv)
+    z = net.in_proj_z(x).reshape(B, L, H, Dv)
+    beta = torch.sigmoid(net.in_proj_b(x))  # [B, L, H]
+    g = -torch.exp(net.A_log) * F.softplus(net.in_proj_a(x) + net.dt_bias)  # [B, L, H]
+    q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6)
+    k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6)
+    q = q[:, :, :, None].expand(B, L, kh, rep, Dk).reshape(B, L, H, Dk) * (Dk**-0.5)
+    k = k[:, :, :, None].expand(B, L, kh, rep, Dk).reshape(B, L, H, Dk)
+    return q, k, v, z, beta, g, pre
+
+
+class StatePass(nn.Module):
+    def __init__(self, cfg: TextConfig, seq_len: int, chunk_size: int = 128):
+        super().__init__()
+        self.cfg, self.S = cfg, seq_len
+        self.decoder = DecoderChunk(cfg, 0, cfg.num_layers, seq_len, chunk_size=chunk_size, with_head=False)
+
+    def _attention(self, attn, x, cos, sin):
+        cfg, S = self.cfg, self.S
+        h, kv, d, r = cfg.num_heads, cfg.num_kv_heads, cfg.head_dim, cfg.rotary_dim
+        qg = attn.q_proj(x).view(1, S, h, 2 * d)
+        q, gate = qg[..., :d], qg[..., d:].reshape(1, S, h * d)
+        q = attn.q_norm(q).transpose(1, 2)
+        k = attn.k_norm(attn.k_proj(x).view(1, S, kv, d)).transpose(1, 2)
+        v = attn.v_proj(x).view(1, S, kv, d).transpose(1, 2)
+        c, s = cos[None, None], sin[None, None]
+        q, k = _rope(q, c, s, r), _rope(k, c, s, r)
+        rep = h // kv
+        kr = k[:, :, None].expand(1, kv, rep, S, d).reshape(1, h, S, d)
+        vr = v[:, :, None].expand(1, kv, rep, S, d).reshape(1, h, S, d)
+        scores = torch.matmul(q, kr.transpose(-1, -2)) * (d**-0.5) + self.decoder.mask
+        out = torch.matmul(torch.softmax(scores, dim=-1), vr).transpose(1, 2).reshape(1, S, h * d)
+        return attn.o_proj(out * torch.sigmoid(gate)), k[0], v[0]
+
+    def _delta(self, net, x, valid, tail_onehot):
+        cfg = self.cfg
+        L, C, N = self.S, net.chunk, net.num_chunks
+        H, Dk, Dv = cfg.lin_v_heads, cfg.lin_k_dim, cfg.lin_v_dim
+        q, k, v, z, beta, g, pre = _delta_inputs(net, x)
+        tail = torch.matmul(tail_onehot, pre[0]).transpose(0, 1)  # [conv, kernel-1]
+        mask = valid[None, :, None]
+        beta, g = beta * mask, g * mask  # padded positions: no write, no decay
+        q = q[0].permute(1, 0, 2).reshape(H, N, C, Dk)
+        k = k[0].permute(1, 0, 2).reshape(H, N, C, Dk)
+        v = v[0].permute(1, 0, 2).reshape(H, N, C, Dv)
+        beta = beta[0].T.reshape(H, N, C, 1)
+        g = g[0].T.reshape(H, N, C)
+        new_v, k_cumdecay, q_dec, k_dec, attn_intra, chunk_decay = _delta_chunk_terms(net, q, k, v, beta, g)
+        state = torch.zeros(H, Dk, Dv, dtype=x.dtype)
+        outs = []
+        for i in range(N):
+            v_new = new_v[:, i] - torch.matmul(k_cumdecay[:, i], state)
+            outs.append(torch.matmul(q_dec[:, i], state) + torch.matmul(attn_intra[:, i], v_new))
+            state = state * chunk_decay[:, i] + torch.matmul(k_dec[:, i].transpose(-1, -2), v_new)
+        core = torch.stack(outs, dim=1).reshape(H, L, Dv).permute(1, 0, 2)
+        core = net.norm(core, z[0])
+        return net.out_proj(core.reshape(1, L, H * Dv)), state, tail
+
+    def forward(self, hidden, cos, sin, valid, tail_onehot):
+        """hidden [1, S, D], cos/sin [S, R], valid [S] (1 = real token), tail_onehot [kernel-1, S] selecting the last
+        real positions in order -> keys, values [A, kv, S, d]; delta states [Ld, H, Dk, Dv]; conv tails [Ld, conv, k-1]."""
+        keys, values, states, tails = [], [], [], []
+        x = hidden
+        for layer in self.decoder.layers:
+            h = layer.input_layernorm(x)
+            if layer.is_linear:
+                out, state, tail = self._delta(layer.linear_attn, h, valid, tail_onehot)
+                states.append(state)
+                tails.append(tail)
+            else:
+                out, key, value = self._attention(layer.self_attn, h, cos, sin)
+                keys.append(key)
+                values.append(value)
+            x = x + out
+            x = x + layer.mlp(layer.post_attention_layernorm(x))
+        return torch.stack(keys), torch.stack(values), torch.stack(states), torch.stack(tails)
+
+
+class QuestionPass(nn.Module):
+    def __init__(self, cfg: TextConfig, state_len: int, question_len: int, batch: int, max_options: int,
+                 pointer_dim: int = 256):
+        super().__init__()
+        self.cfg, self.S, self.Q, self.B, self.K = cfg, state_len, question_len, batch, max_options
+        # chunk = the whole question: one delta-rule chunk on top of the cached state
+        self.decoder = DecoderChunk(cfg, 0, cfg.num_layers, question_len, chunk_size=question_len, with_head=False)
+        self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
+        self.q = nn.Linear(cfg.hidden_size, pointer_dim)
+        self.k = nn.Linear(cfg.hidden_size, pointer_dim)
+        self.register_buffer("pointer_scale", torch.tensor([pointer_dim**-0.5]))
+        self.register_buffer("inverse_temperature", torch.tensor([1.0]))
+        causal = torch.full((question_len, question_len), MASK_VALUE).triu(1)
+        self.register_buffer("causal", causal)
+
+    def _attention(self, attn, x, cos, sin, key_cache, value_cache, cache_bias):
+        cfg, B, Q, S = self.cfg, self.B, self.Q, self.S
+        h, kv, d, r = cfg.num_heads, cfg.num_kv_heads, cfg.head_dim, cfg.rotary_dim
+        qg = attn.q_proj(x).view(B, Q, h, 2 * d)
+        q, gate = qg[..., :d], qg[..., d:].reshape(B, Q, h * d)
+        q = attn.q_norm(q).transpose(1, 2)  # [B, h, Q, d]
+        k = attn.k_norm(attn.k_proj(x).view(B, Q, kv, d)).transpose(1, 2)
+        v = attn.v_proj(x).view(B, Q, kv, d).transpose(1, 2)
+        c, s = cos[None, None], sin[None, None]
+        q, k = _rope(q, c, s, r), _rope(k, c, s, r)
+        k_all = torch.cat([key_cache[None].expand(B, kv, S, d), k], dim=2)  # [B, kv, S+Q, d]
+        v_all = torch.cat([value_cache[None].expand(B, kv, S, d), v], dim=2)
+        rep, T = h // kv, S + Q
+        k_all = k_all[:, :, None].expand(B, kv, rep, T, d).reshape(B, h, T, d)
+        v_all = v_all[:, :, None].expand(B, kv, rep, T, d).reshape(B, h, T, d)
+        bias = torch.cat([cache_bias[None, :].expand(Q, S), self.causal], dim=1)  # [Q, S+Q]
+        scores = torch.matmul(q, k_all.transpose(-1, -2)) * (d**-0.5) + bias
+        out = torch.matmul(torch.softmax(scores, dim=-1), v_all).transpose(1, 2).reshape(B, Q, h * d)
+        return attn.o_proj(out * torch.sigmoid(gate))
+
+    def _delta(self, net, x, state, tail):
+        cfg, B, Q = self.cfg, self.B, self.Q
+        H, Dk, Dv = cfg.lin_v_heads, cfg.lin_k_dim, cfg.lin_v_dim
+        q, k, v, z, beta, g, _ = _delta_inputs(net, x, conv_input_prefix=tail[None].expand(B, -1, -1))
+        G = B * H
+        q = q.permute(0, 2, 1, 3).reshape(G, 1, Q, Dk)
+        k = k.permute(0, 2, 1, 3).reshape(G, 1, Q, Dk)
+        v = v.permute(0, 2, 1, 3).reshape(G, 1, Q, Dv)
+        beta = beta.permute(0, 2, 1).reshape(G, 1, Q, 1)
+        g = g.permute(0, 2, 1).reshape(G, 1, Q)
+        new_v, k_cumdecay, q_dec, _, attn_intra, _ = _delta_chunk_terms(net, q, k, v, beta, g)
+        s0 = state[None].expand(B, H, Dk, Dv).reshape(G, Dk, Dv)
+        v_new = new_v[:, 0] - torch.matmul(k_cumdecay[:, 0], s0)
+        core = torch.matmul(q_dec[:, 0], s0) + torch.matmul(attn_intra[:, 0], v_new)  # [G, Q, Dv]
+        core = core.reshape(B, H, Q, Dv).permute(0, 2, 1, 3)  # [B, Q, H, Dv]
+        core = net.norm(core, z)
+        return net.out_proj(core.reshape(B, Q, H * Dv))
+
+    def forward(self, hidden, cos, sin, keys, values, cache_valid, states, tails, decide_onehot, option_onehot,
+                option_mask):
+        """hidden [B, Q, D]; cos/sin [Q, R] (positions n .. n+Q-1); keys/values [A, kv, S, d]; cache_valid [S];
+        states [Ld, H, Dk, Dv]; tails [Ld, conv, k-1]; decide_onehot [B, Q]; option_onehot [B, K, Q];
+        option_mask [B, K] -> logits [B, K] (temperature applied), probabilities [B, K]."""
+        cache_bias = (1.0 - cache_valid) * MASK_VALUE
+        x = hidden
+        a = d = 0
+        for layer in self.decoder.layers:
+            h = layer.input_layernorm(x)
+            if layer.is_linear:
+                out = self._delta(layer.linear_attn, h, states[d], tails[d])
+                d += 1
+            else:
+                out = self._attention(layer.self_attn, h, cos, sin, keys[a], values[a], cache_bias)
+                a += 1
+            x = x + out
+            x = x + layer.mlp(layer.post_attention_layernorm(x))
+        decide = self.norm(torch.matmul(decide_onehot[:, None, :], x)[:, 0])  # [B, D]
+        options = self.norm(torch.matmul(option_onehot, x))  # [B, K, D]
+        logits = (self.k(options) * self.q(decide)[:, None, :]).sum(-1) * self.pointer_scale * self.inverse_temperature
+        logits = torch.where(option_mask > 0.5, logits, torch.full_like(logits, MASK_VALUE))
+        return logits, torch.softmax(logits, dim=-1)
+
+
+def load_stages(merged, state_len: int, question_len: int, batch: int, max_options: int, chunk_size: int = 128):
+    import json
+    from pathlib import Path
+
+    from safetensors.torch import load_file
+
+    merged = Path(merged)
+    meta = json.loads((merged / "meta.json").read_text())
+    cfg = TextConfig(meta["text_config"])
+    state = load_file(str(merged / "text.safetensors"))
+    head = load_file(str(merged / "head.safetensors"))
+    first = StatePass(cfg, state_len, chunk_size)
+    first.decoder.load_merged(state)
+    second = QuestionPass(cfg, state_len, question_len, batch, max_options)
+    second.decoder.load_merged(state)
+    second.norm.load_state_dict({"weight": state["norm.weight"]})
+    second.q.load_state_dict({"weight": head["q.weight"], "bias": head["q.bias"]})
+    second.k.load_state_dict({"weight": head["k.weight"], "bias": head["k.bias"]})
+    second.pointer_scale.fill_(meta["pointer_scale"])
+    second.inverse_temperature.fill_(1.0 / meta["temperature"])
+    return first.eval(), second.eval(), cfg, meta, state["embed_tokens.weight"]
+
+
+def stage_inputs(cfg, embed, state_ids, question_rows, state_len, question_len, batch, max_options, pad_id):
+    """question_rows: [(branch ids, decide offset, option offsets)] with offsets inside the branch."""
+    from qwen35_export import rope_cos_sin
+
+    n = len(state_ids)
+    if n > state_len or len(question_rows) > batch:
+        raise ValueError(f"state {n} > {state_len} or {len(question_rows)} questions > {batch}")
+    ids = torch.full((state_len,), pad_id, dtype=torch.long)
+    ids[:n] = torch.tensor(state_ids)
+    positions = torch.arange(state_len)
+    cos, sin = rope_cos_sin(cfg, positions.unsqueeze(0).expand(3, -1))
+    valid = torch.zeros(state_len)
+    valid[:n] = 1
+    kernel = cfg.conv_kernel - 1
+    tail = torch.zeros(kernel, state_len)
+    for j in range(kernel):
+        p = n - kernel + j
+        if p >= 0:
+            tail[j, p] = 1
+    state_in = (embed[ids].unsqueeze(0), cos, sin, valid, tail)
+
+    q_ids = torch.full((batch, question_len), pad_id, dtype=torch.long)
+    decide = torch.zeros(batch, question_len)
+    options = torch.zeros(batch, max_options, question_len)
+    mask = torch.zeros(batch, max_options)
+    for b, (branch, d, opts) in enumerate(question_rows):
+        if len(branch) > question_len or len(opts) > max_options:
+            raise ValueError("question exceeds its bucket")
+        q_ids[b, : len(branch)] = torch.tensor(branch)
+        decide[b, d] = 1
+        for i, o in enumerate(opts):
+            options[b, i, o] = 1
+            mask[b, i] = 1
+    q_pos = torch.arange(n, n + question_len)
+    q_cos, q_sin = rope_cos_sin(cfg, q_pos.unsqueeze(0).expand(3, -1))
+    return state_in, (embed[q_ids], q_cos, q_sin, valid, decide, options, mask)
+
+
+class PackedQuestionPass(nn.Module):
+    """Stage B with the questions packed end to end in one sequence of P tokens instead of B padded rows.
+
+    `segment` [P, P] is 1 where j <= i in the same question (padding positions are their own one-token segments):
+    it is the attention mask over the question tokens, the segment-local cumulative-decay operator of the delta rule
+    (every question restarts from the cached state), and its strict part is the WY triangle. A question's first
+    conv_kernel - 1 positions read the state's conv tail instead of the previous question: `lag_keep` [k-1, P] keeps
+    in-segment lags and `lag_tail` [k-1, P, k-1] selects tail entries for the rest.
+    """
+
+    def __init__(self, cfg: TextConfig, state_len: int, packed_len: int, batch: int, max_options: int,
+                 pointer_dim: int = 256):
+        super().__init__()
+        self.cfg, self.S, self.P, self.B, self.K = cfg, state_len, packed_len, batch, max_options
+        self.decoder = DecoderChunk(cfg, 0, cfg.num_layers, packed_len, chunk_size=packed_len, with_head=False)
+        self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
+        self.q = nn.Linear(cfg.hidden_size, pointer_dim)
+        self.k = nn.Linear(cfg.hidden_size, pointer_dim)
+        self.register_buffer("pointer_scale", torch.tensor([pointer_dim**-0.5]))
+        self.register_buffer("inverse_temperature", torch.tensor([1.0]))
+        self.register_buffer("eye", torch.eye(packed_len))
+
+    def _attention(self, attn, x, cos, sin, key_cache, value_cache, bias):
+        cfg, P, S = self.cfg, self.P, self.S
+        h, kv, d, r = cfg.num_heads, cfg.num_kv_heads, cfg.head_dim, cfg.rotary_dim
+        qg = attn.q_proj(x).view(1, P, h, 2 * d)
+        q, gate = qg[..., :d], qg[..., d:].reshape(1, P, h * d)
+        q = attn.q_norm(q).transpose(1, 2)
+        k = attn.k_norm(attn.k_proj(x).view(1, P, kv, d)).transpose(1, 2)
+        v = attn.v_proj(x).view(1, P, kv, d).transpose(1, 2)
+        c, s = cos[None, None], sin[None, None]
+        q, k = _rope(q, c, s, r), _rope(k, c, s, r)
+        k_all = torch.cat([key_cache[None], k], dim=2)  # [1, kv, S+P, d]
+        v_all = torch.cat([value_cache[None], v], dim=2)
+        rep, T = h // kv, S + P
+        k_all = k_all[:, :, None].expand(1, kv, rep, T, d).reshape(1, h, T, d)
+        v_all = v_all[:, :, None].expand(1, kv, rep, T, d).reshape(1, h, T, d)
+        scores = torch.matmul(q, k_all.transpose(-1, -2)) * (d**-0.5) + bias
+        out = torch.matmul(torch.softmax(scores, dim=-1), v_all).transpose(1, 2).reshape(1, P, h * d)
+        return attn.o_proj(out * torch.sigmoid(gate))
+
+    def _delta(self, net, x, state, tail, segment, lag_keep, lag_tail):
+        cfg, P = self.cfg, self.P
+        H, Dk, Dv, kh = cfg.lin_v_heads, cfg.lin_k_dim, cfg.lin_v_dim, cfg.lin_k_heads
+        rep = H // kh
+        pre = net.in_proj_qkv(x)[0]  # [P, conv]
+        weight = net.conv1d.weight[:, 0]  # [conv, kernel]
+        lags = weight.shape[1] - 1
+        conv = pre * weight[:, lags]
+        for s in range(1, lags + 1):
+            shifted = F.pad(pre[: P - s], (0, 0, s, 0)) * lag_keep[s - 1][:, None]
+            shifted = shifted + torch.matmul(lag_tail[s - 1], tail.transpose(0, 1))
+            conv = conv + shifted * weight[:, lags - s]
+        qkv = F.silu(conv)
+        q = qkv[:, : net.key_dim].reshape(P, kh, Dk)
+        k = qkv[:, net.key_dim : 2 * net.key_dim].reshape(P, kh, Dk)
+        v = qkv[:, 2 * net.key_dim :].reshape(P, H, Dv).transpose(0, 1)  # [H, P, Dv]
+        z = net.in_proj_z(x).reshape(1, P, H, Dv)
+        beta = torch.sigmoid(net.in_proj_b(x))[0].transpose(0, 1)[..., None]  # [H, P, 1]
+        g = (-torch.exp(net.A_log) * F.softplus(net.in_proj_a(x) + net.dt_bias))[0].transpose(0, 1)  # [H, P]
+        q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6)
+        k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6)
+        q = (q[:, :, None].expand(P, kh, rep, Dk).reshape(P, H, Dk) * (Dk**-0.5)).transpose(0, 1)  # [H, P, Dk]
+        k = k[:, :, None].expand(P, kh, rep, Dk).reshape(P, H, Dk).transpose(0, 1)
+        cum = torch.matmul(g, segment.transpose(0, 1))  # [H, P]: decay since the segment start
+        pair = (cum[:, :, None] - cum[:, None, :]) * segment + (segment - 1.0) * 1e4
+        pair_decay = torch.exp(pair)
+        k_beta, v_beta = k * beta, v * beta
+        kkt = torch.matmul(k_beta, k.transpose(-1, -2)) * pair_decay
+        attn_intra = torch.matmul(q, k.transpose(-1, -2)) * pair_decay
+        inv = _unit_lower_inverse(self.eye + kkt * (segment - self.eye), P)
+        decay = torch.exp(cum)[..., None]
+        v_new = torch.matmul(inv, v_beta) - torch.matmul(torch.matmul(inv, k_beta * decay), state)
+        core = torch.matmul(q * decay, state) + torch.matmul(attn_intra, v_new)  # [H, P, Dv]
+        core = net.norm(core.transpose(0, 1)[None], z)
+        return net.out_proj(core.reshape(1, P, H * Dv))
+
+    def forward(self, hidden, cos, sin, keys, values, cache_valid, states, tails, segment, lag_keep, lag_tail,
+                decide_onehot, option_onehot, option_mask):
+        """hidden [1, P, D]; cos/sin [P, R] (each question at positions n, n+1, ...); keys/values [A, kv, S, d];
+        cache_valid [S]; states [Ld, H, Dk, Dv]; tails [Ld, conv, k-1]; segment [P, P]; lag_keep [k-1, P];
+        lag_tail [k-1, P, k-1]; decide_onehot [B, P]; option_onehot [B, K, P]; option_mask [B, K]
+        -> logits [B, K] (temperature applied), probabilities [B, K]."""
+        bias = torch.cat([((1.0 - cache_valid) * MASK_VALUE)[None, :].expand(self.P, self.S),
+                          (1.0 - segment) * MASK_VALUE], dim=1)
+        x = hidden
+        a = d = 0
+        for layer in self.decoder.layers:
+            h = layer.input_layernorm(x)
+            if layer.is_linear:
+                out = self._delta(layer.linear_attn, h, states[d], tails[d], segment, lag_keep, lag_tail)
+                d += 1
+            else:
+                out = self._attention(layer.self_attn, h, cos, sin, keys[a], values[a], bias)
+                a += 1
+            x = x + out
+            x = x + layer.mlp(layer.post_attention_layernorm(x))
+        decide = self.norm(torch.matmul(decide_onehot, x[0]))  # [B, D]
+        options = self.norm(torch.matmul(option_onehot, x[0]))  # [B, K, D]
+        logits = (self.k(options) * self.q(decide)[:, None, :]).sum(-1) * self.pointer_scale * self.inverse_temperature
+        logits = torch.where(option_mask > 0.5, logits, torch.full_like(logits, MASK_VALUE))
+        return logits, torch.softmax(logits, dim=-1)
+
+
+def load_packed(merged, state_len: int, packed_len: int, batch: int, max_options: int):
+    import json
+    from pathlib import Path
+
+    from safetensors.torch import load_file
+
+    merged = Path(merged)
+    meta = json.loads((merged / "meta.json").read_text())
+    cfg = TextConfig(meta["text_config"])
+    state = load_file(str(merged / "text.safetensors"))
+    head = load_file(str(merged / "head.safetensors"))
+    model = PackedQuestionPass(cfg, state_len, packed_len, batch, max_options)
+    model.decoder.load_merged(state)
+    model.norm.load_state_dict({"weight": state["norm.weight"]})
+    model.q.load_state_dict({"weight": head["q.weight"], "bias": head["q.bias"]})
+    model.k.load_state_dict({"weight": head["k.weight"], "bias": head["k.bias"]})
+    model.pointer_scale.fill_(meta["pointer_scale"])
+    model.inverse_temperature.fill_(1.0 / meta["temperature"])
+    return model.eval()
+
+
+def packed_inputs(cfg, embed, state_count, question_rows, packed_len, batch, max_options, pad_id, lags=3, lane=None):
+    """question_rows: [(branch ids, decide offset, option offsets)] -> PackedQuestionPass inputs after the cache.
+    With `lane`, a question that would cross a lane boundary starts at the next lane."""
+    from qwen35_export import rope_cos_sin
+
+    P = packed_len
+    ids = torch.full((P,), pad_id, dtype=torch.long)
+    positions = torch.full((P,), state_count, dtype=torch.long)
+    segment = torch.eye(P)
+    lag_keep = torch.zeros(lags, P)
+    lag_tail = torch.zeros(lags, P, lags)
+    decide = torch.zeros(batch, P)
+    options = torch.zeros(batch, max_options, P)
+    mask = torch.zeros(batch, max_options)
+    start = 0
+    for b, (branch, d, opts) in enumerate(question_rows):
+        n = len(branch)
+        if lane and start // lane != (start + n - 1) // lane:
+            start = (start // lane + 1) * lane
+        if start + n > P or b >= batch or len(opts) > max_options:
+            raise ValueError("questions exceed the packed bucket")
+        ids[start : start + n] = torch.tensor(branch)
+        positions[start : start + n] = torch.arange(state_count, state_count + n)
+        segment[start : start + n, start : start + n] = torch.ones(n, n).tril()
+        for p in range(n):
+            for s in range(1, lags + 1):
+                if p >= s:
+                    lag_keep[s - 1, start + p] = 1
+                else:
+                    lag_tail[s - 1, start + p, lags + p - s] = 1
+        decide[b, start + d] = 1
+        for i, o in enumerate(opts):
+            options[b, i, start + o] = 1
+            mask[b, i] = 1
+        start += n
+    cos, sin = rope_cos_sin(cfg, positions.unsqueeze(0).expand(3, -1))
+    return embed[ids].unsqueeze(0), cos, sin, segment, lag_keep, lag_tail, decide, options, mask
+
+
+def _masked_block_inverse(t, C, block=None):
+    """Same blocking as GatedDeltaNet.unit_lower_inverse, kept as one [.., C, C] matrix: with D the inverse of t's
+    diagonal s-blocks and L = t's lower-left s-blocks inside each diagonal 2s-block, the 2s-block inverse is
+    D - D L D (L D L = 0). Four ops per level instead of block extraction, slicing and concatenation."""
+    idx = torch.arange(C)
+    inv = None
+    s = 1
+    while s < min(C, block or C):
+        same_pair = (idx[:, None] // (2 * s)) == (idx[None, :] // (2 * s))
+        lower_left = same_pair & ((idx[:, None] // s) % 2 == 1) & ((idx[None, :] // s) % 2 == 0)
+        mask = lower_left.to(t.dtype)
+        if inv is None:
+            inv = torch.eye(C, dtype=t.dtype) - t * mask
+        else:
+            inv = inv - torch.matmul(torch.matmul(inv, t * mask), inv)
+        s *= 2
+    return inv
+
+
+def _unit_lower_inverse(t, C, block=None):
+    """GatedDeltaNet.unit_lower_inverse for a [G, C, C] unit-lower-triangular batch of any power-of-two C. `block`:
+    t is known block-diagonal in blocks of that size (packed questions never cross a lane)."""
+    if MASKED_INVERSE:
+        return _masked_block_inverse(t, C, block)
+    b = t.shape[0]
+    inv = torch.ones(b, C, 1, 1, dtype=t.dtype)
+    s = 1
+    while s < C:
+        n2 = C // (2 * s)
+        blocks = (t.reshape(b, n2, 2 * s, n2, 2 * s) * torch.eye(n2, dtype=t.dtype)[:, None, :, None]).sum(-2)
+        x = blocks[..., s:, :s]
+        pair = inv.reshape(b, n2, 2, s, s)
+        a_inv, b_inv = pair[..., 0, :, :], pair[..., 1, :, :]
+        lower = -torch.matmul(b_inv, torch.matmul(x, a_inv))
+        inv = torch.cat([torch.cat([a_inv, torch.zeros_like(a_inv)], dim=-1), torch.cat([lower, b_inv], dim=-1)], dim=-2)
+        s *= 2
+    return inv.reshape(b, C, C)
+
+
+class FusedPass(nn.Module):
+    """State pass and packed question pass in one call: sequence = S state tokens (right-padded) + P packed question
+    tokens. Projections, norms, MLPs and attention run once over all S + P positions (attention under one combined
+    mask); only the delta rule splits: chunked over the state, then one segment-masked chunk for the questions on top
+    of the state's final recurrent state and conv tail."""
+
+    def __init__(self, cfg: TextConfig, state_len: int, packed_len: int, batch: int, max_options: int,
+                 chunk_size: int = 64, pointer_dim: int = 256, lane: int = 128):
+        super().__init__()
+        self.cfg, self.S, self.P, self.B, self.K = cfg, state_len, packed_len, batch, max_options
+        self.lane = min(lane, packed_len)
+        self.decoder = DecoderChunk(cfg, 0, cfg.num_layers, state_len, chunk_size=min(chunk_size, state_len),
+                                    with_head=False)
+        self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
+        self.q = nn.Linear(cfg.hidden_size, pointer_dim)
+        self.k = nn.Linear(cfg.hidden_size, pointer_dim)
+        self.register_buffer("pointer_scale", torch.tensor([pointer_dim**-0.5]))
+        self.register_buffer("inverse_temperature", torch.tensor([1.0]))
+        self.register_buffer("eye_p", torch.eye(packed_len))
+        self.register_buffer("state_causal", torch.full((state_len, state_len), MASK_VALUE).triu(1))
+        self.register_buffer("state_block", torch.full((state_len, packed_len), MASK_VALUE))
+
+    def _attention(self, attn, x, cos, sin, bias):
+        cfg, T = self.cfg, self.S + self.P
+        h, kv, d, r = cfg.num_heads, cfg.num_kv_heads, cfg.head_dim, cfg.rotary_dim
+        qg = attn.q_proj(x).view(1, T, h, 2 * d)
+        q, gate = qg[..., :d], qg[..., d:].reshape(1, T, h * d)
+        q = attn.q_norm(q).transpose(1, 2)
+        k = attn.k_norm(attn.k_proj(x).view(1, T, kv, d)).transpose(1, 2)
+        v = attn.v_proj(x).view(1, T, kv, d).transpose(1, 2)
+        c, s = cos[None, None], sin[None, None]
+        q, k = _rope(q, c, s, r), _rope(k, c, s, r)
+        rep = h // kv
+        k = k[:, :, None].expand(1, kv, rep, T, d).reshape(1, h, T, d)
+        v = v[:, :, None].expand(1, kv, rep, T, d).reshape(1, h, T, d)
+        scores = torch.matmul(q, k.transpose(-1, -2)) * (d**-0.5) + bias
+        out = torch.matmul(torch.softmax(scores, dim=-1), v).transpose(1, 2).reshape(1, T, h * d)
+        return attn.o_proj(out * torch.sigmoid(gate))
+
+    def _delta(self, net, x, valid, tail_onehot, segment, lag_keep, lag_tail):
+        cfg, S, P = self.cfg, self.S, self.P
+        T = S + P
+        H, Dk, Dv, kh = cfg.lin_v_heads, cfg.lin_k_dim, cfg.lin_v_dim, cfg.lin_k_heads
+        rep = H // kh
+        pre = net.in_proj_qkv(x)[0]  # [T, conv]
+        pre_s, pre_p = pre[:S], pre[S:]
+        tail = torch.matmul(tail_onehot, pre_s)  # [k-1, conv]
+        weight = net.conv1d.weight[:, 0]  # [conv, kernel]
+        lags = weight.shape[1] - 1
+        conv_s, conv_p = pre_s * weight[:, lags], pre_p * weight[:, lags]
+        for s in range(1, lags + 1):
+            conv_s = conv_s + F.pad(pre_s[: S - s], (0, 0, s, 0)) * weight[:, lags - s]
+            shifted = F.pad(pre_p[: P - s], (0, 0, s, 0)) * lag_keep[s - 1][:, None] + torch.matmul(lag_tail[s - 1], tail)
+            conv_p = conv_p + shifted * weight[:, lags - s]
+        qkv = F.silu(torch.cat([conv_s, conv_p], dim=0))  # [T, conv]
+        q = qkv[:, : net.key_dim].reshape(T, kh, Dk)
+        k = qkv[:, net.key_dim : 2 * net.key_dim].reshape(T, kh, Dk)
+        v = qkv[:, 2 * net.key_dim :].reshape(T, H, Dv).transpose(0, 1)  # [H, T, Dv]
+        z = net.in_proj_z(x).reshape(1, T, H, Dv)
+        beta = torch.sigmoid(net.in_proj_b(x))[0].transpose(0, 1)  # [H, T]
+        g = (-torch.exp(net.A_log) * F.softplus(net.in_proj_a(x) + net.dt_bias))[0].transpose(0, 1)  # [H, T]
+        q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6)
+        k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6)
+        q = (q[:, :, None].expand(T, kh, rep, Dk).reshape(T, H, Dk) * (Dk**-0.5)).transpose(0, 1)  # [H, T, Dk]
+        k = k[:, :, None].expand(T, kh, rep, Dk).reshape(T, H, Dk).transpose(0, 1)
+
+        # state: chunked, padded positions are no-ops
+        C, N = net.chunk, net.num_chunks
+        beta_s, g_s = beta[:, :S] * valid, g[:, :S] * valid
+        new_v, k_cumdecay, q_dec, k_dec, attn_intra, chunk_decay = _delta_chunk_terms(
+            net, q[:, :S].reshape(H, N, C, Dk), k[:, :S].reshape(H, N, C, Dk), v[:, :S].reshape(H, N, C, Dv),
+            beta_s.reshape(H, N, C, 1), g_s.reshape(H, N, C))
+        state = torch.zeros(H, Dk, Dv, dtype=x.dtype)
+        outs = []
+        for i in range(N):
+            v_new = new_v[:, i] - torch.matmul(k_cumdecay[:, i], state)
+            outs.append(torch.matmul(q_dec[:, i], state) + torch.matmul(attn_intra[:, i], v_new))
+            state = state * chunk_decay[:, i] + torch.matmul(k_dec[:, i].transpose(-1, -2), v_new)
+
+        # questions: one segment-masked chunk from the final state
+        qp, kp, vp = q[:, S:], k[:, S:], v[:, S:]
+        beta_p = beta[:, S:][..., None]
+        cum = torch.matmul(g[:, S:], segment.transpose(0, 1))  # [H, P]
+        pair_decay = torch.exp((cum[:, :, None] - cum[:, None, :]) * segment + (segment - 1.0) * 1e4)
+        k_beta = kp * beta_p
+        kkt = torch.matmul(k_beta, kp.transpose(-1, -2)) * pair_decay
+        attn_p = torch.matmul(qp, kp.transpose(-1, -2)) * pair_decay
+        inv = _unit_lower_inverse(self.eye_p + kkt * (segment - self.eye_p), P, self.lane)
+        decay = torch.exp(cum)[..., None]
+        v_new = torch.matmul(inv, vp * beta_p) - torch.matmul(torch.matmul(inv, k_beta * decay), state)
+        outs.append(torch.matmul(qp * decay, state) + torch.matmul(attn_p, v_new))
+
+        core = torch.cat(outs, dim=1).transpose(0, 1)[None]  # [1, T, H, Dv]
+        core = net.norm(core, z)
+        return net.out_proj(core.reshape(1, T, H * Dv))
+
+    def forward(self, hidden, cos, sin, valid, tail_onehot, segment, lag_keep, lag_tail, decide_onehot, option_onehot,
+                option_mask):
+        """hidden [1, S+P, D]; cos/sin [S+P, R]; valid [S]; tail_onehot [k-1, S]; segment [P, P]; lag_keep [k-1, P];
+        lag_tail [k-1, P, k-1]; decide_onehot [B, P]; option_onehot [B, K, P]; option_mask [B, K]
+        -> logits [B, K] (temperature applied), probabilities [B, K]."""
+        S, P = self.S, self.P
+        bias = torch.cat([
+            torch.cat([self.state_causal, self.state_block], dim=1),
+            torch.cat([((1.0 - valid) * MASK_VALUE)[None, :].expand(P, S), (1.0 - segment) * MASK_VALUE], dim=1),
+        ], dim=0)
+        x = hidden
+        for layer in self.decoder.layers:
+            h = layer.input_layernorm(x)
+            if layer.is_linear:
+                out = self._delta(layer.linear_attn, h, valid, tail_onehot, segment, lag_keep, lag_tail)
+            else:
+                out = self._attention(layer.self_attn, h, cos, sin, bias)
+            x = x + out
+            x = x + layer.mlp(layer.post_attention_layernorm(x))
+        xq = x[0, S:]
+        decide = self.norm(torch.matmul(decide_onehot, xq))
+        options = self.norm(torch.matmul(option_onehot, xq))
+        logits = (self.k(options) * self.q(decide)[:, None, :]).sum(-1) * self.pointer_scale * self.inverse_temperature
+        logits = torch.where(option_mask > 0.5, logits, torch.full_like(logits, MASK_VALUE))
+        return logits, torch.softmax(logits, dim=-1)
+
+
+def load_fused(merged, state_len: int, packed_len: int, batch: int, max_options: int, chunk_size: int = 64,
+               lane: int = 128):
+    import json
+    from pathlib import Path
+
+    from safetensors.torch import load_file
+
+    merged = Path(merged)
+    meta = json.loads((merged / "meta.json").read_text())
+    cfg = TextConfig(meta["text_config"])
+    state = load_file(str(merged / "text.safetensors"))
+    head = load_file(str(merged / "head.safetensors"))
+    model = FusedPass(cfg, state_len, packed_len, batch, max_options, chunk_size, lane=lane)
+    model.decoder.load_merged(state)
+    model.norm.load_state_dict({"weight": state["norm.weight"]})
+    model.q.load_state_dict({"weight": head["q.weight"], "bias": head["q.bias"]})
+    model.k.load_state_dict({"weight": head["k.weight"], "bias": head["k.bias"]})
+    model.pointer_scale.fill_(meta["pointer_scale"])
+    model.inverse_temperature.fill_(1.0 / meta["temperature"])
+    return model.eval()
+
+
+def fused_inputs(cfg, embed, state_ids, question_rows, state_len, packed_len, batch, max_options, pad_id, lane=128):
+    state_in, _ = stage_inputs(cfg, embed, state_ids, [], state_len, 32, 2, max_options, pad_id)
+    s_hidden, s_cos, s_sin, valid, tail = state_in
+    p_hidden, p_cos, p_sin, segment, keep, lag_tail, decide, options, mask = packed_inputs(
+        cfg, embed, len(state_ids), question_rows, packed_len, batch, max_options, pad_id, lane=min(lane, packed_len))
+    return (torch.cat([s_hidden, p_hidden], dim=1), torch.cat([s_cos, p_cos]), torch.cat([s_sin, p_sin]), valid, tail,
+            segment, keep, lag_tail, decide, options, mask)
+
+
+def pack_groups(lengths, packed_len, batch, lane=None):
+    """Greedy packing (the order is kept) of question lengths into calls of `packed_len` tokens and `batch` readouts,
+    no question crossing a lane. Returns lists of question indices."""
+    groups, current, start = [], [], 0
+    for i, n in enumerate(lengths):
+        if n > (lane or packed_len):
+            raise ValueError(f"question of {n} tokens exceeds the lane")
+        placed = start
+        if lane and placed // lane != (placed + n - 1) // lane:
+            placed = (placed // lane + 1) * lane
+        if current and (placed + n > packed_len or len(current) == batch):
+            groups.append(current)
+            current, placed = [], 0
+        current.append(i)
+        start = placed + n
+    if current:
+        groups.append(current)
+    return groups

conversion/pack.py

@@ -0,0 +1,121 @@
+"""Host side: System One request -> packed tree inputs, and logits -> answers (vendored code)."""
+
+import sys
+
+import numpy as np
+
+sys.path.insert(0, "kai")
+from decision2._vendor.dev2model.decision_model import encode  # noqa: E402
+from decision2._vendor.dev2model.infer import product_answer, question_to_row  # noqa: E402
+from decision2._vendor.dev2model.score_bias import apply as apply_score_bias  # noqa: E402
+
+NEG = -1e4
+
+
+def rows(tokenizer, state, questions, cap=8192):
+    item = {"id": "request", "state": state}
+    out = []
+    for qid, q in questions.items():
+        row = question_to_row(item, qid, q)
+        out.append((qid, row, encode(row, tokenizer, cap)))
+    return out
+
+
+def pack(jobs, L, N, pad_id):
+    """Shared prefix once, then each suffix; suffix tokens see prefix + own suffix (causal)."""
+    seqs = [e["ids"] for _, _, e in jobs]
+    P = min(min(e["candidate_positions"][0] for _, _, e in jobs), min(len(s) for s in seqs) - 1)
+    for i in range(P):
+        if any(s[i] != seqs[0][i] for s in seqs):
+            P = i
+            break
+    ids, pos, seg = list(seqs[0][:P]), list(range(P)), [-1] * P
+    starts = []
+    for j, s in enumerate(seqs):
+        starts.append(len(ids) - P)
+        ids += s[P:]
+        pos += range(P, len(s))
+        seg += [j] * (len(s) - P)
+    T = len(ids)
+    if T > L:
+        raise ValueError(f"packed length {T} > {L}")
+    cand, qry, owner = [], [], []
+    for j, (_, _, e) in enumerate(jobs):
+        for c in e["candidate_positions"]:
+            cand.append(c if c < P else c + starts[j])
+            qry.append(e["query_position"] + starts[j])
+            owner.append(j)
+    if len(cand) > N:
+        raise ValueError(f"{len(cand)} candidates > {N}")
+    seg = np.array(seg + [-2] * (L - T))
+    p = np.array(pos + [0] * (L - T))
+    i = np.arange(L)
+    causal = i[None, :] <= i[:, None]
+    same = (seg[None, :] == seg[:, None]) | (seg[None, :] == -1)
+    allow = causal & same & (seg[None, :] != -2)
+    allow[np.arange(L), np.arange(L)] = True  # padding rows attend to themselves
+    mask = np.where(allow, 0.0, NEG).astype(np.float32)[None, None]
+    pad = N - len(cand)
+    return {
+        "input_ids": np.array([ids + [pad_id] * (L - T)], dtype=np.int32),
+        "position_ids": p[None].astype(np.int32),
+        "mask": mask,
+        "cand_idx": np.array(cand + [0] * pad, dtype=np.int32),
+        "query_idx": np.array(qry + [0] * pad, dtype=np.int32),
+    }, owner, T
+
+
+def answers(jobs, logits, owner, score_bias=None, temps=None):
+    per = [[] for _ in jobs]
+    for j, v in zip(owner, logits[: len(owner)]):
+        per[j].append(float(v))
+    out = {}
+    for (qid, row, e), values in zip(jobs, per):
+        if score_bias is not None and row["task_type"] == "score":
+            values = apply_score_bias(score_bias, values, len(row["options"]))
+        out[qid] = product_answer(
+            row["task_type"], e["keys"], values, (temps or {}).get(row["task_type"], 1.0),
+            [o["description"] for o in row["options"]],
+        )
+    return out
+
+
+def size(jobs):
+    """(packed tokens, candidates) of one packed call, without building it."""
+    seqs = [e["ids"] for _, _, e in jobs]
+    P = min(min(e["candidate_positions"][0] for _, _, e in jobs), min(len(s) for s in seqs) - 1)
+    for i in range(P):
+        if any(s[i] != seqs[0][i] for s in seqs):
+            P = i
+            break
+    return P + sum(len(s) - P for s in seqs), sum(len(e["keys"]) for _, _, e in jobs)
+
+
+def fits(jobs, L, N):
+    T, C = size(jobs)
+    return T <= L and C <= N
+
+
+def chunks(jobs, L, N):
+    """Greedy split of a request's questions into groups that each pack into one L/N call."""
+    groups, cur = [], []
+    for j in range(len(jobs)):
+        if not cur or fits([jobs[i] for i in cur + [j]], L, N):
+            cur.append(j)
+        else:
+            groups.append(cur)
+            cur = [j]
+    return groups + [cur]
+
+
+def run(model, jobs, L, N, pad_id):
+    """Logits per job (list of lists), over as many packed calls as needed."""
+    per = [None] * len(jobs)
+    for group in chunks(jobs, L, N):
+        sub = [jobs[i] for i in group]
+        x, owner, _ = pack(sub, L, N, pad_id)
+        x["mask"] = x["mask"].astype("float16")
+        out = model.predict(x)["logits"]
+        for k, i in enumerate(group):
+            per[i] = [float(v) for v, o in zip(out, owner) if o == k]
+    return per

conversion/parity_chunked_torch.py

@@ -0,0 +1,26 @@
+import sys, torch
+from transformers import AutoModel, AutoTokenizer
+from cases import CASES
+import pack, q64ctx
+from eos_graph import load_chunked, eos_inputs, with_slots, chunk_inputs
+S, P = int(sys.argv[1]), int(sys.argv[2])
+ref = AutoModel.from_pretrained("eos", trust_remote_code=True, device="mps")
+tok = AutoTokenizer.from_pretrained("eos"); pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
+g = load_chunked("eos", S, P); worst = 0
+cases = CASES + [(q64ctx.state, dict(list(q64ctx.Q.items())[:9]))]
+for state, qs in cases:
+    r = ref.system_one(state=state, questions=qs)["answers"]
+    jobs = pack.rows(tok, state, qs)
+    x, cand, qry, owner, used = eos_inputs(g.cfg, [e for _, _, e in jobs], S, P, pad)
+    x = with_slots(x, cand, qry, 64)
+    x["seg_chunks"], x["cont"], x["last_seg"] = chunk_inputs(x["segment"])
+    names = ("input_ids", "cos", "sin", "valid", "tail_onehot", "segment", "lag_keep", "lag_tail", "seg_chunks", "cont", "last_seg", "cand_idx", "query_idx")
+    with torch.no_grad():
+        lg = g(*[torch.from_numpy(x[k]) for k in names])
+    a = pack.answers(jobs, lg.tolist(), owner, None)
+    d = 0
+    for q in qs:
+        pr = r[q].get("probabilities", {"true": r[q].get("noul")}); pa = a[q].get("probabilities", {"true": a[q].get("noul")})
+        d = max(d, max(abs(pr[k] - pa[k]) for k in pr))
+    worst = max(worst, d); print(f"{len(qs)} questions, packed {used}, max dp {d:.2e}")
+print("WORST", worst)

conversion/q64ctx.py

@@ -0,0 +1,20 @@
+import json, sys, time
+import numpy as np, coremltools as ct
+from transformers import AutoModel, AutoTokenizer
+import pack
+state = ("Subject: URGENT - payment failed for order #88213\n\nHi, I tried to pay for my annual Pro plan renewal "
+ "with my Visa ending 4421 but the charge was declined twice, and now my team of 12 is locked out of the dashboard. "
+ "We have a board demo tomorrow at 9am PST. I already contacted my bank and they say the card is fine. Can you "
+ "please extend access for 48 hours and send me an invoice I can pay by wire instead? Also, the export to CSV "
+ "feature has been timing out since last week. Thanks, Dana (Head of Ops, Brightline Logistics)")
+yn = ["Is this about a failed payment?", "Is the customer locked out?", "Is there a deadline mentioned?", "Does the customer ask for an invoice?",
+ "Is a bug reported?", "Is the customer angry?", "Is this spam?", "Is this a phishing attempt?", "Is the sender an existing customer?",
+ "Does the message mention a team size?", "Does the customer request a refund?", "Does the customer want to cancel?",
+ "Is wire transfer mentioned?", "Is a credit card mentioned?", "Is personal data shared?", "Does it need a reply today?"]
+Q = {}
+for i in range(64):
+    k = i % 4
+    if k == 0: Q[f"q{i}"] = {"type": "noul", "instructions": yn[(i // 4) % 16] + ("" if i < 32 else " (strict)")}
+    elif k == 1: Q[f"q{i}"] = {"type": "choice", "instructions": f"Which team should own item {i}?", "criteria": {"billing": "Payments and invoices", "support": "Product bugs", "sales": "Plans and upgrades", "security": "Fraud and access"}}
+    elif k == 2: Q[f"q{i}"] = {"type": "score", "instructions": f"Urgency rating {i}?", "criteria": ["None", "Low", "Medium", "High", "Critical"]}
+    else: Q[f"q{i}"] = {"type": "choice", "instructions": f"Overall sentiment {i}?", "criteria": {"pos": "Positive", "neu": "Neutral", "neg": "Negative"}}

conversion/qwen35_export.py

@@ -0,0 +1,333 @@
+"""Trace-friendly, prefill-only Qwen3.5 text decoder (copied from models/computer-use/cua-s1-4b/coreml, mobius PR #104).
+
+Cua-S1-4B reads one forward pass: the logits of the answer letters A..Z at the
+last prompt position. Nothing is generated, so this graph has no KV cache, no
+recurrent-state I/O and no full vocabulary head:
+
+- inputs are `inputs_embeds` (host-side embedding gather, so a multimodal
+  caller can splice vision features in) plus host-computed M-RoPE `cos`/`sin`;
+- prompts are right-padded to a fixed bucket; the model is causal and the gated
+  delta rule is a forward scan, so padding after the last real token cannot
+  change anything before it;
+- the last chunk gathers the hidden state at `last_index` (one-hot matmul),
+  applies the final norm and multiplies by the 26 tied-embedding rows of the
+  letter tokens.
+
+The gated delta rule uses the chunked (WY/UT) form. The unit-lower-triangular
+solve is a recursive 2x2 block inverse (log2(chunk) batched matmul levels). The
+nilpotent power series (I + N)(I + N^2)... is exact algebraically but its terms
+grow combinatorially before cancelling and overflow on real Qwen3.5 layers.
+Within-chunk decay differences are computed as masked sums of the per-step
+decays instead of differences of cumulative sums, which keeps them accurate in
+fp16 when the cumulative decay is large.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class TextConfig:
+    def __init__(self, cfg: dict):
+        self.hidden_size = cfg["hidden_size"]
+        self.intermediate_size = cfg["intermediate_size"]
+        self.num_layers = cfg["num_hidden_layers"]
+        self.layer_types = cfg["layer_types"]
+        self.eps = cfg["rms_norm_eps"]
+        self.num_heads = cfg["num_attention_heads"]
+        self.num_kv_heads = cfg["num_key_value_heads"]
+        self.head_dim = cfg["head_dim"]
+        self.lin_k_heads = cfg["linear_num_key_heads"]
+        self.lin_v_heads = cfg["linear_num_value_heads"]
+        self.lin_k_dim = cfg["linear_key_head_dim"]
+        self.lin_v_dim = cfg["linear_value_head_dim"]
+        self.conv_kernel = cfg["linear_conv_kernel_dim"]
+        rope = cfg["rope_parameters"]
+        self.rope_theta = rope["rope_theta"]
+        self.rotary_dim = int(self.head_dim * rope.get("partial_rotary_factor", 1.0))
+        self.mrope_section = rope.get("mrope_section", [11, 11, 10])
+
+
+def rope_cos_sin(cfg: TextConfig, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Interleaved M-RoPE tables. `position_ids` is [3, L] (t, h, w); returns [L, rotary_dim] fp32 each.
+
+    Mirrors `Qwen3_5TextRotaryEmbedding` (computed in float64 then cast, host side)."""
+    dim = cfg.rotary_dim
+    inv_freq = 1.0 / (cfg.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float64) / dim))
+    freqs = position_ids.to(torch.float64)[:, :, None] * inv_freq[None, None, :]  # [3, L, dim/2]
+    thw = freqs[0].clone()
+    for axis, offset in ((1, 1), (2, 2)):
+        length = cfg.mrope_section[axis] * 3
+        thw[:, offset:length:3] = freqs[axis][:, offset:length:3]
+    emb = torch.cat([thw, thw], dim=-1)
+    return emb.cos().float(), emb.sin().float()
+
+
+class RMSNorm(nn.Module):
+    """Qwen3.5 zero-centred RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w)."""
+
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x):
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x * (1.0 + self.weight)
+
+
+class MLP(nn.Module):
+    def __init__(self, cfg: TextConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+def rotate_half(x):
+    half = x.shape[-1] // 2
+    return torch.cat([-x[..., half:], x[..., :half]], dim=-1)
+
+
+class FullAttention(nn.Module):
+    def __init__(self, cfg: TextConfig):
+        super().__init__()
+        self.cfg = cfg
+        h, kv, d = cfg.num_heads, cfg.num_kv_heads, cfg.head_dim
+        self.q_proj = nn.Linear(cfg.hidden_size, h * d * 2, bias=False)
+        self.k_proj = nn.Linear(cfg.hidden_size, kv * d, bias=False)
+        self.v_proj = nn.Linear(cfg.hidden_size, kv * d, bias=False)
+        self.o_proj = nn.Linear(h * d, cfg.hidden_size, bias=False)
+        self.q_norm = RMSNorm(d, cfg.eps)
+        self.k_norm = RMSNorm(d, cfg.eps)
+
+    def forward(self, x, cos, sin, mask):
+        cfg = self.cfg
+        L = x.shape[1]
+        h, kv, d, r = cfg.num_heads, cfg.num_kv_heads, cfg.head_dim, cfg.rotary_dim
+        qg = self.q_proj(x).view(1, L, h, 2 * d)
+        q, gate = qg[..., :d], qg[..., d:]
+        gate = gate.reshape(1, L, h * d)
+        q = self.q_norm(q).transpose(1, 2)  # [1, h, L, d]
+        k = self.k_norm(self.k_proj(x).view(1, L, kv, d)).transpose(1, 2)
+        v = self.v_proj(x).view(1, L, kv, d).transpose(1, 2)
+
+        c, s = cos[None, None], sin[None, None]  # [1, 1, L, r]
+        q = torch.cat([q[..., :r] * c + rotate_half(q[..., :r]) * s, q[..., r:]], dim=-1)
+        k = torch.cat([k[..., :r] * c + rotate_half(k[..., :r]) * s, k[..., r:]], dim=-1)
+
+        rep = h // kv
+        k = k[:, :, None].expand(1, kv, rep, L, d).reshape(1, h, L, d)
+        v = v[:, :, None].expand(1, kv, rep, L, d).reshape(1, h, L, d)
+        scores = torch.matmul(q, k.transpose(-1, -2)) * (d**-0.5) + mask
+        out = torch.matmul(torch.softmax(scores, dim=-1), v)  # [1, h, L, d]
+        out = out.transpose(1, 2).reshape(1, L, h * d)
+        return self.o_proj(out * torch.sigmoid(gate))
+
+
+class GatedRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x, gate):
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return self.weight * x * F.silu(gate)
+
+
+class GatedDeltaNet(nn.Module):
+    def __init__(self, cfg: TextConfig, seq_len: int, chunk_size: int):
+        super().__init__()
+        self.cfg = cfg
+        self.chunk = chunk_size
+        assert seq_len % chunk_size == 0, "bucket length must be a multiple of the delta-rule chunk size"
+        self.num_chunks = seq_len // chunk_size
+        kd = cfg.lin_k_heads * cfg.lin_k_dim
+        vd = cfg.lin_v_heads * cfg.lin_v_dim
+        self.key_dim, self.value_dim = kd, vd
+        self.conv_dim = 2 * kd + vd
+        self.in_proj_qkv = nn.Linear(cfg.hidden_size, self.conv_dim, bias=False)
+        self.in_proj_z = nn.Linear(cfg.hidden_size, vd, bias=False)
+        self.in_proj_b = nn.Linear(cfg.hidden_size, cfg.lin_v_heads, bias=False)
+        self.in_proj_a = nn.Linear(cfg.hidden_size, cfg.lin_v_heads, bias=False)
+        self.conv1d = nn.Conv1d(self.conv_dim, self.conv_dim, cfg.conv_kernel, groups=self.conv_dim, bias=False)
+        self.dt_bias = nn.Parameter(torch.ones(cfg.lin_v_heads))
+        self.A_log = nn.Parameter(torch.zeros(cfg.lin_v_heads))
+        self.norm = GatedRMSNorm(cfg.lin_v_dim, cfg.eps)
+        self.out_proj = nn.Linear(vd, cfg.hidden_size, bias=False)
+
+        C = chunk_size
+        idx = torch.arange(C)
+        # incl[i, j, k] = 1 if j < k <= i : sum_k g_k = cum_i - cum_j (lower triangle incl. diagonal is 0 on i == j)
+        incl = ((idx[None, None, :] > idx[None, :, None]) & (idx[None, None, :] <= idx[:, None, None])).float()
+        self.register_buffer("pair_sum_t", incl.reshape(C * C, C), persistent=False)  # [C*C, C]
+        self.register_buffer("prefix_sum_t", (idx[None, :] <= idx[:, None]).float(), persistent=False)  # [C, C]
+        self.register_buffer("suffix_sum_t", (idx[None, :] > idx[:, None]).float(), persistent=False)
+        self.register_buffer("lower_incl", (idx[None, :] <= idx[:, None]).float(), persistent=False)
+        self.register_buffer("strict_lower", (idx[None, :] < idx[:, None]).float(), persistent=False)
+        self.register_buffer("eye", torch.eye(C), persistent=False)
+        assert C & (C - 1) == 0, "delta-rule chunk size must be a power of two"
+        for n in (C // (2 * s) for s in (2**i for i in range(int(math.log2(C))))):
+            self.register_buffer(f"block_eye_{n}", torch.eye(n), persistent=False)
+
+    def unit_lower_inverse(self, t: torch.Tensor) -> torch.Tensor:
+        """Inverse of unit-lower-triangular [..., C, C] by recursive 2x2 blocking:
+        [[A, 0], [X, B]]^-1 = [[A^-1, 0], [-B^-1 X A^-1, B^-1]].
+
+        Stable like forward substitution (it only forms the true inverse's sub-blocks),
+        but log2(C) batched levels instead of C sequential row updates."""
+        C = self.chunk
+        lead = t.shape[:-2]
+        t = t.reshape(-1, C, C)  # Core ML tensors are rank <= 5
+        b = t.shape[0]
+        inv = torch.ones(b, C, 1, 1, dtype=t.dtype, device=t.device)  # 1x1 diagonal blocks of a unit triangle
+        s = 1
+        while s < C:
+            n2 = C // (2 * s)
+            # diagonal 2s x 2s blocks of t, then their lower-left s x s part
+            blocks = (t.reshape(b, n2, 2 * s, n2, 2 * s) * getattr(self, f"block_eye_{n2}")[:, None, :, None]).sum(-2)
+            x = blocks[..., s:, :s]  # [b, n2, s, s]
+            pair = inv.reshape(b, n2, 2, s, s)
+            a_inv, b_inv = pair[..., 0, :, :], pair[..., 1, :, :]
+            lower = -torch.matmul(b_inv, torch.matmul(x, a_inv))
+            top = torch.cat([a_inv, torch.zeros_like(a_inv)], dim=-1)
+            bottom = torch.cat([lower, b_inv], dim=-1)
+            inv = torch.cat([top, bottom], dim=-2)  # [b, n2, 2s, 2s]
+            s *= 2
+        return inv.reshape(*lead, C, C)
+
+    def forward(self, x):
+        cfg = self.cfg
+        L = x.shape[1]
+        C, N = self.chunk, self.num_chunks
+        H, Dk, Dv = cfg.lin_v_heads, cfg.lin_k_dim, cfg.lin_v_dim
+        rep = cfg.lin_v_heads // cfg.lin_k_heads
+
+        qkv = self.in_proj_qkv(x).transpose(1, 2)  # [1, conv_dim, L]
+        qkv = F.silu(self.conv1d(F.pad(qkv, (cfg.conv_kernel - 1, 0)))).transpose(1, 2)
+        q = qkv[..., : self.key_dim].reshape(1, L, cfg.lin_k_heads, Dk)
+        k = qkv[..., self.key_dim : 2 * self.key_dim].reshape(1, L, cfg.lin_k_heads, Dk)
+        v = qkv[..., 2 * self.key_dim :].reshape(1, L, H, Dv)
+        z = self.in_proj_z(x).reshape(1, L, H, Dv)
+        beta = torch.sigmoid(self.in_proj_b(x))  # [1, L, H]
+        g = -torch.exp(self.A_log) * F.softplus(self.in_proj_a(x) + self.dt_bias)  # [1, L, H], <= 0
+
+        # l2norm (FLA convention) then repeat k heads up to v heads
+        q = q * torch.rsqrt((q * q).sum(-1, keepdim=True) + 1e-6)
+        k = k * torch.rsqrt((k * k).sum(-1, keepdim=True) + 1e-6)
+        q = q[:, :, :, None].expand(1, L, cfg.lin_k_heads, rep, Dk).reshape(1, L, H, Dk)
+        k = k[:, :, :, None].expand(1, L, cfg.lin_k_heads, rep, Dk).reshape(1, L, H, Dk)
+        q = q * (Dk**-0.5)
+
+        # -> [H, N, C, D]
+        q = q[0].permute(1, 0, 2).reshape(H, N, C, Dk)
+        k = k[0].permute(1, 0, 2).reshape(H, N, C, Dk)
+        v = v[0].permute(1, 0, 2).reshape(H, N, C, Dv)
+        beta = beta[0].T.reshape(H, N, C, 1)
+        g = g[0].T.reshape(H, N, C)
+
+        # constant-first matmuls: `x @ const` would lower to Core ML `linear` and be picked up by weight compression
+        g_col = g.unsqueeze(-1)  # [H, N, C, 1]
+        cum = torch.matmul(self.prefix_sum_t, g_col).squeeze(-1)  # cum_i = sum_{k<=i} g_k
+        to_end = torch.matmul(self.suffix_sum_t, g_col).squeeze(-1)  # sum_{k>i} g_k = cum_last - cum_i
+        pair = torch.matmul(self.pair_sum_t, g_col).reshape(H, N, C, C)  # cum_i - cum_j (0 above diagonal)
+        pair_decay = torch.exp(pair) * self.lower_incl
+
+        k_beta = k * beta
+        v_beta = v * beta
+        kkt = torch.matmul(k_beta, k.transpose(-1, -2)) * pair_decay
+        attn_intra = torch.matmul(q, k.transpose(-1, -2)) * pair_decay
+
+        inv = self.unit_lower_inverse(self.eye + kkt * self.strict_lower)
+        new_v = torch.matmul(inv, v_beta)  # [H, N, C, Dv]
+        k_cumdecay = torch.matmul(inv, k_beta * torch.exp(cum)[..., None])  # [H, N, C, Dk]
+
+        q_dec = q * torch.exp(cum)[..., None]
+        k_dec = k * torch.exp(to_end)[..., None]
+        chunk_decay = torch.exp(cum[..., -1:])[..., None]  # [H, N, 1, 1]
+
+        state = torch.zeros(H, Dk, Dv, dtype=x.dtype, device=x.device)
+        outs = []
+        for i in range(N):
+            v_new = new_v[:, i] - torch.matmul(k_cumdecay[:, i], state)
+            outs.append(torch.matmul(q_dec[:, i], state) + torch.matmul(attn_intra[:, i], v_new))
+            if i + 1 < N:
+                state = state * chunk_decay[:, i] + torch.matmul(k_dec[:, i].transpose(-1, -2), v_new)
+        core = torch.stack(outs, dim=1).reshape(H, L, Dv).permute(1, 0, 2)  # [L, H, Dv]
+
+        core = self.norm(core, z[0])
+        return self.out_proj(core.reshape(1, L, H * Dv))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, cfg: TextConfig, layer_idx: int, seq_len: int, chunk_size: int):
+        super().__init__()
+        self.is_linear = cfg.layer_types[layer_idx] == "linear_attention"
+        if self.is_linear:
+            self.linear_attn = GatedDeltaNet(cfg, seq_len, chunk_size)
+        else:
+            self.self_attn = FullAttention(cfg)
+        self.mlp = MLP(cfg)
+        self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.eps)
+        self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.eps)
+
+    def forward(self, x, cos, sin, mask):
+        h = self.input_layernorm(x)
+        h = self.linear_attn(h) if self.is_linear else self.self_attn(h, cos, sin, mask)
+        x = x + h
+        return x + self.mlp(self.post_attention_layernorm(x))
+
+
+class DecoderChunk(nn.Module):
+    """Layers [start, end). The last chunk (with_head) returns letter logits [1, n_letters]."""
+
+    def __init__(
+        self,
+        cfg: TextConfig,
+        start: int,
+        end: int,
+        seq_len: int,
+        chunk_size: int = 64,
+        with_head: bool = False,
+        n_letters: int = 26,
+    ):
+        super().__init__()
+        self.start, self.end, self.with_head = start, end, with_head
+        self.layers = nn.ModuleList([DecoderLayer(cfg, i, seq_len, chunk_size) for i in range(start, end)])
+        mask = torch.full((seq_len, seq_len), -1e4).triu(1)
+        self.register_buffer("mask", mask[None, None], persistent=False)
+        if with_head:
+            self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
+            self.letter_head = nn.Linear(cfg.hidden_size, n_letters, bias=False)
+
+    def forward(self, hidden, cos, sin, last_onehot=None):
+        for layer in self.layers:
+            hidden = layer(hidden, cos, sin, self.mask)
+        if not self.with_head:
+            return hidden
+        last = torch.matmul(last_onehot, hidden[0])  # [1, hidden]
+        return self.letter_head(self.norm(last))
+
+    def load_merged(self, state: dict[str, torch.Tensor], letter_rows: torch.Tensor | None = None):
+        """`state` uses Qwen3_5TextModel names relative to the language model (`layers.N...`, `norm.weight`)."""
+        own = {}
+        for local, global_idx in enumerate(range(self.start, self.end)):
+            prefix = f"layers.{global_idx}."
+            for key, value in state.items():
+                if key.startswith(prefix):
+                    own[f"layers.{local}.{key[len(prefix):]}"] = value
+        if self.with_head:
+            own["norm.weight"] = state["norm.weight"]
+            own["letter_head.weight"] = letter_rows
+        missing, unexpected = self.load_state_dict({k: v.float() for k, v in own.items()}, strict=False)
+        missing = [m for m in missing if not m.endswith(("mask",))]
+        if missing or unexpected:
+            raise RuntimeError(f"weight mismatch: missing={missing[:8]} unexpected={unexpected[:8]}")

conversion/time_eos.py

@@ -0,0 +1,15 @@
+import sys, time
+import numpy as np, coremltools as ct
+for path in sys.argv[1:]:
+    m = ct.models.MLModel(path, compute_units=ct.ComputeUnit.CPU_AND_GPU)
+    spec = {i.name: tuple(i.type.multiArrayType.shape) for i in m.get_spec().description.input}
+    x = {}
+    for k, shp in spec.items():
+        dt = np.int32 if k in ("input_ids", "cand_idx", "query_idx") else np.float16
+        x[k] = np.zeros(shp, dtype=dt)
+    x["valid"][:] = 1; x["segment"][:] = np.eye(spec["segment"][0])
+    if "seg_chunks" in x: x["seg_chunks"][:] = np.eye(64)
+    m.predict(x); ts = []
+    for _ in range(10):
+        t = time.time(); m.predict(x); ts.append((time.time() - t) * 1000)
+    print(f"{path.split('/')[-1]:40s} p50 {np.median(ts):6.1f} ms")

conversion/typed_ref.py

@@ -0,0 +1,13 @@
+"""Upstream answers for typed-decisions TEST (400 requests, 2,000 decisions), cached as JSONL."""
+import json, sys, time
+import pandas as pd
+from transformers import AutoModel
+root, out = sys.argv[1], sys.argv[2]
+d = pd.read_parquet("data/typed/all/test-00000-of-00001.parquet")
+m = AutoModel.from_pretrained(root, trust_remote_code=True, device="mps")
+t = time.time()
+with open(out, "w") as f:
+    for r in d.itertuples():
+        a = m.system_one(state=json.loads(r.state), questions=json.loads(r.questions))
+        f.write(json.dumps({"id": r.id, "answers": a["answers"], "tokens": a["usage"]["input_tokens"]}) + "\n")
+print(f"{len(d)} requests in {time.time()-t:.0f}s")

coreml_config.json

@@ -0,0 +1,24 @@
+{
+  "model_name": "Decision-2.0-Eos-0.8B",
+  "source": "vllm-sr/Decision-2.0-Eos-0.8B",
+  "source_revision": "3594047d69f476f1d01cf84c593e213fc3a4dfe0",
+  "backbone": "qwen3_5",
+  "package": "Decision2EosPacked.mlpackage",
+  "functions": {"S128_C256_N32": 34.3, "S256_C512_N64": 64.1, "S512_C768_N96": 108.6, "S512_C1024_N128": 131.0, "S1024_C1024_N128": 176.0},
+  "functions_note": "function name -> measured predict p50 ms on M5 Pro GPU, used to choose chunking. S = shared-prefix tokens, C = packed question tokens (multiple of 64), N = option slots",
+  "pad_token_id": 248044,
+  "rope": {"rotary_dim": 64, "rope_theta": 10000000},
+  "inputs": {
+    "input_ids": "[1, S+C] int32: shared prefix right-padded to S, then every question's suffix, then padding",
+    "cos, sin": "[S+C, 64] float16: RoPE tables (prefix positions 0.., each suffix continues from the prefix length)",
+    "valid": "[S] float16: 1 for real prefix tokens",
+    "tail_onehot": "[3, S] float16: selects the prefix's last 3 pre-convolution inputs (conv history)",
+    "segment": "[C, C] float16: 1 where j <= i within the same question (padding = own segment)",
+    "lag_keep, lag_tail": "[3, C], [3, C, 3] float16: a question's first 3 conv lags read the prefix tail",
+    "seg_chunks": "[C/64, 64, 64] float16: segment restricted to each 64-token chunk",
+    "cont": "[C] float16: 1 if the token's question started in an earlier chunk",
+    "last_seg": "[C/64, 64] float16: tokens in the question of each chunk's last token",
+    "cand_idx, query_idx": "[N] int32: packed-region index of each option's last token and of its question's final token"
+  },
+  "outputs": {"logits": "[N] float32: one raw logit per option (before softmax)"}
+}

decision2_coreml.py

@@ -0,0 +1,277 @@
+"""Decision 2.0 System One on Core ML: every question of a request in one packed call.
+
+Supports the Qwen3 package (Kai, `L*_N*` functions) and the Qwen3.5 hybrid package (Eos, `S*_C*_N*` functions).
+
+Needs only numpy, tokenizers and coremltools. Prompt rendering, option endpoints, Score
+offsets and answer normalization follow the upstream runtime (vllm-sr Decision 2.0,
+Apache-2.0: decision2/_vendor/dev2model/{decision_model,infer,score_bias}.py).
+
+    from decision2_coreml import Decision2CoreML
+    model = Decision2CoreML(".")  # this repo's folder
+    model.system_one(state="...", questions={...})
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import Any
+
+import coremltools as ct
+import numpy as np
+from tokenizers import Tokenizer
+
+MAX_OPTIONS = 255
+NEG = -1e4
+
+
+def canonical(value: Any) -> str:
+    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False)
+
+
+def _payload(value: Any) -> str:
+    return value if isinstance(value, str) else canonical(value)
+
+
+def question_to_row(state: Any, question: dict[str, Any]) -> dict[str, Any]:
+    kind = question.get("type")
+    if kind not in ("choice", "noul", "score"):
+        raise ValueError("unsupported question type")
+    if not question.get("instructions"):
+        raise ValueError("missing question instructions")
+    criteria = question.get("criteria")
+    if kind == "score":
+        if not isinstance(criteria, list) or not 2 <= len(criteria) <= 10:
+            raise ValueError("score criteria must be an ordered list of 2..10 levels")
+        options = [{"key": str(i), "description": d} for i, d in enumerate(criteria)]
+    else:
+        if kind == "noul":
+            criteria = criteria or {}
+            if not isinstance(criteria, dict) or set(criteria) - {"false", "true"}:
+                raise ValueError("noul requires only false and true criteria")
+            if len(criteria) < 2:
+                criteria = {"false": criteria.get("false", "No"), "true": criteria.get("true", "Yes")}
+        if not isinstance(criteria, dict) or not 2 <= len(criteria) <= MAX_OPTIONS:
+            raise ValueError("choice criteria must be an object with 2..255 options")
+        options = [{"key": k, "description": d} for k, d in criteria.items()]
+    return {"state": state, "instructions": question["instructions"], "options": options, "task_type": kind}
+
+
+def encode(row: dict[str, Any], tokenizer: Tokenizer) -> dict[str, Any]:
+    prefix = (
+        f"Context:\n{_payload(row['state'])}\n\n"
+        f"Task type: {row['task_type']}\nQuestion:\n{_payload(row['instructions'])}\nOptions:"
+    )
+    ids = tokenizer.encode(prefix, add_special_tokens=False).ids
+    endpoints = []
+    for option in row["options"]:
+        text = "\n<option>\n" + canonical({"key": option["key"], "description": option["description"]}) + "\n</option>"
+        ids.extend(tokenizer.encode(text, add_special_tokens=False).ids)
+        endpoints.append(len(ids) - 1)
+    suffix = "\n\nSelect the single option best supported by the context and instructions.\nDecision:"
+    ids.extend(tokenizer.encode(suffix, add_special_tokens=False).ids)
+    return {"ids": ids, "candidate_positions": endpoints, "query_position": len(ids) - 1,
+            "keys": [o["key"] for o in row["options"]]}
+
+
+def _prefix(encoded: list[dict[str, Any]]) -> int:
+    seqs = [e["ids"] for e in encoded]
+    limit = min(min(e["candidate_positions"][0] for e in encoded), min(len(s) for s in seqs) - 1)
+    for i in range(limit):
+        if any(s[i] != seqs[0][i] for s in seqs):
+            return i
+    return limit
+
+
+def packed_size(encoded: list[dict[str, Any]]) -> tuple[int, int]:
+    P = _prefix(encoded)
+    return P + sum(len(e["ids"]) - P for e in encoded), sum(len(e["keys"]) for e in encoded)
+
+
+def pack(encoded: list[dict[str, Any]], L: int, N: int, pad_id: int) -> tuple[dict[str, np.ndarray], list[int]]:
+    """Shared prefix once, then each question's suffix; a suffix sees the prefix and itself (causal)."""
+    P = _prefix(encoded)
+    ids, pos, seg = list(encoded[0]["ids"][:P]), list(range(P)), [-1] * P
+    cand, qry, owner = [], [], []
+    for j, e in enumerate(encoded):
+        shift = len(ids) - P
+        ids += e["ids"][P:]
+        pos += range(P, len(e["ids"]))
+        seg += [j] * (len(e["ids"]) - P)
+        for c in e["candidate_positions"]:
+            cand.append(c if c < P else c + shift)
+            qry.append(e["query_position"] + shift)
+            owner.append(j)
+    T = len(ids)
+    if T > L or len(cand) > N:
+        raise ValueError(f"{T} tokens / {len(cand)} options exceed L{L}_N{N}")
+    s = np.array(seg + [-2] * (L - T))
+    i = np.arange(L)
+    allow = (i[None, :] <= i[:, None]) & ((s[None, :] == s[:, None]) | (s[None, :] == -1)) & (s[None, :] != -2)
+    allow[i, i] = True
+    return {
+        "input_ids": np.array([ids + [pad_id] * (L - T)], dtype=np.int32),
+        "position_ids": np.array([pos + [0] * (L - T)], dtype=np.int32),
+        "mask": np.where(allow, 0.0, NEG).astype(np.float16)[None, None],
+        "cand_idx": np.array(cand + [0] * (N - len(cand)), dtype=np.int32),
+        "query_idx": np.array(qry + [0] * (N - len(cand)), dtype=np.int32),
+    }, owner
+
+
+def pack_hybrid(encoded: list[dict[str, Any]], S: int, P: int, N: int, pad_id: int, rope: dict[str, Any],
+                lags: int = 3, chunk: int = 64) -> tuple[dict[str, np.ndarray], list[int]]:
+    """Qwen3.5 (Gated DeltaNet hybrid) graph inputs: the shared prefix right-padded to S, the questions' suffixes
+    packed in P, plus the masks that make every suffix restart from the prefix's state (see README)."""
+    pre = _prefix(encoded)
+    total = sum(len(e["ids"]) - pre for e in encoded)
+    options = sum(len(e["keys"]) for e in encoded)
+    if not 1 <= pre <= S or total > P or options > N:
+        raise ValueError(f"prefix {pre} / {total} question tokens / {options} options exceed S{S}_C{P}_N{N}")
+    ids = np.full(S + P, pad_id, dtype=np.int32)
+    pos = np.zeros(S + P, dtype=np.float64)
+    ids[:pre] = encoded[0]["ids"][:pre]
+    pos[:pre] = np.arange(pre)
+    pos[pre:] = pre
+    valid = np.zeros(S, dtype=np.float16)
+    valid[:pre] = 1
+    tail = np.zeros((lags, S), dtype=np.float16)
+    for i in range(lags):
+        if pre - lags + i >= 0:
+            tail[i, pre - lags + i] = 1
+    segment = np.eye(P, dtype=np.float16)
+    keep = np.zeros((lags, P), dtype=np.float16)
+    lag_tail = np.zeros((lags, P, lags), dtype=np.float16)
+    cand, qry, owner = [], [], []
+    start = 0
+    for j, e in enumerate(encoded):
+        n = len(e["ids"]) - pre
+        ids[S + start : S + start + n] = e["ids"][pre:]
+        pos[S + start : S + start + n] = np.arange(pre, pre + n)
+        segment[start : start + n, start : start + n] = np.tril(np.ones((n, n)))
+        for p in range(n):
+            for s in range(1, lags + 1):
+                if p >= s:
+                    keep[s - 1, start + p] = 1
+                else:
+                    lag_tail[s - 1, start + p, lags + p - s] = 1
+        for c in e["candidate_positions"]:
+            cand.append(start + c - pre)
+            qry.append(start + e["query_position"] - pre)
+            owner.append(j)
+        start += n
+    dim = rope["rotary_dim"]
+    inv_freq = 1.0 / (rope["rope_theta"] ** (np.arange(0, dim, 2, dtype=np.float64) / dim))
+    freqs = pos[:, None] * inv_freq[None, :]  # text-only M-RoPE: all three axes share the position
+    emb = np.concatenate([freqs, freqs], axis=-1)
+    M = P // chunk
+    starts = np.argmax(segment > 0, axis=1)
+    pad = N - len(cand)
+    return {
+        "input_ids": ids[None], "cos": np.cos(emb).astype(np.float16), "sin": np.sin(emb).astype(np.float16),
+        "valid": valid, "tail_onehot": tail, "segment": segment, "lag_keep": keep, "lag_tail": lag_tail,
+        "seg_chunks": np.stack([segment[i * chunk : (i + 1) * chunk, i * chunk : (i + 1) * chunk] for i in range(M)]),
+        "cont": (starts < (np.arange(P) // chunk) * chunk).astype(np.float16),
+        "last_seg": np.stack([segment[(i + 1) * chunk - 1, i * chunk : (i + 1) * chunk] for i in range(M)]),
+        "cand_idx": np.array(cand + [0] * pad, dtype=np.int32),
+        "query_idx": np.array(qry + [0] * pad, dtype=np.int32),
+    }, owner
+
+
+def answer(kind: str, keys: list[str], logits: list[float], descriptions: list[Any]) -> dict[str, Any]:
+    top = max(logits)
+    exps = [math.exp(v - top) for v in logits]
+    total = sum(exps)
+    probs = [v / total for v in exps]
+    pmap = dict(zip(keys, probs))
+    if kind == "noul":
+        return {"type": "noul", "noul": pmap["true"]}
+    entropy = -sum(p * math.log(p) for p in probs if p > 0)
+    confidence = max(0.0, min(1.0, 1.0 - entropy / math.log(len(keys))))
+    if kind == "score":
+        return {"type": "score", "score": sum(int(k) * pmap[k] for k in keys), "probabilities": pmap,
+                "confidence": confidence,
+                "legend": {k: d if isinstance(d, str) else canonical(d) for k, d in zip(keys, descriptions)}}
+    return {"type": "choice", "choice": keys[probs.index(max(probs))], "probabilities": pmap, "confidence": confidence}
+
+
+def _dims(name: str) -> dict[str, int]:
+    return {part[0]: int(part[1:]) for part in name.split("_")}
+
+
+class Decision2CoreML:
+    def __init__(self, root: str | Path = ".", package: str | None = None,
+                 compute_units: ct.ComputeUnit = ct.ComputeUnit.CPU_AND_GPU):
+        root = Path(root)
+        self.config = json.loads((root / "coreml_config.json").read_text())
+        self.tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
+        self.pad_id = self.config["pad_token_id"]
+        self.hybrid = self.config.get("backbone", "qwen3") == "qwen3_5"
+        bias = root / "score_bias.json"
+        offsets = json.loads(bias.read_text())["offsets"] if bias.exists() else {}
+        self.score_bias = {int(k): v for k, v in offsets.items()}
+        path = str(root / (package or self.config["package"]))
+        self.cost_ms = self.config["functions"]
+        self.buckets = [(name, _dims(name), ct.models.MLModel(path, function_name=name, compute_units=compute_units))
+                        for name in self.config["functions"]]  # smallest first
+
+    def _fits(self, d: dict[str, int], encoded: list[dict[str, Any]]) -> bool:
+        options = sum(len(e["keys"]) for e in encoded)
+        if self.hybrid:
+            pre = _prefix(encoded)
+            return pre <= d["S"] and sum(len(e["ids"]) - pre for e in encoded) <= d["C"] and options <= d["N"]
+        return packed_size(encoded)[0] <= d["L"] and options <= d["N"]
+
+    def _inputs(self, d: dict[str, int], encoded: list[dict[str, Any]]):
+        if self.hybrid:
+            return pack_hybrid(encoded, d["S"], d["C"], d["N"], self.pad_id, self.config["rope"])
+        return pack(encoded, d["L"], d["N"], self.pad_id)
+
+    def _calls(self, encoded: list[dict[str, Any]]) -> list[tuple[Any, dict[str, int], list[int]]]:
+        """Fewest estimated milliseconds: one call in the smallest bucket that fits, or greedy chunks."""
+        best = None
+        for name, d, m in self.buckets:
+            groups, group = [], []
+            for j in range(len(encoded)):
+                if not self._fits(d, [encoded[i] for i in group + [j]]):
+                    if not group or not self._fits(d, [encoded[j]]):
+                        break  # this question alone does not fit this bucket
+                    groups.append(group)
+                    group = []
+                group.append(j)
+            else:
+                groups.append(group)
+                cost = len(groups) * self.cost_ms[name]
+                if best is None or cost < best[0]:
+                    best = (cost, [(m, d, g) for g in groups])
+        if best is None:
+            raise ValueError(f"a question exceeds the largest function {self.buckets[-1][0]}")
+        return best[1]
+
+    def logits(self, encoded: list[dict[str, Any]]) -> list[list[float]]:
+        per: list[list[float]] = [[] for _ in encoded]
+        for m, d, group in self._calls(encoded):
+            x, owner = self._inputs(d, [encoded[i] for i in group])
+            out = m.predict(x)["logits"]
+            for k, o in enumerate(owner):
+                per[group[o]].append(float(out[k]))
+        return per
+    def system_one(self, state: Any, questions: dict[str, dict[str, Any]]) -> dict[str, Any]:
+        rows, encoded, answers = {}, {}, {}
+        for qid, q in questions.items():
+            try:
+                rows[qid] = question_to_row(state, q)
+                encoded[qid] = encode(rows[qid], self.tokenizer)
+            except ValueError:
+                answers[qid] = {"type": q.get("type") if isinstance(q, dict) else None, "error": "invalid_question"}
+        qids = list(encoded)
+        if qids:
+            for qid, values in zip(qids, self.logits([encoded[q] for q in qids])):
+                row = rows[qid]
+                offsets = self.score_bias.get(len(values)) if row["task_type"] == "score" else None
+                if offsets:
+                    values = [v + b for v, b in zip(values, offsets)]
+                answers[qid] = answer(row["task_type"], encoded[qid]["keys"], values,
+                                      [o["description"] for o in row["options"]])
+        return {"model": self.config["model_name"], "answers": {q: answers[q] for q in questions},
+                "usage": {"input_tokens": sum(len(e["ids"]) for e in encoded.values()), "output_tokens": 0}}

decision_config.json

@@ -0,0 +1,58 @@
+{
+  "architecture": "qwen3.5-text-endpoints-global-query-shared-bilinear-mlp",
+  "prompt_version": "decision2-segmented-options-global-query-v1",
+  "base_revision": "2fc06364715b967f1860aea9cf38778875588b17",
+  "source_prompt_version": "structured-segmented-candidate-endpoints-global-query-v2",
+  "head_dim": 256,
+  "max_options": 255,
+  "text_parameter_count": 752393024,
+  "parameter_dtype": "float32",
+  "autocast_dtype": "bfloat16",
+  "head_compute_dtype": "float32",
+  "attention": "sdpa",
+  "checkpoint_format": "full",
+  "full_training_source": {
+    "kind": "decision1",
+    "revision": null,
+    "source_fingerprint": {
+      "source_name": "363c4a5e56afc115b1c78c837633956d0bbb63ab",
+      "files_sha256": {
+        "backbone/config.json": "5cece452e606f8ffca1790b79018fed5b0fde7d1a0791d375ca09492e0d1809b",
+        "backbone/model.safetensors": "613f491da2794c5d22f68e81bd2d41000c5675bc4538180148e6e453e8198abd",
+        "config.json": "37c16ffad2439cc06c7a40de037e59341cb57ae4e948cbaf9b1950c5d38d0d37",
+        "decision_config.json": "71fde3c2ea6f4bf4e781028d7f8f39da73599f9d868979efbfbd8b7f43dd68c3",
+        "decision_head.safetensors": "75f4ee8b512aea8f74fdcb42849b4cd37cfb37ff17c5d0b5fda1bb7b4f6d72f6",
+        "tokenizer.json": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523",
+        "tokenizer_config.json": "bee8eba30f0eb4af73c0fe2cd06d0f89b657d7819941c438157ec42f7c80ea87"
+      }
+    }
+  },
+  "training_mode": "full",
+  "loss_version": "decision2-valid-k-ce-plus-optional-brier-and-replay-kl-v1",
+  "initialization": "linear-interpolation-of-full-checkpoints",
+  "interpolation": {
+    "alpha": 0.75,
+    "formula": "W = (1 - alpha) * release + alpha * arm, per tensor in FP32 (backbone and decision head)",
+    "release": {
+      "path": "/runs/m14/inputs/refs/DEV2.0-0.8B/bede7938a8c209c09f27400b79eed57948d6b75e",
+      "model_sha256": "3f02f0e5fc68a1fd985ce48bf3a3bb2d4ed16d1a1705c2258317b78a4d6514f4"
+    },
+    "arm": {
+      "path": "/runs/m16/inputs/arms/08b-RA",
+      "model_sha256": "f25beedd6b9b80fd8ce92857bfd6a14d77d73624e4fd410cc535357c0a1b41fb",
+      "soup": {
+        "method": "uniform weight average (FP32)",
+        "members": [
+          {
+            "path": "/runs/m12/arms/full/m12-08b-RA-s1/checkpoint-0004115",
+            "model_sha256": "576022450d7e284c33f46f680d1ddaeb419b093353b21f0753efe6f73b120f47"
+          },
+          {
+            "path": "/runs/m12/arms/full/m12-08b-RA-s2/checkpoint-0004109",
+            "model_sha256": "6f254ab5458ef0038ee0528f84d7aaea57335004912d9d23f8a432673311898f"
+          }
+        ]
+      }
+    }
+  }
+}

example.py

@@ -0,0 +1,17 @@
+import json
+from decision2_coreml import Decision2CoreML
+
+model = Decision2CoreML(".")
+result = model.system_one(
+    state="The order arrived damaged yesterday. The customer has a receipt and asks for a replacement today.",
+    questions={
+        "route": {"type": "choice", "instructions": "Which team should handle this request?",
+                  "criteria": {"returns": "Refunds, replacements and damaged deliveries",
+                               "billing": "Payments, invoices and charges",
+                               "technical": "Product setup and faults"}},
+        "receipt": {"type": "noul", "instructions": "Does the customer have a receipt?"},
+        "urgency": {"type": "score", "instructions": "How urgent is this request?",
+                    "criteria": ["Routine", "Soon", "Today"]},
+    },
+)
+print(json.dumps(result["answers"], indent=2))

reports/typed-decisions-test.txt

@@ -0,0 +1,6 @@
+flip margins (upstream top-2) [0.0001, 0.0005, 0.0007, 0.0026]
+  choice n=600 acc upstream 0.430 coreml 0.433
+  noul   n=600 acc upstream 0.598 coreml 0.598
+  score  n=800 acc upstream 0.379 coreml 0.376
+typed: token mismatches 0, 2000 decisions, flips 4, max|dp| 0.0077, request p50 110.1 ms
+64 questions: 13 calls, p50 874 ms end to end

tokenizer.json

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

tokenizer_config.json

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