| """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") |
| 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}} |
|
|
