millaguie/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Actions 31 Packages Projects Releases Wiki Activity

ggml : add AVX-512 DT3 vec_dot

Browse Source 
Decode both planes of a block with VBMI byte permutes: the *3 multiply
chain (wrapping, so it commutes with the permutation) is computed once
on the whole 56-byte block, and masked vpermb picks each element's byte
from the chain vector of its base-3 digit. The qh lanes never see 3^4,
which would read the padding 5th digit of the qh bytes. The trits reach
the integer product as xi in {0, 1, 2} via the same avg trick as
tq1_0, with VNNI dpbusd against the q8_0 bytes and sum(y) subtracted.

The per-q8_0-block sums and the float accumulation keep the exact
operation order of the generic implementation, so the result is
bit-identical to it (checked by test-dt3).

2.2x over the (autovectorized) generic on a Ryzen AI 9 HX 370.
This commit is contained in:
Millaguie
2026-08-10 18:28:15 +02:00
parent 51c6b67e8b
commit b110945afc
2 changed files with 173 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cpu/arch-fallback.h
-1
View File
@@ -87,7 +87,6 @@
#elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64)
// quants.c
#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
#define ggml_vec_dot_dt3_q8_0_generic ggml_vec_dot_dt3_q8_0
// repack.cpp
#define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
#define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
ggml/src/ggml-cpu/arch/x86/quants.c
+173
View File
@@ -1571,6 +1571,179 @@ void ggml_vec_dot_tq2_0_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
#endif
}
#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VBMI__) && defined(__AVX512VNNI__)
// multiply bytes by 3 with wraparound (there is no 8-bit SIMD multiply)
static inline __m512i dt3_mul3_epi8(const __m512i v) {
return _mm512_add_epi8(v, _mm512_add_epi8(v, v));
}
// bring the top base-3 digit of each byte down to xi = ((uint8_t) q * 3) >> 8,
// in {0, 1, 2}, with the same avg trick as ggml_vec_dot_tq1_0_q8_K
static inline __m512i dt3_decode_epi8(__m512i q) {
// cancel the +1 from avg so that it behaves like a halving add
q = _mm512_subs_epu8(q, _mm512_set1_epi8(1));
// multiply by 3 and get the top 2 bits
q = _mm512_avg_epu8(q, _mm512_avg_epu8(q, _mm512_setzero_si512()));
return _mm512_and_si512(_mm512_srli_epi16(q, 6), _mm512_set1_epi8(3));
}
#endif
void ggml_vec_dot_dt3_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
assert(n % QK_DT3 == 0);
assert(nrc == 1);
UNUSED(nrc);
UNUSED(bx);
UNUSED(by);
UNUSED(bs);
const block_dt3 * GGML_RESTRICT x = vx;
const block_q8_0 * GGML_RESTRICT y = vy;
const int nb = n / QK_DT3;
#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VBMI__) && defined(__AVX512VNNI__)
// Source byte of each of the 128 elements of a plane, in element order, as
// offsets into the 56-byte block for plane 0 (see unpack_plane_dt3):
// element m + n*16 (0..79) is digit n of qs[m], m in [0,16)
// element 80 + m + n*8 (80..119) is digit n of qs[16 + m], m in [0,8)
// element 120 + j + n*2 (120..127) is digit n of qh[j], j in [0,2)
// The low vector covers elements 0..63 and the high vector 64..127, so
// that each aligns with two full q8_0 blocks of the other operand.
static const uint8_t kidx_lo[64] = {
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 0 of qs[0..15]
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 1
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 2
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 3
};
static const uint8_t kidx_hi[64] = {
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, // digit 4 of qs[0..15]
16, 17, 18, 19, 20, 21, 22, 23, // digit 0 of qs[16..23]
16, 17, 18, 19, 20, 21, 22, 23, // digit 1
16, 17, 18, 19, 20, 21, 22, 23, // digit 2
16, 17, 18, 19, 20, 21, 22, 23, // digit 3
16, 17, 18, 19, 20, 21, 22, 23, // digit 4
48, 49, 48, 49, 48, 49, 48, 49, // digits 0,0,1,1,2,2,3,3 of qh[0],qh[1]
};
// plane 1 offsets: qs starts 24 bytes later, qh starts 2 bytes later
static const uint8_t koff_hi[64] = {
24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,
24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,
24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24,
24, 24, 24, 24, 24, 24, 24, 24, 2, 2, 2, 2, 2, 2, 2, 2,
};
// digit -> multiplier blend masks, one bit per byte lane (digit n of a
// byte is extracted by multiplying by 3^n and reading the top 2 bits;
// the qh lanes of the high vector never see 3^4, which would read the
// padding 5th digit of the qh bytes)
const __mmask64 mask_lo_3 = (__mmask64) 0x00000000FFFF0000; // lanes 16..31
const __mmask64 mask_lo_9 = (__mmask64) 0x0000FFFF00000000; // lanes 32..47
const __mmask64 mask_lo_27 = (__mmask64) 0xFFFF000000000000; // lanes 48..63
const __mmask64 mask_hi_1 = (__mmask64) 0x0300000000FF0000; // lanes 16..23, 56, 57
const __mmask64 mask_hi_3 = (__mmask64) 0x0C000000FF000000; // lanes 24..31, 58, 59
const __mmask64 mask_hi_9 = (__mmask64) 0x300000FF00000000; // lanes 32..39, 60, 61
const __mmask64 mask_hi_27 = (__mmask64) 0xC000FF0000000000; // lanes 40..47, 62, 63
const __m512i idx_lo0 = _mm512_loadu_si512(kidx_lo);
const __m512i idx_hi0 = _mm512_loadu_si512(kidx_hi);
const __m512i idx_lo1 = _mm512_add_epi8(idx_lo0, _mm512_set1_epi8(24));
const __m512i idx_hi1 = _mm512_add_epi8(idx_hi0, _mm512_loadu_si512(koff_hi));
const __m512i zero = _mm512_setzero_si512();
const __m512i ones = _mm512_set1_epi8(1);
float sumf = 0.0f;
for (int i = 0; i < nb; i++) {
// the block is 56 bytes; the masked load must not read past the end
const __m512i blk = _mm512_maskz_loadu_epi8((__mmask64) ((UINT64_C(1) << sizeof(block_dt3)) - 1), &x[i]);
// multiplying by 3^n (with wraparound) commutes with the byte
// permutation, so the multiply chain is computed once on the whole
// block and shared by both planes; masked vpermb then picks each
// lane's byte from the chain vector of its digit
const __m512i v3 = dt3_mul3_epi8(blk);
const __m512i v9 = dt3_mul3_epi8(v3);
const __m512i v27 = dt3_mul3_epi8(v9);
const __m512i v81 = dt3_mul3_epi8(v27);
__m512i xi_lo[2];
__m512i xi_hi[2];
for (int p = 0; p < 2; p++) {
const __m512i idx_lo = p == 0 ? idx_lo0 : idx_lo1;
const __m512i idx_hi = p == 0 ? idx_hi0 : idx_hi1;
__m512i q_lo = _mm512_permutexvar_epi8(idx_lo, blk);
q_lo = _mm512_mask_permutexvar_epi8(q_lo, mask_lo_3, idx_lo, v3);
q_lo = _mm512_mask_permutexvar_epi8(q_lo, mask_lo_9, idx_lo, v9);
q_lo = _mm512_mask_permutexvar_epi8(q_lo, mask_lo_27, idx_lo, v27);
xi_lo[p] = dt3_decode_epi8(q_lo);
__m512i q_hi = _mm512_permutexvar_epi8(idx_hi, v81);
q_hi = _mm512_mask_permutexvar_epi8(q_hi, mask_hi_1, idx_hi, blk);
q_hi = _mm512_mask_permutexvar_epi8(q_hi, mask_hi_3, idx_hi, v3);
q_hi = _mm512_mask_permutexvar_epi8(q_hi, mask_hi_9, idx_hi, v9);
q_hi = _mm512_mask_permutexvar_epi8(q_hi, mask_hi_27, idx_hi, v27);
xi_hi[p] = dt3_decode_epi8(q_hi);
}
// one DT3 block (128 weights) maps to four q8_0 blocks (4 * 32 = 128)
const __m512i y_lo = _mm512_inserti64x4(_mm512_castsi256_si512(
_mm256_loadu_si256((const __m256i *) y[4*i + 0].qs)),
_mm256_loadu_si256((const __m256i *) y[4*i + 1].qs), 1);
const __m512i y_hi = _mm512_inserti64x4(_mm512_castsi256_si512(
_mm256_loadu_si256((const __m256i *) y[4*i + 2].qs)),
_mm256_loadu_si256((const __m256i *) y[4*i + 3].qs), 1);
// t = xi - 1, so t.y = xi.y - sum(y)
const __m512i sy_lo = _mm512_dpbusd_epi32(zero, ones, y_lo);
const __m512i sy_hi = _mm512_dpbusd_epi32(zero, ones, y_hi);
const __m512i t1_lo = _mm512_sub_epi32(_mm512_dpbusd_epi32(zero, xi_lo[0], y_lo), sy_lo);
const __m512i t1_hi = _mm512_sub_epi32(_mm512_dpbusd_epi32(zero, xi_hi[0], y_hi), sy_hi);
const __m512i t2_lo = _mm512_sub_epi32(_mm512_dpbusd_epi32(zero, xi_lo[1], y_lo), sy_lo);
const __m512i t2_hi = _mm512_sub_epi32(_mm512_dpbusd_epi32(zero, xi_hi[1], y_hi), sy_hi);
// reduce each q8_0 block (8 consecutive int32 lanes) to its sum with
// an hadd tree; integer addition order does not affect the result:
// ab = [A01 A23 B01 B23 | A45 A67 B45 B67]
// abcd = [A0123 B0123 C0123 D0123 | A4567 B4567 C4567 D4567]
// sv = [sum(A) sum(B) sum(C) sum(D)]
const __m256i ab1 = _mm256_hadd_epi32(_mm512_castsi512_si256(t1_lo), _mm512_extracti64x4_epi64(t1_lo, 1));
const __m256i cd1 = _mm256_hadd_epi32(_mm512_castsi512_si256(t1_hi), _mm512_extracti64x4_epi64(t1_hi, 1));
const __m256i ab2 = _mm256_hadd_epi32(_mm512_castsi512_si256(t2_lo), _mm512_extracti64x4_epi64(t2_lo, 1));
const __m256i cd2 = _mm256_hadd_epi32(_mm512_castsi512_si256(t2_hi), _mm512_extracti64x4_epi64(t2_hi, 1));
const __m256i abcd1 = _mm256_hadd_epi32(ab1, cd1);
const __m256i abcd2 = _mm256_hadd_epi32(ab2, cd2);
const __m128i sv1 = _mm_add_epi32(_mm256_castsi256_si128(abcd1), _mm256_extracti128_si256(abcd1, 1));
const __m128i sv2 = _mm_add_epi32(_mm256_castsi256_si128(abcd2), _mm256_extracti128_si256(abcd2, 1));
int sumi1[4];
int sumi2[4];
_mm_storeu_si128((__m128i *) sumi1, sv1);
_mm_storeu_si128((__m128i *) sumi2, sv2);
const float d1 = GGML_CPU_FP16_TO_FP32(x[i].d[0]);
const float d2 = GGML_CPU_FP16_TO_FP32(x[i].d[1]);
// same accumulation order as the generic implementation
for (int k = 0; k < 4; k++) {
const float dy = GGML_CPU_FP16_TO_FP32(y[4*i + k].d);
sumf += dy * (d1*sumi1[k] + d2*sumi2[k]);
}
}
*s = sumf;
#else
UNUSED(x);
UNUSED(y);
UNUSED(nb);
ggml_vec_dot_dt3_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
#endif
}
void ggml_vec_dot_q2_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
assert(nrc == 1);
UNUSED(nrc);