wip

Author: nihui

src/layer/vulkan/sdpa_vulkan.cpp

@@ -22,6 +22,7 @@ SDPA_vulkan::SDPA_vulkan()
 
     pipeline_sdpa_fa = 0;
     use_flash_attention = false;
+    FA_UNROLL_WG_M = 1;
 
     use_cooperative_matrix = false;
     coopmat_M = 0;
@@ -77,19 +78,22 @@ int SDPA_vulkan::create_pipeline(const Option& opt)
 
         // assert coopmat_M != 0 && coopmat_N != 0 && coopmat_K != 0
 
+        FA_UNROLL_WG_M = 2;
+
         // fa
         {
-            std::vector<vk_specialization_type> specializations(1 + 4);
+            std::vector<vk_specialization_type> specializations(1 + 5);
             specializations[0].i = attn_mask;
 
             specializations[1 + 0].u32 = coopmat_M;
             specializations[1 + 1].u32 = coopmat_N;
             specializations[1 + 2].u32 = coopmat_K;
             specializations[1 + 3].u32 = coopmat_subgroup_size;
+            specializations[1 + 4].u32 = FA_UNROLL_WG_M;
 
             pipeline_sdpa_fa = new Pipeline(vkdev);
             pipeline_sdpa_fa->set_subgroup_size(coopmat_subgroup_size);
-            pipeline_sdpa_fa->set_local_size_xyz(coopmat_subgroup_size, 1, 1);
+            pipeline_sdpa_fa->set_local_size_xyz(coopmat_subgroup_size * FA_UNROLL_WG_M, 1, 1);
             pipeline_sdpa_fa->create(LayerShaderType::sdpa_fa_cm, opt, specializations);
         }
     }
@@ -281,6 +285,7 @@ int SDPA_vulkan::destroy_pipeline(const Option& opt)
     }
 
     use_flash_attention = false;
+    FA_UNROLL_WG_M = 1;
 
     use_cooperative_matrix = false;
     coopmat_M = 0;
@@ -430,11 +435,11 @@ int SDPA_vulkan::forward(const std::vector<VkMat>& bottom_blobs, std::vector<VkM
             constants[11].i = top_blob.cstep;
             constants[12].i = attn_mask_blob.cstep;
 
-            const int blocks_x = (src_seqlen + coopmat_M - 1) / (coopmat_M);
-            const int blocks_y = 1; //(out_embed_dim + coopmat_N - 1) / (coopmat_N);
+            const int blocks_x = 1;
+            const int blocks_y = (src_seqlen + coopmat_M * FA_UNROLL_WG_M - 1) / (coopmat_M * FA_UNROLL_WG_M);
 
             VkMat dispatcher;
-            dispatcher.w = (blocks_x * blocks_y) * (coopmat_subgroup_size);
+            dispatcher.w = (blocks_x * blocks_y) * (coopmat_subgroup_size * FA_UNROLL_WG_M);
             dispatcher.h = 1;
             dispatcher.c = num_heads;
 

src/layer/vulkan/sdpa_vulkan.h

@@ -31,6 +31,7 @@ class SDPA_vulkan : public SDPA
     Pipeline* pipeline_sdpa_fa;
 
     bool use_flash_attention;
+    int FA_UNROLL_WG_M;
 
     // cooperative matrix
     bool use_cooperative_matrix;