vulkan: further optimize mul_mat_vec using larger loads #10387
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add some early returns for nonexistent rows in mul_mat_vec shaders. These can only be hit when dispatching a 2D grid of workgroups. Fix the logic for the 2D grid of workgroups to round up. Enable the pipeline robustness extension if it's available, and use it to disable robustness for these pipelines. The instructions to do the bounds checking contend for the same ALU resources as the bit twiddling dequant instructions.
In Vulkan it's not possible to cast pointer types, so instead you have to declare an aliased binding for the memory with a different type. This commit adds aliases for the quant formats using 16b ints, and in a few places where the struct size is a multiple of 4 also using 32b ints. Currently only q4_k's aliases are used, but others will be used in subsequent commits.
Similar to the optimization I did in q4_k recently, this vectorizes some loads and reduces the number of bit twiddling instructions.
Add vec4 dequantization functions, and use them to do K=8 per iteration in mul_mat_vec. This uses 16b loads for the quant values and 128b loads for B which helps reduce the load on the memory system. The K_PER_ITER==2 logic is still there, just for F16/F32, and really only because they support unaligned sizes. Tweak the num_iters/unrolling logic to be simpler and catch a couple missed unrolling opportunities.
4 tasks
|
Wow this is really nice. Here are my numbers on the RX 570 for this PR. Before:
After:
There isn't that much change with the K-quants but Q8 now runs faster than the old Q4! 🎉 Bonus results:
As discussed in #10296 calculating 8 rows at a time helps with the 570. From what I've read online AMD GCN apparently likes longer running shaders. If I optimize the delta calculation using the code below I nearly get 20t/s - and that's all I can come up with for tonight. ------------ ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.comp ------------
index 5fc1ba4a..7f3cb33a 100644
@@ -30,9 +30,9 @@ vec2 dequantize(uint ib, uint iqs, uint a_offset) {
return (vec2(vui & 0xF, vui >> 4) - 8.0f) * d;
}
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
- const float d = float(data_a_packed16[a_offset + ib].d);
+ //const float d = float(data_a_packed16[a_offset + ib].d);
const uint vui = uint(data_a_packed16[a_offset + ib].qs[iqs/2]);
- return (vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, (vui >> 12) & 0xF) - 8.0f) * d;
+ return (vec4(vui & 0xF, (vui >> 4) & 0xF, (vui >> 8) & 0xF, (vui >> 12) & 0xF) - 8.0f);
}
#endif
@@ -95,10 +95,10 @@ vec2 dequantize(uint ib, uint iqs, uint a_offset) {
return vec2(int(data_a[a_offset + ib].qs[iqs]), int(data_a[a_offset + ib].qs[iqs + 1])) * d;
}
vec4 dequantize4(uint ib, uint iqs, uint a_offset) {
- const float d = float(data_a_packed16[a_offset + ib].d);
+ //const float d = float(data_a_packed16[a_offset + ib].d);
uint32_t v0 = data_a_packed16[a_offset + ib].qs[iqs/2];
uint32_t v1 = data_a_packed16[a_offset + ib].qs[iqs/2 + 1];
- return vec4(int8_t(v0 & 0xFF), int8_t((v0 >> 8) & 0xFF), int8_t(v1 & 0xFF), int8_t((v1 >> 8) & 0xFF)) * d;
+ return vec4(int8_t(v0 & 0xFF), int8_t((v0 >> 8) & 0xFF), int8_t(v1 & 0xFF), int8_t((v1 >> 8) & 0xFF));
}
#endif
------------- ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec.comp -------------
index 00807a06..6a15bce5 100644
@@ -73,16 +73,22 @@ void iter(inout FLOAT_TYPE temp[NUM_ROWS], const uint first_row, const uint num_
#if K_PER_ITER == 8
const vec4 v = dequantize4(ib, iqs, a_offset);
const vec4 v2 = dequantize4(ib, iqs+(4/QUANT_R), a_offset);
+ FLOAT_TYPE rowtmp = 0;
// matrix multiplication
- temp[n] = fma(FLOAT_TYPE(v.x), b0, temp[n]);
- temp[n] = fma(FLOAT_TYPE(v.y), b1, temp[n]);
- temp[n] = fma(FLOAT_TYPE(v.z), b2, temp[n]);
- temp[n] = fma(FLOAT_TYPE(v.w), b3, temp[n]);
- temp[n] = fma(FLOAT_TYPE(v2.x), b4, temp[n]);
- temp[n] = fma(FLOAT_TYPE(v2.y), b5, temp[n]);
- temp[n] = fma(FLOAT_TYPE(v2.z), b6, temp[n]);
- temp[n] = fma(FLOAT_TYPE(v2.w), b7, temp[n]);
+ rowtmp = FLOAT_TYPE(v.x) * b0;
+ rowtmp = fma(FLOAT_TYPE(v.y), b1, rowtmp);
+ rowtmp = fma(FLOAT_TYPE(v.z), b2, rowtmp);
+ rowtmp = fma(FLOAT_TYPE(v.w), b3, rowtmp);
+ rowtmp = fma(FLOAT_TYPE(v2.x), b4, rowtmp);
+ rowtmp = fma(FLOAT_TYPE(v2.y), b5, rowtmp);
+ rowtmp = fma(FLOAT_TYPE(v2.z), b6, rowtmp);
+ rowtmp = fma(FLOAT_TYPE(v2.w), b7, rowtmp);
+#if defined(DATA_A_Q4_0) || defined(DATA_A_Q8_0)
+ const float d = float(data_a_packed16[a_offset + ib].d);
+ rowtmp *= d;
+#endif
+ temp[n] += rowtmp;
#else
const vec2 v = dequantize(ib, iqs, a_offset); |
|
I can split out the scale factor multiply in a followon change if there's interest. I haven't seen cases where it's a bottleneck on RTX 4070, but it couldn't hurt. |
|
I'll add some model benchmarks later, but here's the test-backend-ops perf results. Looks like a significant step forward, most effective on Nvidia (not surprising since that's what you're working on), but also significantly positive on AMD. I'm not sure why the difference between AMD Vulkan and ROCm is larger than Nvidia Vulkan and CUDA. My Intel A770 also really liked the changes in q4_0, q4_1, q5_0 and q5_1, but not q8_0 and q6_k, not sure what's going on there.
|
|
Here's the results for Llama3 8B q4_0 and q4_k_s and Mistral Nemo q6_k. Looks good. The difference to CUDA/ROCm is larger than I remembered. Intel didn't like the change in some areas, so at a later point we might want to make the
|
0cc4m
approved these changes
Nov 20, 2024
arthw
pushed a commit
to arthw/llama.cpp
that referenced
this pull request
Dec 20, 2024
) * vulkan: Use pipeline_robustness to disable robustness in mul_mat_vec. Add some early returns for nonexistent rows in mul_mat_vec shaders. These can only be hit when dispatching a 2D grid of workgroups. Fix the logic for the 2D grid of workgroups to round up. Enable the pipeline robustness extension if it's available, and use it to disable robustness for these pipelines. The instructions to do the bounds checking contend for the same ALU resources as the bit twiddling dequant instructions. * vulkan: Add GLSL structure aliases for quant types to allow larger loads In Vulkan it's not possible to cast pointer types, so instead you have to declare an aliased binding for the memory with a different type. This commit adds aliases for the quant formats using 16b ints, and in a few places where the struct size is a multiple of 4 also using 32b ints. Currently only q4_k's aliases are used, but others will be used in subsequent commits. * vulkan: use larger loads in q5_k and q6_k shaders. Similar to the optimization I did in q4_k recently, this vectorizes some loads and reduces the number of bit twiddling instructions. * vulkan: use larger K step per iteration in mul_mat_vec. Add vec4 dequantization functions, and use them to do K=8 per iteration in mul_mat_vec. This uses 16b loads for the quant values and 128b loads for B which helps reduce the load on the memory system. The K_PER_ITER==2 logic is still there, just for F16/F32, and really only because they support unaligned sizes. Tweak the num_iters/unrolling logic to be simpler and catch a couple missed unrolling opportunities.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There are a few things in this PR:
Some performance results on RTX 4070. Note that the directed tests tend to fit in L2 on this system so they are much less memory-limited and don't reflect the real performance gain in networks (which tend to be framebuffer-limited, and benefit less from these optimizations). Also, this "Q4_K" network I tested uses all of Q4_K/Q5_K/Q6_K and is showing some benefit from the optimizations on those shaders (i.e. it's not all from the robustness change).