Zaya 8B Q4NX on the AMD NPU

The lane

HRX is AMD’s IREE-derived runtime — Loom source-first kernels JIT-compiled for the XDNA 2 NPU. We spent this month turning it from a platform-transition watch into an actual execution lane for our own models: the Q4NX tile format (4-bit int4 + BF16 scales, 5120-byte tiles) that our engine uses natively, executed on the NPU.

The final model: Zyphra ZAYA1-8B, 8.84 B params, 6.96 GiB (6.76 BPW). Every quantized op runs on the HRX2 backend of our llama.cpp fork at -ngl 99.

Milestones

milestone	what it took	result
First Q4NX model served	qwen3-0.6B end-to-end on gfx1151; every layer of the chain validated (op-level 3.7e-7)	pipeline proven
Quality fix	the fused mm kernel silently transposed every multi-token matmul (row-major vs ggml col-major)	real tokens
Zaya port — architecture root cause	the graph ran alternating attn/MoE blocks and ignored half the weights; rewritten to HF’s both-blocks-per-layer	8/8 top-1 vs HF
Zaya Q4NX + MoE dispatch	new MUL_MAT_ID_Q4NX op for the stacked experts; root cause: the dispatch read the strided topk view contiguously	corr 0.99999999
7–17× speedup	per-tile 3-copy upload (≈9,200 stream ops/token) collapsed to one raw-blob copy + a tile-major dequant kernel	tg 1.22 → 8.7 t/s
llama-server	the HTTP lane: prompt 18.1 / decode 8.16 t/s	servable

The numbers

measure	value
logits corr vs F32 port	0.99999999
max |dlogit|	0.0041 — top-1 margin is 1,681× the noise
top-5 tokens	identical to F32
F32 port vs HF greedy	8/8 top-1
prefill (pp32)	127.4 t/s (round 23)
decode (tg32)	13.0 t/s (round 23)
decode, pre-optimization	1.22 t/s (≈11× faster now)
prefill, pre-optimization	2.1 t/s (≈61× faster now)
llama-server prompt / decode	18.1 / 8.16 t/s

The bugs that were the work

The two root causes worth remembering:

The architecture was wrong before the math was. The Zaya port ran one block per layer (attention, then MoE, alternating) — HF runs both in every layer. Half the weights sat unused and the logits were uncorrelated with the reference (0.10). Fixing the graph — not the kernels — jumped straight to 8/8 top-1.

The ids are a strided view. In the expert matmul, the per-token expert ids come from an argsort whose output is a strided view (17 sorted experts per token, nb[1] = 17×4). The NPU dispatch copied them contiguously, so token 2 silently received the second-best expert of token 0. The F32 path read the strides correctly — which is exactly why the port “worked” on CPU and broke on the NPU. One stride-aware copy fixed the whole model.

Why 0.99999999 and not 1.0

Because the two runs sum the same math in different orders. The weights are bit-identical between the Q4NX and F32 models; the last 1e-8 of correlation is float32 rounding between the NPU reduction tree and the CPU one — 1,681× below the top-1 margin. It rounds to 1.00000000 and the output is identical.

What’s next (rounds 24–28)

The converter went general. make-q4nx-model.py now reads the architecture from any GGUF, quantizes Q4_K/Q5_K/Q6_K/Q8_0/F16/BF16 sources to Q4NX (numpy dequant ports verified bit-exact vs ggml), handles MoE 3-D expert tensors and per-type keep shapes. The whole local roster was converted and validated on HRX20: Qwen3-0.6B, Qwen2.5-3B/0.5B, MiniCPM5-1B, MiniCPM4-8B generate coherently; the two MoE models (Qwen3-Coder-30B-A3B, GLM-4.7-Flash) now run their expert matmuls on the NPU; Qwen3-Next-80B runs too — it needed the pointwise ncols caps raised to 1M (the recurrent-state SCALE builds 524288-wide views in the worst-case reserve graph), and now says “Paris.” like its Q4_K reference.

The MoE bug that looked like noise. The first MoE runs were complete garbage while every sub-component verified (per-pair mm corr 0.998, down mm corr 1.0000). Layer 0 was fine but layer 1’s input had already drifted to corr 0.63. Root cause: MUL_MAT_ID treated src1 as the shared form [k, ntokens] (gate/up experts) for every dispatch — the down experts feed a per-expert src1 [k, n_expert_used, ntokens], and the old code always read column t·k: every expert silently computed with expert-0’s activation. The “corr 1.0000” down check passed precisely because it replicated the bug. Fixed (fork 49f07b6): per-slot src1 columns (i + t·n_expert_used) through both the grouped tables and the per-pair path, plus a new src1_cols_count config in the two table-scatter kernels. Down MMID now corr 1.0000 vs numpy (L0/L1/L10/L47); the 30B Q4NX matches the Q4_K_M reference nearly token-for-token — “The capital of France is” → “Paris.” for both.

Numbers on the fixed 30B MoE (HRX20, -t 16, ngl 99): pp32 49.2 t/s, pp128 67.5 t/s, tg32 14.3 t/s. An external iGPU reference on the same SoC (Q4_K_M, Vulkan) decodes the 30B-A3B at 86.1 t/s — the iGPU stays ahead at small-batch decode; the NPU’s edge remains memory (6.96 vs 32.93 GiB) and large-batch prefill (zaya pp512 171 t/s).

Round 25d found a silent prefill regression: the per-expert src1 fix added a required src1_cols_count config to the fused tbl kernel, but the dense prefill path never passed it — every dense matmul failed the JIT compile and quietly fell back to the slow dequant+mm slice. One config line restored it: pp32 3B 82 → 121.5 t/s (+47%), 0.6B 208 → 435.8 (2.1×), MiniCPM4-8B 33 → 51. A provider_compile failed trace is what exposes this class of silent fallback.

Final roster (Q4NX on HRX20, sequential — concurrent benches fight over the NPU):

model	pp32	tg32	tg32 (25i zero-copy)
Qwen3-0.6B	435.8	42.3	63.0
Qwen2.5-3B	121.5	23.3	26.7
Qwen2.5-7B	55.2	11.7	13.7
MiniCPM4-8B	51.1	10.4	11.9
GLM-4.7 (30B-A3B MoE)	49.9	12.95	14.8
Qwen3-Coder-30B (MoE)	49–50	14.3–15.5	18.3
Qwen3-Next-80B (MoE)	32.2	7.2	7.9

Round 25h landed the deferred 16-row decode kernel: a workgroup of 16 rows × 1 col with 256 lanes split 16×16 over rows × k-blocks, so each k-step reads 8 contiguous packed bytes instead of the per-pair kernel’s 1-byte stride-8 reads. tg32: 3B 20.6→23.3 (+13%), 8B 9.0→10.4 (+16%), GLM 12.5→12.95; decode stays bandwidth-bound at ~49 GB/s effective.

Round 25i — TRUE zero-DMA-copy. The CPU↔NPU boundary is now copy-free: activations live in host-visible GTT shared by both sides, weights stay device-local (the loader is forced onto a device-local extra buft), and only ~128 bytes of graph input cross per eval. Full-hybrid op claims plus generic JIT rms_norm/soft_max routes keep the CPU out of the shared-memory hot path entirely — if the CPU computes into GTT, the NPU reads those pages slowly (cache-coherency tax over the fabric), so “CPU never touches shared memory” is the design rule. Decode wins across the roster (+10–49%, right column); prefill (pp32) regressed on this NPU (3B 121.5→27.8) because the attention KQ¹/kqv F32 matmuls are still CPU-side and write into shared GTT — the documented cost of shared-memory zero-copy until those move to HRX2 (strided GQA batched kernel, scaffolding already in the fork).

Round 25j — the attention matmuls moved to HRX2 and the prefill coherency tax is gone (pp32 up to 3.4×). The remaining CPU ops in the zero-copy prefill path were the attention KQ¹/kqv matmuls plus ROPE and GLU — every one of them wrote into shared GTT and the NPU paid cross-snoop coherency on read-back (3B pp32 27.8). The matmuls turned out to be F16×F32 (K/V come from the F16 KV cache), and an existing complete kernel already handled exactly that (mul_mat_f16_f32_batched: strided GQA batched views, broadcast) — it just was never claimed. Claiming it exposed a real loom bug: the src0 view bound omitted the (src0_ne2-1)·stride_ne2 term, so with GQA the JIT correctly rejected the shape (“view.load scalar footprint upper bound is not proven”, kqv 32768 vs needed 65536 elements). Fixed in all four kernel variants. The dispatch trace then showed the matmuls were only ~10 ms of 2.9 s — the real tax was uncovered ROPE (K-rope h2 had no routes at all; Q-rope t8/t32 missing) and split GLU (ncols=11008 uncovered). Added generic rope_neox_f32/freq and swiglu_f32_split routes; only 12 embedding lookups stay on CPU. The cols8 batched-mm route domain was widened (cols 8–512) so wide-head models (30B, 32 heads) get the 8-cols-per-workgroup kernel instead of the 1-output-per-workgroup generic (30B pp32 41.8 → 47.1, matching the CPU-attention baseline while staying on NPU), and the stride config decls were widened to 2³⁰ so llama-cli decode shapes JIT-compile. Result: the 3B prefill regression is fully recovered (27.8 → 81.3, vs 121.5 pre-zero-copy) and decode is still up everywhere.

model	pp32 (25i)	pp32 (25j)	tg32 (25i)	tg32 (25j)
Qwen2.5-0.5B	180	617	84.5	80.5
Qwen3-0.6B	211	461	63.0	81.9
Qwen2.5-3B	27.8	81.3	26.7	29.5
MiniCPM5-1B	—	325	—	69.7
Qwen2.5-7B	29.9	30.3	13.7	14.1
GLM-4.7 (30B-A3B MoE)	39.9	43.1	14.8	15.7
MiniCPM4-8B	24.3	24.4	11.9	12.2
Qwen3-Coder-30B (MoE)	47.7	47.1	18.3	18.8
Qwen3-Next-80B (MoE)	32.6	32.7	7.9	8.0

Round 25n — generic norm fusions: prefill is launch-bound, pp32 +33%. Decode-trace forensics corrected a wrong model: the per-dispatch elapsed_us is only host submit time (~1 µs) — the real GPU execution hides inside hrx_stream_synchronize. Fitting time/token vs model size across 0.5B/3B/7B gives decode = 3.7 ms fixed + weights @ 65 GB/s: decode is weight-bandwidth-bound, so kernel count is irrelevant there. Prefill is the opposite — launch-bound at ~31 µs per dispatch — so cutting dispatch count is the lever. The ADD→RMS_NORM→MUL and RMS_NORM→MUL fusion routes existed but never fired on the 3B: they only covered ncols=3072/4096 while n_embd=2048 was uncovered (same route-coverage class as 25j’s ROPE/GLU). Added generic rms_norm_mul + add_rms_norm_mul routes (ncols 1–65536); pp32 dispatches per graph ~2600 → 2028. Two gotchas: the kernels need the vector_width=4 tuning binding (my first generic append omitted it → JIT CONFIG/INVALID), and the route JSONs are hand-formatted with compact arrays — append text-level in the original style, json.dump reformats the whole file.

model	pp32 (25j)	pp32 (25n)	tg32 (25j)	tg32 (25n)
Qwen2.5-0.5B	617	652	80.5	89.7
Qwen3-0.6B	461	523	81.9	85.2
Qwen2.5-3B	81.3	108.3	29.5	30.1
MiniCPM5-1B	325	—	69.7	—
Qwen2.5-7B	30.3	38.5	14.1	14.2
GLM-4.7 (30B-A3B MoE)	43.1	48.9	15.7	16.4
MiniCPM4-8B	24.4	31.9	12.2	12.1
Qwen3-Coder-30B (MoE)	47.1	54.3	18.8	19.7
Qwen3-Next-80B (MoE)	32.7	—	8.0	—

Round 25o — the r16x8 fused prefill matmul: 16 rows × 8 cols per workgroup — pp32 roughly doubled. Prefill-mm forensics found the tbl_tiled kernel does 1 row × 8 cols per workgroup: for the 3B gate mm (11008 rows) that is 44,032 workgroups re-reading scales/zeros per k element at 1-row granularity. Workgroup-size A/B (256/512/1024 → 108/95/58 t/s) and routing prefill through the decode r16 kernel (57 t/s — each column workgroup re-reads all weights) ruled out launch count and pinned the limit on the 1-row-per-workgroup structure. The new hrx2_mul_mat_q4nx_fused_f32_r16x8 kernel adapts the r16 decode kernel (16 rows/workgroup, 256 lanes = 16 rows × 16 k-blocks, 8-byte coalesced packed reads) to process 8 output columns per workgroup: each lane walks its row’s k-chunk once and feeds 8 matmuls — one packed-byte stream, 8 src1 loads + 8 MACs per k — so workgroups drop from rows × cols/8 to rows/16 × cols/8 (gate: 44,032 → 2,752). Per-row reduction over the 16 k-block lanes uses workgroup scratch (phase 2: lanes 0..127 sum 16 kb-partials each). Dense identity columns only (MoE keeps the table-scatter path). Dispatch requires cols % 8 == 0 and rows % 16 == 0 — the column guard is essential: indexing src1 at col_group·8+7 reads past the view on non-multiple-of-8 shapes (llama-cli’s 33-token batch JIT-rejected them until the guard landed).

model	pp32 (25n)	pp32 (25o)	tg32 (25n)	tg32 (25o)
Qwen2.5-0.5B	652	835	89.7	82.7
Qwen3-0.6B	523	786	85.2	84.6
Qwen2.5-3B	108	229	30.1	29.4
Qwen2.5-7B	38.5	98.0	14.2	14.2
GLM-4.7 (30B-A3B MoE)	48.9	53.3	16.4	16.2
MiniCPM4-8B	31.9	78.5	12.2	12.1
Qwen3-Coder-30B (MoE)	54.3	54.2	19.7	18.0
Qwen3-Next-80B (MoE)	—	—	—	—

Round 25p — the MoE boulder moved: the r16x8t table-scatter kernel, grouped prefill +21–33%. Round 25o tripled dense prefill but the MoE models barely moved — GLM-4.7 and Qwen3-Coder-30B route their expert matmuls through the grouped path (slice_grouped), which still used the 1-row-per-workgroup tbl_tiled kernel with i32 table scatter. GLM’s pp32 trace shows why: 9,453 grouped dispatches per graph (vs 652 dense), each ~80 µs serialized — groups are small (rows 1536/2048, cols 2..7), so every group launches rows× workgroups to compute a handful of columns. The new hrx2_mul_mat_q4nx_fused_f32_r16x8t kernel gives the fused tbl the r16x8 structure — 16 rows × 8 columns per workgroup, 256 lanes as 16 rows × 16 k-blocks, one packed-byte stream feeding 8 table-scattered matmuls, per-row scratch reduction — so workgroups drop from rows× to rows/16×. Two real bugs surfaced: the phase-2 reduction indexed its 16 k-block partials kb-major while lanes are row-major (→ garbage on GLM until fixed to row_reduce + kk2·16), and the route-binding validator whitelists only shape.mul_mat_id.* sources while src1_cols_count must ride the dispatch config, not a binding. Result: MoE grouped prefill 30B 54.2→71.9 (+33%), GLM 53.3→64.7 (+21%); dense untouched (r16x8 still serves it) and decode flat as designed. Greedy generations verified identical vs the old path on GLM and 30B (multi-prompt batteries); the 30B poem-prompt degeneration reproduces on the pre-25p build too — a pre-existing chaotic divergence, not a regression.

model	pp32 (25o)	pp32 (25p)	tg32 (25o)	tg32 (25p)
Qwen2.5-3B	229	229.7	29.4	29.8
Qwen2.5-7B	98.0	97.6	14.2	14.1
GLM-4.7 (30B-A3B MoE)	53.3	64.7	16.2	16.3
Qwen3-Coder-30B (MoE)	54.2	71.9	18.0	19.7
MiniCPM4-8B	78.5	80.0	12.1	12.2

Round 25q — the decode “65 GB/s wall” was never the memory path: the r16w wide-read kernel, tg32 3B +72%, 7B +89%. The standing decode model said time/token = 3.7 ms fixed + weights @ 65 GB/s, and 25h guessed the ceiling was “the NPU’s actual memory path”. Round 25q falsified that with a load-only probe: dispatching the dense f32 mm kernel on the same device-local weight buffer reads at 205–219 GB/s cold (8-slice rotating ring), while the q4nx r16 decode kernel reads the same memory at only 64–74 GB/s cold — the gap was the access pattern, not the fabric (~3× headroom existed). Compiler evidence matched: r16’s inner body issues six 1–4 byte scalar loads per lane per k-step; the f32 kernel issues 4B/lane contiguous (~1 KB per instruction). The new hrx2_mul_mat_q4nx_fused_f32_r16w kernel restructures decode around the packed layout’s contiguous axis — 8 bytes at one k-position hold all 16 rows of a half-tile — so each of the 256 lanes owns ONE k-position and ALL 16 output rows of the workgroup: per tile-column the wavefront reads 2 KB contiguous (256 lanes × 8B). 16 f32 row accumulators per lane; scales/zeros staged cooperatively into 512 B of workgroup LDS once per tile-column (cutting 64 × 1B global scale loads to 2 per lane); 16 workgroup.reduces produce the outputs; signed-nibble decode is branchless. One real bug: the dispatch’s workgroup-count expression special-cased r16 (rows/16) but fell through to rows for r16w — 16× too many workgroups reading past their slice (memory fault, caught by the real-model A/B the probe could not see).

model	tg32 (r16)	tg32 (r16w)
Qwen2.5-3B	29.8	51.2 (+72%)
Qwen2.5-7B	14.0	26.5 (+89%)
MiniCPM4-8B	12.1	21.7 (+79%)
GLM-4.7 (30B-A3B MoE)	16.4	19.9 (+21%)
Qwen3-Coder-30B (MoE)	19.9	25.7 (+29%)

Round 26 — the twin: the gap is summation order, now proven. The “Why 0.99999999 and not 1.0” claim above was asserted; round 26 proved it. We built a lossless F32 twin of the Q4NX model: every Q4NX tensor dequantized to F32 holding the same numeric values (W = q4·scale + zp, per-(row, 32-col) bf16 scales/zeros, 35.4 GB output). If the NPU executes the quantized graph exactly, corr(Q4NX, twin) ≈ 1.0; if the gap persists against the twin, it is execution — not quantization. The twin-vs-F32 comparison (corr 0.999999, maxdiff 4.4e-2) confirms the dequant is faithful, yet the NPU-vs-twin gap (corr 0.999973, maxdiff 0.24) is statistically identical to the NPU-vs-F32 gap. Top-1 and top-10 are identical on all three. The residual is pure f32 summation order between the NPU’s 256-lane reduction trees and the CPU order, amplified by the 32-token recurrence — not the 4-bit weights. Deterministic corr(32-tok) = 1.0 remains the fixed-point pipeline’s property (the 1BP reader reconstructs W = q4·s + zp byte-exact: 0/3,145,728 mismatches vs the engine’s own dequant; the (q4·16·rq + zpq) >> 22 re-quant matches int8 to rounding boundaries only).

Round 27 — the real model on the NPU. The actual ROCm/FastFlowLM runtime (Qwen3-0.6B, the same Q4NX tile format) now runs end-to-end on the XDNA 2 NPU: prompt “Hello” → 146-token reply, prefill 27.8→29.6 t/s, decode 96.6→92.8 t/s, zero IO_PAGE_FAULTs and zero TDRs. Two things made the hand-rolled launcher deadlock while the real runtime works: (1) the runtime generates the layer instruction stream with the real weights loaded and sizes the BOs to match — our weight-less dump and guessed BOs meant the shim-DMA sync tokens never fired; (2) the amdxdna kernel ABI is (opcode, instr_bo, ninstr, bo0..boN) — without the instruction buffer the ERT command reports COMPLETED but the AIE never executes (a silent no-op, proven with sentinel buffers). Feeding the insts and binding each BO to its kernel-argument slot via the BO group id makes the same xclbin execute deterministically, input- and weight-offset-sensitively. We also reverse-engineered the tile dequant against the runtime’s own library: W = (q − zp) · scale (group-major bf16 scales/zeros, lane-swizzled nibbles) — maxdiff 0.0 over the whole projection, the only residual one BF16 ULP from the weight BO.

Round 28 — the engine matched the real runtime byte-for-byte: corr 1.000000. The hand-rolled NPU launcher’s prefill logits are now byte-identical to the real FastFlowLM runtime on the same xclbin: for the token-1000 input, corr 1.000000, argmax 397 == 397, top-5 identical [397, 3219, 144370, 42044, 255], verified by dumping the engine’s prefill logits (NPU_LOGITS_DUMP) and diffing against the runtime’s logits_1000.bin. Decode runs at ~15 ms/tok on Qwen3-0.6B/XDNA 2. This closes the “why 0.99999999 and not 1.0” question from round 26 in the strongest way: the engine executes the same instruction stream as AMD’s own runtime and produces the same bytes — the residual corr gap that remained in the llama.cpp lane is summation order between reduction trees, not the Q4NX math. The fix trail that got here: the amdxdna ABI needs the (opcode, instr_bo, ninstr, bo0..boN) instruction buffer (without it the ERT reports COMPLETED but the AIE never executes — a silent no-op proven with sentinel buffers), each host BO must be created with its kernel-argument group id (group-0 BOs are ignored and can wedge the NPU with IO_PAGE_FAULTs), and the Q4NX dequant is W = (q − zp) · scale (group-major bf16 scales/zeros, lane-swizzled nibbles) — maxdiff 0.0 over the whole projection.

One pool, one API call, zero copies

We reverse-engineered how the Windows NPU stack does memory — then matched it on Linux. The control plane was already unified (one process, one OpenAI-compatible API, every model resident); what was missing was the Windows-style memory + dispatch model underneath. Windows (MCDM, ipustack.sys) hands every process one device heap and submits with one chained command message (CHAIN_EXEC_DPU) — a front-end to the same NPU firmware as Linux amdxdna, opcode tables matching. Portable, not proprietary: the Linux kernel already had the device-heap carve (amdxdna_drm_create_dev_heap) and the chained EXEC_CMD; the gap was userspace. We closed it — one 64 MB heap (every tensor a slice carved by the kernel’s own drm_mm) and one EXEC_CMD per layer carrying the whole command chain (QKV, O, GU, D) as a single mailbox message. Probe-verified end-to-end on kernel 7.1.5; the real-kernel path is built against Qwen3-0.6B’s actual instruction stream and q4nx weights.

And the pages are shared, not copied. Models sit in a unified memory pool (SharedBO) — NPU-owned, GPU-imported, host-coherent: one set of physical pages, three views, no staging DMA. Proven both directions (dma-buf round trip, 0 mismatches), and in production form by a Vulkan shader reading the NPU’s KV pages straight out of the pool (rel err 2e-4). The rule that makes it safe: the NPU owns the pages — reverse the ownership and AMD-Vi faults. GPU import is Vulkan (VK_KHR_external_memory_fd); the installed TheRock HIP has no dma-buf import path, so that route is the one that works.