Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
am17an merged 10 commits into from
Sep 7, 2026
Merged

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

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
10d708a
Update Q4_K and Q5_K to use branchless computation, which stops the s…
praneshgo Aug 7, 2026
b686848
Gating the change off from DGX Spark due to no gain
praneshgo Aug 13, 2026
6754cee
Adding prefetch gated to Spark, making branchless change in Q4_K and …
praneshgo Aug 27, 2026
9b9cf68
Guard the mmvq L2 prefetch against MUSA as well as HIP
praneshgo Aug 31, 2026
63daa18
Define the mmvq L2 prefetch only under the Spark guard
praneshgo Aug 31, 2026
99e1393
Update switch point for Q4_K to accommodate more models
praneshgo Sep 2, 2026
2097c8c
Remove stale comments
praneshgo Sep 4, 2026
9c2afe6
Add block_size to ggml_cuda_type_traits and create a separate mmvq_sh…
praneshgo Sep 7, 2026
22c832a
Rename block_size to bs for cleaner indentation
praneshgo Sep 7, 2026
aa17a1b
Fix build error on non-Spark CUDA arch with appropriate conditional a…
praneshgo Sep 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 40 additions & 4 deletions ggml/src/ggml-cuda/mmvq.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,23 @@
#include <cstdint>
#include <type_traits>

// block size in bytes; 0 compiles the prefetch out for that type
template <ggml_type type> struct mmvq_pf { static constexpr int bytes = 0; };
template <> struct mmvq_pf<GGML_TYPE_Q4_0> { static constexpr int bytes = sizeof(block_q4_0); };
template <> struct mmvq_pf<GGML_TYPE_Q8_0> { static constexpr int bytes = sizeof(block_q8_0); };
template <> struct mmvq_pf<GGML_TYPE_Q3_K> { static constexpr int bytes = sizeof(block_q3_K); };
template <> struct mmvq_pf<GGML_TYPE_Q4_K> { static constexpr int bytes = sizeof(block_q4_K); };
template <> struct mmvq_pf<GGML_TYPE_Q5_K> { static constexpr int bytes = sizeof(block_q5_K); };
template <> struct mmvq_pf<GGML_TYPE_Q6_K> { static constexpr int bytes = sizeof(block_q6_K); };
template <> struct mmvq_pf<GGML_TYPE_IQ4_XS> { static constexpr int bytes = sizeof(block_iq4_xs); };

@am17an am17an Sep 7, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

block_size should be added to ggml_cuda_type_traits, create a separate mmvq_should_prefetch function for excluding q2_k

@praneshgo praneshgo Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Acknowledged. I will commit the necessary changes for review in a short bit.

@praneshgo praneshgo Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Made the necessary changes in the latest commit. Please review.

// Q2_K is left out: prefetching does not raise its L2 hit rate, so the requests only add pressure

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
static __device__ __forceinline__ void mmvq_prefetch_l2(const void * p) {
asm volatile("prefetch.global.L2 [%0];" :: "l"(p));
}
#endif

@am17an am17an Sep 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

leave a comment or link to PR discussion where DGX sparks numbers are displayed

@praneshgo praneshgo Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a comment with link to the PR comment for DGX Spark results


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) {
Expand Down Expand Up @@ -298,9 +315,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;
}
Expand All @@ -310,8 +324,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:
Expand Down Expand Up @@ -663,6 +678,27 @@ 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. only pays where the kernel is
// latency-bound rather than bandwidth-bound, which on this part it is
if constexpr (mmvq_pf<type>::bytes > 0) {
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) * mmvq_pf<type>::bytes;
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
Expand Down
41 changes: 25 additions & 16 deletions ggml/src/ggml-cuda/vecdotq.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;

Expand Down
Loading