CUDA: branchless Q4_K/Q5_K unpack to speed up mmvq, L2 prefetch on DGX Spark (#26705)

* Update Q4_K and Q5_K to use branchless computation, which stops the scale unpack being re-executed for every column in mmvq, improving perf at batch sizes > 1

* Gating the change off from DGX Spark due to no gain

* Adding prefetch gated to Spark, making branchless change in Q4_K and Q5_K general and modifying switch points based on latest perf data

* Guard the mmvq L2 prefetch against MUSA as well as HIP

* Define the mmvq L2 prefetch only under the Spark guard

* Update switch point for Q4_K to accommodate more models

* Remove stale comments

* Add block_size to ggml_cuda_type_traits and create a separate mmvq_should_prefetch function

* Rename block_size to bs for cleaner indentation

* Fix build error on non-Spark CUDA arch with appropriate conditional around new function added

---------

Co-authored-by: praneshgo <227579474+praneshgo@users.noreply.github.com>
This commit is contained in:
Pranesh Gonegandla
2026-09-07 19:36:58 +08:00
committed by GitHub
co-authored by praneshgo
parent 0cae43063c
commit 73ab7599b5
3 changed files with 100 additions and 20 deletions
+24
View File
@@ -976,6 +976,7 @@ template<>
struct ggml_cuda_type_traits<GGML_TYPE_F16> {
static constexpr int qk = 1;
static constexpr int qr = 1;
static constexpr int bs = sizeof(ggml_half);
};
template<>
@@ -983,6 +984,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q1_0> {
static constexpr int qk = QK1_0;
static constexpr int qr = QR1_0;
static constexpr int qi = QI1_0;
static constexpr int bs = sizeof(block_q1_0);
};
template<>
@@ -990,6 +992,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q2_0> {
static constexpr int qk = QK2_0;
static constexpr int qr = QR2_0;
static constexpr int qi = QI2_0;
static constexpr int bs = sizeof(block_q2_0);
};
template<>
@@ -997,6 +1000,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q4_0> {
static constexpr int qk = QK4_0;
static constexpr int qr = QR4_0;
static constexpr int qi = QI4_0;
static constexpr int bs = sizeof(block_q4_0);
};
template<>
@@ -1004,6 +1008,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q4_1> {
static constexpr int qk = QK4_1;
static constexpr int qr = QR4_1;
static constexpr int qi = QI4_1;
static constexpr int bs = sizeof(block_q4_1);
};
template<>
@@ -1011,6 +1016,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q5_0> {
static constexpr int qk = QK5_0;
static constexpr int qr = QR5_0;
static constexpr int qi = QI5_0;
static constexpr int bs = sizeof(block_q5_0);
};
template<>
@@ -1018,6 +1024,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q5_1> {
static constexpr int qk = QK5_1;
static constexpr int qr = QR5_1;
static constexpr int qi = QI5_1;
static constexpr int bs = sizeof(block_q5_1);
};
template<>
@@ -1025,6 +1032,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q8_0> {
static constexpr int qk = QK8_0;
static constexpr int qr = QR8_0;
static constexpr int qi = QI8_0;
static constexpr int bs = sizeof(block_q8_0);
};
template<>
@@ -1032,6 +1040,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_MXFP4> {
static constexpr int qk = QK_MXFP4;
static constexpr int qr = QR_MXFP4;
static constexpr int qi = QI_MXFP4;
static constexpr int bs = sizeof(block_mxfp4);
};
template<>
@@ -1039,6 +1048,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_NVFP4> {
static constexpr int qk = QK_NVFP4;
static constexpr int qr = QR_NVFP4;
static constexpr int qi = QI_NVFP4;
static constexpr int bs = sizeof(block_nvfp4);
};
template<>
@@ -1046,6 +1056,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q2_K> {
static constexpr int qk = QK_K;
static constexpr int qr = QR2_K;
static constexpr int qi = QI2_K;
static constexpr int bs = sizeof(block_q2_K);
};
template<>
@@ -1053,6 +1064,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q3_K> {
static constexpr int qk = QK_K;
static constexpr int qr = QR3_K;
static constexpr int qi = QI3_K;
static constexpr int bs = sizeof(block_q3_K);
};
template<>
@@ -1060,6 +1072,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q4_K> {
static constexpr int qk = QK_K;
static constexpr int qr = QR4_K;
static constexpr int qi = QI4_K;
static constexpr int bs = sizeof(block_q4_K);
};
template<>
@@ -1067,6 +1080,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q5_K> {
static constexpr int qk = QK_K;
static constexpr int qr = QR5_K;
static constexpr int qi = QI5_K;
static constexpr int bs = sizeof(block_q5_K);
};
template<>
@@ -1074,6 +1088,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_Q6_K> {
static constexpr int qk = QK_K;
static constexpr int qr = QR6_K;
static constexpr int qi = QI6_K;
static constexpr int bs = sizeof(block_q6_K);
};
template<>
@@ -1081,6 +1096,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ2_XXS> {
static constexpr int qk = QK_K;
static constexpr int qr = QR2_XXS;
static constexpr int qi = QI2_XXS;
static constexpr int bs = sizeof(block_iq2_xxs);
};
template<>
@@ -1088,6 +1104,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ2_XS> {
static constexpr int qk = QK_K;
static constexpr int qr = QR2_XS;
static constexpr int qi = QI2_XS;
static constexpr int bs = sizeof(block_iq2_xs);
};
template<>
@@ -1095,6 +1112,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ2_S> {
static constexpr int qk = QK_K;
static constexpr int qr = QR2_S;
static constexpr int qi = QI2_S;
static constexpr int bs = sizeof(block_iq2_s);
};
template<>
@@ -1102,6 +1120,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ3_XXS> {
static constexpr int qk = QK_K;
static constexpr int qr = QR3_XXS;
static constexpr int qi = QI3_XXS;
static constexpr int bs = sizeof(block_iq3_xxs);
};
template<>
@@ -1109,6 +1128,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ1_S> {
static constexpr int qk = QK_K;
static constexpr int qr = QR1_S;
static constexpr int qi = QI1_S;
static constexpr int bs = sizeof(block_iq1_s);
};
template<>
@@ -1116,6 +1136,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ1_M> {
static constexpr int qk = QK_K;
static constexpr int qr = QR1_M;
static constexpr int qi = QI1_M;
static constexpr int bs = sizeof(block_iq1_m);
};
template<>
@@ -1123,6 +1144,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ4_NL> {
static constexpr int qk = QK4_NL;
static constexpr int qr = QR4_NL;
static constexpr int qi = QI4_NL;
static constexpr int bs = sizeof(block_iq4_nl);
};
template<>
@@ -1130,6 +1152,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ4_XS> {
static constexpr int qk = QK_K;
static constexpr int qr = QR4_XS;
static constexpr int qi = QI4_XS;
static constexpr int bs = sizeof(block_iq4_xs);
};
template<>
@@ -1137,6 +1160,7 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ3_S> {
static constexpr int qk = QK_K;
static constexpr int qr = QR3_S;
static constexpr int qi = QI3_S;
static constexpr int bs = sizeof(block_iq3_s);
};
//////////////////////
+51 -4
View File
@@ -6,6 +6,35 @@
#include <cstdint>
#include <type_traits>
// only enabled on DGX Spark, where it is a gain on every type below. On the higher-bandwidth parts the kernel
// has little exposed latency left to hide and the extra requests cost more than they save.
// For perf data, see https://github.com/ggml-org/llama.cpp/pull/26705#issuecomment-5569335031
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
// returns true only for those quants that benefit from prefetch and false otherwise
static constexpr __host__ __device__ bool mmvq_should_prefetch(ggml_type type) {
switch (type) {
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q8_0:
case GGML_TYPE_MXFP4:
case GGML_TYPE_Q3_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
case GGML_TYPE_Q6_K:
case GGML_TYPE_IQ1_M:
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_IQ4_XS:
return true;
default:
return false;
}
}
static __device__ __forceinline__ void mmvq_prefetch_l2(const void * p) {
asm volatile("prefetch.global.L2 [%0];" :: "l"(p));
}
#endif
typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs);
static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) {
@@ -298,9 +327,6 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) {
return ne11 <= 4;
case GGML_TYPE_Q3_K:
return ne11 <= 6;
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
return ne11 <= 7;
default:
return ne11 <= MMVQ_MAX_BATCH_SIZE;
}
@@ -310,8 +336,9 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) {
case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
return ne11 <= 5;
case GGML_TYPE_Q5_K:
return ne11 <= 6;
case GGML_TYPE_Q6_K:
return ne11 <= 7;
default:
@@ -675,6 +702,26 @@ static __global__ void mul_mat_vec_q(
// x block quant index when casting the quants to int
const int kqs = vdr * (tid % (qi/vdr));
#if defined(__CUDA_ARCH__) && __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
const int kbx_pf = kbx + pf_dist*blocks_per_iter;
if (kbx_pf < blocks_per_row_x) {
#pragma unroll
for (int i = 0; i < rows_per_cuda_block; ++i) {
const size_t off = (size_t)(kbx_offset + i*stride_row_x + kbx_pf) * ggml_cuda_type_traits<type>::bs;
mmvq_prefetch_l2((const char *) vx + off);
if constexpr (has_fusion) {
if (use_gate) {
mmvq_prefetch_l2((const char *) vgate + off);
}
}
}
}
}
#endif
#pragma unroll
for (int j = 0; j < ncols_dst; ++j) {
#pragma unroll
+25 -16
View File
@@ -936,16 +936,20 @@ static __device__ __forceinline__ float vec_dot_q4_K_q8_1(
v[0] = q4[0];
v[1] = q4[4];
// branchless so nvcc can hoist this out of the ncols_dst loop
const uint16_t * scales = (const uint16_t *)bq4_K->scales;
const int j = bq8_offset/2;
const int jm = j & 1;
const uint32_t s0 = scales[jm + 0];
const uint32_t s2 = scales[jm + 2];
const uint32_t s4 = scales[jm + 4];
const uint32_t hi = (uint32_t) -(int32_t) (j >= 2);
uint16_t aux[2];
const int j = bq8_offset/2;
if (j < 2) {
aux[0] = scales[j+0] & 0x3f3f;
aux[1] = scales[j+2] & 0x3f3f;
} else {
aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2);
aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2);
}
aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi));
aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi));
const uint8_t * sc = (const uint8_t *)aux;
const uint8_t * m = sc + 2;
@@ -981,16 +985,21 @@ static __device__ __forceinline__ float vec_dot_q5_K_q8_1(
vh[0] = qh[0] >> bq8_offset;
vh[1] = qh[4] >> bq8_offset;
// same as q4_K
const uint16_t * scales = (const uint16_t *)bq5_K->scales;
const int j = bq8_offset/2;
const int jm = j & 1;
const uint32_t s0 = scales[jm + 0];
const uint32_t s2 = scales[jm + 2];
const uint32_t s4 = scales[jm + 4];
const uint32_t hi = (uint32_t) -(int32_t) (j >= 2);
uint16_t aux[2];
const int j = bq8_offset/2;
if (j < 2) {
aux[0] = scales[j+0] & 0x3f3f;
aux[1] = scales[j+2] & 0x3f3f;
} else {
aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2);
aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2);
}
aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi));
aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi));
const uint8_t * sc = (const uint8_t *)aux;
const uint8_t * m = sc + 2;