vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-05-09 11:00:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

vulkan: load each Q value once. optimize O reduction. more tuning

Browse source
This commit is contained in:
Jeff Bolz 2025-05-07 15:57:38 -05:00
parent c747227a57
commit 989bfb18fc
2 changed files with 19 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

6
ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file

@ -1590,7 +1590,7 @@ static void ggml_vk_wait_events(vk_context& ctx, std::vector<vk::Event>&& events
// number of rows/cols for flash attention shader
static constexpr uint32_t flash_attention_num_small_rows = 32;
static constexpr uint32_t scalar_flash_attention_num_small_rows = 4;
static constexpr uint32_t scalar_flash_attention_num_small_rows = 8;
static uint32_t get_fa_num_small_rows(bool scalar) {
return scalar ? scalar_flash_attention_num_small_rows : flash_attention_num_small_rows;
@ -1601,7 +1601,7 @@ static std::array<uint32_t, 2> fa_rows_cols(bool scalar, uint32_t D, uint32_t cl
// small rows, large cols
if (small_rows || scalar) {
return {get_fa_num_small_rows(scalar), 64};
return {get_fa_num_small_rows(scalar), 32};
}
// small cols to reduce register count
@ -1913,7 +1913,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
// D_split can't be larger than a subgroup because we use subgroupShuffle to reduce it.
// D_split can't be larger than the LSB of D divided by 4 due to vectorization in the shader.
const uint32_t D_lsb = D ^ (D & (D-1));
uint32_t D_split = std::min(std::min(device->subgroup_size, 16u), D_lsb / 4);
uint32_t D_split = std::min(std::min(device->subgroup_size, 8u), D_lsb / 4);
// mask dim1 is padded to 64, we rely on this to avoid clamping mask loads
GGML_ASSERT((GGML_KQ_MASK_PAD % rows_cols[0]) == 0);

31
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
View file

@ -105,6 +105,7 @@ ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const i
}
shared FLOAT_TYPE tmpsh[gl_WorkGroupSize.x];
shared vec4 tmpshv4[gl_WorkGroupSize.x];
shared float16_t masksh[Bc][Br];
shared vec4 Qf[Br][D / 4];
@ -168,13 +169,15 @@ void main() {
uint32_t q_offset = (iq2*p.nb02+iq3*p.nb03) / 4;
[[unroll]] for (uint32_t r = 0; r < Br; ++r) {
if (i * Br + r < N) {
[[unroll]] for (uint32_t d = 0; d < D_per_thread / 4; ++d) {
Qf[r][d * D_split + d_tid] = vec4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d * D_split + d_tid]) * p.scale;
}
[[unroll]] for (uint32_t idx = 0; idx < Br * D / 4; idx += gl_WorkGroupSize.x) {
uint32_t d = (idx + tid) % (D / 4);
uint32_t r = (idx + tid) / (D / 4);
if (r < Br && d < D / 4 &&
i * Br + r < N) {
Qf[r][d] = vec4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d]) * p.scale;
}
}
barrier();
vec4 Of[Br][D_per_thread / 4];
[[unroll]] for (uint32_t d = 0; d < D_per_thread / 4; ++d) {
@ -350,20 +353,18 @@ void main() {
[[unroll]] for (uint32_t d = 0; d < D_per_thread / 4; ++d) {
Of[r][d] = eMf * Of[r][d];
[[unroll]] for (uint32_t c = 0; c < 4; ++c) {
tmpsh[tid] = Of[r][d][c];
tmpshv4[tid] = Of[r][d];
barrier();
[[unroll]] for (int s = int(gl_WorkGroupSize.x) / 2; s >= D_split; s >>= 1) {
if (tid < s) {
Of[r][d][c] += tmpsh[tid + s];
tmpsh[tid] = Of[r][d][c];
}
barrier();
barrier();
[[unroll]] for (int s = int(gl_WorkGroupSize.x) / 2; s >= D_split; s >>= 1) {
if (tid < s) {
Of[r][d] += tmpshv4[tid + s];
tmpshv4[tid] = Of[r][d];
}
Of[r][d][c] = tmpsh[d_tid];
barrier();
}
Of[r][d] = tmpshv4[d_tid];
barrier();
}
}