docs: add megakernel-rmsnorm-qmv.md -- setup, testing, and full design history

Reproducibility pass per Brian's request: someone with no prior context
should be able to build, enable, correctness-test, and benchmark this
branch from this doc alone, without reading commit messages.

Covers: build instructions, the embedded-metallib build-cache gotcha that
cost real time in this session (cmake sometimes silently skips
recompiling a .metal-only edit), exact correctness-check and benchmark
commands (verified to work verbatim before committing), the debug env var
table, the structural constraint on contiguous fusion (with the exact
verification method to re-check it on a different model), the full
measured design history (v1 through the tiny-ne01 fix, each with its real
number and root cause -- not just "we tried some things"), and an honest
"known limitations" section (parity is not yet a demonstrated speedup,
untested under concurrent serving load).

Author: bri-prism

docs/megakernel-rmsnorm-qmv.md

@@ -0,0 +1,235 @@
+# RMSNorm + Q1_0/Q2_0 mat-vec fusion (Metal, batch=1 decode)
+
+Branch: `megakernel/rmsnorm-qmv-fuse`. Status: **parity with baseline, opt-in, not default.**
+
+This fuses ggml-metal's existing two-step pattern —
+`kernel_rms_norm_mul_f32` (writes `normed = rmsnorm(x)*w` to device memory) followed by one or
+more separate `kernel_mul_mv_{q1_0,q2_0}_f32` dispatches that read it back — into a single
+fusion that never materializes the normed vector in device memory. It only applies at batch=1
+(single decode token); prefill/batched matmul is untouched.
+
+Built for Bonsai-27B (Qwen3-Next hybrid architecture) on Apple M-series GPUs. Read this whole
+file before touching the code — the design went through several measured, wrong-then-fixed
+iterations, and the reasons matter for anyone extending it.
+
+## TL;DR for someone picking this up
+
+1. Build normally (`cmake -B build -DGGML_METAL=ON -DCMAKE_BUILD_TYPE=Release && cmake --build
+   build --target llama-bench llama-simple`). The fused kernels compile in unconditionally;
+   they just don't run unless you opt in.
+2. Set `GGML_METAL_RMSNORM_QMV_FUSE=1` to enable the fusion. Unset (or `=0`) is the default —
+   baseline behavior, byte-for-byte.
+3. Run `llama-bench` with and without the env var set to A/B it. As of the last commit on this
+   branch, the two are statistically indistinguishable (see "Current status" below) — this is
+   NOT yet a demonstrated speedup, just no longer a regression.
+4. If you change anything in `ggml-metal.metal`, see "Gotcha: the embedded-metallib build cache"
+   below before you trust a `cmake --build` that reports no work to do.
+
+## Why this exists
+
+Real per-decode-step tracing on this model found that RMSNorm's output almost never feeds
+exactly one matmul — it typically feeds 2, 3, or 4 (gate+up MLP share one norm, Q/K/V/gate-style
+GDN-input projections share another). A naive single-consumer fusion (v0, not on this branch in
+its original form) fires essentially never on this architecture. The shipped design generalizes
+to N-consumer fusion (N ∈ {2,3,4}, matching the real, live-traced fan-outs) — see
+`ggml_metal_op_rmsnorm_qmv_multi_try` in `ggml/src/ggml-metal/ggml-metal-ops.cpp`.
+
+## What's actually in this branch
+
+- `ggml/src/ggml-metal/ggml-metal.metal`:
+  - `kernel_rmsnorm_scale_f32` — single-threadgroup kernel, computes the RMSNorm scale once,
+    writes it to a 1-float scratch slot.
+  - `kernel_rmsnorm_mv_{q1_0,q2_0}_f32` — single-consumer fused kernel (kept for completeness;
+    **does not fire on Bonsai-27B**, see below — every RMSNorm here has ≥2 consumers).
+  - `kernel_rmsnorm_mv{2,3,4}_{q1_0,q2_0}_f32` and the `_small` variants — the N-consumer fused
+    kernels that actually fire, in two NR0 tunings (see "Tiny-ne01 corner case" below).
+- `ggml/src/ggml-metal/ggml-metal-ops.cpp`:
+  - `ggml_metal_op_rmsnorm_qmv_try` — single-consumer detector/dispatcher (dead code on this
+    model, kept because it's cheap to check and harmless).
+  - `ggml_metal_op_rmsnorm_qmv_multi_try` — the real one. Pattern-matches
+    `RMS_NORM -> MUL -> {MUL_MAT(q1_0|q2_0)} × N` at graph-build time, dispatches the fused
+    kernels. Returns `0` (decline to fuse, caller falls back to the normal path) on any
+    mismatch — this is the safe path, not an error path.
+  - `has_single_use()` / `use_count()` — new methods added to the `ggml_metal_op` struct for
+    this work.
+- `ggml/src/ggml-metal/ggml-metal-common.cpp`: added `GGML_OP_GATED_DELTA_NET` and
+  `GGML_OP_FLASH_ATTN_EXT` to the graph-reorder pass's `h_safe()` allowlist (a separate,
+  independently-useful scheduling fix — see "Scheduling fix" below).
+- `docs/megakernel-rmsnorm-qmv.md` — this file.
+
+## How to build
+
+Nothing special — the fused kernels are always compiled in, gated at runtime by an env var, not
+a build flag.
+
+```sh
+cmake -B build -DGGML_METAL=ON -DCMAKE_BUILD_TYPE=Release \
+  -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=ON -DLLAMA_BUILD_SERVER=OFF
+cmake --build build --target llama-simple llama-bench -j 8
+```
+
+### Gotcha: the embedded-metallib build cache
+
+`ggml-metal.metal`'s source gets embedded into the binary as text and JIT-compiled at runtime
+(you'll see `ggml_metal_library_compile_pipeline: compiling pipeline: ...` in stderr on first
+use of each kernel — this is normal, not an error). CMake's dependency tracking on **pure
+`.metal`-file-only edits** was unreliable in testing on this box — a `cmake --build` sometimes
+reports `Built target ggml-metal` with no compile step even after a real source change. If your
+edit doesn't seem to take effect, force it:
+
+```sh
+touch ggml/src/ggml-metal/ggml-metal.metal
+rm -f build/ggml/src/ggml-metal/autogenerated/ggml-metal-embed.s \
+      build/ggml/src/ggml-metal/autogenerated/ggml-metal-embed.metal.tmp \
+      build/ggml/src/ggml-metal/autogenerated/ggml-metal-embed.metal \
+      build/ggml/src/ggml-metal/CMakeFiles/ggml-metal.dir/autogenerated/ggml-metal-embed.s.o
+cmake --build build --target ggml-metal llama-simple llama-bench -j 8
+```
+
+If you see `[50%] Building ASM object .../ggml-metal-embed.s.o` in the output, your edit was
+picked up. If you don't see that line, it wasn't — force it as above.
+
+You can sanity-check a `.metal` edit compiles at all, independent of the full build, with:
+
+```sh
+xcrun metal -S -O3 -I ggml/src/ggml-metal -I ggml/include -I ggml/src \
+  ggml/src/ggml-metal/ggml-metal.metal -o /tmp/check.air
+```
+
+This only checks syntax/type-correctness (AIR codegen), not runtime correctness — always follow
+it with a real build + the correctness check below.
+
+## How to test correctness
+
+Greedy output must be byte-identical with the fusion on vs off, on every checkpoint you care
+about. This is the only correctness bar that matters here — the kernel is designed so that any
+bug can only produce a wrong dot product (caught immediately by this diff), not a crash or a
+plausible-looking-but-wrong output that slips through.
+
+```sh
+MODEL=~/models/bonsai-27B/bonsai-27B-q1_0.gguf   # binary; swap for a Q2_0 (ternary) checkpoint too
+
+./build/bin/llama-simple -m "$MODEL" -n 64 -ngl 999 "The capital of France is" 2>/dev/null > /tmp/off.txt
+GGML_METAL_RMSNORM_QMV_FUSE=1 ./build/bin/llama-simple -m "$MODEL" -n 64 -ngl 999 "The capital of France is" 2>/dev/null > /tmp/on.txt
+diff /tmp/off.txt /tmp/on.txt && echo "BYTE-IDENTICAL"
+```
+
+Note: a raw-text (non-chat-templated) prompt on an instruct model like this often degenerates
+into a repeating token ("the the the the...") — that's a known, unrelated prompting artifact,
+not a bug in this fusion. It's still a valid correctness check (both runs must produce the
+exact same degenerate output), just don't mistake the degeneracy itself for a problem.
+
+Confirm the fusion is actually firing (easy to silently get 0 matches if a future model/arch
+change breaks the pattern match):
+
+```sh
+GGML_METAL_RMSNORM_QMV_FUSE=1 GGML_METAL_RMSNORM_QMV_TRACE=1 \
+  ./build/bin/llama-simple -m "$MODEL" -n 8 -ngl 999 "Hi" 2>&1 | grep -c MATCHED
+```
+
+On Bonsai-27B this should print `126` per forward pass (64 layers × ~2 fusion sites/layer). If
+you see `0`, or a much smaller number, something about the graph structure changed — see
+`GGML_METAL_RMSNORM_QMV_TRACE`'s `REJECT: <reason>` lines (below) to find out why.
+
+## How to benchmark
+
+```sh
+./build/bin/llama-bench -m "$MODEL" -n 128 -ngl 999 -r 5              # fusion off
+GGML_METAL_RMSNORM_QMV_FUSE=1 ./build/bin/llama-bench -m "$MODEL" -n 128 -ngl 999 -r 5   # fusion on
+```
+
+Compare `tg128` between the two. Run it 3+ times each way before drawing any conclusion — this
+box's baseline tok/s drifts several percent run-to-run from background system load (confirmed by
+watching `pp512`, which never touches this code path, drift by the same amount as `tg128`
+between unrelated runs). A single A/B pair is not a result.
+
+## Debug / diagnostic env vars
+
+| Var | Effect |
+|---|---|
+| `GGML_METAL_RMSNORM_QMV_FUSE=1` | Enable the fusion (default off). |
+| `GGML_METAL_RMSNORM_QMV_TRACE=1` | Print one line per RMS_NORM node: `MATCHED n=<N> qtype=<...> ne01=[...]` if fused, `REJECT: <reason>` if not (`op mismatch`, `multi-use`, `weight not q1_0/q2_0`, `batch != 1 or broadcast dims`, etc). Start here if match count looks wrong. |
+| `GGML_METAL_FUSION_DEBUG=2` | ggml-metal's own generic fusion logging — also prints for the *unrelated* built-in `RMS_NORM + MUL` (2-op) fusion, which still fires as a fallback wherever this branch's N-consumer fusion doesn't apply. |
+| `GGML_METAL_GRAPH_DEBUG=1` | Full per-node op listing in execution order — useful for checking whether sibling matmuls are actually contiguous in the graph (see "Structural constraint" below) on a model/arch you haven't traced before. |
+| `GGML_METAL_CONCURRENCY_STATS=1` | Prints `[concurrency-stats] nodes=<N> barriers=<N> avg_group=<X>` once per graph-compute call — quantifies how fragmented the barrier/concurrency schedule is. Unrelated to the fusion itself; added while investigating the scheduling fix below. |
+
+## Structural constraint you must respect if you extend this
+
+ggml-metal's fusion mechanism (the `n_fuse` return value from an op-dispatch function) only
+supports **contiguous** node ranges — the generic caller advances the encode cursor by exactly
+`n_fuse` positions, it cannot "skip forward" over interleaved unrelated nodes. This branch's
+N-consumer fusion only works because the real fan-out matmuls happen to land contiguously
+immediately after their shared `RMS_NORM`/`MUL` pair, in real model-construction order
+(`build_ffn`/`build_qkvz` in `src/models/qwen3next.cpp` emit sibling matmuls back-to-back). This
+was verified empirically via `GGML_METAL_GRAPH_DEBUG=1` before writing any dispatch code, for
+both an N=2 and an N=4 real site — **don't assume it holds on a different architecture without
+re-checking.** If a future model interleaves unrelated nodes between siblings, this fusion will
+just decline to fire there (safe), not silently miscompute.
+
+## Design history (why it's built this way, not some other way)
+
+Every one of these was root-caused via a real `llama-bench` measurement on real hardware, not
+guessed — several "obvious" fixes made things worse before the actual fix was found. If you're
+tempted to "simplify" the design, re-read this section first.
+
+1. **v1 (superseded): stage the normed vector in threadgroup memory.** Fuses norm+mul+matvec by
+   computing `normed[i] = x[i]*scale*norm_w[i]` for the whole hidden dimension into a
+   threadgroup-memory buffer, then all output rows read from there instead of device memory.
+   Measured: **-38% tg128** (41.35 → 25.62 t/s). Root cause: ~20KB/threadgroup footprint caps
+   Apple-GPU occupancy to ~1 resident threadgroup/core, which is exactly wrong for a
+   bandwidth-bound batch-1 decode kernel where occupancy is what hides DRAM latency while
+   streaming the (much larger) quantized weight matrix. The underlying premise was also wrong:
+   x and norm_w are tiny (~20KB) and stay SLC-cache-resident after the first touch, so
+   "redundant" per-threadgroup device reads of them were never a real bandwidth cost.
+2. **v2: drop the staging, read x/norm_w straight from device memory inline** (matching where
+   the baseline `kernel_mul_mv_{q1_0,q2_0}_f32_impl` reads its own activation), and match
+   upstream's proven `NSG=2` (had been an arbitrary `4`). **-23%.**
+3. **v3: fixed a second bug — grid-imbalance waste.** The multi-consumer kernel dispatched one
+   uniform `(max_tg, N)` grid sized to the *largest* sibling for every N-slot. Real fan-outs are
+   wildly imbalanced (observed `ne01=[10240,48,48,6144]` on one real site) — small siblings got
+   >99% of their dispatched threadgroups wasted, each still paying a full norm reduction before
+   exiting on the row-bounds check. Fixed: dispatch once per sibling with a grid sized to *that*
+   sibling's own `ne01`. **-10%.**
+4. **v4: broadcast-scale kernel.** Diagnostic (scale hardcoded to `1.0`, wrong output, used only
+   to isolate cost) showed the remaining gap was the norm reduction itself being redone by every
+   threadgroup in every sibling's dispatch, where the baseline computes it exactly once. Added
+   `kernel_rmsnorm_scale_f32` — one threadgroup, computes the scale once, writes it to a 1-float
+   scratch slot (reuses the `RMS_NORM` node's own now-unused output allocation, no new buffer
+   allocation needed — see the code comment in `ggml_metal_op_rmsnorm_qmv_try` for the safety
+   argument). Matvec dispatches read the precomputed scale instead of recomputing it. **~0%** —
+   parity, confirmed stable across repeated runs on both binary and ternary checkpoints.
+5. **Tiny-ne01 corner case (separate from the regression-chasing above):** found while auditing
+   for low-GPU-utilization patterns, not from a throughput regression. Some real sibling
+   matmuls have `ne01` as low as 48 — at the standard tiling that's only 3 threadgroups total,
+   leaving most of the GPU idle regardless of fusion. Added an `nr0=1` "_small" kernel variant,
+   selected per-sibling when the standard tiling would produce fewer than 16 threadgroups. Not
+   yet isolated as its own measured win (the tiny matmuls' own share of aggregate decode time is
+   too small for `llama-bench`'s aggregate `tg128` to resolve it) — would need a targeted
+   microbenchmark of just those dispatches to confirm.
+
+### Scheduling fix (separate from the fusion work, same branch)
+
+While investigating occupancy, found that `ggml_metal_graph_optimize_reorder`'s `h_safe()`
+allowlist (which ops the concurrency-scheduling lookahead is allowed to look past) was missing
+`GGML_OP_GATED_DELTA_NET` and `GGML_OP_FLASH_ATTN_EXT` — the lookahead does a hard `break` (not
+skip) the instant it hits an unlisted op, and this hybrid model hits one of those two on nearly
+every layer. Added both (their state is fully explicit in src/dst tensors, same as the
+already-safe-listed `GGML_OP_SSM_SCAN` — the actual hazard check, `h_check` against real buffer
+ranges, still gates correctness regardless of this list). Verified bit-exact. Measured own
+isolated impact (vs a true pre-branch baseline build, fusion off in both): **~0.6%, within
+noise** — the barrier-count reduction (~14% fewer) didn't translate to a comparable tok/s
+change; the per-token decode critical path is mostly inherently serial.
+
+## Known limitations / not done here
+
+- Single-consumer fusion (`kernel_rmsnorm_mv_{q1_0,q2_0}_f32`) is dead code on Bonsai-27B —
+  every RMSNorm here has ≥2 consumers. Kept because it's a strict subset check, costs nothing to
+  leave in, and might matter on a different architecture.
+- Not tested under real concurrent serving load (multiple simultaneous requests) — only
+  single-stream `llama-bench`. Dispatch-count reduction might matter more there than it does
+  here, but that's untested.
+- Not validated on prefill/batched decode (deliberately out of scope — batch=1 only).
+- Parity is not yet a demonstrated *speedup*. If you're picking this up hoping for a quick win,
+  it isn't one yet — it's a structurally-sound fusion that no longer costs anything, on the
+  hypothesis (untested end-to-end) that fewer dispatches matter more under real serving
+  conditions than in this isolated benchmark.