mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-27 21:46:57 +02:00
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 <[email protected]>
* 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 0688cf43d5.
* 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 <[email protected]>
* Cleanup merge-duplicated declaration of test-backend-ops target
* Undo accidental line deletion in ggml.c
---------
Co-authored-by: Georgi Gerganov <[email protected]>
This commit is contained in:
jbooth
co-authored by
Georgi Gerganov
Georgi Gerganov
parent
9f70b2cecd
commit
d834d44e64
+1
-1
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 <malloc.h> // 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:
|
||||
{
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -0,0 +1,719 @@
|
||||
// mulmat microtile kernels
|
||||
|
||||
#include "tiled-kernel.h"
|
||||
|
||||
#include "ggml.h"
|
||||
|
||||
#include <cstring>
|
||||
|
||||
#if defined(__AVX512VNNI__) || defined(__AVX2__) || defined(__AVX__)
|
||||
#include <immintrin.h>
|
||||
#endif
|
||||
|
||||
|
||||
// Reference implementation, slower than existing vec_dot approach
|
||||
template <int SUBBLK, bool HAS_MIN, int BIAS, int NK>
|
||||
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 <int SUBBLK, bool HAS_MIN, int BIAS, int NK>
|
||||
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 <int SUBBLK, bool HAS_MIN, int BIAS, int NK>
|
||||
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<SUBBLK, HAS_MIN, BIAS, NK>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
|
||||
tiled_run_micro_vnni_8x16<SUBBLK, HAS_MIN, BIAS, NK>(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 <int SUBBLK, bool HAS_MIN, int BIAS, int NK, bool ACTBIAS>
|
||||
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 <int SUBBLK, bool HAS_MIN, int BIAS, int NK>
|
||||
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 <int SUBBLK, bool HAS_MIN, int BIAS, bool ACTBIAS>
|
||||
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<SUBBLK, HAS_MIN, BIAS, 1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
|
||||
} else {
|
||||
tiled_run_microtile_vnni<SUBBLK, HAS_MIN, BIAS, 0>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
|
||||
}
|
||||
#elif defined(__AVX2__)
|
||||
if (standard) {
|
||||
tiled_run_microtile_avx2<SUBBLK, HAS_MIN, BIAS, 1, ACTBIAS>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
|
||||
} else {
|
||||
tiled_run_microtile_avx2<SUBBLK, HAS_MIN, BIAS, 0, ACTBIAS>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
|
||||
}
|
||||
#elif defined(__AVX__)
|
||||
if (standard) {
|
||||
tiled_run_microtile_avx<SUBBLK, HAS_MIN, BIAS, 1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
|
||||
} else {
|
||||
tiled_run_microtile_avx<SUBBLK, HAS_MIN, BIAS, 0>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
|
||||
}
|
||||
#else
|
||||
if (standard) {
|
||||
tiled_run_microtile_scalar<SUBBLK, HAS_MIN, BIAS, 1>(src0, src1, i0, j0, num_k, slab, buf, buf_stride);
|
||||
} else {
|
||||
tiled_run_microtile_scalar<SUBBLK, HAS_MIN, BIAS, 0>(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 <int SUBBLK>
|
||||
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);
|
||||
|
||||
|
||||
@@ -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 <stddef.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#if defined(__AVX2__)
|
||||
#include <immintrin.h>
|
||||
#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 <int S> 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 <int S, int D, int M>
|
||||
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 <int S>
|
||||
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 <int S, int D, int M>
|
||||
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 <int SUBBLK, bool HAS_MIN, int BIAS, bool ACTBIAS>
|
||||
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 <int SUBBLK>
|
||||
void tiled_repack_src0(tiled_tile_src0 * tile, int n_rows, int num_k, int BIAS, bool corr);
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,31 @@
|
||||
#pragma once
|
||||
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#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
|
||||
@@ -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")
|
||||
|
||||
@@ -0,0 +1,852 @@
|
||||
#include "ggml-alloc.h"
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
|
||||
#include <time.h>
|
||||
#include <math.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
|
||||
#if defined(_WIN32)
|
||||
#include <windows.h>
|
||||
#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;
|
||||
}
|
||||
Reference in New Issue
Block a user