From d834d44e643681f7b046a22d357335f6f4ff6107 Mon Sep 17 00:00:00 2001 From: jbooth Date: Sat, 26 Sep 2026 13:38:57 +0800 Subject: [PATCH] ggml-cpu: tiled mul_mat for k-quants (#27851) * Added tiled mul_mat. For each mul_mat_one_chunk, quants are unpacked into (max) 256x256 tiles of int8, one routine per quent. Then microkernel computes 16x16 tiles before writing out 256x256 float reults to main memory. Tests/benches in tests/test-tiled-mulmat.cpp. 3-6x speed improvement for large matmul, break even at 4096x64 * 64x4096, 80% performance (net loss) for GEMV. Error rates trivial (order of 1-e04 max, 1-e05 rmse). * Fixes for ARM/windows builds * more windows fixes, ggml-cpu.h isn't visible in MSVC for some reason * unified iqp + tiled on the Q5_K, IQ4_XS set for benchmarking, updated benchmark * Fixed accidental removal of llama_build_and_test(test-backend-ops.cpp) * First integration of iqp code Co-authored-by Bartowski <3266127+bartowski1182@users.noreply.github.com> * Cleaning up declaration of iq unpacking helpers to align with the bit unpackers * Removed iqp path * Fix cross-platform warnings * Disabling benchmarks unless explicitly enabled * Fix backend_init for DLL-based builds, add self and bartowski to CODEOWNERS for tiled * Put benchmarks behind a flag * kernel fix for AVX2, iq quants * Fix for asan, leaking memory in test-tiled-mulmat and avoid stack use after return * guarding env flags with std::call_once * Simplified repacking for VNNI to a single call per macrotile * No threadlocals anymore, aligned wdata access * Doing aligned reads since we ensure alignment with padding in wdata * Eliminated per-thread gather of Q8_K rows in mul_mat_id, we now gather/repack in a single pass. Repack method now takes pointer array to support both dense/normal and mmid paths. Interface with ggml-cpu.c simplified as a result * Unified/simplified dispatch and support checks. Put details on wdata needed inside the kernel.h body, simplified interactions with ggml-cpu.c. * Cleanup includes and whitespace, update src1_repack to return false if we don't need a special repack, so the common case is handled by driver * Better detection of win32 and additional whitespace fixes * Gating fuzz tests behind a parameter and some extra prints to try and fix slow CI hosts * Optimized AVX2 kernel * Changed interleave format and added ability to interleave in-place after dequant * Repacks now happen in-place, 16x64 microtiles are independent of each other * Only repack rows in groups of 16 as they're needed. Save work in low n_rows cases and optimize L1 usage in other cases * Use long panels for memory-bound regime (M <= 16), reintroduce IQP path for benchmarks * Fix unused warnings and cleanup. Improved IQ dequantization speed. * Removed separate process benchmarks * Revert "Removed separate process benchmarks" This reverts commit 0688cf43d5e37ec7a52828c5576df563514fe760. * AVX2 optimizations and guards for tests on windows * Removed temp perf harness * Remove perf-mulmat from build * Removed IQP path, simplified tests to not use sub processes * Cleaning up alignment of wdata * Whitespace fixes and aligning L2 workspace to clean 512kb boundaries * Update ggml/src/ggml-cpu/tiled/tiled-kernel.cpp Co-authored-by: Georgi Gerganov * Cleanup merge-duplicated declaration of test-backend-ops target * Undo accidental line deletion in ggml.c --------- Co-authored-by: Georgi Gerganov --- CODEOWNERS | 2 +- ggml/src/ggml-cpu/CMakeLists.txt | 6 +- ggml/src/ggml-cpu/ggml-cpu.c | 48 +- ggml/src/ggml-cpu/iqp.cpp | 1253 ---------------------- ggml/src/ggml-cpu/iqp.h | 39 - ggml/src/ggml-cpu/tiled/tiled-kernel.cpp | 719 +++++++++++++ ggml/src/ggml-cpu/tiled/tiled-kernel.h | 180 ++++ ggml/src/ggml-cpu/tiled/tiled.cpp | 1245 +++++++++++++++++++++ ggml/src/ggml-cpu/tiled/tiled.h | 31 + tests/CMakeLists.txt | 1 + tests/test-tiled-mulmat.cpp | 852 +++++++++++++++ 11 files changed, 3050 insertions(+), 1326 deletions(-) delete mode 100644 ggml/src/ggml-cpu/iqp.cpp delete mode 100644 ggml/src/ggml-cpu/iqp.h create mode 100644 ggml/src/ggml-cpu/tiled/tiled-kernel.cpp create mode 100644 ggml/src/ggml-cpu/tiled/tiled-kernel.h create mode 100644 ggml/src/ggml-cpu/tiled/tiled.cpp create mode 100644 ggml/src/ggml-cpu/tiled/tiled.h create mode 100644 tests/test-tiled-mulmat.cpp diff --git a/CODEOWNERS b/CODEOWNERS index 0fbfcfdd12..36d817fdf9 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -57,7 +57,7 @@ /ggml/src/ggml-cann/ @ggml-org/ggml-cann /ggml/src/ggml-common.h @ggerganov /ggml/src/ggml-cpu/ @ggerganov -/ggml/src/ggml-cpu/iqp.* @bartowski1182 +/ggml/src/ggml-cpu/tiled/ @jbooth @bartowski1182 /ggml/src/ggml-cpu/spacemit/ @alex-spacemit /ggml/src/ggml-cuda/ @ggml-org/ggml-cuda /ggml/src/ggml-cuda/vendors/hip.h @IMbackK diff --git a/ggml/src/ggml-cpu/CMakeLists.txt b/ggml/src/ggml-cpu/CMakeLists.txt index 1c7338eea4..2765c27276 100644 --- a/ggml/src/ggml-cpu/CMakeLists.txt +++ b/ggml/src/ggml-cpu/CMakeLists.txt @@ -31,8 +31,6 @@ function(ggml_add_cpu_backend_variant_impl tag_name) ggml-cpu/ggml-cpu.cpp ggml-cpu/repack.cpp ggml-cpu/repack.h - ggml-cpu/iqp.cpp - ggml-cpu/iqp.h ggml-cpu/hbm.cpp ggml-cpu/hbm.h ggml-cpu/quants.c @@ -43,6 +41,10 @@ function(ggml_add_cpu_backend_variant_impl tag_name) ggml-cpu/amx/amx.h ggml-cpu/amx/mmq.cpp ggml-cpu/amx/mmq.h + ggml-cpu/tiled/tiled.cpp + ggml-cpu/tiled/tiled.h + ggml-cpu/tiled/tiled-kernel.cpp + ggml-cpu/tiled/tiled-kernel.h ggml-cpu/ggml-cpu-impl.h ggml-cpu/common.h ggml-cpu/binary-ops.h diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 8bb0ff7bc3..046f17ca09 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -4,7 +4,6 @@ #include "ggml-backend-impl.h" #include "ggml-backend.h" #include "traits.h" -#include "iqp.h" #include "ggml-cpu-impl.h" #include "ggml-impl.h" #include "quants.h" @@ -15,6 +14,7 @@ #include "ops.h" #include "ggml.h" #include "common.h" +#include "tiled/tiled.h" #if defined(_MSC_VER) || defined(__MINGW32__) #include // using malloc.h with MSC/MINGW @@ -1265,6 +1265,11 @@ void ggml_compute_forward_mul_mat( return; } + // If tiled is supported, it will execute the full op here and we return + if (ggml_compute_forward_mul_mat_tiled(params, dst)) { + return; + } + GGML_TENSOR_BINARY_OP_LOCALS const int ith = params->ith; @@ -1372,13 +1377,6 @@ UseGgmlGemm1:; ggml_barrier(params->threadpool); - // IQ panel gemm (see iqp.h) - must come after the barrier above, it consumes the q8_K rows - // of src1 from the work buffer - if (ggml_cpu_iqp_supports_mul_mat(dst) && !params->use_ref) { - ggml_compute_forward_mul_mat_iqp(params, dst); - return; - } - #if GGML_USE_LLAMAFILE if (src1->type != vec_dot_type) { const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; @@ -1596,15 +1594,9 @@ static void ggml_compute_forward_mul_mat_id( char (*atomic_current_chunk)[CACHE_LINE_SIZE] = // [n_as] incr_ptr_aligned(&wdata_cur, CACHE_LINE_SIZE * n_as, CACHE_LINE_SIZE); - // IQ panel gemm (see iqp.h); per expert eligibility is decided below, but the work buffer is - // reserved for the whole node (ggml_graph_plan sizes it without params, use_ref only skips the dispatch) - const bool iqp = ggml_cpu_iqp_supports_mul_mat_id(dst) && !params->use_ref; - - char * iqp_panels = NULL; - - if (iqp) { - iqp_panels = incr_ptr_aligned(&wdata_cur, nth * ggml_cpu_iqp_scratch_size(dst), 64); - } + // Tiled matmul (see tiled.h); per-thread work buffers, 0 bytes when disabled. The + // reservation is unconditional, the per expert eligibility is decided at dispatch time + char * tiled_scratch = incr_ptr_aligned(&wdata_cur, ggml_tiled_wdata_size(nth, dst), 64); GGML_ASSERT(params->wsize >= (size_t)((char *) wdata_cur - (char *) params->wdata)); @@ -1677,10 +1669,8 @@ static void ggml_compute_forward_mul_mat_id( continue; } - if (iqp && ggml_cpu_iqp_mul_mat_id_min_batch(cne1)) { - ggml_compute_forward_mul_mat_id_iqp(params, dst, cur_a, cne1, (const int32_t *) &MMID_MATRIX_ROW(cur_a, 0), - iqp_panels); - + // tiled takes over if profitable for this expert (see tiled.h) + if (ggml_compute_forward_mul_mat_id_tiled(params, dst, cur_a, cne1, (const int32_t *) &MMID_MATRIX_ROW(cur_a, 0), tiled_scratch)) { continue; } @@ -2887,15 +2877,12 @@ struct ggml_cplan ggml_graph_plan( case GGML_OP_MUL_MAT: { const enum ggml_type vec_dot_type = type_traits_cpu[node->src[0]->type].vec_dot_type; - if (node->src[1]->type != vec_dot_type) { cur = ggml_row_size(vec_dot_type, ggml_nelements(node->src[1])); } - - // the IQ panel path needs one scratch panel per thread past the q8_K rows - if (ggml_cpu_iqp_supports_mul_mat(node)) { - cur = GGML_PAD(cur, 64) + n_tasks * ggml_cpu_iqp_scratch_size(node); - } + // Workspace for tiled (see tiled.h) + cur = GGML_PAD(cur, 64); + cur += ggml_tiled_wdata_size(n_tasks, node); } break; case GGML_OP_MUL_MAT_ID: { @@ -2915,10 +2902,9 @@ struct ggml_cplan ggml_graph_plan( cur += n_as*ids->ne[0]*ids->ne[1]*sizeof(struct mmid_row_mapping) + sizeof(int64_t); // atomic_current_chunk cur += CACHE_LINE_SIZE*n_as + CACHE_LINE_SIZE; - // the IQ panel path needs one scratch panel per thread on top of that - if (ggml_cpu_iqp_supports_mul_mat_id(node)) { - cur += n_tasks * ggml_cpu_iqp_scratch_size(node) + 64; - } + // Workspace for tiled (see tiled.h) + cur = GGML_PAD(cur, 64); + cur += ggml_tiled_wdata_size(n_tasks, node); } break; case GGML_OP_OUT_PROD: { diff --git a/ggml/src/ggml-cpu/iqp.cpp b/ggml/src/ggml-cpu/iqp.cpp deleted file mode 100644 index b9201db381..0000000000 --- a/ggml/src/ggml-cpu/iqp.cpp +++ /dev/null @@ -1,1253 +0,0 @@ -#define GGML_COMMON_IMPL_CPP -#define GGML_COMMON_DECL_CPP -#include "ggml-common.h" - -#include "ggml-impl.h" -#include "ggml-cpu.h" -#include "ggml-cpu-impl.h" -#include "simd-mappings.h" -#include "traits.h" - -#include -#include -#include - -#include "iqp.h" - -#define UNUSED GGML_UNUSED - -// smallest src1 batch for which the decode pays for itself -#define GGML_IQP_MIN_BATCH 8 - -// same, per expert, for MUL_MAT_ID -#define GGML_IQP_MIN_BATCH_ID 8 - -bool ggml_cpu_iqp_mul_mat_id_min_batch(int64_t cne1) { - return cne1 >= GGML_IQP_MIN_BATCH_ID; -} - -// src0 rows interleaved per panel -#define IQP_NB_ROWS 8 - -#define IQP_SB_SIZE 16 // weights per sub-block -#define IQP_NSB (QK_K / IQP_SB_SIZE) // sub-blocks per super-block - -// one super-block of a grid based IQ type decoded to int8, 8 rows interleaved: -// dfac[row] * iscales[sb*8 + row] * qs is bit identical to dequantize_row_iq* -struct block_iqp_x8 { - float dfac[8]; // f32 super-block scale, d * 2^-k - int32_t bias[8]; // 128 * sum(qs * iscale), see GGML_IQP_USE_BIAS - int8_t iscales[IQP_NSB * 8]; // integer sub-block scales, in [-32, 31] - int8_t qs[QK_K * 8]; // qs[sb*128 + g*32 + row*4 + k] = column sb*16 + g*4 + k -}; - -static_assert(sizeof(block_iqp_x8) == 8 * sizeof(float) + 8 * sizeof(int32_t) + IQP_NSB * 8 + QK_K * 8, - "wrong iqp_x8 block size/padding"); - -// feed the activations to VNNI as unsigned bytes (y + 128) and correct with bias[]; without VNNI the kernels use the maddubs sign trick instead and bias[] is not filled -#if defined(__AVX2__) && ((defined(__AVX512VNNI__) && defined(__AVX512VL__)) || defined(__AVXVNNI__)) -# define GGML_IQP_USE_BIAS 1 -#else -# define GGML_IQP_USE_BIAS 0 -#endif - -static inline size_t ggml_cpu_iqp_row_size(const struct ggml_tensor * dst) { - return ggml_row_size(GGML_TYPE_Q8_K, dst->src[1]->ne[0]); -} - -// the low 7 bits of v are the first 7 signs and the 8th is their parity (cf. unpack_ksigns in the CUDA backend) -static inline uint8_t iqp_unpack_ksigns(uint32_t v) { - uint32_t p = v ^ (v >> 4); - - p ^= p >> 2; - p ^= p >> 1; - - return (uint8_t) (v ^ ((p & 1) << 7)); -} - -#if defined(__AVX2__) - -// 0xFF in every byte whose sign bit is set; sv holds each sign byte broadcast over the 8 bytes it governs -static inline __m256i iqp_sign_mask(__m256i sv) { - const __m256i sel = _mm256_set1_epi64x((int64_t) 0x8040201008040201ULL); - -# if defined(__GFNI__) - // computes the and + compare in one instruction - return _mm256_gf2p8affine_epi64_epi8(sel, sv, 0); -# else - return _mm256_cmpeq_epi8(_mm256_and_si256(sv, sel), sel); -# endif -} - -// signs holds four sign bytes, byte l governing values 8*l .. 8*l+7 - spread each over its 8 lanes -static inline __m256i iqp_sign_bytes(uint32_t signs) { - const __m256i bcast = _mm256_setr_epi8(0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, // - 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3); - - return _mm256_shuffle_epi8(_mm256_set1_epi32((int32_t) signs), bcast); -} - -// x ^ m - m negates the lanes where m is 0xFF -static inline __m256i iqp_apply_signs(__m256i x, __m256i m) { - return _mm256_sub_epi8(_mm256_xor_si256(x, m), m); -} - -#endif - -// 32 values from four 8 byte grid entries, sign byte l of signs applied to group l -static inline void iqp_store_signed_x8(int8_t * GGML_RESTRICT dst, - uint64_t g0, - uint64_t g1, - uint64_t g2, - uint64_t g3, - uint32_t signs) { -#if defined(__AVX2__) - const __m256i g = _mm256_set_epi64x((int64_t) g3, (int64_t) g2, (int64_t) g1, (int64_t) g0); - const __m256i m = iqp_sign_mask(iqp_sign_bytes(signs)); - - _mm256_storeu_si256((__m256i *) dst, iqp_apply_signs(g, m)); -#else - const uint64_t g[4] = { g0, g1, g2, g3 }; - - for (int l = 0; l < 4; ++l) { - const uint8_t * grid = (const uint8_t *) &g[l]; - const uint8_t s = (uint8_t) (signs >> 8 * l); - - for (int j = 0; j < 8; ++j) { - dst[8 * l + j] = s & kmask_iq2xs[j] ? -grid[j] : grid[j]; - } - } -#endif -} - -// same, but the eight values of group l come from two 4 byte grid entries -static inline void iqp_store_signed_x4(int8_t * GGML_RESTRICT dst, - uint32_t g0a, - uint32_t g0b, - uint32_t g1a, - uint32_t g1b, - uint32_t g2a, - uint32_t g2b, - uint32_t g3a, - uint32_t g3b, - uint32_t signs) { -#if defined(__AVX2__) - const __m256i g = _mm256_setr_epi32((int32_t) g0a, (int32_t) g0b, (int32_t) g1a, (int32_t) g1b, (int32_t) g2a, - (int32_t) g2b, (int32_t) g3a, (int32_t) g3b); - const __m256i m = iqp_sign_mask(iqp_sign_bytes(signs)); - - _mm256_storeu_si256((__m256i *) dst, iqp_apply_signs(g, m)); -#else - const uint32_t ga[4] = { g0a, g1a, g2a, g3a }; - const uint32_t gb[4] = { g0b, g1b, g2b, g3b }; - - for (int l = 0; l < 4; ++l) { - const uint8_t * grid1 = (const uint8_t *) &ga[l]; - const uint8_t * grid2 = (const uint8_t *) &gb[l]; - const uint8_t s = (uint8_t) (signs >> 8 * l); - - for (int j = 0; j < 4; ++j) { - dst[8 * l + j + 0] = s & kmask_iq2xs[j + 0] ? -grid1[j] : grid1[j]; - dst[8 * l + j + 4] = s & kmask_iq2xs[j + 4] ? -grid2[j] : grid2[j]; - } - } -#endif -} - -// 32 values of 8 * grid + delta from four 8 byte grid entries (grid bytes are in {-1, 0, 1}), byte l of deltas applying to group l -static inline void iqp_store_iq1_x8(int8_t * GGML_RESTRICT dst, - uint64_t g0, - uint64_t g1, - uint64_t g2, - uint64_t g3, - uint32_t deltas) { -#if defined(__AVX2__) - __m256i g = _mm256_set_epi64x((int64_t) g3, (int64_t) g2, (int64_t) g1, (int64_t) g0); - - // no byte shift in AVX2 - g = _mm256_add_epi8(g, g); - g = _mm256_add_epi8(g, g); - g = _mm256_add_epi8(g, g); - - _mm256_storeu_si256((__m256i *) dst, _mm256_add_epi8(g, iqp_sign_bytes(deltas))); -#else - const uint64_t g[4] = { g0, g1, g2, g3 }; - - for (int l = 0; l < 4; ++l) { - const int8_t * grid = (const int8_t *) &g[l]; - const int8_t delta = (int8_t) (deltas >> 8 * l); - - for (int j = 0; j < 8; ++j) { - dst[8 * l + j] = 8 * grid[j] + delta; - } - } -#endif -} - -// 32 values from 16 packed nibbles through the kvalues_iq4nl lookup: low nibbles first, then high -static inline void iqp_store_iq4_x32(int8_t * GGML_RESTRICT dst, const uint8_t * GGML_RESTRICT qs) { -#if defined(__AVX2__) - const __m128i q = _mm_loadu_si128((const __m128i *) qs); - const __m128i lut = _mm_loadu_si128((const __m128i *) kvalues_iq4nl); - const __m128i m4 = _mm_set1_epi8(0xf); - - _mm_storeu_si128((__m128i *) (dst + 0), _mm_shuffle_epi8(lut, _mm_and_si128(q, m4))); - _mm_storeu_si128((__m128i *) (dst + 16), _mm_shuffle_epi8(lut, _mm_and_si128(_mm_srli_epi16(q, 4), m4))); -#else - for (int j = 0; j < 16; ++j) { - dst[j + 0] = kvalues_iq4nl[qs[j] & 0xf]; - dst[j + 16] = kvalues_iq4nl[qs[j] >> 4]; - } -#endif -} - -#if GGML_IQP_USE_BIAS - -// sum of qs * iscale over one super-block, at most 256 * 127 * 32 = 1.04e6 -static inline int32_t iqp_weighted_sum(const int8_t * GGML_RESTRICT vals, const int8_t * GGML_RESTRICT iscales) { -#if defined(__AVX2__) - static_assert(IQP_SB_SIZE == 16, "the vector path folds two sub-blocks per 32 byte load"); - - const __m256i ones8 = _mm256_set1_epi8(1); - const __m256i ones16 = _mm256_set1_epi16(1); - - __m256i acc = _mm256_setzero_si256(); - - for (int i = 0; i < QK_K / 32; ++i) { - // sum groups of 4 bytes into int32, the low four lanes cover sub-block 2*i and the high four 2*i + 1 - const __m256i v = _mm256_loadu_si256((const __m256i *) (vals + 32 * i)); - const __m256i p = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, v), ones16); - - const __m256i s = _mm256_set_m128i(_mm_set1_epi32(iscales[2 * i + 1]), _mm_set1_epi32(iscales[2 * i + 0])); - - acc = _mm256_add_epi32(acc, _mm256_mullo_epi32(p, s)); - } - - __m128i sum = _mm_add_epi32(_mm256_castsi256_si128(acc), _mm256_extracti128_si256(acc, 1)); - - sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, _MM_SHUFFLE(1, 0, 3, 2))); - sum = _mm_add_epi32(sum, _mm_shuffle_epi32(sum, _MM_SHUFFLE(2, 3, 0, 1))); - - return _mm_cvtsi128_si32(sum); -#else - int32_t wsum = 0; - - for (int sb = 0; sb < IQP_NSB; ++sb) { - int32_t vsum = 0; - - for (int k = 0; k < IQP_SB_SIZE; ++k) { - vsum += vals[sb * IQP_SB_SIZE + k]; - } - - wsum += iscales[sb] * vsum; - } - - return wsum; -#endif -} - -#endif // GGML_IQP_USE_BIAS - -static void iqp_decode_iq2_xxs(const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - const block_iq2_xxs * x = (const block_iq2_xxs *) vx; - - // db = d * (0.5 + ls) * 0.25 = (d / 8) * (2 * ls + 1), ls 4 bit - *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.125f; - - uint32_t aux32[2]; - const uint8_t * aux8 = (const uint8_t *) aux32; - - for (int ib32 = 0; ib32 < QK_K / 32; ++ib32) { - memcpy(aux32, x->qs + 4 * ib32, 2 * sizeof(uint32_t)); - const int8_t ls = (int8_t) (2 * (aux32[1] >> 28) + 1); - - iscales[2 * ib32 + 0] = ls; - iscales[2 * ib32 + 1] = ls; - - const uint32_t signs = (uint32_t) iqp_unpack_ksigns((aux32[1] >> 0) & 127) | - (uint32_t) iqp_unpack_ksigns((aux32[1] >> 7) & 127) << 8 | - (uint32_t) iqp_unpack_ksigns((aux32[1] >> 14) & 127) << 16 | - (uint32_t) iqp_unpack_ksigns((aux32[1] >> 21) & 127) << 24; - - iqp_store_signed_x8(vals + 32 * ib32, iq2xxs_grid[aux8[0]], iq2xxs_grid[aux8[1]], iq2xxs_grid[aux8[2]], - iq2xxs_grid[aux8[3]], signs); - } -} - -static void iqp_decode_iq2_xs(const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - const block_iq2_xs * x = (const block_iq2_xs *) vx; - - *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.125f; - - for (int ib32 = 0; ib32 < QK_K / 32; ++ib32) { - iscales[2 * ib32 + 0] = (int8_t) (2 * (x->scales[ib32] & 0xf) + 1); - iscales[2 * ib32 + 1] = (int8_t) (2 * (x->scales[ib32] >> 4) + 1); - - const uint16_t * q = x->qs + 4 * ib32; - - const uint32_t signs = (uint32_t) iqp_unpack_ksigns(q[0] >> 9) | (uint32_t) iqp_unpack_ksigns(q[1] >> 9) << 8 | - (uint32_t) iqp_unpack_ksigns(q[2] >> 9) << 16 | - (uint32_t) iqp_unpack_ksigns(q[3] >> 9) << 24; - - iqp_store_signed_x8(vals + 32 * ib32, iq2xs_grid[q[0] & 511], iq2xs_grid[q[1] & 511], iq2xs_grid[q[2] & 511], - iq2xs_grid[q[3] & 511], signs); - } -} - -static void iqp_decode_iq2_s(const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - const block_iq2_s * x = (const block_iq2_s *) vx; - - const uint8_t * qs = x->qs; - const uint8_t * qh = x->qh; - const uint8_t * signs = qs + QK_K / 8; - - *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.125f; - - for (int ib32 = 0; ib32 < QK_K / 32; ++ib32) { - iscales[2 * ib32 + 0] = (int8_t) (2 * (x->scales[ib32] & 0xf) + 1); - iscales[2 * ib32 + 1] = (int8_t) (2 * (x->scales[ib32] >> 4) + 1); - - const uint32_t sbits = - (uint32_t) signs[0] | (uint32_t) signs[1] << 8 | (uint32_t) signs[2] << 16 | (uint32_t) signs[3] << 24; - - iqp_store_signed_x8(vals + 32 * ib32, iq2s_grid[qs[0] | (qh[ib32] << 8 & 0x300)], - iq2s_grid[qs[1] | (qh[ib32] << 6 & 0x300)], iq2s_grid[qs[2] | (qh[ib32] << 4 & 0x300)], - iq2s_grid[qs[3] | (qh[ib32] << 2 & 0x300)], sbits); - qs += 4; - signs += 4; - } -} - -static void iqp_decode_iq3_xxs(const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - const block_iq3_xxs * x = (const block_iq3_xxs *) vx; - - const uint8_t * qs = x->qs; - const uint8_t * scales_and_signs = qs + QK_K / 4; - - // db = d * (0.5 + ls) * 0.5 = (d / 4) * (2 * ls + 1), ls 4 bit - *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.25f; - - uint32_t aux32; - - for (int ib32 = 0; ib32 < QK_K / 32; ++ib32) { - memcpy(&aux32, scales_and_signs + 4 * ib32, sizeof(uint32_t)); - const int8_t ls = (int8_t) (2 * (aux32 >> 28) + 1); - - iscales[2 * ib32 + 0] = ls; - iscales[2 * ib32 + 1] = ls; - - const uint32_t signs = (uint32_t) iqp_unpack_ksigns((aux32 >> 0) & 127) | - (uint32_t) iqp_unpack_ksigns((aux32 >> 7) & 127) << 8 | - (uint32_t) iqp_unpack_ksigns((aux32 >> 14) & 127) << 16 | - (uint32_t) iqp_unpack_ksigns((aux32 >> 21) & 127) << 24; - - iqp_store_signed_x4(vals + 32 * ib32, iq3xxs_grid[qs[0]], iq3xxs_grid[qs[1]], iq3xxs_grid[qs[2]], - iq3xxs_grid[qs[3]], iq3xxs_grid[qs[4]], iq3xxs_grid[qs[5]], iq3xxs_grid[qs[6]], - iq3xxs_grid[qs[7]], signs); - qs += 8; - } -} - -static void iqp_decode_iq3_s(const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - const block_iq3_s * x = (const block_iq3_s *) vx; - - const uint8_t * qs = x->qs; - const uint8_t * qh = x->qh; - const uint8_t * signs = x->signs; - - // db = d * (1 + 2 * ls), ls 4 bit - *dfac = GGML_CPU_FP16_TO_FP32(x->d); - - int k = 0; - - for (int ib32 = 0; ib32 < QK_K / 32; ib32 += 2) { - const int8_t db1 = (int8_t) (1 + 2 * (x->scales[ib32 / 2] & 0xf)); - const int8_t db2 = (int8_t) (1 + 2 * (x->scales[ib32 / 2] >> 4)); - - iscales[2 * ib32 + 0] = db1; - iscales[2 * ib32 + 1] = db1; - iscales[2 * ib32 + 2] = db2; - iscales[2 * ib32 + 3] = db2; - - for (int h = 0; h < 2; ++h) { - const uint32_t sbits = - (uint32_t) signs[0] | (uint32_t) signs[1] << 8 | (uint32_t) signs[2] << 16 | (uint32_t) signs[3] << 24; - - iqp_store_signed_x4(vals + k, iq3s_grid[qs[0] | ((qh[h] << 8) & 256)], - iq3s_grid[qs[1] | ((qh[h] << 7) & 256)], iq3s_grid[qs[2] | ((qh[h] << 6) & 256)], - iq3s_grid[qs[3] | ((qh[h] << 5) & 256)], iq3s_grid[qs[4] | ((qh[h] << 4) & 256)], - iq3s_grid[qs[5] | ((qh[h] << 3) & 256)], iq3s_grid[qs[6] | ((qh[h] << 2) & 256)], - iq3s_grid[qs[7] | ((qh[h] << 1) & 256)], sbits); - - k += 32; - qs += 8; - signs += 4; - } - qh += 2; - } -} - -// dequantize_row_iq1_* computes y = dl * (grid[j] + delta) with delta = +-1/8, so the panel stores 8 * grid[j] +- 1 and folds the /8 into dfac -static void iqp_decode_iq1_s(const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - const block_iq1_s * x = (const block_iq1_s *) vx; - - const uint8_t * qs = x->qs; - const uint16_t * qh = x->qh; - - // dl = d * (2 * ls + 1) * 0.125, ls 3 bit - *dfac = GGML_CPU_FP16_TO_FP32(x->d) * 0.125f; - - for (int ib = 0; ib < QK_K / 32; ++ib) { - const int8_t dl = (int8_t) (2 * ((qh[ib] >> 12) & 7) + 1); - const int8_t delta = qh[ib] & 0x8000 ? -1 : 1; - - iscales[2 * ib + 0] = dl; - iscales[2 * ib + 1] = dl; - - iqp_store_iq1_x8(vals + 32 * ib, iq1s_grid[qs[0] | (((qh[ib] >> 0) & 7) << 8)], - iq1s_grid[qs[1] | (((qh[ib] >> 3) & 7) << 8)], iq1s_grid[qs[2] | (((qh[ib] >> 6) & 7) << 8)], - iq1s_grid[qs[3] | (((qh[ib] >> 9) & 7) << 8)], ((uint8_t) delta) * 0x01010101u); - qs += 4; - } -} - -static void iqp_decode_iq1_m(const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - const block_iq1_m * x = (const block_iq1_m *) vx; - - // block_iq1_m has no d field - the fp16 super-block scale is spread over the top nibbles of the four scale words - const uint16_t * sc = (const uint16_t *) x->scales; - - iq1m_scale_t scale; - scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); - - *dfac = GGML_CPU_FP16_TO_FP32(scale.f16) * 0.125f; - - const uint8_t * qs = x->qs; - const uint8_t * qh = x->qh; - - for (int ib = 0; ib < QK_K / 32; ++ib) { - iscales[2 * ib + 0] = (int8_t) (2 * ((sc[ib / 2] >> (6 * (ib % 2) + 0)) & 0x7) + 1); - iscales[2 * ib + 1] = (int8_t) (2 * ((sc[ib / 2] >> (6 * (ib % 2) + 3)) & 0x7) + 1); - - const uint16_t idx[4] = { - (uint16_t) (qs[0] | ((qh[0] << 8) & 0x700)), - (uint16_t) (qs[1] | ((qh[0] << 4) & 0x700)), - (uint16_t) (qs[2] | ((qh[1] << 8) & 0x700)), - (uint16_t) (qs[3] | ((qh[1] << 4) & 0x700)), - }; - const uint32_t deltas = (uint32_t) (qh[0] & 0x08 ? 0xff : 0x01) | (uint32_t) (qh[0] & 0x80 ? 0xff : 0x01) << 8 | - (uint32_t) (qh[1] & 0x08 ? 0xff : 0x01) << 16 | - (uint32_t) (qh[1] & 0x80 ? 0xff : 0x01) << 24; - - iqp_store_iq1_x8(vals + 32 * ib, iq1s_grid[idx[0]], iq1s_grid[idx[1]], iq1s_grid[idx[2]], iq1s_grid[idx[3]], - deltas); - qs += 4; - qh += 2; - } -} - -static void iqp_decode_iq4_xs(const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - const block_iq4_xs * x = (const block_iq4_xs *) vx; - - const uint8_t * qs = x->qs; - - // dl = d * (ls - 32), ls 6 bit, so the integer scale is in [-32, 31] - *dfac = GGML_CPU_FP16_TO_FP32(x->d); - - for (int ib = 0; ib < QK_K / 32; ++ib) { - const int ls = ((x->scales_l[ib / 2] >> 4 * (ib % 2)) & 0xf) | (((x->scales_h >> 2 * ib) & 3) << 4); - const int8_t dl = (int8_t) (ls - 32); - - iscales[2 * ib + 0] = dl; - iscales[2 * ib + 1] = dl; - - iqp_store_iq4_x32(vals + 32 * ib, qs); - qs += 16; - } -} - -// expanded by the eligibility test and the decode dispatch -#define IQP_TYPE_LIST(T) \ - T(IQ2_XXS, iq2_xxs) \ - T(IQ2_XS, iq2_xs) \ - T(IQ2_S, iq2_s) \ - T(IQ3_XXS, iq3_xxs) \ - T(IQ3_S, iq3_s) \ - T(IQ1_S, iq1_s) \ - T(IQ1_M, iq1_m) \ - T(IQ4_XS, iq4_xs) - -static bool iqp_decode_superblock(enum ggml_type type, - const void * GGML_RESTRICT vx, - int8_t * GGML_RESTRICT vals, - int8_t * GGML_RESTRICT iscales, - float * GGML_RESTRICT dfac) { - switch (type) { -#define IQP_CASE(E, name) \ - case GGML_TYPE_##E: \ - iqp_decode_##name(vx, vals, iscales, dfac); \ - return true; - IQP_TYPE_LIST(IQP_CASE) -#undef IQP_CASE - default: - return false; - } -} - -#if defined(__AVX2__) - -// 8x8 int32 transpose of the 32 column group starting at column off -static inline void iqp_interleave_x8(int8_t * GGML_RESTRICT dst, const int8_t (*vals)[QK_K], int off) { - static_assert(IQP_NB_ROWS == 8, "the transpose is 8x8"); - - __m256i v[IQP_NB_ROWS]; - - for (int r = 0; r < IQP_NB_ROWS; ++r) { - v[r] = _mm256_loadu_si256((const __m256i *) (vals[r] + off)); - } - - // pair rows into dword couples, then into qword quadruples, then swap the 128 bit lanes - const __m256i a0 = _mm256_unpacklo_epi32(v[0], v[1]); - const __m256i a1 = _mm256_unpackhi_epi32(v[0], v[1]); - const __m256i a2 = _mm256_unpacklo_epi32(v[2], v[3]); - const __m256i a3 = _mm256_unpackhi_epi32(v[2], v[3]); - const __m256i a4 = _mm256_unpacklo_epi32(v[4], v[5]); - const __m256i a5 = _mm256_unpackhi_epi32(v[4], v[5]); - const __m256i a6 = _mm256_unpacklo_epi32(v[6], v[7]); - const __m256i a7 = _mm256_unpackhi_epi32(v[6], v[7]); - - const __m256i b0 = _mm256_unpacklo_epi64(a0, a2); - const __m256i b1 = _mm256_unpackhi_epi64(a0, a2); - const __m256i b2 = _mm256_unpacklo_epi64(a1, a3); - const __m256i b3 = _mm256_unpackhi_epi64(a1, a3); - const __m256i b4 = _mm256_unpacklo_epi64(a4, a6); - const __m256i b5 = _mm256_unpackhi_epi64(a4, a6); - const __m256i b6 = _mm256_unpacklo_epi64(a5, a7); - const __m256i b7 = _mm256_unpackhi_epi64(a5, a7); - - _mm256_storeu_si256((__m256i *) (dst + 0 * 32), _mm256_permute2x128_si256(b0, b4, 0x20)); - _mm256_storeu_si256((__m256i *) (dst + 1 * 32), _mm256_permute2x128_si256(b1, b5, 0x20)); - _mm256_storeu_si256((__m256i *) (dst + 2 * 32), _mm256_permute2x128_si256(b2, b6, 0x20)); - _mm256_storeu_si256((__m256i *) (dst + 3 * 32), _mm256_permute2x128_si256(b3, b7, 0x20)); - _mm256_storeu_si256((__m256i *) (dst + 4 * 32), _mm256_permute2x128_si256(b0, b4, 0x31)); - _mm256_storeu_si256((__m256i *) (dst + 5 * 32), _mm256_permute2x128_si256(b1, b5, 0x31)); - _mm256_storeu_si256((__m256i *) (dst + 6 * 32), _mm256_permute2x128_si256(b2, b6, 0x31)); - _mm256_storeu_si256((__m256i *) (dst + 7 * 32), _mm256_permute2x128_si256(b3, b7, 0x31)); -} - -#endif - -// decode IQP_NB_ROWS consecutive source rows (starting at src, row stride nb01) into a panel of nblocks block_iqp_x8 -static void iqp_decode_panel_8(enum ggml_type type, - const char * GGML_RESTRICT src, - size_t nb01, - int64_t nblocks, - block_iqp_x8 * GGML_RESTRICT dst) { - const size_t bsize = ggml_type_size(type); - - int8_t vals[IQP_NB_ROWS][QK_K]; - int8_t iscales[IQP_NB_ROWS][IQP_NSB]; - float dfac[IQP_NB_ROWS]; - - for (int64_t x = 0; x < nblocks; x++) { - for (int r = 0; r < IQP_NB_ROWS; r++) { - const char * blk = src + r * nb01 + x * bsize; - - const bool ok = iqp_decode_superblock(type, blk, vals[r], iscales[r], &dfac[r]); - GGML_ASSERT(ok); - -#ifdef GGML_IQP_VERIFY - // check that the panel reproduces the reference dequantization bit exactly - float ref[QK_K]; - ggml_get_type_traits(type)->to_float(blk, ref, QK_K); - for (int j = 0; j < QK_K; j++) { - const float scale = dfac[r] * iscales[r][j / IQP_SB_SIZE]; - GGML_ASSERT(scale * vals[r][j] == ref[j]); - } -#endif - } - - for (int r = 0; r < IQP_NB_ROWS; r++) { - dst->dfac[r] = dfac[r]; - - for (int sb = 0; sb < IQP_NSB; sb++) { - dst->iscales[sb * IQP_NB_ROWS + r] = iscales[r][sb]; - } - -#if GGML_IQP_USE_BIAS - dst->bias[r] = 128 * iqp_weighted_sum(vals[r], iscales[r]); -#endif - } - -#if defined(__AVX2__) - for (int grp = 0; grp < QK_K / 32; grp++) { - iqp_interleave_x8(dst->qs + grp * 256, vals, grp * 32); - } -#else - for (int r = 0; r < IQP_NB_ROWS; r++) { - for (int sb = 0; sb < IQP_NSB; sb++) { - for (int g = 0; g < IQP_SB_SIZE / 4; g++) { - memcpy(dst->qs + sb * 128 + g * 32 + r * 4, vals[r] + sb * IQP_SB_SIZE + g * 4, 4); - } - } - } -#endif - - dst++; - } -} - -// gemm/gemv kernels: vx points at block_iqp_x8, vy at plain (non interleaved) block_q8_K rows - -static void iqp_gemv_8x8_q8_K_generic(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - const int nb = n / QK_K; - const int ncols_interleaved = 8; - - assert(n % QK_K == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - - const block_iqp_x8 * b_ptr_start = (const block_iqp_x8 *) vx; - const block_q8_K * a_ptr = (const block_q8_K *) vy; - - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iqp_x8 * b_ptr = b_ptr_start + x * nb; - - float sumf[8] = { 0 }; - - for (int l = 0; l < nb; l++) { - int32_t sumi[8] = { 0 }; - - for (int sb = 0; sb < IQP_NSB; sb++) { - int32_t isum[8] = { 0 }; - - for (int g = 0; g < 4; g++) { - for (int j = 0; j < ncols_interleaved; j++) { - for (int k = 0; k < 4; k++) { - isum[j] += b_ptr[l].qs[sb * 128 + g * 32 + j * 4 + k] * a_ptr[l].qs[sb * 16 + g * 4 + k]; - } - } - } - - for (int j = 0; j < ncols_interleaved; j++) { - sumi[j] += isum[j] * b_ptr[l].iscales[sb * 8 + j]; - } - } - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[j] += (float) sumi[j] * (b_ptr[l].dfac[j] * a_ptr[l].d); - } - } - - for (int j = 0; j < ncols_interleaved; j++) { - s[x * ncols_interleaved + j] = sumf[j]; - } - } -} - -// one 4 row x nc column tile; s points at the first of the four output rows, bs floats apart -static void iqp_gemm_tile_4_generic(int nb, - float * GGML_RESTRICT s, - size_t bs, - const block_iqp_x8 * GGML_RESTRICT b_ptr_start, - const block_q8_K * const a_ptr[4], - int nc) { - const int ncols_interleaved = 8; - - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iqp_x8 * b_ptr = b_ptr_start + x * nb; - - float sumf[4][8]; - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] = 0.0f; - } - } - - for (int l = 0; l < nb; l++) { - for (int m = 0; m < 4; m++) { - int32_t sumi[8] = { 0 }; - - for (int sb = 0; sb < IQP_NSB; sb++) { - int32_t isum[8] = { 0 }; - - for (int g = 0; g < 4; g++) { - for (int j = 0; j < ncols_interleaved; j++) { - for (int k = 0; k < 4; k++) { - isum[j] += - b_ptr[l].qs[sb * 128 + g * 32 + j * 4 + k] * a_ptr[m][l].qs[sb * 16 + g * 4 + k]; - } - } - } - - for (int j = 0; j < ncols_interleaved; j++) { - sumi[j] += isum[j] * b_ptr[l].iscales[sb * 8 + j]; - } - } - - for (int j = 0; j < ncols_interleaved; j++) { - sumf[m][j] += (float) sumi[j] * (b_ptr[l].dfac[j] * a_ptr[m][l].d); - } - } - } - - for (int m = 0; m < 4; m++) { - for (int j = 0; j < ncols_interleaved; j++) { - s[m * bs + x * ncols_interleaved + j] = sumf[m][j]; - } - } - } -} - -static void iqp_gemm_8x8_q8_K_generic(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - const int nb = n / QK_K; - - assert(n % QK_K == 0); - assert(nr % 4 == 0); - assert(nc % 8 == 0); - - const block_iqp_x8 * b_ptr_start = (const block_iqp_x8 *) vx; - const block_q8_K * a_ptr_start = (const block_q8_K *) vy; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_K * a_ptr[4]; - for (int m = 0; m < 4; m++) { - a_ptr[m] = a_ptr_start + (y * 4 + m) * nb; - } - - iqp_gemm_tile_4_generic(nb, s + y * 4 * bs, bs, b_ptr_start, a_ptr, nc); - } -} - -static void iqp_gemm_8x8_q8_K_p4_generic(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * const * GGML_RESTRICT vy, - int nc) { - const int nb = n / QK_K; - - assert(n % QK_K == 0); - assert(nc % 8 == 0); - - const block_q8_K * a_ptr[4]; - for (int m = 0; m < 4; m++) { - a_ptr[m] = (const block_q8_K *) vy[m]; - } - - iqp_gemm_tile_4_generic(nb, s, bs, (const block_iqp_x8 *) vx, a_ptr, nc); -} - -#if defined(__AVX2__) - -// add int16_t pairwise and return as 256 bit int vector, then add the accumulator -static inline __m256i sum_i16_pairs_acc_int32x8(const __m256i acc, const __m256i x) { - const __m256i ones = _mm256_set1_epi16(1); - return _mm256_add_epi32(acc, _mm256_madd_epi16(ones, x)); -} - -static inline __m256i mul_sum_us8_pairs_acc_int32x8(const __m256i acc, const __m256i ax, const __m256i sy) { -# if defined(__AVX512VNNI__) && defined(__AVX512VL__) - return _mm256_dpbusd_epi32(acc, ax, sy); -# elif defined(__AVXVNNI__) - return _mm256_dpbusd_avx_epi32(acc, ax, sy); -# else - // Perform multiplication and create 16-bit values - const __m256i dot = _mm256_maddubs_epi16(ax, sy); - return sum_i16_pairs_acc_int32x8(acc, dot); -# endif -} - -// Integer variant of the function defined in ggml-quants.c -// multiply int8_t, add results pairwise twice and return as 256 bit int vector, then add the accumulator -static inline __m256i mul_sum_i8_pairs_acc_int32x8(const __m256i acc, const __m256i x, const __m256i y) { -# if defined(__AVXVNNIINT8__) - return _mm256_dpbssd_epi32(acc, x, y); -# else - // Get absolute values of x vectors - const __m256i ax = _mm256_sign_epi8(x, x); - // Sign the values of the y vectors - const __m256i sy = _mm256_sign_epi8(y, x); - return mul_sum_us8_pairs_acc_int32x8(acc, ax, sy); -# endif -} - -// load the 16 activations of one sub-block, offset by 128 when they are fed to dpbusd as unsigned bytes -static inline __m256i iqp_load_y(const int8_t * GGML_RESTRICT qs) { - __m128i y = _mm_loadu_si128((const __m128i *) qs); -# if GGML_IQP_USE_BIAS - y = _mm_xor_si128(y, _mm_set1_epi8((char) 0x80)); -# endif - return _mm256_broadcastsi128_si256(y); -} - -// xv: 8 rows x 4 signed weights, yb: the matching 4 activation bytes broadcast to all 8 lanes -static inline __m256i iqp_dot4(const __m256i acc, const __m256i xv, const __m256i yb) { -# if GGML_IQP_USE_BIAS - return mul_sum_us8_pairs_acc_int32x8(acc, yb, xv); -# else - return mul_sum_i8_pairs_acc_int32x8(acc, xv, yb); -# endif -} - -static inline __m256i iqp_load_iscales(const int8_t * GGML_RESTRICT iscales) { - return _mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *) iscales)); -} - -// accumulate one super-block of 8 interleaved rows against one q8_K row in int32; worst case 16 * 32 * 16 * 255 * 127 = 2.65e8 plus a bias of at most 1.33e8 does not overflow -static inline __m256i iqp_acc_block(const block_iqp_x8 * GGML_RESTRICT b, const block_q8_K * GGML_RESTRICT a) { - __m256i sumi = _mm256_setzero_si256(); - - for (int sb = 0; sb < IQP_NSB; sb++) { - const int8_t * qs = b->qs + sb * 128; - - const __m256i yv = iqp_load_y(a->qs + sb * 16); - - __m256i isum = _mm256_setzero_si256(); - - isum = iqp_dot4(isum, _mm256_loadu_si256((const __m256i *) (qs + 0)), _mm256_shuffle_epi32(yv, 0x00)); - isum = iqp_dot4(isum, _mm256_loadu_si256((const __m256i *) (qs + 32)), _mm256_shuffle_epi32(yv, 0x55)); - isum = iqp_dot4(isum, _mm256_loadu_si256((const __m256i *) (qs + 64)), _mm256_shuffle_epi32(yv, 0xAA)); - isum = iqp_dot4(isum, _mm256_loadu_si256((const __m256i *) (qs + 96)), _mm256_shuffle_epi32(yv, 0xFF)); - - sumi = _mm256_add_epi32(sumi, _mm256_mullo_epi32(isum, iqp_load_iscales(b->iscales + sb * 8))); - } - -# if GGML_IQP_USE_BIAS - sumi = _mm256_sub_epi32(sumi, _mm256_loadu_si256((const __m256i *) b->bias)); -# endif - - return sumi; -} - -// one 4 row x nc column tile; s points at the first of the four output rows, bs floats apart -static inline void iqp_gemm_tile_4(int nb, - float * GGML_RESTRICT s, - size_t bs, - const block_iqp_x8 * GGML_RESTRICT b_ptr_start, - const block_q8_K * const a_ptr[4], - int nc) { - const int ncols_interleaved = 8; - - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iqp_x8 * b_ptr = b_ptr_start + x * nb; - - __m256 sumf[4]; - for (int m = 0; m < 4; m++) { - sumf[m] = _mm256_setzero_ps(); - } - - for (int l = 0; l < nb; l++) { - __m256i sumi[4]; - for (int m = 0; m < 4; m++) { - sumi[m] = _mm256_setzero_si256(); - } - - for (int sb = 0; sb < IQP_NSB; sb++) { - const int8_t * qs = b_ptr[l].qs + sb * 128; - - __m256i yv[4]; - __m256i isum[4]; - for (int m = 0; m < 4; m++) { - yv[m] = iqp_load_y(a_ptr[m][l].qs + sb * 16); - isum[m] = _mm256_setzero_si256(); - } - - const __m256i xv0 = _mm256_loadu_si256((const __m256i *) (qs + 0)); - const __m256i xv1 = _mm256_loadu_si256((const __m256i *) (qs + 32)); - const __m256i xv2 = _mm256_loadu_si256((const __m256i *) (qs + 64)); - const __m256i xv3 = _mm256_loadu_si256((const __m256i *) (qs + 96)); - - for (int m = 0; m < 4; m++) { - isum[m] = iqp_dot4(isum[m], xv0, _mm256_shuffle_epi32(yv[m], 0x00)); - isum[m] = iqp_dot4(isum[m], xv1, _mm256_shuffle_epi32(yv[m], 0x55)); - isum[m] = iqp_dot4(isum[m], xv2, _mm256_shuffle_epi32(yv[m], 0xAA)); - isum[m] = iqp_dot4(isum[m], xv3, _mm256_shuffle_epi32(yv[m], 0xFF)); - } - - const __m256i isc = iqp_load_iscales(b_ptr[l].iscales + sb * 8); - for (int m = 0; m < 4; m++) { - sumi[m] = _mm256_add_epi32(sumi[m], _mm256_mullo_epi32(isum[m], isc)); - } - } - -# if GGML_IQP_USE_BIAS - const __m256i bias = _mm256_loadu_si256((const __m256i *) b_ptr[l].bias); - for (int m = 0; m < 4; m++) { - sumi[m] = _mm256_sub_epi32(sumi[m], bias); - } -# endif - - const __m256 dfac = _mm256_loadu_ps(b_ptr[l].dfac); - for (int m = 0; m < 4; m++) { - sumf[m] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(sumi[m]), - _mm256_mul_ps(dfac, _mm256_set1_ps(a_ptr[m][l].d)), sumf[m]); - } - } - - for (int m = 0; m < 4; m++) { - _mm256_storeu_ps(s + m * bs + x * ncols_interleaved, sumf[m]); - } - } -} - -#endif // __AVX2__ - -static void iqp_gemv_8x8_q8_K(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - const int nb = n / QK_K; - const int ncols_interleaved = 8; - - assert(n % QK_K == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(bs); - UNUSED(nr); - UNUSED(nb); - UNUSED(ncols_interleaved); - -#if defined(__AVX2__) - const block_iqp_x8 * b_ptr_start = (const block_iqp_x8 *) vx; - const block_q8_K * a_ptr = (const block_q8_K *) vy; - - for (int x = 0; x < nc / ncols_interleaved; x++) { - const block_iqp_x8 * b_ptr = b_ptr_start + x * nb; - - __m256 sumf = _mm256_setzero_ps(); - - for (int l = 0; l < nb; l++) { - const __m256 dv = _mm256_mul_ps(_mm256_loadu_ps(b_ptr[l].dfac), _mm256_set1_ps(a_ptr[l].d)); - - sumf = _mm256_fmadd_ps(_mm256_cvtepi32_ps(iqp_acc_block(b_ptr + l, a_ptr + l)), dv, sumf); - } - - _mm256_storeu_ps(s + x * ncols_interleaved, sumf); - } - - return; -#endif - - iqp_gemv_8x8_q8_K_generic(n, s, bs, vx, vy, nr, nc); -} - -static void iqp_gemm_8x8_q8_K(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * GGML_RESTRICT vy, - int nr, - int nc) { - const int nb = n / QK_K; - const int ncols_interleaved = 8; - - assert(n % QK_K == 0); - assert(nr % 4 == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(nb); - UNUSED(ncols_interleaved); - -#if defined(__AVX2__) - const block_iqp_x8 * b_ptr_start = (const block_iqp_x8 *) vx; - const block_q8_K * a_ptr_start = (const block_q8_K *) vy; - - for (int y = 0; y < nr / 4; y++) { - const block_q8_K * a_ptr[4]; - for (int m = 0; m < 4; m++) { - a_ptr[m] = a_ptr_start + (y * 4 + m) * nb; - } - - iqp_gemm_tile_4(nb, s + y * 4 * bs, bs, b_ptr_start, a_ptr, nc); - } - - return; -#endif - - iqp_gemm_8x8_q8_K_generic(n, s, bs, vx, vy, nr, nc); -} - -// same as iqp_gemm_8x8_q8_K with nr = 4, but the activation rows are passed as separate pointers (for the scattered rows of MUL_MAT_ID) -static void iqp_gemm_8x8_q8_K_p4(int n, - float * GGML_RESTRICT s, - size_t bs, - const void * GGML_RESTRICT vx, - const void * const * GGML_RESTRICT vy, - int nc) { - const int nb = n / QK_K; - const int ncols_interleaved = 8; - - assert(n % QK_K == 0); - assert(nc % ncols_interleaved == 0); - - UNUSED(nb); - UNUSED(ncols_interleaved); - -#if defined(__AVX2__) - const block_q8_K * a_ptr[4]; - for (int m = 0; m < 4; m++) { - a_ptr[m] = (const block_q8_K *) vy[m]; - } - - iqp_gemm_tile_4(nb, s, bs, (const block_iqp_x8 *) vx, a_ptr, nc); - - return; -#endif - - iqp_gemm_8x8_q8_K_p4_generic(n, s, bs, vx, vy, nc); -} - -static bool iqp_type_supported(enum ggml_type type) { - switch (type) { -#define IQP_CASE(E, name) case GGML_TYPE_##E: - IQP_TYPE_LIST(IQP_CASE) -#undef IQP_CASE - return true; - default: - return false; - } -} - -static bool iqp_supported_common(const struct ggml_tensor * dst) { - const struct ggml_tensor * src0 = dst->src[0]; - const struct ggml_tensor * src1 = dst->src[1]; - - if (!iqp_type_supported(src0->type)) { - return false; - } - - // the path assumes the src1 conversion type is q8_K - if (ggml_get_type_traits_cpu(src0->type)->vec_dot_type != GGML_TYPE_Q8_K) { - return false; - } - - // escape hatch to A/B the panel against the plain vec_dot path without rebuilding (--no-repack does not cover this path) - static const bool disabled = getenv("GGML_NO_IQ_PANEL") != nullptr; - if (disabled) { - return false; - } - - if (!ggml_cpu_has_avx2()) { - return false; - } - - if (src1->type != GGML_TYPE_F32) { - return false; - } - - if (src0->ne[0] % QK_K != 0 || src0->ne[1] % IQP_NB_ROWS != 0) { - return false; - } - - if (src0->ne[3] != 1 || src1->ne[3] != 1 || !ggml_is_contiguous(src0)) { - return false; - } - - if (dst->type != GGML_TYPE_F32 || dst->nb[0] != sizeof(float)) { - return false; - } - - return true; -} - -bool ggml_cpu_iqp_supports_mul_mat(const struct ggml_tensor * dst) { - const struct ggml_tensor * src0 = dst->src[0]; - const struct ggml_tensor * src1 = dst->src[1]; - - if (!iqp_supported_common(dst)) { - return false; - } - - if (src1->ne[1] < GGML_IQP_MIN_BATCH) { - return false; - } - - // plain 2D weight matmuls only (src1 may still be batched over ne12) - if (src0->ne[2] != 1) { - return false; - } - - return true; -} - -bool ggml_cpu_iqp_supports_mul_mat_id(const struct ggml_tensor * dst) { - const struct ggml_tensor * ids = dst->src[2]; - - if (!iqp_supported_common(dst)) { - return false; - } - - // skip the node entirely (work buffer included) if no expert can reach the per expert threshold - if (!ggml_cpu_iqp_mul_mat_id_min_batch(ids->ne[0] * ids->ne[1])) { - return false; - } - - return true; -} - -void ggml_compute_forward_mul_mat_id_iqp(const struct ggml_compute_params * params, - struct ggml_tensor * dst, - int64_t cur_a, - int64_t cne1, - const int32_t * expert_rows, - void * panels) { - const struct ggml_tensor * src0 = dst->src[0]; - const struct ggml_tensor * src1 = dst->src[1]; - - GGML_TENSOR_BINARY_OP_LOCALS - - const int ith = params->ith; - const int nth = params->nth; - - const int64_t nblocks = ne00 / QK_K; - - const size_t nbw1 = ggml_cpu_iqp_row_size(dst); - - block_iqp_x8 * panel = (block_iqp_x8 *) ((char *) panels + (size_t) ith * ggml_cpu_iqp_scratch_size(dst)); - - const char * src0_cur = (const char *) src0->data + cur_a * nb02; - - const int64_t ngroups = ne01 / IQP_NB_ROWS; - - const int64_t g0 = (ngroups * ith) / nth; - const int64_t g1 = (ngroups * (ith + 1)) / nth; - - for (int64_t g = g0; g < g1; g++) { - const int64_t r = g * IQP_NB_ROWS; - - iqp_decode_panel_8(src0->type, src0_cur + r * nb01, nb01, nblocks, panel); - - // the dst rows are scattered, so the gemm writes into tmp and it is copied out row by row - float tmp[4 * IQP_NB_ROWS]; - - for (int64_t k = 0; k < cne1; k += 4) { - const int64_t nrows = MIN(4, cne1 - k); - - // a short tail tile duplicates its last row into the unused slots; the padding is never copied out - const void * rows[4]; - - for (int64_t m = 0; m < 4; m++) { - const int64_t kk = k + MIN(m, nrows - 1); - - rows[m] = (const char *) params->wdata + - ((expert_rows[2 * kk + 0] % ne11) + expert_rows[2 * kk + 1] * ne11) * nbw1; - } - - iqp_gemm_8x8_q8_K_p4(ne00, tmp, IQP_NB_ROWS, panel, rows, IQP_NB_ROWS); - - for (int64_t m = 0; m < nrows; m++) { - float * dst_col = (float *) ((char *) dst->data + expert_rows[2 * (k + m) + 0] * nb1 + - expert_rows[2 * (k + m) + 1] * nb2); - memcpy(dst_col + r, tmp + m * IQP_NB_ROWS, IQP_NB_ROWS * sizeof(float)); - } - } - } -} - -size_t ggml_cpu_iqp_scratch_size(const struct ggml_tensor * dst) { - return GGML_PAD((dst->src[0]->ne[0] / QK_K) * sizeof(block_iqp_x8), 64); -} - -void ggml_compute_forward_mul_mat_iqp(const struct ggml_compute_params * params, struct ggml_tensor * dst) { - const struct ggml_tensor * src0 = dst->src[0]; - const struct ggml_tensor * src1 = dst->src[1]; - - GGML_TENSOR_BINARY_OP_LOCALS - - const int ith = params->ith; - const int nth = params->nth; - - const int64_t nblocks = ne00 / QK_K; - - const size_t nbw1 = ggml_row_size(GGML_TYPE_Q8_K, ne10); - const size_t nbw2 = nbw1 * ne11; - - const size_t scratch_size = ggml_cpu_iqp_scratch_size(dst); - - const size_t scratch_offset = GGML_PAD(nbw2 * ne12, 64); - - GGML_ASSERT(scratch_offset + (size_t) nth * scratch_size <= params->wsize); - - block_iqp_x8 * panel = (block_iqp_x8 *) ((char *) params->wdata + scratch_offset + (size_t) ith * scratch_size); - - const int64_t nrows = ne11; - - const int64_t ngroups = ne01 / IQP_NB_ROWS; - - // aim for 4 chunks per thread; the caller has already reset the chunk counter - // on NUMA systems fall back to one chunk per thread - const int64_t chunks_per_thread = ggml_is_numa() ? 1 : 4; - const int64_t groups_per_chunk = MAX(1, (ngroups + nth * chunks_per_thread - 1) / (nth * chunks_per_thread)); - const int64_t nchunk = (ngroups + groups_per_chunk - 1) / groups_per_chunk; - - int current_chunk = ith; - - while (current_chunk < nchunk) { - const int64_t g0 = current_chunk * groups_per_chunk; - const int64_t g1 = MIN(g0 + groups_per_chunk, ngroups); - - for (int64_t g = g0; g < g1; g++) { - const int64_t r = g * IQP_NB_ROWS; - - iqp_decode_panel_8(src0->type, (const char *) src0->data + r * nb01, nb01, nblocks, panel); - - for (int64_t i12 = 0; i12 < ne12; i12++) { - const char * src1_ptr = (const char *) params->wdata + i12 * nbw2; - char * dst_ptr = (char *) dst->data + i12 * nb2; - - if (nrows > 3) { - iqp_gemm_8x8_q8_K(ne00, (float *) dst_ptr + r, nb1 / nb0, panel, src1_ptr, nrows - (nrows % 4), - IQP_NB_ROWS); - } - for (int64_t iter = nrows - (nrows % 4); iter < nrows; iter++) { - iqp_gemv_8x8_q8_K(ne00, (float *) (dst_ptr + iter * nb1) + r, ne01, panel, src1_ptr + nbw1 * iter, - 1 /* nrows */, IQP_NB_ROWS); - } - } - } - - current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); - } -} diff --git a/ggml/src/ggml-cpu/iqp.h b/ggml/src/ggml-cpu/iqp.h deleted file mode 100644 index 017b03fb43..0000000000 --- a/ggml/src/ggml-cpu/iqp.h +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once - -#include "ggml-cpu-impl.h" -#include "ggml.h" - -// GGML internal header - -// batched mul_mat path for the grid based IQ types: decode 8 src0 rows at a time into per thread scratch -// (block_iqp_x8, see iqp.cpp) and run an integer gemm over them against all src1 columns - -#ifdef __cplusplus -extern "C" { -#endif - -// whether cne1 rows of src1 are enough for the decode to pay for itself, per expert, for MUL_MAT_ID -bool ggml_cpu_iqp_mul_mat_id_min_batch(int64_t cne1); - -bool ggml_cpu_iqp_supports_mul_mat(const struct ggml_tensor * dst); - -// node level test only - per expert eligibility is decided with ggml_cpu_iqp_mul_mat_id_min_batch -bool ggml_cpu_iqp_supports_mul_mat_id(const struct ggml_tensor * dst); - -// per thread panel scratch bytes, padded -size_t ggml_cpu_iqp_scratch_size(const struct ggml_tensor * dst); - -// must be called after src1 has been converted to q8_K into params->wdata and the threads have synchronized on it -void ggml_compute_forward_mul_mat_iqp(const struct ggml_compute_params * params, struct ggml_tensor * dst); - -// one expert: expert_rows points at its row of the matrix_rows table of (i1, i2) int32 pairs, panels at the base of the per thread panel scratches -void ggml_compute_forward_mul_mat_id_iqp(const struct ggml_compute_params * params, - struct ggml_tensor * dst, - int64_t cur_a, - int64_t cne1, - const int32_t * expert_rows, - void * panels); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp b/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp new file mode 100644 index 0000000000..bffa737c7c --- /dev/null +++ b/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp @@ -0,0 +1,719 @@ +// mulmat microtile kernels + +#include "tiled-kernel.h" + +#include "ggml.h" + +#include + +#if defined(__AVX512VNNI__) || defined(__AVX2__) || defined(__AVX__) +#include +#endif + + +// Reference implementation, slower than existing vec_dot approach +template +static void tiled_run_microtile_scalar(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + constexpr int NB = TILED_TILE_K / SUBBLK; // subblocks per 256-K block + constexpr int NS = SUBBLK / 16; // per-16 bsums per subblock + + // num_k/slab: the narrow path holds num_k slabs at row stride num_k*256 (num_k=1, slab=0 = standard) + // NK: 1 = standard single-slab (offsets 0, strides compile-time); 0 = runtime + // (num_k/slab from the args). The narrow path is memory-bound, so one runtime + // version serves all of it. + const int nkr = (NK > 0) ? NK : num_k; + const int seff = (NK > 0) ? 0 : slab; + const int qk_stride = nkr * TILED_TILE_K; + const int qk_off = seff * TILED_TILE_K; + const int nb_stride = NB * nkr; + const int nb_off = seff * NB; + const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO; + const int bs_off = seff * TILED_MICRO; + + float acc[TILED_MICRO][TILED_MICRO]; + memset(acc, 0, sizeof(acc)); + + // subdots at subblock granularity over the 256-K block, exact integer math + for (int s = 0; s < NB; s++) { + for (int j = 0; j < TILED_MICRO; j++) { + const int br = j0 + j; + const int8_t * q1 = &src1.q[br * qk_stride + qk_off + s * SUBBLK]; + int32_t bsum = 0; + for (int u = 0; u < NS; u++) { + bsum += src1.bsums[(bs_off + s * NS + u) * bs_stride + br]; + } + + for (int i = 0; i < TILED_MICRO; i++) { + const int ar = i0 + i; + const int d_off = seff * TILED_MICRO + ar; + const uint8_t * q0 = &src0.q[ar * qk_stride + qk_off + s * SUBBLK]; + + int32_t raw = 0; + for (int e = 0; e < SUBBLK; e++) { + raw += (int32_t) q0[e] * (int32_t) q1[e]; + } + + // BIAS: subtract BIAS*bsum (src1's per-subblock code sum) from the exact int raw + int32_t corr = raw; + if constexpr (BIAS != 0) { + corr -= BIAS * bsum; + } + const int32_t scales_raw = (int32_t) src0.scales[ar * nb_stride + nb_off + s] * corr; + // d is NOT applied here: it is constant over the s-loop, so we apply it last before write-out + if constexpr (HAS_MIN) { + const int32_t mins_bsum = (int32_t) src0.mins[ar * nb_stride + nb_off + s] * bsum; + acc[i][j] += (float) src0.d[d_off] * (float) scales_raw + - (float) src0.dmin[d_off] * (float) mins_bsum; + } else { + acc[i][j] += (float) src0.d[d_off] * (float) scales_raw; + } + } + } + } + + // Apply d and write out to buf + for (int i = 0; i < TILED_MICRO; i++) { + for (int j = 0; j < TILED_MICRO; j++) { + buf[(i0 + i) * buf_stride + (j0 + j)] += src1.d[seff * TILED_MICRO + j0 + j] * acc[i][j]; + } + } +} + +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512DQ__) + + +// VNNI microkernel +// 8x16 band pass: 8 src0 rows (i0, i0+8) x 16 src1 cols (j0, j0+16) +// Math: s1_acc += scales_s*raw_s, s2_acc += mins_s*bsums_s per subblock; the row result +// is d0*s1_acc - dmin*s2_acc (f32, exact in range), applied once per row. +// +// Register pressure: the band pass holds acc16 + s1_acc = 16 zmm for the +// 8-row band (s2_acc has no dpbusd dependency, so it is computed after the +// band pass and adds no register pressure); +template +static void tiled_run_micro_vnni_8x16(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + constexpr int NB = TILED_TILE_K / SUBBLK; + constexpr int NS = SUBBLK / 16; + constexpr int NG = SUBBLK / 4; + + // band width, see the register-pressure note above + constexpr int NUM_ROWS = 8; + + // num_k = K-blocks per row (the narrow path holds num_k slabs at row stride num_k*256 so + // each weight row is one long stream); slab = this call's slab. num_k=1, slab=0 = standard. + // NK: 1 = standard single-slab (offsets 0, strides compile-time); 0 = runtime + // (num_k/slab from the args). The narrow path is memory-bound, so one runtime + // version serves all of it. + const int nkr = (NK > 0) ? NK : num_k; + const int seff = (NK > 0) ? 0 : slab; + const int qk_stride = nkr * TILED_TILE_K; + const int qk_off = seff * TILED_TILE_K; + const int nb_stride = NB * nkr; + const int nb_off = seff * NB; + const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO; + const int bs_off = seff * TILED_MICRO; + + const __m512 d1_vec = _mm512_load_ps(&src1.d[seff * TILED_MICRO + j0]); + + __m512i s1_acc[NUM_ROWS]; + for (int t = 0; t < NUM_ROWS; t++) { s1_acc[t] = _mm512_setzero_si512(); } + __m512i s2_acc[NUM_ROWS]; + if constexpr (HAS_MIN) { + for (int t = 0; t < NUM_ROWS; t++) { s2_acc[t] = _mm512_setzero_si512(); } + } + + const int32_t * bsums = src1.bsums; // int32 per-16 sums; NS > 1 combines NS lanes per subblock + for (int s = 0; s < NB; s++) { + __m512i bsums32 = _mm512_load_si512((const __m512i *) &bsums[(bs_off + s * NS) * bs_stride + j0]); + for (int u = 1; u < NS; u++) { + bsums32 = _mm512_add_epi32(bsums32, _mm512_load_si512((const __m512i *) &bsums[(bs_off + s * NS + u) * bs_stride + j0])); + } + + __m512i bias32 = _mm512_setzero_si512(); + if constexpr (BIAS != 0) { + bias32 = _mm512_mullo_epi32(bsums32, _mm512_set1_epi32(BIAS)); + } + + __m512i acc16[NUM_ROWS]; + for (int t = 0; t < NUM_ROWS; t++) { acc16[t] = _mm512_setzero_si512(); } + +#ifdef __GNUC__ + #pragma GCC unroll 8 // pragma unrolled justified by measuing with/without +#endif + for (int g = 0; g < NG; g++) { + const int kg = s * NG + g; + // in-place interleave layout: [kg%16 @ TILED_TILE_K][kg/16 @ 64][row @ 4] + const __m512i codes = _mm512_load_si512((const __m512i *) &src1.q[(j0 / TILED_MICRO) * (TILED_MICRO * qk_stride) + + qk_off + (kg % TILED_MICRO) * qk_stride + (kg / TILED_MICRO) * (TILED_MICRO * 4)]); + for (int t = 0; t < NUM_ROWS; t++) { + const uint32_t u4 = *(const uint32_t *) &src0.q[(i0 + t) * qk_stride + qk_off + kg * 4]; + const __m512i u4b = _mm512_set1_epi32((int) u4); + acc16[t] = _mm512_dpbusd_epi32(acc16[t], u4b, codes); + } + } + + // int correction: s1_acc += scales*(raw-BIAS*bsums) + for (int t = 0; t < NUM_ROWS; t++) { + const int ar = i0 + t; + __m512i rawi = acc16[t]; + if constexpr (BIAS != 0) { + rawi = _mm512_sub_epi32(rawi, bias32); + } + s1_acc[t] = _mm512_add_epi32(s1_acc[t], + _mm512_mullo_epi32(rawi, _mm512_set1_epi32(src0.scales[ar * nb_stride + nb_off + s]))); + } + } + + // s2_acc += mins*bsums; independent of the dpbusd results, so it runs here instead of in the + // band pass: the band pass keeps the register budget for the 8-row band + if constexpr (HAS_MIN) { + for (int s = 0; s < NB; s++) { + __m512i bsums32 = _mm512_load_si512((const __m512i *) &bsums[(bs_off + s * NS) * bs_stride + j0]); + for (int u = 1; u < NS; u++) { + bsums32 = _mm512_add_epi32(bsums32, _mm512_load_si512((const __m512i *) &bsums[(bs_off + s * NS + u) * bs_stride + j0])); + } + for (int t = 0; t < NUM_ROWS; t++) { + s2_acc[t] = _mm512_add_epi32(s2_acc[t], + _mm512_mullo_epi32(bsums32, _mm512_set1_epi32(src0.mins[(i0 + t) * nb_stride + nb_off + s]))); + } + } + } + + // epilogue: int->float, apply per-row scales, store to buf + for (int t = 0; t < NUM_ROWS; t++) { + const int ar = i0 + t; + const int d_off = seff * TILED_MICRO + ar; + __m512 f1 = _mm512_cvtepi32_ps(s1_acc[t]); + __m512 result = _mm512_mul_ps(f1, _mm512_set1_ps(src0.d[d_off])); + if constexpr (HAS_MIN) { + __m512 f2 = _mm512_cvtepi32_ps(s2_acc[t]); + result = _mm512_fnmadd_ps(_mm512_set1_ps(src0.dmin[d_off]), f2, result); + } + float * p = &buf[(i0 + t) * buf_stride + j0]; + _mm512_store_ps(p, _mm512_add_ps(_mm512_load_ps(p), _mm512_mul_ps(result, d1_vec))); + } +} + +// 16x16 microtile as two explicit 8x16 band passes +template +static void tiled_run_microtile_vnni(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + tiled_run_micro_vnni_8x16(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + tiled_run_micro_vnni_8x16(src0, src1, i0 + 8, j0, num_k, slab, buf, buf_stride); +} + +#endif // __AVX512VNNI__ && __AVX512VL__ + +#if defined(__AVX2__) + +// AVX2 kernel. +// We're effectively applying the existing vec_dot algorithms to an 8x16 block here. +// Different paths based on subblock size as it affects when/where we multiply in scales and apply mins +// SUBBLK=32: one 256-bit maddubs per (s, column), a 256-bit set1 scale, +// one 8-lane i32 accumulator per column. +// SUBBLK=16: two subblocks per 32B load. The 256-bit maddubs product +// still gives 16 i16 lanes (0..7 = subblock sp, 8..15 = sp+1), but the +// scale is applied per 128-bit half. +// With BIAS != 0 the q0 codes are biased (up to 241) and a maddubs i16 pair +// overflows, so the weight is debiased (w = q0 - BIAS, small). Two MACs: +// ACTBIAS: activation pre-biased (+128) by repack, one maddubs(q1b, w); the +128 is +// corrected out in the epilogue (128*sum_s scales[s]*w_bsum[s], precomputed in repack_src0). +// Fits only when the debiased weight magnitude Wmax <= 64 (2*255*Wmax < 32767). +// else (iq4_xs, Wmax = 127): the signed-int8 sign trick, ax = |w|, sy = q1*sign(w), +// one maddubs (2*127*127 < 32767) so it always fits. +template +static void tiled_run_microtile_avx2(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + constexpr int NB = TILED_TILE_K / SUBBLK; + constexpr int GROUP = 8; // Process 8 src1 columns concurrently (1 YMM vector) + + static_assert(SUBBLK == 16 || SUBBLK == 32, "unsupported SUBBLK"); + + // num_k/slab: the narrow path holds num_k slabs at row stride num_k*256 (num_k=1, slab=0 = standard) + // NK: 1 = standard single-slab (offsets 0, strides compile-time); 0 = runtime + // (num_k/slab from the args). The narrow path is memory-bound, so one runtime + // version serves all of it. + const int nkr = (NK > 0) ? NK : num_k; + const int seff = (NK > 0) ? 0 : slab; + const int qk_stride = nkr * TILED_TILE_K; + const int qk_off = seff * TILED_TILE_K; + const int nb_stride = NB * nkr; + const int nb_off = seff * NB; + const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO; + const int bs_off = seff * TILED_MICRO; + + // Horizontal sum helper: converts 8x32-bit int lane inside a YMM register to scalar int32 + auto hsum256_epi32 = [](const __m256i v) -> int32_t { + __m128i low = _mm256_castsi256_si128(v); + __m128i high = _mm256_extracti128_si256(v, 1); + __m128i sum128 = _mm_add_epi32(low, high); + sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_SHUFFLE(2, 3, 0, 1))); + sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_SHUFFLE(1, 0, 3, 2))); + return _mm_cvtsi128_si32(sum128); + }; + + for (int i = 0; i < TILED_MICRO; i++) { + const int ar = i0 + i; + const int d_off = seff * TILED_MICRO + ar; + const float d0 = src0.d[d_off]; + const float dmin0 = src0.dmin[d_off]; + const uint8_t * q0 = &src0.q[ar * qk_stride + qk_off]; + const int32_t * scales_row = &src0.scales[ar * nb_stride + nb_off]; + const int32_t * mins_row = &src0.mins[ar * nb_stride + nb_off]; + + for (int g = 0; g < TILED_MICRO; g += GROUP) { + // Pre-calculate weight pointers for this column group + const int8_t * q1_ptr[GROUP]; + for (int t = 0; t < GROUP; t++) { + q1_ptr[t] = &src1.q[(j0 + g + t) * qk_stride + qk_off]; + } + + __m256i s2 = _mm256_setzero_si256(); + + __m256i acc[GROUP]; + for (int t = 0; t < GROUP; t++) { + acc[t] = _mm256_setzero_si256(); + } + + if constexpr (SUBBLK == 32) { + for (int s = 0; s < NB; s++) { + // OPTIMIZATION 1: Use aligned loads (_mm256_load_si256) + const __m256i q0_32 = _mm256_load_si256((const __m256i *) &q0[s * SUBBLK]); + const __m256i scales16 = _mm256_set1_epi16(scales_row[s]); + + if constexpr (BIAS != 0) { + const __m256i w = _mm256_sub_epi8(q0_32, _mm256_set1_epi8((int8_t) BIAS)); // debiased weight + if constexpr (ACTBIAS) { + // activation pre-biased (+128) by repack; q1b unsigned, w signed + #pragma GCC unroll 8 + for (int t = 0; t < GROUP; t++) { + const __m256i q1b_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][s * SUBBLK]); + acc[t] = _mm256_add_epi32(acc[t], + _mm256_madd_epi16(scales16, _mm256_maddubs_epi16(q1b_32, w))); + } + } else { + // iq4_xs (Wmax > 64 overflows act-bias): sign trick, both operands <= 127 + const __m256i ax = _mm256_sign_epi8(w, w); + #pragma GCC unroll 8 + for (int t = 0; t < GROUP; t++) { + const __m256i q1_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][s * SUBBLK]); + acc[t] = _mm256_add_epi32(acc[t], + _mm256_madd_epi16(scales16, _mm256_maddubs_epi16(ax, _mm256_sign_epi8(q1_32, w)))); + } + } + } else { + #pragma GCC unroll 8 + for (int t = 0; t < GROUP; t++) { + const __m256i q1_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][s * SUBBLK]); + acc[t] = _mm256_add_epi32(acc[t], + _mm256_madd_epi16(scales16, _mm256_maddubs_epi16(q0_32, q1_32))); + } + } + + if constexpr (HAS_MIN) { + const __m256i bsums_v = _mm256_add_epi32( + _mm256_load_si256((const __m256i *) &src1.bsums[(bs_off + s * 2) * bs_stride + j0 + g]), + _mm256_load_si256((const __m256i *) &src1.bsums[(bs_off + s * 2 + 1) * bs_stride + j0 + g])); + s2 = _mm256_add_epi32(s2, _mm256_mullo_epi32(bsums_v, _mm256_set1_epi32(mins_row[s]))); + } + } + } else { // SUBBLK == 16 + for (int sp = 0; sp < NB; sp += 2) { + const __m256i q0_32 = _mm256_load_si256((const __m256i *) &q0[sp * SUBBLK]); + const __m256i scalesv = _mm256_set_m128i(_mm_set1_epi16(scales_row[sp + 1]), _mm_set1_epi16(scales_row[sp])); + if constexpr (BIAS != 0) { + const __m256i w = _mm256_sub_epi8(q0_32, _mm256_set1_epi8((int8_t) BIAS)); // debiased weight + if constexpr (ACTBIAS) { + // activation pre-biased (+128) by repack; q1b unsigned, w signed + #pragma GCC unroll 8 + for (int t = 0; t < GROUP; t++) { + const __m256i q1b_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][sp * SUBBLK]); + acc[t] = _mm256_add_epi32(acc[t], _mm256_madd_epi16( + scalesv, _mm256_maddubs_epi16(q1b_32, w))); + } + } else { + // iq4_xs (Wmax > 64 overflows act-bias): sign trick, both operands <= 127 + const __m256i ax = _mm256_sign_epi8(w, w); + #pragma GCC unroll 8 + for (int t = 0; t < GROUP; t++) { + const __m256i q1_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][sp * SUBBLK]); + acc[t] = _mm256_add_epi32(acc[t], _mm256_madd_epi16( + scalesv, _mm256_maddubs_epi16(ax, _mm256_sign_epi8(q1_32, w)))); + } + } + } else { + #pragma GCC unroll 8 + for (int t = 0; t < GROUP; t++) { + const __m256i q1_32 = _mm256_load_si256((const __m256i *) &q1_ptr[t][sp * SUBBLK]); + acc[t] = _mm256_add_epi32(acc[t], _mm256_madd_epi16( + scalesv, _mm256_maddubs_epi16(q0_32, q1_32))); + } + } + + if constexpr (HAS_MIN) { + const __m256i bsums0_v = _mm256_load_si256((const __m256i *) &src1.bsums[(bs_off + sp) * bs_stride + j0 + g]); + const __m256i bsums1_v = _mm256_load_si256((const __m256i *) &src1.bsums[(bs_off + sp + 1) * bs_stride + j0 + g]); + s2 = _mm256_add_epi32(s2, _mm256_add_epi32( + _mm256_mullo_epi32(bsums0_v, _mm256_set1_epi32(mins_row[sp])), + _mm256_mullo_epi32(bsums1_v, _mm256_set1_epi32(mins_row[sp + 1])))); + } + } + } + + // OPTIMIZATION 2: Vectorized Epilogue with FMA3 gating + // Reduce accumulators into a 256-bit vector across 8 columns + alignas(32) int32_t s1_vals[GROUP]; + #pragma GCC unroll 8 + for (int t = 0; t < GROUP; t++) { + s1_vals[t] = hsum256_epi32(acc[t]); + } + __m256i s1_vec = _mm256_load_si256((const __m256i *) s1_vals); + + if constexpr (BIAS != 0 && ACTBIAS) { + // act-bias correction: raw = desired + 128*sum_s scales[s]*w_bsum[s]; corr precomputed in repack_src0 + s1_vec = _mm256_sub_epi32(s1_vec, _mm256_set1_epi32(mins_row[0])); + } + + const __m256 src1_d_vec = _mm256_load_ps(&src1.d[seff * TILED_MICRO + j0 + g]); + __m256 res_vec = _mm256_mul_ps(_mm256_set1_ps(d0), _mm256_cvtepi32_ps(s1_vec)); + + if constexpr (HAS_MIN) { +#if defined(__FMA__) + // FMA3 implementation: res_vec = res_vec - (dmin0 * s2) + res_vec = _mm256_fnmadd_ps(_mm256_set1_ps(dmin0), _mm256_cvtepi32_ps(s2), res_vec); +#else + // Non-FMA fallback + res_vec = _mm256_sub_ps(res_vec, _mm256_mul_ps(_mm256_set1_ps(dmin0), _mm256_cvtepi32_ps(s2))); +#endif + } + + float * buf_ptr = &buf[ar * buf_stride + j0 + g]; + __m256 current_buf = _mm256_loadu_ps(buf_ptr); + +#if defined(__FMA__) + // FMA3 accumulation: current_buf + (res_vec * src1_d_vec) + __m256 updated_buf = _mm256_fmadd_ps(res_vec, src1_d_vec, current_buf); +#else + __m256 updated_buf = _mm256_add_ps(current_buf, _mm256_mul_ps(res_vec, src1_d_vec)); +#endif + + _mm256_storeu_ps(buf_ptr, updated_buf); + } + } +} + +#endif // __AVX2__ + +#if defined(__AVX__) && !defined(__AVX2__) +template +static void tiled_run_microtile_avx(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + constexpr int NB = TILED_TILE_K / SUBBLK; + constexpr int NS = SUBBLK / 16; + constexpr int GROUP = 4; // src1 columns per group: one acc32 per column + + // num_k/slab: the narrow path holds num_k slabs at row stride num_k*256 (num_k=1, slab=0 = standard) + // NK: 1 = standard single-slab (offsets 0, strides compile-time); 0 = runtime + // (num_k/slab from the args). The narrow path is memory-bound, so one runtime + // version serves all of it. + const int nkr = (NK > 0) ? NK : num_k; + const int seff = (NK > 0) ? 0 : slab; + const int qk_stride = nkr * TILED_TILE_K; + const int qk_off = seff * TILED_TILE_K; + const int nb_stride = NB * nkr; + const int nb_off = seff * NB; + const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO; + const int bs_off = seff * TILED_MICRO; + + for (int i = 0; i < TILED_MICRO; i++) { + const int ar = i0 + i; + const int d_off = seff * TILED_MICRO + ar; + const float d0 = src0.d[d_off]; + const float dmin0 = src0.dmin[d_off]; + const uint8_t * q0 = &src0.q[ar * qk_stride + qk_off]; + const int32_t * scales_row = &src0.scales[ar * nb_stride + nb_off]; + const int32_t * mins_row = &src0.mins[ar * nb_stride + nb_off]; + + for (int g = 0; g < TILED_MICRO; g += GROUP) { + const int8_t * q1g[GROUP]; + __m128i acc[GROUP]; + __m128i s2 = _mm_setzero_si128(); // sum_s mins_s * bsums_s + for (int t = 0; t < GROUP; t++) { + q1g[t] = &src1.q[(j0 + g + t) * qk_stride + qk_off]; + acc[t] = _mm_setzero_si128(); + } + + for (int s = 0; s < NB; s++) { + const __m128i scales16 = _mm_set1_epi16(scales_row[s]); + if constexpr (HAS_MIN) { + __m128i bsums_v = _mm_setzero_si128(); // per-col bsums as a 4-lane i32 vector + for (int u = 0; u < NS; u++) { + bsums_v = _mm_add_epi32(bsums_v, _mm_loadu_si128( + (const __m128i *) &src1.bsums[(bs_off + s * NS + u) * bs_stride + j0 + g])); + } + s2 = _mm_add_epi32(s2, _mm_mullo_epi32(bsums_v, _mm_set1_epi32(mins_row[s]))); + } + for (int u = 0; u < NS; u++) { + const __m128i a16 = _mm_loadu_si128((const __m128i *) &q0[s * SUBBLK + u * 16]); + if constexpr (BIAS != 0) { + // debiased weight w = a16 - BIAS; signed dot via the sign trick (see header) + const __m128i a16s = _mm_sub_epi8(a16, _mm_set1_epi8((int8_t) BIAS)); + const __m128i ax = _mm_sign_epi8(a16s, a16s); + for (int t = 0; t < GROUP; t++) { + const __m128i q1_16 = _mm_loadu_si128((const __m128i *) &q1g[t][s * SUBBLK + u * 16]); + acc[t] = _mm_add_epi32(acc[t], + _mm_madd_epi16(scales16, _mm_maddubs_epi16(ax, _mm_sign_epi8(q1_16, a16s)))); + } + } else { + for (int t = 0; t < GROUP; t++) { + acc[t] = _mm_add_epi32(acc[t], _mm_madd_epi16(scales16, + _mm_maddubs_epi16(a16, + _mm_loadu_si128((const __m128i *) &q1g[t][s * SUBBLK + u * 16])))); + } + } + } + } + + int32_t s2_s[GROUP] = { 0, 0, 0, 0 }; + if constexpr (HAS_MIN) { + _mm_storeu_si128((__m128i *) s2_s, s2); + } + + for (int t = 0; t < GROUP; t++) { + // 4 i32 lanes -> scalar + __m128i v = _mm_shuffle_epi32(acc[t], _MM_SHUFFLE(2, 3, 0, 1)); + acc[t] = _mm_add_epi32(acc[t], v); + v = _mm_shuffle_epi32(acc[t], _MM_SHUFFLE(1, 0, 3, 2)); + acc[t] = _mm_add_epi32(acc[t], v); + const int32_t s1 = _mm_cvtsi128_si32(acc[t]); + + float res = d0 * (float) s1; + if constexpr (HAS_MIN) { + res -= dmin0 * (float) s2_s[t]; + } + buf[ar * buf_stride + j0 + g + t] += src1.d[seff * TILED_MICRO + j0 + g + t] * res; + } + } + } +} + +#endif // __AVX__ && !__AVX2__ + +// main microtile entry point. num_k = K-blocks per row (the narrow path reads num_k slabs at row +// stride num_k*256 so each weight row is one long stream); slab = this call's slab. num_k=1, +// slab=0 is the standard single-slab path. +// num_k==1 dispatches to the NK=1 instantiation, whose strides/offsets are compile-time +// constants (shifts/scaled-leas). num_k>1 (the narrow path) uses the NK=0 runtime-strided +// version; it is memory-bound, so one version serves all of it. +template +void tiled_run_microtile(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride) { + const bool standard = (num_k == 1); +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512DQ__) + if (standard) { + tiled_run_microtile_vnni(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } else { + tiled_run_microtile_vnni(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } +#elif defined(__AVX2__) + if (standard) { + tiled_run_microtile_avx2(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } else { + tiled_run_microtile_avx2(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } +#elif defined(__AVX__) + if (standard) { + tiled_run_microtile_avx(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } else { + tiled_run_microtile_avx(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } +#else + if (standard) { + tiled_run_microtile_scalar(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } else { + tiled_run_microtile_scalar(src0, src1, i0, j0, num_k, slab, buf, buf_stride); + } +#endif +} + +// explicit instantiations for the in-use formats (q4_K and q5_K share the constants) +// ACTBIAS selects the act-bias MAC (Wmax <= 64); iq4_xs (Wmax 127) keeps the sign trick +// q4_K / q5_K: BIAS = 0 + template void tiled_run_microtile<32, true, 0, false>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); +// iq4_xs: sign trick (Wmax 127 overflows act-bias) +template void tiled_run_microtile<32, false, 128, false>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); +// iq2_xxs, iq3_xxs, iq3_s, iq1_s: act-bias (Wmax <= 62) +template void tiled_run_microtile<32, false, 128, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); +// iq2_xs, iq2_s, iq1_m: act-bias +template void tiled_run_microtile<16, false, 128, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); +// q6_K: act-bias (Wmax 32) +template void tiled_run_microtile<16, false, 32, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); +// q3_K: act-bias (Wmax 4) +template void tiled_run_microtile<16, false, 4, true>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); +// q2_K: BIAS = 0 +template void tiled_run_microtile<16, true, 0, false>(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); + + +#define MIN(a, b) ((a) < (b) ? (a) : (b)) + +// block_q8_K in 4-byte words, for the int32 gather indices +static_assert(sizeof(block_q8_K) == 292 && offsetof(block_q8_K, qs) == 4, + "block_q8_K layout changed, fix the src1 repack"); + +// VNNI: interleave the natural [row][k] codes in-place into the group-local +// [kg%16][kg/16][row][4] layout. base points at row 0 (row r at base + r*row_stride); +// n_tiles 16x16 int32 tiles are transposed to the right (tile t occupies bytes +// [t*64, t*64+64) per row). Each 16-row group is self-contained in row_stride*16 bytes. +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512DQ__) +void tiled_repack_src1(tiled_tile_src1 * src1, int row0, int num_k, bool bias) { + GGML_UNUSED(bias); + const int n_tiles = num_k * (TILED_TILE_K / (TILED_MICRO * 4)); + const int row_stride = num_k * TILED_TILE_K; + uint8_t * base = (uint8_t *) src1->q + row0 * row_stride; + for (int t = 0; t < n_tiles; t++) { + uint8_t * cbase = base + t * (TILED_MICRO * 4); + __m512i v[16]; + for (int r = 0; r < 16; r++) { + v[r] = _mm512_load_si512((const __m512i *) (cbase + r * row_stride)); + } + + // 16x16 int32 transpose, 4 butterfly phases + // Phase 1: 1-element interleave, pairs (0,1), (2,3), ..., (14,15) + { + const __m512i idx_a = _mm512_setr_epi32(0, 16, 1, 17, 2, 18, 3, 19, 4, 20, 5, 21, 6, 22, 7, 23); + const __m512i idx_b = _mm512_setr_epi32(8, 24, 9, 25, 10, 26, 11, 27, 12, 28, 13, 29, 14, 30, 15, 31); + for (int i = 0; i < 16; i += 2) { + __m512i a = v[i], b = v[i+1]; + v[i] = _mm512_permutex2var_epi32(a, idx_a, b); + v[i+1] = _mm512_permutex2var_epi32(a, idx_b, b); + } + } + // Phase 2: 2-element interleave, pairs (0,2), (1,3), (4,6), (5,7), ... + { + const __m512i idx_a = _mm512_setr_epi32(0, 1, 16, 17, 4, 5, 20, 21, 8, 9, 24, 25, 12, 13, 28, 29); + const __m512i idx_b = _mm512_setr_epi32(2, 3, 18, 19, 6, 7, 22, 23, 10, 11, 26, 27, 14, 15, 30, 31); + for (int i = 0; i < 16; i += 4) { + __m512i a = v[i], b = v[i+2]; + v[i] = _mm512_permutex2var_epi32(a, idx_a, b); + v[i+2] = _mm512_permutex2var_epi32(a, idx_b, b); + a = v[i+1], b = v[i+3]; + v[i+1] = _mm512_permutex2var_epi32(a, idx_a, b); + v[i+3] = _mm512_permutex2var_epi32(a, idx_b, b); + } + } + // Phase 3: 4-element interleave, pairs (0,4), (1,5), ..., (7,11), (8,12), ... + { + const __m512i idx_a = _mm512_setr_epi32(0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23); + const __m512i idx_b = _mm512_setr_epi32(8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31); + for (int i = 0; i < 16; i += 8) { + for (int j = 0; j < 4; j++) { + __m512i a = v[i+j], b = v[i+4+j]; + v[i+j] = _mm512_permutex2var_epi32(a, idx_a, b); + v[i+4+j] = _mm512_permutex2var_epi32(a, idx_b, b); + } + } + } + // Phase 4: 8-element interleave, pairs (0,8), (1,9), ..., (7,15) + { + const __m512i idx_a = _mm512_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23); + const __m512i idx_b = _mm512_setr_epi32(8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31); + for (int i = 0; i < 8; i++) { + __m512i a = v[i], b = v[i+8]; + v[i] = _mm512_permutex2var_epi32(a, idx_a, b); + v[i+8] = _mm512_permutex2var_epi32(a, idx_b, b); + } + } + + // store: v[g] holds 16 int32s for k-group col_order[g], rows 0..15 + static const int col_order[16] = {0, 8, 1, 9, 4, 12, 5, 13, 2, 10, 3, 11, 6, 14, 7, 15}; + for (int g = 0; g < 16; g++) { + _mm512_store_si512((void *) (cbase + col_order[g] * row_stride), v[g]); + } + } +} +#elif defined(__AVX2__) +// act-bias: pre-bias the activation in place (+128) so the MAC's maddubs(q1b, w) sees unsigned bytes; +// no transpose on AVX2 (natural layout). 16 rows x n_tiles*64 bytes, row r at base + r*row_stride +void tiled_repack_src1(tiled_tile_src1 * src1, int row0, int num_k, bool bias) { + if (!bias) { + return; + } + const int row_stride = num_k * TILED_TILE_K; + uint8_t * base = (uint8_t *) src1->q + row0 * row_stride; + const int nbytes = row_stride; // full row width + const __m256i b128 = _mm256_set1_epi8((int8_t) 0x80); + for (int r = 0; r < TILED_MICRO; r++) { + uint8_t * row = base + r * row_stride; + for (int off = 0; off < nbytes; off += 32) { + const __m256i v = _mm256_loadu_si256((const __m256i *) (row + off)); + _mm256_storeu_si256((__m256i *) (row + off), _mm256_xor_si256(v, b128)); + } + } +} +#else +void tiled_repack_src1(tiled_tile_src1 * src1, int row0, int num_k, bool bias) { + GGML_UNUSED(src1); GGML_UNUSED(row0); GGML_UNUSED(num_k); GGML_UNUSED(bias); +} +#endif + +#if defined(__AVX2__) +// sum 32 unsigned bytes to int32 (via maddubs -> int16 pairs -> hsum) +static int32_t tiled_byte_sum_32(const uint8_t * p) { + const __m256i dot = _mm256_maddubs_epi16(_mm256_loadu_si256((const __m256i *) p), _mm256_set1_epi8(1)); + const __m256i sum = _mm256_madd_epi16(_mm256_set1_epi16(1), dot); // 8 int32 + // reduce 8 lanes: lo+hi -> 4 lanes, hadd -> 2, hadd -> 1 (a single hadd pair only folds 4 lanes) + const __m128i t = _mm_add_epi32(_mm256_castsi256_si128(sum), _mm256_extracti128_si256(sum, 1)); + const __m128i h = _mm_hadd_epi32(t, t); + return _mm_cvtsi128_si32(_mm_hadd_epi32(h, h)); +} +// sum 16 unsigned bytes to int32 +static int32_t tiled_byte_sum_16(const uint8_t * p) { + const __m128i dot = _mm_maddubs_epi16(_mm_loadu_si128((const __m128i *) p), _mm_set1_epi8(1)); + const __m128i sum = _mm_madd_epi16(_mm_set1_epi16(1), dot); // 4 int32 + const __m128i h = _mm_hadd_epi32(sum, sum); + return _mm_cvtsi128_si32(_mm_hadd_epi32(h, h)); +} +#endif + +// act-bias: precompute the per-weight-row correction corr = 128*sum_s scales[s]*w_bsum[s] where +// w_bsum[s] = sum of the debiased weight bytes in subblock s. Stored in mins[r][0] (unused for +// HAS_MIN = false, which is exactly the act-bias case). Reused across all act bands and weight groups. +// No-op off AVX2 (the act-bias MAC is AVX2-only) and when corr is false (iq4_xs / BIAS = 0). +template +void tiled_repack_src0(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr) { +#if defined(__AVX2__) + if (!corr) { + return; + } + constexpr int NB = TILED_TILE_K / SUBBLK; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const uint8_t * q = &tile->q[r * (num_k * TILED_TILE_K) + slab * TILED_TILE_K]; + const int32_t * scales = &tile->scales[r * nb_stride + slab * NB]; + int32_t corr_val = 0; + for (int s = 0; s < NB; s++) { + const int32_t qsum = (SUBBLK == 32) ? tiled_byte_sum_32(&q[s * SUBBLK]) : tiled_byte_sum_16(&q[s * SUBBLK]); + corr_val += 128 * scales[s] * (qsum - SUBBLK * BIAS); + } + tile->mins[r * nb_stride + slab * NB] = corr_val; + } + } +#else + GGML_UNUSED(tile); GGML_UNUSED(n_rows); GGML_UNUSED(num_k); GGML_UNUSED(BIAS); GGML_UNUSED(corr); +#endif +} + +template void tiled_repack_src0<16>(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr); +template void tiled_repack_src0<32>(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr); + + diff --git a/ggml/src/ggml-cpu/tiled/tiled-kernel.h b/ggml/src/ggml-cpu/tiled/tiled-kernel.h new file mode 100644 index 0000000000..b9df2c93c6 --- /dev/null +++ b/ggml/src/ggml-cpu/tiled/tiled-kernel.h @@ -0,0 +1,180 @@ +#pragma once + +// Tiled matmul kernel API: tile structs, kernel definitions + +// Currently only optimized for x86, new architectures should implement: +// tiled_run_microtile: 16x16 microkernel +// tiled_repack_src0: Optional repack/recalculation of src0, per macrotile +// tiled_repack_src1: Optional repack/recalculation of src1, per microtile-band +// bit unpacking routines: tiled_unpk_nib4, tiled_unpk_2bit, tiled_unpk_or +// LUT value expansion routines: tiled_lut8, tiled_unpk_sign32, tiled_unpk_tern8 + +#define GGML_COMMON_DECL_C +#include "ggml-common.h" + +#include +#include + +#if defined(__AVX2__) +#include +#endif + +#define TILED_TILE_K 256 // one QK_K block +#define TILED_TILE_ROWS 256 // max window rows, ragged at edges +#define TILED_MICRO 16 // microtile edge (also the bsums code-sum granularity) +#define TILED_WS_SLOT (512 * 1024) // per-thread workspace slot, a clean 512KB (multiple of 64B) + +// src0 tile: weight side, shared by all formats. +// scales/mins are sized for the max subblock count (SUBBLK=16); +// SUBBLK=32 formats index at stride 8 and leave the slack unused. +struct tiled_tile_src0 { + static constexpr int NB_MAX = TILED_TILE_K / 16; // max subblocks per 256-elem block + + alignas(64) uint8_t q[TILED_TILE_ROWS * TILED_TILE_K]; // unsigned quants, widened to uint8 + float d[TILED_TILE_ROWS]; // One d from each input block, widened to f32 + float dmin[TILED_TILE_ROWS]; // dmin from each input block (if applicable), widened to F32 + int32_t scales[TILED_TILE_ROWS * NB_MAX]; // per-subblock scale, stored as int32_t + int32_t mins[TILED_TILE_ROWS * NB_MAX]; // per-subblock min, used when HAS_MIN +}; + +// src1 tile: built from q8_K (wdata) +struct tiled_tile_src1 { + // q8 codes, one byte per element. Note for VNNI these are reshaped + transposed to be suitable for dpbusd. + alignas(64) int8_t q[TILED_TILE_ROWS * TILED_TILE_K]; + // per-16 code sums from q8_k (int16), widened to int32 so the kernels load them directly, no per-use cvt + alignas(64) int32_t bsums[(TILED_TILE_K / 16) * TILED_TILE_ROWS]; + // f32 (not f16): q8_k stores fp16, the unpack converts once + float d[TILED_TILE_ROWS]; +}; + +// per-thread workspace: all tiled state lives here, allocated in wdata (one slot per thread) +struct tiled_ws { + tiled_tile_src0 src0; + tiled_tile_src1 src1; + alignas(64) float acc[TILED_TILE_ROWS * TILED_TILE_ROWS]; +}; + +static_assert(sizeof(tiled_ws) <= TILED_WS_SLOT, "tiled workspace exceeds the 512KB per-thread slot"); +// the slot base is 64B-aligned and TILED_WS_SLOT is a multiple of 64B, so each per-thread slot +// is 64B-aligned; these pin the tile fields and the acc buffer at aligned offsets within a slot +static_assert(offsetof(tiled_ws, src0) % 64 == 0, "src0 not 64B-aligned in the workspace"); +static_assert(offsetof(tiled_ws, src1) % 64 == 0, "src1 not 64B-aligned in the workspace"); +static_assert(offsetof(tiled_ws, acc) % 64 == 0, "acc not 64B-aligned in the workspace"); + +// unpack primitives for reading quants, defined as inline here to keep arch-specific code in kernel.h/.cpp +// If this section gets too hairy later, we can break up into separate includes. +#if defined(__AVX2__) +// packed 4-bit codes -> low nibbles (lo) + high nibbles (hi) +inline void tiled_unpk_nib4(const uint8_t * src, uint8_t * lo, uint8_t * hi) { + const __m256i v = _mm256_loadu_si256((const __m256i *) src); + // mask before the lane shift so bits do not cross byte boundaries + _mm256_storeu_si256((__m256i *) lo, _mm256_and_si256(v, _mm256_set1_epi8(0x0F))); + _mm256_storeu_si256((__m256i *) hi, _mm256_srli_epi32(_mm256_and_si256(v, _mm256_set1_epi8((int8_t) 0xF0)), 4)); +} +// 2-bit values at bit offset S +template inline void tiled_unpk_2bit(const uint8_t * src, uint8_t * dst) { + _mm256_storeu_si256((__m256i *) dst, _mm256_and_si256( + _mm256_srli_epi32(_mm256_loadu_si256((const __m256i *) src), S), _mm256_set1_epi8(0x03))); +} +// OR the M-bit value at bit offset S of src into bit offset D of dst +template +inline void tiled_unpk_or(uint8_t * dst, const uint8_t * src) { + const __m256i v = _mm256_slli_epi32(_mm256_and_si256( + _mm256_srli_epi32(_mm256_loadu_si256((const __m256i *) src), S), _mm256_set1_epi8((uint8_t) M)), D); + _mm256_storeu_si256((__m256i *) dst, _mm256_or_si256(_mm256_loadu_si256((const __m256i *) dst), v)); +} + + +// Unpacking kernels for IQ quants + +// LUT value expansion for the LUT-based formats (iq4_xs, iq grids): the bit unpackers +// above give the indices, these expand 8/16 of them to widened codes in one pass +// 16-entry byte LUT: dst[j] = lut[src[j]] (16 bytes) +inline void tiled_lut8(const uint8_t * lut, const uint8_t * src, uint8_t * dst) { + _mm_storeu_si128((__m128i *) dst, _mm_shuffle_epi8(_mm_loadu_si128((const __m128i *) lut), + _mm_loadu_si128((const __m128i *) src))); +} + + +// 32 grid magnitudes (4 x 64-bit groups g0..g3, 8 values each) + 4 sign bytes +// (byte l signs values 8*l .. 8*l+7) -> codes stored as (value + 128) +inline void tiled_unpk_sign32(uint64_t g0, uint64_t g1, uint64_t g2, uint64_t g3, + const uint8_t signs[4], uint8_t * dst32) { + const __m256i v = _mm256_set_epi64x((int64_t) g3, (int64_t) g2, (int64_t) g1, (int64_t) g0); + const __m256i sv = _mm256_shuffle_epi8(_mm256_set1_epi32((int32_t) (signs[0] | signs[1] << 8 | signs[2] << 16 | signs[3] << 24)), + _mm256_setr_epi8(0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, + 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3)); + const __m256i sel = _mm256_set1_epi64x((int64_t) 0x8040201008040201ULL); + // 0xFF in every lane whose sign bit is set; GFNI does the and+compare in one instruction +#ifdef __GFNI__ + const __m256i mask = _mm256_gf2p8affine_epi64_epi8(sel, sv, 0); +#else + const __m256i mask = _mm256_cmpeq_epi8(_mm256_and_si256(sv, sel), sel); +#endif + // v ^ mask - mask negates the signed lanes (v elsewhere); + 128 gives the biased code + const __m256i sgn = _mm256_sub_epi8(_mm256_xor_si256(v, mask), mask); + _mm256_storeu_si256((__m256i *) dst32, _mm256_add_epi8(sgn, _mm256_set1_epi8((int8_t) 128))); +} +// 8 ternary grid bytes (0 = 0, 1 = +1, 0xFF = -1): dst[j] = 128 + delta + 8 * (int8_t) src[j] +inline void tiled_unpk_tern8(const uint8_t * src, int8_t delta, uint8_t * dst) { + const __m128i v = _mm_cvtepi8_epi16(_mm_loadl_epi64((const __m128i *) src)); + const __m128i p = _mm_add_epi16(_mm_slli_epi16(v, 3), _mm_set1_epi16(128 + (int) delta)); + _mm_storel_epi64((__m128i *) dst, _mm_packus_epi16(p, _mm_setzero_si128())); +} +#else + +// Scalar definitions for unpackers. + + +inline void tiled_unpk_nib4(const uint8_t * src, uint8_t * lo, uint8_t * hi) { + for (int l = 0; l < 32; l++) { lo[l] = (uint8_t) (src[l] & 0xF); hi[l] = (uint8_t) (src[l] >> 4); } +} +template +inline void tiled_unpk_2bit(const uint8_t * src, uint8_t * dst) { + for (int l = 0; l < 32; l++) { dst[l] = (uint8_t) ((src[l] >> S) & 3); } +} +template +inline void tiled_unpk_or(uint8_t * dst, const uint8_t * src) { + for (int l = 0; l < 32; l++) { dst[l] = (uint8_t) (dst[l] | (((src[l] >> S) & M) << D)); } +} +inline void tiled_lut8(const uint8_t * lut, const uint8_t * src, uint8_t * dst) { + for (int j = 0; j < 16; j++) { dst[j] = lut[src[j]]; } +} +inline void tiled_unpk_sign32(uint64_t g0, uint64_t g1, uint64_t g2, uint64_t g3, + const uint8_t signs[4], uint8_t * dst32) { + const uint64_t g[4] = { g0, g1, g2, g3 }; + for (int l = 0; l < 4; l++) { + const uint8_t * v = (const uint8_t *) &g[l]; + const uint8_t s = signs[l]; + for (int j = 0; j < 8; j++) { + dst32[8 * l + j] = (s & (1 << j)) ? (uint8_t) (128 - v[j]) : (uint8_t) (128 + v[j]); + } + } +} +// 8 ternary grid bytes (0 = 0, 1 = +1, 0xFF = -1): dst[j] = 128 + delta + 8 * (int8_t) src[j] +inline void tiled_unpk_tern8(const uint8_t * src, int8_t delta, uint8_t * dst) { + for (int j = 0; j < 8; j++) { dst[j] = (uint8_t) (128 + (int) delta + 8 * (int8_t) src[j]); } +} +#endif + +// Accumulate one 16x16 microtile (src0 rows [i0, i0+16), src1 cols [j0, j0+16)) +// over one 256-K slab held in the tiles into a j-major float buffer +// (row width buf_stride): buf[i*buf_stride + j] += partial. +// SUBBLK/HAS_MIN/BIAS are the src0 format constants (see tiled_tile_src0). +// ACTBIAS (AVX2 only): the activation is pre-biased +128 by tiled_repack_src1, +// num_k = K-blocks per row: the tile holds num_k slabs at row stride num_k*256 +// (Default case is num_k=1, 256x256 tiles, we go to longer num_k to improve memory bandwidth when num_rows is small) +template +void tiled_run_microtile(const tiled_tile_src0 & src0, const tiled_tile_src1 & src1, + int i0, int j0, int num_k, int slab, float * buf, int buf_stride); + +// Optional repack, if profitable for the kernel. +// Repacks one 16-row band of src1 codes, called by driver as we reach each 16-row band in outer loop +void tiled_repack_src1(tiled_tile_src1 * src1, int row0, int num_k, bool bias); + +// Optional repack, if profitable for the kernel +// Repack the entire src0 panel in place, called by the driver immediately after dequant +template +void tiled_repack_src0(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr); + + diff --git a/ggml/src/ggml-cpu/tiled/tiled.cpp b/ggml/src/ggml-cpu/tiled/tiled.cpp new file mode 100644 index 0000000000..15a1d0c011 --- /dev/null +++ b/ggml/src/ggml-cpu/tiled/tiled.cpp @@ -0,0 +1,1245 @@ +#include "tiled.h" +#include "tiled-kernel.h" + +#include "ggml-cpu-impl.h" +#include "ggml-cpu.h" +#include "ggml.h" + +// kvalues table (impl section) for the iq4_xs unpack +#define GGML_COMMON_IMPL_CPP +#include "ggml-common.h" + +#include +#include +#include + +#include + +#define UNUSED GGML_UNUSED + +// unpack routines for various quant types src0 +static void tiled_unpack_src0(const block_q4_K * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 8; // 32-wide subblocks + // 12-byte packed scale/min decode, same extraction as the reference kernels + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_q4_K & x = rows[r * row_stride + slab]; + + tile->d[d_off] = ggml_fp16_to_fp32(x.d); + tile->dmin[d_off] = ggml_fp16_to_fp32(x.dmin); + + uint32_t utmp[4]; + memcpy(utmp, x.scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + const uint8_t * scales = (const uint8_t *) &utmp[0]; + const uint8_t * mins = (const uint8_t *) &utmp[2]; + for (int s = 0; s < NB; s++) { + tile->scales[s_off + s] = (int32_t) scales[s]; + tile->mins[s_off + s] = (int32_t) mins[s]; + } + + // extract the 4-bit codes (low 4 + high 4), same extraction as the reference kernels + uint8_t * q = &tile->q[q_off]; + tiled_unpk_nib4(x.qs + 0, q + 0, q + 32); + tiled_unpk_nib4(x.qs + 32, q + 64, q + 96); + tiled_unpk_nib4(x.qs + 64, q + 128, q + 160); + tiled_unpk_nib4(x.qs + 96, q + 192, q + 224); + } + } +} + +static void tiled_unpack_src0(const block_q5_K * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 8; // 32-wide subblocks + // 12-byte packed scale/min decode, same extraction as the reference kernels + static const uint32_t kmask1 = 0x3f3f3f3f; + static const uint32_t kmask2 = 0x0f0f0f0f; + static const uint32_t kmask3 = 0x03030303; + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_q5_K & x = rows[r * row_stride + slab]; + + tile->d[d_off] = ggml_fp16_to_fp32(x.d); + tile->dmin[d_off] = ggml_fp16_to_fp32(x.dmin); + + uint32_t utmp[4]; + memcpy(utmp, x.scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + + const uint8_t * scales = (const uint8_t *) &utmp[0]; + const uint8_t * mins = (const uint8_t *) &utmp[2]; + for (int s = 0; s < NB; s++) { + tile->scales[s_off + s] = (int32_t) scales[s]; + tile->mins[s_off + s] = (int32_t) mins[s]; + } + + // extract the 5-bit codes (4 low bits + 1 high bit), same as the generic kernels: + // 64-element chunk j uses qh bits 2j (low 32) and 2j+1 (high 32); OR adds the 5th bit + // (no overlap with the 4-bit codes, identical to the reference ADD) + uint8_t * q = &tile->q[q_off]; + tiled_unpk_nib4(x.qs + 0, q + 0, q + 32); + tiled_unpk_nib4(x.qs + 32, q + 64, q + 96); + tiled_unpk_nib4(x.qs + 64, q + 128, q + 160); + tiled_unpk_nib4(x.qs + 96, q + 192, q + 224); + tiled_unpk_or<0, 4, 1>(q + 0, x.qh); + tiled_unpk_or<1, 4, 1>(q + 32, x.qh); + tiled_unpk_or<2, 4, 1>(q + 64, x.qh); + tiled_unpk_or<3, 4, 1>(q + 96, x.qh); + tiled_unpk_or<4, 4, 1>(q + 128, x.qh); + tiled_unpk_or<5, 4, 1>(q + 160, x.qh); + tiled_unpk_or<6, 4, 1>(q + 192, x.qh); + tiled_unpk_or<7, 4, 1>(q + 224, x.qh); + } + } +} + +static void tiled_unpack_src0(const block_q6_K * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 16; // 16-wide subblocks + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_q6_K & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d); + + // 6-bit code = 4 low bits (ql) | 2 high bits (qh); see ggml_vec_dot_q6_K_q8_K_generic + // per half the lanes are [ql lo(0:32)] [ql lo(32:64)] [ql hi(0:32)] [ql hi(32:64)] + uint8_t * q = &tile->q[q_off]; + for (int half = 0; half < 2; half++) { + uint8_t * out = q + 128 * half; + tiled_unpk_nib4(x.ql + 64 * half + 0, out + 0, out + 64); + tiled_unpk_nib4(x.ql + 64 * half + 32, out + 32, out + 96); + tiled_unpk_or<0, 4, 3>(out + 0, x.qh + 32 * half); + tiled_unpk_or<2, 4, 3>(out + 32, x.qh + 32 * half); + tiled_unpk_or<4, 4, 3>(out + 64, x.qh + 32 * half); + tiled_unpk_or<6, 4, 3>(out + 96, x.qh + 32 * half); + } + // scale is a plain int8 per 16-element subblock (16 per 256-K) + for (int s = 0; s < NB; s++) { tile->scales[s_off + s] = (int32_t) (int8_t) x.scales[s]; } + } + } +} + +static void tiled_unpack_src0(const block_q3_K * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 16; // 16-wide subblocks + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_q3_K & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d); + + // 3-bit code = 2 low bits (qs) | (1 high bit from hmask << 2) + // element e (0..255): half=e>>7, el=e&127, group=el>>5, l=el&31 + // low2 = (qs[half*32 + l] >> 2*group) & 3 + // high = (hmask[l] >> (half*4 + group)) & 1 + // see ggml_vec_dot_q3_K_q8_K_generic + uint8_t * q = &tile->q[q_off]; + const uint8_t * s0 = x.qs; + const uint8_t * s1 = x.qs + 32; + uint8_t * o0 = q; + uint8_t * o1 = q + 128; + tiled_unpk_2bit<0>(s0, o0); tiled_unpk_or<0, 2, 1>(o0, x.hmask); + tiled_unpk_2bit<2>(s0, o0 + 32); tiled_unpk_or<1, 2, 1>(o0 + 32, x.hmask); + tiled_unpk_2bit<4>(s0, o0 + 64); tiled_unpk_or<2, 2, 1>(o0 + 64, x.hmask); + tiled_unpk_2bit<6>(s0, o0 + 96); tiled_unpk_or<3, 2, 1>(o0 + 96, x.hmask); + tiled_unpk_2bit<0>(s1, o1); tiled_unpk_or<4, 2, 1>(o1, x.hmask); + tiled_unpk_2bit<2>(s1, o1 + 32); tiled_unpk_or<5, 2, 1>(o1 + 32, x.hmask); + tiled_unpk_2bit<4>(s1, o1 + 64); tiled_unpk_or<6, 2, 1>(o1 + 64, x.hmask); + tiled_unpk_2bit<6>(s1, o1 + 96); tiled_unpk_or<7, 2, 1>(o1 + 96, x.hmask); + // 6-bit scale decode (same kmask trick as the reference), stored as (scales - 32) + static const uint32_t kmask1 = 0x03030303; + static const uint32_t kmask2 = 0x0f0f0f0f; + uint32_t auxs[4]; + memcpy(auxs, x.scales, 12); + const uint32_t tmp = auxs[2]; + auxs[2] = ((auxs[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + auxs[3] = ((auxs[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + auxs[0] = (auxs[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + auxs[1] = (auxs[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + const int8_t * scales = (const int8_t *) &auxs[0]; + for (int s = 0; s < NB; s++) { tile->scales[s_off + s] = (int32_t) scales[s] - 32; } + } + } +} + +static void tiled_unpack_src0(const block_q2_K * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 16; // 16-wide subblocks + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_q2_K & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d); + tile->dmin[d_off] = ggml_fp16_to_fp32(x.dmin); + + // 2-bit code: element e -> half=e>>7, el=e&127 + // byte = half*32 + (el & 31), shift = 2*(el >> 5) + // see ggml_vec_dot_q2_K_q8_K_generic + uint8_t * q = &tile->q[q_off]; + for (int half = 0; half < 2; half++) { + const uint8_t * s = x.qs + 32 * half; + uint8_t * out = q + 128 * half; + tiled_unpk_2bit<0>(s, out + 0); + tiled_unpk_2bit<2>(s, out + 32); + tiled_unpk_2bit<4>(s, out + 64); + tiled_unpk_2bit<6>(s, out + 96); + } + // scale/min packed in one byte per 16-element subblock: low 4 bits = scale, high 4 = min + for (int s = 0; s < NB; s++) { + tile->scales[s_off + s] = (int32_t) (x.scales[s] & 0xF); + tile->mins[s_off + s] = (int32_t) (x.scales[s] >> 4); + } + } + } +} + +// iq4_xs: 4-bit codes through the kvalues_iq4nl LUT, 6-bit scales per 32, no min. +// LUT values are stored as kvalues + 128 so the +128 shift the kernel applies to +// the activations cancels against the bsums correction (BIAS = 128); the +128 is +// folded into the table up front so the expansion is a plain lookup +static void tiled_unpack_src0(const block_iq4_xs * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 8; // 32-wide subblocks + + static const uint8_t lut[16] = { + (uint8_t) (kvalues_iq4nl[0] + 128), (uint8_t) (kvalues_iq4nl[1] + 128), + (uint8_t) (kvalues_iq4nl[2] + 128), (uint8_t) (kvalues_iq4nl[3] + 128), + (uint8_t) (kvalues_iq4nl[4] + 128), (uint8_t) (kvalues_iq4nl[5] + 128), + (uint8_t) (kvalues_iq4nl[6] + 128), (uint8_t) (kvalues_iq4nl[7] + 128), + (uint8_t) (kvalues_iq4nl[8] + 128), (uint8_t) (kvalues_iq4nl[9] + 128), + (uint8_t) (kvalues_iq4nl[10] + 128), (uint8_t) (kvalues_iq4nl[11] + 128), + (uint8_t) (kvalues_iq4nl[12] + 128), (uint8_t) (kvalues_iq4nl[13] + 128), + (uint8_t) (kvalues_iq4nl[14] + 128), (uint8_t) (kvalues_iq4nl[15] + 128), + }; + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_iq4_xs & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d); + + // 6-bit scale per 32, stored as (ls - 32); same extraction as dequantize_row_iq4_xs + for (int s = 0; s < NB; s++) { + const int ls = ((x.scales_l[s / 2] >> 4 * (s % 2)) & 0xf) | (((x.scales_h >> 2 * s) & 3) << 4); + tile->scales[s_off + s] = (int32_t) ls - 32; + } + + // 4-bit codes through the LUT: low nibbles of a byte pair come first + uint8_t * q = &tile->q[q_off]; + uint8_t lo[32], hi[32]; + for (int u = 0; u < 4; u++) { + tiled_unpk_nib4(x.qs + 32 * u, lo, hi); + tiled_lut8(lut, lo + 0, q + 64 * u + 0); + tiled_lut8(lut, hi + 0, q + 64 * u + 16); + tiled_lut8(lut, lo + 16, q + 64 * u + 32); + tiled_lut8(lut, hi + 16, q + 64 * u + 48); + } + } + } +} + +// iq2_xxs: 2-bit grids through the iq2xxs_grid LUT, 4-bit scale per 32, no min +// codes stored as (value + 128), matching BIAS = 128 +static void tiled_unpack_src0(const block_iq2_xxs * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 8; // 32-wide subblocks + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_iq2_xxs & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d) * 0.125f; + + uint32_t aux32[2]; + const uint8_t * aux8 = (const uint8_t *) aux32; + uint64_t g[4]; + uint8_t signs4[4]; + for (int ib32 = 0; ib32 < NB; ib32++) { + memcpy(aux32, x.qs + 4 * ib32, 2 * sizeof(uint32_t)); + tile->scales[s_off + ib32] = (int32_t) (2 * (aux32[1] >> 28) + 1); + for (int l = 0; l < 4; l++) { + g[l] = iq2xxs_grid[aux8[l]]; + signs4[l] = ksigns_iq2xs[(aux32[1] >> (7 * l)) & 127]; + } + tiled_unpk_sign32(g[0], g[1], g[2], g[3], signs4, &tile->q[q_off + 32 * ib32]); + } + } + } +} + +// iq2_xs: 2-bit grids through the iq2xs_grid LUT, 4-bit scale per 16 (two per 32), no min +// codes stored as (value + 128), matching BIAS = 128 +static void tiled_unpack_src0(const block_iq2_xs * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 16; // 16-wide subblocks + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_iq2_xs & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d) * 0.125f; + + uint64_t g[4]; + uint8_t signs4[4]; + for (int ib32 = 0; ib32 < QK_K / 32; ib32++) { + tile->scales[s_off + 2 * ib32 + 0] = (int32_t) (2 * (x.scales[ib32] & 0xf) + 1); + tile->scales[s_off + 2 * ib32 + 1] = (int32_t) (2 * (x.scales[ib32] >> 4) + 1); + const uint16_t * q = x.qs + 4 * ib32; + for (int l = 0; l < 4; l++) { + g[l] = iq2xs_grid[q[l] & 511]; + signs4[l] = ksigns_iq2xs[q[l] >> 9]; + } + tiled_unpk_sign32(g[0], g[1], g[2], g[3], signs4, &tile->q[q_off + 32 * ib32]); + } + } + } +} + +// iq2_s: 2-bit grids through the iq2s_grid LUT, 4-bit scale per 16 (two per 32), no min +// codes stored as (value + 128), matching BIAS = 128 +static void tiled_unpack_src0(const block_iq2_s * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 16; // 16-wide subblocks + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_iq2_s & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d) * 0.125f; + + const uint8_t * qs = x.qs; + const uint8_t * signs = x.qs + QK_K / 8; // packed sign bytes share the qs array, same as the reference + uint64_t g[4]; + for (int ib32 = 0; ib32 < QK_K / 32; ib32++) { + tile->scales[s_off + 2 * ib32 + 0] = (int32_t) (2 * (x.scales[ib32] & 0xf) + 1); + tile->scales[s_off + 2 * ib32 + 1] = (int32_t) (2 * (x.scales[ib32] >> 4) + 1); + for (int l = 0; l < 4; l++) { + g[l] = iq2s_grid[qs[l] | (x.qh[ib32] << (8 - 2 * l) & 0x300)]; + } + tiled_unpk_sign32(g[0], g[1], g[2], g[3], signs, &tile->q[q_off + 32 * ib32]); + qs += 4; + signs += 4; + } + } + } +} + +// iq3_xxs: 3-bit grids through the iq3xxs_grid LUT, 4-bit scale per 32, no min +// codes stored as (value + 128), matching BIAS = 128 +static void tiled_unpack_src0(const block_iq3_xxs * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 8; // 32-wide subblocks + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_iq3_xxs & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d) * 0.25f; + + const uint8_t * qs = x.qs; + const uint8_t * scales_and_signs = x.qs + QK_K / 4; // 4 bytes per 32: code bits in the top nibble, signs in 7-bit chunks + uint64_t g[4]; + uint8_t signs4[4]; + for (int ib32 = 0; ib32 < QK_K / 32; ib32++) { + const uint32_t aux32 = *(const uint32_t *) (scales_and_signs + 4 * ib32); + tile->scales[s_off + ib32] = (int32_t) (2 * (aux32 >> 28) + 1); + for (int l = 0; l < 4; l++) { + // bytes 8*l .. 8*l+7: low entry e1 (4 values) then high entry e2 (4 values) + g[l] = (uint64_t) iq3xxs_grid[qs[2 * l + 1]] << 32 | iq3xxs_grid[qs[2 * l + 0]]; + signs4[l] = ksigns_iq2xs[(aux32 >> (7 * l)) & 127]; + } + tiled_unpk_sign32(g[0], g[1], g[2], g[3], signs4, &tile->q[q_off + 32 * ib32]); + qs += 8; + } + } + } +} + +// iq3_s: 3-bit grids through the iq3s_grid LUT, 4-bit scale per 32 (two per scale byte), no min +// codes stored as (value + 128), matching BIAS = 128 +static void tiled_unpack_src0(const block_iq3_s * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 8; // 32-wide subblocks + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_iq3_s & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d); + + const uint8_t * qs = x.qs; + const uint8_t * qh = x.qh; + const uint8_t * signs = x.signs; + for (int ib32 = 0; ib32 < QK_K / 32; ib32 += 2) { + tile->scales[s_off + ib32 + 0] = (int32_t) (1 + 2 * (x.scales[ib32 / 2] & 0xf)); + tile->scales[s_off + ib32 + 1] = (int32_t) (1 + 2 * (x.scales[ib32 / 2] >> 4)); + for (int h = 0; h < 2; h++) { + uint64_t g[4]; + for (int l = 0; l < 4; l++) { + // bytes 8*l .. 8*l+7: low entry e1 (4 values) then high entry e2 (4 values) + g[l] = (uint64_t) iq3s_grid[qs[2 * l + 1] | ((qh[h] << (7 - 2 * l)) & 256)] << 32 + | iq3s_grid[qs[2 * l + 0] | ((qh[h] << (8 - 2 * l)) & 256)]; + } + tiled_unpk_sign32(g[0], g[1], g[2], g[3], signs, &tile->q[q_off + (ib32 + h) * 32]); + qs += 8; + signs += 4; + } + qh += 2; + } + } + } +} + +// iq1_s: ternary grid (values +-1/0) scaled by 8 to leave room for the +-1 delta offset, +// the /8 folds into d; 3-bit scale per 32, no min +// codes stored as (8 * grid + delta + 128), matching BIAS = 128 +static void tiled_unpack_src0(const block_iq1_s * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 8; // 32-wide subblocks + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_iq1_s & x = rows[r * row_stride + slab]; + tile->d[d_off] = ggml_fp16_to_fp32(x.d) * 0.125f; + + const uint8_t * qs = x.qs; + for (int ib = 0; ib < QK_K / 32; ib++) { + const uint16_t hw = x.qh[ib]; + tile->scales[s_off + ib] = (int32_t) (2 * ((hw >> 12) & 7) + 1); + const int8_t delta = (hw & 0x8000) ? -1 : 1; + for (int l = 0; l < 4; l++) { + const uint64_t entry = iq1s_grid[qs[l] | (((hw >> (3 * l)) & 7) << 8)]; + tiled_unpk_tern8((const uint8_t *) &entry, delta, &tile->q[q_off + 32 * ib + 8 * l]); + } + qs += 4; + } + } + } +} + +// iq1_m: like iq1_s but the fp16 scale is packed across the 4 scale bytes (no d field) and the +// delta offset is per 8, giving one scale per 16; 3-bit scale per 16, no min +// codes stored as (8 * grid + delta + 128), matching BIAS = 128 +static void tiled_unpack_src0(const block_iq1_m * rows, int64_t row_stride, int n_rows, tiled_tile_src0 * tile, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + constexpr int NB = 16; // 16-wide subblocks + + const int qk_stride = num_k * TILED_TILE_K; + const int nb_stride = NB * num_k; + for (int slab = 0; slab < num_k; slab++) { + for (int r = 0; r < n_rows; r++) { + const int d_off = slab * TILED_MICRO + r; + const int q_off = r * qk_stride + slab * TILED_TILE_K; + const int s_off = r * nb_stride + slab * NB; + const block_iq1_m & x = rows[r * row_stride + slab]; + + const uint16_t * sc = (const uint16_t *) x.scales; + iq1m_scale_t scale; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); + tile->d[d_off] = ggml_fp16_to_fp32(scale.f16) * 0.125f; + + const uint8_t * qs = x.qs; + const uint8_t * qh = x.qh; + for (int ib = 0; ib < QK_K / 32; ib++) { + const uint16_t hw = sc[ib / 2]; + const int sh = 6 * (ib % 2); + tile->scales[s_off + 2 * ib + 0] = (int32_t) (2 * ((hw >> sh) & 7) + 1); + tile->scales[s_off + 2 * ib + 1] = (int32_t) (2 * ((hw >> (sh + 3)) & 7) + 1); + + const uint16_t idx[4] = { + (uint16_t) (qs[0] | ((qh[0] << 8) & 0x700)), + (uint16_t) (qs[1] | ((qh[0] << 4) & 0x700)), + (uint16_t) (qs[2] | ((qh[1] << 8) & 0x700)), + (uint16_t) (qs[3] | ((qh[1] << 4) & 0x700)), + }; + const int8_t delta[4] = { + (int8_t) ((qh[0] & 0x08) ? -1 : 1), (int8_t) ((qh[0] & 0x80) ? -1 : 1), + (int8_t) ((qh[1] & 0x08) ? -1 : 1), (int8_t) ((qh[1] & 0x80) ? -1 : 1), + }; + for (int l = 0; l < 4; l++) { + const uint64_t entry = iq1s_grid[idx[l]]; + tiled_unpk_tern8((const uint8_t *) &entry, delta[l], &tile->q[q_off + 32 * ib + 8 * l]); + } + qs += 4; + qh += 2; + } + } + } +} + +// unpack src1 tile from q8_K rows. num_k = K-blocks per row: each row decodes num_k consecutive +// source blocks (rows[r][kblk + slab]) into the tile at row stride num_k*256, so each activation +// row is one long stream. num_k=1 is the standard single-slab path (byte-identical to before). +static void tiled_unpack_src1_q8_K(const block_q8_K * const * rows, int n_rows, tiled_tile_src1 * tile, + int kblk, int num_k) { + GGML_ASSERT(n_rows <= TILED_TILE_ROWS); + const int n_padded = (n_rows + TILED_MICRO - 1) & ~(TILED_MICRO - 1); + const int bs_stride = (num_k == 1) ? TILED_TILE_ROWS : TILED_MICRO; + const int n16 = TILED_TILE_K / TILED_MICRO; + // natural [row][k] fill, zero-pad ragged tail; codes stay natural here, the driver + // repacks them just-in-time via tiled_repack_src1 (per 16-row band on the standard path, + // the whole chunk on the narrow path) so each repacked region is L1-hot for its uses + for (int slab = 0; slab < num_k; slab++) { + const int q_off = slab * TILED_TILE_K; + for (int r = 0; r < n_padded; r++) { + if (r < n_rows) { + memcpy(&tile->q[r * (num_k * TILED_TILE_K) + q_off], rows[r][kblk + slab].qs, TILED_TILE_K); + } else { + memset(&tile->q[r * (num_k * TILED_TILE_K) + q_off], 0, TILED_TILE_K); + } + } + // d and bsums (ISA-independent) + for (int r = 0; r < n_padded; r++) { + if (r < n_rows) { + const block_q8_K & x = rows[r][kblk + slab]; + for (int s = 0; s < n16; s++) { tile->bsums[(slab * n16 + s) * bs_stride + r] = (int32_t) x.bsums[s]; } + tile->d[slab * TILED_MICRO + r] = x.d; + } else { + for (int s = 0; s < n16; s++) { tile->bsums[(slab * n16 + s) * bs_stride + r] = 0; } + tile->d[slab * TILED_MICRO + r] = 0.0f; + } + } + } +} + +// GGML_CPU_TILED_MM: master switch, on by default. If off, we fast return false and normal vec_dot mul_mat resumes +static bool ggml_tiled_matmul_enabled(void) { + static bool enabled = true; + static std::once_flag flag; + std::call_once(flag, []() { + const char * env = getenv("GGML_CPU_TILED_MM"); + enabled = env == NULL || atoi(env) != 0; + }); + return enabled; +} + +// GGML_CPU_TILED_MM_FORCE: test/bench only, take the tiled path even when unprofitable +static bool ggml_tiled_matmul_forced(void) { + static bool forced = false; + static std::once_flag flag; + std::call_once(flag, []() { + const char * env = getenv("GGML_CPU_TILED_MM_FORCE"); + forced = env != NULL && atoi(env) == 1; + }); + return forced; +} + +// hard constraints shared by the MUL_MAT and MUL_MAT_ID entries; the src0 type gate is the +// entries' type switch + +#if !defined(__AVX512VNNI__) && !defined(__AVX2__) && !defined(__AVX__) +static bool ggml_tiled_supported(const struct ggml_tensor * src0, + const struct ggml_tensor * src1) { + UNUSED(src0); + UNUSED(src1); + return false; +} +#else +static bool ggml_tiled_supported(const struct ggml_tensor * src0, + const struct ggml_tensor * src1) { + if (!ggml_tiled_matmul_enabled()) { + return false; + } + + // repack-buffer weights hold a repacked layout, let that kernel handle + if (src0->extra != NULL) { + return false; + } + // Supported quant types for src0 + switch (src0->type) { + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q2_K: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + return true; + default: + return false; + } + if (src1->type != GGML_TYPE_F32 && src1->type != GGML_TYPE_Q8_K) { + return false; + } + + if (src1->type == GGML_TYPE_Q8_K && !ggml_is_contiguous(src1)) { + // We can handle noncontiguous floats because we're repacking to q8_k anyways + return false; + } + return true; +} +#endif + +// per-thread workspace slot size (0 when tiled is disabled or unsupported on this arch) +static size_t ggml_tiled_ws_size(void) { + if (!ggml_tiled_matmul_enabled()) { + return 0; + } + return TILED_WS_SLOT; // clean 512KB slot, rounded up from sizeof(tiled_ws) +} + +size_t ggml_tiled_wdata_size(int n_tasks, struct ggml_tensor * dst) { + if (! ggml_tiled_supported(dst->src[0], dst->src[1])) { + return 0; // unsupported, don't allocate + } + return 64 + n_tasks * ggml_tiled_ws_size(); // 64 for alignment plus one 512KB slot per thread +} + + +// narrow-path K chunk: how many K's worth of weights to read contiguously into long/skinny tiles +// Assumes 32kb L1 cache budget +static int ggml_tiled_narrow_k_extent(int64_t n_rows, int64_t ne00) { + int l1 = 31 * 1024; // fit weights in 31k to leave room for scales + + int rows = (int) n_rows; + if (rows < TILED_MICRO) rows = TILED_MICRO; + int ke = l1 / (TILED_MICRO + rows); + ke &= ~(TILED_TILE_K - 1); // floor to a multiple of 256 + const int ke_max = TILED_TILE_ROWS * TILED_TILE_K / TILED_MICRO; // buffer ceiling + if (ke > ke_max) ke = ke_max; + while (ke >= TILED_TILE_K && ne00 % ke != 0) { + ke -= TILED_TILE_K; + } + return ke < TILED_TILE_K ? 0 : ke; +} + +// Writeback of the 256x256 window: buf is j-major (row stride buf_stride), dst is i-major (column stride dst_stride). +static void tiled_store_window(const float * buf, int n_src0, int n_src1, int buf_stride, + float * dst, size_t dst_stride) { + int ri = 0; + for (; ri + 16 <= n_src0; ri += 16) { + int rj = 0; + for (; rj + 8 <= n_src1; rj += 8) { + float r[16][8]; + for (int t = 0; t < 16; t++) { + for (int u = 0; u < 8; u++) { + r[t][u] = buf[(ri + t) * buf_stride + rj + u]; + } + } + for (int u = 0; u < 8; u++) { + for (int t = 0; t < 16; t++) { + dst[(ri + t) + (size_t) (rj + u) * dst_stride] = r[t][u]; + } + } + } + // ragged j tail + for (; rj < n_src1; rj++) { + for (int t = 0; t < 16; t++) { + dst[(ri + t) + (size_t) rj * dst_stride] = buf[(ri + t) * buf_stride + rj]; + } + } + } + // ragged i tail + for (; ri < n_src0; ri++) { + for (int j = 0; j < n_src1; j++) { + dst[ri + (size_t) j * dst_stride] = buf[ri * buf_stride + j]; + } + } +} + +// MUL_MAT_ID (MoE): src1 rows are gathered per output row (expert dispatch) and dst rows are +// scattered back. The expert's cne1 gathered q8_K rows are staged once per expert into +// thread-local scratch (contiguous rows, plus the [k/4][row][4] interleave on VNNI) so the +// unpack and microtile are reused unchanged from mul_mat. + +// src0 rows per g group; finer than the TILED_TILE_K kernel tile so all threads stay busy at small ne01 +#define TILED_MMID_GROUP 64 + +// like tiled_store_window, but the dst columns are not contiguous: col_ptrs[j] points at the +// start of dst column j, whose rows are contiguous (dim 0, nb0 == 4 bytes). +static void tiled_store_window_scatter(const float * buf, int n_src0, int n_src1, int buf_stride, + float * const * col_ptrs) { + int ri = 0; + for (; ri + 16 <= n_src0; ri += 16) { + int rj = 0; + for (; rj + 8 <= n_src1; rj += 8) { + float r[16][8]; + for (int t = 0; t < 16; t++) { + for (int u = 0; u < 8; u++) { + r[t][u] = buf[(ri + t) * buf_stride + rj + u]; + } + } + for (int u = 0; u < 8; u++) { + for (int t = 0; t < 16; t++) { + col_ptrs[rj + u][ri + t] = r[t][u]; + } + } + } + // ragged j tail + for (; rj < n_src1; rj++) { + for (int t = 0; t < 16; t++) { + col_ptrs[rj][ri + t] = buf[(ri + t) * buf_stride + rj]; + } + } + } + // ragged i tail + for (; ri < n_src0; ri++) { + for (int j = 0; j < n_src1; j++) { + col_ptrs[j][ri] = buf[ri * buf_stride + j]; + } + } +} + + +// one (g, k) macrotile: zero the acc window, sweep K in 256 element slabs, scatter the result rows +// rows points at this k window's routed src1 rows, row r at its base block +template +static void tiled_mmid_gemm_window(struct ggml_tensor * dst, const struct ggml_tensor * src0, + const char * src0_cur, int64_t r, int64_t k, int64_t nrows, + const int32_t * expert_rows, + const block_q8_K * const * rows, + tiled_ws * ws) { + const int64_t ne00 = src0->ne[0]; + const int64_t ne01 = src0->ne[1]; + + const int64_t r_end = MIN(r + TILED_MMID_GROUP, ne01); + const int n_src0 = (int) (r_end - r); + + const size_t src0_bs = ggml_type_size(src0->type); + const int64_t src0_stride = src0->nb[1] / src0_bs; + + // the window is at most TILED_MMID_GROUP x TILED_TILE_K, so zero only that region of acc + for (int64_t i = 0; i < n_src0; i++) { + memset(&ws->acc[i * TILED_TILE_ROWS], 0, nrows * sizeof(float)); + } + + // scattered writeback: column m goes to its routed dst row; r * nb[0] is the window row offset + float * col_ptrs[TILED_TILE_ROWS]; + for (int64_t m = 0; m < nrows; m++) { + col_ptrs[m] = (float *) ((char *) dst->data + r * dst->nb[0] + expert_rows[2 * (k + m) + 0] * dst->nb[1] + + expert_rows[2 * (k + m) + 1] * dst->nb[2]); + } + + // narrow path (small nrows): we're memory bound in this case so do longer, skinny tiled in order + // to have contiguous reads and improve memory bandwidth + if (nrows <= TILED_MICRO) { + const int k_extent = ggml_tiled_narrow_k_extent(nrows, ne00); + const int num_k = k_extent / TILED_TILE_K; + for (int64_t k0 = 0; k0 < ne00; k0 += k_extent) { + const int kstart = (int) (k0 / TILED_TILE_K); + // activation: one 16-row band (nrows <= 16). Unpack the whole k_extent chunk (all + // slabs) so each activation row is a long stream, then repack each slab in place + tiled_unpack_src1_q8_K(rows, nrows, &ws->src1, kstart, num_k); + // repack the whole k_extent chunk in place: num_k slabs x 4 tiles, 16 rows, row stride num_k*256 + tiled_repack_src1(&ws->src1, 0, num_k, ACTBIAS); + // weight groups (16 at a time): unpack the k_extent chunk of the group (the long per-row + // read), then one standard MAC per slab accumulating into the same acc rows + for (int64_t ir0 = r; ir0 < r_end; ir0 += TILED_MICRO) { + const int n0 = (int) MIN(TILED_MICRO, r_end - ir0); + const B * wbase = (const B *) (src0_cur + ir0 * src0->nb[1] + kstart * src0_bs); + tiled_unpack_src0(wbase, src0_stride, n0, &ws->src0, num_k); + tiled_repack_src0(&ws->src0, n0, num_k, BIAS, ACTBIAS); + float * buf = ws->acc + (ir0 - r) * TILED_TILE_ROWS; + for (int slab = 0; slab < num_k; slab++) { + tiled_run_microtile(ws->src0, ws->src1, 0, 0, num_k, slab, buf, TILED_TILE_ROWS); + } + } + } + } else { + // standard per-slab path: K is stepped in 256-element slabs + for (int64_t ib = 0; ib < ne00; ib += TILED_TILE_K) { + const int kblk = (int) (ib / TILED_TILE_K); + tiled_unpack_src0((const B *) (src0_cur + r * src0->nb[1] + kblk * src0_bs), src0_stride, n_src0, &ws->src0, 1); + tiled_repack_src0(&ws->src0, n_src0, 1, BIAS, ACTBIAS); + tiled_unpack_src1_q8_K(rows, nrows, &ws->src1, kblk, 1); + // 16x16 microtiles sweeping the window; repack each src1 band just-in-time + // (j0-outer) so only the bands actually used are repacked and each is L1-hot + for (int64_t ir1 = 0; ir1 < nrows; ir1 += TILED_MICRO) { + tiled_repack_src1(&ws->src1, (int) ir1, 1, ACTBIAS); + for (int64_t ir0 = r; ir0 < r_end; ir0 += TILED_MICRO) { + tiled_run_microtile(ws->src0, ws->src1, + (int) (ir0 - r), (int) ir1, 1, 0, + ws->acc, TILED_TILE_ROWS); + } + } + } + } + + tiled_store_window_scatter(ws->acc, n_src0, nrows, TILED_TILE_ROWS, col_ptrs); +} + +// one expert of MUL_MAT_ID: each k window's dispatched rows are pointed at by a per-window +// row pointer list and swept over the thread's row windows, the dst rows are scattered back +// per window +template +static void ggml_compute_forward_mul_mat_id_tiled_one_expert( + const struct ggml_compute_params * params, + struct ggml_tensor * dst, + int64_t cur_a, + int64_t cne1, + const int32_t * expert_rows, + char * scratch) { + + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + + GGML_TENSOR_BINARY_OP_LOCALS + + const int ith = params->ith; + const int nth = params->nth; + + const enum ggml_type vec_dot_type = ggml_get_type_traits_cpu(src0->type)->vec_dot_type; + + const size_t nbw1 = (src1->type == vec_dot_type) ? nb11 : ggml_row_size(vec_dot_type, ne10); + const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; + + const char * src0_cur = (const char *) src0->data + cur_a * nb02; + + // this thread's workspace + tiled_ws * ws = (tiled_ws *) (scratch + (size_t) ith * ggml_tiled_ws_size()); + + // groups of TILED_MMID_GROUP rows; rounded up, the window tail is clamped in the gemm + const int64_t ngroups = (ne01 + TILED_MMID_GROUP - 1) / TILED_MMID_GROUP; + + const int64_t g0 = (ngroups * ith) / nth; + const int64_t g1 = (ngroups * (ith + 1)) / nth; + + // no rows for this thread; nothing to do + if (g0 >= g1) { + return; + } + + // MUL_MAT_ID scatters the src1 rows per routed row: point each k window's rows + // at the routed src1 rows, then sweep the row windows + for (int64_t k = 0; k < cne1; k += TILED_TILE_K) { + const int64_t nrows = MIN(TILED_TILE_K, cne1 - k); + + const block_q8_K * rows[TILED_TILE_ROWS]; + for (int64_t i = 0; i < nrows; i++) { + const int64_t i11 = expert_rows[2 * (k + i) + 0] % ne11; + const int64_t i12 = expert_rows[2 * (k + i) + 1]; + rows[i] = (const block_q8_K *) ((const char *) wdata + (i11 + i12 * ne11) * nbw1); + } + + for (int64_t g = g0; g < g1; g++) { + const int64_t r = g * TILED_MMID_GROUP; + + tiled_mmid_gemm_window(dst, src0, src0_cur, r, k, nrows, expert_rows, rows, ws); + } + } +} + +template +static void ggml_compute_forward_mul_mat_tiled_one_chunk( + const struct ggml_compute_params * params, + struct ggml_tensor * dst, + const int64_t ir0_start, + const int64_t ir0_end, + const int64_t ir1_start, + const int64_t ir1_end, + tiled_ws * ws) { + + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + + GGML_TENSOR_BINARY_OP_LOCALS + + const enum ggml_type vec_dot_type = ggml_get_type_traits_cpu(src0->type)->vec_dot_type; + + // broadcast factors + const int64_t r2 = ne12 / ne02; + const int64_t r3 = ne13 / ne03; + + if (ir0_start >= ir0_end || ir1_start >= ir1_end) { + return; + } + + const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; + const size_t row_size = ggml_row_size(vec_dot_type, ne10); + const size_t src0_bs = ggml_type_size(src0->type); + const size_t src1_bs = ggml_type_size(vec_dot_type); + + GGML_ASSERT(ne00 % 256 == 0); + assert(ne12 % ne02 == 0); + assert(ne13 % ne03 == 0); + + const int64_t src0_stride = nb01 / src0_bs; // blocks between src0 rows + const int64_t src1_stride = (src1->type == vec_dot_type ? src1->nb[1] : row_size) / src1_bs; + + const int64_t TILE = 256; + const int64_t MICRO = 16; + + // 256-wide windows over the chunk. The iir1 window is additionally clamped at the + // src1 batch (ne11) boundary: the tiles require a constant batch index (i12/i13) + // within a window, so advance by the clamped end, not a fixed 256. + for (int64_t iir1 = ir1_start; iir1 < ir1_end; ) { + int64_t iir1_end = MIN(iir1 + TILE, ir1_end); + const int64_t bnd = (iir1 / ne11 + 1) * ne11; + if (bnd < iir1_end) { + iir1_end = bnd; + } + + const int n_src1 = (int) (iir1_end - iir1); + + // batch coords, constant within the clamped window + const int64_t i13 = iir1 / (ne12 * ne11); + const int64_t i12 = (iir1 - i13 * ne12 * ne11) / ne11; + // within-batch row; dst_col below holds the i12/i13 batch offset, so the store + // applies i11 * nb1 (not the flattened iir1, which spans all batch dims) + const int64_t i11 = iir1 - i13 * ne12 * ne11 - i12 * ne11; + + // dst batches == src1 batches (ggml_mul_mat), so the loop batch coords are in + // src1 space; src0 batches are broadcast over them, map down into src0's batch + const int64_t i02 = i12 / r2; + const int64_t i03 = i13 / r3; + + const char * src0_row = (const char *) src0->data + i02 * src0->nb[2] + i03 * src0->nb[3]; + char * dst_col = (char *) dst->data + i12 * nb2 + i13 * nb3; + + // rows[r] is the window's row r base block; the k-slab offset is applied in the unpack + const block_q8_K * rows[TILED_TILE_ROWS]; + for (int r = 0; r < n_src1; r++) { + rows[r] = (const block_q8_K *) ((const char *) wdata + (iir1 + r) * src1_stride * src1_bs); + } + + for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += TILE) { + int64_t iir0_end = MIN(iir0 + TILE, ir0_end); + const int n_src0 = (int) (iir0_end - iir0); + + // result buffer zeroed once per macrotile + memset(ws->acc, 0, (size_t)TILED_TILE_ROWS * TILED_TILE_ROWS * sizeof(float)); + + if (n_src1 <= TILED_MICRO) { + // narrow path: if we have few enough rows, we're memory bound + // Do longer, skinner tiles so we have longer contiguous reads and improve bandwidth + const int k_extent = ggml_tiled_narrow_k_extent(n_src1, ne00); + const int num_k = k_extent / TILED_TILE_K; + for (int64_t k0 = 0; k0 < ne00; k0 += k_extent) { + const int kstart = (int) (k0 / TILED_TILE_K); + // activation: one 16-row band (n_src1 <= 16). Unpack the whole k_extent chunk + // (all slabs) so each activation row is a long stream, then repack each slab in place + tiled_unpack_src1_q8_K(rows, n_src1, &ws->src1, kstart, num_k); + // repack the whole k_extent chunk in place, if kernel wants to: num_k slabs x 4 tiles, 16 rows, row stride num_k*256 + tiled_repack_src1(&ws->src1, 0, num_k, ACTBIAS); + // weight groups: unpack the k_extent chunk of each 16-row group (the long per-row + // read), then one standard MAC per slab accumulating into the same acc rows + for (int64_t ir0 = iir0; ir0 < iir0_end; ir0 += MICRO) { + const int n0 = (int) MIN(MICRO, iir0_end - ir0); + const B * wbase = (const B *) (src0_row + ir0 * nb01 + kstart * src0_bs); + tiled_unpack_src0(wbase, src0_stride, n0, &ws->src0, num_k); + tiled_repack_src0(&ws->src0, n0, num_k, BIAS, ACTBIAS); + float * buf = ws->acc + (ir0 - iir0) * TILED_TILE_ROWS; + for (int slab = 0; slab < num_k; slab++) { + tiled_run_microtile(ws->src0, ws->src1, 0, 0, num_k, slab, buf, TILED_TILE_ROWS); + } + } + } + } else { + // standard per-slab path: K is stepped in 256-element chunks + for (int64_t ib = 0; ib < ne00; ib += TILE) { + const int kblk = (int) (ib / TILE); + tiled_unpack_src0((const B *) (src0_row + iir0 * nb01 + kblk * src0_bs), src0_stride, n_src0, &ws->src0, 1); + tiled_repack_src0(&ws->src0, n_src0, 1, BIAS, ACTBIAS); + tiled_unpack_src1_q8_K(rows, n_src1, &ws->src1, kblk, 1); + + // 16x16 microtiles sweeping the window; repack each src1 band before first use + for (int64_t ir1 = iir1; ir1 < iir1_end; ir1 += MICRO) { + tiled_repack_src1(&ws->src1, (int) (ir1 - iir1), 1, ACTBIAS); + for (int64_t ir0 = iir0; ir0 < iir0_end; ir0 += MICRO) { + tiled_run_microtile(ws->src0, ws->src1, + (int) (ir0 - iir0), (int) (ir1 - iir1), 1, 0, + ws->acc, TILED_TILE_ROWS); + } + } + } + } + // write acc back out from L2 to main memory + tiled_store_window(ws->acc, n_src0, n_src1, TILED_TILE_ROWS, + (float *) (dst_col + iir0 * nb0 + i11 * nb1), nb1 / nb0); + } + iir1 = iir1_end; + } +} + +template +static void ggml_compute_forward_mul_mat_tiled_driver( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + + GGML_TENSOR_BINARY_OP_LOCALS + + const int ith = params->ith; + const int nth = params->nth; + + enum ggml_type const vec_dot_type = ggml_get_type_traits_cpu(src0->type)->vec_dot_type; + ggml_from_float_t const from_float = ggml_get_type_traits_cpu(vec_dot_type)->from_float; + + GGML_ASSERT(ne0 == ne01); + GGML_ASSERT(ne1 == ne11); + GGML_ASSERT(ne2 == ne12); + GGML_ASSERT(ne3 == ne13); + + // we don't support permuted src0 or src1 + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + if (src1->type != vec_dot_type) { + char * wdata = (char *) params->wdata; + + const size_t nbw0 = ggml_type_size(vec_dot_type); + const size_t nbw1 = ggml_row_size(vec_dot_type, ne10); + const size_t nbw2 = nbw1*ne11; + const size_t nbw3 = nbw2*ne12; + + assert(params->wsize >= ne13*nbw3); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + + for (int64_t i13 = 0; i13 < ne13; ++i13) { + for (int64_t i12 = 0; i12 < ne12; ++i12) { + for (int64_t i11 = 0; i11 < ne11; ++i11) { + size_t bs = ggml_blck_size(vec_dot_type); + int64_t ne10_block_start = (ith * ne10/bs) / nth; + int64_t ne10_block_end = ((ith + 1) * ne10/bs) / nth; + from_float((float *)((char *) src1->data + i13*src1->nb[3] + i12*src1->nb[2] + i11*src1->nb[1] + ne10_block_start*bs*src1->nb[0]), + (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), + (ne10_block_end - ne10_block_start) * bs); + } + } + } + } + + if (ith == 0) { + // Every thread starts at ith, so the first unprocessed chunk is nth. This saves a bit of coordination right at the start. + ggml_threadpool_chunk_set(params->threadpool, nth); + } + + ggml_barrier(params->threadpool); + + // This is the size of the first dimension of the result, so we can iterate that way. (see the ASSERT above, these are the same numbers) + const int64_t nr0 = ne0; + + // This is the size of the rest of the dimensions of the result + const int64_t nr1 = ne1 * ne2 * ne3; + + // Now select a reasonable chunk size. + int chunk_size = 256; //TILED_TILE_ROWS; + + // distribute the work across the inner or outer loop based on which one is larger + // The number of chunks in the 0/1 dim. CEIL(nr/chunk_size) + int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size; + int64_t nchunk1 = (nr1 + chunk_size - 1) / chunk_size; + + // Step down chunk size if too few chunks to saturate cores, minimum is microtile size + while (nchunk0 * nchunk1 < nth * 4 && chunk_size > 16) { + chunk_size = chunk_size / 2; + nchunk0 = (nr0 + chunk_size - 1) / chunk_size; + nchunk1 = (nr1 + chunk_size - 1) / chunk_size; + } + + // The number of elements in each chunk + const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; + const int64_t dr1 = (nr1 + nchunk1 - 1) / nchunk1; + + // The first chunk comes from our thread_id, the rest will get auto-assigned. + int current_chunk = ith; + + // per-thread workspace: after any converted src1 data in wdata, aligned to 64 + char * ws_base = (char *) params->wdata; + if (src1->type != vec_dot_type) { + ws_base += GGML_PAD(ggml_row_size(vec_dot_type, ggml_nelements(src1)), 64); + } + ws_base = (char *) (((uintptr_t) ws_base + 63) & ~(uintptr_t) 63); + tiled_ws * ws = (tiled_ws *) (ws_base + (size_t) ith * ggml_tiled_ws_size()); + + // TODO: if we KNOW we're on a machine where all cores are equal, we could skip the coordination/work-stealing and just assign chunks deterministically + while (current_chunk < nchunk0 * nchunk1) { + const int64_t ith0 = current_chunk % nchunk0; + const int64_t ith1 = current_chunk / nchunk0; + + const int64_t ir0_start = dr0 * ith0; + const int64_t ir0_end = MIN(ir0_start + dr0, nr0); + + const int64_t ir1_start = dr1 * ith1; + const int64_t ir1_end = MIN(ir1_start + dr1, nr1); + + ggml_compute_forward_mul_mat_tiled_one_chunk(params, dst, ir0_start, ir0_end, ir1_start, ir1_end, ws); + + if (nth >= nchunk0 * nchunk1) { + break; + } + + current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1); + } +} + +// src0 type dispatch, shared by the MUL_MAT and MUL_MAT_ID entries: one expert for +// MUL_MAT_ID (expert_rows != NULL), the full op for MUL_MAT +template +static bool tiled_matmul_dispatch(const struct ggml_compute_params * params, + struct ggml_tensor * dst, + const int32_t * expert_rows, + int64_t cur_a, + int64_t cne1, + char * scratch) { + if (expert_rows == NULL) { + ggml_compute_forward_mul_mat_tiled_driver(params, dst); + } else { + ggml_compute_forward_mul_mat_id_tiled_one_expert(params, dst, cur_a, cne1, expert_rows, scratch); + } + return true; +} + +// the supported src0 types, one list for both ops; false to fall through to vec_dot +static bool ggml_tiled_matmul_type_dispatch(const struct ggml_compute_params * params, + struct ggml_tensor * dst, + const int32_t * expert_rows = NULL, + int64_t cur_a = 0, + int64_t cne1 = 0, + char * scratch = NULL) { + switch (dst->src[0]->type) { + case GGML_TYPE_Q6_K: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_Q5_K: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_Q4_K: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_Q3_K: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_Q2_K: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_IQ4_XS: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_IQ2_XXS: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_IQ2_XS: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_IQ2_S: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_IQ3_XXS: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_IQ3_S: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_IQ1_S: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + case GGML_TYPE_IQ1_M: + return tiled_matmul_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); + default: + return false; + } +} + +static bool ggml_tiled_min_batch(int64_t rows) { + // Profitable at rows >= 8, take even when unprofitable if we're forced + return rows >= 8 || ggml_tiled_matmul_forced(); +} + +// tiled K-quant matmul; returns true if the op was computed here, +// false to fall through to the stock path +bool ggml_compute_forward_mul_mat_tiled( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + // --use-ref means bail out and go back to vec_dot reference impl + if (params->use_ref) { + return false; + } + if (!ggml_tiled_supported(dst->src[0], dst->src[1])) { + return false; + } + if (!ggml_tiled_min_batch(dst->src[1]->ne[1])) { + return false; + } + return ggml_tiled_matmul_type_dispatch(params, dst); +} + +// MUL_MAT_ID (MoE), one expert; returns true if the expert was computed here, +// per expert eligibility (type gate, batch floor) is decided here +bool ggml_compute_forward_mul_mat_id_tiled( + const struct ggml_compute_params * params, + struct ggml_tensor * dst, + int64_t cur_a, + int64_t cne1, + const int32_t * expert_rows, + char * scratch) { + if (params->use_ref) { + return false; + } + if (!ggml_tiled_supported(dst->src[0], dst->src[1])) { + return false; + } + // profitability is per expert: the rows routed to this expert + if (!ggml_tiled_min_batch(cne1)) { + return false; + } + return ggml_tiled_matmul_type_dispatch(params, dst, expert_rows, cur_a, cne1, scratch); +} diff --git a/ggml/src/ggml-cpu/tiled/tiled.h b/ggml/src/ggml-cpu/tiled/tiled.h new file mode 100644 index 0000000000..6f3e8dc13e --- /dev/null +++ b/ggml/src/ggml-cpu/tiled/tiled.h @@ -0,0 +1,31 @@ +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Amount of wdata to reserve for tiled workspaces +size_t ggml_tiled_wdata_size(int n_tasks, struct ggml_tensor * dst); + +// tiled K-quant matmul; returns true if the op was computed here, +// false to fall through to the stock path +bool ggml_compute_forward_mul_mat_tiled(const struct ggml_compute_params * params, + struct ggml_tensor * dst); + +// MUL_MAT_ID (MoE) path, one expert; returns true if the expert was computed here, per expert +// eligibility (type gate, batch floor) is decided inside. expert_rows points at the expert's +// row of the matrix_rows table of (expert slot, batch row) int32 pairs; scratch is the +// per-thread tiled_ws region reserved in wdata, n_tasks of ggml_tiled_ws_size() bytes +bool ggml_compute_forward_mul_mat_id_tiled(const struct ggml_compute_params * params, + struct ggml_tensor * dst, + int64_t cur_a, + int64_t cne1, + const int32_t * expert_rows, + char * scratch); + +#ifdef __cplusplus +} +#endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 9b3a4fcc4b..996e7f58f3 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -309,6 +309,7 @@ endif() # TODO: make this test (and others) not link `libllama` as it is not needed [TAG_TESTS_LLAMA_LINK] llama_build(test-backend-ops.cpp) +llama_build_and_test(test-tiled-mulmat.cpp) llama_build_and_test(test-model-load-cancel.cpp LABEL "model") llama_build_and_test(test-autorelease.cpp LABEL "model") llama_build_and_test(test-backend-sampler.cpp LABEL "model") diff --git a/tests/test-tiled-mulmat.cpp b/tests/test-tiled-mulmat.cpp new file mode 100644 index 0000000000..30cfdd1979 --- /dev/null +++ b/tests/test-tiled-mulmat.cpp @@ -0,0 +1,852 @@ +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml.h" + +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#endif + +// failed checks; main exits non-zero on any failures +static int n_failed = 0; + +static float * gen_rand_f32(int64_t n) { + float * data = (float *) malloc(n * sizeof(float)); + for (int64_t i = 0; i < n; ++i) { + data[i] = (float)rand() / (float)RAND_MAX - 0.5f; + data[i] *= 5.0f; + } + return data; +} + +// updates pointers to max_err and rms_err +static void compare_f32(const float * ref, const float * out, int64_t n, float * max_err, float * rms_err) { + *max_err = 0.0f; + double sum_sq_err = 0.0; + for (int64_t i = 0; i < n; ++i) { + float err = fabsf(ref[i] - out[i]); + if (err > *max_err) { + *max_err = err; + } + sum_sq_err += (double)err * err; + } + *rms_err = sqrt(sum_sq_err / n); +} + +// Fill from a flat row-major f32 source: rows * cols floats. +static void fill_tensor(struct ggml_tensor * t, const float * src, int64_t rows, int64_t cols, ggml_type qtype) { + GGML_ASSERT(t->ne[0] == cols); + GGML_ASSERT(rows * cols == ggml_nelements(t)); + if (qtype == GGML_TYPE_F32) { + ggml_backend_tensor_set(t, src, 0, ggml_nbytes(t)); + return; + } + void * q = malloc(ggml_nbytes(t)); + // some iq types (iq2_xxs, iq2_xs, iq1_s) steer their code choice with an imatrix; the + // quantized bytes are shared by the reference and tiled paths, so a dummy imatrix is enough + const float * imatrix = NULL; + float * im = NULL; + if (ggml_quantize_requires_imatrix(qtype)) { + im = (float *) malloc(rows * cols * sizeof(float)); + for (int64_t i = 0; i < rows * cols; i++) im[i] = 1.0f; + imatrix = im; + } + ggml_quantize_chunk(qtype, src, q, 0, rows, cols, imatrix); + ggml_backend_tensor_set(t, q, 0, ggml_nbytes(t)); + free(q); + free(im); +} + +// The CPU-specific control API is resolved through the backend registry: with +// GGML_BACKEND_DL the CPU backend is a runtime-loaded module and its symbols +// are not available at link time +static void cpu_set_use_ref(ggml_backend_t backend, bool use_ref) { + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + void (* set_use_ref)(ggml_backend_t, bool) = + (void (*)(ggml_backend_t, bool)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cpu_set_use_ref"); + if (!set_use_ref) { + fprintf(stderr, "ggml_backend_cpu_set_use_ref not available\n"); + exit(1); + } + set_use_ref(backend, use_ref); +} + +static void cpu_set_n_threads(ggml_backend_t backend, int n_threads) { + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + void (* set_n_threads)(ggml_backend_t, int) = + (void (*)(ggml_backend_t, int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads"); + if (!set_n_threads) { + fprintf(stderr, "ggml_backend_set_n_threads not available\n"); + exit(1); + } + set_n_threads(backend, n_threads); +} + +static void test_matmul(ggml_backend_t backend, int64_t M, int64_t N, int64_t K, ggml_type quant_type) { + srand(0xBEEF); + + float * src1_ref = gen_rand_f32(M * N); + float * src0_ref = gen_rand_f32(N * K); + float * dst_out = (float *) malloc(M * K * sizeof(float)); + float * dst_tiled = (float *) malloc(M * K * sizeof(float)); + + struct ggml_init_params ip = { 1024*1024*1024, nullptr, true }; + struct ggml_context * ctx = ggml_init(ip); + + // ggml_mul_mat(t0, t1) computes t1 * t0^T: t0 (src0, quant) is N wide x + // K high, t1 (src1, f32) is N wide x M high, dst = (K, M). + struct ggml_tensor * src1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, N, M); + struct ggml_tensor * src0 = ggml_new_tensor_2d(ctx, quant_type, N, K); + + struct ggml_cgraph * gf = ggml_new_graph(ctx); + struct ggml_tensor * dst = ggml_mul_mat(ctx, src0, src1); + ggml_build_forward_expand(gf, dst); + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + fill_tensor(src1, src1_ref, M, N, GGML_TYPE_F32); + fill_tensor(src0, src0_ref, K, N, quant_type); + + cpu_set_use_ref(backend, true); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, dst_out, 0, ggml_nbytes(dst)); + cpu_set_use_ref(backend, false); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, dst_tiled, 0, ggml_nbytes(dst)); + + // std vs tiled: identical quantized inputs, so any large difference here is a bug in the tiled kernel + float max_err, rms_err; + compare_f32(dst_out, dst_tiled, M*K, &max_err, &rms_err); + float tol = 1e-3f; + + printf("TEST %lldx%lld * %lldx%lld (%s): %s (max_err: %f, rms: %f, tolerance: %f)\n", + (long long)M, (long long)N, (long long)N, (long long)K, + ggml_type_name(quant_type), (max_err <= tol) ? "PASS" : "FAIL", max_err, rms_err, tol); + + // if the tiled kernel deviates from std beyond quantization tolerance, + // dump a few offenders + if (max_err > tol) { + int64_t shown = 0; + for (int64_t i = 0; i < M*K && shown < 8; ++i) { + float err = fabsf(dst_out[i] - dst_tiled[i]); + if (err > tol) { + printf(" tiled vs std: i=%lld (m=%lld k=%lld) std=%f tiled=%f err=%f\n", + (long long)i, (long long)(i/K), (long long)(i%K), dst_out[i], dst_tiled[i], err); + ++shown; + } + } + ++n_failed; + } + + ggml_backend_buffer_free(buf); + ggml_free(ctx); + free(src1_ref); free(src0_ref); free(dst_out); free(dst_tiled); +} + +// Higher-dim (ne[2], ne[3] > 1) check, std (ggml_mul_mat) is trusted as the reference +// +// src1 (F32) : [N, M, src1_2, src1_3] ne0=N (reduction), ne1=M (out0) +// src0 (quant): [N, K, src0_2, src0_3] ne0=N (reduction), ne1=K (out1) +// dst = ggml_mul_mat(src0, src1) : [K, M, src1_2, src1_3] +static void test_matmul_highdim(ggml_backend_t backend, int64_t M, int64_t N, int64_t K, + int64_t src1_2, int64_t src1_3, + int64_t src0_2, int64_t src0_3, + ggml_type quant_type) { + srand(0xBEEF); + + const int64_t n_src1 = N*M*src1_2*src1_3; + const int64_t n_src0 = N*K*src0_2*src0_3; + const int64_t n_dst = K*M*src1_2*src1_3; + + float * src1_ref = gen_rand_f32(n_src1); + float * src0_ref = gen_rand_f32(n_src0); + float * dst_std = (float *) malloc(n_dst * sizeof(float)); + float * dst_tiled = (float *) malloc(n_dst * sizeof(float)); + + struct ggml_init_params ip = { 1024*1024*1024, nullptr, true }; + struct ggml_context * ctx = ggml_init(ip); + + int64_t ne_src1[4] = { N, M, src1_2, src1_3 }; + int64_t ne_src0[4] = { N, K, src0_2, src0_3 }; + struct ggml_tensor * src1 = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_src1); + struct ggml_tensor * src0 = ggml_new_tensor(ctx, quant_type, 4, ne_src0); + + struct ggml_cgraph * gf = ggml_new_graph(ctx); + struct ggml_tensor * dst = ggml_mul_mat(ctx, src0, src1); + ggml_build_forward_expand(gf, dst); + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + // the batch dims fold into the row count (rows are contiguous, ne0 fastest) + fill_tensor(src1, src1_ref, M * src1_2 * src1_3, N, GGML_TYPE_F32); + fill_tensor(src0, src0_ref, K * src0_2 * src0_3, N, quant_type); + + // reference = stock path (use_ref keeps the op off the tiled hook), tiled + // = the gated path; same weights, same op, run twice into separate buffers + cpu_set_use_ref(backend, true); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, dst_std, 0, ggml_nbytes(dst)); + cpu_set_use_ref(backend, false); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, dst_tiled, 0, ggml_nbytes(dst)); + + // max |dst|: scale for the quantization tolerance + float scale = 0.0f; + for (int64_t i = 0; i < n_dst; ++i) { + scale = fmaxf(scale, fabsf(dst_std[i])); + } + + // std vs tiled: identical quantized inputs, so a large difference is a bug + float max_err, rms_err; + compare_f32(dst_std, dst_tiled, n_dst, &max_err, &rms_err); + float tol = (quant_type == GGML_TYPE_F32) ? 1e-4f : fmaxf(1e-3f, 1e-3f*scale); + + // which path did the tiled op take (mirrors the hard-constraint gate; the + // harness runs with force on, so the profitability check is bypassed) + bool tiled_kernel = (N % 256 == 0) && (src0_2 > 0) && (src0_3 > 0) + && (src1_2 % src0_2 == 0) && (src1_3 % src0_3 == 0); + + printf("TEST %lldx%lldx%lldx%lld * %lldx%lldx%lldx%lld (%s, %s): %s (max_err: %f, rms: %f, scale: %f)\n", + (long long)M, (long long)N, (long long)src1_2, (long long)src1_3, + (long long)K, (long long)N, (long long)src0_2, (long long)src0_3, + ggml_type_name(quant_type), tiled_kernel ? "tiled-kernel" : "stock", + (max_err <= tol) ? "PASS" : "FAIL", max_err, rms_err, scale); + + if (max_err > tol) { + int64_t shown = 0; + for (int64_t i = 0; i < n_dst && shown < 8; ++i) { + float err = fabsf(dst_std[i] - dst_tiled[i]); + if (err > tol) { + printf(" tiled vs std: i=%lld std=%f tiled=%f err=%f\n", + (long long)i, dst_std[i], dst_tiled[i], err); + ++shown; + } + } + ++n_failed; + } + + ggml_backend_buffer_free(buf); + ggml_free(ctx); + free(src1_ref); free(src0_ref); free(dst_std); free(dst_tiled); +} + +// MUL_MAT_ID (MoE) check, std (vec_dot) is trusted as the reference +// +// src0 (as, quant): [K, R, n_experts] ne0 = K (reduction), ne1 = R rows per expert, ne2 = experts +// src1 (b, F32) : [K, b_slots, batch] ne0 = K (reduction), ne1 = b_slots (b rows, broadcast over k columns) +// ids : [k, batch] expert picked per (slot, batch row); k % b_slots == 0 +// dst = ggml_mul_mat_id(as, b, ids) : [R, k, batch]; column (id, t) = GEMV of src1 row (id % b_slots + t*b_slots) +// against expert ids[t*k+id]. b_slots = 1 is the common MoE case (all top-k experts see one input row). +static void test_mul_mat_id(ggml_backend_t backend, int64_t K, int64_t R, int64_t n_experts, + int64_t k, int64_t b_slots, int64_t batch, ggml_type quant_type, + bool src1_strided = false) { + srand(0xBEEF); + + const int64_t n_as = K * R * n_experts; + const int64_t n_b = K * b_slots * batch; + const int64_t n_dst = R * k * batch; + + float * as_ref = gen_rand_f32(n_as); + float * b_ref = gen_rand_f32(n_b); + int32_t * ids = (int32_t *) malloc(k * batch * sizeof(int32_t)); + float * dst_ref = (float *) malloc(n_dst * sizeof(float)); + float * dst_tiled = (float *) malloc(n_dst * sizeof(float)); + + struct ggml_init_params ip = { 1024*1024*1024, nullptr, true }; + struct ggml_context * ctx = ggml_init(ip); + + int64_t ne_as[4] = { K, R, n_experts, 1 }; + int64_t ne_b[4] = { K, b_slots, batch, 1 }; + int64_t ne_ids[4] = { k, batch, 1, 1 }; + struct ggml_tensor * src0 = ggml_new_tensor(ctx, quant_type, 4, ne_as); + struct ggml_tensor * ids_t = ggml_new_tensor(ctx, GGML_TYPE_I32, 4, ne_ids); + struct ggml_tensor * src1; + struct ggml_tensor * src1_base = NULL; + if (src1_strided) { + // strided view: nb[1] = 4 bytes, the tiled path reads src1 through strides into wdata + src1_base = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, b_slots * batch); + src1 = ggml_view_3d(ctx, src1_base, K, b_slots, batch, 4, (size_t) b_slots * K * 4, 0); + } else { + src1 = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_b); + } + + struct ggml_cgraph * gf = ggml_new_graph(ctx); + struct ggml_tensor * dst = ggml_mul_mat_id(ctx, src0, src1, ids_t); + ggml_build_forward_expand(gf, dst); + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + // deterministic balanced routing: each expert gets ~ k*batch/n_experts rows + for (int64_t t = 0; t < batch; t++) { + for (int64_t e = 0; e < k; e++) { + ids[t * k + e] = (int32_t) ((t * k + e) % n_experts); + } + } + + fill_tensor(src1, b_ref, b_slots * batch, K, GGML_TYPE_F32); + fill_tensor(src0, as_ref, R * n_experts, K, quant_type); + ggml_backend_tensor_set(ids_t, ids, 0, ggml_nbytes(ids_t)); + + cpu_set_use_ref(backend, true); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, dst_ref, 0, ggml_nbytes(dst)); + cpu_set_use_ref(backend, false); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, dst_tiled, 0, ggml_nbytes(dst)); + + // same quantized inputs, so a large difference is a bug in the tiled path + float max_err, rms_err; + compare_f32(dst_ref, dst_tiled, n_dst, &max_err, &rms_err); + float tol = 1e-3f; + + printf("TEST mul_mat_id %lldx%lldx%lld * %lldx%lldx%lld ids[%lld,%lld] (%s%s): %s (max_err: %f, rms: %f, tolerance: %f)\n", + (long long)K, (long long)R, (long long)n_experts, + (long long)K, (long long)b_slots, (long long)batch, + (long long)k, (long long)batch, + ggml_type_name(quant_type), src1_strided ? ", strided-b" : "", + (max_err <= tol) ? "PASS" : "FAIL", max_err, rms_err, tol); + + if (max_err > tol) { + int64_t shown = 0; + for (int64_t i = 0; i < n_dst && shown < 8; ++i) { + float err = fabsf(dst_ref[i] - dst_tiled[i]); + if (err > tol) { + printf(" tiled vs std: i=%lld (row=%lld slot=%lld batch=%lld) std=%f tiled=%f err=%f\n", + (long long)i, (long long)(i % R), (long long)((i / R) % k), (long long)(i / (R * k)), + dst_ref[i], dst_tiled[i], err); + ++shown; + } + } + ++n_failed; + } + + ggml_backend_buffer_free(buf); + ggml_free(ctx); + free(as_ref); free(b_ref); free(ids); free(dst_ref); free(dst_tiled); +} + +static double time_graph_compute(ggml_backend_t backend, struct ggml_cgraph * gf) { +#if defined(_WIN32) + // high-res timer; clock_gettime is not portable across Windows toolchains + LARGE_INTEGER freq, t0, t1; + QueryPerformanceFrequency(&freq); + QueryPerformanceCounter(&t0); + ggml_backend_graph_compute(backend, gf); + QueryPerformanceCounter(&t1); + return (double) (t1.QuadPart - t0.QuadPart) / (double) freq.QuadPart; +#else + struct timespec t0, t1; + clock_gettime(CLOCK_MONOTONIC, &t0); + ggml_backend_graph_compute(backend, gf); + clock_gettime(CLOCK_MONOTONIC, &t1); + return (t1.tv_sec - t0.tv_sec) + (t1.tv_nsec - t0.tv_nsec) / 1e9; +#endif +} + +// Warm up, then time n runs and return the best +static double time_graph_compute_best(ggml_backend_t backend, struct ggml_cgraph * gf, int n, + void * flush_buf = NULL, size_t flush_size = 0) { + + time_graph_compute(backend, gf); // warmup + double best = 1e30; + for (int i = 0; i < n; ++i) { + // Evict L3 cache first + if (flush_buf) { memset(flush_buf, 0xAB, flush_size); } + const double t = time_graph_compute(backend, gf); + if (t < best) best = t; + } + return best; +} + +// Fetch the repack extra buffer type through the public proc-address API +static ggml_backend_buffer_type_t get_cpu_repack_buft(void) { + ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + if (!cpu_dev) { + return NULL; + } + ggml_backend_reg_t cpu_reg = ggml_backend_dev_backend_reg(cpu_dev); + ggml_backend_dev_get_extra_bufts_t get_extra = + (ggml_backend_dev_get_extra_bufts_t) ggml_backend_reg_get_proc_address(cpu_reg, "ggml_backend_dev_get_extra_bufts"); + if (!get_extra) { + return NULL; + } + ggml_backend_buffer_type_t * bufts = get_extra(cpu_dev); + // the only extra buffer type the CPU backend exposes is CPU_REPACK + return (bufts && *bufts) ? *bufts : NULL; +} + +struct bench_row { + const char * name; + double time_std, time_repack, time_tiled; + float max_err_repack, rmse_repack; + float max_err_tiled, rmse_tiled; + bool have_repack; +}; + +// MUL_MAT three-way bench: std (default optimized GEMM) vs repack (CPU_REPACK buffer) vs +// tiled kernel. std is timed in-process with use_ref=true (bypasses the tiled gate); tiled +// and repack with use_ref=false (repack's tiled gate declines on the buffer's extra). repack +// is n/a where no repack kernel exists for the type. Errors are vs the vec-only reference (use_ref). +static bench_row bench_three_way(ggml_backend_t backend, int64_t M, int64_t N, int64_t K, ggml_type quant_type) { + bench_row row; + row.name = ggml_type_name(quant_type); + row.time_std = row.time_repack = row.time_tiled = 0.0; + row.max_err_repack = row.rmse_repack = row.max_err_tiled = row.rmse_tiled = 0.0f; + row.have_repack = false; + + struct ggml_init_params ip = { 1024*1024*1024, nullptr, true }; + struct ggml_context * ctx = ggml_init(ip); + + struct ggml_tensor * src1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, N, M); + struct ggml_tensor * src0_std = ggml_new_tensor_2d(ctx, quant_type, N, K); + struct ggml_tensor * src0_rep = ggml_new_tensor_2d(ctx, quant_type, N, K); + + struct ggml_cgraph * gf = ggml_new_graph(ctx); + struct ggml_tensor * dst = ggml_mul_mat(ctx, src0_std, src1); + ggml_build_forward_expand(gf, dst); + + struct ggml_cgraph * gf_repack = ggml_new_graph(ctx); + struct ggml_tensor * dst_repack = ggml_mul_mat(ctx, src0_rep, src1); + ggml_build_forward_expand(gf_repack, dst_repack); + + ggml_backend_buffer_type_t repack_buft = get_cpu_repack_buft(); + ggml_backend_buffer_t buf_rep = NULL; + if (repack_buft && N % 8 == 0 && K % 8 == 0) { + buf_rep = ggml_backend_buft_alloc_buffer(repack_buft, ggml_nbytes(src0_rep)); + src0_rep->buffer = buf_rep; + src0_rep->data = ggml_backend_buffer_get_base(buf_rep); + ggml_backend_buffer_init_tensor(buf_rep, src0_rep); + row.have_repack = (src0_rep->extra != NULL); + } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + srand(0xBEEF); + float * src1_data = gen_rand_f32(M * N); + float * src0_data = gen_rand_f32(N * K); + fill_tensor(src1, src1_data, M, N, GGML_TYPE_F32); + fill_tensor(src0_std, src0_data, K, N, quant_type); + if (row.have_repack) { + fill_tensor(src0_rep, src0_data, K, N, quant_type); + } + free(src1_data); free(src0_data); + + // one warmup + N timings; best (min) wins. The flush buffer is larger than the + // whole-chip cache, so its memset evicts L3 and each timed iteration streams DRAM. + const size_t flush_size = 256 * 1024 * 1024; + void * flush_buf = malloc(flush_size); + const int n_reps = 5; + + // std in-process (use_ref bypasses the tiled gate; the vec_dot is the same optimized one) + cpu_set_use_ref(backend, true); + row.time_std = time_graph_compute_best(backend, gf, n_reps, flush_buf, flush_size); + + // tiled in-process (the parent holds TILED_MM=1 forced on) + cpu_set_use_ref(backend, false); + row.time_tiled = time_graph_compute_best(backend, gf, n_reps, flush_buf, flush_size); + + // repack in-process (the tiled gate declines on the buffer's extra) + if (row.have_repack) { + row.time_repack = time_graph_compute_best(backend, gf_repack, n_reps, flush_buf, flush_size); + } + free(flush_buf); + + // errors vs the vec-only reference (use_ref) + float * out_ref = (float *) malloc(M * K * sizeof(float)); + float * out_tiled = (float *) malloc(M * K * sizeof(float)); + float * out_repack = (float *) malloc(M * K * sizeof(float)); + cpu_set_use_ref(backend, true); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, out_ref, 0, ggml_nbytes(dst)); + cpu_set_use_ref(backend, false); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, out_tiled, 0, ggml_nbytes(dst)); + compare_f32(out_ref, out_tiled, M * K, &row.max_err_tiled, &row.rmse_tiled); + if (row.have_repack) { + ggml_backend_graph_compute(backend, gf_repack); + ggml_backend_tensor_get(dst_repack, out_repack, 0, ggml_nbytes(dst_repack)); + compare_f32(out_ref, out_repack, M * K, &row.max_err_repack, &row.rmse_repack); + } + free(out_ref); free(out_tiled); free(out_repack); + + ggml_backend_buffer_free(buf_rep); + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return row; +} + +static void print_bench_table(int64_t M, int64_t N, int64_t K, const bench_row * rows, size_t n_types) { + const double flops = 2.0 * M * N * K; + + printf("\nBENCH dense %lldx%lld * %lldx%lld (std vs repack vs tiled), min of 5 timings, 8 threads\n", + (long long)M, (long long)N, (long long)N, (long long)K); + printf("%-8s %10s %12s %12s %11s %11s %17s %17s %17s %17s\n", + "type", "std TF", "repack TF", "tiled TF", "repack/std", "tiled/std", + "max_err(repack)", "rmse(repack)", "max_err(tiled)", "rmse(tiled)"); + for (size_t i = 0; i < n_types; ++i) { + const bench_row * r = &rows[i]; + if (r->have_repack) { + printf("%-8s %10.3f %12.3f %12.3f %11.2f %11.2f %17.5e %17.5e %17.5e %17.5e\n", + r->name, + flops / (r->time_std * 1e12), + flops / (r->time_repack * 1e12), flops / (r->time_tiled * 1e12), + r->time_std / r->time_repack, r->time_std / r->time_tiled, + r->max_err_repack, r->rmse_repack, r->max_err_tiled, r->rmse_tiled); + } else { + printf("%-8s %10.3f %12s %12.3f %11s %11.2f %17s %17s %17.5e %17.5e\n", + r->name, + flops / (r->time_std * 1e12), + "n/a", flops / (r->time_tiled * 1e12), "n/a", + r->time_std / r->time_tiled, + "n/a", "n/a", r->max_err_tiled, r->rmse_tiled); + } + } +} + +struct bench_row_mmid { + const char * name; + double time_std, time_repack, time_tiled; + float max_err_repack, rmse_repack; + float max_err_tiled, rmse_tiled; + bool have_repack; +}; + +// MUL_MAT_ID (MoE) three-way bench: std (default GEMM, use_ref) vs repack (CPU_REPACK buffer) +// vs tiled kernel (default; TILED_MM_FORCE is set in main so ragged experts take it too). +// Column (id, t) = GEMV of src1 row (id % b_slots + t * b_slots) against expert ids[t*k+id]; +// cne1 = k * batch / n_experts rows per expert on average. repack is n/a where no repack +// kernel exists. Errors are vs the vec-only reference (use_ref). +static bench_row_mmid bench_mul_mat_id(ggml_backend_t backend, int64_t K, int64_t R, int64_t n_experts, + int64_t k, int64_t b_slots, int64_t batch, ggml_type quant_type) { + bench_row_mmid row; + row.name = ggml_type_name(quant_type); + row.time_std = row.time_repack = row.time_tiled = 0.0; + row.max_err_repack = row.rmse_repack = row.max_err_tiled = row.rmse_tiled = 0.0f; + row.have_repack = false; + + struct ggml_init_params ip = { 1024*1024*1024, nullptr, true }; + struct ggml_context * ctx = ggml_init(ip); + + int64_t ne_as[4] = { K, R, n_experts, 1 }; + int64_t ne_b[4] = { K, b_slots, batch, 1 }; + int64_t ne_ids[4] = { k, batch, 1, 1 }; + struct ggml_tensor * src1 = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_b); + struct ggml_tensor * src0_std = ggml_new_tensor(ctx, quant_type, 4, ne_as); + struct ggml_tensor * src0_rep = ggml_new_tensor(ctx, quant_type, 4, ne_as); + struct ggml_tensor * ids_t = ggml_new_tensor(ctx, GGML_TYPE_I32, 4, ne_ids); + + struct ggml_cgraph * gf = ggml_new_graph(ctx); + struct ggml_tensor * dst = ggml_mul_mat_id(ctx, src0_std, src1, ids_t); + ggml_build_forward_expand(gf, dst); + + struct ggml_cgraph * gf_repack = ggml_new_graph(ctx); + struct ggml_tensor * dst_repack = ggml_mul_mat_id(ctx, src0_rep, src1, ids_t); + ggml_build_forward_expand(gf_repack, dst_repack); + + ggml_backend_buffer_type_t repack_buft = get_cpu_repack_buft(); + ggml_backend_buffer_t buf_rep = NULL; + if (repack_buft && R % 8 == 0) { + buf_rep = ggml_backend_buft_alloc_buffer(repack_buft, ggml_nbytes(src0_rep)); + src0_rep->buffer = buf_rep; + src0_rep->data = ggml_backend_buffer_get_base(buf_rep); + ggml_backend_buffer_init_tensor(buf_rep, src0_rep); + row.have_repack = (src0_rep->extra != NULL); + } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + // deterministic balanced routing: each expert gets ~ k*batch/n_experts rows + int32_t * ids = (int32_t *) malloc(k * batch * sizeof(int32_t)); + for (int64_t t = 0; t < batch; t++) { + for (int64_t e = 0; e < k; e++) { + ids[t * k + e] = (int32_t) ((t * k + e) % n_experts); + } + } + srand(0xBEEF); + float * b_data = gen_rand_f32(K * b_slots * batch); + float * as_data = gen_rand_f32(K * R * n_experts); + fill_tensor(src1, b_data, b_slots * batch, K, GGML_TYPE_F32); + fill_tensor(src0_std, as_data, R * n_experts, K, quant_type); + if (row.have_repack) { + fill_tensor(src0_rep, as_data, R * n_experts, K, quant_type); + } + ggml_backend_tensor_set(ids_t, ids, 0, ggml_nbytes(ids_t)); + free(b_data); free(as_data); + + const size_t flush_size = 256 * 1024 * 1024; + void * flush_buf = malloc(flush_size); + const int n_reps = 5; + + // std in-process (use_ref bypasses the tiled gate; the vec_dot is the same optimized one) + cpu_set_use_ref(backend, true); + row.time_std = time_graph_compute_best(backend, gf, n_reps, flush_buf, flush_size); + + // tiled in-process (the parent holds TILED_MM=1 forced on) + cpu_set_use_ref(backend, false); + row.time_tiled = time_graph_compute_best(backend, gf, n_reps, flush_buf, flush_size); + + // repack in-process (the tiled gate declines on the buffer's extra) + if (row.have_repack) { + row.time_repack = time_graph_compute_best(backend, gf_repack, n_reps, flush_buf, flush_size); + } + free(flush_buf); + + // errors vs the vec-only reference (use_ref) + float * out_ref = (float *) malloc(R * k * batch * sizeof(float)); + float * out_tiled = (float *) malloc(R * k * batch * sizeof(float)); + float * out_repack = (float *) malloc(R * k * batch * sizeof(float)); + cpu_set_use_ref(backend, true); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, out_ref, 0, ggml_nbytes(dst)); + cpu_set_use_ref(backend, false); + ggml_backend_graph_compute(backend, gf); + ggml_backend_tensor_get(dst, out_tiled, 0, ggml_nbytes(dst)); + compare_f32(out_ref, out_tiled, R * k * batch, &row.max_err_tiled, &row.rmse_tiled); + if (row.have_repack) { + ggml_backend_graph_compute(backend, gf_repack); + ggml_backend_tensor_get(dst_repack, out_repack, 0, ggml_nbytes(dst_repack)); + compare_f32(out_ref, out_repack, R * k * batch, &row.max_err_repack, &row.rmse_repack); + } + free(out_ref); free(out_tiled); free(out_repack); free(ids); + + ggml_backend_buffer_free(buf_rep); + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return row; +} + +static void print_mmid_table(int64_t K, int64_t R, int64_t n_experts, int64_t k, + int64_t b_slots, int64_t batch, const bench_row_mmid * rows, size_t n_types) { + const double flops = 2.0 * (double) K * R * k * batch; + + printf("\nBENCH mmid %lldx%lldx%lld k=%lld b_slots=%lld batch=%lld (cne1=%lld, std vs repack vs tiled), min of 5 timings, 8 threads\n", + (long long)K, (long long)R, (long long)n_experts, (long long)k, + (long long)b_slots, (long long)batch, (long long)(k * batch / n_experts)); + printf("%-8s %10s %12s %12s %11s %11s %17s %17s %17s %17s\n", + "type", "std TF", "repack TF", "tiled TF", "repack/std", "tiled/std", + "max_err(repack)", "rmse(repack)", "max_err(tiled)", "rmse(tiled)"); + for (size_t i = 0; i < n_types; ++i) { + const bench_row_mmid * r = &rows[i]; + if (r->have_repack) { + printf("%-8s %10.3f %12.3f %12.3f %11.2f %11.2f %17.5e %17.5e %17.5e %17.5e\n", + r->name, + flops / (r->time_std * 1e12), + flops / (r->time_repack * 1e12), flops / (r->time_tiled * 1e12), + r->time_std / r->time_repack, r->time_std / r->time_tiled, + r->max_err_repack, r->rmse_repack, r->max_err_tiled, r->rmse_tiled); + } else { + printf("%-8s %10.3f %12s %12.3f %11s %11.2f %17s %17s %17.5e %17.5e\n", + r->name, + flops / (r->time_std * 1e12), + "n/a", flops / (r->time_tiled * 1e12), "n/a", + r->time_std / r->time_tiled, + "n/a", "n/a", r->max_err_tiled, r->rmse_tiled); + } + } +} + +int main(int argc, char ** argv) { + bool run_bench = false; + bool run_fuzz = false; + for (int i = 1; i < argc; ++i) { + if (strcmp(argv[i], "--bench") == 0) { + run_bench = true; + } else if (strcmp(argv[i], "--fuzz") == 0) { + run_fuzz = true; + } else { + fprintf(stderr, "error: unknown argument: %s\n", argv[i]); + return 1; + } + } + + // Enable tiled MM, also force tiled MM even when unprofitable for benchmarks +#if defined(_WIN32) + _putenv("GGML_CPU_TILED_MM=1"); + _putenv("GGML_CPU_TILED_MM_FORCE=1"); +#else + setenv("GGML_CPU_TILED_MM", "1", 1); + setenv("GGML_CPU_TILED_MM_FORCE", "1", 1); +#endif + + ggml_backend_load_all(); + ggml_backend_t backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, NULL); + if (!backend) { + fprintf(stderr, "failed to initialize the CPU backend\n"); + return 1; + } + cpu_set_n_threads(backend, 8); + // base case, smoke each quant type + test_matmul(backend, 256, 512, 256, GGML_TYPE_Q6_K); + test_matmul(backend, 256, 512, 256, GGML_TYPE_Q5_K); + test_matmul(backend, 256, 512, 256, GGML_TYPE_Q4_K); + test_matmul(backend, 256, 512, 256, GGML_TYPE_Q3_K); + test_matmul(backend, 256, 512, 256, GGML_TYPE_Q2_K); + test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ4_XS); + test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ2_XXS); + test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ2_XS); + test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ2_S); + test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ3_XXS); + test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ3_S); + test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ1_S); + test_matmul(backend, 256, 512, 256, GGML_TYPE_IQ1_M); + test_mul_mat_id(backend, 256, 256, 2, 8, 1, 64, GGML_TYPE_Q4_K); // single block, cne1 = 256 + test_mul_mat_id(backend, 512, 300, 4, 2, 1, 150, GGML_TYPE_IQ2_XXS); // iq type, imatrix quant + + if (run_fuzz) { + printf("Running fuzz tests.\n"); + // ragged edges, both subblock lengths (Q5_K = 32, Q6_K = 16) + test_matmul(backend, 256, 1024, 8192, GGML_TYPE_Q6_K); // long K, int32 accumulation + test_matmul(backend, 357, 1024, 137, GGML_TYPE_Q6_K); // ragged M and K + test_matmul(backend, 16, 1024, 16, GGML_TYPE_Q5_K); + test_matmul(backend, 18, 1024, 7, GGML_TYPE_Q5_K); // ragged M tail, ragged K in first subblock + test_matmul(backend, 8, 256, 8, GGML_TYPE_Q5_K); + test_matmul(backend, 18, 1024, 256, GGML_TYPE_Q5_K); + test_matmul(backend, 9, 256, 256, GGML_TYPE_Q5_K); + test_matmul(backend, 8, 256, 8, GGML_TYPE_Q6_K); // tiny M, single QK_K block + test_matmul(backend, 256, 1024, 8192, GGML_TYPE_Q3_K); + test_matmul(backend, 357, 1024, 137, GGML_TYPE_Q3_K); + test_matmul(backend, 8, 256, 8, GGML_TYPE_Q3_K); + test_matmul(backend, 17, 1024, 257, GGML_TYPE_Q3_K); // K past a full 256 tile + test_matmul(backend, 257, 1024, 17, GGML_TYPE_Q3_K); // K past a full microtile + test_matmul(backend, 256, 1024, 8192, GGML_TYPE_Q2_K); + test_matmul(backend, 357, 1024, 137, GGML_TYPE_Q2_K); + test_matmul(backend, 8, 256, 8, GGML_TYPE_Q2_K); + test_matmul(backend, 17, 1024, 257, GGML_TYPE_Q2_K); + test_matmul(backend, 257, 1024, 17, GGML_TYPE_Q2_K); + test_matmul(backend, 357, 1024, 137, GGML_TYPE_IQ4_XS); // ragged M and K + test_matmul(backend, 8, 256, 8, GGML_TYPE_IQ4_XS); // tiny M, single QK_K block + test_matmul(backend, 17, 1024, 257, GGML_TYPE_IQ4_XS); // K past a full 256 tile + + // fuzz: M/K around the microtile (16) and tile (256) boundaries + // Q4_K = subblock 32, Q6_K = subblock 16 + test_matmul(backend, 1, 1024, 1, GGML_TYPE_Q4_K); + test_matmul(backend, 2, 1024, 3, GGML_TYPE_Q4_K); + test_matmul(backend, 15, 1024, 15, GGML_TYPE_Q4_K); // M and K in the first microtile + test_matmul(backend, 17, 1024, 17, GGML_TYPE_Q4_K); // M and K past the first microtile + test_matmul(backend, 255, 1024, 255, GGML_TYPE_Q4_K); // M and K one short of a tile + test_matmul(backend, 257, 1024, 257, GGML_TYPE_Q4_K); // M and K one past a tile + test_matmul(backend, 271, 1024, 271, GGML_TYPE_Q4_K); // 1 tile + 15 + test_matmul(backend, 272, 1024, 272, GGML_TYPE_Q4_K); // 1 tile + 16 + test_matmul(backend, 511, 1024, 511, GGML_TYPE_Q4_K); // 2 tiles - 1 + test_matmul(backend, 513, 1024, 513, GGML_TYPE_Q4_K); // 2 tiles + 1 + test_matmul(backend, 17, 512, 257, GGML_TYPE_Q4_K); // ragged M and K, N = 2 blocks + + test_matmul(backend, 1, 1024, 1, GGML_TYPE_Q6_K); + test_matmul(backend, 2, 1024, 3, GGML_TYPE_Q6_K); + test_matmul(backend, 15, 1024, 15, GGML_TYPE_Q6_K); + test_matmul(backend, 17, 1024, 17, GGML_TYPE_Q6_K); + test_matmul(backend, 255, 1024, 255, GGML_TYPE_Q6_K); + test_matmul(backend, 257, 1024, 257, GGML_TYPE_Q6_K); + test_matmul(backend, 271, 1024, 271, GGML_TYPE_Q6_K); + test_matmul(backend, 272, 1024, 272, GGML_TYPE_Q6_K); + test_matmul(backend, 511, 1024, 511, GGML_TYPE_Q6_K); + test_matmul(backend, 513, 1024, 513, GGML_TYPE_Q6_K); + test_matmul(backend, 17, 512, 257, GGML_TYPE_Q6_K); + + // MUL_MAT_ID (MoE): K = reduction (tiled gate needs K % 256 == 0), R = output rows per expert, + // k = top-k slots, b_slots = b rows (1 = broadcast MoE, k = i11-diverse gather), cne1 = k*batch/E + test_mul_mat_id(backend, 512, 300, 4, 2, 1, 150, GGML_TYPE_Q4_K); // ragged R, broadcast, cne1 = 75 + test_mul_mat_id(backend, 512, 256, 4, 2, 2, 150, GGML_TYPE_Q5_K); // i11-diverse gather, HAS_MIN + test_mul_mat_id(backend, 512, 256, 4, 2, 1, 100, GGML_TYPE_Q6_K); // ragged cne1 = 50 + test_mul_mat_id(backend, 768, 256, 4, 2, 1, 150, GGML_TYPE_Q3_K); // K = 3 slabs + test_mul_mat_id(backend, 1024, 2048, 128, 8, 1, 2048, GGML_TYPE_Q4_K); // large, cne1 = 128 + test_mul_mat_id(backend, 512, 256, 32, 2, 1, 8, GGML_TYPE_Q4_K); // tiny cne1 (~0.5), forced only + + // cne1 > 256: the k-outer ring sweeps multiple 256-row windows per expert (the multi-window path) + test_mul_mat_id(backend, 256, 256, 8, 8, 1, 512, GGML_TYPE_Q4_K); // cne1 = 512, two full windows + test_mul_mat_id(backend, 512, 256, 8, 8, 1, 300, GGML_TYPE_Q4_K); // cne1 = 300, 2nd window ragged (44 rows) + test_mul_mat_id(backend, 256, 64, 8, 8, 1, 257, GGML_TYPE_Q4_K); // cne1 = 257, 1-row tail window, R = 1 group + test_mul_mat_id(backend, 512, 128, 8, 8, 1, 511, GGML_TYPE_Q4_K); // cne1 = 511, 15-row ragged tail + test_mul_mat_id(backend, 768, 300, 8, 8, 1, 300, GGML_TYPE_Q4_K); // 3 K slabs x 2 windows, ragged R and cne1 + test_mul_mat_id(backend, 768, 300, 8, 2, 2, 300, GGML_TYPE_Q5_K); // cne1 = 600, 3 slabs, 3 windows (last ragged), i11-diverse, HAS_MIN + + // cne1 = 255: one short of a full ring, last ring row zero-padded by the unpack + test_mul_mat_id(backend, 512, 256, 4, 4, 1, 255, GGML_TYPE_Q4_K); + + // long K: 8 slabs per ring row + test_mul_mat_id(backend, 2048, 300, 8, 2, 1, 400, GGML_TYPE_Q4_K); // cne1 = 100 + + // R group edges: 64 = TILED_MMID_GROUP (ceil div, ragged group tail, single-group early return) + test_mul_mat_id(backend, 512, 64, 4, 2, 1, 128, GGML_TYPE_Q4_K); // ngroups = 1, most threads idle + test_mul_mat_id(backend, 512, 65, 4, 2, 1, 128, GGML_TYPE_Q4_K); // ceil-div boundary, 1-row tail group + test_mul_mat_id(backend, 512, 63, 4, 2, 1, 128, GGML_TYPE_Q4_K); // ragged group tail + test_mul_mat_id(backend, 512, 1, 4, 1, 1, 256, GGML_TYPE_Q4_K); // R = 1 row, cne1 = 256 + + // strided (non-contiguous) F32 src1: nb[1] = 4 bytes, the tiled path reads through strides + test_mul_mat_id(backend, 512, 256, 8, 8, 1, 128, GGML_TYPE_Q4_K, true); // cne1 = 128 + + // narrow path (cne1 in [8,16], K a multiple of k_extent): the weight is read as long per-row K streams + test_mul_mat_id(backend, 1024, 256, 2, 2, 1, 16, GGML_TYPE_Q4_K); // cne1 = 16 + test_mul_mat_id(backend, 1024, 256, 2, 1, 1, 16, GGML_TYPE_Q4_K); // cne1 = 8 + test_mul_mat_id(backend, 512, 128, 2, 3, 1, 8, GGML_TYPE_Q5_K); // cne1 = 12 + + // higher-dim (ne[2], ne[3] > 1) + // equal-batch and broadcast shapes (src1_2=2*src0_2, src1_3=2*src0_3) + test_matmul_highdim(backend, 1024, 1024, 1024, 2, 1, 2, 1, GGML_TYPE_Q4_K); // 3D, tiled + test_matmul_highdim(backend, 1024, 1024, 1024, 1, 2, 1, 2, GGML_TYPE_Q4_K); // 3D, tiled + test_matmul_highdim(backend, 512, 1024, 512, 2, 2, 2, 2, GGML_TYPE_Q4_K); // 4D, tiled + test_matmul_highdim(backend, 1024, 1024, 1024, 2, 1, 1, 1, GGML_TYPE_Q4_K); // broadcast r2=2 (src1_2>src0_2) + test_matmul_highdim(backend, 1024, 1024, 1024, 2, 2, 2, 1, GGML_TYPE_Q4_K); // broadcast r3=2 (src1_3>src0_3) + test_matmul_highdim(backend, 1024, 1024, 1024, 2, 1, 2, 1, GGML_TYPE_Q6_K); // 3D, tiled, q6_K + test_matmul_highdim(backend, 1024, 1024, 1024, 2, 2, 2, 2, GGML_TYPE_Q3_K); // 4D, tiled, q3_K + test_matmul_highdim(backend, 1024, 1024, 1024, 2, 2, 2, 2, GGML_TYPE_Q2_K); // 4D, tiled, q2_K + } + + + // benchmark: std (use_ref) vs repack (CPU_REPACK buffer) vs tiled kernel, over all tiled + // types. std is timed with use_ref=true (bypasses the tiled gate); tiled and repack with + // use_ref=false. Errors are vs the vec-only reference (use_ref). + if (run_bench) { + printf("Starting benchmark runs.\n"); + const ggml_type bench_types[] = { GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, + GGML_TYPE_IQ4_XS, GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, + GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S, GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M }; + const size_t n_types = sizeof(bench_types) / sizeof(bench_types[0]); + + struct { int64_t M, N, K; } shapes[] = { + { 4096, 4096, 4096 }, + { 4096, 4096, 64 }, + { 4096, 4096, 32 }, + { 4096, 4096, 24 }, + { 4096, 4096, 16 }, + { 4096, 4096, 8 }, + { 4096, 4096, 1 }, + }; + for (size_t s = 0; s < sizeof(shapes) / sizeof(shapes[0]); ++s) { + bench_row rows[n_types]; + for (size_t i = 0; i < n_types; ++i) { + rows[i] = bench_three_way(backend, shapes[s].M, shapes[s].N, shapes[s].K, bench_types[i]); + } + print_bench_table(shapes[s].M, shapes[s].N, shapes[s].K, rows, n_types); + } + + // MUL_MAT_ID (MoE): K must be a multiple of 256 (the tiled slab). cne1 = k*batch/E. + struct { int64_t K, R, E, k, b_slots, batch; } mmid_shapes[] = { + { 1024, 1024, 8, 8, 1, 1 }, // cne1 = 1, single-token decode + { 1024, 1024, 4, 2, 1, 16 }, // cne1 = 8, narrow path + { 1024, 1024, 4, 2, 1, 32 }, // cne1 = 16, narrow path + { 512, 512, 8, 2, 1, 128 }, // cne1 = 32 + { 1024, 1024, 16, 8, 1, 64 }, // cne1 = 32 + { 1024, 1024, 16, 8, 1, 256 }, // cne1 = 128 + { 1024, 2048, 32, 8, 1, 512 }, // cne1 = 128, wide experts + { 2048, 1024, 16, 8, 1, 1024 }, // cne1 = 512, long dot + }; + for (size_t s = 0; s < sizeof(mmid_shapes) / sizeof(mmid_shapes[0]); ++s) { + bench_row_mmid rows[n_types]; + for (size_t i = 0; i < n_types; ++i) { + rows[i] = bench_mul_mat_id(backend, mmid_shapes[s].K, mmid_shapes[s].R, mmid_shapes[s].E, + mmid_shapes[s].k, mmid_shapes[s].b_slots, mmid_shapes[s].batch, bench_types[i]); + } + print_mmid_table(mmid_shapes[s].K, mmid_shapes[s].R, mmid_shapes[s].E, mmid_shapes[s].k, + mmid_shapes[s].b_slots, mmid_shapes[s].batch, rows, n_types); + } + } + + ggml_backend_free(backend); + return n_failed ? 1 : 0; +}