nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Optimize mul_mat_q8_1_r8_q8_2 with AVX-512 for faster Q4_K/Q5_K prompt processing (#1578) 

Add HAVE_FANCY_SIMD path that processes 16 rows at a time using 512-bit
operations, combining two R8 groups via _mm512_inserti32x8. Reuses the
existing qx_r8_q8_dot_product 512-bit overload for the inner dot product.
Also updates num_rows for Q8_1 to 16 under HAVE_FANCY_SIMD.

Co-authored-by: Adam Caldwell <accaldwell@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Adam Caldwell 2026-04-06 00:07:23 -07:00 committed by GitHub
parent 5e8bb724ce
commit 6d4cdef511
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 108 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

107
ggml/src/iqk/iqk_gemm_legacy_quants.cpp
View File

@ -1760,6 +1760,112 @@ typedef struct {
uint8_t qs[256];
} block_q8_1_r8;
#ifdef HAVE_FANCY_SIMD
template <int nrc_y>
static void mul_mat_q8_1_r8_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) {
GGML_ASSERT(nrc_x%16 == 0);
Q8<nrc_y, block_q8_2_x4> q8(info);
int nb = n / QK8_0;
if constexpr (nrc_y == 1) {
__m256 acc[1] = {};
float d8[4];
__m256i qx[4];
auto dot = [&qx] (const int8_t * qy) {
auto y128 = _mm_loadu_si128((const __m128i*)qy);
auto y = MM256_SET_M128I(y128, y128);
auto sumi = _mm256_setzero_si256();
sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00));
sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55));
sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa));
sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff));
return sumi;
};
for (int ix = 0; ix < nrc_x; ix += 8) {
const block_q8_1_r8 * iq8 = (const block_q8_1_r8 *)((const char *)vx + ix*bx);
for (int i4 = 0; i4 < nb/4; ++i4) {
{
__m256 mx[4];
for (int ib32 = 0; ib32 < 4; ++ib32) mx[ib32] = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[4*i4+ib32].d+1));
auto scales = _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)q8.y[0][i4].d)), 16));
_mm_storeu_ps(d8, scales);
auto bsums4 = _mm_cvtepi32_ps(_mm_cvtepi16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[0][i4].d+4))));
bsums4 = _mm_mul_ps(bsums4, scales);
auto bsums = _mm256_set_m128(bsums4, bsums4);
acc[0] = _mm256_fmadd_ps(mx[0], _mm256_shuffle_ps(bsums, bsums, 0x00), acc[0]);
acc[0] = _mm256_fmadd_ps(mx[1], _mm256_shuffle_ps(bsums, bsums, 0x55), acc[0]);
acc[0] = _mm256_fmadd_ps(mx[2], _mm256_shuffle_ps(bsums, bsums, 0xaa), acc[0]);
acc[0] = _mm256_fmadd_ps(mx[3], _mm256_shuffle_ps(bsums, bsums, 0xff), acc[0]);
}
for (int ib32 = 0; ib32 < 4; ++ib32) {
auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[4*i4+ib32].d));
for (int j = 0; j < 4; ++j) {
qx[j] = _mm256_loadu_si256((const __m256i *)iq8[4*i4+ib32].qs+j);
}
auto sumi = dot(q8.y[0][i4].qs+32*ib32);
auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[ib32]));
acc[0] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[0]);
for (int j = 0; j < 4; ++j) {
qx[j] = _mm256_loadu_si256((const __m256i *)iq8[4*i4+ib32].qs+4+j);
}
sumi = dot(q8.y[0][i4].qs+32*ib32+16);
d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[ib32]));
acc[0] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[0]);
}
}
info.store(ix, 0, acc[0]);
acc[0] = _mm256_setzero_ps();
}
} else {
__m512 acc[nrc_y] = {};
__m512i qx[8];
float d8[4*nrc_y];
for (int ix = 0; ix < nrc_x; ix += 16) {
const block_q8_1_r8 * q8l = (const block_q8_1_r8 *)((const char *)vx + (ix+0)*bx);
const block_q8_1_r8 * q8h = (const block_q8_1_r8 *)((const char *)vx + (ix+8)*bx);
for (int i4 = 0; i4 < nb/4; ++i4) {
{
__m512 mx[4];
for (int ib32 = 0; ib32 < 4; ++ib32) {
auto mx_l = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8l[4*i4+ib32].d+1));
auto mx_h = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8h[4*i4+ib32].d+1));
mx[ib32] = _mm512_insertf32x8(_mm512_castps256_ps512(mx_l), mx_h, 1);
}
for (int iy = 0; iy < nrc_y; ++iy) {
auto scales128 = _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)q8.y[iy][i4].d)), 16));
_mm_storeu_ps(d8 + 4*iy, scales128);
auto bsums4 = _mm_cvtepi32_ps(_mm_cvtepi16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][i4].d+4))));
bsums4 = _mm_mul_ps(bsums4, scales128);
auto bsums256 = _mm256_set_m128(bsums4, bsums4);
auto bsums = _mm512_insertf32x8(_mm512_castps256_ps512(bsums256), bsums256, 1);
acc[iy] = _mm512_fmadd_ps(mx[0], _mm512_shuffle_ps(bsums, bsums, 0x00), acc[iy]);
acc[iy] = _mm512_fmadd_ps(mx[1], _mm512_shuffle_ps(bsums, bsums, 0x55), acc[iy]);
acc[iy] = _mm512_fmadd_ps(mx[2], _mm512_shuffle_ps(bsums, bsums, 0xaa), acc[iy]);
acc[iy] = _mm512_fmadd_ps(mx[3], _mm512_shuffle_ps(bsums, bsums, 0xff), acc[iy]);
}
}
for (int ib32 = 0; ib32 < 4; ++ib32) {
auto scales_l = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8l[4*i4+ib32].d));
auto scales_h = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8h[4*i4+ib32].d));
auto scales = _mm512_insertf32x8(_mm512_castps256_ps512(scales_l), scales_h, 1);
for (int j = 0; j < 8; ++j) {
qx[j] = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)q8l[4*i4+ib32].qs+j)),
_mm256_loadu_si256((const __m256i *)q8h[4*i4+ib32].qs+j), 1);
}
for (int iy = 0; iy < nrc_y; ++iy) {
auto sumi = qx_r8_q8_dot_product(qx, q8.y[iy][i4].qs+32*ib32);
auto dy = _mm512_set1_ps(d8[4*iy+ib32]);
acc[iy] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[iy]);
}
}
}
for (int iy = 0; iy < nrc_y; ++iy) {
info.store(ix, iy, acc[iy]);
acc[iy] = _mm512_setzero_ps();
}
}
}
}
#else
template <int nrc_y>
static void mul_mat_q8_1_r8_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) {
GGML_ASSERT(nrc_x%8 == 0);
@ -1830,6 +1936,7 @@ static void mul_mat_q8_1_r8_q8_2(int n, const void * vx, size_t bx, const DataIn
}
}
}
#endif
void iqk_convert_q80_q80_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) {
static_assert(QK4_0 == QK8_0);

2
ggml/src/iqk/iqk_mul_mat.cpp
View File

@ -352,10 +352,10 @@ struct MulMat {
case GGML_TYPE_Q5_K_R4:
case GGML_TYPE_Q8_KV:
case GGML_TYPE_Q8_KV_R8:
case GGML_TYPE_Q8_1:
case GGML_TYPE_Q8_K_R8: return 8;
case GGML_TYPE_Q4_0_R8:
case GGML_TYPE_Q8_0_R8:
case GGML_TYPE_Q8_1:
case GGML_TYPE_Q8_K_R16:
case GGML_TYPE_BF16_R16: return 16;
default: return 1;