Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-07 22:59:00 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

E-2: prefetch the i-quant expert rows ahead of the AVX-512 decode; an opt-in gather decode

Browse Source 
Measured on a Ryzen 5 7600 (IQ3_S, 4K context, CLI decode timeline, outputs identical):
gate/up 37.1 -> 36.5 ms/round with a 2 KB prefetch; the gather decode (same results) is 3-5%
slower than the scalar lookups, so it is opt-in (STRATA_IQ_GATHER=1). More pool threads (SMT)
do not help: the pool reads its rows at ~30 GB/s, near DDR5's practical rate through 4 KB pages
- memory-bound, so decode speed now comes from VRAM hits (E-1) rather than CPU kernels.
This commit is contained in:
Niko1221 committed 2026-09-29 07:59:37 +02:00
1 parent 07b107403a
commit 24d4013c87
1 file changed
+29 -2
Download Patch File Download Diff File Expand all files Collapse all files
src/kernels/cpu/iq_avx512.cpp
+29 -2
View File
@@ -6,7 +6,8 @@
// `maddubs` and a `madd`. The arithmetic is ggml's (ggml-cpu/quants.c, the `_generic` references): integer sums
// per block, times d_x * d_y * the format's constant - only the order of the float additions differs.
//
// Formats: IQ2_XXS (16), IQ2_XS (17), IQ3_XXS (18), IQ3_S (21), IQ2_S (22). IQ1_M stays on ggml-cpu.
// Formats: IQ2_XXS (16), IQ2_XS (17), IQ3_XXS (18), IQ3_S (21), IQ2_S (22). IQ1_M stays on ggml-cpu (no shipped
// model has IQ1_M expert rows: the 'Coder IQ1_M' pack's gate/up are IQ2_S / IQ3_XXS / IQ3_S).
#include "strata/kernels/cpu/iq_avx512.hpp"
#define GGML_COMMON_DECL_CPP
@@ -16,6 +17,7 @@
#include <immintrin.h>
#include <cmath>
#include <cstdlib>
#include <cstring>
namespace strata::kernels::cpu {
@@ -26,6 +28,16 @@ inline uint32_t u32(const uint8_t* p) { uint32_t v; std::memcpy(&v, p, 4); retur
inline uint16_t u16(const uint8_t* p) { uint16_t v; std::memcpy(&v, p, 2); return v; }
inline uint64_t u64(const uint8_t* p) { uint64_t v; std::memcpy(&v, p, 8); return v; }
// E-2: the IQ3 grids by one AVX-512 gather (the same words, the same results) - measured 3-5% SLOWER than the scalar
// lookups on a Ryzen 5 7600 (Zen 4 gathers are microcoded), so opt-in: STRATA_IQ_GATHER=1
const bool gather = std::getenv("STRATA_IQ_GATHER") != nullptr;
// E-2: software prefetch distance in bytes (STRATA_IQ_PREFETCH; 0 = off): -2-3% gate/up time on IQ3_S decode
// (the rows are read at ~30 GB/s by 6 cores, near what DDR5 with 4 KB pages gives: the pool is memory-bound)
const int prefetch_ahead = [] {
const char* v = std::getenv("STRATA_IQ_PREFETCH");
return v ? std::atoi(v) : 2048;
}();
// lanes 0-7 -> s0, 8-15 -> s1, 16-23 -> s2, 24-31 -> s3 (int16 lanes of a maddubs result: two values each)
inline __m512i scales4(int s0, int s1, int s2, int s3) {
const __m512i idx = _mm512_set_epi16(3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1,
@@ -98,7 +110,9 @@ template <> struct Fmt<18> { // IQ3_XXS: d, qs[64] grid bytes, 8 x u32 (4 x 7-
static constexpr float K = 0.25f;
static inline void decode(const uint8_t* b, int j, __m512i& g, __mmask64& m, __m512i& sc) {
const uint8_t* q = b + 2 + 16 * j;
g = _mm512_set_epi32((int) iq3xxs_grid[q[15]], (int) iq3xxs_grid[q[14]], (int) iq3xxs_grid[q[13]], (int) iq3xxs_grid[q[12]],
// E-2: one gather of the 16 grid words (the scalar lookups assembled with set_epi32 cost ~30 uops)
g = gather ? _mm512_i32gather_epi32(_mm512_cvtepu8_epi32(_mm_loadu_si128((const __m128i*) q)), iq3xxs_grid, 4)
: _mm512_set_epi32((int) iq3xxs_grid[q[15]], (int) iq3xxs_grid[q[14]], (int) iq3xxs_grid[q[13]], (int) iq3xxs_grid[q[12]],
(int) iq3xxs_grid[q[11]], (int) iq3xxs_grid[q[10]], (int) iq3xxs_grid[q[9]], (int) iq3xxs_grid[q[8]],
(int) iq3xxs_grid[q[7]], (int) iq3xxs_grid[q[6]], (int) iq3xxs_grid[q[5]], (int) iq3xxs_grid[q[4]],
(int) iq3xxs_grid[q[3]], (int) iq3xxs_grid[q[2]], (int) iq3xxs_grid[q[1]], (int) iq3xxs_grid[q[0]]);
@@ -119,12 +133,18 @@ template <> struct Fmt<21> { // IQ3_S: d, qs[64], qh[8], signs[32], scales[4]
static inline void decode(const uint8_t* b, int j, __m512i& g, __mmask64& m, __m512i& sc) {
const uint8_t* q = b + 2 + 16 * j;
const uint32_t h0 = b[66 + 2 * j], h1 = b[66 + 2 * j + 1];
if (gather) { // E-2: the 9-bit indices (the high bit from qh under a mask), one gather
__m512i idx = _mm512_cvtepu8_epi32(_mm_loadu_si128((const __m128i*) q));
idx = _mm512_mask_add_epi32(idx, (__mmask16) (h0 | (h1 << 8)), idx, _mm512_set1_epi32(256));
g = _mm512_i32gather_epi32(idx, iq3s_grid, 4);
} else {
#define G3(k, h, kk) (int) iq3s_grid[q[k] | (((h >> kk) & 1) << 8)]
g = _mm512_set_epi32(G3(15, h1, 7), G3(14, h1, 6), G3(13, h1, 5), G3(12, h1, 4),
G3(11, h1, 3), G3(10, h1, 2), G3(9, h1, 1), G3(8, h1, 0),
G3(7, h0, 7), G3(6, h0, 6), G3(5, h0, 5), G3(4, h0, 4),
G3(3, h0, 3), G3(2, h0, 2), G3(1, h0, 1), G3(0, h0, 0));
#undef G3
}
m = _cvtu64_mask64(u64(b + 74 + 8 * j));
const uint8_t s = b[106 + j];
const int sa = 2 * (s & 15) + 1, sb = 2 * (s >> 4) + 1;
@@ -139,6 +159,13 @@ inline void row_dot(const uint8_t* row, int nblocks, const block_q8_K* const* y,
const __m512i zero = _mm512_setzero_si512();
for (int i = 0; i < nblocks; ++i) {
const uint8_t* blk = row + (size_t) i * Fmt<TY>::bytes;
// E-2: the weight rows stream from DRAM (4 KB pages when large pages are refused): ask for the bytes a few
// blocks ahead (the next row's start included) before the decode needs them
if (prefetch_ahead > 0) {
const uint8_t* ahead = blk + (size_t) prefetch_ahead;
_mm_prefetch((const char*) ahead, _MM_HINT_T0);
_mm_prefetch((const char*) ahead + 64, _MM_HINT_T0);
}
__m512i acci[NT];
for (int t = 0; t < NT; ++t) acci[t] = _mm512_setzero_si512();
for (int j = 0; j < 4; ++j) {