SME2 acceleration for CQ4 matmul, prefill attention, and orthogonal embeddings#697
SME2 acceleration for CQ4 matmul, prefill attention, and orthogonal embeddings#697ncylich wants to merge 6 commits into
Conversation
…mbeddings Add an ARM SME2 backend for the hot CPU inference paths, behind runtime feature dispatch. Measured end-to-end on Gemma 4 E2B IT (Apple M4 Pro, 1068-token prompt, 32 generated tokens, temperature 0, 6-cycle alternating A/B against the production NEON path in the same binary): decode 35.8 -> 48.0 tok/s (+34%) TTFT 3413 -> 2249 ms (-34%) prefill 319 -> 477 tok/s (+49%) All SME paths fall back to the existing NEON kernels on hardware, toolchains, or shapes they do not cover. Quantized matmul (cactus-kernels): - New translation units matmul_sme2.cpp (streaming SMOPA leaf kernels) and matmul_sme2_gemv.cpp (LUTI4/ZT0 in-engine 4-bit codebook expansion; compiled at -O1 to work around an Apple clang 17 mis-compile of vg1x4 ZA-array dots at -O2/-O3 that SIGILLs at runtime), bridged by matmul_simd.h. Only these two files get -march=...+sme2 (armv9-a is avoided deliberately: it implies non-streaming SVE2, which is absent on Apple Silicon). Scalar stubs keep all symbols defined when the toolchain lacks SME2. - Hybrid fused GEMV (M=1): one pool dispatch runs phase A (parallel Hadamard transform + int8 activation quantization, spin barrier), then phase B steals 64-channel super-blocks dynamically - SME workers stream packed nibbles through the LUTI4 partials leaf (half the weight bytes, no ZA readback into the core), remaining workers run the production NEON kernels on the native packed layouts. The main thread participates as SME worker 0 and spin-joins instead of sleeping in wait_all; the condition-variable sleep alone was a 20-30% tax on sub-100us calls. Auto-dispatched for K*N >= 3M elements. Measured vs production NEON at Gemma shapes: down_proj 2.8x, o_proj 3.2x, gate_up 1.6x, q_proj 1.4x, lm_head 1.1x (DRAM-bound). - Fused SME2 GEMM (M>1): (M-tile x super-block) pair stealing over the packed weight cache, four ZA tiles in flight, raw int32 partials rescaled in NEON per pair; 2.4-3.6x vs NEON across every M at K=N=1024. - The orthogonal-rotation CQ4 lm_head maps onto the same hybrid via a mathematically exact virtual 128-wide group view (the rotation is applied during phase A). - The packed SME weight cache (CactusQuantMatrix::expanded_sme / norm_sme) builds once per weight from either row-major or INTERLEAVED_4ROW production layouts, with byte-identical output for both (covered by a test). Prefill attention (attention_hybrid.cpp): - New SME2 path for the cached-int8 KV segment of chunked prefill, which is fully causally visible when position_offset >= cache_len. Per 16-query-row tile: SMOPA QK over 64-kv blocks using a flat scale per Q row and per K vector (per-group scale ratios are requantized into the int8 values so ZA accumulates the full head_dim with one 4 KB readout per block); vectorized masked softmax; then u8 x s8 USMOPA for the AV product with per-(row, v-group, block) P quantization folded into the operand (softmax probabilities are non-negative, so unsigned doubles the resolution). The fp16 new-token segment continues through the existing per-row flash-softmax loop seeded from the cached statistics. Sliding-window and attention-sink masking replicate the incumbent exactly, and fully-masked block runs skip SME work entirely. - 4.6x (global layers) / 2.7x (sliding-window layers) vs the NEON kernel at Gemma shapes; CACTUS_SME_ATTENTION=0 disables the path. Accuracy is differential-tested against a double-precision oracle (worst-case relative error ~0.6-0.7% on adversarial uniform attention vs ~0.1% for the fp16 incumbent). Orthogonal embedding dequantization (cactus-graph + cactus-kernels): - compute_embedding_node previously un-rotated every unique token row with a serial scalar K^2 matvec (~2.4M scalar FMAs per row at K=1536), repeated for each per-layer embedding lookup per chunk - one of the largest prefill terms. The new cactus_quant_dequantize_orthogonal_embedding_rows dedups indices and runs a single parallel 8-wide pass with identical fp32-accumulate math. This speeds up both backends. Runtime dispatch and threading: - cpu_has_sme()/cpu_has_sme2() runtime detection (Apple sysctl, Android HWCAP2). cpu_has_sme2() additionally requires the streaming vector length to be 64 bytes (svcntb probe): every shipped SME layout is pinned to 512-bit SVL, so SME2 implementations with other vector lengths fall back to NEON for correctness rather than computing garbage. - CACTUS_QUANT_BACKEND=0/1/2 (auto / force NEON / force SME2) enables in-binary A/B testing; CACTUS_SME_GEMV_WORKERS tunes the SME/NEON worker split; CACTUS_SME_PROF=1 dumps fused-GEMM phase timings at exit. - Coarse parallel_for work items now use strided round-robin partitioning where it matters: the static splitter floors work_per_thread and gives the last worker the entire remainder (64 items / 14 workers = 13x4 + 1x12, measured as ~60% pool idle on attention tiles). - BufferDesc move construction/assignment now transfers the SME cache pointer instead of silently dropping a built cache. Tests (61 across 7 suites, all passing): - Forced-backend CQ1-4 GEMV/GEMM correctness against the fp32 oracle through the real cactus_quant_matmul dispatch, including M- and N-tail tile cases. - Production INTERLEAVED_4ROW fixtures built as the exact inverse of the shipped decoder; SME cache layout-invariance (row-major vs interleaved builds byte-identical); orthogonal lm_head vs oracle; batched embedding rows vs the per-row reference on both layouts. - Attention differential + double-precision oracle across four cases: global layer mid-prefill, sliding window with rolled cache (sink exemption live), window at zero cache offset, and batch=2 with partial query tiles and partial KV blocks. - Alternating best-of benchmark harnesses at Gemma shapes (thermal drift exceeds kernel deltas, so single-sided timing is unreliable). docs/sme/ adds the engineering notebook for this work: a status board, verified intrinsic and API notes, hardware findings (per-cluster in-order SME command queue, ZA readout costs, toolchain traps), a dated experiment log, runnable probe examples, and research notes. Signed-off-by: Noah Cylich <noahcylich@gmail.com>
There was a problem hiding this comment.
Pull request overview
This PR introduces an ARM SME2-accelerated CPU backend for Cactus’s hottest inference paths (quantized matmul, chunked-prefill attention, and orthogonal embedding dequantization), with runtime feature dispatch and NEON fallbacks to preserve correctness and portability.
Changes:
- Add SME2 leaf kernels (SMOPA/LUTI4/USMOPA) plus runtime SME/SME2 detection and build wiring to isolate
+sme2codepaths. - Extend CQ weight representation with an optional SME2 packed cache and integrate lazy cache building into graph execution (including an SME fast-path for orthogonal
lm_head). - Expand test coverage with forced-backend correctness checks and oracle-based differential testing for hybrid int8 prefill attention; add SME engineering notebook docs and runnable probes.
Reviewed changes
Copilot reviewed 31 out of 31 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| cactus-kernels/src/matmul_sme2.cpp | SME2 streaming-mode leaf kernels (SMOPA) + SVL probe for runtime gating. |
| cactus-kernels/src/matmul_sme2_gemv.cpp | SME2 GEMV/GEMM LUTI4/ZT0-based kernels (compiled at -O1 per toolchain constraint). |
| cactus-kernels/src/matmul_simd.h | Shared SME leaf-kernel declarations bridging NEON and SME2 TUs. |
| cactus-kernels/src/threading.h | Add cpu_has_sme() / cpu_has_sme2() runtime detection and SVL==64 correctness gate. |
| cactus-kernels/CMakeLists.txt | Add SME2 sources and per-TU -march ... +sme2 (and -O1) compile options with feature probe. |
| cactus-kernels/cactus_kernels.h | Extend CactusQuantMatrix with optional SME cache pointers; add SME cache builder + SME control APIs + batched orth embedding API. |
| cactus-kernels/tests/test_matmul.cpp | Add forced-backend SME2/NEON validation, interleaved fixtures, orth lm_head SME test, cache invariance checks, and benchmark printouts. |
| cactus-kernels/tests/test_attention.cpp | Add DP oracle + differential test for SME hybrid int8 prefill attention and benchmark comparisons. |
| cactus-graph/cactus_graph.h | Add CactusSmeCache and lazy BufferDesc::ensure_sme_cache(); plumb cache pointers into to_cq_matrix(). |
| cactus-graph/src/ops_nn.cpp | Ensure SME cache is built before CQ matmuls; add SME orth lm_head GEMV fast-path. |
| cactus-graph/src/ops_tensor.cpp | Add batched+deduped orthogonal embedding-row dequantization path. |
| cactus-graph/src/core.cpp | Move sme_cache with BufferDesc moves. |
| cactus-kernels/docs/sme/** | Add SME knowledge base, synthesis/research notes, and runnable working examples/probes. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| const double rel_neon = e_neon / std::max(ref_amax, 1e-12); | ||
| const double rel_sme = e_sme / std::max(ref_amax, 1e-12); | ||
| const bool engaged_expected = cactus_quant_sme_available() != 0; | ||
| bool ok = rel_sme < 0.05 && rel_sme < 5.0 * std::max(rel_neon, 0.002); | ||
| if (engaged_expected && d_ns == 0.0) ok = false; // SME path silently not taken |
| // CQ GEMV (M=1) per-group SMOPA accumulate, 4 output channels per call-block. | ||
| // act_i8 : [num_groups * gs] int8, per-group quantized activations (group g at +g*gs). | ||
| // w_block : Cactus 'expanded' INT8 weights for ONE 4-channel N-block; group g panel at | ||
| // w_block + (size_t)g*gs*4, with each kg-th 16 bytes = 4 channels x 4 K | ||
| // (identical layout to the NEON SDOT path). Up to `valid_cols` (1..4) channels valid. | ||
| // partials_out : [num_groups * 4] int32 raw dot products; partials_out[g*4 + c] = | ||
| // sum_k act[g][k] * weight[channel c][k]. Caller applies act_scale * norm rescale. |
The orthogonal-rotation CQ4 lm_head ships in the INTERLEAVED_4ROW packed layout; row-major orthogonal bundles came from a transpiler bug and are being retired. The SME lm_head path previously required row-major nibbles and silently disengaged on production bundles (verified by symbol profiling: zero orth-SME samples on a current compile, with the NEON interleaved kernel running instead). Key observation, proven by panel arithmetic and by the oracle test on an interleaved fixture: the virtual 128-wide group regrouping used by the SME lm_head path is byte-exact on INTERLEAVED_4ROW as well. Interleaved panel bytes are ordered by k-chunk, so re-viewing the single gs=K group as vng=K/128 groups of 128 lands each virtual group on a contiguous 256-byte sub-panel at (nb*vng + g)*256 — identical to the physical nb*2K + g*256. No repacking is needed; the existing interleaved preexpander and block kernels consume the virtual view natively. Changes: - ensure_sme_cache (cactus-graph): the orthogonal branch now REQUIRES INTERLEAVED_4ROW (gate flipped), keeps the IL flag on the virtual view, and replicates per-row norms in the interleaved [(nb*ng + g)*4 + ni] layout that the IL preexpander and NEON co-workers read. - Both orthogonal matmul dispatch sites (ops_nn.cpp) build the virtual view with flags = INTERLEAVED_4ROW. - cactus_quant_orth_sme_gemv: NEON co-workers switched from the row-major from-packed SDOT reader to the production interleaved block processor (cactus_quant_interleaved4_gemv_blocks), consuming the virtual view directly. - Row-major orthogonal support removed from the new code paths: the batched embedding dequant (cactus_quant_dequantize_orthogonal_ embedding_rows) is interleaved-only, and compute_embedding_node routes legacy row-major orthogonal bundles to the per-row fallback. Legacy bundles still load and run through the incumbent NEON paths; they should be recompiled. - test_orth_sme converted to an interleaved fixture (encoder is the exact inverse of the shipped panel decoder), exercising both the production NEON interleaved lm_head kernel and the SME path against the fp32 oracle; orth_embed_rows_batched is interleaved-only. Validation: 61 tests across 7 suites pass; SME engagement restored on a current production compile (orth-SME hot, incumbent at zero samples); temperature-0 output identical across backends. Isolated lm_head op cost: SME-interleaved best 0.98 ms vs the NEON interleaved kernel's best 1.20 ms per call (DRAM floor ~0.77 ms for the 192 MB nibble stream — the op is bandwidth-capped). Note the corrected format also speeds up the NEON baseline itself (~+14% decode: its old row-major orthogonal lm_head path cost ~4.8 ms/step vs 1.33 ms interleaved), so backend-vs-backend deltas on corrected bundles are smaller than against the buggy-format baseline while absolute SME throughput is unchanged. Signed-off-by: Noah Cylich <noahcylich@gmail.com>
|
Follow-up commit 9f30254: the SME orthogonal lm_head path now targets the production INTERLEAVED_4ROW format (row-major orthogonal bundles were a transpiler bug and are retired from the new code paths). The virtual 128-wide group regrouping turns out to be byte-exact on the interleaved layout too (panel bytes are k-chunk-ordered, so the gs=128 re-view lands on contiguous 256-byte sub-panels) — no repacking; the existing interleaved preexpander and block kernels consume the virtual view natively. Validation: 61/61 tests (orth test converted to an interleaved fixture against the fp32 oracle), SME engagement verified on a current production compile, temperature-0 output identical across backends. Isolated lm_head op: SME best 0.98 ms vs NEON interleaved best 1.20 ms (DRAM floor |
…le pages Demand paging kept the packed .weights bytes resident (the hybrid's NEON co-workers and small-shape GEMVs streamed them every token) on top of the materialized SME cache - roughly 2x the weight RAM on SME devices. This change makes the SME cache the single runtime weight format for both engines and drops the file-backed packed pages after the one-time build. Measured on Gemma 4 E2B IT (M4 Pro, auto backend): physical footprint at the same late-decode moment 4.9 GB -> 3.7 GB (-1.2 GB); max RSS 5992 -> 5018 MB. Decode/TTFT at parity with the previous binary in alternating runs; 61 tests across 7 suites pass; temperature-0 output coherent on both backends. - New cactus_quant_esme_gemv_blocks: a NEON SDOT kernel that consumes the SME cache layout directly. The cache's nibbles expand (low nibble first, pinned by LUTI4's identity nibble->byte mapping) to the classic [4 vec][16 ch][4 K] SDOT panel order, so unpacking is and/shr + 2x vqtbl1q + 2x vzip + 2x vdotq_laneq per 16 bytes, with 16 independent accumulators for ILP. The two zips are the structural price of sharing one format (a zip-free planar nibble packing would push the shuffles into the SME command queue instead - the scarce resource). All the tuned production-kernel techniques carry over: codebook tbl lookup, lane-broadcast SDOT with one activation load per four k-quads, single fp32 rescale with cb_scale folded into the cached norms, fp16 stores clipped to valid channels. Kernel-level: parity-or-better with the file-layout kernels everywhere, including the pure-NEON small shape (1536x512: 71.6 vs 55.6 GF, 1.29x) and unchanged hybrid ratios (o_proj 3.17x, down 2.93x, gate_up ~1.5x). - Dispatch: with the cache present the fused driver is now always the M=1 dispatch; it sizes its SME worker share by shape (>= 3M weight elements -> CACTUS_SME_GEMV_WORKERS, below -> 0, i.e. pure NEON co-workers over the cache). The hybrid and orthogonal-lm_head drivers' co-workers consume the cache instead of the file panels; the M>1 fused GEMM already did. No file-layout matmul kernel runs in auto mode once a weight's cache exists. - Page release: BufferDesc::mmap_backed (set by the io.cpp weight loaders) gates cactus_quant_release_packed_pages - an madvise(MADV_DONTNEED) on the page-aligned interior of the packed region only (norms/codebook/rotation stay resident). Clean file pages refault from disk if a file-layout path runs again (forced-NEON A/B, embedding row gathers). CACTUS_RAM_DEBUG=1 logs each release. - Cache dedup across graph components: every transpiled component maps the same weight file independently, so per-BufferDesc caches were built once PER COMPONENT (measured 2x cache RAM and 554 builds for 277 weights, with the second build refaulting just-released pages). ensure_sme_cache now shares one cache per weight through a process-wide weak_ptr registry keyed by BufferDesc::weight_key, an FNV-1a hash of the resolved weight path set at load time - mmap addresses do not identify a weight across components. - Cache builds are skipped entirely under CACTUS_QUANT_BACKEND=1 (cactus_quant_sme_enabled), keeping the force-NEON A/B baseline a pure file-kernel run with no cache RAM; non-SME devices are unaffected (no cache, no release, file path as before). - Removed the now-unused row-major from-packed co-worker wrapper; the orthogonal driver no longer needs the replicated norms at runtime (freed after the cache build). Known transient: the cache builder's int8 expansion temp still spikes peak RSS during the first prefill (~400 MB on the lm_head); streaming the build is a noted follow-up alongside eager build at model load. Signed-off-by: Noah Cylich <noahcylich@gmail.com>
|
Follow-up commit 1c2077d: single runtime weight format — no more 2× weight RAM on SME devices. Demand paging was keeping the packed
Measured (Gemma 4 E2B, auto backend): physical footprint at the same late-decode moment 4.9 GB → 3.7 GB (−1.2 GB); max RSS 5992 → 5018 MB. Decode/TTFT at parity in alternating binary-vs-binary runs; 61/61 tests; temperature-0 output coherent on both backends. Known transient: the cache builder's int8 temp still spikes first-prefill peak (~400 MB on the lm_head) — streaming the build is a noted follow-up alongside eager build at load. |
… SME GEMV workers With NEON now consuming the cache layout, the honest baseline for the SME GEMV hybrid changed - and the measurements with it. Added a runtime worker-count override (cactus_quant_set_sme_gemv_workers / CACTUS_SME_GEMV_WORKERS) and re-ran everything three ways: legacy file-layout NEON vs pure NEON over the cache vs the SME hybrid, alternating, kernel-level and end-to-end. Findings (Gemma 4 E2B IT, M4 Pro): - The long-standing "NEON collapses on K-heavy GEMV" behavior was a file-layout artifact, not a NEON limit: o_proj goes 57-67 GF on the interleaved file panels to ~200-235 GF on the cache layout. - Against the same-format NEON baseline the SME GEMV leaf keeps only modest bursty kernel wins (down 1.1-1.3x, ffn/gate_up ~1.0-1.2x), loses on o_proj (0.82-0.94x across runs) and lm_head (~0.97x), and nets NEGATIVE end-to-end under sustained decode: pure NEON-over-cache decodes 51.4 tok/s vs 44.8 for the hybrid configuration. Nibble streaming through the shared in-order SME queue (~64-85 MACs per queue operation) cannot beat ten-plus unobstructed NEON cores once the layout is right. - SME keeps earning its place where it genuinely wins: the single cache format itself, the fused M>1 GEMM (2.4-3.6x), and prefill attention (~100 ms TTFT in direct A/B). Shipped default is therefore: NEON GEMVs over the cache + SME fused GEMM + SME prefill attention. Measured end-to-end (alternating runs, same windows): decode ~50-52 tok/s, TTFT ~2.0-2.2 s - versus the legacy file-layout NEON baseline at 42.0 tok/s / 2.9 s, i.e. decode +19%, TTFT -27.5%, the best configuration measured in this effort, on top of the -1.2 GB physical footprint from the single-format change. Mechanics: - CACTUS_SME_GEMV_WORKERS now defaults to 0 (pure NEON co-workers for every M=1 GEMV, including the orthogonal lm_head driver); the env var or the new setter re-enables SME GEMV workers for experiments and for re-tuning on other SME2 silicon. - Forced-SME (backend 2) clamps the worker count to >= 1 so the [sme2] correctness tests keep covering the streaming leaf; test_orth_sme pins 4 workers explicitly for the same reason. - The interleaved-comparison benchmark now reports all three columns (file-NEON / cache-NEON / hybrid with 4 workers) with the speedup quoted against the same-format NEON baseline. - docs/sme: dated log entry plus a correction to the K-heavy gotcha; the lesson recorded is to gate wins on the baseline's best format, not just your own. 61 tests across 7 suites pass; temperature-0 output coherent. Signed-off-by: Noah Cylich <noahcylich@gmail.com>
|
Follow-up commit 9e6ac5c: re-measured everything against the same-format NEON baseline, and retuned the dispatch defaults accordingly. Three-way comparison (legacy file-layout NEON vs pure NEON over the cache vs SME hybrid; alternating, kernel-level + E2E on Gemma 4 E2B):
Two findings worth highlighting:
Net vs the legacy production baseline: decode +19%, TTFT −27.5%, physical footprint −1.2 GB, with 61/61 tests passing. The benchmark harness now reports all three columns with speedups quoted against the same-format NEON baseline. |
The low thread counts in the legacy kernels were deliberate mobile/battery policy, not an accident - saturating cores is wasteful on phones. This restores the original budgets in the new fused drivers and re-tunes the SME worker policy within them. - GEMV driver: thread budget = ceil(N/256), the exact legacy formula (kv_proj 2 threads, o_proj/down 6, gate_up/lm_head pool-capped). - GEMM driver: one thread per 16-row M-tile, matching the legacy M>1 path. Embedding lookups cap at ~3 threads for decode-time single-row gathers. - Latency improvements that add no threads are kept: the fused single dispatch, the activation preamble running parallel on the same woken workers (the legacy kernel ran it serially), and the spin-join. - SME GEMV workers REPLACE NEON workers inside the budget, never add threads. Measured at budgeted thread counts, k=2 SME workers win K-heavy shapes (o_proj +20%, down +23%) and gate_up (+11%), are par on ffn/q_proj, and lose on tiny shapes and the DRAM-bound lm_head - the shipped auto policy is k=2 when the budget is >= 4 threads and N < 65536, else 0. CACTUS_SME_GEMV_WORKERS or the runtime setter overrides; forced-SME clamps to >= 1 so the [sme2] tests keep covering the streaming leaf. - The interleaved-comparison benchmark reports the k-sweep at budgeted threads. End-to-end (Gemma 4 E2B IT, alternating runs): the budgeted configuration BEATS core saturation - decode 52.7 tok/s and TTFT ~1990 ms vs ~50-51 tok/s / 2100-2200 ms for the all-cores configs - and the intermittent run-to-run collapses observed under saturation disappear entirely. Versus the legacy production baseline at the same thread counts: decode 40.6 -> 52.7 tok/s (+29.7%), TTFT 2875 -> 1993 ms (-30.7%), prefill +44%, on top of the -1.2 GB physical footprint. At matched budgets the SME GEMV contribution is end-to-end neutral on this chip (kernel-level wins wash out); the format, the fused M>1 GEMM, and prefill attention carry the wins. 61 tests across 7 suites pass. Signed-off-by: Noah Cylich <noahcylich@gmail.com>
|
Follow-up commit (thread budget): the low thread counts in the legacy kernels were deliberate mobile/battery policy — restored them in the fused drivers (GEMV Surprise result: the budgeted configuration beats core saturation end-to-end — decode 52.7 tok/s / TTFT ~1990 ms vs ~50–51 / 2100–2200 for the all-cores configs — and the intermittent run collapses seen under saturation disappear. Final shipped numbers vs the legacy production baseline at the same thread counts: decode 40.6 → 52.7 tok/s (+29.7%), TTFT 2875 → 1993 ms (−30.7%), prefill +44%, physical footprint −1.2 GB. At matched budgets the SME GEMV leaf is kernel-positive on K-heavy shapes (+20–23%) but E2E-neutral on M4; the durable SME wins are the single cache format, the fused M>1 GEMM, and prefill attention. 61/61 tests. |
Measured the decode power/speed frontier directly (powermetrics CPU
power over decode-only windows, 420-token decodes, two passes in
opposite order to cancel thermal drift) across a grid of thread
budgets (CACTUS_GEMV_SB_PER_THREAD in {2,4,8}) and SME worker counts
(CACTUS_SME_GEMV_WORKERS in {0,1,2}) on Gemma 4 E2B IT / M4 Pro:
legacy budget, no SME (spt4 k0): 48.7 tok/s @ 21.0 W = 430 mJ/tok
legacy budget, 1 SME (spt4 k1): 50.0 tok/s @ 18.3 W = 367 mJ/tok
half budget, 2 SME (spt8 k2): 49.1 tok/s @ 16.1 W = 328 mJ/tok
SME workers strictly dominate zero-SME at every thread budget -
faster AND lower power. The earlier "SME GEMV is end-to-end neutral"
verdict was an artifact of measuring tokens/sec only: the leaf's
value is on the power axis, replacing NEON-saturated cores with one
queue-feeding core plus the matrix block. A flat worker count also
beat the per-shape gating policy, which is removed.
Shipped defaults: GEMV thread budget halved to one thread per eight
super-blocks, with a flat k = min(2, nt-1) SME workers everywhere
including the orthogonal lm_head driver. Versus the legacy budget
without SME this is +0.8% speed, -23% CPU power, -24% energy per
token; TTFT is unchanged (~2.0 s) and verification runs at the
shipped defaults hold 49-51 tok/s at 14.8-15.8 W. Both knobs remain
runtime-tunable for per-device retuning (the Galaxy-class SoC probed
alongside this work exposes SVE2 + SME1 but no SME2, so it runs the
NEON-over-cache path and is unaffected by the worker default).
61 tests across 7 suites pass; temperature-0 output coherent.
Signed-off-by: Noah Cylich <noahcylich@gmail.com>
|
Follow-up commit (power frontier): measured the decode power/speed frontier directly with
SME workers strictly dominate zero-SME at every thread budget — faster and lower power. The earlier "SME GEMV is E2E-neutral" reading came from measuring tok/s only; the leaf's value is on the power axis. Shipped defaults: half the legacy GEMV thread budget with a flat 2 SME workers (+0.8% speed, −24% energy/token vs the legacy budget without SME; TTFT unchanged; 61/61 tests). Both knobs ( |
|
Superseded by the two-PR split: #709 moves the packed-panel layout into the .cact weight files themselves (no runtime cache — eliminates the lazy build, the page-release machinery, and the 2× RAM transient by construction) with NEON kernels for all aarch64, and #710 stacks the SME2 streaming layer on top reading the mmap'd panels zero-copy. Together they land everything this branch contained, with the format now explicit per file (.cq4p.weights) and validated E2E on Gemma 4 E2B IT (cumulative vs legacy: prefill +35%, decode +21%, TTFT −26%, decode energy/token −18%). |
Summary
Adds an ARM SME2 backend for the hot CPU inference paths — quantized matmul (GEMV + GEMM), chunked-prefill attention, and orthogonal embedding dequantization — behind runtime feature dispatch with NEON fallback everywhere SME2 is absent or doesn't win.
End-to-end on Gemma 4 E2B IT (Apple M4 Pro, 1068-token prompt, 32 tokens, temperature 0, MTP off; 6-cycle alternating A/B vs the production NEON path in the same binary, cycle 1 discarded):
CACTUS_QUANT_BACKEND=1)What's in here
Quantized matmul (
cactus-kernels)matmul_sme2.cpp(streaming SMOPA leaves) andmatmul_sme2_gemv.cpp(LUTI4/ZT0 in-engine 4-bit codebook expansion; pinned to-O1— Apple clang 17 mis-compiles vg1x4 ZA-array dots at-O2/-O3into a runtime SIGILL). Only these two files get+sme2flags (armv9-adeliberately avoided: it implies non-streaming SVE2, absent on Apple Silicon). Scalar stubs keep the build green on non-SME2 toolchains; base library flags unchanged.wait_allcondition-variable sleep alone was a 20–30% tax on sub-100µs calls. Kernel-level vs production NEON at Gemma shapes: down_proj 2.8×, o_proj 3.2×, gate_up 1.6×, q_proj 1.4×, lm_head 1.1× (DRAM-bound). Auto-dispatch forK*N ≥ 3Melements.INTERLEAVED_4ROWlayouts, byte-identical for both (tested).Prefill attention (
attention_hybrid.cpp)New SME2 path for the cached-int8 KV segment of chunked prefill (fully causally visible when
position_offset ≥ cache_len): 16-query-row tiles, SMOPA QK with flat per-row/per-kv scales (group-scale ratios requantized into the int8 values so ZA accumulates the whole head_dim with one 4 KB readout per 64-kv block), vectorized masked softmax, then u8×s8 USMOPA AV with per-(row, v-group, block) P quantization (softmax probabilities are non-negative — unsigned doubles the resolution). The fp16 new-token segment continues through the existing flash-softmax loop seeded from the cached statistics; sliding-window + attention-sink masking matches the incumbent exactly. 4.6× (global) / 2.7× (sliding-window) at Gemma shapes.CACTUS_SME_ATTENTION=0disables.Orthogonal embedding dequantization (both backends benefit)
compute_embedding_nodewas un-rotating each unique token row with a serial scalar K² matvec (~2.4M scalar FMAs per row at K=1536), repeated per per-layer embedding lookup per chunk — one of the largest prefill terms. The new batched path dedups indices and runs one parallel 8-wide NEON pass with identical fp32-accumulate math.Runtime / threading
cpu_has_sme2()(Apple sysctl / Android HWCAP2) additionally requires streamingsvcntb() == 64: all layouts are pinned to 512-bit SVL, so other SME2 implementations fall back to NEON for correctness until SVL-parametric variants exist.CACTUS_QUANT_BACKEND=0/1/2(auto / force NEON / force SME2) for in-binary A/B;CACTUS_SME_GEMV_WORKERS,CACTUS_SME_PROF=1knobs.parallel_foritems (the static splitter floorswork_per_threadand dumps the remainder on the last worker — measured ~60% pool idle on attention tiles).BufferDescmove construction/assignment now transfers the SME cache.Docs
cactus-kernels/docs/sme/is the engineering notebook: status board, verified intrinsic/API notes, hardware findings (per-cluster in-order SME command queue, ZA readout costs, toolchain traps), dated experiment log, runnable probes (working-examples/), and research notes.Correctness / test plan
cactus-kernels/test.sh), including new:INTERLEAVED_4ROWfixtures built as the exact inverse of the shipped decoder; SME-cache layout invariance (row-major vs interleaved byte-identical);Notes for reviewers
docs/sme/debug-log.md: eager-parallel SME cache build at model load (first-call lazy builds cost ~8.5 s on cold start, currently hidden by warmup), tail-sized prefill graph components, SME decode attention for long contexts.