Skip to content

Commit 1b847b1

Browse files
nihuinihui
authored
sdpa vulkan flash attention optimization without cooperative matrix (Tencent#6538)
1 parent d3a83b0 commit 1b847b1

4 files changed

Lines changed: 577 additions & 49 deletions

File tree

‎src/layer/vulkan/sdpa_vulkan.cpp‎

Lines changed: 107 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -67,59 +67,95 @@ int SDPA_vulkan::create_pipeline(const Option& opt)
6767
use_bf16_cooperative_matrix = true;
6868
}
6969

70-
use_flash_attention = use_cooperative_matrix && (vkdev->info.support_subgroup_ops() & (VK_SUBGROUP_FEATURE_BASIC_BIT | VK_SUBGROUP_FEATURE_SHUFFLE_BIT));
70+
use_flash_attention = (opt.use_fp16_storage || opt.use_fp16_packed || opt.use_bf16_storage || opt.use_bf16_packed);
7171

7272
if (use_flash_attention)
7373
{
74-
int M = 1024;
75-
int N = 1024;
76-
int K = 1024;
77-
78-
if (use_bf16_cooperative_matrix)
79-
{
80-
vkdev->info.get_optimal_cooperative_matrix_mnk(M, N, K, VK_COMPONENT_TYPE_BFLOAT16_KHR, VK_COMPONENT_TYPE_FLOAT32_KHR, VK_SCOPE_SUBGROUP_KHR, FA_coopmat_M, FA_coopmat_N, FA_coopmat_K, FA_coopmat_subgroup_size);
81-
}
82-
else
74+
if (use_cooperative_matrix)
8375
{
84-
vkdev->info.get_optimal_cooperative_matrix_mnk(M, N, K, VK_COMPONENT_TYPE_FLOAT16_KHR, VK_COMPONENT_TYPE_FLOAT32_KHR, VK_SCOPE_SUBGROUP_KHR, FA_coopmat_M, FA_coopmat_N, FA_coopmat_K, FA_coopmat_subgroup_size);
85-
}
76+
int M = 1024;
77+
int N = 1024;
78+
int K = 1024;
8679

87-
// assert FA_coopmat_M != 0 && FA_coopmat_N != 0 && FA_coopmat_K != 0
80+
if (use_bf16_cooperative_matrix)
81+
{
82+
vkdev->info.get_optimal_cooperative_matrix_mnk(M, N, K, VK_COMPONENT_TYPE_BFLOAT16_KHR, VK_COMPONENT_TYPE_FLOAT32_KHR, VK_SCOPE_SUBGROUP_KHR, FA_coopmat_M, FA_coopmat_N, FA_coopmat_K, FA_coopmat_subgroup_size);
83+
}
84+
else
85+
{
86+
vkdev->info.get_optimal_cooperative_matrix_mnk(M, N, K, VK_COMPONENT_TYPE_FLOAT16_KHR, VK_COMPONENT_TYPE_FLOAT32_KHR, VK_SCOPE_SUBGROUP_KHR, FA_coopmat_M, FA_coopmat_N, FA_coopmat_K, FA_coopmat_subgroup_size);
87+
}
8888

89-
if (FA_coopmat_N != FA_coopmat_K || FA_coopmat_subgroup_size < FA_coopmat_N)
90-
{
91-
// not implemented yet
92-
use_flash_attention = false;
89+
// assert FA_coopmat_M != 0 && FA_coopmat_N != 0 && FA_coopmat_K != 0
90+
91+
if (FA_coopmat_N != FA_coopmat_K || FA_coopmat_subgroup_size < FA_coopmat_N)
92+
{
93+
// not implemented yet
94+
use_flash_attention = false;
95+
}
96+
else
97+
{
98+
// fa
99+
FA_UNROLL_SG_M = 2;
100+
101+
FA_UNROLL_WG_M = 2;
102+
103+
std::vector<vk_specialization_type> specializations(1 + 8);
104+
specializations[0].i = attn_mask;
105+
106+
specializations[1 + 0].u32 = FA_coopmat_M;
107+
specializations[1 + 1].u32 = FA_coopmat_N;
108+
specializations[1 + 2].u32 = FA_coopmat_K;
109+
specializations[1 + 3].u32 = FA_coopmat_subgroup_size;
110+
specializations[1 + 4].u32 = FA_UNROLL_SG_M;
111+
specializations[1 + 5].u32 = FA_UNROLL_WG_M;
112+
113+
for (int i = 0; i < 8; i++)
114+
{
115+
int MAX_OUT_CHUNKS = i + 1;
116+
int UNROLL_P_N = std::min(4, FA_coopmat_subgroup_size / FA_coopmat_N);
117+
118+
specializations[1 + 6].u32 = MAX_OUT_CHUNKS;
119+
specializations[1 + 7].u32 = UNROLL_P_N;
120+
121+
pipeline_sdpa_fa[i] = new Pipeline(vkdev);
122+
pipeline_sdpa_fa[i]->set_subgroup_size(FA_coopmat_subgroup_size);
123+
pipeline_sdpa_fa[i]->set_local_size_xyz(FA_coopmat_subgroup_size * FA_UNROLL_WG_M, 1, 1);
124+
pipeline_sdpa_fa[i]->create(LayerShaderType::sdpa_fa_cm, opt, specializations);
125+
}
126+
}
93127
}
94128
else
95129
{
96-
// fa
97-
FA_UNROLL_SG_M = 2;
130+
FA_coopmat_M = 4;
131+
FA_coopmat_N = 32;
132+
FA_coopmat_K = 32;
133+
FA_UNROLL_WG_M = 4;
134+
const int subgroup_size = vkdev->info.subgroup_size();
98135

99-
FA_UNROLL_WG_M = 2;
136+
// assert FA_coopmat_N == FA_coopmat_K
137+
// assert local_size % FA_coopmat_M == 0
100138

101-
std::vector<vk_specialization_type> specializations(1 + 8);
139+
// fa
140+
std::vector<vk_specialization_type> specializations(1 + 6);
102141
specializations[0].i = attn_mask;
103142

104143
specializations[1 + 0].u32 = FA_coopmat_M;
105144
specializations[1 + 1].u32 = FA_coopmat_N;
106145
specializations[1 + 2].u32 = FA_coopmat_K;
107-
specializations[1 + 3].u32 = FA_coopmat_subgroup_size;
108-
specializations[1 + 4].u32 = FA_UNROLL_SG_M;
109-
specializations[1 + 5].u32 = FA_UNROLL_WG_M;
146+
specializations[1 + 3].u32 = subgroup_size;
147+
specializations[1 + 4].u32 = FA_UNROLL_WG_M;
110148

111149
for (int i = 0; i < 8; i++)
112150
{
113151
int MAX_OUT_CHUNKS = i + 1;
114-
int UNROLL_P_N = std::min(4, FA_coopmat_subgroup_size / FA_coopmat_N);
115152

116-
specializations[1 + 6].u32 = MAX_OUT_CHUNKS;
117-
specializations[1 + 7].u32 = UNROLL_P_N;
153+
specializations[1 + 5].u32 = MAX_OUT_CHUNKS;
118154

119155
pipeline_sdpa_fa[i] = new Pipeline(vkdev);
120-
pipeline_sdpa_fa[i]->set_subgroup_size(FA_coopmat_subgroup_size);
121-
pipeline_sdpa_fa[i]->set_local_size_xyz(FA_coopmat_subgroup_size * FA_UNROLL_WG_M, 1, 1);
122-
pipeline_sdpa_fa[i]->create(LayerShaderType::sdpa_fa_cm, opt, specializations);
156+
pipeline_sdpa_fa[i]->set_subgroup_size(subgroup_size);
157+
pipeline_sdpa_fa[i]->set_local_size_xyz(subgroup_size * FA_UNROLL_WG_M, 1, 1);
158+
pipeline_sdpa_fa[i]->create(LayerShaderType::sdpa_fa, opt, specializations);
123159
}
124160
}
125161
}
@@ -404,7 +440,7 @@ int SDPA_vulkan::forward(const std::vector<VkMat>& bottom_blobs, std::vector<VkM
404440
if (top_blob.empty())
405441
return -100;
406442

407-
// fa
443+
if (use_cooperative_matrix)
408444
{
409445
std::vector<VkMat> bindings(5);
410446
bindings[0] = query;
@@ -442,6 +478,46 @@ int SDPA_vulkan::forward(const std::vector<VkMat>& bottom_blobs, std::vector<VkM
442478

443479
cmd.record_pipeline(pipeline, bindings, constants, dispatcher);
444480
}
481+
else
482+
{
483+
std::vector<VkMat> bindings(5);
484+
bindings[0] = query;
485+
bindings[1] = key;
486+
bindings[2] = value;
487+
bindings[3] = top_blob;
488+
bindings[4] = attn_mask_blob;
489+
490+
std::vector<vk_constant_type> constants(13);
491+
constants[0].f = _scale;
492+
constants[1].i = src_seqlen;
493+
constants[2].i = dst_seqlen;
494+
constants[3].i = embed_dim;
495+
constants[4].i = out_embed_dim;
496+
constants[5].i = num_heads;
497+
constants[6].i = attn_mask_blob.dims && attn_mask_blob.c > 1 ? 3 : attn_mask_blob.dims;
498+
constants[7].i = num_heads_per_group;
499+
constants[8].i = query.cstep;
500+
constants[9].i = key.cstep;
501+
constants[10].i = value.cstep;
502+
constants[11].i = top_blob.cstep;
503+
constants[12].i = attn_mask_blob.cstep;
504+
505+
const int subgroup_size = vkdev->info.subgroup_size();
506+
507+
const int blocks_x = 1;
508+
const int blocks_y = (src_seqlen + FA_coopmat_M - 1) / FA_coopmat_M;
509+
510+
VkMat dispatcher;
511+
dispatcher.w = (blocks_x * blocks_y) * (subgroup_size * FA_UNROLL_WG_M);
512+
dispatcher.h = 1;
513+
dispatcher.c = num_heads;
514+
515+
const int MAX_OUT_CHUNKS = (out_embed_dim + FA_coopmat_N - 1) / FA_coopmat_N;
516+
517+
const Pipeline* pipeline = pipeline_sdpa_fa[MAX_OUT_CHUNKS - 1];
518+
519+
cmd.record_pipeline(pipeline, bindings, constants, dispatcher);
520+
}
445521

446522
if (kv_cache)
447523
{

‎src/layer/vulkan/sdpa_vulkan.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ class SDPA_vulkan : public SDPA
2828
Pipeline* pipeline_sdpa_qk_cross;
2929
Pipeline* pipeline_sdpa_qkv_cross;
3030

31+
// flash attention
3132
Pipeline* pipeline_sdpa_fa[8];
3233

3334
bool use_flash_attention;

0 commit comments

Comments
 (0)