Skip to content

SME2 streaming kernels over the packed-panel weight format - #710

Merged
ncylich merged 1 commit into
panel-format-neonfrom
sme2-integration
Jun 11, 2026
Merged

ncylich merged 1 commit into
panel-format-neonfrom
sme2-integration

Conversation

@ncylich

@ncylich ncylich commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

Second PR of the panel-format split (stacked on #709 — retarget to main once that merges). The SME2 layer consumes the file-borne panels zero-copy: the panel bytes ARE the streaming layout, so there is no cache builder, no madvise machinery, and no second resident weight copy. Non-SME2 silicon (and any SME2 implementation with SVL ≠ 512-bit) falls back to the #709 NEON panel path bit-for-bit.

What's here

  • Streaming leaves (matmul_sme2_gemv.cpp, matmul_sme2.cpp): LUTI4/ZT0 nibble-expanding GEMV + fused SMOPA GEMM reading mmap'd panels directly, and the prefill-attention QK/AV SMOPA passes. All leaves emit raw int32 partials; the fp32 rescale stays in NEON, so SME numerics are identical to the NEON path (locked by an agreement test).
  • Dispatch: SME workers replace NEON workers inside the existing mobile thread budget (flat k = min(2, nt−1), the measured power/speed frontier; CACTUS_SME_GEMV_WORKERS overrides). The fused SMOPA GEMM takes all M>1 panel matmuls. CACTUS_QUANT_BACKEND/cactus_quant_set_backend: 0 auto, 1 force NEON, 2 force the streaming leaves.
  • Gating: cpu_has_sme2() requires streaming svcntb() == 64 — every layout is pinned to 512-bit SVL; other SVLs fall back to NEON for correctness. The SME TUs build with armv8.2-a+…+sme2 (armv9-a implies non-streaming SVE2 that Apple Silicon lacks) and the GEMV TU at -O1 (clang mis-compiles the vg1x4 ZA dots at -O2+ → SIGILL).
  • Docs: cactus-kernels/docs/sme/ knowledge base (design rules, gotchas, verified working examples for every layout).

Tests (TDD: suite written first, implemented to green)

[sme2] variants of every panel test through the real dispatch — CQ1–4 GEMV, GEMM incl. M/N tails, orthogonal lm_head, interleaved fixtures — plus a NEON↔SME agreement gate on identical fixtures and the SME prefill-attention differential. 35/35 matmul + 6/6 attention + full graph suite green.

Measured (M4 Pro, Gemma 4 E2B IT panel bundle, ~1k-token prompt, 32 decode tokens, 6 alternating cycles vs the #709 NEON panel engine, cycle 1 dropped; coherent output both sides)

prefill tok/s decode tok/s TTFT ms decode energy
#709 (panel NEON) 300.5 49.25 3278 358 mJ/token
this PR (panel + SME2) 353.3 50.68 2788 294 mJ/token
delta +17.6% +2.9% −14.9% −17.9%

Cumulative vs the legacy interleaved format on the same protocol: prefill +35%, decode +21%, TTFT −26%.

Kernel-level: fused SMOPA GEMM 1.5–3.1× the NEON panel GEMM (M16: 872 vs 285 GFLOPS); GEMV with SME workers up to 2.1× on K-heavy decode shapes (kv_proj 278 vs 130 GF).

@ncylich

ncylich commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

Generalization + warm-prefill parity

Other architectures, same protocol (legacy = production interleaved bundle on the NEON panel engine; panel+SME2 = this PR):

model config prefill tok/s decode tok/s TTFT ms
qwen3-1.7b legacy / NEON 239.4 36.96 4123
panel + SME2 399.5 (+66.9%) 53.01 (+43.4%) 2471 (−40.1%)
lfm2-350m legacy / NEON 143.6 139.7 7249
panel + SME2 195.0 (+35.8%) 185.0 (+32.5%) 5339 (−26.4%)
gemma-4-e2b-it legacy / NEON 261.2 42.02 3772
panel + SME2 353.3 (+35.3%) 50.68 (+20.6%) 2788 (−26.1%)

Coherent on all three; qwen panel NEON-vs-SME outputs are identical.

Warm-process prefill: measured on a second in-process message (weights paged in, caches warm), this PR sustains ~500 tok/s prefill on Gemma 4 E2B — run side-by-side with the pre-split SME branch it matches it exactly (480–502 vs 478–502 tok/s), confirming the file-borne panels lose nothing vs the old runtime-cache design while eliminating its in-process cache build (which costs that branch 1.4–3.4 s of first-message TTFT per process).

Stacks the SME2 layer on the panel file format: the LUTI4/SMOPA leaves read
the mmap'd panels directly (the panel bytes ARE the streaming layout), so
there is no cache builder and no second weight copy. Non-SME2 silicon and
non-512-bit SVL fall back to the NEON panel path bit-for-bit.

Kernels
- matmul_sme2_gemv.cpp: LUTI4/ZT0 nibble-expanding GEMV + fused GEMM leaves
  emitting raw int32 partials (fp32 rescale stays in NEON, numerics identical
  to the NEON path). Compiled at -O1: clang mis-compiles the vg1x4 ZA dots at
  -O2+ into encodings that SIGILL on M4.
- matmul_sme2.cpp: SMOPA/USMOPA leaves incl. prefill-attention QK/AV passes
  and the streaming-SVL probe; compiled with armv8.2-a+...+sme2 (armv9-a
  implies non-streaming SVE2 Apple Silicon lacks).
- attention_hybrid.cpp: SME prefill attention over the cached int8 KV segment
  (flat-scale QK requantize at pre-pack, u8 USMOPA AV with folded V scales).

Dispatch
- SME workers replace NEON workers inside the existing thread budget (flat
  k = min(2, nt-1), the measured power/speed frontier; CACTUS_SME_GEMV_WORKERS
  overrides). The fused SMOPA GEMM takes all M>1 panel matmuls.
- CACTUS_QUANT_BACKEND / cactus_quant_set_backend: 0 auto, 1 force NEON,
  2 force the streaming leaves (test coverage).
- cpu_has_sme2() gates on streaming svcntb() == 64: every layout is pinned to
  512-bit SVL; other SME2 implementations fall back to NEON for correctness.

Tests
- [sme2] variants of every panel test (GEMV/GEMM incl. M/N tails, orthogonal
  lm_head, interleaved fixtures, CQ1-4) through the real dispatch, plus a
  NEON<->SME agreement gate on identical fixtures and the SME prefill
  attention differential.

Measured (M4 Pro, Gemma 4 E2B IT panel bundle, ~1k-token prompt, 32 tokens,
6 alternating cycles vs the NEON panel engine): prefill +17.6% (300->353
tok/s), TTFT -14.9% (3278->2788 ms), decode +2.9%, decode energy/token -17.9%
(358->294 mJ). Kernel-level: fused SMOPA GEMM 1.5-3.1x the NEON panel GEMM;
GEMV with SME workers up to 2.1x on K-heavy shapes.

Signed-off-by: Noah Cylich <noahcylich@gmail.com>
@ncylich
ncylich force-pushed the sme2-integration branch from a6e04b6 to ff35ac9 Compare June 10, 2026 22:21
@ncylich
ncylich merged commit ff35ac9 into panel-format-neon Jun 11, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant