edge/python/d3_graph.py

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


class D3RowPart(nn.Module):
    """Layers [start, stop) of D3RowEmbeds, for checkpoints too large to trace or load as one graph (d3-flash 9B).
    hidden [1, L, H] in; the last part ends with the gather at the last real token and the final norm ([1, H]),
    the others return hidden [1, L, H] for the next part."""

    def __init__(self, cfg: TextConfig, seq_len: int, start: int, stop: int, chunk: int = 64):
        super().__init__()
        self.cfg, self.L, self.start, self.stop = cfg, seq_len, start, stop
        self.last = stop == cfg.num_layers
        self.layers = nn.ModuleList([DecoderLayer(cfg, i, seq_len, chunk) for i in range(start, stop)])
        if self.last:
            self.norm = RMSNorm(cfg.hidden_size, cfg.eps)

    def forward(self, hidden, cos, sin, key_valid, last_onehot):
        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])) if self.last else x


def load_row_part(root: str, seq_len: int, start: int, stop: int, chunk: int = 64, dtype=torch.float16) -> D3RowPart:
    """Only this part's weights are read (layer i -> layers[i - start])."""
    cfg = load_text_config(root)
    prefix = "language_model."
    state = {}
    from safetensors import safe_open
    for path in sorted(__import__("glob").glob(f"{root}/model*.safetensors")):
        with safe_open(path, "pt") as f:
            for k in f.keys():
                if not k.startswith(prefix):
                    continue
                name = k[len(prefix):]
                if name.startswith("layers."):
                    i = int(name.split(".")[1])
                    if start <= i < stop:
                        state[f"layers.{i - start}." + name.split(".", 2)[2]] = f.get_tensor(k).to(dtype)
                elif name.startswith("norm.") and stop == cfg.num_layers:
                    state[name] = f.get_tensor(k).to(dtype)
    model = D3RowPart(cfg, seq_len, start, stop, 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()