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

cactus-graph/cactus_graph.h

@@ -223,6 +223,21 @@ class BufferPool {
     size_t round_up_size(size_t size) const;
 };
 
+// One-time SME2 packed-weight cache for a CQ buffer (see CactusQuantMatrix::expanded_sme).
+// Held by shared_ptr so buffer copies share it; built lazily at first matmul via
+// cactus_quant_build_sme_cache. ~K*N/2 bytes per weight matrix.
+struct CactusSmeCache {
+    std::vector<uint8_t> packed;
+    std::vector<float> norms;
+    // Orthogonal (lm_head) extras: K x K fp16 rotation transposed + per-row norms replicated
+    // across the K/128 virtual groups (see cactus_quant_orth_sme_gemv).
+    std::vector<__fp16> rot_t;
+    std::vector<__fp16> norms_rep;
+    uint32_t virt_ng = 0;
+    std::once_flag once;
+    bool ready = false;
+};
+
 struct BufferDesc {
     std::vector<size_t> shape;
     size_t total_size;
@@ -237,6 +252,7 @@ struct BufferDesc {
 
     void* activation_scales_data = nullptr;
     std::unique_ptr<char[]> owned_activation_scales;
+    mutable std::shared_ptr<CactusSmeCache> sme_cache;
     size_t num_rows_for_activation_scales = 0;
 
     BufferDesc();
@@ -264,8 +280,95 @@ struct BufferDesc {
     const uint32_t* cq_permutation = nullptr;
     const __fp16* cq_rotation = nullptr;
     uint32_t cq_flags = 0;
+    // True when external data points into a read-only file mapping (weights loaded via
+    // MappedFile) — required before releasing packed pages after the SME cache build.
+    bool mmap_backed = false;
+    // Stable per-weight-file identity (hash of resolved path; 0 = none). Graph components map
+    // the same weight file independently, so data pointers do NOT identify a weight.
+    uint64_t weight_key = 0;
+
+    // Lazily build the SME2 packed cache for this weight buffer (no-op on non-SME hardware or
+    // unsupported group sizes; orthogonal weights get a virtual-group view + transposed rotation
+    // for the SME lm_head path). Thread-safe; heavy work runs once.
+    void ensure_sme_cache() const {
+        if (sme_cache && sme_cache->ready) return;
+        if (!cactus_quant_sme_enabled()) return;   // also skips builds under forced-NEON A/B
+        if (group_size <= 0 || (group_size % 32) != 0) return;
+        const uint32_t bits = PrecisionTraits::cq_bits(precision);
+        if (bits < 1 || bits > 4) return;
+        const bool orth = (cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) != 0;
+        if (orth) {
+            // lm_head shape only: single group spanning K, fp16 rotation, INTERLEAVED_4ROW
+            // nibbles (the production format — row-major orthogonal lm_heads came from a
+            // transpiler bug and are not SME-accelerated).
+            CactusQuantMatrix W = to_cq_matrix();
+            if (bits != 4 || !cq_rotation || W.num_groups != 1 || W.group_size != W.K ||
+                (W.K % 128) != 0 || (W.N % 4) != 0 ||
+                (cq_flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) == 0) return;
+        }
+        if (!sme_cache) {
+            // Process-wide registry keyed by the mmap'd weight bytes: every graph component
+            // holds its own BufferDesc for the same weight file, and per-desc caches would
+            // duplicate the entire SME cache once per component (measured 2x on Gemma) and
+            // re-fault the released file pages. weak_ptr keeps lifetime with the buffers.
+            static std::mutex reg_mu;
+            static std::unordered_map<const void*, std::weak_ptr<CactusSmeCache>> registry;
+            const void* key = weight_key
+                ? reinterpret_cast<const void*>(weight_key) : get_data();
+            std::lock_guard<std::mutex> lk(reg_mu);
+            auto& slot = registry[key];
+            sme_cache = slot.lock();
+            if (!sme_cache) {
+                sme_cache = std::make_shared<CactusSmeCache>();
+                slot = sme_cache;
+            }
+        }
+        std::call_once(sme_cache->once, [&] {
+            CactusQuantMatrix W = to_cq_matrix();
+            if (orth) {
+                // Virtual 128-wide regrouping — exact for INTERLEAVED_4ROW: panel bytes are
+                // ordered by k-chunk, so the gs=128 re-view lands on contiguous 256-byte
+                // sub-panels at (nb*vng + g)*256 == the physical nb*2K + g*256. Norms are
+                // constant across subgroups; replicate them in the interleaved norms layout
+                // [(nb*ng + g)*4 + ni] that the IL preexpander and NEON co-worker read.
+                const uint32_t vng = W.K / 128;
+                sme_cache->virt_ng = vng;
+                sme_cache->norms_rep.resize(static_cast<size_t>(W.N) * vng);
+                for (size_t nb = 0; nb < W.N / 4; ++nb)
+                    for (uint32_t g = 0; g < vng; ++g)
+                        for (size_t ni = 0; ni < 4; ++ni)
+                            sme_cache->norms_rep[(nb * vng + g) * 4 + ni] = cq_norms[nb * 4 + ni];
+                sme_cache->rot_t.resize(static_cast<size_t>(W.K) * W.K);
+                for (size_t k = 0; k < W.K; ++k)
+                    for (size_t i = 0; i < W.K; ++i)
+                        sme_cache->rot_t[i * W.K + k] = cq_rotation[k * W.K + i];
+                W.flags = CACTUS_QUANT_FLAG_INTERLEAVED_4ROW;
+                W.group_size = 128; W.num_groups = vng;
+                W.norms = sme_cache->norms_rep.data();
+            }
+            const size_t SB64 = (static_cast<size_t>(W.N) + 63) / 64;
+            sme_cache->packed.resize(SB64 * W.num_groups * static_cast<size_t>(W.group_size / 4) * 128);
+            sme_cache->norms.resize(SB64 * W.num_groups * 64);
+            cactus_quant_build_sme_cache(&W, sme_cache->packed.data(), sme_cache->norms.data());
+            if (orth) {
+                // Builder input only — runtime kernels read norm_sme (norms.data()) instead.
+                sme_cache->norms_rep.clear();
+                sme_cache->norms_rep.shrink_to_fit();
+            }
+            if (mmap_backed) {
+                // Single-format RAM policy: the cache is now the runtime weight layout for BOTH
+                // the SME leaves and the NEON co-workers, so drop the demand-paged packed file
+                // bytes (clean pages; they refault from disk only on file-layout paths such as
+                // forced-NEON A/B or embedding row gathers).
+                CactusQuantMatrix W0 = to_cq_matrix();
+                cactus_quant_release_packed_pages(&W0);
+            }
+            sme_cache->ready = true;
+        });
+    }
 
     CactusQuantMatrix to_cq_matrix() const {
+        const bool sme_ready = sme_cache && sme_cache->ready;
         return CactusQuantMatrix{
             .bits = PrecisionTraits::cq_bits(precision),
             .K = static_cast<uint32_t>(shape.size() >= 2 ? shape[1] : shape[0]),
@@ -284,6 +387,8 @@ struct BufferDesc {
             .rotation = cq_rotation,
             .expanded = nullptr,
             .norm_f32 = nullptr,
+            .expanded_sme = sme_ready ? sme_cache->packed.data() : nullptr,
+            .norm_sme = sme_ready ? sme_cache->norms.data() : nullptr,
         };
     }
 

cactus-graph/src/core.cpp

@@ -96,7 +96,10 @@ BufferDesc::BufferDesc(BufferDesc&& other) noexcept
       cq_right_signs(other.cq_right_signs),
       cq_permutation(other.cq_permutation),
       cq_rotation(other.cq_rotation),
-      cq_flags(other.cq_flags) {
+      cq_flags(other.cq_flags),
+      mmap_backed(other.mmap_backed),
+      weight_key(other.weight_key),
+      sme_cache(std::move(other.sme_cache)) {
     other.total_size = 0;
     other.byte_size = 0;
     other.external_data = nullptr;
@@ -143,6 +146,9 @@ BufferDesc& BufferDesc::operator=(BufferDesc&& other) noexcept {
         cq_permutation = other.cq_permutation;
         cq_rotation = other.cq_rotation;
         cq_flags = other.cq_flags;
+        mmap_backed = other.mmap_backed;
+        weight_key = other.weight_key;
+        sme_cache = std::move(other.sme_cache);
 
         other.total_size = 0;
         other.byte_size = 0;

cactus-graph/src/io.cpp

@@ -420,6 +420,14 @@ size_t CactusGraph::mmap_embeddings(const std::string& filename) {
         if (mapped_file->is_interleaved_4row()) {
             buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW;
         }
+        buffer.mmap_backed = true;
+        // Stable identity across graph components mapping the same weight file: lets the SME
+        // cache registry share ONE cache per weight (FNV-1a of the resolved path).
+        {
+            uint64_t h = 1469598103934665603ull;
+            for (char c : resolved_filename) { h ^= static_cast<unsigned char>(c); h *= 1099511628211ull; }
+            buffer.weight_key = h;
+        }
     } else if (precision == Precision::INT8 && mapped_file->group_size() > 0) {
         buffer.group_size = mapped_file->group_size();
         buffer.num_groups = mapped_file->num_groups();
@@ -485,6 +493,14 @@ size_t CactusGraph::mmap_weights(const std::string& filename) {
         if (mapped_file->is_interleaved_4row()) {
             buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW;
         }
+        buffer.mmap_backed = true;
+        // Stable identity across graph components mapping the same weight file: lets the SME
+        // cache registry share ONE cache per weight (FNV-1a of the resolved path).
+        {
+            uint64_t h = 1469598103934665603ull;
+            for (char c : resolved_filename) { h ^= static_cast<unsigned char>(c); h *= 1099511628211ull; }
+            buffer.weight_key = h;
+        }
     }
 
     size_t file_idx = mapped_files_.size();
@@ -568,6 +584,14 @@ void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename)
         if (mapped_file->is_interleaved_4row()) {
             buffer.cq_flags |= CACTUS_QUANT_FLAG_INTERLEAVED_4ROW;
         }
+        buffer.mmap_backed = true;
+        // Stable identity across graph components mapping the same weight file: lets the SME
+        // cache registry share ONE cache per weight (FNV-1a of the resolved path).
+        {
+            uint64_t h = 1469598103934665603ull;
+            for (char c : resolved_filename) { h ^= static_cast<unsigned char>(c); h *= 1099511628211ull; }
+            buffer.weight_key = h;
+        }
     } else if (precision == Precision::INT8 && mapped_file->group_size() > 0) {
         buffer.group_size = mapped_file->group_size();
         buffer.num_groups = mapped_file->num_groups();

cactus-graph/src/ops_nn.cpp

@@ -53,11 +53,27 @@ void compute_matmul_node(GraphNode& node, const std::vector<std::unique_ptr<Grap
         const __fp16* lhs = lhs_buffer.data_as<__fp16>();
         __fp16* output = node.output_buffer.data_as<__fp16>();
 
+        rhs_buffer.ensure_sme_cache();
         CactusQuantMatrix mat = rhs_buffer.to_cq_matrix();
-        if (rhs_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL)
-            cactus_quant_orthogonal_matmul(&mat, lhs, static_cast<uint32_t>(M), output);
-        else
+        if (rhs_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) {
+            auto& sc = rhs_buffer.sme_cache;
+            if (M == 1 && sc && sc->ready && sc->virt_ng > 0
+                && cactus_quant_sme_orth_enabled()) {
+                CactusQuantMatrix W2 = mat;
+                W2.flags = CACTUS_QUANT_FLAG_INTERLEAVED_4ROW;
+                W2.group_size = 128;
+                W2.num_groups = sc->virt_ng;
+                W2.norms = mat.norms;   // co-workers read norm_sme; norms_rep is freed post-build
+                W2.expanded_sme = sc->packed.data();
+                W2.norm_sme = sc->norms.data();
+                cactus_quant_orth_sme_gemv(&W2, sc->rot_t.data(), mat.input_scale_recip,
+                                           lhs, output);
+            } else {
+                cactus_quant_orthogonal_matmul(&mat, lhs, static_cast<uint32_t>(M), output);
+            }
+        } else {
             cactus_quant_matmul(&mat, lhs, static_cast<uint32_t>(M), output);
+        }
     } else {
         if (lhs_buffer.precision != Precision::FP16) {
             throw std::runtime_error("FP16 matmul requires FP16 activations (got precision " + std::to_string(static_cast<int>(lhs_buffer.precision)) + ")");
@@ -115,10 +131,25 @@ namespace {
         }
 
         if (PrecisionTraits::is_cq(rhs_buffer.precision) && rhs_buffer.group_size > 0) {
+            rhs_buffer.ensure_sme_cache();
             CactusQuantMatrix mat = rhs_buffer.to_cq_matrix();
-            if (rhs_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL)
-                cactus_quant_orthogonal_matmul(&mat, lhs, static_cast<uint32_t>(M), output);
-            else
+            if (rhs_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) {
+                auto& sc = rhs_buffer.sme_cache;
+                if (M == 1 && sc && sc->ready && sc->virt_ng > 0
+                    && cactus_quant_sme_orth_enabled()) {
+                    CactusQuantMatrix W2 = mat;
+                    W2.flags = CACTUS_QUANT_FLAG_INTERLEAVED_4ROW;
+                    W2.group_size = 128;
+                    W2.num_groups = sc->virt_ng;
+                    W2.norms = mat.norms;   // co-workers read norm_sme; norms_rep is freed post-build
+                    W2.expanded_sme = sc->packed.data();
+                    W2.norm_sme = sc->norms.data();
+                    cactus_quant_orth_sme_gemv(&W2, sc->rot_t.data(), mat.input_scale_recip,
+                                               lhs, output);
+                } else {
+                    cactus_quant_orthogonal_matmul(&mat, lhs, static_cast<uint32_t>(M), output);
+                }
+            } else
                 cactus_quant_matmul(&mat, lhs, static_cast<uint32_t>(M), output);
             return;
         }
@@ -342,8 +373,11 @@ void compute_dense_mlp_tq_fused_node(GraphNode& node, const std::vector<std::uni
     __fp16* prod = prod_buf.data();
     __fp16* output = node.output_buffer.data_as<__fp16>();
 
+    gate_buffer.ensure_sme_cache();
     CactusQuantMatrix gate_mat = gate_buffer.to_cq_matrix();
+    up_buffer.ensure_sme_cache();
     CactusQuantMatrix up_mat = up_buffer.to_cq_matrix();
+    down_buffer.ensure_sme_cache();
     CactusQuantMatrix down_mat = down_buffer.to_cq_matrix();
     const bool use_safe_product_scale = node.params.scalar != 0.0f && node.params.scalar != 1.0f;
     const bool trace_dense_mlp = std::getenv("CACTUS_TRACE_DENSE_MLP") != nullptr;

cactus-graph/src/ops_tensor.cpp

@@ -157,6 +157,47 @@ void compute_embedding_node(GraphNode& node, const std::vector<std::unique_ptr<G
     Precision emb_prec = embeddings_buffer.precision;
     if (PrecisionTraits::is_cq(emb_prec) && embeddings_buffer.group_size > 0) {
         bool orthogonal = (embeddings_buffer.cq_flags & CACTUS_QUANT_FLAG_ORTHOGONAL) != 0;
+        bool interleaved = (embeddings_buffer.cq_flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) != 0;
+        if (orthogonal && interleaved && (hidden_dim % 8) == 0) {
+            // Batched path (production INTERLEAVED_4ROW orthogonal embeddings): dedup indices,
+            // then ONE parallel dequant+un-rotation pass over the unique rows (the per-row
+            // fallback below is a serial scalar K^2 matvec per row — it was a top prefill term
+            // for per-layer embeddings). Legacy row-major orthogonal bundles (a transpiler bug)
+            // take the per-row fallback.
+            std::unordered_map<size_t, uint32_t> slot;
+            slot.reserve(std::min<size_t>(num_indices, 256));
+            std::vector<uint32_t> uniq;
+            uniq.reserve(std::min<size_t>(num_indices, 256));
+            std::vector<uint32_t> idx_slot(num_indices);
+            for (size_t i = 0; i < num_indices; i++) {
+                size_t idx = static_cast<size_t>(indices_ptr[i]);
+                if (idx >= vocab_size) {
+                    throw std::runtime_error("Embedding index out of bounds: " + std::to_string(idx) + " >= " + std::to_string(vocab_size));
+                }
+                auto it = slot.emplace(idx, static_cast<uint32_t>(uniq.size()));
+                if (it.second) uniq.push_back(static_cast<uint32_t>(idx));
+                idx_slot[i] = it.first->second;
+            }
+            std::vector<__fp16> rows_out(uniq.size() * hidden_dim);
+            cactus_quant_dequantize_orthogonal_embedding_rows(
+                PrecisionTraits::cq_bits(emb_prec),
+                static_cast<uint32_t>(hidden_dim),
+                uniq.data(),
+                static_cast<uint32_t>(uniq.size()),
+                embeddings_buffer.data_as<uint8_t>(),
+                embeddings_buffer.cq_codebook,
+                embeddings_buffer.cq_norms,
+                embeddings_buffer.cq_input_scale_recip,
+                embeddings_buffer.cq_rotation,
+                embeddings_buffer.cq_flags,
+                rows_out.data());
+            for (size_t i = 0; i < num_indices; i++) {
+                std::memcpy(output + i * hidden_dim,
+                            rows_out.data() + static_cast<size_t>(idx_slot[i]) * hidden_dim,
+                            hidden_dim * sizeof(__fp16));
+            }
+            return;
+        }
         std::unordered_map<size_t, size_t> row_cache;
         std::vector<__fp16> cached_rows;
         if (num_indices > 16) {

cactus-kernels/CMakeLists.txt

@@ -15,6 +15,8 @@ add_library(cactus_kernels STATIC
     src/image.cpp
     src/lstm.cpp
     src/matmul.cpp
+    src/matmul_sme2.cpp
+    src/matmul_sme2_gemv.cpp
     src/nn.cpp
     src/norms_rope.cpp
     src/quants.cpp
@@ -46,6 +48,28 @@ if(APPLE)
     target_link_libraries(cactus_kernels PUBLIC "-framework Accelerate")
 endif()
 
+# SME2 path: compile only matmul_sme2.cpp with +sme2 so the streaming/ZA code is isolated and the
+# base library keeps targeting armv8.2-a (runs on non-SME devices). Use armv8.2-a+...+sme2, NOT
+# armv9-a+sme2 — armv9-a implies non-streaming SVE2 that Apple Silicon lacks (SIGILLs at runtime).
+include(CheckCXXSourceCompiles)
+set(_CACTUS_SME2_FLAG "-march=armv8.2-a+fp16+simd+dotprod+i8mm+sme2")
+set(CMAKE_REQUIRED_FLAGS "${_CACTUS_SME2_FLAG}")
+check_cxx_source_compiles("
+#include <arm_sme.h>
+__arm_locally_streaming __arm_new(\"za\") static void f(){ svzero_za(); }
+int main(){ return 0; }
+" CACTUS_HAS_SME2)
+unset(CMAKE_REQUIRED_FLAGS)
+if(CACTUS_HAS_SME2)
+    set_source_files_properties(src/matmul_sme2.cpp PROPERTIES COMPILE_OPTIONS "${_CACTUS_SME2_FLAG}")
+    # vg1x4 GEMV TU: -O1 ONLY — clang 17 mis-compiles the SME2 multi-vector ZA-array dot at -O2/-O3
+    # into an illegal encoding that SIGILLs on M4 (correct + fast at -O0/-O1). Isolate it here.
+    set_source_files_properties(src/matmul_sme2_gemv.cpp PROPERTIES COMPILE_OPTIONS "${_CACTUS_SME2_FLAG};-O1")
+    message(STATUS "cactus_kernels: SME2 enabled (matmul_sme2.cpp; matmul_sme2_gemv.cpp at -O1)")
+else()
+    message(STATUS "cactus_kernels: SME2 not supported by toolchain; SME TUs use scalar stubs")
+endif()
+
 if((BUILD_TESTING OR CACTUS_BUILD_TESTS) AND CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
     file(GLOB TEST_SOURCES "tests/test_*.cpp")
     foreach(TEST_SRC ${TEST_SOURCES})