Skip to content

SME2 acceleration for CQ4 matmul, prefill attention, and orthogonal embeddings#697

Closed
ncylich wants to merge 6 commits into
mainfrom
sme2-cq4-acceleration
Closed

SME2 acceleration for CQ4 matmul, prefill attention, and orthogonal embeddings#697
ncylich wants to merge 6 commits into
mainfrom
sme2-cq4-acceleration

Conversation

@ncylich

@ncylich ncylich commented Jun 9, 2026

Copy link
Copy Markdown
Collaborator

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):

metric NEON (CACTUS_QUANT_BACKEND=1) SME2 (auto) delta
decode 35.8 tok/s 48.0 tok/s +34% (worst cycle +28.6%)
TTFT 3413 ms 2249 ms −34%
prefill 319 tok/s 477 tok/s +49%

What's in here

Quantized matmul (cactus-kernels)

  • Two new SME2 translation units: matmul_sme2.cpp (streaming SMOPA leaves) and matmul_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/-O3 into a runtime SIGILL). Only these two files get +sme2 flags (armv9-a deliberately avoided: it implies non-streaming SVE2, absent on Apple Silicon). Scalar stubs keep the build green on non-SME2 toolchains; base library flags unchanged.
  • Hybrid fused GEMV (M=1): one pool dispatch — phase A does parallel Hadamard + int8 activation quantization with a spin barrier, phase B dynamically steals 64-channel super-blocks: SME workers stream packed nibbles through the LUTI4 partials leaf (half the weight bytes, ZA never read back into the core), the remaining workers run the production NEON kernels on the native packed layouts. The main thread participates as SME worker 0 with a spin-join — the wait_all condition-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 for K*N ≥ 3M elements.
  • Fused SME2 GEMM (M>1): (M-tile × super-block) work stealing over the packed weight cache, four ZA tiles in flight, raw int32 partials with NEON rescale — 2.4–3.6× vs NEON at every M (K=N=1024).
  • The orthogonal-rotation CQ4 lm_head maps onto the same hybrid through an exact virtual 128-wide-group view.
  • The packed SME weight cache builds once per weight from row-major or production INTERLEAVED_4ROW layouts, 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=0 disables.

Orthogonal embedding dequantization (both backends benefit)

compute_embedding_node was 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 streaming svcntb() == 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=1 knobs.
  • Strided round-robin partitioning for coarse parallel_for items (the static splitter floors work_per_thread and dumps the remainder on the last worker — measured ~60% pool idle on attention tiles).
  • BufferDesc move 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

  • 61 tests across 7 suites pass (cactus-kernels/test.sh), including new:
    • forced-backend CQ1–4 GEMV/GEMM vs the fp32 oracle through the real dispatch, with M/N tile-tail cases;
    • production INTERLEAVED_4ROW fixtures built as the exact inverse of the shipped decoder; SME-cache layout invariance (row-major vs interleaved byte-identical);
    • orthogonal lm_head vs oracle; batched embedding rows vs the per-row reference on both layouts;
    • attention differential + double-precision oracle over 4 cases (global mid-prefill, rolled sliding window with sink exemption, zero-offset window, batch=2 with partial tiles/blocks). Worst-case attention relative error ~0.6–0.7% on adversarial uniform synthetic vs ~0.1% for the fp16 incumbent.
  • End-to-end temperature-0 output verified coherent on both backends; SME-path engagement verified by symbol-level profiling.
  • Benchmarks use alternating best-of runs (thermal drift exceeds kernel deltas).

Notes for reviewers

  • The attention path is the first SME path with real numeric divergence from NEON (int8 Q/K, u8 P quantization) — bounded and oracle-tested as above; all matmul SME paths produce integer partials identical to the NEON SDOT inner loop.
  • Perf policies (worker splits, 3M-element dispatch threshold, queue model) are measured on M4 Pro; they are heuristics to re-validate on other SME2 silicon — the SVL gate keeps such devices correct (NEON) in the meantime.
  • Known follow-ups documented in 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.

…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>
Copilot AI review requested due to automatic review settings June 9, 2026 23:05

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 +sme2 codepaths.
  • 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.

Comment on lines +212 to +216
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
Comment on lines +17 to +23
// 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>
@ncylich

ncylich commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

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. ensure_sme_cache now requires IL4ROW for orthogonal weights, the orth driver's NEON co-workers use the interleaved block processor, and the batched embedding dequant is interleaved-only (legacy row-major bundles fall back to the incumbent per-row path and should be recompiled).

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 0.77 ms). Note for benchmark interpretation: the corrected format also speeds up the NEON baseline (+14% decode — its old row-major orth lm_head cost ~4.8 ms/step vs 1.33 ms interleaved), so backend-vs-backend deltas measured on corrected bundles are smaller than the headline numbers above (which were measured against the buggy-format vintage), while absolute SME throughput is unchanged (~47.6 tok/s decode).

…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>
@ncylich

ncylich commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

Follow-up commit 1c2077d: single runtime weight format — no more 2× weight RAM on SME devices.

Demand paging was keeping the packed .weights bytes resident (NEON co-workers and small-shape GEMVs streamed them every token) on top of the materialized SME cache. Now the SME cache is the one runtime format for both engines:

  • New cactus_quant_esme_gemv_blocks: NEON SDOT directly over the cache layout (its expanded byte order is the classic [4ch×4K] SDOT panel). All tuned production-kernel techniques carried over; kernel-level parity-or-better vs the file kernels everywhere (pure-NEON small shape 1.29×, hybrid ratios unchanged).
  • With a cache present, the fused driver is always the M=1 dispatch (SME share sized by shape, down to 0 = pure NEON over the cache); no file-layout matmul kernel runs in auto mode.
  • Packed file pages are madvise(MADV_DONTNEED)'d after the one-time build (gated on a new mmap_backed flag; norms/codebook/rotation stay resident; forced-NEON A/B refaults cleanly). CACTUS_QUANT_BACKEND=1 now skips cache builds entirely.
  • Found and fixed a pre-existing 2× on top: every graph component maps the same weight file independently, so per-BufferDesc caches were built once per component (554 builds for 277 weights). A process-wide weak_ptr registry keyed by weight-file identity dedupes them.

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>
@ncylich

ncylich commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

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):

config prefill tok/s decode tok/s TTFT
legacy file-NEON (CACTUS_QUANT_BACKEND=1) 368 42.0 2902 ms
pure NEON over the cache (no SME GEMV/attn) 487 51.4 2192 ms
shipped default (NEON GEMVs + SME GEMM + SME attention) ~510 ~50–52 ~2.0–2.2 s

Two findings worth highlighting:

  1. The long-standing "NEON collapses on K-heavy GEMV" behavior was a file-layout artifact: o_proj jumps 57–67 → ~200–235 GF on the cache layout. With that fixed, the queue-limited SME GEMV leaf nets negative E2E under sustained decode (51.4 vs 44.8 tok/s) despite bursty per-shape kernel wins — so CACTUS_SME_GEMV_WORKERS now defaults to 0 (env/setter re-enables it for experiments and other SME2 silicon; forced-SME clamps to ≥1 so the leaf stays test-covered).
  2. SME's durable wins are the single cache format itself, the fused M>1 GEMM (2.4–3.6×), and prefill attention (~100 ms TTFT in direct A/B) — all retained.

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>
@ncylich

ncylich commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

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 ceil(N/256), GEMM one thread per M-tile), with SME workers living inside the budget (replacing NEON workers, never adding threads; auto k=2 on ≥4-thread non-giant shapes, override via CACTUS_SME_GEMV_WORKERS).

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>
@ncylich

ncylich commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

Follow-up commit (power frontier): measured the decode power/speed frontier directly with powermetrics (decode-only windows, thermally-paired passes) over a thread-budget × SME-worker grid:

config decode CPU power energy/token
legacy budget, no SME 48.7 tok/s 21.0 W 430 mJ
legacy budget, 1 SME worker 50.0 tok/s 18.3 W 367 mJ
half budget, 2 SME workers (shipped) 49.1 tok/s 16.1 W 328 mJ

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 (CACTUS_GEMV_SB_PER_THREAD, CACTUS_SME_GEMV_WORKERS) stay runtime-tunable for per-device retuning — the Galaxy-class SoC probed alongside (SVE2 + SME1, no SME2) runs the NEON-over-cache path and is unaffected.

@ncylich

ncylich commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

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%).

@ncylich ncylich closed this Jun 10, 2026
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.

2 participants