Shared Kai/Eos runtime, release parity report, link to Eos

README.md

@@ -52,9 +52,9 @@ Against the upstream runtime (fp32, Transformers 5.18) on the **typed-decisions
 
 | | Upstream | Core ML fp16 |
 |---|---:|---:|
-| Choice accuracy (600) | 0.423 | 0.420 |
+| Choice accuracy (600) | 0.423 | 0.422 |
 | Yes/No accuracy (600) | 0.628 | 0.628 |
-| Score accuracy (800) | 0.398 | 0.396 |
+| Score accuracy (800) | 0.398 | 0.398 |
 | Answers that differ from upstream | — | 5 / 2,000 (all near-ties: upstream top-2 margin ≤ 0.009) |
 | Max probability difference | — | 0.0082 |
 
@@ -67,7 +67,7 @@ comparison with upstream: they are not the model card's benchmark protocol.
 |---|---:|---:|---:|
 | Card quickstart, 3 questions | 263 ms | 89 ms | **28 ms** |
 | typed-decisions request, 5 questions (median) | — | 418 ms | **61 ms** |
-| One email, 64 questions | 9,514 ms | 3,650 ms | **437 ms** (7 calls) |
+| One email, 64 questions | 9,514 ms | 3,650 ms | **442 ms** (7 calls) |
 
 Times are full requests (tokenize, pack, predict, answer) and exclude model loading. Per call: L256 14 ms, L512 28 ms,
 L1024 58 ms, L2048 142 ms. The upstream runtime on a Mac has none of its GPU fast paths (HIP graphs, fused kernels),
@@ -81,6 +81,8 @@ See `coreml_config.json`. `input_ids`, `position_ids` [1, L] int32; `mask` [1, 1
 
 ## Conversion
 
+Decision-2.0-Eos-0.8B (Qwen3.5 hybrid) is at [FluidInference/decision-2.0-eos-coreml](https://huggingface.co/FluidInference/decision-2.0-eos-coreml), with the same runtime.
+
 `conversion/`: `kai_graph.py` (traceable Qwen3 backbone + decision head, loads the source safetensors),
 `convert.py L N` (torch 2.7.0, coremltools 9.0), `combine.py` (multifunction package), `pack.py` + `typed_eval.py` +
 `typed_ref.py` (parity against the upstream runtime; `pack.py` imports the source repo's vendored encoder from `kai/`).

conversion/check_release.py

@@ -1,15 +1,17 @@
 import json, sys, time
 import numpy as np, pandas as pd
-sys.path.insert(0, "release")
+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 = sys.argv[1], sys.argv[2]
+ref_root, refpath, rel = sys.argv[1], sys.argv[2], sys.argv[3]
 tok = AutoTokenizer.from_pretrained(ref_root)
-m = Decision2CoreML("release")
+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)
@@ -20,10 +22,16 @@ for r in d.itertuples():
     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 = []

decision2_coreml.py

@@ -1,5 +1,7 @@
 """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).
@@ -117,6 +119,65 @@ def pack(encoded: list[dict[str, Any]], L: int, N: int, pad_id: int) -> tuple[di
     }, 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]
@@ -134,6 +195,10 @@ def answer(kind: str, keys: list[str], logits: list[float], descriptions: list[A
     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):
@@ -141,50 +206,56 @@ class Decision2CoreML:
         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"]
-        bias = json.loads((root / "score_bias.json").read_text())["offsets"]
-        self.score_bias = {int(k): v for k, v in bias.items()}
+        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 = []  # (L, N, model), smallest first
-        for name in self.config["functions"]:
-            L, N = (int(x[1:]) for x in name.split("_"))
-            self.buckets.append((L, N, ct.models.MLModel(path, function_name=name, compute_units=compute_units)))
+        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 _calls(self, encoded: list[dict[str, Any]]) -> list[tuple[int, int, Any, list[int]]]:
+    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 L, N, m in self.buckets:
+        for name, d, m in self.buckets:
             groups, group = [], []
             for j in range(len(encoded)):
-                T, C = packed_size([encoded[i] for i in group + [j]])
-                if T > L or C > N:
-                    if not group:
+                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 = []
-                    T, C = packed_size([encoded[j]])
-                    if T > L or C > N:
-                        break
                 group.append(j)
             else:
                 groups.append(group)
-                cost = len(groups) * self.cost_ms[f"L{L}_N{N}"]
+                cost = len(groups) * self.cost_ms[name]
                 if best is None or cost < best[0]:
-                    best = (cost, [(L, N, m, g) for g in groups])
+                    best = (cost, [(m, d, g) for g in groups])
         if best is None:
-            L, N, _ = self.buckets[-1]
-            raise ValueError(f"a question exceeds the largest bucket L{L}_N{N}")
+            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 L, N, m, group in self._calls(encoded):
-            x, owner = pack([encoded[i] for i in group], L, N, self.pad_id)
+        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():

reports/typed-decisions-test.txt

@@ -1,10 +1,6 @@
-  flip agent_trace_observability_000049/outcome choice upstream success coreml partial upstream top-2 margin 0.0004
-  flip agent_trace_observability_000087/urgency score upstream 0 coreml 1 upstream top-2 margin 0.0007
-  flip customer_service_000075/action choice upstream escalate_to_human coreml request_information upstream top-2 margin 0.0003
-  flip invoice_processing_000020/duplicate noul upstream true coreml false upstream top-2 margin 0.0008
-  flip security_incidents_000023/true_positive noul upstream false coreml true upstream top-2 margin 0.0088
-ALL: 2000 decisions, flips 5, max|dp| 0.0082, request p50 61.2 ms p95 63.8 ms, calls 407
-  bucket use {1024: 400}
-  choice n=600 acc upstream 0.423 coreml 0.420
+flip margins (upstream top-2) [0.0011, 0.0004, 0.0003, 0.0008, 0.0088]
+  choice n=600 acc upstream 0.423 coreml 0.422
   noul   n=600 acc upstream 0.628 coreml 0.628
-  score  n=800 acc upstream 0.398 coreml 0.396
+  score  n=800 acc upstream 0.398 coreml 0.398
+typed: token mismatches 0, 2000 decisions, flips 5, max|dp| 0.0082, request p50 60.9 ms
+64 questions: 7 calls, p50 442 ms end to end