From 5c02052d9f76708bac09d19b4eed0f03170af8e7 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Thu, 11 Jun 2026 01:59:26 -0700 Subject: [PATCH] Fuse dispatch for the legacy interleaved CQ4 GEMV Replace the legacy IL GEMV's three-pass dispatch (pool-parallel Hadamard + cv wait, serial int8 quantize, static parallel_ranges + cv wait) with a single fused dispatch: group-stolen phase A behind a spin barrier, dynamic 16-block-chunk stealing in phase B, main thread as worker 0 with a spin-join, thread budget ceil(chunks / CACTUS_GEMV_SB_PER_THREAD). The inner micro-kernel is unchanged and factored into cactus_quant_interleaved4_gemv_blocks. M4 kernel: kv_proj 44.7 -> 158.2 GF, o_proj 59.2 -> 188.4 GF. E2E decode on production bundles: gemma-4-e2b-it +19%, qwen3-1.7b +42%, lfm2-350m +26% (M4); Samsung 8-Elite +6%/+18%/+53%; Pixel Tensor G4 parity (DRAM roofline) with lfm2 +12-16%. Signed-off-by: Noah Cylich --- cactus-kernels/src/matmul.cpp | 279 ++++++++++++++++----------- cactus-kernels/tests/test_matmul.cpp | 79 ++++++++ 2 files changed, 250 insertions(+), 108 deletions(-) diff --git a/cactus-kernels/src/matmul.cpp b/cactus-kernels/src/matmul.cpp index ffc88560b..41daa4bde 100644 --- a/cactus-kernels/src/matmul.cpp +++ b/cactus-kernels/src/matmul.cpp @@ -1,6 +1,8 @@ #include "../cactus_kernels.h" #include "threading.h" #include +#include +#include #include #include #include @@ -318,6 +320,47 @@ static void cactus_quant_parallel_ranges(size_t total_work, size_t work_per_thre pool.wait_all(); } +static size_t cactus_quant_gemv_sb_per_thread() { + static const size_t v = [] { + const char* e = getenv("CACTUS_GEMV_SB_PER_THREAD"); + const int i = e ? atoi(e) : 8; + return static_cast(i > 0 ? i : 8); + }(); + return v; +} + +template +static void cactus_quant_two_phase_run(size_t nt, uint32_t num_groups, uint32_t n_items, + PhaseA&& phase_a, PhaseB&& phase_b) { + std::atomic next_group{0}; + std::atomic groups_done{0}; + std::atomic next_item{0}; + std::atomic workers_done{0}; + auto worker = [&](size_t wid) { + for (uint32_t g; (g = next_group.fetch_add(1, std::memory_order_relaxed)) < num_groups; ) { + phase_a(g); + groups_done.fetch_add(1, std::memory_order_release); + } + while (groups_done.load(std::memory_order_acquire) < num_groups) {} + for (;;) { + const uint32_t seen = next_item.load(std::memory_order_relaxed); + if (seen >= n_items) break; + const uint32_t want = (n_items - seen > 4u * nt) ? 4u : 1u; + const uint32_t item = next_item.fetch_add(want, std::memory_order_relaxed); + if (item >= n_items) break; + phase_b(wid, item, std::min(want, n_items - item)); + } + }; + auto& pool = CactusThreading::get_thread_pool(); + // Main joins as worker 0 and spins: a cv sleep costs ~5-10us per call, material at decode rates. + pool.enqueue_n_threads(nt - 1, nt - 1, [&](size_t wid, size_t) { + worker(wid + 1); + workers_done.fetch_add(1, std::memory_order_release); + }); + worker(0); + while (workers_done.load(std::memory_order_acquire) < nt - 1) {} +} + static void cactus_quant_matmul_f32_segment_accum( const __fp16* __restrict__ A, size_t a_stride, @@ -2160,6 +2203,102 @@ void cactus_quant_orthogonal_matmul( }); } +static void cactus_quant_interleaved4_gemv_blocks( + const CactusQuantMatrix* W, const uint8_t* packed_interleaved, + const __fp16* norms_interleaved, const int8_t* act_i8, const float* act_scales, + const int8x16_t cb_lut, const float cb_scale, + size_t block_start, size_t block_end, __fp16* y) { + const uint32_t gs = W->group_size; + const uint32_t pgb = cactus_quant_packed_group_bytes(4, gs); + const uint32_t num_groups = W->num_groups; + const size_t panel_bytes = 4 * pgb; + const uint8x16_t lo_mask = vdupq_n_u8(0x0F); + size_t nb = block_start; + size_t aligned_end = (block_end / 2) * 2; + + for (; nb + 2 <= aligned_end; nb += 2) { + float32x4_t acc0 = vdupq_n_f32(0.f); + float32x4_t acc1 = vdupq_n_f32(0.f); + + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* p0 = packed_interleaved + ((nb + 0) * num_groups + g) * panel_bytes; + const uint8_t* p1 = packed_interleaved + ((nb + 1) * num_groups + g) * panel_bytes; + const int8_t* a_grp = act_i8 + static_cast(g) * gs; + const float a_sc = act_scales[g]; + + int32x4_t d0a = vdupq_n_s32(0), d0b = vdupq_n_s32(0); + int32x4_t d1a = vdupq_n_s32(0), d1b = vdupq_n_s32(0); + + for (uint32_t kb = 0; kb < gs; kb += 16) { + int8x16_t a_v = vld1q_s8(a_grp + kb); + + #define PROC_PANEL(P, DA, DB) do { \ + const uint8_t* c0 = (P) + (kb / 8 + 0) * 16; \ + const uint8_t* c1 = (P) + (kb / 8 + 1) * 16; \ + uint8x16_t b0 = vld1q_u8(c0); \ + int8x16_t wl0 = vqtbl1q_s8(cb_lut, vandq_u8(b0, lo_mask)); \ + int8x16_t wh0 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b0, 4)); \ + DA = vdotq_laneq_s32(DA, wl0, a_v, 0); \ + DB = vdotq_laneq_s32(DB, wh0, a_v, 1); \ + uint8x16_t b1 = vld1q_u8(c1); \ + int8x16_t wl1 = vqtbl1q_s8(cb_lut, vandq_u8(b1, lo_mask)); \ + int8x16_t wh1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b1, 4)); \ + DA = vdotq_laneq_s32(DA, wl1, a_v, 2); \ + DB = vdotq_laneq_s32(DB, wh1, a_v, 3); \ + } while (0) + + PROC_PANEL(p0, d0a, d0b); + PROC_PANEL(p1, d1a, d1b); + #undef PROC_PANEL + } + + int32x4_t dot0 = vaddq_s32(d0a, d0b); + int32x4_t dot1 = vaddq_s32(d1a, d1b); + + float32x4_t n0 = vcvt_f32_f16(vld1_f16(norms_interleaved + ((nb + 0) * num_groups + g) * 4)); + float32x4_t n1 = vcvt_f32_f16(vld1_f16(norms_interleaved + ((nb + 1) * num_groups + g) * 4)); + float scale_grp = cb_scale * a_sc; + acc0 = vfmaq_f32(acc0, vcvtq_f32_s32(dot0), vmulq_n_f32(n0, scale_grp)); + acc1 = vfmaq_f32(acc1, vcvtq_f32_s32(dot1), vmulq_n_f32(n1, scale_grp)); + } + + vst1_f16(y + (nb + 0) * 4, vcvt_f16_f32(acc0)); + vst1_f16(y + (nb + 1) * 4, vcvt_f16_f32(acc1)); + } + + for (; nb < block_end; ++nb) { + float32x4_t acc = vdupq_n_f32(0.f); + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* p_base = packed_interleaved + (nb * num_groups + g) * panel_bytes; + const int8_t* a_grp = act_i8 + static_cast(g) * gs; + const float a_sc = act_scales[g]; + + int32x4_t dot_a = vdupq_n_s32(0); + int32x4_t dot_b = vdupq_n_s32(0); + + for (uint32_t kb = 0; kb < gs; kb += 16) { + int8x16_t a_v = vld1q_s8(a_grp + kb); + uint8x16_t b0 = vld1q_u8(p_base + (kb / 8 + 0) * 16); + int8x16_t wl0 = vqtbl1q_s8(cb_lut, vandq_u8(b0, lo_mask)); + int8x16_t wh0 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b0, 4)); + dot_a = vdotq_laneq_s32(dot_a, wl0, a_v, 0); + dot_b = vdotq_laneq_s32(dot_b, wh0, a_v, 1); + uint8x16_t b1 = vld1q_u8(p_base + (kb / 8 + 1) * 16); + int8x16_t wl1 = vqtbl1q_s8(cb_lut, vandq_u8(b1, lo_mask)); + int8x16_t wh1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b1, 4)); + dot_a = vdotq_laneq_s32(dot_a, wl1, a_v, 2); + dot_b = vdotq_laneq_s32(dot_b, wh1, a_v, 3); + } + + int32x4_t dot = vaddq_s32(dot_a, dot_b); + float32x4_t norm = vcvt_f32_f16(vld1_f16(norms_interleaved + (nb * num_groups + g) * 4)); + norm = vmulq_n_f32(norm, cb_scale * a_sc); + acc = vfmaq_f32(acc, vcvtq_f32_s32(dot), norm); + } + vst1_f16(y + nb * 4, vcvt_f16_f32(acc)); + } +} + void cactus_quant_4bit_gemv_interleaved( const CactusQuantMatrix* W, const uint8_t* packed_interleaved, @@ -2173,123 +2312,47 @@ void cactus_quant_4bit_gemv_interleaved( if (W->group_size > 256) return; const uint32_t gs = W->group_size; - const uint32_t pgb = cactus_quant_packed_group_bytes(4, gs); const uint32_t num_groups = W->num_groups; + const size_t N_blocks = W->N / 4; + const size_t n_chunks = (N_blocks + 15) / 16; + auto& pool = CactusThreading::get_thread_pool(); + const size_t sb_per_thread = cactus_quant_gemv_sb_per_thread(); + const size_t nt_budget = std::max(1, (n_chunks + sb_per_thread - 1) / sb_per_thread); + const size_t nt = std::min(pool.num_workers(), std::min(nt_budget, n_chunks)); - thread_local std::vector<__fp16> code_basis_buf; - if (code_basis_buf.size() < W->K) code_basis_buf.resize(W->K); - cactus_quant_transform_hadamard_activations(*W, x, 1, code_basis_buf.data()); - const __fp16* code_basis = code_basis_buf.data(); - - thread_local std::vector act_i8_buf; - thread_local std::vector act_scales_buf; - if (act_i8_buf.size() < W->K) act_i8_buf.resize(W->K); - if (act_scales_buf.size() < num_groups) act_scales_buf.resize(num_groups); - for (uint32_t g = 0; g < num_groups; ++g) { - act_scales_buf[g] = tq_quantize_group_i8( - code_basis + static_cast(g) * gs, - act_i8_buf.data() + static_cast(g) * gs, gs); - } - const int8_t* act_i8 = act_i8_buf.data(); - const float* act_scales = act_scales_buf.data(); + static thread_local std::vector tl_il_act_i8; + static thread_local std::vector tl_il_act_scales; + if (tl_il_act_i8.size() < W->K) tl_il_act_i8.resize(W->K); + if (tl_il_act_scales.size() < num_groups) tl_il_act_scales.resize(num_groups); + int8_t* act_i8 = tl_il_act_i8.data(); + float* act_scales = tl_il_act_scales.data(); int8_t cb_i8[16] = {}; const float cb_scale = tq_quantize_codebook_i8(W->codebook, cb_i8, 16); const int8x16_t cb_lut = vld1q_s8(cb_i8); - const size_t panel_bytes = 4 * pgb; - const size_t N_blocks = W->N / 4; - const size_t N_tiles_4 = N_blocks / 4; // groups of 4 n_blocks = 16 outputs - - const uint8x16_t lo_mask = vdupq_n_u8(0x0F); - - cactus_quant_parallel_ranges(N_blocks, 64, [&](size_t block_start, size_t block_end) { - size_t nb = block_start; - size_t aligned_end = (block_end / 2) * 2; - - for (; nb + 2 <= aligned_end; nb += 2) { - float32x4_t acc0 = vdupq_n_f32(0.f); - float32x4_t acc1 = vdupq_n_f32(0.f); - - for (uint32_t g = 0; g < num_groups; ++g) { - const uint8_t* p0 = packed_interleaved + ((nb + 0) * num_groups + g) * panel_bytes; - const uint8_t* p1 = packed_interleaved + ((nb + 1) * num_groups + g) * panel_bytes; - const int8_t* a_grp = act_i8 + static_cast(g) * gs; - const float a_sc = act_scales[g]; - - int32x4_t d0a = vdupq_n_s32(0), d0b = vdupq_n_s32(0); - int32x4_t d1a = vdupq_n_s32(0), d1b = vdupq_n_s32(0); - - for (uint32_t kb = 0; kb < gs; kb += 16) { - int8x16_t a_v = vld1q_s8(a_grp + kb); - - #define PROC_PANEL(P, DA, DB) do { \ - const uint8_t* c0 = (P) + (kb / 8 + 0) * 16; \ - const uint8_t* c1 = (P) + (kb / 8 + 1) * 16; \ - uint8x16_t b0 = vld1q_u8(c0); \ - int8x16_t wl0 = vqtbl1q_s8(cb_lut, vandq_u8(b0, lo_mask)); \ - int8x16_t wh0 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b0, 4)); \ - DA = vdotq_laneq_s32(DA, wl0, a_v, 0); \ - DB = vdotq_laneq_s32(DB, wh0, a_v, 1); \ - uint8x16_t b1 = vld1q_u8(c1); \ - int8x16_t wl1 = vqtbl1q_s8(cb_lut, vandq_u8(b1, lo_mask)); \ - int8x16_t wh1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b1, 4)); \ - DA = vdotq_laneq_s32(DA, wl1, a_v, 2); \ - DB = vdotq_laneq_s32(DB, wh1, a_v, 3); \ - } while (0) - - PROC_PANEL(p0, d0a, d0b); - PROC_PANEL(p1, d1a, d1b); - #undef PROC_PANEL - } - - int32x4_t dot0 = vaddq_s32(d0a, d0b); - int32x4_t dot1 = vaddq_s32(d1a, d1b); - - float32x4_t n0 = vcvt_f32_f16(vld1_f16(norms_interleaved + ((nb + 0) * num_groups + g) * 4)); - float32x4_t n1 = vcvt_f32_f16(vld1_f16(norms_interleaved + ((nb + 1) * num_groups + g) * 4)); - float scale_grp = cb_scale * a_sc; - acc0 = vfmaq_f32(acc0, vcvtq_f32_s32(dot0), vmulq_n_f32(n0, scale_grp)); - acc1 = vfmaq_f32(acc1, vcvtq_f32_s32(dot1), vmulq_n_f32(n1, scale_grp)); - } - - vst1_f16(y + (nb + 0) * 4, vcvt_f16_f32(acc0)); - vst1_f16(y + (nb + 1) * 4, vcvt_f16_f32(acc1)); - } - - for (; nb < block_end; ++nb) { - float32x4_t acc = vdupq_n_f32(0.f); - for (uint32_t g = 0; g < num_groups; ++g) { - const uint8_t* p_base = packed_interleaved + (nb * num_groups + g) * panel_bytes; - const int8_t* a_grp = act_i8 + static_cast(g) * gs; - const float a_sc = act_scales[g]; - - int32x4_t dot_a = vdupq_n_s32(0); - int32x4_t dot_b = vdupq_n_s32(0); + auto phase_a_group = [&](uint32_t g) { + __fp16 basis[256]; + cactus_quant_transform_hadamard_group(*W, x + static_cast(g) * gs, g, basis); + act_scales[g] = tq_quantize_group_i8(basis, act_i8 + static_cast(g) * gs, gs); + }; - for (uint32_t kb = 0; kb < gs; kb += 16) { - int8x16_t a_v = vld1q_s8(a_grp + kb); - uint8x16_t b0 = vld1q_u8(p_base + (kb / 8 + 0) * 16); - int8x16_t wl0 = vqtbl1q_s8(cb_lut, vandq_u8(b0, lo_mask)); - int8x16_t wh0 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b0, 4)); - dot_a = vdotq_laneq_s32(dot_a, wl0, a_v, 0); - dot_b = vdotq_laneq_s32(dot_b, wh0, a_v, 1); - uint8x16_t b1 = vld1q_u8(p_base + (kb / 8 + 1) * 16); - int8x16_t wl1 = vqtbl1q_s8(cb_lut, vandq_u8(b1, lo_mask)); - int8x16_t wh1 = vqtbl1q_s8(cb_lut, vshrq_n_u8(b1, 4)); - dot_a = vdotq_laneq_s32(dot_a, wl1, a_v, 2); - dot_b = vdotq_laneq_s32(dot_b, wh1, a_v, 3); - } + if (nt <= 1) { + for (uint32_t g = 0; g < num_groups; ++g) phase_a_group(g); + cactus_quant_interleaved4_gemv_blocks(W, packed_interleaved, norms_interleaved, + act_i8, act_scales, cb_lut, cb_scale, + 0, N_blocks, y); + return; + } - int32x4_t dot = vaddq_s32(dot_a, dot_b); - float32x4_t norm = vcvt_f32_f16(vld1_f16(norms_interleaved + (nb * num_groups + g) * 4)); - norm = vmulq_n_f32(norm, cb_scale * a_sc); - acc = vfmaq_f32(acc, vcvtq_f32_s32(dot), norm); - } - vst1_f16(y + nb * 4, vcvt_f16_f32(acc)); - } - }); - (void)N_tiles_4; + cactus_quant_two_phase_run(nt, num_groups, static_cast(n_chunks), phase_a_group, + [&](size_t, uint32_t ck, uint32_t cnt) { + const size_t b0 = static_cast(ck) * 16; + const size_t b1 = std::min(N_blocks, b0 + static_cast(cnt) * 16); + cactus_quant_interleaved4_gemv_blocks(W, packed_interleaved, norms_interleaved, + act_i8, act_scales, cb_lut, cb_scale, + b0, b1, y); + }); } void cactus_quant_3bit_gemv_interleaved( diff --git a/cactus-kernels/tests/test_matmul.cpp b/cactus-kernels/tests/test_matmul.cpp index 7e765b177..74f917974 100644 --- a/cactus-kernels/tests/test_matmul.cpp +++ b/cactus-kernels/tests/test_matmul.cpp @@ -6,6 +6,8 @@ using namespace TestUtils; +static uint8_t unpack_index(const uint8_t* base, uint32_t bits, uint32_t k); + struct SyntheticCQ { uint32_t bits, K, N, group_size, num_groups; std::vector<__fp16> codebook; @@ -150,6 +152,56 @@ struct SyntheticCQ { .norm_f32 = norm_f32_buf.empty() ? nullptr : norm_f32_buf.data(), }; } + + // INTERLEAVED_4ROW encoding: exact inverse of the shipped decoder. Per 8-byte half, + // low nibbles = k 0-3 and high nibbles = k 4-7 of the half's K-range. + std::vector packed_il; + std::vector<__fp16> norms_il; + + void make_interleaved() { + if (bits != 4 || (N % 4) != 0) return; + const uint32_t pgb = cactus_quant_packed_group_bytes(4, group_size); + const size_t NB = N / 4; + packed_il.assign(NB * num_groups * 4 * (size_t)pgb, 0); + norms_il.resize(NB * num_groups * 4); + for (size_t nb = 0; nb < NB; ++nb) + for (uint32_t g = 0; g < num_groups; ++g) { + uint8_t* panel = packed_il.data() + (nb * num_groups + g) * 4 * (size_t)pgb; + for (uint32_t r = 0; r < 4; ++r) { + const size_t n = nb * 4 + r; + const uint8_t* row = packed.data() + (n * num_groups + g) * pgb; + for (uint32_t v = 0; v < group_size / 16; ++v) + for (uint32_t b = 0; b < 4; ++b) { + auto idx = [&](uint32_t k) { return unpack_index(row, 4, k); }; + panel[(2 * v) * 16 + r * 4 + b] = + (uint8_t)(idx(16 * v + b) | (idx(16 * v + 4 + b) << 4)); + panel[(2 * v + 1) * 16 + r * 4 + b] = + (uint8_t)(idx(16 * v + 8 + b) | (idx(16 * v + 12 + b) << 4)); + } + norms_il[(nb * num_groups + g) * 4 + r] = norms[n * num_groups + g]; + } + } + } + + CactusQuantMatrix matrix_interleaved() { + if (packed_il.empty()) make_interleaved(); + return CactusQuantMatrix{ + .bits = bits, .K = K, .N = N, + .group_size = group_size, .num_groups = num_groups, + .flags = CACTUS_QUANT_FLAG_INTERLEAVED_4ROW, + .codebook = codebook.data(), + .input_scale = input_scale.data(), + .input_scale_recip = input_scale_recip.data(), + .norms = norms_il.data(), + .packed_indices = packed_il.data(), + .left_signs = left_signs.data(), + .right_signs = right_signs.data(), + .permutation = permutation.data(), + .rotation = nullptr, + .expanded = nullptr, + .norm_f32 = nullptr, + }; + } }; static void fwht_f32(float* x, uint32_t n) { @@ -502,6 +554,26 @@ void print_mse_report() { } } +// Legacy IL CQ4 vs the FP32 oracle through the real dispatch. +static bool test_cq4_interleaved(double& mse_out, + uint32_t K = 1024, uint32_t N = 192, uint32_t gs = 128) { + SyntheticCQ cq(4, K, N, gs, 777); + CactusQuantMatrix mat = cq.matrix_interleaved(); + + std::mt19937 gen(31); + std::uniform_real_distribution dist(-1.f, 1.f); + std::vector x_f32(K); + for (auto& v : x_f32) v = dist(gen); + std::vector ref(N, 0.f); + cq_reference_gemv_f32(cq, x_f32.data(), ref.data()); + + std::vector<__fp16> x_f16(K), y(N, (__fp16)0); + for (size_t i = 0; i < K; i++) x_f16[i] = (__fp16)x_f32[i]; + cactus_quant_matmul(&mat, x_f16.data(), 1, y.data()); + mse_out = compute_mse(ref.data(), y.data(), N); + return mse_out <= 0.1; +} + int main() { TestRunner runner("Matrix Multiplication"); runner.run_test("matmul_f16", test_matmul_f16()); @@ -509,6 +581,13 @@ int main() { runner.run_test("matmul_cq2", test_cq_correctness(2)); runner.run_test("matmul_cq3", test_cq_correctness(3)); runner.run_test("matmul_cq4", test_cq_correctness(4)); + { + double m1 = 0; + runner.run_test("matmul_cq4_il", test_cq4_interleaved(m1)); + // N=4164 -> 66 chunks incl. a 1-block tail: exercises the multi-thread fused driver. + double m_mt = 0; + runner.run_test("matmul_cq4_il_mt", test_cq4_interleaved(m_mt, 1024, 4164, 128)); + } runner.print_benchmarks_header(); runner.run_bench("benchmarks", run_benchmarks()); print_mse_report();