Repository navigation
SME2 streaming kernels over the packed-panel weight format - #710
Merged
Merged
Conversation
ncylich
force-pushed
the
sme2-integration
branch
from
June 10, 2026 10:34
46446a5 to
a6e04b6
Compare
Collaborator
Author
Generalization + warm-prefill parityOther architectures, same protocol (legacy = production interleaved bundle on the NEON panel engine; panel+SME2 = this PR):
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
force-pushed
the
sme2-integration
branch
from
June 10, 2026 22:21
a6e04b6 to
ff35ac9
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
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).k = min(2, nt−1), the measured power/speed frontier;CACTUS_SME_GEMV_WORKERSoverrides). 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.cpu_has_sme2()requires streamingsvcntb() == 64— every layout is pinned to 512-bit SVL; other SVLs fall back to NEON for correctness. The SME TUs build witharmv8.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).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)
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).