d3 Core ML: lite (vllm-sr Decision 3.0 v3.1.0)

.gitattributes

@@ -33,3 +33,7 @@ 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
+lite/embeddings.f16 filter=lfs diff=lfs merge=lfs -text
+lite/pos_embed.f32 filter=lfs diff=lfs merge=lfs -text
+lite/readout.f32 filter=lfs diff=lfs merge=lfs -text
+lite/tokenizer.json filter=lfs diff=lfs merge=lfs -text

lite/README.md

@@ -0,0 +1,50 @@
+---
+license: apache-2.0
+base_model: vllm-sr/d3-lite
+library_name: coreml
+pipeline_tag: zero-shot-classification
+tags:
+- coreml
+- apple-silicon
+- decision-model
+- vision
+- video
+---
+
+# d3-lite — Core ML
+
+[vllm-sr/d3-lite](https://huggingface.co/vllm-sr/d3-lite) (0.85B, part of [Decision 3.0](https://huggingface.co/collections/vllm-sr/decision-30) by the vLLM Semantic Router team) converted to Core ML for Apple silicon: text, image and video requests on device. Same contract as upstream: one forward pass per question, the 255-way FP32 readout at the last prompt token, a probability for every option, no text generated.
+
+## Files
+
+| File | What |
+|---|---|
+| `d3_E256` / `d3_E640` / `d3_E1536.mlpackage` | text graph (fp16 compute, int8 weights) per prompt-length bucket; input = token embeddings with image/video tokens spliced in |
+| `vision_P320` / `vision_P1536` / `vision_P3072.mlpackage` | vision tower (fp32) per patch bucket; video frames go through it one frame pair at a time (block mask) |
+| `embeddings.f16`, `pos_embed.f32`, `readout.f32`, `tokenizer.json`, `config.json`, `coreml_config.json` | host-side tables and config |
+| `python/` | Python host (coremltools): `d3_coreml_text.py`, `d3_coreml_image.py`, `d3_coreml_video.py` |
+
+Swift: `D3VisionManager` in [FluidUse](https://github.com/FluidInference/FluidUse) loads this folder.
+
+## Checked against upstream
+
+Same weights as upstream v3.1.0. Answers that differ from the upstream PyTorch runtime (bf16, MPS), same requests:
+
+| | flips |
+|---|---|
+| Text (typed-decisions TEST, 20 requests) | 2 / 100 |
+| Images (receipt checks) | 1 / 24 |
+| Video (4-second street clips, 2 fps) | 0 / 28 |
+
+Speed on an M5 Pro (24 GB), `CPU_AND_GPU`: ~46 ms per text question, 0.32 s per 4-second video look. 
+
+## Python
+
+```python
+from d3_coreml_text import D3CoreMLText      # d3_coreml_image.D3CoreMLImage, d3_coreml_video.D3CoreMLVideo
+m = D3CoreMLText("<vllm-sr/d3-lite snapshot>", "<this repo>")
+m.system_one(state="...", questions={"route": {"type": "choice", "criteria": {...}}})
+```
+The Python host builds prompts with upstream's own runtime (`d3_runtime.py` in the upstream snapshot, run without weights), so prompts match upstream exactly.
+
+License: Apache-2.0, as upstream.

lite/config.json

@@ -0,0 +1,116 @@
+{
+  "architectures": [
+    "Qwen3_5Model"
+  ],
+  "dtype": "bfloat16",
+  "image_token_id": 248056,
+  "model_type": "qwen3_5",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_output_gate": true,
+    "bos_token_id": null,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": 3584,
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 16,
+    "linear_value_head_dim": 128,
+    "mamba_ssm_dtype": "float32",
+    "max_position_embeddings": 262144,
+    "mlp_only_layers": [],
+    "model_type": "qwen3_5_text",
+    "mtp_num_hidden_layers": 1,
+    "mtp_use_dedicated_embeddings": false,
+    "num_attention_heads": 8,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 2,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        11,
+        11,
+        10
+      ],
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 10000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": true,
+    "use_cache": true,
+    "vocab_size": 248320
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.17.0",
+  "video_token_id": 248057,
+  "vision_config": {
+    "deepstack_visual_indexes": [],
+    "depth": 12,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 768,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "model_type": "qwen3_5_vision",
+    "num_heads": 12,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 1024,
+    "patch_size": 16,
+    "rope_parameters": {
+      "rope_theta": 10000.0,
+      "rope_type": "axial"
+    },
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 248054,
+  "vision_start_token_id": 248053,
+  "auto_map": {
+    "AutoModel": "modeling_d3.D3Model"
+  },
+  "custom_pipelines": {
+    "decision": {
+      "impl": "pipeline_d3.D3Pipeline",
+      "pt": [
+        "AutoModel"
+      ],
+      "type": "text"
+    }
+  }
+}

lite/coreml_config.json

@@ -0,0 +1,312 @@
+{
+ "model_name": "lite",
+ "source": "vllm-sr/lite",
+ "codes": [
+  "A",
+  "B",
+  "C",
+  "D",
+  "E",
+  "F",
+  "G",
+  "H",
+  "I",
+  "J",
+  "K",
+  "L",
+  "M",
+  "N",
+  "O",
+  "P",
+  "Q",
+  "R",
+  "S",
+  "T",
+  "U",
+  "V",
+  "W",
+  "X",
+  "Y",
+  "Z",
+  "AA",
+  "AB",
+  "AC",
+  "AD",
+  "AE",
+  "AF",
+  "AG",
+  "AH",
+  "AI",
+  "AJ",
+  "AK",
+  "AL",
+  "AM",
+  "AN",
+  "AO",
+  "AP",
+  "AQ",
+  "AR",
+  "AS",
+  "AT",
+  "AU",
+  "AV",
+  "AW",
+  "AX",
+  "AY",
+  "AZ",
+  "BA",
+  "BB",
+  "BC",
+  "BD",
+  "BE",
+  "BF",
+  "BG",
+  "BH",
+  "BI",
+  "BJ",
+  "BK",
+  "BL",
+  "BM",
+  "BN",
+  "BO",
+  "BP",
+  "BR",
+  "BS",
+  "BT",
+  "BU",
+  "BV",
+  "BW",
+  "BX",
+  "BY",
+  "CA",
+  "CB",
+  "CC",
+  "CD",
+  "CE",
+  "CF",
+  "CG",
+  "CH",
+  "CI",
+  "CK",
+  "CL",
+  "CM",
+  "CN",
+  "CO",
+  "CP",
+  "CR",
+  "CS",
+  "CT",
+  "CU",
+  "CV",
+  "CW",
+  "CX",
+  "CY",
+  "DA",
+  "DB",
+  "DC",
+  "DD",
+  "DE",
+  "DF",
+  "DG",
+  "DH",
+  "DI",
+  "DJ",
+  "DK",
+  "DL",
+  "DM",
+  "DN",
+  "DO",
+  "DP",
+  "DR",
+  "DS",
+  "DT",
+  "DU",
+  "DV",
+  "DW",
+  "DX",
+  "DY",
+  "EA",
+  "EB",
+  "EC",
+  "ED",
+  "EE",
+  "EF",
+  "EG",
+  "EH",
+  "EI",
+  "EK",
+  "EL",
+  "EM",
+  "EN",
+  "EO",
+  "EP",
+  "EQ",
+  "ER",
+  "ES",
+  "ET",
+  "EU",
+  "EV",
+  "EW",
+  "EX",
+  "EZ",
+  "FA",
+  "FB",
+  "FC",
+  "FD",
+  "FE",
+  "FF",
+  "FG",
+  "FH",
+  "FI",
+  "FK",
+  "FL",
+  "FM",
+  "FN",
+  "FO",
+  "FP",
+  "FR",
+  "FS",
+  "FT",
+  "FU",
+  "FW",
+  "FX",
+  "FY",
+  "GA",
+  "GB",
+  "GC",
+  "GD",
+  "GE",
+  "GF",
+  "GG",
+  "GH",
+  "GI",
+  "GL",
+  "GM",
+  "GN",
+  "GO",
+  "GP",
+  "GR",
+  "GS",
+  "GT",
+  "GU",
+  "GV",
+  "GW",
+  "GX",
+  "GY",
+  "HA",
+  "HB",
+  "HC",
+  "HD",
+  "HE",
+  "HF",
+  "HG",
+  "HH",
+  "HI",
+  "HK",
+  "HL",
+  "HM",
+  "HN",
+  "HO",
+  "HP",
+  "HQ",
+  "HR",
+  "HS",
+  "HT",
+  "HU",
+  "HV",
+  "HW",
+  "HX",
+  "HY",
+  "HZ",
+  "IA",
+  "IB",
+  "IC",
+  "ID",
+  "IE",
+  "IF",
+  "IG",
+  "IH",
+  "II",
+  "IJ",
+  "IK",
+  "IL",
+  "IM",
+  "IN",
+  "IO",
+  "IP",
+  "IQ",
+  "IR",
+  "IS",
+  "IT",
+  "IU",
+  "IV",
+  "IW",
+  "IX",
+  "IZ",
+  "JA",
+  "JB",
+  "JC",
+  "JD",
+  "JE",
+  "JI",
+  "JJ",
+  "JK",
+  "JM",
+  "JO",
+  "JP",
+  "JR",
+  "JS",
+  "JT"
+ ],
+ "temperature": 1.0,
+ "pad_token_id": 248044,
+ "image_token_id": 248056,
+ "vision_start_token_id": 248053,
+ "vision_end_token_id": 248054,
+ "hidden_size": 1024,
+ "vocab_size": 248320,
+ "rope": {
+  "theta": 10000000,
+  "rotary_dim": 64,
+  "mrope_section": [
+   11,
+   11,
+   10
+  ]
+ },
+ "vision": {
+  "patch_size": 16,
+  "merge_size": 2,
+  "temporal_patch_size": 2,
+  "hidden_size": 768,
+  "num_heads": 12,
+  "out_hidden_size": 1024,
+  "num_position_embeddings": 2304,
+  "rope_theta": 10000.0
+ },
+ "image": {
+  "min_pixels": 65536,
+  "max_pixels": 1638400,
+  "mean": [
+   0.5,
+   0.5,
+   0.5
+  ],
+  "std": [
+   0.5,
+   0.5,
+   0.5
+  ],
+  "rescale": 0.00392156862745098
+ },
+ "vision_buckets": [
+  320,
+  1536,
+  3072
+ ],
+ "text_buckets": [
+  256,
+  640,
+  1536
+ ]
+}

lite/d3_E1536.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

lite/d3_E1536.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

lite/d3_E1536.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "0F786AA6-669F-4BB1-9F0F-9BE9E09F20EF": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "78F6A973-8C96-4F9B-B548-288DFCBD0406": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "0F786AA6-669F-4BB1-9F0F-9BE9E09F20EF"
+}

lite/d3_E256.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

lite/d3_E256.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

lite/d3_E256.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "5A657C48-D6B0-4EC9-A96E-CE6DDEFEC813": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "AFC1813F-EB9F-47B9-A90B-AAF702E6BD57": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "5A657C48-D6B0-4EC9-A96E-CE6DDEFEC813"
+}

lite/d3_E640.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

lite/d3_E640.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

lite/d3_E640.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "176FF80A-94E3-4B6C-82DB-0642EF7BE052": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "F528A27F-912C-4218-ACEE-AE0622862F5D": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "176FF80A-94E3-4B6C-82DB-0642EF7BE052"
+}

lite/embeddings.f16

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

lite/pos_embed.f32

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

lite/python/d3_coreml_image.py

@@ -0,0 +1,128 @@
+"""d3 image requests on Core ML (Python host): upstream prompt, processor and answer code; Core ML vision tower
+(fp32) + text graph (fp16, host-spliced embeddings, host M-RoPE); FP32 readout on the host.
+
+    m = D3CoreMLImage("<d3 checkpoint dir>", "<coreml dir: Vision_P*, d3_E*.mlpackage, embeddings.f16>")
+    m.system_one(state=..., images=[...], questions={...})   # same answers dict as upstream
+
+Positions come from transformers' Qwen3_5Model.compute_3d_position_ids on a weightless (meta) model."""
+import json
+import sys
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+import torch
+
+HERE = Path(__file__).parent
+
+
+class D3CoreMLImage:
+    def __init__(self, root, coreml_dir, compute_units=ct.ComputeUnit.CPU_AND_GPU):
+        root, cdir = Path(root), Path(coreml_dir)
+        sys.path.insert(0, str(root))
+        sys.path.insert(0, str(HERE.parent / "vision"))
+        sys.path.insert(0, str(HERE.parent))
+        import d3_runtime as up
+        from qwen35_export import rope_cos_sin
+        from transformers import AutoConfig
+        from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Model
+        from vision_export import VisionConfig
+        from d3_graph import load_text_config
+
+        self.up, self.rope_cos_sin = up, rope_cos_sin
+        config = json.loads((root / "decision_config.json").read_text())
+        self.codes, self.temperature = config["codes"], float(config.get("temperature", 1.0))
+        self.processor = up.load_processor(root)
+        self.vcfg = VisionConfig(json.loads((root / "config.json").read_text())["vision_config"])
+        self.tcfg = load_text_config(str(root))
+        with torch.device("meta"):
+            self.positions = Qwen3_5Model(AutoConfig.from_pretrained(str(root)))
+        from safetensors import safe_open
+        self.pos_table = None
+        for path in sorted(root.glob("model*.safetensors")):
+            with safe_open(str(path), "pt") as f:
+                if "visual.pos_embed.weight" in f.keys():
+                    self.pos_table = f.get_tensor("visual.pos_embed.weight").float()
+        from safetensors.torch import load_file
+        self.readout = load_file(str(root / "readout.safetensors"))["weight"].float().numpy()
+        H = self.tcfg.hidden_size
+        self.embed = np.memmap(cdir / "embeddings.f16", dtype=np.float16, mode="r").reshape(-1, H)
+        self.image_token = AutoConfig.from_pretrained(str(root)).image_token_id
+        self.vision = {int(p.name[8:]): ct.models.MLModel(str(p / "VisionTower_fp32.mlpackage"), compute_units=compute_units)
+                       for p in sorted(cdir.glob("Vision_P*"))}
+        multi = cdir / "d3_text.mlpackage"  # one package, functions E{L} sharing the weights
+        if multi.exists():
+            names = json.loads((cdir / "d3_text.functions.json").read_text())
+            self.text = {int(n[1:]): ct.models.MLModel(str(multi), compute_units=compute_units, function_name=n)
+                         for n in names}
+        else:
+            self.text = {int(p.stem[4:]): ct.models.MLModel(str(p), compute_units=compute_units)
+                         for p in sorted(cdir.glob("d3_E*.mlpackage"))}
+        self.model_name = root.name + "-coreml"
+
+    def warm(self):
+        """Run every vision and text bucket once: the first call per bucket pays the GPU setup (seconds)."""
+        for N, model in self.vision.items():
+            model.predict({"patches": np.zeros((N, self.vcfg.patch_dim), np.float32),
+                           "pos_embeds": np.zeros((N, self.vcfg.hidden), np.float32),
+                           "cos": np.zeros((N, self.vcfg.head_dim), np.float32),
+                           "sin": np.zeros((N, self.vcfg.head_dim), np.float32),
+                           "mask": np.zeros((1, 1, N, N), np.float32)})
+        H, R = self.tcfg.hidden_size, self.tcfg.rotary_dim
+        for L, model in self.text.items():
+            model.predict({"hidden": np.zeros((1, L, H), np.float16), "cos": np.ones((L, R), np.float16),
+                           "sin": np.zeros((L, R), np.float16), "key_valid": np.ones(L, np.float16),
+                           "last_onehot": np.eye(1, L, L - 1, dtype=np.float16)})
+
+    def _tokens(self, pv, grid):
+        from vision_export import prepare_inputs
+        N = next(b for b in sorted(self.vision) if b >= pv.shape[0])
+        patches, pe, cs, sn, mask, n = prepare_inputs(self.vcfg, self.pos_table, pv, grid, N)
+        out = self.vision[N].predict({"patches": patches.numpy(), "pos_embeds": pe.numpy(), "cos": cs.numpy(),
+                                      "sin": sn.numpy(), "mask": mask.numpy()})["tokens"]
+        return out[: n // 4]
+
+    def _probs(self, text, images, count, image_tokens):
+        enc = self.processor(text=[text], images=images, return_tensors="pt")
+        ids = enc["input_ids"][0].numpy()
+        hidden = self.embed[ids].astype(np.float32)
+        rows = np.where(ids == self.image_token)[0]
+        hidden[rows] = np.concatenate(image_tokens)
+        with torch.no_grad():
+            pos = self.positions.compute_3d_position_ids(
+                input_ids=enc["input_ids"], inputs_embeds=None, image_grid_thw=enc["image_grid_thw"],
+                attention_mask=enc["attention_mask"], mm_token_type_ids=enc["mm_token_type_ids"])[:, 0]
+        n = len(ids)
+        L = next((b for b in sorted(self.text) if b >= n), None)
+        if L is None:
+            raise ValueError(f"{n} tokens exceed the largest bucket {max(self.text)}")
+        cos, sin = self.rope_cos_sin(self.tcfg, torch.cat([pos, pos[:, -1:].expand(3, L - n)], dim=1))
+        x = np.zeros((1, L, hidden.shape[1]), dtype=np.float16); x[0, :n] = hidden
+        valid = np.zeros(L, dtype=np.float16); valid[:n] = 1
+        last = np.zeros((1, L), dtype=np.float16); last[0, n - 1] = 1
+        h = self.text[L].predict({"hidden": x, "cos": cos.numpy().astype(np.float16), "sin": sin.numpy().astype(np.float16),
+                                  "key_valid": valid, "last_onehot": last})["final"].reshape(-1)
+        logits = self.readout[:count].astype(np.float64) @ h.astype(np.float64) / self.temperature
+        p = np.exp(logits - logits.max())
+        return (p / p.sum()).tolist(), n
+
+    def system_one(self, *, state, questions, images):
+        up = self.up
+        decoded = [up.load_image(v) for v in images]
+        enc = self.processor.image_processor(images=decoded, return_tensors="pt")
+        # one vision call per image; every question reuses the image tokens
+        tokens, start = [], 0
+        for k, grid in enumerate(enc["image_grid_thw"]):
+            n = int(grid.prod())
+            tokens.append(self._tokens(enc["pixel_values"][start:start + n], grid[None]))
+            start += n
+        answers, used = {}, 0
+        for key, q in questions.items():
+            question = up.normalize_question(q)
+            text = self.processor.apply_chat_template(
+                up.image_messages("d3", state, question.rendered, self.codes, len(decoded)),
+                tokenize=False, add_generation_prompt=True, enable_thinking=False)
+            probs, n = self._probs(text, decoded, len(question.keys), tokens)
+            used += n
+            answers[key] = up.product_answer(question, probs)
+        return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}}

lite/python/d3_coreml_text.py

@@ -0,0 +1,80 @@
+"""d3 text requests on Core ML (Python host), for the checkpoints whose text graph takes embeddings (d3-mini).
+
+Upstream's own request preparation runs on a weightless (meta-device) d3, so prompts and token ids are exactly
+upstream's; Core ML runs the text graph per question (host embedding gather, text positions on all three M-RoPE axes);
+FP32 readout on the host.
+
+    m = D3CoreMLText("<d3 checkpoint>", "<coreml dir>", text_buckets=[256, 640])
+    m.system_one(state=..., questions={...})"""
+import sys
+import time
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+import torch
+
+HERE = Path(__file__).parent
+
+
+class D3CoreMLText:
+    def __init__(self, root, coreml_dir, text_buckets=None, compute_units=ct.ComputeUnit.CPU_AND_GPU):
+        root, cdir = Path(root), Path(coreml_dir)
+        sys.path.insert(0, str(root))
+        sys.path.insert(0, str(HERE.parent))
+        import d3_runtime as up
+        from d3_graph import load_text_config
+        from qwen35_export import rope_cos_sin
+        from safetensors.torch import load_file
+
+        self.up, self.rope_cos_sin = up, rope_cos_sin
+        self.prep = up.D3.from_pretrained(str(root), device="meta", verify="none")  # no weights: prompts only
+        self.temperature = self.prep.temperature
+        self.tcfg = load_text_config(str(root))
+        self.readout = load_file(str(root / "readout.safetensors"))["weight"].float().numpy()
+        self.embed = np.memmap(cdir / "embeddings.f16", dtype=np.float16, mode="r").reshape(-1, self.tcfg.hidden_size)
+        self.text = {int(p.stem[4:]): ct.models.MLModel(str(p), compute_units=compute_units)
+                     for p in sorted(cdir.glob("d3_E*.mlpackage")) if not text_buckets or int(p.stem[4:]) in text_buckets}
+        self.pad = self.prep.pad_id
+        self.model_name = root.name + "-coreml"
+        self.timing = []
+
+    def warm(self):
+        for L, model in self.text.items():
+            for _ in range(2):
+                self._final(list(range(10)), L)
+
+    def _final(self, ids, L):
+        n = len(ids)
+        x = np.zeros((1, L, self.tcfg.hidden_size), dtype=np.float16)
+        x[0, :n] = self.embed[np.asarray(ids)]
+        x[0, n:] = self.embed[self.pad]
+        pos = torch.arange(L).clamp(max=n - 1)[None].expand(3, -1)
+        cos, sin = self.rope_cos_sin(self.tcfg, pos)
+        valid = np.zeros(L, dtype=np.float16); valid[:n] = 1
+        last = np.zeros((1, L), dtype=np.float16); last[0, n - 1] = 1
+        return self.text[L].predict({"hidden": x, "cos": cos.numpy().astype(np.float16),
+                                     "sin": sin.numpy().astype(np.float16), "key_valid": valid,
+                                     "last_onehot": last})["final"].reshape(-1)
+
+    def system_one(self, *, state, questions):
+        prepared = self.prep.prepare(state, questions)
+        answers, used, self.timing = {}, 0, []
+        for key in prepared.keys:
+            if key in prepared.errors:
+                answers[key] = prepared.errors[key]
+                continue
+            ids, question = prepared.sequences[key], prepared.questions[key]
+            L = next((b for b in sorted(self.text) if b >= len(ids)), None)
+            if L is None:
+                answers[key] = {"type": question.kind, "error": "max_length_exceeded",
+                                "message": f"{len(ids)} tokens exceed the largest Core ML bucket {max(self.text)}"}
+                continue
+            t = time.perf_counter()
+            h = self._final(ids, L)
+            self.timing.append((len(ids), L, round(time.perf_counter() - t, 3)))
+            logits = self.readout[:len(question.keys)].astype(np.float64) @ h.astype(np.float64) / self.temperature
+            p = np.exp(logits - logits.max())
+            answers[key] = self.up.product_answer(question, (p / p.sum()).tolist())
+            used += len(ids)
+        return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}}

lite/python/d3_coreml_video.py

@@ -0,0 +1,122 @@
+"""d3 video requests on Core ML (Python host).
+
+Upstream's own request preparation runs on a weightless (meta-device) d3, so frame sampling, the video processor,
+the timestamped placeholders and the 3-D positions are exactly upstream's. Core ML runs the compute: the fp32 vision
+tower once per video (attention within each frame pair, as upstream: one block per temporal group), then the text
+graph per question with the video tokens spliced in; FP32 readout on the host.
+
+    m = D3CoreMLVideo("<d3 checkpoint>", "<coreml dir>", max_pixels=100_352)
+    m.system_one(state=..., questions={...}, videos=[frames])     # frames: uint8 [n, h, w, 3] at 2 fps"""
+import json
+import sys
+import time
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+import torch
+
+HERE = Path(__file__).parent
+
+
+class D3CoreMLVideo:
+    def __init__(self, root, coreml_dir, max_pixels=None, vision_buckets=None, text_buckets=None,
+                 compute_units=ct.ComputeUnit.CPU_AND_GPU):
+        """`vision_buckets` / `text_buckets`: load only these sizes (each loaded model costs memory)."""
+        root, cdir = Path(root), Path(coreml_dir)
+        sys.path.insert(0, str(root))
+        sys.path.insert(0, str(HERE.parent / "vision"))
+        sys.path.insert(0, str(HERE.parent))
+        import d3_runtime as up
+        if max_pixels:
+            up.VIDEO_MAX_PIXELS = max_pixels
+        from d3_graph import load_text_config
+        from qwen35_export import rope_cos_sin
+        from safetensors import safe_open
+        from safetensors.torch import load_file
+        from vision_export import VisionConfig
+
+        self.up, self.rope_cos_sin = up, rope_cos_sin
+        self.prep = up.D3.from_pretrained(str(root), device="meta", verify="none")  # no weights: prompts only
+        self.processor, self.temperature = self.prep.processor, self.prep.temperature
+        self.backbone = self.prep.backbone
+        self.vcfg = VisionConfig(json.loads((root / "config.json").read_text())["vision_config"])
+        self.tcfg = load_text_config(str(root))
+        for path in sorted(root.glob("model*.safetensors")):
+            with safe_open(str(path), "pt") as f:
+                if "visual.pos_embed.weight" in f.keys():
+                    self.pos_table = f.get_tensor("visual.pos_embed.weight").float()
+        self.readout = load_file(str(root / "readout.safetensors"))["weight"].float().numpy()
+        H = self.tcfg.hidden_size
+        self.embed = np.memmap(cdir / "embeddings.f16", dtype=np.float16, mode="r").reshape(-1, H)
+        self.video_token = self.backbone.config.video_token_id
+        self.vision = {int(p.name[8:]): ct.models.MLModel(str(p / "VisionTower_fp32.mlpackage"), compute_units=compute_units)
+                       for p in sorted(cdir.glob("Vision_P*")) if not vision_buckets or int(p.name[8:]) in vision_buckets}
+        self.text = {int(p.stem[4:]): ct.models.MLModel(str(p), compute_units=compute_units)
+                     for p in sorted(cdir.glob("d3_E*.mlpackage")) if not text_buckets or int(p.stem[4:]) in text_buckets}
+        self.model_name = root.name + "-coreml"
+        self.timing = {}
+
+    def warm(self, frames, state, questions):
+        for _ in range(2):
+            self.system_one(state=state, questions=questions, videos=[frames])
+
+    def _video_tokens(self, pixel_values, grid):
+        """Vision tower for one video: attention within each temporal group (frame pair), padding masked."""
+        from vision_export import prepare_inputs
+        n = pixel_values.shape[0]
+        N = next(b for b in sorted(self.vision) if b >= n)
+        patches, pe, cs, sn, mask, n = prepare_inputs(self.vcfg, self.pos_table, pixel_values, grid, N)
+        t, h, w = (int(x) for x in grid.reshape(-1))
+        block = np.full((N, N), -1e4, dtype=np.float32)
+        for g in range(t):  # rows are t-major: each frame pair is h*w consecutive patches
+            s = g * h * w
+            block[s:s + h * w, s:s + h * w] = 0
+        out = self.vision[N].predict({"patches": patches.numpy(), "pos_embeds": pe.numpy(), "cos": cs.numpy(),
+                                      "sin": sn.numpy(), "mask": block[None, None]})["tokens"]
+        return out[: n // 4]
+
+    def system_one(self, *, state, questions, videos):
+        up, proc = self.up, self.processor
+        t0 = time.perf_counter()
+        prepared = self.prep.prepare(state, questions, videos=videos)
+        media = prepared.media
+        grids = media["video_grid_thw"]
+        tokens, start = [], 0
+        for grid in grids:
+            n = int(grid.prod())
+            tokens.append(self._video_tokens(media["pixel_values_videos"][start:start + n], grid[None]))
+            start += n
+        tokens = np.concatenate(tokens)
+        t1 = time.perf_counter()
+        answers, used = {}, 0
+        for key in prepared.runnable:
+            question = prepared.questions[key]
+            expanded = proc.get_text_with_replacements([prepared.texts[key]], [], prepared.video_texts)[0][0]
+            enc = proc.tokenizer([expanded], return_token_type_ids=False, return_tensors="pt")
+            ids = enc["input_ids"]
+            types = torch.as_tensor(proc.create_mm_token_type_ids(ids))
+            row = ids[0].numpy()
+            hidden = self.embed[row].astype(np.float32)
+            hidden[np.where(row == self.video_token)[0]] = tokens
+            with torch.no_grad():
+                pos = self.backbone.compute_3d_position_ids(
+                    input_ids=ids, inputs_embeds=None, image_grid_thw=None, video_grid_thw=grids,
+                    attention_mask=enc["attention_mask"], past_key_values=None, mm_token_type_ids=types)[:, 0]
+            n = len(row)
+            L = next((b for b in sorted(self.text) if b >= n), None)
+            if L is None:
+                raise ValueError(f"{n} tokens exceed the largest text bucket {max(self.text)}")
+            cos, sin = self.rope_cos_sin(self.tcfg, torch.cat([pos, pos[:, -1:].expand(3, L - n)], dim=1))
+            x = np.zeros((1, L, hidden.shape[1]), dtype=np.float16); x[0, :n] = hidden
+            valid = np.zeros(L, dtype=np.float16); valid[:n] = 1
+            last = np.zeros((1, L), dtype=np.float16); last[0, n - 1] = 1
+            h = self.text[L].predict({"hidden": x, "cos": cos.numpy().astype(np.float16),
+                                      "sin": sin.numpy().astype(np.float16), "key_valid": valid,
+                                      "last_onehot": last})["final"].reshape(-1)
+            logits = self.readout[:len(question.keys)].astype(np.float64) @ h.astype(np.float64) / self.temperature
+            p = np.exp(logits - logits.max())
+            answers[key] = up.product_answer(question, (p / p.sum()).tolist())
+            used += n
+        self.timing = {"prepare+vision": round(t1 - t0, 3), "text": round(time.perf_counter() - t1, 3)}
+        return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}}

lite/python/d3_graph.py

@@ -0,0 +1,120 @@
+"""d3-lite (Qwen3.5-0.8B, code readout v1) as one Core ML row graph.
+
+One question per call, as upstream: the row is right-padded to the bucket length L. The 6 full-attention layers are
+noncausal (decision_config `noncausal_full_attention`): they see every real token and mask only the padding keys. The
+Gated DeltaNet layers are causal scans, so padding after the last real token cannot reach it. Output is the final-normed
+hidden state at the last real token; the 255 x 1024 FP32 readout runs on the host, as upstream.
+"""
+
+import json
+
+import numpy as np
+import torch
+from safetensors.torch import load_file
+from torch import nn
+
+from qwen35_export import DecoderLayer, RMSNorm, TextConfig, rope_cos_sin
+
+MASK_VALUE = -1e4
+
+
+class D3Row(nn.Module):
+    def __init__(self, cfg: TextConfig, seq_len: int, chunk: int = 64):
+        super().__init__()
+        self.cfg, self.L = cfg, seq_len
+        self.embed_tokens = nn.Embedding(1, cfg.hidden_size)  # replaced on load
+        self.layers = nn.ModuleList([DecoderLayer(cfg, i, seq_len, chunk) for i in range(cfg.num_layers)])
+        self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
+
+    def forward(self, input_ids, cos, sin, key_valid, last_onehot):
+        """input_ids [1, L] int32; cos/sin [L, R]; key_valid [L] (1 real, 0 pad); last_onehot [1, L] -> hidden [1, H]."""
+        mask = ((1.0 - key_valid) * MASK_VALUE)[None, None, None, :]
+        x = self.embed_tokens(input_ids)
+        for layer in self.layers:
+            x = layer(x, cos, sin, mask)
+        return self.norm(torch.matmul(last_onehot, x[0]))
+
+
+class D3RowEmbeds(nn.Module):
+    """D3Row for image requests: the host gathers token embeddings and splices the vision tower's tokens in at the
+    image-pad positions, and builds cos/sin from the 3-D M-RoPE positions (text and image tokens differ)."""
+
+    def __init__(self, cfg: TextConfig, seq_len: int, chunk: int = 64):
+        super().__init__()
+        self.cfg, self.L = cfg, seq_len
+        self.layers = nn.ModuleList([DecoderLayer(cfg, i, seq_len, chunk) for i in range(cfg.num_layers)])
+        self.norm = RMSNorm(cfg.hidden_size, cfg.eps)
+
+    def forward(self, hidden, cos, sin, key_valid, last_onehot):
+        """hidden [1, L, H]; cos/sin [L, R]; key_valid [L]; last_onehot [1, L] -> final hidden [1, H]."""
+        mask = ((1.0 - key_valid) * MASK_VALUE)[None, None, None, :]
+        x = hidden
+        for layer in self.layers:
+            x = layer(x, cos, sin, mask)
+        return self.norm(torch.matmul(last_onehot, x[0]))
+
+
+def load_row_embeds(root: str, seq_len: int, chunk: int = 64, dtype=torch.float32) -> tuple[D3RowEmbeds, torch.Tensor]:
+    """(graph, embedding table [V, H]). dtype=float16 halves peak memory for tracing the larger checkpoints (d3-mini:
+    fp32 weights plus their copy do not fit in 24 GB); weights are assigned in place, not copied."""
+    cfg = load_text_config(root)
+    prefix = "language_model."
+    state = {}
+    from safetensors import safe_open
+    files = sorted(__import__("glob").glob(f"{root}/model*.safetensors"))
+    for path in files:
+        with safe_open(path, "pt") as f:
+            for k in f.keys():
+                if k.startswith(prefix):
+                    state[k[len(prefix):]] = f.get_tensor(k).to(dtype)
+    embed = state.pop("embed_tokens.weight")
+    model = D3RowEmbeds(cfg, seq_len, chunk).to(dtype)
+    missing, unexpected = model.load_state_dict(state, strict=False, assign=True)
+    if missing or unexpected:
+        raise RuntimeError(f"weight mismatch: missing={missing[:8]} unexpected={unexpected[:8]}")
+    return model.eval(), embed
+
+
+def load_text_config(root: str) -> TextConfig:
+    return TextConfig(json.load(open(f"{root}/config.json"))["text_config"])
+
+
+def load_row(root: str, seq_len: int, chunk: int = 64) -> D3Row:
+    cfg = load_text_config(root)
+    weights = load_file(f"{root}/model.safetensors")
+    prefix = "language_model."
+    state = {k[len(prefix):]: v.float() for k, v in weights.items() if k.startswith(prefix)}
+    model = D3Row(cfg, seq_len, chunk)
+    model.embed_tokens = nn.Embedding.from_pretrained(state.pop("embed_tokens.weight"), freeze=True)
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    missing = [m for m in missing if not m.startswith("embed_tokens")]
+    if missing or unexpected:
+        raise RuntimeError(f"weight mismatch: missing={missing[:8]} unexpected={unexpected[:8]}")
+    return model.eval()
+
+
+def load_readout(root: str) -> np.ndarray:
+    return load_file(f"{root}/readout.safetensors")["weight"].float().numpy()
+
+
+def row_inputs(cfg: TextConfig, ids, L: int, pad_id: int) -> dict[str, np.ndarray]:
+    n = len(ids)
+    if not 1 <= n <= L:
+        raise ValueError(f"{n} tokens do not fit L{L}")
+    input_ids = np.full((1, L), pad_id, dtype=np.int32)
+    input_ids[0, :n] = ids
+    pos = torch.arange(L)
+    pos[n:] = n - 1
+    cos, sin = rope_cos_sin(cfg, pos[None].expand(3, -1))
+    valid = np.zeros(L, dtype=np.float32)
+    valid[:n] = 1
+    last = np.zeros((1, L), dtype=np.float32)
+    last[0, n - 1] = 1
+    return {"input_ids": input_ids, "cos": cos.numpy(), "sin": sin.numpy(), "key_valid": valid, "last_onehot": last}
+
+
+def code_probabilities(hidden: np.ndarray, readout: np.ndarray, count: int, temperature: float = 1.0) -> np.ndarray:
+    logits = hidden.astype(np.float64).reshape(-1) @ readout[:count].astype(np.float64).T / temperature
+    logits -= logits.max()
+    p = np.exp(logits)
+    return p / p.sum()

lite/python/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]}")

lite/python/vision_export.py

@@ -0,0 +1,230 @@
+"""Trace-friendly Qwen3.5 vision tower for Core ML (from clef-vision-0.8b; d3 weight names): one image per call, patches right-padded to a fixed count.
+
+Inputs (all fp32, host-prepared from the HF processor's ``pixel_values`` / ``image_grid_thw``):
+  patches     [N, 1536]  flattened (C, T, P, P) patches, zero padded after the image's n patches
+  pos_embeds  [N, 768]   bilinear-resampled learned position grid (host: HF interpolation indices/weights)
+  cos, sin    [N, 64]    axial 2D rope tables (host)
+  mask        [1, 1, N, N] additive key mask: 0 for real patches, -1e4 for padding
+Output: merged tokens [N/4, 1024]; rows [0, n/4) are the image's tokens in the order the language model expects.
+
+Padding is appended after the image, so with the key mask real patches never attend to it; the merger groups four
+consecutive patches, and n is always a multiple of 4 (grid h, w are multiples of the merge size), so no group mixes
+real and padded patches. Parity vs ``Qwen3_5VisionModel`` is checked in ``vision_parity``.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+VISION_PREFIX = "visual."
+
+
+class VisionConfig:
+    def __init__(self, cfg: dict):
+        self.depth = cfg["depth"]
+        self.hidden = cfg["hidden_size"]
+        self.heads = cfg["num_heads"]
+        self.head_dim = self.hidden // self.heads
+        self.intermediate = cfg["intermediate_size"]
+        self.out_hidden = cfg["out_hidden_size"]
+        self.patch = cfg["patch_size"]
+        self.temporal = cfg["temporal_patch_size"]
+        self.channels = cfg["in_channels"]
+        self.merge = cfg["spatial_merge_size"]
+        self.num_pos = cfg["num_position_embeddings"]
+        self.act = cfg["hidden_act"]  # gelu_pytorch_tanh for the blocks; the merger uses exact GELU
+        self.rope_theta = cfg.get("rope_parameters", {}).get("rope_theta", 10000.0)
+        self.patch_dim = self.channels * self.temporal * self.patch * self.patch
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: VisionConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.norm1 = nn.LayerNorm(cfg.hidden, eps=1e-6)
+        self.norm2 = nn.LayerNorm(cfg.hidden, eps=1e-6)
+        self.qkv = nn.Linear(cfg.hidden, cfg.hidden * 3)
+        self.proj = nn.Linear(cfg.hidden, cfg.hidden)
+        self.fc1 = nn.Linear(cfg.hidden, cfg.intermediate)
+        self.fc2 = nn.Linear(cfg.intermediate, cfg.hidden)
+        self.scale = cfg.head_dim**-0.5
+
+    def forward(self, x, cos, sin, mask):
+        n = x.shape[0]
+        h = self.norm1(x)
+        qkv = self.qkv(h).reshape(n, 3, self.cfg.heads, self.cfg.head_dim)
+        q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]  # [N, H, d]
+        c, s = cos[:, None, :], sin[:, None, :]
+        q = q * c + rotate_half(q) * s
+        k = k * c + rotate_half(k) * s
+        q = q.permute(1, 0, 2)[None]  # [1, H, N, d]
+        k = k.permute(1, 0, 2)[None]
+        v = v.permute(1, 0, 2)[None]
+        scores = torch.matmul(q, k.transpose(-1, -2)) * self.scale + mask
+        attn = torch.matmul(torch.softmax(scores, dim=-1), v)  # [1, H, N, d]
+        attn = attn[0].permute(1, 0, 2).reshape(n, self.cfg.hidden)
+        x = x + self.proj(attn)
+        h = self.norm2(x)
+        act = F.gelu(self.fc1(h), approximate="tanh" if self.cfg.act == "gelu_pytorch_tanh" else "none")
+        return x + self.fc2(act)
+
+
+class VisionTower(nn.Module):
+    def __init__(self, cfg: VisionConfig, n_patches: int):
+        super().__init__()
+        assert n_patches % (cfg.merge**2) == 0
+        self.cfg = cfg
+        self.n_patches = n_patches
+        self.patch_embed = nn.Linear(cfg.patch_dim, cfg.hidden)  # Conv3d with kernel == stride == patch
+        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.depth)])
+        merged = cfg.hidden * cfg.merge**2
+        self.merger_norm = nn.LayerNorm(cfg.hidden, eps=1e-6)
+        self.merger_fc1 = nn.Linear(merged, merged)
+        self.merger_fc2 = nn.Linear(merged, cfg.out_hidden)
+
+    def forward(self, patches, pos_embeds, cos, sin, mask):
+        x = self.patch_embed(patches) + pos_embeds
+        for block in self.blocks:
+            x = block(x, cos, sin, mask)
+        x = self.merger_norm(x).reshape(self.n_patches // self.cfg.merge**2, -1)
+        return self.merger_fc2(F.gelu(self.merger_fc1(x)))
+
+    def load_hf(self, state: dict[str, torch.Tensor]) -> None:
+        """``state`` holds the HF vision weights with the ``visual.`` prefix stripped."""
+        own = {
+            "patch_embed.weight": state["patch_embed.proj.weight"].reshape(self.cfg.hidden, -1),
+            "patch_embed.bias": state["patch_embed.proj.bias"],
+            "merger_norm.weight": state["merger.norm.weight"],
+            "merger_norm.bias": state["merger.norm.bias"],
+            "merger_fc1.weight": state["merger.linear_fc1.weight"],
+            "merger_fc1.bias": state["merger.linear_fc1.bias"],
+            "merger_fc2.weight": state["merger.linear_fc2.weight"],
+            "merger_fc2.bias": state["merger.linear_fc2.bias"],
+        }
+        for i in range(self.cfg.depth):
+            src, dst = f"blocks.{i}.", f"blocks.{i}."
+            for a, b in (("norm1", "norm1"), ("norm2", "norm2"), ("attn.qkv", "qkv"), ("attn.proj", "proj"),
+                         ("mlp.linear_fc1", "fc1"), ("mlp.linear_fc2", "fc2")):
+                own[f"{dst}{b}.weight"] = state[f"{src}{a}.weight"]
+                own[f"{dst}{b}.bias"] = state[f"{src}{a}.bias"]
+        missing, unexpected = self.load_state_dict({k: v.float() for k, v in own.items()}, strict=False)
+        if missing or unexpected:
+            raise RuntimeError(f"vision weight mismatch: missing={missing[:6]} unexpected={unexpected[:6]}")
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+# Host-side preparation (numpy/torch; mirrors transformers.vision_utils + Qwen3_5VisionRotaryEmbedding)
+
+def axial_rope(cfg: VisionConfig, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """``position_ids`` [n, 2] (h, w) -> cos/sin [n, head_dim]: freqs over head_dim//4 per axis, concat (h, w), x2."""
+    spatial = cfg.head_dim // 2
+    inv_freq = 1.0 / (cfg.rope_theta ** (torch.arange(0, spatial, 2, dtype=torch.float32) / spatial))
+    freqs = position_ids[..., None].float() * inv_freq  # [n, 2, spatial/2]
+    hw = torch.cat([freqs[:, 0], freqs[:, 1]], dim=-1)
+    emb = torch.cat([hw, hw], dim=-1)
+    return emb.cos(), emb.sin()
+
+
+def prepare_inputs(cfg: VisionConfig, pos_table: torch.Tensor, pixel_values: torch.Tensor, grid_thw: torch.Tensor,
+                   n_patches: int):
+    """One image. Uses the HF helpers for position ids and pos-embed interpolation so the host matches the reference."""
+    from transformers.vision_utils import get_vision_interpolation_indices_and_weights, get_vision_position_ids
+
+    n = int(pixel_values.shape[0])
+    if n > n_patches:
+        raise ValueError(f"image has {n} patches; bucket holds {n_patches}")
+    grid = grid_thw.reshape(1, 3)
+    idx, w = get_vision_interpolation_indices_and_weights(
+        grid, num_grid_per_side=int(math.isqrt(cfg.num_pos)), mode="bilinear", align_corners=True,
+        spatial_merge_size=cfg.merge, kwargs={})
+    pos_embeds = (pos_table[idx] * w[:, :, None]).sum(1)  # [n, hidden]
+    position_ids = get_vision_position_ids(grid, cfg.merge, kwargs={})  # [n, 2]
+    cos, sin = axial_rope(cfg, position_ids)
+    patches = torch.zeros(n_patches, cfg.patch_dim)
+    patches[:n] = pixel_values.float()
+    pe = torch.zeros(n_patches, cfg.hidden)
+    pe[:n] = pos_embeds.float()
+    cs = torch.zeros(n_patches, cfg.head_dim)
+    sn = torch.zeros(n_patches, cfg.head_dim)
+    cs[:n], sn[:n] = cos, sin
+    mask = torch.zeros(1, 1, n_patches, n_patches)
+    mask[..., n:] = -1e4
+    return patches, pe, cs, sn, mask, n
+
+
+def load_vision(merged_backbone: Path, n_patches: int):
+    from safetensors.torch import load_file
+
+    config = json.loads((merged_backbone / "config.json").read_text())
+    cfg = VisionConfig(config["vision_config"])
+    state = {}
+    for path in sorted(merged_backbone.glob("model*.safetensors")):  # one file or shards
+        state.update({k[len(VISION_PREFIX):]: v for k, v in load_file(str(path)).items() if k.startswith(VISION_PREFIX)})
+    tower = VisionTower(cfg, n_patches)
+    tower.load_hf(state)
+    return tower.eval(), cfg, state["pos_embed.weight"].float()
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--merged", type=Path, required=True, help="d3 checkpoint dir")
+    ap.add_argument("--patches", type=int, default=784, help="patch bucket (784 = 448x448 cap -> 196 tokens)")
+    ap.add_argument("--precision", choices=("fp16", "fp32", "mixed"), default="fp32",
+                    help="fp32 is the only exact option: fp16 (0.27 max err) and mixed with fp32 norms/softmax "
+                    "(0.33) both break on this ViT's activation range; fp32 runs 35 ms on the M5 Pro GPU")
+    ap.add_argument("--out", type=Path, default=Path("build/coreml"))
+    args = ap.parse_args()
+    import coremltools as ct
+
+    tower, cfg, pos_table = load_vision(args.merged, args.patches)
+    N = args.patches
+    example = (torch.zeros(N, cfg.patch_dim), torch.zeros(N, cfg.hidden), torch.zeros(N, cfg.head_dim),
+               torch.zeros(N, cfg.head_dim), torch.zeros(1, 1, N, N))
+    started = time.time()
+    with torch.no_grad():
+        traced = torch.jit.trace(tower, example, check_trace=False)
+    names = ["patches", "pos_embeds", "cos", "sin", "mask"]
+    if args.precision == "fp32":
+        precision = ct.precision.FLOAT32
+    elif args.precision == "fp16":
+        precision = ct.precision.FLOAT16
+    else:
+        keep_fp32 = {"layer_norm", "softmax", "reduce_mean", "reduce_sum", "pow", "sqrt", "rsqrt", "real_div"}
+        precision = ct.transform.FP16ComputePrecision(op_selector=lambda op: op.op_type not in keep_fp32)
+    model = ct.convert(
+        traced, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS17,
+        compute_precision=precision,
+        compute_units=ct.ComputeUnit.CPU_ONLY,
+        inputs=[ct.TensorType(name=n, shape=tuple(t.shape), dtype=np.float32) for n, t in zip(names, example)],
+        outputs=[ct.TensorType(name="tokens", dtype=np.float32)],
+    )
+    model.short_description = "d3 vision tower (vLLM-SR Decision 3.0): one image, fixed patch bucket"
+    model.license = "Apache-2.0"
+    model.user_defined_metadata.update({"patches": str(N), "tokens": str(N // cfg.merge**2)})
+    out = args.out / f"Vision_P{N}"
+    out.mkdir(parents=True, exist_ok=True)
+    package = out / f"VisionTower_{args.precision}.mlpackage"
+    model.save(str(package))
+    pos_table.numpy().astype(np.float32).tofile(out / "pos_embed.f32")  # the host resamples this per image
+    (out / "config.json").write_text(json.dumps({"patches": N, "tokens": N // cfg.merge**2, "patch_dim": cfg.patch_dim,
+                                                  "hidden": cfg.hidden, "head_dim": cfg.head_dim,
+                                                  "out_hidden": cfg.out_hidden, "precision": args.precision},
+                                                 indent=2) + "\n")
+    print(f"saved {package} in {time.time() - started:.0f} s")
+
+
+if __name__ == "__main__":
+    main()

lite/readout.f32

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

lite/tokenizer.json

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

lite/vision_P1536.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

lite/vision_P1536.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

lite/vision_P1536.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "ABBDA214-38F9-44F0-A089-B636CD88D5C9": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "E4544D1C-B3E4-4D3A-A03E-147542E6B2CD": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "ABBDA214-38F9-44F0-A089-B636CD88D5C9"
+}

lite/vision_P3072.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

lite/vision_P3072.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

lite/vision_P3072.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "77FBF011-9CB1-4E8C-8C83-A7CA7E3C7916": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "A8B8926F-B70F-4353-AE83-938EC417390B": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "77FBF011-9CB1-4E8C-8C83-A7CA7E3C7916"
+}

lite/vision_P320.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

lite/vision_P320.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

lite/vision_P320.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "02E1986B-12A5-4C92-8F54-8B6A3CADE0DF": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "AB4963EB-278E-43AA-B7F3-BAC7CACD9427": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "AB4963EB-278E-43AA-B7F3-BAC7CACD9427"
+}