CUDA: faster MMVQ for low-bit types (#30517)

* CUDA: faster MMVQ for low-bit types

* ggml-cuda: hip: small nwarps GCN fix on Q3_K(ncols=2,3), Q6_K(ncols=4)

---------

Co-authored-by: thelittlefireman <thelittlefireman3@gmail.com>
This commit is contained in:
Aman Guptaandthelittlefireman authored and GitHub committed 2026-10-17 15:35:30 +02:00
1 parent 10b8b75534
commit 2744e48d1e
2 files changed
+125 -59

No files matched your search

+29 -6
View File
@@ -468,10 +468,12 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d
} else if (table_id == MMVQ_PARAMETERS_GCN) {
switch (ncols_dst) {
case 1:
return 2;
case 2:
case 3:
return type == GGML_TYPE_Q3_K ? 1 : 2;
case 4:
return 2;
return type == GGML_TYPE_Q6_K? 1 : 2;
case 5:
case 6:
case 7:
@@ -576,11 +578,31 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d
return 1;
}
static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) {
// Low bit types load few weight bytes per thread, a second row per block shares the q8_1 loads and adds memory parallelism.
static constexpr __host__ __device__ bool mmvq_two_rows_ncols1(ggml_type type) {
switch (type) {
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_IQ1_S:
case GGML_TYPE_IQ1_M:
return true;
default:
return false;
}
}
static constexpr __host__ __device__ int calc_rows_per_block(ggml_type type, int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) {
if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) {
switch (ncols_dst) {
case 1:
return small_k ? nwarps : 1;
if (small_k) {
return nwarps;
}
if (table_id == MMVQ_PARAMETERS_GB10) {
return type == GGML_TYPE_IQ2_S || type == GGML_TYPE_IQ3_S ? 2 : 1;
}
return table_id == MMVQ_PARAMETERS_GENERIC && mmvq_two_rows_ncols1(type) ? 2 : 1;
case 2:
case 3:
case 4:
@@ -615,7 +637,7 @@ static __global__ void mul_mat_vec_q(
constexpr int vdr = get_vdr_mmvq(type);
constexpr mmvq_parameter_table_id table_id = get_device_table_id();
constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters);
constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps);
constexpr int rows_per_cuda_block = calc_rows_per_block(type, ncols_dst, table_id, small_k, nwarps);
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);
@@ -727,7 +749,8 @@ static __global__ void mul_mat_vec_q(
#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
// start the next iterations' weight loads early
if constexpr (mmvq_should_prefetch(type)) {
constexpr int pf_dist = 2; // loop iterations, not blocks
// loop iterations, not blocks; q3_K and q6_K keep the distance in blocks they had with vdr 1
constexpr int pf_dist = type == GGML_TYPE_Q3_K || type == GGML_TYPE_Q6_K ? 1 : 2;
const int kbx_pf = kbx + pf_dist*blocks_per_iter;
if (kbx_pf < blocks_per_row_x) {
#pragma unroll
@@ -1013,7 +1036,7 @@ static std::pair<dim3, dim3> calc_launch_params(
const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens,
const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false, const bool halve_iters = false) {
const int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters);
const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps);
const int rpb = calc_rows_per_block(type, ncols_dst, table_id, small_k, nwarps);
const int64_t nblocks = (nrows_x + rpb - 1) / rpb;
const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens);
const dim3 block_dims(warp_size, nwarps, 1);
+96 -53
View File
@@ -360,12 +360,12 @@ static __device__ __forceinline__ float vec_dot_nvfp4_q8_1(
return sum;
}
#define VDR_Q2_K_Q8_1_MMVQ 1
#define VDR_Q2_K_Q8_1_MMVQ 2
#define VDR_Q2_K_Q8_1_MMQ 4
// contiguous v/x values
static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq(
const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales,
const int * __restrict__ v, const int * __restrict__ u, const int * __restrict__ scales,
const half2 & dm2, const float * __restrict__ d8) {
float sumf_d = 0.0f;
@@ -373,17 +373,20 @@ static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq(
#pragma unroll
for (int i = 0; i < QR2_K; ++i) {
const int sc = scales[2*i];
const int sc = scales[i/2] >> (16*(i%2));
const int vi = (v >> (2*i)) & 0x03030303;
int sumi_d = 0;
int sumi_m = 0;
#pragma unroll
for (int j = 0; j < VDR_Q2_K_Q8_1_MMVQ; ++j) {
const int vi = (v[j] >> (2*i)) & 0x03030303;
sumf_d += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product
sumi_d = ggml_cuda_dp4a(vi, u[VDR_Q2_K_Q8_1_MMVQ*i + j], sumi_d); // SIMD dot product
sumi_m = ggml_cuda_dp4a(0x01010101, u[VDR_Q2_K_Q8_1_MMVQ*i + j], sumi_m); // sum of q8_1 values
}
// fill int with 4x m
int m = sc >> 4;
m |= m << 8;
m |= m << 16;
sumf_m += d8[i] * ggml_cuda_dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values
sumf_d += d8[i] * (sumi_d * (sc & 0xF));
sumf_m += d8[i] * (sumi_m * ((sc >> 4) & 0xF)); // multiply constant q2_K part with sum of q8_1 values
}
const float2 dm2f = __half22float2(dm2);
@@ -443,37 +446,35 @@ static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq(
return sumf + d8*sumf_d8;
}
#define VDR_Q3_K_Q8_1_MMVQ 1
#define VDR_Q3_K_Q8_1_MMVQ 2
#define VDR_Q3_K_Q8_1_MMQ 2
// contiguous v/x values
static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq(
const int & vl, const int & vh, const int * __restrict__ u, const uint8_t * __restrict__ scales,
const int & scale_offset, const float & d3, const float * __restrict__ d8) {
const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const int * __restrict__ scales,
const float & d3, const float * __restrict__ d8) {
float sumf = 0.0f;
#pragma unroll
for (int i = 0; i < QR3_K; ++i) {
const int isc = scale_offset + 2*i;
// 4 low bits from scales[0..7], 2 high bits from scales[8..11]
const int sc_low = (scales[i/2] >> (16*(i%2))) & 0x0F;
const int sc_high = (scales[2] >> (16*(i%2) + 2*(i/2))) & 0x03;
const int sc = (sc_low | (sc_high << 4)) - 32;
const int isc_low = isc % (QK_K/32);
const int sc_shift_low = 4 * (isc / (QK_K/32));
const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF;
int sumi = 0;
#pragma unroll
for (int j = 0; j < VDR_Q3_K_Q8_1_MMVQ; ++j) {
const int vil = (vl[j] >> (2*i)) & 0x03030303;
const int vih = ((vh[j] >> i) << 2) & 0x04040404;
const int isc_high = isc % (QK_K/64);
const int sc_shift_high = 2 * (isc / (QK_K/64));
const int sc_high = ((scales[(QK_K/32) + isc_high] >> sc_shift_high) & 3) << 4;
// (vil | vih) - 4 for each byte, the - 4 is done by the second dp4a
sumi = ggml_cuda_dp4a(vil | vih, u[VDR_Q3_K_Q8_1_MMVQ*i + j], sumi); // SIMD dot product
sumi = ggml_cuda_dp4a(0xFCFCFCFC, u[VDR_Q3_K_Q8_1_MMVQ*i + j], sumi);
}
const int sc = (sc_low | sc_high) - 32;
const int vil = (vl >> (2*i)) & 0x03030303;
const int vih = ((vh >> i) << 2) & 0x04040404;
const int vi = __vsub4(vil, vih);
sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product
sumf += d8[i] * (sumi * sc);
}
return d3 * sumf;
@@ -620,12 +621,12 @@ static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq(
return dm4f.x*sumf_d - dm4f.y*sumf_m;
}
#define VDR_Q6_K_Q8_1_MMVQ 1
#define VDR_Q6_K_Q8_1_MMVQ 2
#define VDR_Q6_K_Q8_1_MMQ 8
// contiguous v/x values
static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq(
const int & vl, const int & vh, const int * __restrict__ u, const int8_t * __restrict__ scales,
const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const int8_t * __restrict__ scales,
const float & d, const float * __restrict__ d8) {
float sumf = 0.0f;
@@ -634,13 +635,18 @@ static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq(
for (int i = 0; i < QR6_K; ++i) {
const int sc = scales[4*i];
const int vil = (vl >> (4*i)) & 0x0F0F0F0F;
int sumi = 0;
#pragma unroll
for (int j = 0; j < VDR_Q6_K_Q8_1_MMVQ; ++j) {
const int vil = (vl[j] >> (4*i)) & 0x0F0F0F0F;
const int vih = ((vh[j] >> (4*i)) << 4) & 0x30303030;
const int vih = ((vh >> (4*i)) << 4) & 0x30303030;
// (vil | vih) - 32 for each byte, the - 32 is done by the second dp4a
sumi = ggml_cuda_dp4a(vil | vih, u[VDR_Q6_K_Q8_1_MMVQ*i + j], sumi); // SIMD dot product
sumi = ggml_cuda_dp4a(0xE0E0E0E0, u[VDR_Q6_K_Q8_1_MMVQ*i + j], sumi);
}
const int vi = __vsub4((vil | vih), 0x20202020); // vi = (vil | vih) - 32
sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product
sumf += d8[i] * (sumi * sc);
}
return d*sumf;
@@ -870,18 +876,31 @@ static __device__ __forceinline__ float vec_dot_q2_K_q8_1(
const block_q2_K * bq2_K = (const block_q2_K *) vbq + kbx;
// iqs is in 0,2..14, each 32 bytes of qs hold 4 q8_1 blocks (one per 2 bit shift)
const int bq8_offset = QR2_K * (iqs / QI8_1);
const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
const int iqs8 = iqs % QI8_1;
const uint8_t * scales = bq2_K->scales + scale_offset;
int v[VDR_Q2_K_Q8_1_MMVQ];
#pragma unroll
for (int j = 0; j < VDR_Q2_K_Q8_1_MMVQ; ++j) {
v[j] = get_int_b4(bq2_K->qs, iqs + j);
}
const int v = get_int_b4(bq2_K->qs, iqs);
int u[QR2_K];
// scale of shift i is byte 2*i + (iqs8 >= QI8_1/2) of the 8 bytes for this half
const int sc_shift = 8 * (iqs8 / (QI8_1/2));
int scales[2];
scales[0] = get_int_b4(bq2_K->scales, 2*(iqs / QI8_1) + 0) >> sc_shift;
scales[1] = get_int_b4(bq2_K->scales, 2*(iqs / QI8_1) + 1) >> sc_shift;
int u[VDR_Q2_K_Q8_1_MMVQ*QR2_K];
float d8[QR2_K];
#pragma unroll
for (int i = 0; i < QR2_K; ++ i) {
u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
for (int i = 0; i < QR2_K; ++i) {
#pragma unroll
for (int j = 0; j < VDR_Q2_K_Q8_1_MMVQ; ++j) {
u[VDR_Q2_K_Q8_1_MMVQ*i + j] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs8 + j);
}
d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
}
@@ -893,26 +912,42 @@ static __device__ __forceinline__ float vec_dot_q3_K_q8_1(
const block_q3_K * bq3_K = (const block_q3_K *) vbq + kbx;
// iqs is in 0,2..14, each 32 bytes of qs hold 4 q8_1 blocks (one per 2 bit shift)
const int bq8_offset = QR3_K * (iqs / (QI3_K/2));
const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2);
const int iqs8 = iqs % QI8_1;
const float d = bq3_K->d;
const int vl = get_int_b2(bq3_K->qs, iqs);
int vl[VDR_Q3_K_Q8_1_MMVQ];
int vh[VDR_Q3_K_Q8_1_MMVQ];
#pragma unroll
for (int j = 0; j < VDR_Q3_K_Q8_1_MMVQ; ++j) {
vl[j] = get_int_b2(bq3_K->qs, iqs + j);
vh[j] = get_int_b2(bq3_K->hmask, iqs8 + j) >> bq8_offset;
}
// invert the mask with ~ so that a 0/1 results in 4/0 being subtracted
const int vh = ~get_int_b2(bq3_K->hmask, iqs % (QI3_K/2)) >> bq8_offset;
// sub-block 8*(iqs/8) + 2*i + (iqs8 >= QI8_1/2) for shift i:
// 4 low bits in nibble iqs/8 of bytes, 2 high bits at bit 4*(iqs/8) + 2*(i/2) of bytes 8..11
const int sc_shift = 8 * (iqs8 / (QI8_1/2)) + 4 * (iqs / (QI3_K/2));
int scales[3];
#pragma unroll
for (int l = 0; l < 3; ++l) {
scales[l] = get_int_b2(bq3_K->scales, l) >> sc_shift;
}
int u[QR3_K];
int u[VDR_Q3_K_Q8_1_MMVQ*QR3_K];
float d8[QR3_K];
#pragma unroll
for (int i = 0; i < QR3_K; ++i) {
u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1);
#pragma unroll
for (int j = 0; j < VDR_Q3_K_Q8_1_MMVQ; ++j) {
u[VDR_Q3_K_Q8_1_MMVQ*i + j] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs8 + j);
}
d8[i] = __low2float(bq8_1[bq8_offset + i].ds);
}
return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8);
return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, scales, d, d8);
}
static __device__ __forceinline__ float vec_dot_q4_K_q8_1(
@@ -1025,17 +1060,25 @@ static __device__ __forceinline__ float vec_dot_q6_K_q8_1(
const int scale_offset = (QI6_K/4) * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/8);
const int vh_shift = 2 * ((iqs % (QI6_K/2)) / (QI6_K/4));
const int vl = get_int_b2(bq6_K->ql, iqs);
const int vh = get_int_b2(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4)) >> vh_shift;
int vl[VDR_Q6_K_Q8_1_MMVQ];
int vh[VDR_Q6_K_Q8_1_MMVQ];
#pragma unroll
for (int j = 0; j < VDR_Q6_K_Q8_1_MMVQ; ++j) {
vl[j] = get_int_b2(bq6_K->ql, iqs + j);
vh[j] = get_int_b2(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4) + j) >> vh_shift;
}
const int8_t * scales = bq6_K->scales + scale_offset;
int u[QR6_K];
int u[VDR_Q6_K_Q8_1_MMVQ*QR6_K];
float d8[QR6_K];
#pragma unroll
for (int i = 0; i < QR6_K; ++i) {
u[i] = get_int_b4(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1);
#pragma unroll
for (int j = 0; j < VDR_Q6_K_Q8_1_MMVQ; ++j) {
u[VDR_Q6_K_Q8_1_MMVQ*i + j] = get_int_b4(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1 + j);
}
d8[i] = __low2float(bq8_1[bq8_offset + 2*i].ds);
}