Runtime MPSGraph null-deref (mlir::FloatType::getWidth) executing 2+ GatedDeltaUpdate layers + an attention layer with dynamic KV context

[scndls]

Runtime crash (MPSGraph mlir::FloatType::getWidth() null-deref) executing a hybrid model with 2+ GatedDeltaUpdate layers + an attention layer with dynamic KV context

Summary

A model that stacks 2 or more GatedDeltaUpdate (Gated DeltaNet) layers together with one
attention layer reading a dynamic-length KV context (the Qwen3.5 / Qwen3-Next hybrid shape)
exports and loads fine, then segfaults on the first execution:

thread #23, stop reason = EXC_BAD_ACCESS (code=1, address=0x0)
  frame #0: MetalPerformanceShadersGraph`mlir::FloatType::getWidth() + 16

A null FloatType is dereferenced during MPSGraph JIT compilation of the executed function.

The crash is governed precisely by the number of GatedDeltaUpdate layers:

Model	Result
1 GatedDeltaUpdate layer + 1 attention layer (dynamic context)	runs (exit 0)
2 GatedDeltaUpdate layers + 1 attention layer (dynamic context)	EXC_BAD_ACCESS at execute
N stacked GatedDeltaUpdate layers, no attention, fully static export	runs (verified to N=3)

So it is not the multi-layer DeltaNet state handling (that stacks fine on its own), and
not attention alone (a single attention layer with a static query + dynamic context runs).
It is the combination: 2+ DeltaNet scf.while scans plus a dynamic KV-context dimension in
the same exported function.

Environment

• coreai-torch 0.4.0, coreai-models @ b1cb71b
• torch 2.9.0, Python 3.11.15
• macOS 27.0 (build 26A5353q), arm64 (Apple M1 Max, 32 GB)

Repro

Full self-contained script: toy_hybrid.py (≈180 lines; attached / available on request).
It builds embed → N×[RMSNorm + GatedDeltaNet(conv1d+GatedDeltaUpdate+gated RMSNorm) + MLP] → [RMSNorm + gpt_oss-style attention + MLP] → RMSNorm → lm_head, with a static query
(input_ids, len 12) and a dynamic KV context (position_ids len + KV-cache seq dim).

Shape policy (the crux): the query length is static so each DeltaNet GatedDeltaUpdate's
scf.while lowers; the attention reads a dynamic context the gpt_oss.py way
(sequence_length = position_ids.shape[-1], offset = sequence_length - query_len,
cache.update_and_fetch(..., seq_len=sequence_length, query_len=query_len)).

dynamic_shapes = {
    "input_ids": None,                                              # STATIC query chunk
    "position_ids": {1: torch.export.Dim("ctx",  min=13, max=64)},  # DYNAMIC context
    "k_cache":      {3: torch.export.Dim("kseq", min=13, max=64)},
    "v_cache":      {3: torch.export.Dim("vseq", min=13, max=64)},
    "conv_state": None, "ssm_state": None,
}
prog = export_to_coreai(model, inputs, input_names=("input_ids","position_ids"),
                        output_names=("logits",),
                        state_names=("k_cache","v_cache","conv_state","ssm_state"),
                        dynamic_shapes=dynamic_shapes)
prog.optimize(); prog.save_asset(path)
m = await AIModel.load(path)
fn = m.load_function(m.function_names[0])
await fn(inputs=..., state=...)          # <-- EXC_BAD_ACCESS for NDELTA=2, OK for NDELTA=1

Run:

NDELTA=1 uv run python toy_hybrid.py    # RAN — exit 0
NDELTA=2 uv run python toy_hybrid.py    # EXC_BAD_ACCESS in MPSGraph FloatType::getWidth()

Notes

• SDPA is de-externalized in the repro (removed from _EXTERNALIZE_SPECS so it decomposes
  to primitive ops) purely to work around a separate coreai-torch externalize bug that blocks
  exporting a static-query / dynamic-context model with externalized SDPA (filed separately —
  the d_20 unbounded-key-dim issue). If that export bug is fixed and SDPA stays fused, this
  MPSGraph crash should be re-checked with the fused attention kernel.
• The model is tiny (hidden 256, 2 layers), so this is not a memory-pressure issue.
• This blocks running native Qwen3.5 / Qwen3-Next style hybrids (e.g. a 27B with 48 DeltaNet +
  16 attention layers) on Core AI, since they necessarily stack many DeltaNet layers alongside
  attention with a dynamic context.

[Lewis300]

Hello, could you please provide your repro script, toy_hybrid.py?

[scndls]

Sure — full script inline below. Self-contained apart from a coreai-models checkout (it imports the macos primitives from there).

NDELTA=1 uv run python toy_hybrid.py   # runs, exit 0
NDELTA=2 uv run python toy_hybrid.py   # EXC_BAD_ACCESS in MPSGraph FloatType::getWidth() at first execute

Re-verified both today on coreai-torch 0.4.0 / coreai-models b1cb71b / macOS 27.0 (26A5353q).

Two things you'll see when running it, neither of which is this bug:

• The NDELTA=1 run prints an ANECCompile() FAILED ... MLIR MPS to ANEC conversion failed error on the first execute. The runtime falls back and the call still completes (exit 0).
• NDELTA=1 reports rel err ~0.09 ("NUMERIC-OFF"). The toy's attention is deliberately simplified (full-dim RoPE, no output gating), so its numerics are loose; a faithful build of the same layer stack passes numeric checks. The issue here is the NDELTA=2 segfault, not the numbers.

One correction to the environment section in the issue: the machine is an Apple M1 Max, 32 GB — not M3. Same crash regardless.

toy_hybrid.py

"""Repro for apple/coreai-torch#2: N GatedDeltaUpdate (DeltaNet) layers + 1 attention
layer in one exported function, with a STATIC query chunk and a DYNAMIC KV context
(the qwen3.5 / qwen3-next hybrid shape), minimal assembly.

  NDELTA=1  exports, loads, executes (exit 0)
  NDELTA=2  exports and loads fine, then EXC_BAD_ACCESS (null deref) in MPSGraph
            mlir::FloatType::getWidth() on the first execution

SDPA is removed from _EXTERNALIZE_SPECS so it decomposes to primitive ops; with SDPA
externalized this model doesn't export at all (issue #1).

Run from a coreai-models checkout:

  NDELTA=1 uv run python toy_hybrid.py
  NDELTA=2 uv run python toy_hybrid.py
"""

import asyncio
import os
import sys
import tempfile
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from coreai.runtime import AIModel, NDArray
import coreai_models.export.macos as _macos
from coreai_models.export.macos import export_to_coreai
from coreai_models.primitives._ops import mutable_slice_update
from coreai_models.primitives.macos.cache import KVCache
from coreai_models.primitives.macos.mlp import MLP
from coreai_models.primitives.macos.rms_norm import RMSNorm, RMSNormGated
from coreai_models.primitives.macos.rope import initialize_rope
from coreai_models.primitives.macos.sdpa import SDPA
from coreai_torch.composite_ops import GatedDeltaUpdate

# decompose SDPA instead of externalizing it (works around issue #1)
_macos._EXTERNALIZE_SPECS = [
    s for s in _macos._EXTERNALIZE_SPECS
    if getattr(s, "composite_op_name", "") != "scaled_dot_product_attention"
]

HID, VOCAB = 256, 512
NK, NV, DK, DV, K = 2, 4, 32, 32, 4
KEY_DIM, VALUE_DIM = DK * NK, DV * NV
CONV_DIM = KEY_DIM * 2 + VALUE_DIM
NH, NKV, HD = 8, 2, 64
NDELTA = int(os.environ.get("NDELTA", "2"))


def repeat_kv(x, n):
    b, h, s, d = x.shape
    return x[:, :, None, :, :].expand(b, h, n, s, d).reshape(b, h * n, s, d)


def _slice_write(stack, slot, value):
    begin = torch.tensor([slot] + [0] * (stack.dim() - 1), dtype=torch.int32)
    end = torch.tensor([slot + 1] + [stack.size(i) for i in range(1, stack.dim())], dtype=torch.int32)
    mutable_slice_update(x=stack, update=value.unsqueeze(0), begin=begin, end=end)


class DeltaNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.in_qkv = nn.Linear(HID, CONV_DIM, bias=False)
        self.in_z = nn.Linear(HID, VALUE_DIM, bias=False)
        self.in_b = nn.Linear(HID, NV, bias=False)
        self.in_a = nn.Linear(HID, NV, bias=False)
        self.dt_bias = nn.Parameter(torch.ones(NV))
        self.A_log = nn.Parameter(torch.log(torch.rand(NV) * 16 + 1e-3))
        self.conv1d = nn.Conv1d(CONV_DIM, CONV_DIM, K, groups=CONV_DIM, bias=False)
        self.norm = RMSNormGated(DV, eps=1e-6)
        self.out_proj = nn.Linear(VALUE_DIM, HID, bias=False)
        self.gdu = GatedDeltaUpdate(use_qk_l2_norm=True)

    def forward(self, x, ss, cs):
        b, s, _ = x.shape
        qkv = self.in_qkv(x)
        z = self.in_z(x).reshape(b, s, NV, DV)
        beta = torch.sigmoid(self.in_b(x))
        a = self.in_a(x)
        g = -torch.exp(self.A_log.float()) * F.softplus(a.float() + self.dt_bias.float())
        conv_in = torch.cat([cs, qkv.transpose(1, 2)], dim=2)
        new_conv = conv_in.narrow(2, conv_in.size(2) - (K - 1), K - 1)
        qkv = F.silu(self.conv1d(conv_in)).transpose(1, 2)
        q, k, v = torch.split(qkv, [KEY_DIM, KEY_DIM, VALUE_DIM], dim=-1)
        q = repeat_kv(q.reshape(b, s, NK, DK).permute(0, 2, 1, 3), NV // NK)
        k = repeat_kv(k.reshape(b, s, NK, DK).permute(0, 2, 1, 3), NV // NK)
        v = v.reshape(b, s, NV, DV).permute(0, 2, 1, 3)
        out, new_ssm = self.gdu(q, k, v, g.permute(0, 2, 1), beta.permute(0, 2, 1), ss)
        out = self.norm(out, z)
        out = self.out_proj(out.reshape(b, s, VALUE_DIM))
        return out, new_ssm, new_conv


class Attn(nn.Module):
    def __init__(self):
        super().__init__()
        self.q_proj = nn.Linear(HID, NH * HD, bias=False)
        self.k_proj = nn.Linear(HID, NKV * HD, bias=False)
        self.v_proj = nn.Linear(HID, NKV * HD, bias=False)
        self.o_proj = nn.Linear(NH * HD, HID, bias=False)
        self.q_norm = RMSNorm(HD, eps=1e-6)
        self.k_norm = RMSNorm(HD, eps=1e-6)
        self.rope = initialize_rope(dims=HD, base=1e5)
        self.sdpa = SDPA(is_causal=True)

    def forward(self, x, position_ids, cache):
        b, query_len, _ = x.shape
        sequence_length = position_ids.shape[-1]
        torch._check_is_size(sequence_length)
        torch._check(sequence_length <= cache._k_cache.size(KVCache.seq_len_dim()))
        offset = sequence_length - query_len
        torch._check_is_size(offset)
        q = self.q_proj(x).reshape(b, query_len, NH, HD).permute(0, 2, 1, 3)
        k = self.k_proj(x).reshape(b, query_len, NKV, HD).permute(0, 2, 1, 3)
        v = self.v_proj(x).reshape(b, query_len, NKV, HD).permute(0, 2, 1, 3)
        q = self.q_norm(q)
        k = self.k_norm(k)
        rp = position_ids.narrow(-1, offset, query_len)
        q = self.rope(q, position_ids=rp)
        k = self.rope(k, position_ids=rp)
        k, v = cache.update_and_fetch(0, offset, k, v, seq_len=sequence_length, query_len=query_len)
        k = repeat_kv(k, NH // NKV)
        v = repeat_kv(v, NH // NKV)
        cap = cache._k_cache.size(KVCache.seq_len_dim())
        for t in (k, v):
            torch._check(t.size(-2) >= query_len)
            torch._check(t.size(-2) <= cap)
        out = self.sdpa(q, k, v).permute(0, 2, 1, 3).reshape(b, query_len, NH * HD)
        return self.o_proj(out)


class Model(nn.Module):
    def __init__(self, ndelta):
        super().__init__()
        self.embed = nn.Embedding(VOCAB, HID)
        self.deltas = nn.ModuleList([DeltaNet() for _ in range(ndelta)])
        self.d_ln = nn.ModuleList([RMSNorm(HID, eps=1e-6) for _ in range(ndelta)])
        self.d_mln = nn.ModuleList([RMSNorm(HID, eps=1e-6) for _ in range(ndelta)])
        self.d_mlp = nn.ModuleList([MLP(HID, 4 * HID) for _ in range(ndelta)])
        self.attn = Attn()
        self.a_ln = RMSNorm(HID, eps=1e-6)
        self.a_mln = RMSNorm(HID, eps=1e-6)
        self.a_mlp = MLP(HID, 4 * HID)
        self.norm = RMSNorm(HID, eps=1e-6)
        self.lm_head = nn.Linear(HID, VOCAB, bias=False)

    def forward(self, input_ids, position_ids, k_cache, v_cache, conv_state, ssm_state):
        x = self.embed(input_ids)
        kv = KVCache(k_cache, v_cache)
        for i, layer in enumerate(self.deltas):
            ss = ssm_state.narrow(0, i, 1).squeeze(0)
            cs = conv_state.narrow(0, i, 1).squeeze(0)
            r, new_ssm, new_conv = layer(self.d_ln[i](x), ss, cs)
            x = x + r
            x = x + self.d_mlp[i](self.d_mln[i](x))
            _slice_write(ssm_state, i, new_ssm)
            _slice_write(conv_state, i, new_conv)
        x = x + self.attn(self.a_ln(x), position_ids, kv)
        x = x + self.a_mlp(self.a_mln(x))
        return self.lm_head(self.norm(x))


def main() -> int:
    torch.manual_seed(0)
    s = 12
    ctx = int(os.environ.get("CTX", "20"))   # context length; CTX=s -> offset 0 (clean first chunk)
    cap = 64
    model = Model(NDELTA).to(torch.float16).eval()
    inputs = {
        "input_ids": torch.randint(1, VOCAB, (1, s), dtype=torch.int32),
        "position_ids": torch.arange(ctx, dtype=torch.int32).unsqueeze(0),
        "k_cache": torch.zeros(1, 1, NKV, cap, HD, dtype=torch.float16),
        "v_cache": torch.zeros(1, 1, NKV, cap, HD, dtype=torch.float16),
        "conv_state": torch.zeros(NDELTA, 1, CONV_DIM, K - 1, dtype=torch.float16),
        "ssm_state": torch.zeros(NDELTA, 1, NV, DK, DV, dtype=torch.float16),
    }
    with torch.no_grad():
        ref = model(**{k: v.clone() for k, v in inputs.items()})
    print(f"[NDELTA={NDELTA}] torch ref {tuple(ref.shape)}", flush=True)

    state_names = ("k_cache", "v_cache", "conv_state", "ssm_state")
    input_names = ("input_ids", "position_ids")
    # STATIC=1 exports a fully fixed-shape function (no dynamic dims at all). The crash is
    # triggered by the dynamic context dim, so a fixed query + fixed context window runs fine.
    if os.environ.get("STATIC", "0") == "1":
        dynamic_shapes = None
    else:
        dynamic_shapes = {
            "input_ids": None,                                              # STATIC query
            "position_ids": {1: torch.export.Dim("ctx", min=s + 1, max=cap)},
            "k_cache": {3: torch.export.Dim("kseq", min=s + 1, max=cap)},
            "v_cache": {3: torch.export.Dim("vseq", min=s + 1, max=cap)},
            "conv_state": None,
            "ssm_state": None,
        }
    prog = export_to_coreai(model, {k: v.clone() for k, v in inputs.items()},
                            input_names=input_names, output_names=("logits",),
                            state_names=state_names, dynamic_shapes=dynamic_shapes)
    prog.optimize()
    asset = Path(tempfile.mkdtemp(prefix="toy_hybrid_")) / "m.aimodel"
    prog.save_asset(asset)
    print(f"[NDELTA={NDELTA}] EXPORT+SAVE OK", flush=True)

    async def run():
        m = await AIModel.load(asset)
        print("LOADED", flush=True)
        fn = m.load_function(m.function_names[0])
        ins = {n: NDArray(inputs[n].clone()) for n in input_names}
        st = {n: NDArray(inputs[n].clone()) for n in state_names}
        print("EXECUTING...", flush=True)
        r = await fn(inputs=ins, state=st)
        print("EXECUTED", flush=True)
        return r

    res = asyncio.run(run())
    got = res["logits"].numpy().astype(np.float32)
    want = ref.detach().numpy().astype(np.float32)
    rel = np.abs(got - want).max() / (np.abs(want).max() + 1e-6)
    print(f"[NDELTA={NDELTA}] RAN — rel err {rel:.4f}  {'GO' if rel < 0.06 else 'NUMERIC-OFF'}")
    return 0


if __name__ == "__main__":
    sys.exit(main())