CUDA: branchless Q4_K/Q5_K unpack to speed up mmvq, L2 prefetch on DGX Spark - #26705
Conversation
|
@gaugarg-nv @ORippler can you please review this PR? Thanks. |
This should be determined/validated/invalidated based on nsight compute. Feel free to ping me offline if you need help interpreting. Intuitively, hoisting=less work, and less work should almost never incur slowdowns |
|
I tested PR #26705 on Qwen3.8-27B, using my custom 5.01 BPW iMatrix/NVFP4 hybrid quant on an RTX PRO 4000 Blackwell SFF 24 GB. On the same 40K-context benchmark, current master averaged 41.76 tok/s, while master + #26705 reached 43.11 tok/s, a +3.22% end-to-end improvement. Acceptance rate and generated output were identical, so in my test this looks like a clean kernel-level speedup with no sampling-quality trade-off |
|
A follow-up to my earlier 40K result: I kept commit 7061f25 from this PR in the final runtime and released the model used for the test. The isolated five-run A/B was: master: 41.76 tok/s Acceptance remained 78.73% and all response hashes matched, so this was a kernel-level gain rather than a different speculative continuation. The result was smaller than the RTX 5090 number reported here, which makes sense because only part of my mixed quant uses Q5_K. The remaining tensors are mainly NVFP4, Q6_K and Q8_0. In the later 256K production-profile sweep, the #26001 + #26048 + #26705 bundle reached 45.866 tok/s versus 45.422 for clean b10454, or +0.98%. I am keeping that number separate because it used a different runtime gate. Hardware: RTX PRO 4000 Blackwell SFF, 24,467 MiB, sm120a, CUDA 12.9.86, driver 610.57.04. Qwen3.8-27B, 5.01 BPW iMatrix/NVFP4 hybrid, embedded MTP, Q4_0 target KV. Full measurements and arguments: GGUF: |
I looked at q4k_n5_base (no changes) and q4k_n5_brchless (changes made in this PR) ncu-rep files as an example to better identify reasons for no gain in perf on DGX Spark. n here denotes number of matmul cols (also same as ne11). For this testing, I did not use the CUDA_ARCH == GGML_CUDA_CC_DGX_SPARK conditional used in the PR, so as to see the true branchless behavior on DGX Spark. The number of executed instructions go down from ~32 mil (base) to ~23 mil (branchless) as expected. However, the number of elapsed cycles as per the ncu-rep files largely remained the same: 465,637 (base) → 453,280 (branchless) in spite of close to 30% fewer instructions. That 2.7% difference is also about the launch-to-launch spread I see on the same binary, so it is within noise. The issue slots busy % (% of peak sustained active) reduced from ~36% to ~26% and warp cycles per issued instruction go from 9.59 (base) to 14.69 (branchless). This means kernel issues less often and waits longer per instruction. Even if we free some cycles by optimizing the unpack, they just become part of the stall. |
So if we are not regressing on DGX/RTX Spark, we should enable this as it will increase power-efficiency of the system (stalling is cheaper than issuing an instruction, aka every instruction counts) |
I am in fact seeing a very slight decline in perf if we add this change on DGX Spark (around 2-3%). I ran a few experiments by locking and unlocking clocks, checked by enabling/disabling graphs and checked if coalescing is responsible for this drop, seems like none of them explain it completely. So, until this is figured out, it may be better to keep this change gated for DGX Spark. |
7061f25 to
3d77a86
Compare
|
To avoid the dip due to branchless change alone and ensure an overall gain, we make Spark use prefetch. This in our results ensured increased decode throughput across several quant types which can be evident from the performance numbers below. Prefetch while improving decode throughput on DGX Spark, seemed to negatively contributed in case of GPUs with larger bandwidth like RTX 5090. We analyzed ncu traces with prefetch enabled on RTX 5090, profiling Q8_0 on Llama-3.1-8B across ne11 1-8, with Q2_K as a control:
Improvement in mmvq performance means that switch points in #26079 needs changes, which are also done in this PR. |
gaugarg-nv
left a comment
There was a problem hiding this comment.
The change looks good to me. Can you see how this impacts SpecDec perf, specifically measure MTP and Dflash2 with Qwen3.6-35 and Qwen3.8-27B?
| // Q2_K is left out: prefetching does not raise its L2 hit rate, so the requests only add pressure | ||
|
|
||
| static __device__ __forceinline__ void mmvq_prefetch_l2(const void * p) { | ||
| #if !defined(GGML_USE_HIP) |
There was a problem hiding this comment.
We should also add check for MUSA.
There was a problem hiding this comment.
Good catch. Added the MUSA check in the latest commit.
…cale unpack being re-executed for every column in mmvq, improving perf at batch sizes > 1
…Q5_K general and modifying switch points based on latest perf data
| // Q2_K is left out: prefetching does not raise its L2 hit rate, so the requests only add pressure | ||
|
|
||
| static __device__ __forceinline__ void mmvq_prefetch_l2(const void * p) { | ||
| #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) |
There was a problem hiding this comment.
Non-spark CUDA build fails too with warning #177-D: function "mmvq_prefetch_l2" was declared but never referenced.
Moving the helper under the same NVIDIA Spark architecture guard would remove the warning and also avoid exposing PTX to HIP/MUSA.
There was a problem hiding this comment.
Checking the errors. Will fix soon.
There was a problem hiding this comment.
Updated the code by defining the mmvq L2 prefetch only under the Spark guard as per the suggestion. This should ideally fix the build warning/failure.
134360d to
63daa18
Compare
Measured the speculative decoding impact as asked. End-to-end decode throughput through Performance
Update: the RTX 5090 numbers below were measured with Q4_K switch point <= 6. That switch point is now back to 5 — see comment-5507374664 for the updated 5090 table. GB10 and 4090 are unaffected, as the change is inside the Blackwell block only. RTX 5090
DGX Spark GB10
RTX 4090
The three architectures do not gain in the same place, and each one tracks its own switch points. Blackwell peaks at Ada is flat through 4 and jumps at 8. The per-type sweeps put the branchless gain at +0.7 to +2.7% at batch 4 but +12 to +24% at batch 8, so nothing shows until verification reaches the top of the range. DFlash2 runs there by default - its trained block size is 8, so that is where it verifies unless the draft length is capped below it. GB10 climbs steadily and never jumps, since no switch point moves on that part. An isolation build with only the branchless unpack gives
mmvq was the default to begin with, for @gaugarg-nv please review. |
|
Do I understand correctly that ne11 == draft_length + 1? What is causing perf regression in dflash2 for ne11 =6? |
|
On RTX 5090, after moving the Q4_K switch point back to 5 (was 6 earlier in the PR), the 7% dip seen in case of Qwen3.8-27B(dflash2) now becomes close to 0. Perf numbers on other tested HW, including DGX Spark and 4090, remain unchanged and are mentioned in the description as well as comment-5492301911 End-to-end decode throughput, this PR vs master, measured through
Detailed Performance results
RTX 5090
Mentioning DGX Spark and 4090 results, which are the same as those from comment-5492301911, so that all data is in one place. DGX Spark GB10
RTX 4090
|
|
Performance improvements are solid for Nvidia GPUs. |
|
@IMbackK Possible for you test it on AMD GPUs? |
| case GGML_TYPE_Q4_K: | ||
| case GGML_TYPE_Q5_K: | ||
| return ne11 <= 7; | ||
| // branchless unpack pushes Q4_K and Q5_K past the cap, so they now take the default |
There was a problem hiding this comment.
this kind of comment should be avoided
There was a problem hiding this comment.
Removed this and another comment of the same kind in the latest commit.
|
Its perf positve on HIP Details
|
Thanks @IMbackK |
|
Should have looked here, but found this independently a few days ago. My version Pyrolistical@ecd06de TG/s results over master on rocm r9700 running qwen3.8 27b q4_k_xl:
|
|
@ggml-org/ggml-cuda Could I please get another review? |
| 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); }; |
There was a problem hiding this comment.
block_size should be added to ggml_cuda_type_traits, create a separate mmvq_should_prefetch function for excluding q2_k
There was a problem hiding this comment.
Acknowledged. I will commit the necessary changes for review in a short bit.
There was a problem hiding this comment.
Made the necessary changes in the latest commit. Please review.
| #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 |
There was a problem hiding this comment.
leave a comment or link to PR discussion where DGX sparks numbers are displayed
There was a problem hiding this comment.
Added a comment with link to the PR comment for DGX Spark results
|
Pulling all the perf sweeps done of DGX Spark over multiple models, quants and HW in one place for reference. DGX Spark Perf
DGX Spark GB10 - Llama-3.2-3B
DGX Spark GB10 - Llama-3.1-8B
DGX Spark GB10 - gemma-4-12B-it
DGX Spark GB10 - Qwen3.6-27B
DGX Spark GB10 - few other models
DGX Spark GB10 - more quantsThis table measures the gain/loss when a quant type is allowed to use prefetch, with every other change in the PR added and kept as-is. This is used to determine which quants benefit from prefetch and which do not. DGX Spark GB10 - Llama-3.2-3B
DGX Spark GB10 - Llama-3.1-8B
DGX Spark GB10 - gemma-4-12B-it
DGX Spark GB10 - Qwen3.6-27B
|
…ould_prefetch function
|
@am17an |
| static constexpr int qk = QK_K; | ||
| static constexpr int qr = QR3_S; | ||
| static constexpr int qi = QI3_S; | ||
| static constexpr int block_size = sizeof(block_iq3_s); |
There was a problem hiding this comment.
you can probably name this bs and keep the alignment the same
There was a problem hiding this comment.
Renamed block_size to bs for cleaner indentation in the latest commit.
@am17an can you please reapprove the PR? Thank you.
|
Also please re-name the PR before merging to something concise but informative. |
…round new function added
Updated the PR title to be concise and informative. |
…05 commits, 4 reworks) [unreviewed]
…X Spark (ggml-org#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>
…X Spark (ggml-org#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>
Overview
The Q4_K/Q5_K scale unpack branches on a runtime value to pick between two layouts. nvcc can't fold that, so it predicates it and then repeats the whole thing for every column of the ncols_dst loop instead of doing it once. Swapping the branch for a mask select fixes it: predicated instructions at ncols_dst=8 go from 226 to 34, and the Q4_K kernel at 16 columns from 2256 to 1648. None of the other quant types have this problem, they already unpack unconditionally (Q3_K does the same 6-bit layout with plain shifts).
We see perf gains from batch size >=4 on RTX 5090 and >=5 on RTX 4090 (more details in the Performance drop down), however we do not see a gain with DGX Spark with branchless alone, and in fact a small regression at batch sizes 4-6. The most probable reason for this can be that the memory bandwidth on DGX Spark is much lower than that of RTX 4090 or 5090, and this means even with this optimization there is no effective time to save, since in the original code the unpack was already running while the kernel waits on the weight loads and so was not adding to the total time.
To avoid this dip and ensure an overall gain, we make Spark use prefetch. This in our results ensured increased decode throughput across several quant types which can be evident from the performance numbers below. Prefetch while improving decode throughput on DGX Spark, seemed to negatively contributed in case of GPUs with larger bandwidth like RTX 5090.
We analyzed ncu traces with prefetch enabled on RTX 5090, profiling Q8_0 on Llama-3.1-8B across ne11 1-8, with Q2_K as a control:
Improvement in mmvq performance means that switch points in #26079 needs changes, which are also done in this PR.
Performance
Decode throughput in tokens/s, measured with
llama-bench -p <n> -n 0 -embd 1 -r 30 -o csvon single-type--purequantisations, so each number reflects the kernel under test rather than a model's tensor-type mix.The mvq -> MMQ switch points are disabled in every build here, so batch sizes 1-8 all reach
mul_mat_vec_q. Without that, master routes K-quants to MMQ above a per-architecture cap and any change to the vector kernel measures as exactly zero above it.branchless + prefetcharm, withmul_mat_vec_qdisabled so everything runs MMQMMQ is built from the shipping source but is unaffected by it:
mmq.cuhreferences neither the modified vector dot products nor the prefetch helper, so that row measures unmodified MMQ and shows where the crossover sits.The two
vs basecolumns are that arm againstbaseat the same batch size. MMQ is a different kernel rather than a variant of base, so it is shown in absolute terms only; comparing it against theshipcolumn is what locates the mvq -> MMQ crossover.Each cell is the mean of two ABBA-ordered passes at that batch size. No cell is a median or an average across types, models or batches.
Q2_K doubles as a control. Neither change can reach it - it is excluded from the prefetch by
mmvq_pf<GGML_TYPE_Q2_K>::bytes == 0, and branchless rewrites only the Q4_K/Q5_K scale unpack - so its two percentage columns should read zero, and what they actually read is the measurement floor for the table around them.DGX Spark GB10
sm_121, ~273 GB/s unified LPDDR5, driver 610.43.02, CUDA 13.3.
DGX Spark GB10 - Llama-3.2-3B
DGX Spark GB10 - Llama-3.1-8B
DGX Spark GB10 - gemma-4-12B-it
DGX Spark GB10 - Qwen3.6-27B
RTX 5090
sm_120, 1792 GB/s, driver 620.23, CUDA 12.9.
RTX 5090 - Llama-3.2-3B
RTX 5090 - Llama-3.1-8B
RTX 5090 - gemma-4-12B-it
RTX 5090 - Qwen3.6-27B
RTX 4090
sm_89, 1008 GB/s, driver 591, CUDA 12.9.
RTX 4090 - Llama-3.2-3B
RTX 4090 - Llama-3.1-8B
RTX 4090 - gemma-4-12B-it
RTX 4090 - Qwen3.6-27B
Q6_K, Q8_0 not measured here - they exceed this card's memory.
Requirements