From 2744e48d1e88d7416224e92bf58d4ffedda4a205 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Sat, 17 Oct 2026 19:05:30 +0530 Subject: [PATCH] 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 --- ggml/src/ggml-cuda/mmvq.cu | 35 ++++++-- ggml/src/ggml-cuda/vecdotq.cuh | 149 +++++++++++++++++++++------------ 2 files changed, 125 insertions(+), 59 deletions(-) diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 8068ba364..930f1adc4 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -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 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); diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index f65fe530e..14c1737c0 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -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); }