magpie: §6.3 host-owned-cache decoder_step — ANE-safe + ~40% faster/step

[Alex-Wengg]

Follow-up to the top open item in models/tts/magpie/SWIFT_PORT_FINDINGS.md ("try rank-3 K/V + int32 position to unblock decoder_step on ANE"). Applies the Surgical Inference §6.3 boundary to the autoregressive decoder_step.

The change

The current decoder_step does two ANE-hostile things per the paper: an in-graph full-cache read-modify-write blend (new_k = kv_k*(1-mask) + k_new*mask) and runtime-position-driven mask compares (positions_range == position), ×12 layers. The §6.3 rewrite (HostCacheDecoderStep):

• reads the past cache read-only, concatenates the current-step K/V, attends against a host-supplied additive mask, and outputs only the current-step K/V slice [1,1,H,D];
• the host owns the cache append and the mask — no in-graph state write, no position compares.

This mirrors the falsification-ladder rung that compiled ANE-clean in the paper: cache reads + current-token update outputs are admissible; the in-graph state write is the cliff.

Results (M5 Pro, coremltools 9.0; random weights — ANE admission is structural)

Equivalence (torch fp32): max |Δlogits| across 8 decode steps = 9.5e-7 → pure restructuring, not a new model.

Graph	CPU_AND_NE p50	Placement (device-assigned ops)
OLD (in-graph blend + pos compares)	10.2 ms	100% ANE (752/752)
NEW (§6.3 host-owned cache)	6.0 ms	100% ANE (555/555)

1. The documented M2 ANECCompile -14 does not reproduce on M5 Pro — the current graph already runs 100% on the ANE here through the full incrementing-position loop. The block was older-compiler-specific.
2. The §6.3 rewrite is ~40% faster/step anyway, still 100% ANE, fewer ops, and far smaller I/O (current-slice outputs vs the [1,512,12,64]×24 full-cache re-emit the old graph does every step — the ~19 MB cache-out PERF.md flags). Decode runs 50–200×/utterance, so this is a direct RTFx lever.

Scope

Proof-of-concept under experiments/ane_host_cache/ (no NeMo / gated download): proves ANE admission, placement, latency, and torch equivalence — not CoreML fp16 end-to-end parity with real weights or the Swift host-side plumbing. Productionizing = port the boundary into traceable_decoder_step.py/convert_decoder_step.py with real weights, move the cache append + mask into MagpieModelStore.swift, and validate audio parity.

Build artifacts + logs gitignored.