From e85961a570614e208161eafec6c2ed2d1b21b603 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Tue, 20 Jan 2026 12:28:28 +0100 Subject: [PATCH 01/30] vulkan: use coopmat for flash attention p*v matrix multiplication --- .../vulkan-shaders/flash_attn_cm1.comp | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 0eb50fe58f9d..4661b82cbe33 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -275,6 +275,107 @@ void main() { Lf[r] = eMf[r]*Lf[r]; } +#if 1 + const uint psh_stride = Br / 4; // No space for padding, maybe increase Qf size? + + // Calculate and store Pf in Qf + [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { + const uint col = c * cols_per_iter + col_tid; + + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { + const uint row = tile_row(r); + f16vec4 Pfvec; + [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { + const float16_t Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); + Lf[r + vec_idx] += Pf; + Pfvec[vec_idx] = Pf; + } + Qf[col * psh_stride + row / 4] = Pfvec; + } + } + + const uint num_hsv_tiles = HSV / MatBc / row_split; + + // Each subgroup handles HSV/4 columns + [[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) { + const uint hsv_offset = (hsv_tile * row_split + gl_SubgroupID) * 16; + + SfMat = coopmat(0); + +#if BLOCK_SIZE > 1 + // Preload V tiles for [Bc, 16 * num subgroups] + const uint v_rows = Bc; + const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups + const uint v_total = v_rows * v_cols; + const uint v_loads_per_thread = v_total / gl_WorkGroupSize.x; + + const uint vsh_stride = v_cols; + + [[unroll]] for (uint32_t i = 0; i < v_loads_per_thread; ++i) { + const uint idx = i * gl_WorkGroupSize.x + tid; + const uint row = idx / v_cols; + const uint col = idx % v_cols; + + const uint v_row = j * Bc + row; + const uint v_col = hsv_tile * MatBc * row_split + col * 4; + + const uint coord = v_row * v_stride * BLOCK_SIZE + v_col; + const uint ib = coord / BLOCK_SIZE; + const uint iqs = coord % BLOCK_SIZE; + + ksh[row * vsh_stride + col] = f16vec4(dequantize4(ib, iqs, v_offset, BINDING_IDX_V)); + } + + barrier(); +#endif + + [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { + coopMatLoad(KMat, Qf, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); + +#if BLOCK_SIZE > 1 + const uint v_tile_offset = bc_chunk * MatBr * v_cols + gl_SubgroupID * (MatBc / 4); + coopMatLoad(QMat, ksh, v_tile_offset, vsh_stride, gl_CooperativeMatrixLayoutRowMajor); +#else + // F16 values can be loaded directly from global memory + const uint v_tile_row = j * Bc + bc_chunk * MatBc; + const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4; + coopMatLoad(QMat, data_vv4, v_tile_offset, v_stride / 4, gl_CooperativeMatrixLayoutRowMajor); +#endif + + SfMat = coopMatMulAdd(KMat, QMat, SfMat); + } + + // Store SfMat to sfsh and load into Of + const uint osh_stride = row_split * MatBc + 8; + const uint o_offset = gl_SubgroupID * MatBc; + coopMatStore(SfMat, sfsh, o_offset, osh_stride, gl_CooperativeMatrixLayoutRowMajor); + + barrier(); + + const uint hsv_per_tile = row_split * MatBc; + const uint hsv_base = hsv_tile * hsv_per_tile; + const uint d_values_per_tile = hsv_per_tile / 4 / D_split; + + const uint d_start = hsv_tile * d_values_per_tile; + const uint d_end = d_start + d_values_per_tile; + + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + const uint row = tile_row(r); + + [[unroll]] for (uint32_t d = d_start; d < d_end; ++d) { + const uint hsv_col = 4 * (d * D_split + d_tid); + const uint local_hsv = hsv_col - hsv_base; + + Of[r][d] += ACC_TYPEV4( + sfsh[row * osh_stride + local_hsv], + sfsh[row * osh_stride + local_hsv + 1], + sfsh[row * osh_stride + local_hsv + 2], + sfsh[row * osh_stride + local_hsv + 3] + ); + } + } + } +#else [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { continue; @@ -298,6 +399,7 @@ void main() { } } } +#endif barrier(); } From 7c46dbf190c21cffe5813ee996d7be715e5b0876 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Wed, 21 Jan 2026 09:53:02 +0100 Subject: [PATCH 02/30] fix P loading issue --- .../vulkan-shaders/flash_attn_cm1.comp | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 4661b82cbe33..d2b71faea063 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -279,18 +279,20 @@ void main() { const uint psh_stride = Br / 4; // No space for padding, maybe increase Qf size? // Calculate and store Pf in Qf - [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - const uint col = c * cols_per_iter + col_tid; - - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { - const uint row = tile_row(r); - f16vec4 Pfvec; - [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { - const float16_t Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); - Lf[r + vec_idx] += Pf; - Pfvec[vec_idx] = Pf; + if (d_tid == 0) { // All other threads have the same P values + [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { + const uint col = c * cols_per_iter + col_tid; + + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { + const uint row = tile_row(r); + f16vec4 Pfvec; + [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { + const float16_t Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); + Lf[r + vec_idx] += Pf; + Pfvec[vec_idx] = Pf; + } + Qf[col * psh_stride + row / 4] = Pfvec; } - Qf[col * psh_stride + row / 4] = Pfvec; } } From b58acd0bdd0bd5160ab0a905abe4d80c00293286 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Wed, 21 Jan 2026 10:31:23 +0100 Subject: [PATCH 03/30] fix barrier position --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index d2b71faea063..f9b1185235ed 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -327,9 +327,9 @@ void main() { ksh[row * vsh_stride + col] = f16vec4(dequantize4(ib, iqs, v_offset, BINDING_IDX_V)); } +#endif barrier(); -#endif [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { coopMatLoad(KMat, Qf, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); From 379608fe5b91042c9fbe0e8ba6ef0b6121f7584a Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Wed, 21 Jan 2026 11:11:40 +0100 Subject: [PATCH 04/30] remove reduction that is no longer needed --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index f9b1185235ed..7e48709a706f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -453,6 +453,7 @@ void main() { barrier(); } +#if 0 [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { @@ -470,6 +471,7 @@ void main() { barrier(); } } +#endif // If there is split_k, then the split_k resolve shader does the final // division by L. Store the intermediate O value and per-row m and L values. From 4cfc48d4e7f1997441de5bf17472851bd09fe144 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Wed, 21 Jan 2026 13:03:07 +0100 Subject: [PATCH 05/30] move max thread reduction into loop --- .../vulkan-shaders/flash_attn_cm1.comp | 95 ++++++------------- 1 file changed, 27 insertions(+), 68 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 7e48709a706f..546b45568109 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -259,11 +259,25 @@ void main() { } float Moldf = Mf[r]; + tmpsh[tid] = rowmaxf; + // Compute max across the row + barrier(); + [[unroll]] for (int s = int(gl_WorkGroupSize.x / row_split) / 2; s >= D_split; s >>= 1) { + rowmaxf = max(rowmaxf, tmpsh[tid ^ s]); + barrier(); + tmpsh[tid] = rowmaxf; + barrier(); + } + rowmaxf = tmpsh[d_tid + row_tid * threads_per_rowgroup]; + barrier(); + // M = max(rowmax, Mold) // P = e^(S - M) // eM = e^(Mold - M) Mf[r] = max(rowmaxf, Moldf); eMf[r] = exp(Moldf - Mf[r]); + + Lf[r] = eMf[r]*Lf[r]; } [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { @@ -271,26 +285,23 @@ void main() { Of[r][d] = ACC_TYPE(eMf[r]) * Of[r][d]; } } - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Lf[r] = eMf[r]*Lf[r]; - } #if 1 const uint psh_stride = Br / 4; // No space for padding, maybe increase Qf size? // Calculate and store Pf in Qf - if (d_tid == 0) { // All other threads have the same P values - [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - const uint col = c * cols_per_iter + col_tid; - - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { - const uint row = tile_row(r); - f16vec4 Pfvec; - [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { - const float16_t Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); - Lf[r + vec_idx] += Pf; - Pfvec[vec_idx] = Pf; - } + [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { + const uint col = c * cols_per_iter + col_tid; + + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { + const uint row = tile_row(r); + f16vec4 Pfvec; + [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { + const float16_t Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); + Lf[r + vec_idx] += Pf; + Pfvec[vec_idx] = Pf; + } + if (d_tid == 0) { // All other threads have the same P values Qf[col * psh_stride + row / 4] = Pfvec; } } @@ -406,40 +417,8 @@ void main() { barrier(); } - // prevent race on tmpsh - barrier(); - - // reduce across threads - - float rowmaxf[rows_per_thread], eMf[rows_per_thread], Moldf[rows_per_thread]; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - FLOAT_TYPE M = Mf[r]; - tmpsh[tid] = M; - // Compute max across the row - barrier(); - [[unroll]] for (int s = int(gl_WorkGroupSize.x / row_split) / 2; s >= D_split; s >>= 1) { - M = max(M, tmpsh[tid ^ s]); - barrier(); - tmpsh[tid] = M; - barrier(); - } - rowmaxf[r] = tmpsh[d_tid + row_tid * threads_per_rowgroup]; - barrier(); - } - - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Moldf[r] = Mf[r]; - - // M = max(rowmax, Mold) - // eM = e^(Mold - M) - Mf[r] = max(rowmaxf[r], Moldf[r]); - eMf[r] = exp(Moldf[r] - Mf[r]); - - Lf[r] = eMf[r]*Lf[r]; - } - - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - FLOAT_TYPE L = Lf[r]; + float L = Lf[r]; tmpsh[tid] = L; // Compute sum across the row barrier(); @@ -453,26 +432,6 @@ void main() { barrier(); } -#if 0 - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { - - Of[r][d] = ACC_TYPE(eMf[r]) * Of[r][d]; - tmpshv4[tid] = Of[r][d]; - - barrier(); - [[unroll]] for (int s = int(gl_WorkGroupSize.x / row_split) / 2; s >= D_split; s >>= 1) { - Of[r][d] += tmpshv4[tid ^ s]; - barrier(); - tmpshv4[tid] = Of[r][d]; - barrier(); - } - Of[r][d] = tmpshv4[d_tid + row_tid * threads_per_rowgroup]; - barrier(); - } - } -#endif - // If there is split_k, then the split_k resolve shader does the final // division by L. Store the intermediate O value and per-row m and L values. if (p.k_num > 1) { From b4e96f54f51481c3409ffcaf634d6c237ce425d7 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Thu, 22 Jan 2026 10:30:38 +0100 Subject: [PATCH 06/30] remove osh padding --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 546b45568109..4499ac0637ad 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -313,7 +313,7 @@ void main() { [[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) { const uint hsv_offset = (hsv_tile * row_split + gl_SubgroupID) * 16; - SfMat = coopmat(0); + SfMat = coopmat(0); #if BLOCK_SIZE > 1 // Preload V tiles for [Bc, 16 * num subgroups] @@ -359,7 +359,7 @@ void main() { } // Store SfMat to sfsh and load into Of - const uint osh_stride = row_split * MatBc + 8; + const uint osh_stride = row_split * MatBc; const uint o_offset = gl_SubgroupID * MatBc; coopMatStore(SfMat, sfsh, o_offset, osh_stride, gl_CooperativeMatrixLayoutRowMajor); From 1ed1f354c7ea758bf973d689b01efef9f62507b0 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Thu, 22 Jan 2026 15:04:07 +0100 Subject: [PATCH 07/30] add bounds checks and padding --- .../vulkan-shaders/flash_attn_base.glsl | 4 ++ .../vulkan-shaders/flash_attn_cm1.comp | 52 +++++++++++++------ .../vulkan-shaders/flash_attn_cm2.comp | 2 +- 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 29b5c7c3a418..36ea9200dec6 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -74,6 +74,10 @@ layout (binding = 1) readonly buffer K_PACKED16 {A_TYPE_PACKED16 k_data_packed16 layout (binding = 2) readonly buffer V_PACKED16 {A_TYPE_PACKED16 v_data_packed16[];} v_packed; #endif +#ifndef BLOCK_SIZE +#define BLOCK_SIZE 1 +#endif + #if defined(DATA_A_F32) #undef BLOCK_SIZE #define BLOCK_SIZE 4 diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 4499ac0637ad..9fce2b0ba03c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -307,7 +307,7 @@ void main() { } } - const uint num_hsv_tiles = HSV / MatBc / row_split; + const uint num_hsv_tiles = max(HSV / MatBc / row_split, 1); // Each subgroup handles HSV/4 columns [[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) { @@ -315,7 +315,6 @@ void main() { SfMat = coopmat(0); -#if BLOCK_SIZE > 1 // Preload V tiles for [Bc, 16 * num subgroups] const uint v_rows = Bc; const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups @@ -324,6 +323,10 @@ void main() { const uint vsh_stride = v_cols; +#if BLOCK_SIZE == 1 + // For f16, only preload if not aligned + if (KV_bounds_check) { +#endif [[unroll]] for (uint32_t i = 0; i < v_loads_per_thread; ++i) { const uint idx = i * gl_WorkGroupSize.x + tid; const uint row = idx / v_cols; @@ -336,7 +339,17 @@ void main() { const uint ib = coord / BLOCK_SIZE; const uint iqs = coord % BLOCK_SIZE; - ksh[row * vsh_stride + col] = f16vec4(dequantize4(ib, iqs, v_offset, BINDING_IDX_V)); + if (!KV_bounds_check || v_col < HSV) { +#if BLOCK_SIZE > 1 + ksh[row * vsh_stride + col] = f16vec4(dequantize4(ib, iqs, v_offset, BINDING_IDX_V)); +#else + ksh[row * vsh_stride + col] = data_vv4[(v_offset + v_row * v_stride + v_col) / 4]; +#endif + } else { + ksh[row * vsh_stride + col] = f16vec4(0.0f); + } + } +#if BLOCK_SIZE == 1 } #endif @@ -345,15 +358,18 @@ void main() { [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { coopMatLoad(KMat, Qf, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); -#if BLOCK_SIZE > 1 - const uint v_tile_offset = bc_chunk * MatBr * v_cols + gl_SubgroupID * (MatBc / 4); - coopMatLoad(QMat, ksh, v_tile_offset, vsh_stride, gl_CooperativeMatrixLayoutRowMajor); -#else - // F16 values can be loaded directly from global memory - const uint v_tile_row = j * Bc + bc_chunk * MatBc; - const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4; - coopMatLoad(QMat, data_vv4, v_tile_offset, v_stride / 4, gl_CooperativeMatrixLayoutRowMajor); +#if BLOCK_SIZE == 1 + if (!KV_bounds_check) { + // F16 values can be loaded directly from global memory + const uint v_tile_row = j * Bc + bc_chunk * MatBc; + const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4; + coopMatLoad(QMat, data_vv4, v_tile_offset, v_stride / 4, gl_CooperativeMatrixLayoutRowMajor); + } else #endif + { + const uint v_tile_offset = bc_chunk * MatBr * v_cols + gl_SubgroupID * (MatBc / 4); + coopMatLoad(QMat, ksh, v_tile_offset, vsh_stride, gl_CooperativeMatrixLayoutRowMajor); + } SfMat = coopMatMulAdd(KMat, QMat, SfMat); } @@ -379,12 +395,14 @@ void main() { const uint hsv_col = 4 * (d * D_split + d_tid); const uint local_hsv = hsv_col - hsv_base; - Of[r][d] += ACC_TYPEV4( - sfsh[row * osh_stride + local_hsv], - sfsh[row * osh_stride + local_hsv + 1], - sfsh[row * osh_stride + local_hsv + 2], - sfsh[row * osh_stride + local_hsv + 3] - ); + if (!KV_bounds_check || hsv_col < HSV) { + Of[r][d] += ACC_TYPEV4( + sfsh[row * osh_stride + local_hsv], + sfsh[row * osh_stride + local_hsv + 1], + sfsh[row * osh_stride + local_hsv + 2], + sfsh[row * osh_stride + local_hsv + 3] + ); + } } } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp index d49a8da65fbe..376cb7860e43 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -55,7 +55,7 @@ ACC_TYPE Max(const in uint32_t row, const in uint32_t col, const in ACC_TYPE ele return max(elem0, elem1); } -#if defined(BLOCK_SIZE) +#if BLOCK_SIZE > 1 #define DECODEFUNC , DEQUANTFUNC #else #define DECODEFUNC From 0bbee6d95dba8bd525f2a30d879bef5de94e3892 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Fri, 23 Jan 2026 12:11:07 +0100 Subject: [PATCH 08/30] remove unused code --- .../vulkan-shaders/flash_attn_cm1.comp | 26 ------------------- 1 file changed, 26 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 9fce2b0ba03c..9f2069bbde3e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -286,7 +286,6 @@ void main() { } } -#if 1 const uint psh_stride = Br / 4; // No space for padding, maybe increase Qf size? // Calculate and store Pf in Qf @@ -406,31 +405,6 @@ void main() { } } } -#else - [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { - continue; - } - float Pf[rows_per_thread]; - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Pf[r] = exp(sfsh[tile_row(r) + (c * cols_per_iter + col_tid) * sfshstride] - Mf[r]); - Lf[r] += Pf[r]; - } - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { -#if BLOCK_SIZE > 1 - uint coord = (j * Bc + c * cols_per_iter + col_tid) * v_stride * BLOCK_SIZE + 4 * (d * D_split + d_tid); - uint ib = coord / BLOCK_SIZE; - uint iqs = (coord % BLOCK_SIZE); - vec4 Vf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); -#else - vec4 Vf = vec4(data_vv4[v_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * v_stride / 4 + d * D_split + d_tid]); -#endif - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Of[r][d] += ACC_TYPE(Pf[r]) * ACC_TYPEV4(Vf); - } - } - } -#endif barrier(); } From f435f34a6ac00bfa67b949cd7d367d3e728d9b90 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Fri, 23 Jan 2026 14:53:02 +0100 Subject: [PATCH 09/30] fix shmem sizes, loop duration and accesses --- .../vulkan-shaders/flash_attn_cm1.comp | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 9f2069bbde3e..7b08e54d3dee 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -48,14 +48,17 @@ shared FLOAT_TYPE tmpsh[gl_WorkGroupSize.x]; shared ACC_TYPEV4 tmpshv4[gl_WorkGroupSize.x]; const uint32_t qstride = HSK_pad / 4 + 2; // in units of f16vec4 -shared f16vec4 Qf[Br * qstride]; +const uint psh_stride = Br / 4 + 2; +shared f16vec4 Qf[(Br * qstride >= Bc * psh_stride) ? (Br * qstride) : (Bc * psh_stride)]; // Avoid padding for hsk==256 to make it fit in 48KB shmem. const uint32_t sfshstride = (HSK <= 128) ? (Br + 8) : Br; shared ACC_TYPE sfsh[Bc * sfshstride]; const uint32_t kshstride = HSK_pad / 4 + 2; // in units of f16vec4 -shared f16vec4 ksh[Bc * kshstride]; +const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups +const uint vsh_stride = v_cols; +shared f16vec4 ksh[(kshstride >= vsh_stride) ? (Bc * kshstride) : (Bc * vsh_stride)]; shared float slope[Br]; @@ -286,8 +289,6 @@ void main() { } } - const uint psh_stride = Br / 4; // No space for padding, maybe increase Qf size? - // Calculate and store Pf in Qf [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { const uint col = c * cols_per_iter + col_tid; @@ -306,7 +307,7 @@ void main() { } } - const uint num_hsv_tiles = max(HSV / MatBc / row_split, 1); + const uint num_hsv_tiles = (HSV + MatBc * row_split - 1) / (MatBc * row_split); // round up // Each subgroup handles HSV/4 columns [[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) { @@ -316,12 +317,9 @@ void main() { // Preload V tiles for [Bc, 16 * num subgroups] const uint v_rows = Bc; - const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups const uint v_total = v_rows * v_cols; const uint v_loads_per_thread = v_total / gl_WorkGroupSize.x; - const uint vsh_stride = v_cols; - #if BLOCK_SIZE == 1 // For f16, only preload if not aligned if (KV_bounds_check) { @@ -385,7 +383,7 @@ void main() { const uint d_values_per_tile = hsv_per_tile / 4 / D_split; const uint d_start = hsv_tile * d_values_per_tile; - const uint d_end = d_start + d_values_per_tile; + const uint d_end = min(d_start + d_values_per_tile, HSV_per_thread / 4); [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); From 350313307877bd3850d3c96bdc0cdceb2fb1123b Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Fri, 23 Jan 2026 16:21:32 +0100 Subject: [PATCH 10/30] don't overwrite Qf, add new shared psh buffer instead --- .../src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 7b08e54d3dee..fe3a587d9f0e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -48,8 +48,10 @@ shared FLOAT_TYPE tmpsh[gl_WorkGroupSize.x]; shared ACC_TYPEV4 tmpshv4[gl_WorkGroupSize.x]; const uint32_t qstride = HSK_pad / 4 + 2; // in units of f16vec4 +shared f16vec4 Qf[Br * qstride]; + const uint psh_stride = Br / 4 + 2; -shared f16vec4 Qf[(Br * qstride >= Bc * psh_stride) ? (Br * qstride) : (Bc * psh_stride)]; +shared f16vec4 Psh[Bc * psh_stride]; // Avoid padding for hsk==256 to make it fit in 48KB shmem. const uint32_t sfshstride = (HSK <= 128) ? (Br + 8) : Br; @@ -289,7 +291,7 @@ void main() { } } - // Calculate and store Pf in Qf + // Calculate and store Pf in Psh [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { const uint col = c * cols_per_iter + col_tid; @@ -302,7 +304,7 @@ void main() { Pfvec[vec_idx] = Pf; } if (d_tid == 0) { // All other threads have the same P values - Qf[col * psh_stride + row / 4] = Pfvec; + Psh[col * psh_stride + row / 4] = Pfvec; } } } @@ -353,7 +355,7 @@ void main() { barrier(); [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { - coopMatLoad(KMat, Qf, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad(KMat, Psh, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); #if BLOCK_SIZE == 1 if (!KV_bounds_check) { From 1226ba0bbf8fbe5e2e163ef1054fd87c5f11d5d0 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Fri, 23 Jan 2026 20:01:34 +0100 Subject: [PATCH 11/30] add missing bounds checks --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index fe3a587d9f0e..fcac4943b502 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -299,7 +299,12 @@ void main() { const uint row = tile_row(r); f16vec4 Pfvec; [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { - const float16_t Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); + float16_t Pf; + if (KV_bounds_check && j * Bc + col >= KV) { + Pf = float16_t(0.0f); + } else { + Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); + } Lf[r + vec_idx] += Pf; Pfvec[vec_idx] = Pf; } @@ -338,7 +343,7 @@ void main() { const uint ib = coord / BLOCK_SIZE; const uint iqs = coord % BLOCK_SIZE; - if (!KV_bounds_check || v_col < HSV) { + if (!KV_bounds_check || (v_row < KV && v_col < HSV)) { #if BLOCK_SIZE > 1 ksh[row * vsh_stride + col] = f16vec4(dequantize4(ib, iqs, v_offset, BINDING_IDX_V)); #else From 44cd07a22bc40208cd143c6960f3696c4dc8cf56 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 06:18:04 +0100 Subject: [PATCH 12/30] use subgroup reductions --- .../vulkan-shaders/flash_attn_cm1.comp | 28 ++++--------------- 1 file changed, 6 insertions(+), 22 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index fcac4943b502..e3b6a15a07ed 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -8,6 +8,7 @@ #extension GL_KHR_shader_subgroup_basic : enable #extension GL_KHR_shader_subgroup_vote : enable +#extension GL_KHR_shader_subgroup_shuffle : enable #extension GL_KHR_memory_scope_semantics : enable #extension GL_KHR_cooperative_matrix : enable @@ -44,8 +45,7 @@ D_TYPE perElemOpGqaStore(const in uint32_t r, const in uint32_t c, const in D_TY const uint32_t MatBr = 16; const uint32_t MatBc = 16; -shared FLOAT_TYPE tmpsh[gl_WorkGroupSize.x]; -shared ACC_TYPEV4 tmpshv4[gl_WorkGroupSize.x]; +shared FLOAT_TYPE tmpsh[WorkGroupSize / row_split]; const uint32_t qstride = HSK_pad / 4 + 2; // in units of f16vec4 shared f16vec4 Qf[Br * qstride]; @@ -264,17 +264,10 @@ void main() { } float Moldf = Mf[r]; - tmpsh[tid] = rowmaxf; // Compute max across the row - barrier(); - [[unroll]] for (int s = int(gl_WorkGroupSize.x / row_split) / 2; s >= D_split; s >>= 1) { - rowmaxf = max(rowmaxf, tmpsh[tid ^ s]); - barrier(); - tmpsh[tid] = rowmaxf; - barrier(); + [[unroll]] for (uint s = D_split; s < 32; s *= 2) { + rowmaxf = max(rowmaxf, subgroupShuffleXor(rowmaxf, s)); } - rowmaxf = tmpsh[d_tid + row_tid * threads_per_rowgroup]; - barrier(); // M = max(rowmax, Mold) // P = e^(S - M) @@ -415,18 +408,9 @@ void main() { } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - float L = Lf[r]; - tmpsh[tid] = L; - // Compute sum across the row - barrier(); - [[unroll]] for (int s = int(gl_WorkGroupSize.x / row_split) / 2; s >= D_split; s >>= 1) { - L += tmpsh[tid ^ s]; - barrier(); - tmpsh[tid] = L; - barrier(); + [[unroll]] for (uint s = D_split; s < 32; s *= 2) { + Lf[r] += subgroupShuffleXor(Lf[r], s); } - Lf[r] = tmpsh[d_tid + row_tid * threads_per_rowgroup]; - barrier(); } // If there is split_k, then the split_k resolve shader does the final From 3037c75e3c48bfbaa6e063bdbf3ff08ce4cb02a2 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 06:25:25 +0100 Subject: [PATCH 13/30] optimize --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index e3b6a15a07ed..780c50eef60a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -45,7 +45,7 @@ D_TYPE perElemOpGqaStore(const in uint32_t r, const in uint32_t c, const in D_TY const uint32_t MatBr = 16; const uint32_t MatBc = 16; -shared FLOAT_TYPE tmpsh[WorkGroupSize / row_split]; +shared FLOAT_TYPE tmpsh[row_split]; const uint32_t qstride = HSK_pad / 4 + 2; // in units of f16vec4 shared f16vec4 Qf[Br * qstride]; @@ -172,7 +172,6 @@ void main() { } // skip the block if the mask is entirely -inf bool all_less = subgroupAll(max_mask <= NEG_FLT_MAX_OVER_2); - barrier(); if (gl_SubgroupInvocationID == 0) { tmpsh[gl_SubgroupID] = all_less ? NEG_FLT_MAX_OVER_2 : 0.0f; } @@ -213,7 +212,7 @@ void main() { coopmat KMat; coopmat QMat; - for (uint32_t d = 0; d < HSK_pad / 16; ++d) { + [[unroll]] for (uint32_t d = 0; d < HSK_pad / 16; ++d) { coopMatLoad(QMat, Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor); uint coord = (gl_SubgroupID * MatBc) * kshstride + d * 16 / 4; From 4b3b6a6de31c87bf9ef8e54bf271de9b6a529365 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 06:52:02 +0100 Subject: [PATCH 14/30] move bounds check, reduce barriers --- .../vulkan-shaders/flash_attn_cm1.comp | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 780c50eef60a..df885e41af6e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -130,13 +130,11 @@ void main() { uint r = tid; slope[r] = perElemOpComputeSlope(r, col_tid, ACC_TYPE(0), iq2); } - barrier(); } else { if (tid < Br) { uint r = tid; slope[r] = 1.0; } - barrier(); } #if BLOCK_SIZE > 1 @@ -290,15 +288,16 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { const uint row = tile_row(r); f16vec4 Pfvec; - [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { - float16_t Pf; - if (KV_bounds_check && j * Bc + col >= KV) { - Pf = float16_t(0.0f); - } else { - Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); + if (KV_bounds_check && j * Bc + col >= KV) { + [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { + Pfvec[vec_idx] = float16_t(0.0f); + } + } else { + [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { + const float16_t Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); + Lf[r + vec_idx] += Pf; + Pfvec[vec_idx] = Pf; } - Lf[r + vec_idx] += Pf; - Pfvec[vec_idx] = Pf; } if (d_tid == 0) { // All other threads have the same P values Psh[col * psh_stride + row / 4] = Pfvec; From 1c40f9e1e9732248eb898a1e31e18fcb8d5a4714 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 08:55:58 +0100 Subject: [PATCH 15/30] support other Bc values and other subgroup sizes --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 31 +++++++++++++------ .../vulkan-shaders/flash_attn_base.glsl | 1 + .../vulkan-shaders/flash_attn_cm1.comp | 14 ++++----- 3 files changed, 30 insertions(+), 16 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index b5e5dba95fe4..2635d373c66f 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2763,7 +2763,10 @@ static std::array fa_rows_cols(FaCodePath path, uint32_t hsk, uint3 if (small_rows) { return {scalar_flash_attention_num_small_rows, scalar_flash_attention_Bc}; } else { - return {coopmat1_flash_attention_num_large_rows, scalar_flash_attention_Bc}; + if ((ggml_is_quantized(type) && (hsk >= 256 || hsv >= 256)) || hsk >= 512 || hsv >= 512) { + return {16, 32}; + } + return {16, 64}; } } @@ -3162,17 +3165,27 @@ static void ggml_vk_load_shaders(vk_device& device) { // For scalar, use 128 (arbitrary) // The same D_split value is used for both HSK and HSV, so just base it on the union of the LSBs. const uint32_t D = (hsk|hsv); - uint32_t wg_size = (path == FA_SCALAR || path == FA_COOPMAT1) - ? scalar_flash_attention_workgroup_size - : ((small_rows && (D % 32) == 0) ? 256 : 128); auto rows_cols = fa_rows_cols(path, hsk, hsv, clamp, type, small_rows, small_cache); + uint32_t wg_size; + switch (path) { + case FA_COOPMAT2: + wg_size = ((small_rows && (D % 32) == 0) ? 256 : 128); + break; + case FA_COOPMAT1: + wg_size = (rows_cols[1] / 16) * device->subgroup_size; // enough subgroups for Bc/MatBc + break; + default: + wg_size = scalar_flash_attention_workgroup_size; + break; + } + // 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, 8u), D_lsb / 4); - return {wg_size, rows_cols[0], rows_cols[1], hsk, hsv, clamp, D_split}; + return {wg_size, rows_cols[0], rows_cols[1], hsk, hsv, clamp, D_split, device->subgroup_size}; }; #define CREATE_FA(TYPE, NAMELC, FAPATH, SUFFIX) \ @@ -3187,15 +3200,15 @@ static void ggml_vk_load_shaders(vk_device& device) { if (path == FAPATH) { \ if (aligned) { \ if (f32acc) { \ - ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_aligned_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## SUFFIX ## _len, flash_attn_f32_f16_ ## NAMELC ## SUFFIX ## _data, "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_align(FAPATH,HSK,HSV,TYPE,small_rows,small_cache), true, FAPATH==FA_COOPMAT1, (FAPATH==FA_COOPMAT1 ? 32 : 0)); \ + ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_aligned_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## SUFFIX ## _len, flash_attn_f32_f16_ ## NAMELC ## SUFFIX ## _data, "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_align(FAPATH,HSK,HSV,TYPE,small_rows,small_cache), true, FAPATH==FA_COOPMAT1, (FAPATH==FA_COOPMAT1 ? device->subgroup_size : 0)); \ } else { \ - ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_aligned_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _len, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _data, "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_align(FAPATH,HSK,HSV,TYPE,small_rows,small_cache), true, FAPATH==FA_COOPMAT1, (FAPATH==FA_COOPMAT1 ? 32 : 0)); \ + ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_aligned_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _len, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _data, "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,0,TYPE,small_rows,small_cache), fa_align(FAPATH,HSK,HSV,TYPE,small_rows,small_cache), true, FAPATH==FA_COOPMAT1, (FAPATH==FA_COOPMAT1 ? device->subgroup_size : 0)); \ } \ } else { \ if (f32acc) { \ - ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## SUFFIX ## _len, flash_attn_f32_f16_ ## NAMELC ## SUFFIX ## _data, "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), 1, true, FAPATH==FA_COOPMAT1, (FAPATH==FA_COOPMAT1 ? 32 : 0)); \ + ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_f32acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## SUFFIX ## _len, flash_attn_f32_f16_ ## NAMELC ## SUFFIX ## _data, "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), 1, true, FAPATH==FA_COOPMAT1, (FAPATH==FA_COOPMAT1 ? device->subgroup_size : 0)); \ } else { \ - ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _len, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _data, "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), 1, true, FAPATH==FA_COOPMAT1, (FAPATH==FA_COOPMAT1 ? 32 : 0)); \ + ggml_vk_create_pipeline(device, fa.second, "flash_attn_f32_f16_f16acc" #NAMELC, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _len, flash_attn_f32_f16_ ## NAMELC ## _f16acc ## SUFFIX ## _data, "main", 6, sizeof(vk_flash_attn_push_constants), fa_wg_denoms(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), fa_spec_constants(FAPATH, HSK,HSV,1,TYPE,small_rows,small_cache), 1, true, FAPATH==FA_COOPMAT1, (FAPATH==FA_COOPMAT1 ? device->subgroup_size : 0)); \ } \ } \ } \ diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 36ea9200dec6..6b5b2b248ca4 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -8,6 +8,7 @@ layout (constant_id = 3) const uint32_t HSK = 32; layout (constant_id = 4) const uint32_t HSV = 32; layout (constant_id = 5) const uint32_t Clamp = 0; layout (constant_id = 6) const uint32_t D_split = 16; +layout (constant_id = 7) const uint32_t SubGroupSize = 32; // Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths const uint32_t HSK_pad = (HSK + 15) & ~15; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index df885e41af6e..e40341e754eb 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -15,10 +15,14 @@ #include "types.glsl" #include "flash_attn_base.glsl" +// These need to be supported N,M values for a MatBc x MatBr x 16 coopmatmuladd +const uint32_t MatBr = 16; +const uint32_t MatBc = 16; + const uint32_t HSK_per_thread = HSK / D_split; const uint32_t HSV_per_thread = HSV / D_split; -const uint32_t row_split = 4; +const uint32_t row_split = Bc / MatBc; const uint32_t rows_per_thread = Br / row_split; const uint32_t cols_per_iter = gl_WorkGroupSize.x / D_split / row_split; const uint32_t cols_per_thread = Bc / cols_per_iter; @@ -41,10 +45,6 @@ D_TYPE perElemOpGqaStore(const in uint32_t r, const in uint32_t c, const in D_TY return elem; } -// These need to be supported N,M values for a MatBc x MatBr x 16 coopmatmuladd -const uint32_t MatBr = 16; -const uint32_t MatBc = 16; - shared FLOAT_TYPE tmpsh[row_split]; const uint32_t qstride = HSK_pad / 4 + 2; // in units of f16vec4 @@ -262,7 +262,7 @@ void main() { float Moldf = Mf[r]; // Compute max across the row - [[unroll]] for (uint s = D_split; s < 32; s *= 2) { + [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) { rowmaxf = max(rowmaxf, subgroupShuffleXor(rowmaxf, s)); } @@ -406,7 +406,7 @@ void main() { } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - [[unroll]] for (uint s = D_split; s < 32; s *= 2) { + [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) { Lf[r] += subgroupShuffleXor(Lf[r], s); } } From 61745fd71db869bcc7f50ce03163886375888f07 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 13:08:47 +0100 Subject: [PATCH 16/30] remove D_split --- .../vulkan-shaders/flash_attn_cm1.comp | 50 ++++++++----------- 1 file changed, 20 insertions(+), 30 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index e40341e754eb..2c665642c509 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -7,8 +7,8 @@ #extension GL_EXT_shader_explicit_arithmetic_types_int32 : require #extension GL_KHR_shader_subgroup_basic : enable +#extension GL_KHR_shader_subgroup_arithmetic : enable #extension GL_KHR_shader_subgroup_vote : enable -#extension GL_KHR_shader_subgroup_shuffle : enable #extension GL_KHR_memory_scope_semantics : enable #extension GL_KHR_cooperative_matrix : enable @@ -19,12 +19,9 @@ const uint32_t MatBr = 16; const uint32_t MatBc = 16; -const uint32_t HSK_per_thread = HSK / D_split; -const uint32_t HSV_per_thread = HSV / D_split; - const uint32_t row_split = Bc / MatBc; const uint32_t rows_per_thread = Br / row_split; -const uint32_t cols_per_iter = gl_WorkGroupSize.x / D_split / row_split; +const uint32_t cols_per_iter = gl_WorkGroupSize.x / row_split; const uint32_t cols_per_thread = Bc / cols_per_iter; @@ -75,8 +72,7 @@ void main() { const uint32_t threads_per_rowgroup = gl_WorkGroupSize.x / row_split; const uint32_t row_tid = gl_LocalInvocationIndex / threads_per_rowgroup; - const uint32_t d_tid = gl_LocalInvocationIndex % D_split; - const uint32_t col_tid = (gl_LocalInvocationIndex % threads_per_rowgroup) / D_split; + const uint32_t col_tid = gl_LocalInvocationIndex % threads_per_rowgroup; #define tile_row(r) (row_tid * rows_per_thread + (r)) @@ -107,8 +103,8 @@ void main() { } barrier(); - ACC_TYPEV4 Of[rows_per_thread][HSV_per_thread / 4]; - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { + ACC_TYPEV4 Of[rows_per_thread][HSV / 4]; + [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Of[r][d] = ACC_TYPEV4(0.0); } @@ -262,9 +258,7 @@ void main() { float Moldf = Mf[r]; // Compute max across the row - [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) { - rowmaxf = max(rowmaxf, subgroupShuffleXor(rowmaxf, s)); - } + rowmaxf = subgroupMax(rowmaxf); // M = max(rowmax, Mold) // P = e^(S - M) @@ -275,7 +269,7 @@ void main() { Lf[r] = eMf[r]*Lf[r]; } - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { + [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Of[r][d] = ACC_TYPE(eMf[r]) * Of[r][d]; } @@ -299,9 +293,7 @@ void main() { Pfvec[vec_idx] = Pf; } } - if (d_tid == 0) { // All other threads have the same P values - Psh[col * psh_stride + row / 4] = Pfvec; - } + Psh[col * psh_stride + row / 4] = Pfvec; } } @@ -378,16 +370,16 @@ void main() { const uint hsv_per_tile = row_split * MatBc; const uint hsv_base = hsv_tile * hsv_per_tile; - const uint d_values_per_tile = hsv_per_tile / 4 / D_split; + const uint d_values_per_tile = hsv_per_tile / 4; const uint d_start = hsv_tile * d_values_per_tile; - const uint d_end = min(d_start + d_values_per_tile, HSV_per_thread / 4); + const uint d_end = min(d_start + d_values_per_tile, HSV / 4); [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); [[unroll]] for (uint32_t d = d_start; d < d_end; ++d) { - const uint hsv_col = 4 * (d * D_split + d_tid); + const uint hsv_col = 4 * d; const uint local_hsv = hsv_col - hsv_base; if (!KV_bounds_check || hsv_col < HSV) { @@ -406,9 +398,7 @@ void main() { } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) { - Lf[r] += subgroupShuffleXor(Lf[r], s); - } + Lf[r] = subgroupAdd(Lf[r]); } // If there is split_k, then the split_k resolve shader does the final @@ -419,9 +409,9 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { + [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - perElemOpGqaStore(tile_row(r), 4*(d * D_split + d_tid) + comp, float(Of[r][d][comp]), o_offset, iq2, N); + perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Of[r][d][comp]), o_offset, iq2, N); } } } @@ -448,7 +438,7 @@ void main() { if (sink > Mf[r]) { ms = exp(Mf[r] - sink); - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { + [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { Of[r][d] *= ACC_TYPE(ms); } } else { @@ -464,7 +454,7 @@ void main() { Lfrcp[r] = (Lf[r] == 0.0) ? 0.0 : (1.0 / Lf[r]); } - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { + [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Of[r][d] *= ACC_TYPE(Lfrcp[r]); #if defined(ACC_TYPE_MAX) @@ -478,9 +468,9 @@ void main() { if (p.gqa_ratio > 1) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { + [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - perElemOpGqaStore(tile_row(r), 4*(d * D_split + d_tid) + comp, float(Of[r][d][comp]), o_offset, iq2, N); + perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Of[r][d][comp]), o_offset, iq2, N); } } } @@ -488,9 +478,9 @@ void main() { } else { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (i * Br + tile_row(r) < N) { - [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { + [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - data_o[o_offset + iq2 * HSV + (i * Br + tile_row(r)) * p.ne1 * HSV + 4*(d * D_split + d_tid) + comp] = D_TYPE(Of[r][d][comp]); + data_o[o_offset + iq2 * HSV + (i * Br + tile_row(r)) * p.ne1 * HSV + 4 * d + comp] = D_TYPE(Of[r][d][comp]); } } } From 74d3246c238118d490d4d284137aaa5f8dd68ca7 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 13:38:08 +0100 Subject: [PATCH 17/30] replace Of register array with shared memory Ofsh array --- .../vulkan-shaders/flash_attn_cm1.comp | 32 +++++++++++-------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 2c665642c509..70cead3f5e36 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -59,6 +59,10 @@ const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * nu const uint vsh_stride = v_cols; shared f16vec4 ksh[(kshstride >= vsh_stride) ? (Bc * kshstride) : (Bc * vsh_stride)]; +// Output buffer +const uint32_t ofsh_stride = HSV / 4 + 2; +shared ACC_TYPEV4 Ofsh[Br * ofsh_stride]; + shared float slope[Br]; void main() { @@ -101,14 +105,13 @@ void main() { Qf[r * qstride + d] = f16vec4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale); } } - barrier(); - ACC_TYPEV4 Of[rows_per_thread][HSV / 4]; - [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Of[r][d] = ACC_TYPEV4(0.0); + [[unroll]] for (uint idx = 0; idx < Br * ofsh_stride; idx += gl_WorkGroupSize.x) { + if (idx + tid < Br * ofsh_stride) { + Ofsh[idx + tid] = ACC_TYPEV4(0.0); } } + barrier(); float Lf[rows_per_thread], Mf[rows_per_thread]; @@ -271,7 +274,7 @@ void main() { [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Of[r][d] = ACC_TYPE(eMf[r]) * Of[r][d]; + Ofsh[tile_row(r) * ofsh_stride + d] = ACC_TYPE(eMf[r]) * Ofsh[tile_row(r) * ofsh_stride + d]; } } @@ -361,7 +364,7 @@ void main() { SfMat = coopMatMulAdd(KMat, QMat, SfMat); } - // Store SfMat to sfsh and load into Of + // Store SfMat to sfsh and load into Ofsh const uint osh_stride = row_split * MatBc; const uint o_offset = gl_SubgroupID * MatBc; coopMatStore(SfMat, sfsh, o_offset, osh_stride, gl_CooperativeMatrixLayoutRowMajor); @@ -383,7 +386,7 @@ void main() { const uint local_hsv = hsv_col - hsv_base; if (!KV_bounds_check || hsv_col < HSV) { - Of[r][d] += ACC_TYPEV4( + Ofsh[tile_row(r) * ofsh_stride + d] += ACC_TYPEV4( sfsh[row * osh_stride + local_hsv], sfsh[row * osh_stride + local_hsv + 1], sfsh[row * osh_stride + local_hsv + 2], @@ -411,7 +414,7 @@ void main() { if (tile_row(r) < N) { [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Of[r][d][comp]), o_offset, iq2, N); + perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Ofsh[tile_row(r) * ofsh_stride + d][comp]), o_offset, iq2, N); } } } @@ -439,7 +442,7 @@ void main() { ms = exp(Mf[r] - sink); [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { - Of[r][d] *= ACC_TYPE(ms); + Ofsh[tile_row(r) * ofsh_stride + d] *= ACC_TYPE(ms); } } else { vs = exp(sink - Mf[r]); @@ -456,10 +459,11 @@ void main() { [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Of[r][d] *= ACC_TYPE(Lfrcp[r]); + ACC_TYPEV4 Of = Ofsh[tile_row(r) * ofsh_stride + d] * ACC_TYPE(Lfrcp[r]); #if defined(ACC_TYPE_MAX) - Of[r][d] = clamp(Of[r][d], -ACC_TYPE_MAX, ACC_TYPE_MAX); + Of = clamp(Of, -ACC_TYPE_MAX, ACC_TYPE_MAX); #endif + Ofsh[tile_row(r) * ofsh_stride + d] = Of; } } @@ -470,7 +474,7 @@ void main() { if (tile_row(r) < N) { [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Of[r][d][comp]), o_offset, iq2, N); + perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Ofsh[tile_row(r) * ofsh_stride + d][comp]), o_offset, iq2, N); } } } @@ -480,7 +484,7 @@ void main() { if (i * Br + tile_row(r) < N) { [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - data_o[o_offset + iq2 * HSV + (i * Br + tile_row(r)) * p.ne1 * HSV + 4 * d + comp] = D_TYPE(Of[r][d][comp]); + data_o[o_offset + iq2 * HSV + (i * Br + tile_row(r)) * p.ne1 * HSV + 4 * d + comp] = D_TYPE(Ofsh[tile_row(r) * ofsh_stride + d][comp]); } } } From e0c414dd2163b6000c93c34013ca95f415eeff04 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 14:08:42 +0100 Subject: [PATCH 18/30] parallelize HSV across the rowgroups --- .../vulkan-shaders/flash_attn_cm1.comp | 29 ++++++++++++++----- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 70cead3f5e36..c316ddb18a7f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -272,7 +272,9 @@ void main() { Lf[r] = eMf[r]*Lf[r]; } - [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { + [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { + const uint d = d0 + col_tid; + if (d >= HSV / 4) break; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Ofsh[tile_row(r) * ofsh_stride + d] = ACC_TYPE(eMf[r]) * Ofsh[tile_row(r) * ofsh_stride + d]; } @@ -381,7 +383,10 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); - [[unroll]] for (uint32_t d = d_start; d < d_end; ++d) { + [[unroll]] for (uint32_t d0 = d_start; d0 < d_end; d0 += threads_per_rowgroup) { + const uint d = d0 + col_tid; + if (d >= d_end) break; + const uint hsv_col = 4 * d; const uint local_hsv = hsv_col - hsv_base; @@ -412,7 +417,9 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { - [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { + [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { + const uint d = d0 + col_tid; + if (d >= HSV/4) break; [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Ofsh[tile_row(r) * ofsh_stride + d][comp]), o_offset, iq2, N); } @@ -441,7 +448,9 @@ void main() { if (sink > Mf[r]) { ms = exp(Mf[r] - sink); - [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { + [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { + const uint d = d0 + col_tid; + if (d >= HSV / 4) break; Ofsh[tile_row(r) * ofsh_stride + d] *= ACC_TYPE(ms); } } else { @@ -457,7 +466,9 @@ void main() { Lfrcp[r] = (Lf[r] == 0.0) ? 0.0 : (1.0 / Lf[r]); } - [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { + [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { + const uint d = d0 + col_tid; + if (d >= HSV / 4) break; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { ACC_TYPEV4 Of = Ofsh[tile_row(r) * ofsh_stride + d] * ACC_TYPE(Lfrcp[r]); #if defined(ACC_TYPE_MAX) @@ -472,7 +483,9 @@ void main() { if (p.gqa_ratio > 1) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { - [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { + [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { + const uint d = d0 + col_tid; + if (d >= HSV / 4) break; [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Ofsh[tile_row(r) * ofsh_stride + d][comp]), o_offset, iq2, N); } @@ -482,7 +495,9 @@ void main() { } else { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (i * Br + tile_row(r) < N) { - [[unroll]] for (uint32_t d = 0; d < HSV / 4; ++d) { + [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { + const uint d = d0 + col_tid; + if (d >= HSV / 4) break; [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { data_o[o_offset + iq2 * HSV + (i * Br + tile_row(r)) * p.ne1 * HSV + 4 * d + comp] = D_TYPE(Ofsh[tile_row(r) * ofsh_stride + d][comp]); } From ed11a955af99ddeb3ccdd074e7811f3d0529a443 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 15:10:06 +0100 Subject: [PATCH 19/30] go back to Of in registers, not shmem --- .../vulkan-shaders/flash_attn_cm1.comp | 46 +++++++++---------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index c316ddb18a7f..49509fe21321 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -59,10 +59,6 @@ const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * nu const uint vsh_stride = v_cols; shared f16vec4 ksh[(kshstride >= vsh_stride) ? (Bc * kshstride) : (Bc * vsh_stride)]; -// Output buffer -const uint32_t ofsh_stride = HSV / 4 + 2; -shared ACC_TYPEV4 Ofsh[Br * ofsh_stride]; - shared float slope[Br]; void main() { @@ -75,6 +71,7 @@ void main() { const uint32_t tid = gl_LocalInvocationIndex; const uint32_t threads_per_rowgroup = gl_WorkGroupSize.x / row_split; + const uint32_t d_per_thread = (HSV/4 + threads_per_rowgroup - 1) / threads_per_rowgroup; const uint32_t row_tid = gl_LocalInvocationIndex / threads_per_rowgroup; const uint32_t col_tid = gl_LocalInvocationIndex % threads_per_rowgroup; @@ -105,13 +102,14 @@ void main() { Qf[r * qstride + d] = f16vec4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale); } } + barrier(); - [[unroll]] for (uint idx = 0; idx < Br * ofsh_stride; idx += gl_WorkGroupSize.x) { - if (idx + tid < Br * ofsh_stride) { - Ofsh[idx + tid] = ACC_TYPEV4(0.0); + ACC_TYPEV4 Of[rows_per_thread][d_per_thread]; + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + [[unroll]] for (uint32_t d = 0; d < d_per_thread; ++d) { + Of[r][d] = ACC_TYPEV4(0.0); } } - barrier(); float Lf[rows_per_thread], Mf[rows_per_thread]; @@ -273,10 +271,9 @@ void main() { } [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { - const uint d = d0 + col_tid; - if (d >= HSV / 4) break; + const uint d_local = d0 / threads_per_rowgroup; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Ofsh[tile_row(r) * ofsh_stride + d] = ACC_TYPE(eMf[r]) * Ofsh[tile_row(r) * ofsh_stride + d]; + Of[r][d_local] = ACC_TYPE(eMf[r]) * Of[r][d_local]; } } @@ -366,7 +363,7 @@ void main() { SfMat = coopMatMulAdd(KMat, QMat, SfMat); } - // Store SfMat to sfsh and load into Ofsh + // Store SfMat to sfsh and load into Of const uint osh_stride = row_split * MatBc; const uint o_offset = gl_SubgroupID * MatBc; coopMatStore(SfMat, sfsh, o_offset, osh_stride, gl_CooperativeMatrixLayoutRowMajor); @@ -391,7 +388,8 @@ void main() { const uint local_hsv = hsv_col - hsv_base; if (!KV_bounds_check || hsv_col < HSV) { - Ofsh[tile_row(r) * ofsh_stride + d] += ACC_TYPEV4( + const uint d_local = d0 / threads_per_rowgroup; + Of[r][d_local] += ACC_TYPEV4( sfsh[row * osh_stride + local_hsv], sfsh[row * osh_stride + local_hsv + 1], sfsh[row * osh_stride + local_hsv + 2], @@ -420,8 +418,9 @@ void main() { [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { const uint d = d0 + col_tid; if (d >= HSV/4) break; + const uint d_local = d0 / threads_per_rowgroup; [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Ofsh[tile_row(r) * ofsh_stride + d][comp]), o_offset, iq2, N); + perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Of[r][d_local][comp]), o_offset, iq2, N); } } } @@ -449,9 +448,8 @@ void main() { ms = exp(Mf[r] - sink); [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { - const uint d = d0 + col_tid; - if (d >= HSV / 4) break; - Ofsh[tile_row(r) * ofsh_stride + d] *= ACC_TYPE(ms); + const uint d_local = d0 / threads_per_rowgroup; + Of[r][d_local] *= ACC_TYPE(ms); } } else { vs = exp(sink - Mf[r]); @@ -467,14 +465,12 @@ void main() { } [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { - const uint d = d0 + col_tid; - if (d >= HSV / 4) break; + const uint d_local = d0 / threads_per_rowgroup; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - ACC_TYPEV4 Of = Ofsh[tile_row(r) * ofsh_stride + d] * ACC_TYPE(Lfrcp[r]); + Of[r][d_local] *= ACC_TYPE(Lfrcp[r]); #if defined(ACC_TYPE_MAX) - Of = clamp(Of, -ACC_TYPE_MAX, ACC_TYPE_MAX); + Of[r][d_local] = clamp(Of[r][d_local], -ACC_TYPE_MAX, ACC_TYPE_MAX); #endif - Ofsh[tile_row(r) * ofsh_stride + d] = Of; } } @@ -486,8 +482,9 @@ void main() { [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { const uint d = d0 + col_tid; if (d >= HSV / 4) break; + const uint d_local = d0 / threads_per_rowgroup; [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Ofsh[tile_row(r) * ofsh_stride + d][comp]), o_offset, iq2, N); + perElemOpGqaStore(tile_row(r), 4 * d + comp, float(Of[r][d_local][comp]), o_offset, iq2, N); } } } @@ -498,8 +495,9 @@ void main() { [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { const uint d = d0 + col_tid; if (d >= HSV / 4) break; + const uint d_local = d0 / threads_per_rowgroup; [[unroll]] for (uint32_t comp = 0; comp < 4; ++comp) { - data_o[o_offset + iq2 * HSV + (i * Br + tile_row(r)) * p.ne1 * HSV + 4 * d + comp] = D_TYPE(Ofsh[tile_row(r) * ofsh_stride + d][comp]); + data_o[o_offset + iq2 * HSV + (i * Br + tile_row(r)) * p.ne1 * HSV + 4 * d + comp] = D_TYPE(Of[r][d_local][comp]); } } } From 9384172c475d1015f00ae0ca7282932b07f2e596 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 16:49:53 +0100 Subject: [PATCH 20/30] vectorize sfsh --- .../vulkan-shaders/flash_attn_cm1.comp | 103 +++++++++++------- 1 file changed, 61 insertions(+), 42 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 49509fe21321..413d355f5cb1 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -51,15 +51,15 @@ const uint psh_stride = Br / 4 + 2; shared f16vec4 Psh[Bc * psh_stride]; // Avoid padding for hsk==256 to make it fit in 48KB shmem. -const uint32_t sfshstride = (HSK <= 128) ? (Br + 8) : Br; -shared ACC_TYPE sfsh[Bc * sfshstride]; +const uint32_t sfshstride = (HSK <= 128) ? (Br / 4 + 2) : Br / 4; +shared ACC_TYPEV4 sfsh[Bc * sfshstride]; const uint32_t kshstride = HSK_pad / 4 + 2; // in units of f16vec4 const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups const uint vsh_stride = v_cols; shared f16vec4 ksh[(kshstride >= vsh_stride) ? (Bc * kshstride) : (Bc * vsh_stride)]; -shared float slope[Br]; +shared ACC_TYPE slope[Br]; void main() { #ifdef NEEDS_INIT_IQ_SHMEM @@ -130,7 +130,7 @@ void main() { } else { if (tid < Br) { uint r = tid; - slope[r] = 1.0; + slope[r] = ACC_TYPE(1.0); } } @@ -149,19 +149,22 @@ void main() { [[dont_unroll]] for (uint32_t j = start_j; j < end_j; ++j) { - float mask_cache[Bc * Br / WorkGroupSize]; + f16vec4 mask_cache[Bc * Br / 4 / WorkGroupSize]; if ((p.mask_n_head_log2 & MASK_ENABLE_BIT) != 0) { bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; float max_mask = NEG_FLT_MAX_OVER_2; - [[unroll]] for (uint32_t idx = 0; idx < Bc * Br; idx += gl_WorkGroupSize.x) { - uint32_t c = (idx + tid) % Bc; - uint32_t r = (idx + tid) / Bc; - if (idx + tid < Bc * Br || idx + gl_WorkGroupSize.x <= Bc * Br) { - if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r < p.nem1)) { - float m = float(data_m[m_offset + (i * Br + r) * m_stride + (j * Bc + c)]); + [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { + uint32_t c = (idx + tid) / (Br / 4); + uint32_t r = (idx + tid) % (Br / 4); + if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { + if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r * 4 < p.nem1)) { + const f16vec4 m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 3) * m_stride + (j * Bc + c)]); mask_cache[idx / WorkGroupSize] = m; - max_mask = max(max_mask, m); + max_mask = max(max(max(max(max_mask, float(m[0])), float(m[1])), float(m[2])), float(m[3])); } } } @@ -221,11 +224,11 @@ void main() { barrier(); if (p.logit_softcap != 0.0f) { - [[unroll]] for (uint32_t idx = 0; idx < Bc * Br; idx += gl_WorkGroupSize.x) { - uint32_t c = (idx + tid) / Br; - uint32_t r = (idx + tid) % Br; - if (idx + tid < Bc * Br || idx + gl_WorkGroupSize.x <= Bc * Br) { - sfsh[c * sfshstride + r] = ACC_TYPE(p.logit_softcap * tanh(sfsh[c * sfshstride + r])); + [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { + uint32_t c = (idx + tid) / (Br / 4); + uint32_t r = (idx + tid) % (Br / 4); + if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { + sfsh[c * sfshstride + r] = ACC_TYPEV4(p.logit_softcap * tanh(sfsh[c * sfshstride + r])); } } barrier(); @@ -234,13 +237,34 @@ void main() { if ((p.mask_n_head_log2 & MASK_ENABLE_BIT) != 0) { bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; - [[unroll]] for (uint32_t idx = 0; idx < Bc * Br; idx += gl_WorkGroupSize.x) { - uint32_t c = (idx + tid) % Bc; - uint32_t r = (idx + tid) / Bc; - if (idx + tid < Bc * Br || idx + gl_WorkGroupSize.x <= Bc * Br) { - if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r < p.nem1)) { - float f = mask_cache[idx / WorkGroupSize]; - sfsh[c * sfshstride + r] += ACC_TYPE(slope[r] * f); + [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { + uint32_t c = (idx + tid) / (Br / 4); + uint32_t r = (idx + tid) % (Br / 4); + if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { + if (!KV_bounds_check || j * Bc + c < KV) { + if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { + ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); + ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); + sfsh[c * sfshstride + r] += slopes * masks; + } else if (i * Br + r * 4 + 2 < p.nem1) { + ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); + masks.w = ACC_TYPE(0.0f); + ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); + sfsh[c * sfshstride + r] += slopes * masks; + } else if (i * Br + r * 4 + 1 < p.nem1) { + ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); + masks.z = ACC_TYPE(0.0f); + masks.w = ACC_TYPE(0.0f); + ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); + sfsh[c * sfshstride + r] += slopes * masks; + } else if (i * Br + r * 4 < p.nem1) { + ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); + masks.y = ACC_TYPE(0.0f); + masks.z = ACC_TYPE(0.0f); + masks.w = ACC_TYPE(0.0f); + ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); + sfsh[c * sfshstride + r] += slopes * masks; + } } } } @@ -249,12 +273,15 @@ void main() { float eMf[rows_per_thread]; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + const uint r_vec = tile_row(r) / 4; + const uint r_comp = tile_row(r) % 4; + float rowmaxf = NEG_FLT_MAX_OVER_2; [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { continue; } - rowmaxf = max(rowmaxf, float(sfsh[tile_row(r) + (c * cols_per_iter + col_tid) * sfshstride])); + rowmaxf = max(rowmaxf, float(sfsh[r_vec + (c * cols_per_iter + col_tid) * sfshstride][r_comp])); } float Moldf = Mf[r]; @@ -283,19 +310,16 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { const uint row = tile_row(r); - f16vec4 Pfvec; if (KV_bounds_check && j * Bc + col >= KV) { - [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { - Pfvec[vec_idx] = float16_t(0.0f); - } + Psh[col * psh_stride + row / 4] = f16vec4(0.0f); } else { + const f16vec4 mfvec = f16vec4(Mf[r], Mf[r + 1], Mf[r + 2], Mf[r + 3]); + const f16vec4 Pf = f16vec4(exp(sfsh[row / 4 + col * sfshstride] - mfvec)); [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { - const float16_t Pf = float16_t(exp(sfsh[row + vec_idx + col * sfshstride] - Mf[r + vec_idx])); - Lf[r + vec_idx] += Pf; - Pfvec[vec_idx] = Pf; + Lf[r + vec_idx] += Pf[vec_idx]; } + Psh[col * psh_stride + row / 4] = Pf; } - Psh[col * psh_stride + row / 4] = Pfvec; } } @@ -364,8 +388,8 @@ void main() { } // Store SfMat to sfsh and load into Of - const uint osh_stride = row_split * MatBc; - const uint o_offset = gl_SubgroupID * MatBc; + const uint osh_stride = row_split * MatBc / 4; + const uint o_offset = gl_SubgroupID * MatBc / 4; coopMatStore(SfMat, sfsh, o_offset, osh_stride, gl_CooperativeMatrixLayoutRowMajor); barrier(); @@ -385,16 +409,11 @@ void main() { if (d >= d_end) break; const uint hsv_col = 4 * d; - const uint local_hsv = hsv_col - hsv_base; + const uint local_hsv = (hsv_col - hsv_base) / 4; if (!KV_bounds_check || hsv_col < HSV) { const uint d_local = d0 / threads_per_rowgroup; - Of[r][d_local] += ACC_TYPEV4( - sfsh[row * osh_stride + local_hsv], - sfsh[row * osh_stride + local_hsv + 1], - sfsh[row * osh_stride + local_hsv + 2], - sfsh[row * osh_stride + local_hsv + 3] - ); + Of[r][d_local] += ACC_TYPEV4(sfsh[row * osh_stride + local_hsv]); } } } From a875cc20fe5c78b7f65d929ceb1ff9da7e84084f Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 18:42:25 +0100 Subject: [PATCH 21/30] don't store entire K tile in shmem --- .../vulkan-shaders/flash_attn_cm1.comp | 67 ++++++++++++------- 1 file changed, 42 insertions(+), 25 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 413d355f5cb1..dc4db55a13a9 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -54,7 +54,7 @@ shared f16vec4 Psh[Bc * psh_stride]; const uint32_t sfshstride = (HSK <= 128) ? (Br / 4 + 2) : Br / 4; shared ACC_TYPEV4 sfsh[Bc * sfshstride]; -const uint32_t kshstride = HSK_pad / 4 + 2; // in units of f16vec4 +const uint32_t kshstride = MatBr / 4 + 2; // in units of f16vec4 const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups const uint vsh_stride = v_cols; shared f16vec4 ksh[(kshstride >= vsh_stride) ? (Bc * kshstride) : (Bc * vsh_stride)]; @@ -182,27 +182,6 @@ void main() { } } - [[unroll]] for (uint32_t idx = 0; idx < Bc * HSK / 4; idx += gl_WorkGroupSize.x) { - uint32_t d = (idx + tid) % (HSK / 4); - uint32_t c = (idx + tid) / (HSK / 4); - if (c < Bc && d < HSK / 4) { - f16vec4 K_Tf = f16vec4(0); - if (!KV_bounds_check || j * Bc + c < KV) { -#if BLOCK_SIZE > 1 - uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE + 4 * d; - uint ib = coord / BLOCK_SIZE; - uint iqs = (coord % BLOCK_SIZE); - K_Tf = f16vec4(dequantize4(ib, iqs, k_offset, BINDING_IDX_K)); -#else - K_Tf = f16vec4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]); -#endif - } - - ksh[c * kshstride + d] = K_Tf; - } - } - barrier(); - // K * Q^T -> S^T: Bc x HSK_pad * HSK_pad x Br -> Bc x Br // Bc split across workgroup (four subgroups), loop over HSK in chunks of 16: 16 x 16 * 16 x 16 -> 16 x 16 // This is written transposed in order to allow for N being 8 if implementations need it @@ -211,10 +190,48 @@ void main() { coopmat QMat; [[unroll]] for (uint32_t d = 0; d < HSK_pad / 16; ++d) { - coopMatLoad(QMat, Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor); +#if BLOCK_SIZE == 1 + if (KV_bounds_check) { +#endif + [[unroll]] for (uint32_t idx = 0; idx < Bc * MatBr / 4; idx += gl_WorkGroupSize.x) { + uint32_t col_vec = (idx + tid) % (MatBr / 4); + uint32_t row = (idx + tid) / (MatBr / 4); + if (idx + tid < Bc * MatBr / 4) { + f16vec4 K_Tf = f16vec4(0); + if (!KV_bounds_check || j * Bc + row < KV) { +#if BLOCK_SIZE > 1 + uint coord = (j * Bc + row) * k_stride * BLOCK_SIZE + d * 16 + col_vec * 4; + uint ib = coord / BLOCK_SIZE; + uint iqs = (coord % BLOCK_SIZE); + K_Tf = f16vec4(dequantize4(ib, iqs, k_offset, BINDING_IDX_K)); +#else + K_Tf = f16vec4(data_kv4[k_offset / 4 + (j * Bc + row) * k_stride / 4 + d * 16 / 4 + col_vec]); +#endif + } + + ksh[row * kshstride + col_vec] = K_Tf; + } + } + barrier(); +#if BLOCK_SIZE == 1 + } +#endif - uint coord = (gl_SubgroupID * MatBc) * kshstride + d * 16 / 4; - coopMatLoad(KMat, ksh, coord, kshstride, gl_CooperativeMatrixLayoutRowMajor); +#if BLOCK_SIZE == 1 + if (KV_bounds_check) +#endif + { + uint coord = (gl_SubgroupID * MatBc) * kshstride; + coopMatLoad(KMat, ksh, coord, kshstride, gl_CooperativeMatrixLayoutRowMajor); + } +#if BLOCK_SIZE == 1 + else { + const uint coord = k_offset / 4 + (j * Bc + gl_SubgroupID * MatBc) * k_stride / 4 + d * 16 / 4; + coopMatLoad(KMat, data_kv4, coord, k_stride / 4, gl_CooperativeMatrixLayoutRowMajor); + } +#endif + + coopMatLoad(QMat, Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor); SfMat = coopMatMulAdd(KMat, QMat, SfMat); } From 7d75bea05513781b50e189698739383ee1698f4c Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 24 Jan 2026 21:22:47 +0100 Subject: [PATCH 22/30] fixes --- .../vulkan-shaders/flash_attn_cm1.comp | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index dc4db55a13a9..80ee536b4956 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -191,14 +191,15 @@ void main() { [[unroll]] for (uint32_t d = 0; d < HSK_pad / 16; ++d) { #if BLOCK_SIZE == 1 - if (KV_bounds_check) { + if (KV_bounds_check || d * 16 + 16 > HSK) { #endif + barrier(); [[unroll]] for (uint32_t idx = 0; idx < Bc * MatBr / 4; idx += gl_WorkGroupSize.x) { uint32_t col_vec = (idx + tid) % (MatBr / 4); uint32_t row = (idx + tid) / (MatBr / 4); if (idx + tid < Bc * MatBr / 4) { f16vec4 K_Tf = f16vec4(0); - if (!KV_bounds_check || j * Bc + row < KV) { + if ((!KV_bounds_check || j * Bc + row < KV) && (HSK == HSK_pad || d * 16 + col_vec * 4 < HSK)) { #if BLOCK_SIZE > 1 uint coord = (j * Bc + row) * k_stride * BLOCK_SIZE + d * 16 + col_vec * 4; uint ib = coord / BLOCK_SIZE; @@ -218,7 +219,7 @@ void main() { #endif #if BLOCK_SIZE == 1 - if (KV_bounds_check) + if (KV_bounds_check || d * 16 + 16 > HSK) #endif { uint coord = (gl_SubgroupID * MatBc) * kshstride; @@ -421,15 +422,12 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); - [[unroll]] for (uint32_t d0 = d_start; d0 < d_end; d0 += threads_per_rowgroup) { - const uint d = d0 + col_tid; - if (d >= d_end) break; - + [[unroll]] for (uint32_t d_local = 0; d_local < d_per_thread; ++d_local) { + const uint d = d_local * threads_per_rowgroup + col_tid; const uint hsv_col = 4 * d; - const uint local_hsv = (hsv_col - hsv_base) / 4; - if (!KV_bounds_check || hsv_col < HSV) { - const uint d_local = d0 / threads_per_rowgroup; + if (hsv_col >= hsv_base && hsv_col < hsv_base + hsv_per_tile && hsv_col < HSV) { + const uint local_hsv = (hsv_col - hsv_base) / 4; Of[r][d_local] += ACC_TYPEV4(sfsh[row * osh_stride + local_hsv]); } } From a0c9f40821ff635fb7b7ecdc9b40f3d8afe7f1c0 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sun, 25 Jan 2026 07:51:58 +0100 Subject: [PATCH 23/30] load large k tiles to shmem on Nvidia --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 6 +++- .../vulkan-shaders/flash_attn_base.glsl | 1 + .../vulkan-shaders/flash_attn_cm1.comp | 30 ++++++++++++++++++- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 2635d373c66f..7ca1efd9596c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3185,7 +3185,11 @@ static void ggml_vk_load_shaders(vk_device& device) { const uint32_t D_lsb = D ^ (D & (D-1)); uint32_t D_split = std::min(std::min(device->subgroup_size, 8u), D_lsb / 4); - return {wg_size, rows_cols[0], rows_cols[1], hsk, hsv, clamp, D_split, device->subgroup_size}; + // Nvidia prefers shared memory use to load large tiles of K + // AMD prefers loading K directly from global memory + const uint32_t k_load_shmem = device->vendor_id == VK_VENDOR_ID_NVIDIA ? 1 : 0; + + return {wg_size, rows_cols[0], rows_cols[1], hsk, hsv, clamp, D_split, device->subgroup_size, k_load_shmem}; }; #define CREATE_FA(TYPE, NAMELC, FAPATH, SUFFIX) \ diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 6b5b2b248ca4..23a4d2c00586 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -9,6 +9,7 @@ layout (constant_id = 4) const uint32_t HSV = 32; layout (constant_id = 5) const uint32_t Clamp = 0; layout (constant_id = 6) const uint32_t D_split = 16; layout (constant_id = 7) const uint32_t SubGroupSize = 32; +layout (constant_id = 8) const uint32_t K_LOAD_SHMEM = 0; // Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths const uint32_t HSK_pad = (HSK + 15) & ~15; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 80ee536b4956..904991720553 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -54,7 +54,7 @@ shared f16vec4 Psh[Bc * psh_stride]; const uint32_t sfshstride = (HSK <= 128) ? (Br / 4 + 2) : Br / 4; shared ACC_TYPEV4 sfsh[Bc * sfshstride]; -const uint32_t kshstride = MatBr / 4 + 2; // in units of f16vec4 +const uint32_t kshstride = (K_LOAD_SHMEM != 0 ? HSK_pad : MatBr) / 4 + 2; // in units of f16vec4 const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups const uint vsh_stride = v_cols; shared f16vec4 ksh[(kshstride >= vsh_stride) ? (Bc * kshstride) : (Bc * vsh_stride)]; @@ -182,6 +182,29 @@ void main() { } } + if (K_LOAD_SHMEM != 0) { + [[unroll]] for (uint32_t idx = 0; idx < Bc * HSK / 4; idx += gl_WorkGroupSize.x) { + uint32_t d = (idx + tid) % (HSK / 4); + uint32_t c = (idx + tid) / (HSK / 4); + if (c < Bc && d < HSK / 4) { + f16vec4 K_Tf = f16vec4(0); + if (!KV_bounds_check || j * Bc + c < KV) { +#if BLOCK_SIZE > 1 + uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE + 4 * d; + uint ib = coord / BLOCK_SIZE; + uint iqs = (coord % BLOCK_SIZE); + K_Tf = f16vec4(dequantize4(ib, iqs, k_offset, BINDING_IDX_K)); +#else + K_Tf = f16vec4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]); +#endif + } + + ksh[c * kshstride + d] = K_Tf; + } + } + barrier(); + } + // K * Q^T -> S^T: Bc x HSK_pad * HSK_pad x Br -> Bc x Br // Bc split across workgroup (four subgroups), loop over HSK in chunks of 16: 16 x 16 * 16 x 16 -> 16 x 16 // This is written transposed in order to allow for N being 8 if implementations need it @@ -190,6 +213,7 @@ void main() { coopmat QMat; [[unroll]] for (uint32_t d = 0; d < HSK_pad / 16; ++d) { + if (K_LOAD_SHMEM == 0) { #if BLOCK_SIZE == 1 if (KV_bounds_check || d * 16 + 16 > HSK) { #endif @@ -231,6 +255,10 @@ void main() { coopMatLoad(KMat, data_kv4, coord, k_stride / 4, gl_CooperativeMatrixLayoutRowMajor); } #endif + } else { + uint coord = (gl_SubgroupID * MatBc) * kshstride + d * 16 / 4; + coopMatLoad(KMat, ksh, coord, kshstride, gl_CooperativeMatrixLayoutRowMajor); + } coopMatLoad(QMat, Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor); From 1e576ba8b45afc82ee9fae42c29f5bc934939b61 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sun, 25 Jan 2026 08:38:05 +0100 Subject: [PATCH 24/30] adapt shared memory host check function to shader changes --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 35 ++++++++++++------- .../vulkan-shaders/flash_attn_cm1.comp | 2 +- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 7ca1efd9596c..cb5265c61973 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -8351,41 +8351,50 @@ static bool ggml_vk_flash_attn_scalar_shmem_support(const vk_device& device, con const uint32_t total_size = tmpsh + tmpshv4 + masksh + Qf; const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; - VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", total_size=" << total_size << ", supported=" << supported); + VK_LOG_DEBUG("ggml_vk_flash_attn_scalar_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", total_size=" << total_size << ", supported=" << supported); return supported; } -static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const uint32_t hsk, uint32_t hsv, bool f32acc) { +static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const uint32_t hsk, uint32_t hsv, bool f32acc, ggml_type kv_type) { // Needs to be kept up to date on shader changes GGML_UNUSED(hsv); - const uint32_t wg_size = scalar_flash_attention_workgroup_size; - const uint32_t Br = coopmat1_flash_attention_num_large_rows; - const uint32_t Bc = scalar_flash_attention_Bc; + const auto rows_cols = fa_rows_cols(FA_COOPMAT1, hsk, hsv, 0, kv_type, false, false); + const uint32_t Br = rows_cols[0]; + const uint32_t Bc = rows_cols[1]; + + const uint32_t MatBr = 16, MatBc = 16; + + const uint32_t row_split = Bc / MatBc; + const uint32_t wg_size = row_split * device->subgroup_size; const uint32_t hsk_pad = ROUNDUP_POW2(hsk, 16); const uint32_t acctype = f32acc ? 4 : 2; const uint32_t f16vec4 = 8; - const uint32_t tmpsh = wg_size * sizeof(float); - const uint32_t tmpshv4 = wg_size * 4 * acctype; + const uint32_t tmpsh = (Bc / MatBc) * sizeof(float); const uint32_t qstride = hsk_pad / 4 + 2; const uint32_t Qf = Br * qstride * f16vec4; + const uint32_t psh_stride = Br / 4 + 2; + const uint32_t Psh = Bc * psh_stride * f16vec4; + const uint32_t sfshstride = (hsk <= 128) ? (Br + 8) : Br; const uint32_t sfsh = Bc * sfshstride * acctype; - const uint32_t kshstride = hsk_pad / 4 + 2; - const uint32_t ksh = Bc * kshstride * f16vec4; + const bool k_load_shmem = device->vendor_id == VK_VENDOR_ID_NVIDIA; + const uint32_t kshstride = (k_load_shmem ? hsk_pad : MatBr) / 4 + 2; + const uint32_t vsh_stride = MatBc / 4 * row_split; + const uint32_t ksh = ((kshstride >= vsh_stride) ? (Bc * kshstride) : (Bc * vsh_stride)) * f16vec4; - const uint32_t slope = Br * sizeof(float); + const uint32_t slope = Br * acctype; - const uint32_t total_size = tmpsh + tmpshv4 + Qf + sfsh + ksh + slope; + const uint32_t total_size = tmpsh + Qf + Psh + sfsh + ksh + slope; const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; - VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", total_size=" << total_size << ", supported=" << supported); + VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", kv_type=" << kv_type << ", total_size=" << total_size << ", supported=" << supported); return supported; } @@ -8449,7 +8458,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const bool coopmat_shape_supported = (dst->op_params[3] == GGML_PREC_F32 && ctx->device->coopmat_support_16x16x16_f32acc) || (dst->op_params[3] != GGML_PREC_F32 && ctx->device->coopmat_support_16x16x16_f16acc); - const bool coopmat_shmem_supported = ggml_vk_flash_attn_coopmat_shmem_support(ctx->device, HSK, HSV, dst->op_params[3] == GGML_PREC_F32); + const bool coopmat_shmem_supported = ggml_vk_flash_attn_coopmat_shmem_support(ctx->device, HSK, HSV, dst->op_params[3] == GGML_PREC_F32, k->type); if (!coopmat_shape_supported || !coopmat_shmem_supported) { path = FA_SCALAR; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 904991720553..d55be1cf1fa8 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -42,7 +42,7 @@ D_TYPE perElemOpGqaStore(const in uint32_t r, const in uint32_t c, const in D_TY return elem; } -shared FLOAT_TYPE tmpsh[row_split]; +shared float tmpsh[row_split]; const uint32_t qstride = HSK_pad / 4 + 2; // in units of f16vec4 shared f16vec4 Qf[Br * qstride]; From ef86d3adf8ae19fe648accc70c727a06ec9bf7c4 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sun, 25 Jan 2026 12:59:26 +0100 Subject: [PATCH 25/30] remove Bc 32 case --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index cb5265c61973..2bb72b9822a3 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2763,10 +2763,7 @@ static std::array fa_rows_cols(FaCodePath path, uint32_t hsk, uint3 if (small_rows) { return {scalar_flash_attention_num_small_rows, scalar_flash_attention_Bc}; } else { - if ((ggml_is_quantized(type) && (hsk >= 256 || hsv >= 256)) || hsk >= 512 || hsv >= 512) { - return {16, 32}; - } - return {16, 64}; + return {coopmat1_flash_attention_num_large_rows, scalar_flash_attention_Bc}; } } From fcd3a2628434a3296a8eaf3e3afdb0c410d2a70f Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sun, 25 Jan 2026 13:09:22 +0100 Subject: [PATCH 26/30] remove unused variable --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 2bb72b9822a3..30092fde3445 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -8363,7 +8363,6 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co const uint32_t MatBr = 16, MatBc = 16; const uint32_t row_split = Bc / MatBc; - const uint32_t wg_size = row_split * device->subgroup_size; const uint32_t hsk_pad = ROUNDUP_POW2(hsk, 16); From 80a4ac05dfe2593682fe614b39e9110c3d4cca8d Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Mon, 26 Jan 2026 07:34:36 +0100 Subject: [PATCH 27/30] fix missing mask reduction tmspsh barrier --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 1 + 1 file changed, 1 insertion(+) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index d55be1cf1fa8..5d9b7ee36faa 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -170,6 +170,7 @@ void main() { } // skip the block if the mask is entirely -inf bool all_less = subgroupAll(max_mask <= NEG_FLT_MAX_OVER_2); + barrier(); if (gl_SubgroupInvocationID == 0) { tmpsh[gl_SubgroupID] = all_less ? NEG_FLT_MAX_OVER_2 : 0.0f; } From 92051ec92632c376359238c241c3e179447087f8 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Tue, 27 Jan 2026 08:18:56 +0100 Subject: [PATCH 28/30] fix mask bounds check --- .../vulkan-shaders/flash_attn_cm1.comp | 64 ++++++++++--------- 1 file changed, 33 insertions(+), 31 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 5d9b7ee36faa..118f78e4a96d 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -158,13 +158,36 @@ void main() { uint32_t c = (idx + tid) / (Br / 4); uint32_t r = (idx + tid) % (Br / 4); if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { - if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r * 4 < p.nem1)) { - const f16vec4 m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 3) * m_stride + (j * Bc + c)]); + if ((!KV_bounds_check || j * Bc + c < KV)) { + f16vec4 m; + if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { + m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 3) * m_stride + (j * Bc + c)]); + max_mask = max(max(max(max(max_mask, float(m[0])), float(m[1])), float(m[2])), float(m[3])); + } else if (i * Br + r * 4 + 2 < p.nem1) { + m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], + 0.0); + max_mask = max(max(max(max_mask, float(m[0])), float(m[1])), float(m[2])); + } else if (i * Br + r * 4 + 1 < p.nem1) { + m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], + data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], + 0.0, + 0.0); + max_mask = max(max(max_mask, float(m[0])), float(m[1])); + } else if (i * Br + r * 4 < p.nem1) { + m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], + 0.0, + 0.0, + 0.0); + max_mask = max(max_mask, float(m[0])); + } else { + m = f16vec4(0.0); + } mask_cache[idx / WorkGroupSize] = m; - max_mask = max(max(max(max(max_mask, float(m[0])), float(m[1])), float(m[2])), float(m[3])); } } } @@ -282,36 +305,15 @@ void main() { } if ((p.mask_n_head_log2 & MASK_ENABLE_BIT) != 0) { - bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; - [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { uint32_t c = (idx + tid) / (Br / 4); uint32_t r = (idx + tid) % (Br / 4); if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { if (!KV_bounds_check || j * Bc + c < KV) { - if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { - ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); - ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); - sfsh[c * sfshstride + r] += slopes * masks; - } else if (i * Br + r * 4 + 2 < p.nem1) { - ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); - masks.w = ACC_TYPE(0.0f); - ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); - sfsh[c * sfshstride + r] += slopes * masks; - } else if (i * Br + r * 4 + 1 < p.nem1) { - ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); - masks.z = ACC_TYPE(0.0f); - masks.w = ACC_TYPE(0.0f); - ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); - sfsh[c * sfshstride + r] += slopes * masks; - } else if (i * Br + r * 4 < p.nem1) { - ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); - masks.y = ACC_TYPE(0.0f); - masks.z = ACC_TYPE(0.0f); - masks.w = ACC_TYPE(0.0f); - ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); - sfsh[c * sfshstride + r] += slopes * masks; - } + // Mask nem1 bounds check is handled when loading masks + ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); + ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); + sfsh[c * sfshstride + r] += slopes * masks; } } } From 32465b87d1779d9536a64624e53c75d5ef0c2058 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Wed, 28 Jan 2026 10:13:58 +0100 Subject: [PATCH 29/30] fix rowmax f16 under/overflow to inf --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 118f78e4a96d..83d52d19d672 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -362,8 +362,8 @@ void main() { if (KV_bounds_check && j * Bc + col >= KV) { Psh[col * psh_stride + row / 4] = f16vec4(0.0f); } else { - const f16vec4 mfvec = f16vec4(Mf[r], Mf[r + 1], Mf[r + 2], Mf[r + 3]); - const f16vec4 Pf = f16vec4(exp(sfsh[row / 4 + col * sfshstride] - mfvec)); + const vec4 mfvec = vec4(Mf[r], Mf[r + 1], Mf[r + 2], Mf[r + 3]); + const f16vec4 Pf = f16vec4(exp(vec4(sfsh[row / 4 + col * sfshstride]) - mfvec)); [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { Lf[r + vec_idx] += Pf[vec_idx]; } From 7056c66fa7a9bb0595908866027da669f4bb7ae4 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Wed, 28 Jan 2026 10:17:15 +0100 Subject: [PATCH 30/30] fix flash_attn_cm2 BLOCK_SIZE preprocessor directives --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp index 376cb7860e43..54f1b0b62267 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -85,7 +85,7 @@ void main() { tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); -#if defined(BLOCK_SIZE) +#if BLOCK_SIZE > 1 tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, BLOCK_SIZE); tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, BLOCK_SIZE); #endif @@ -98,7 +98,7 @@ void main() { if (Clamp != gl_CooperativeMatrixClampModeConstantNV) { q_stride &= ~7; -#if !defined(BLOCK_SIZE) +#if BLOCK_SIZE == 1 k_stride &= ~7; v_stride &= ~7; #endif