D1A-E4B v0.6 Core ML: text trunk, 7 int8 chunks, L=512 (#224)

.gitattributes

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

README.md

@@ -0,0 +1,49 @@
+---
+license: apache-2.0
+base_model: google/gemma-4-E4B
+tags: [coreml, apple-neural-engine, decision-model, d1a]
+---
+# D1A-E4B v0.6 · Core ML (text, ≤ 512 tokens)
+
+A Core ML conversion of D1A-E4B v0.6 (`JohnP1/d1a-e4b-mlx-q8@v0.6`, the model the Mac mini serves). It is the text path
+only, for the Apple Neural Engine. The details are in jonpol01/d1a#224.
+
+## Status
+
+It is **correct at ≤ 512 tokens**. On 28 real requests (42 questions), 0 clear-cut answers changed against the
+fp32 torch rebuild, on an M1 Max and an M4 Neural Engine.
+
+**Speed (median):**
+- M4: 1.8 s per request, against 0.63 s for MLX on the GPU.
+- M1 Max: 2.9 s, against 0.8 s.
+
+Use it where the GPU should stay free or power matters, not for speed.
+
+**Not included yet:**
+- Photo, voice and video (the vision and audio encoders are separate models, still to convert).
+- Inputs longer than 512 tokens: 2,048-token programs give wrong answers on the M4 Neural Engine.
+
+## Files
+
+| Path | What |
+|---|---|
+| `coreml/L512_<a>-<b>.mlpackage` | 7 int8 chunks (per-block 64), layers a..b of the 42-layer trunk. Chunk 18-24 outputs the shared keys/values that chunks 24-42 read. |
+| `io.json` | each chunk's input names |
+| `embeddings.safetensors` | the token and per-layer embedding tables (MLX affine-quantized), looked up on the host, not in Core ML |
+| `head.safetensors`, `d1a_config.json` | pointer head (q, k), temperature 1.954, pad and delimiter ids |
+| `tokenizer.json` | Gemma 4 tokenizer |
+| `reference/` | Python reference: host lookups, masks, the chunk chain (one process per chunk), the pointer head |
+
+## How a request runs
+
+1. Encode with D1A's layout (`d1a.core.encoding.encode`) and right-pad to 512 with the pad token.
+2. Look up the embeddings (`emb` scaled by √2560) and the per-layer embeddings (`ple_tok` scaled by √256).
+3. Build the block-causal masks (full and sliding-512), clamped to -1e4.
+4. Run the 7 chunks in order, each in its own process.
+5. Read the hidden states at each question's `<decide>` and `</opt>` positions.
+6. Score them with the head, divide by the temperature, and take the softmax.
+
+## Requirements
+
+macOS 15 / iOS 18 or later. Loading 7 large Neural Engine programs in one process can crash the runtime, so the
+reference runs one process per chunk.

config.json

@@ -0,0 +1,171 @@
+{
+    "architectures": [
+        "Gemma4ForConditionalGeneration"
+    ],
+    "audio_config": {
+        "_name_or_path": "",
+        "architectures": null,
+        "attention_chunk_size": 12,
+        "attention_context_left": 13,
+        "attention_context_right": 0,
+        "attention_invalid_logits_value": -1000000000.0,
+        "attention_logit_cap": 50.0,
+        "chunk_size_feed_forward": 0,
+        "conv_kernel_size": 5,
+        "dtype": "bfloat16",
+        "gradient_clipping": 10000000000.0,
+        "hidden_act": "silu",
+        "hidden_size": 1024,
+        "id2label": {
+            "0": "LABEL_0",
+            "1": "LABEL_1"
+        },
+        "initializer_range": 0.02,
+        "is_encoder_decoder": false,
+        "label2id": {
+            "LABEL_0": 0,
+            "LABEL_1": 1
+        },
+        "model_type": "gemma4_audio",
+        "num_attention_heads": 8,
+        "num_hidden_layers": 12,
+        "output_attentions": false,
+        "output_hidden_states": false,
+        "output_proj_dims": 1536,
+        "problem_type": null,
+        "residual_weight": 0.5,
+        "return_dict": true,
+        "rms_norm_eps": 1e-06,
+        "subsampling_conv_channels": [
+            128,
+            32
+        ],
+        "use_clipped_linears": true
+    },
+    "audio_token_id": 258881,
+    "boa_token_id": 256000,
+    "boi_token_id": 255999,
+    "dtype": "bfloat16",
+    "eoa_token_id": 258883,
+    "eoa_token_index": 258883,
+    "eoi_token_id": 258882,
+    "eos_token_id": 1,
+    "image_token_id": 258880,
+    "initializer_range": 0.02,
+    "model_type": "gemma4",
+    "quantization": {
+        "group_size": 64,
+        "bits": 8,
+        "mode": "affine",
+        "language_model.model.embed_tokens_per_layer": {
+            "group_size": 64,
+            "bits": 4,
+            "mode": "affine"
+        }
+    },
+    "quantization_config": {
+        "group_size": 64,
+        "bits": 8,
+        "mode": "affine",
+        "language_model.model.embed_tokens_per_layer": {
+            "group_size": 64,
+            "bits": 4,
+            "mode": "affine"
+        }
+    },
+    "text_config": {
+        "attention_bias": false,
+        "attention_dropout": 0.0,
+        "attention_k_eq_v": false,
+        "bos_token_id": 2,
+        "dtype": "bfloat16",
+        "enable_moe_block": false,
+        "eos_token_id": 1,
+        "expert_intermediate_size": null,
+        "final_logit_softcapping": 30.0,
+        "global_head_dim": 512,
+        "head_dim": 256,
+        "hidden_activation": "gelu_pytorch_tanh",
+        "hidden_size": 2560,
+        "hidden_size_per_layer_input": 256,
+        "initializer_range": 0.02,
+        "intermediate_size": 10240,
+        "layer_types": [
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention"
+        ],
+        "max_position_embeddings": 131072,
+        "model_type": "gemma4_text",
+        "num_attention_heads": 8,
+        "num_experts": null,
+        "num_global_key_value_heads": null,
+        "num_hidden_layers": 42,
+        "num_key_value_heads": 2,
+        "num_kv_shared_layers": 18,
+        "pad_token_id": 0,
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {
+            "full_attention": {
+                "partial_rotary_factor": 0.25,
+                "rope_theta": 1000000.0,
+                "rope_type": "proportional"
+            },
+            "sliding_attention": {
+                "rope_theta": 10000.0,
+                "rope_type": "default"
+            }
+        },
+        "sliding_window": 512,
+        "tie_word_embeddings": true,
+        "top_k_experts": null,
+        "use_bidirectional_attention": null,
+        "use_cache": true,
+        "use_double_wide_mlp": false,
+        "vocab_size": 262144,
+        "vocab_size_per_layer_input": 262144
+    },
+    "tie_word_embeddings": true,
+    "transformers_version": "5.5.0.dev0",
+    "video_token_id": 258884,
+    "vision_soft_tokens_per_image": 280
+}

coreml/L512_00-06.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

coreml/L512_00-06.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_00-06.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "3C09B49F-607E-44C3-8A4A-BE61559AAC82": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "E68B5681-3E21-4324-A4E2-C98CC1AFAF15": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "E68B5681-3E21-4324-A4E2-C98CC1AFAF15"
+}

coreml/L512_06-12.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

coreml/L512_06-12.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_06-12.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "1E669D90-60B5-4734-AC55-7740F8D01BA4": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "B0AFD2A8-B841-4662-B82D-D2CCB4B83093": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "B0AFD2A8-B841-4662-B82D-D2CCB4B83093"
+}

coreml/L512_12-18.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

coreml/L512_12-18.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_12-18.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "2C8B7345-8589-4596-85E7-BEDDD4493047": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "82121673-49E3-483C-995A-8FA0699BF674": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "82121673-49E3-483C-995A-8FA0699BF674"
+}

coreml/L512_18-24.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

coreml/L512_18-24.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_18-24.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "07C10B3B-C0D2-493A-98F2-873F1FA95521": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "703DA848-64F6-498C-BCAC-B4990539ACC0": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "07C10B3B-C0D2-493A-98F2-873F1FA95521"
+}

coreml/L512_24-30.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

coreml/L512_24-30.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_24-30.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "481F0249-E92D-463B-9B6B-6F8606E597ED": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "F73B99D0-DC7C-4319-9AB4-2A3E3EFB7838": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "F73B99D0-DC7C-4319-9AB4-2A3E3EFB7838"
+}

coreml/L512_30-36.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

coreml/L512_30-36.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_30-36.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "92E60EB2-8586-46A5-AAFB-175CFD413BA4": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "F9C38AD8-3C08-487B-930C-0E61ED5BD177": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "92E60EB2-8586-46A5-AAFB-175CFD413BA4"
+}

coreml/L512_36-42.mlpackage/Data/com.apple.CoreML/model.mlmodel

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

coreml/L512_36-42.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

coreml/L512_36-42.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "2E3E8772-B27F-4A92-98C4-57320668E2BE": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "D218135C-CC81-4D22-A335-A01164CC36C5": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "2E3E8772-B27F-4A92-98C4-57320668E2BE"
+}

d1a_config.json

@@ -0,0 +1,38 @@
+{
+  "format": "d1a-mlx",
+  "format_version": 1,
+  "base": "google/gemma-4-E4B",
+  "base_revision": "411aa17b749aa952df1359d2dcea73917a544d9a",
+  "source": "JohnP1/d1a-e4b",
+  "source_revision": "3ca01b824280cacf027826b0ded8632f242a7e8b",
+  "adapter_sha256": "739ce962c3ae630243ee53b212a822e8b2feee41080b68d5b460fd0ebb19e4ea",
+  "lora": 16,
+  "head_dim": 256,
+  "hidden_size": 2560,
+  "temperature": 1.9543199368684914,
+  "leading_ids": [
+    2
+  ],
+  "bos_id": 2,
+  "delimiter_ids": [
+    6,
+    7,
+    8,
+    9,
+    10
+  ],
+  "pad_id": 0,
+  "dtype": "bfloat16",
+  "quantization": {
+    "group_size": 64,
+    "bits": 8,
+    "mode": "affine",
+    "language_model.model.embed_tokens_per_layer": {
+      "group_size": 64,
+      "bits": 4,
+      "mode": "affine"
+    },
+    "embeddings": true
+  },
+  "head": "head.safetensors"
+}

embeddings.safetensors

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

head.safetensors

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

io.json

@@ -0,0 +1 @@
+{"chunks": ["00-06", "06-12", "12-18", "18-24", "24-30", "30-36", "36-42"], "inputs": {"00-06": ["emb", "ple_tok", "pos", "mask_full", "mask_slide"], "06-12": ["h_in", "ple", "pos", "mask_full", "mask_slide"], "12-18": ["h_in", "ple", "pos", "mask_full", "mask_slide"], "18-24": ["h_in", "ple", "pos", "mask_full", "mask_slide"], "24-30": ["h_in", "ple", "pos", "mask_full", "mask_slide", "k_full_attention", "v_full_attention", "k_sliding_attention", "v_sliding_attention"], "30-36": ["h_in", "ple", "pos", "mask_full", "mask_slide", "k_full_attention", "v_full_attention", "k_sliding_attention", "v_sliding_attention"], "36-42": ["h_in", "ple", "pos", "mask_full", "mask_slide", "k_full_attention", "v_full_attention", "k_sliding_attention", "v_sliding_attention", "picks"]}}

reference/d1a_coreml.py

@@ -0,0 +1,183 @@
+"""D1A (Gemma 4 text) -> Core ML pilot (jonpol01/d1a#224): rebuild a D1A MLX export as a torch module, split into
+  - host side: the two embedding tables (token embedding, per-layer embedding: 2.75B of E2B's 4.6B parameters) read
+    row by row from the export's safetensors (dequantized per request, never held whole), and
+  - Core: everything that computes (per-layer projection, 35 decoder layers, final norm, pointer head), which is what
+    gets converted to Core ML.
+Inputs follow D1A's packed encoding (d1a.core.encoding.encode): one record, block-causal mask, positions restarting per
+question after the state, readout at each question's <decide> and </opt> tokens."""
+import json
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors import safe_open
+from torch import nn
+from transformers import AutoConfig
+from transformers.models.gemma4 import modeling_gemma4 as G
+
+PREFIX = "language_model.model."
+
+
+def dequant(packed, scales, biases, bits, group=64):
+    """MLX affine dequantization: packed uint32 [out, in*bits/32] (lowest bits first) -> float32 [out, in]."""
+    p = packed.view(torch.int32).to(torch.int64) & 0xFFFFFFFF
+    per = 32 // bits
+    shifts = torch.arange(per, dtype=torch.int64) * bits
+    q = ((p[..., None] >> shifts) & ((1 << bits) - 1)).reshape(p.shape[0], -1).to(torch.float32)
+    s, b = scales.to(torch.float32), biases.to(torch.float32)
+    return (q.view(q.shape[0], -1, group) * s[..., None] + b[..., None]).reshape(q.shape[0], -1)
+
+
+class _Shards:
+    """model.safetensors or model-0000N-of-0000M.safetensors (via model.safetensors.index.json), read like one safe_open."""
+
+    def __init__(self, path):
+        idx = Path(path) / "model.safetensors.index.json"
+        if idx.exists():
+            wm = json.loads(idx.read_text())["weight_map"]
+            files = {f: safe_open(str(Path(path) / f), "pt") for f in sorted(set(wm.values()))}
+            self.where = {k: files[f] for k, f in wm.items()}
+        else:
+            f = safe_open(str(Path(path) / "model.safetensors"), "pt")
+            self.where = {k: f for k in f.keys()}
+
+    def keys(self):
+        return self.where.keys()
+
+    def get_tensor(self, k):
+        return self.where[k].get_tensor(k)
+
+    def get_slice(self, k):
+        return self.where[k].get_slice(k)
+
+
+class Export:
+    """A D1A MLX export directory (config.json, d1a_config.json, model.safetensors, head.safetensors)."""
+
+    def __init__(self, path):
+        self.path = Path(path)
+        self.cfg = AutoConfig.from_pretrained(self.path)
+        self.tc = self.cfg.text_config
+        self.tc._attn_implementation = "eager"
+        self.d1a = json.loads((self.path / "d1a_config.json").read_text())
+        q = self.d1a["quantization"]
+        self.bits = {"default": (q["bits"], q["group_size"]), "embed_tokens_per_layer": (q[PREFIX.replace("model.", "") + "model.embed_tokens_per_layer"]["bits"] if False else q["language_model.model.embed_tokens_per_layer"]["bits"], q["language_model.model.embed_tokens_per_layer"]["group_size"])}
+        self.f = _Shards(self.path)
+
+    def bits_of(self, name):
+        return self.bits["embed_tokens_per_layer"] if name.endswith("embed_tokens_per_layer") else self.bits["default"]
+
+    def tensor(self, name):
+        """A float32 tensor for `name` (without PREFIX), dequantized if it is stored quantized."""
+        full = PREFIX + name
+        keys = set(self.f.keys())
+        if full + ".weight" in keys and full + ".scales" in keys:
+            bits, group = self.bits_of(name)
+            return dequant(self.f.get_tensor(full + ".weight"), self.f.get_tensor(full + ".scales"), self.f.get_tensor(full + ".biases"), bits, group)
+        return self.f.get_tensor(full).to(torch.float32)
+
+    def text_model(self, dtype=torch.float32):
+        """The Gemma 4 text trunk with every non-embedding weight loaded (embedding tables stay empty: host side)."""
+        with torch.device("meta"):
+            m = G.Gemma4TextModel(self.tc)
+        m = m.to_empty(device="cpu")
+        state = {}
+        keys = set(self.f.keys())
+        for name, p in m.named_parameters():
+            if name.startswith("embed_tokens"):
+                continue
+            base = name[: -len(".weight")] if name.endswith(".weight") else name
+            if PREFIX + base + ".scales" in keys:
+                state[name] = self.tensor(base)
+            else:
+                state[name] = self.f.get_tensor(PREFIX + name).to(torch.float32)
+            assert state[name].shape == p.shape, (name, state[name].shape, p.shape)
+        for name, b in m.named_buffers():
+            if PREFIX + name in keys:
+                state[name] = self.f.get_tensor(PREFIX + name).to(torch.float32)
+        missing, unexpected = m.load_state_dict(state, strict=False)
+        missing = [k for k in missing if not k.startswith("embed_tokens")]
+        assert not missing and not unexpected, (missing, unexpected)
+        # buffers transformers computes itself (rotary inv_freq, embed_scale): re-run their init on CPU
+        m.rotary_emb = G.Gemma4TextRotaryEmbedding(self.tc)
+        return m.to(dtype).eval()
+
+    def head(self):
+        h = safe_open(str(self.path / "head.safetensors"), "pt")
+        dp, d = h.get_tensor("q.weight").shape                     # pointer width, hidden size (E2B 1536, E4B 2560)
+        q, k = nn.Linear(d, dp), nn.Linear(d, dp)
+        q.weight.data, q.bias.data = h.get_tensor("q.weight").float(), h.get_tensor("q.bias").float()
+        k.weight.data, k.bias.data = h.get_tensor("k.weight").float(), h.get_tensor("k.bias").float()
+        return q, k
+
+
+class HostEmbed:
+    """Row lookups of the two embedding tables, dequantized per request: emb = E[ids] * sqrt(hidden) and
+    ple = E_pl[ids] * sqrt(ple_dim), reshaped [L, layers, ple_dim] (transformers' get_per_layer_inputs)."""
+
+    def __init__(self, ex):
+        self.ex, self.tc = ex, ex.tc
+        self.f = ex.f
+
+    def _rows(self, name, ids):
+        full = PREFIX + name
+        bits, group = self.ex.bits_of(name)
+        w, s, b = (self.f.get_slice(full + suf) for suf in (".weight", ".scales", ".biases"))
+        uniq, inv = np.unique(np.asarray(ids), return_inverse=True)
+        rows = [dequant(w[int(i): int(i) + 1], s[int(i): int(i) + 1], b[int(i): int(i) + 1], bits, group) for i in uniq]
+        return torch.cat(rows)[torch.as_tensor(inv)]
+
+    def __call__(self, ids):
+        emb = self._rows("embed_tokens", ids) * (self.tc.hidden_size ** 0.5)
+        ple = self._rows("embed_tokens_per_layer", ids) * (self.tc.hidden_size_per_layer_input ** 0.5)
+        return emb[None], ple.view(1, len(ids), self.tc.num_hidden_layers, self.tc.hidden_size_per_layer_input)
+
+
+def masks_for(enc, window, L=None, dtype=torch.float32):
+    """D1A's packed masks (d1a.backends.torch.branch_masks) for one record, padded to L: additive [1, 1, L, L]."""
+    seg, pos = enc["seg"], enc["pos"]
+    n = len(seg)
+    L = L or n
+    s = torch.full((L,), -1)
+    s[:n] = torch.tensor(seg)
+    p = torch.zeros(L, dtype=torch.long)
+    p[:n] = torch.tensor(pos)
+    causal = torch.tril(torch.ones(L, L, dtype=torch.bool))
+    same = (s[None, :] == s[:, None]) | (s[None, :] == 0)
+    allow = (causal & same & (s != -1)[None, :]) | torch.eye(L, dtype=torch.bool)
+    neg = torch.finfo(dtype).min
+    full = torch.zeros(L, L, dtype=dtype).masked_fill(~allow, neg)
+    far = (p[:, None] - p[None, :] >= window) & ~torch.eye(L, dtype=torch.bool)
+    slide = full.masked_fill(far, neg)
+    return full[None, None], slide[None, None], p[None]
+
+
+class Core(nn.Module):
+    """Everything that computes: per-layer projection, decoder layers, final norm, pointer head.
+    forward(emb [1,L,d], ple_tok [1,L,layers,256], pos [1,L] int, mask_full, mask_slide [1,1,L,L],
+            decide [Q] int, opts [Q,K] int, opt_valid [Q,K] float (1 used, 0 padding)) -> logits [Q,K]
+    (raw, temperature not applied; padded options get a large negative)."""
+
+    def __init__(self, text, head_q, head_k, silu_tanh=True):
+        super().__init__()
+        self.text, self.hq, self.hk = text, head_q, head_k
+        self.types = text.config.layer_types
+        self.scale = 1 / math.sqrt(head_q.out_features)
+
+    def forward(self, emb, ple_tok, pos, mask_full, mask_slide, decide, opts, opt_valid):
+        t = self.text
+        ple = t.project_per_layer_inputs(emb, ple_tok)
+        rope = {lt: t.rotary_emb(emb, pos, lt) for lt in t.unique_layer_types}
+        masks = {"full_attention": mask_full, "sliding_attention": mask_slide}
+        from collections import UserDict
+        kv = UserDict()
+        h = emb
+        for i, layer in enumerate(t.layers):
+            h = layer(h, ple[:, :, i, :], shared_kv_states=kv, position_embeddings=rope[self.types[i]],
+                      attention_mask=masks[self.types[i]], position_ids=pos)
+        h = t.norm(h)[0]                                    # [L, d]
+        q = self.hq(h[decide])                              # [Q, dp]
+        k = self.hk(h[opts.reshape(-1)]).view(*opts.shape, -1)   # [Q, K, dp]
+        z = (k * q[:, None, :]).sum(-1) * self.scale
+        return z + (opt_valid - 1.0) * 1e4

reference/pipeline.py

@@ -0,0 +1,79 @@
+"""Run a chunked D1A Core ML trunk end to end: host embeddings -> chunk 0 .. chunk N (h, per-layer inputs and the shared
+keys/values passed along) -> hidden states at the picks -> pointer head on the host.
+Usage as a library: Pipeline(export_dir, chunk_dir, L, units).probs(enc) -> [per-question probabilities]."""
+import glob
+import re
+import sys
+import time
+
+import coremltools as ct
+import numpy as np
+import torch
+
+sys.path.insert(0, "/Users/john/d1a-runs/coreml")
+from d1a_coreml import Export, HostEmbed, masks_for  # noqa: E402
+
+P_PICKS = 128
+
+
+class Pipeline:
+    def __init__(self, export_dir, chunk_dir, L, units="CPU_AND_NE"):
+        self.ex = Export(export_dir)
+        self.host = HostEmbed(self.ex)
+        self.hq, self.hk = self.ex.head()
+        self.T = self.ex.d1a["temperature"]
+        self.L = L
+        cu = getattr(ct.ComputeUnit, units)
+        self.chunks = []
+        if chunk_dir.endswith(".mlpackage"):   # one multifunction package: functions c00_07, c07_15, ...
+            names = ["c00_07", "c07_15", "c15_22", "c22_29", "c29_35"]
+            for name in names:
+                a, b = map(int, name[1:].split("_"))
+                m = ct.models.MLModel(chunk_dir, compute_units=cu, function_name=name)
+                fn = next(f for f in m.get_spec().description.functions if f.name == name)
+                self.chunks.append((a, b, m, [i.name for i in fn.input], [o.name for o in fn.output]))
+            return
+        files = sorted(glob.glob(f"{chunk_dir}/L{L}_*.mlpackage"))
+        for f in files:
+            a, b = map(int, re.search(r"_(\d+)-(\d+)\.mlpackage$", f).groups())
+            m = ct.models.MLModel(f, compute_units=cu)
+            spec = m.get_spec()
+            self.chunks.append((a, b, m, [i.name for i in spec.description.input], [o.name for o in spec.description.output]))
+
+    def hidden(self, enc):
+        L, n = self.L, len(enc["ids"])
+        # right-pad with the pad token's own embeddings (as d1a's torch backend pads ids), never zero rows: an all-zero row
+        # makes every RMSNorm compute 0 / sqrt(eps), and eps (1e-6) is subnormal in fp16, which the Neural Engine flushes
+        ids = list(enc["ids"]) + [self.ex.d1a.get("pad_id", 0)] * (L - n)
+        emb, ple = self.host(ids)
+        e = emb.numpy().astype(np.float16)
+        p = ple.numpy().astype(np.float16)
+        mf, ms, pos = masks_for(enc, self.ex.tc.sliding_window, L=L)
+        picks = list(enc["decide_idx"]) + [o for oo in enc["opt_idx"] for o in oo]
+        pk = np.zeros(P_PICKS, np.int32); pk[: len(picks)] = picks
+        common = {"pos": pos.numpy().astype(np.int32), "mask_full": mf.clamp(min=-1e4).numpy().astype(np.float16),
+                  "mask_slide": ms.clamp(min=-1e4).numpy().astype(np.float16)}
+        state = {"emb": e, "ple_tok": p, "picks": pk}
+        for a, b, m, ins, outs in self.chunks:
+            feed = {k: np.ascontiguousarray(common[k] if k in common else state[k]) for k in ins}
+            feed = {k: (v.astype(np.float16) if v.dtype == np.float32 else v) for k, v in feed.items()}   # chunk outputs come back fp32
+            out = {k: np.array(v, copy=True) for k, v in m.predict(feed).items()}   # Core ML may reuse its output buffers on the next call
+            if "hidden" in out:
+                return torch.tensor(out["hidden"][: len(picks)], dtype=torch.float32)
+            state["h_in"] = out["h"]
+            if "ple_out" in out:
+                state["ple"] = out["ple_out"]
+            for k, v in out.items():
+                if k.endswith("_out") and k != "ple_out":
+                    state[k[: -len("_out")]] = v
+        raise RuntimeError("no chunk returned hidden")
+
+    def probs(self, enc):
+        h = self.hidden(enc)
+        Q = len(enc["decide_idx"])
+        out, j = [], Q
+        with torch.no_grad():
+            for qi, oo in enumerate(enc["opt_idx"]):
+                z = (self.hk(h[j: j + len(oo)]) @ self.hq(h[qi])) / (self.hq.out_features ** 0.5)
+                out.append(torch.softmax(z / self.T, -1).numpy()); j += len(oo)
+        return out

reference/workers.py

@@ -0,0 +1,99 @@
+"""One process per Core ML chunk (workaround: several long-context chunks in one process crash the Neural Engine runtime
+on the second request on this Mac, while each chunk alone, or all on the GPU, runs fine). The parent runs the host side
+(embeddings, masks, pointer head) and passes tensors along the chain through pipes.
+  WorkerPipeline(export_dir, chunk_dir, L, units).probs(enc)"""
+import glob
+import multiprocessing as mp
+import re
+import sys
+
+import numpy as np
+import torch
+
+sys.path.insert(0, "/Users/john/d1a-runs/coreml")
+
+
+import os
+BUCKETS = tuple(int(x) for x in os.environ.get("D1A_COREML_BUCKETS", "512,2048").split(","))   # each length is its own ANE program; fewer keeps the ANE within its memory
+
+
+def _worker(path, units, conn):
+    import coremltools as ct
+    cu = getattr(ct.ComputeUnit, units)
+    # a precompiled .mlmodelc (compiled once by coremltools.models.utils.compile_model): loading each length function
+    # from it never recompiles, where MLModel(.mlpackage) compiles the whole package again for every function
+    models = {L: ct.models.CompiledMLModel(path, compute_units=cu, function_name=f"L{L}") for L in BUCKETS}
+    spec = ct.utils.load_spec(path.replace("compiled/", "buckets/").replace(".mlmodelc", ".mlpackage"))
+    fn = next(f for f in spec.description.functions if f.name == f"L{BUCKETS[0]}")
+    conn.send(([i.name for i in fn.input], [o.name for o in fn.output]))
+    while True:
+        msg = conn.recv()
+        if msg is None:
+            break
+        L, feed = msg
+        import time as _t
+        t0 = _t.perf_counter()
+        out = models[L].predict(feed)
+        t1 = _t.perf_counter()
+        res = {k: np.array(v, dtype=np.float16 if v.dtype == np.float32 else v.dtype, copy=True) for k, v in out.items()}
+        res["_predict_ms"] = np.array((t1 - t0) * 1000)
+        conn.send(res)
+
+
+class WorkerPipeline:
+    def __init__(self, export_dir, chunk_dir, L=None, units="CPU_AND_NE"):
+        from d1a_coreml import Export, HostEmbed
+        self.ex = Export(export_dir)
+        self.host = HostEmbed(self.ex)
+        self.hq, self.hk = self.ex.head()
+        self.T, self.L = self.ex.d1a["temperature"], L
+        ctx = mp.get_context("spawn")
+        self.chunks = []
+        for f in sorted(glob.glob(f"{chunk_dir}/*.mlmodelc")):
+            a, b = map(int, re.search(r"(\d+)-(\d+)\.mlmodelc$", f).groups())
+            parent, child = ctx.Pipe()
+            p = ctx.Process(target=_worker, args=(f, units, child), daemon=True)
+            p.start()
+            self.chunks.append((a, b, p, parent))
+            self.io = getattr(self, "io", []) + [parent.recv()]   # one at a time: concurrent first loads race in Core ML's compiler
+
+    def close(self):
+        for _, _, p, conn in self.chunks:
+            conn.send(None); p.join(5)
+
+    def hidden(self, enc):
+        from d1a_coreml import masks_for
+        self.last_predict_ms = []
+        n = len(enc["ids"])
+        L = next(b for b in BUCKETS if b >= n)   # the smallest compiled length that fits
+        ids = list(enc["ids"]) + [self.ex.d1a.get("pad_id", 0)] * (L - n)
+        emb, ple = self.host(ids)
+        mf, ms, pos = masks_for(enc, self.ex.tc.sliding_window, L=L)
+        picks = list(enc["decide_idx"]) + [o for oo in enc["opt_idx"] for o in oo]
+        pk = np.zeros(128, np.int32); pk[: len(picks)] = picks
+        common = {"pos": pos.numpy().astype(np.int32), "mask_full": mf.clamp(min=-1e4).numpy().astype(np.float16),
+                  "mask_slide": ms.clamp(min=-1e4).numpy().astype(np.float16)}
+        state = {"emb": emb.numpy().astype(np.float16), "ple_tok": ple.numpy().astype(np.float16), "picks": pk}
+        for (a, b, p, conn), (ins, outs) in zip(self.chunks, self.io):
+            conn.send((L, {k: (common[k] if k in common else state[k]) for k in ins}))
+            out = conn.recv()
+            self.last_predict_ms = getattr(self, "last_predict_ms", []) + [float(out.pop("_predict_ms"))]
+            if "hidden" in out:
+                return torch.tensor(out["hidden"][: len(picks)].astype(np.float32))
+            state["h_in"] = out["h"]
+            if "ple_out" in out:
+                state["ple"] = out["ple_out"]
+            for k, v in out.items():
+                if k.endswith("_out") and k != "ple_out":
+                    state[k[: -len("_out")]] = v
+        raise RuntimeError("no chunk returned hidden")
+
+    def probs(self, enc):
+        h = self.hidden(enc)
+        Q = len(enc["decide_idx"])
+        out, j = [], Q
+        with torch.no_grad():
+            for qi, oo in enumerate(enc["opt_idx"]):
+                z = (self.hk(h[j: j + len(oo)]) @ self.hq(h[qi])) / (self.hq.out_features ** 0.5)
+                out.append(torch.softmax(z / self.T, -1).numpy()); j += len(oo)
+        return out

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,55 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}