ggml-webgpu: tune subgroup split (d_split) in flash_attn_vec (#25418)

diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
index 80a16d16d23662666943fc4baa3b26bf1cff0d73..d7692363a1dedf93119da6c9c351e9f5b78306bd 100644 (file)
--- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
@@ -2821,23 +2821,16 @@ class ggml_webgpu_shader_lib {
             variant.resize(variant.size() - (sizeof("_mask") - 1));
             variant += "_mask_blk";
         }
-        uint32_t vec_ne = 1u;
-        if (key.common.k_type == GGML_TYPE_F16 && key.common.v_type == GGML_TYPE_F16 &&
-            key.common.head_dim_qk == key.common.head_dim_v) {
-            switch (key.common.head_dim_qk) {
-                case 64:
-                case 192:
-                case 576:
-                    vec_ne = 2u;
-                    break;
-                case 96:
-                    vec_ne = 4u;
-                    break;
-                default:
-                    break;
-            }
+
+        uint32_t d_split = context.min_subgroup_size;
+        if (key.common.k_type == GGML_TYPE_F16 && key.common.v_type == GGML_TYPE_F16) {
+            const uint32_t D     = key.common.head_dim_qk | key.common.head_dim_v;
+            const uint32_t D_lsb = D & (~(D - 1u));
+            d_split              = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));
         }
-        defines.push_back(std::string("VEC_NE=") + std::to_string(vec_ne) + "u");
+
+        defines.push_back(std::string("D_SPLIT=") + std::to_string(d_split));
+        variant += "_dsplit" + std::to_string(d_split);
 
         auto            pipeline_decisions = std::make_shared<ggml_webgpu_flash_attn_vec_decisions>(decisions);
         webgpu_pipeline pipeline =

diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl
index 30ed97cca0c356f7f5ddc2233fe5985592fe8c4c..d5127624196b9d844e9190c7771fce9c137d07ab 100644 (file)
--- a/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl
+++ b/ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl
@@ -39,9 +39,6 @@ enable subgroups;
 #define KV_GRANULARITY 8
 #define KV_TILE 16
 #define WG_SIZE 64
-#ifndef VEC_NE
-#define VEC_NE 4u
-#endif
 
 #define KV_BLOCKS (KV_TILE / KV_GRANULARITY)
 
@@ -367,11 +364,11 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
 
       // accumulate q block * k block into registers across the entire KV tile
       if (!skip_tile) {
-        let num_of_threads = subgroup_size / VEC_NE;
+        let num_of_threads:u32 = D_SPLIT;
         let tx = sg_inv_id % num_of_threads;
         let ty = sg_inv_id / num_of_threads;
           if (subgroup_id == 0u && q_row_start < params.seq_len_q) {
-              for (var kv_base : u32 = 0u; kv_base < KV_TILE; kv_base += VEC_NE) {
+              for (var kv_base : u32 = 0u; kv_base < KV_TILE; kv_base += subgroup_size / D_SPLIT) {
                   let kv_idx = kv_base + ty;
                   var partial_sum: f32 = 0.0;
                   let kv_valid = kv_idx < KV_TILE && (kv_tile + kv_idx) < params.seq_len_kv;
@@ -486,15 +483,18 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
       if (!skip_tile) {
           // we have P (KV_TILE) in inter_shmem and V (KV_TILE x head_dim_v) in kv_shmem
           // we want to compute O += P * V across the full KV tile
-          let ne_threads : u32 = VEC_NE;
+          let ne_threads : u32 = subgroup_size / D_SPLIT;
           let nl_threads = max(1u, subgroup_size / ne_threads);
           let tx_pv = sg_inv_id % nl_threads;
           let ty_pv = sg_inv_id / nl_threads;
           if (subgroup_id == 0u && q_row_start < params.seq_len_q) {
               for (var vec_col = tx_pv; vec_col < (HEAD_DIM_V / 4u); vec_col += nl_threads) {
                   var lo = vec4<f32>(0.0, 0.0, 0.0, 0.0);
-                  for (var cc = 0u; cc < KV_TILE / ne_threads; cc += 1u) {
+                  for (var cc = 0u; cc * ne_threads < KV_TILE; cc += 1u) {
                       let kv_idx = cc * ne_threads + ty_pv;
+                      if (kv_idx >= KV_TILE) {
+                          continue;
+                      }
                       let v_row = kv_tile + kv_idx;
                       if (v_row >= params.seq_len_kv) {
                           continue;