webgpu: Normalize indirect-dispatch buffer in flash attention

onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc

@@ -16,6 +16,48 @@ namespace onnxruntime {
 namespace contrib {
 namespace webgpu {
 
+// GPU twin of ProgramManager::NormalizeDispatchGroupSize for indirect-dispatch
+// programs that compute their group count on the device. Mirrors the host
+// helper's three tiers (1D, 2D sqrt, 3D cbrt). The caller passes the intended
+// (x, y, z) layout; the helper writes the chosen layout into a storage output
+// named `indirect_buffer`.
+//
+// Caller requirement: declare a storage output named exactly `indirect_buffer`
+// of type `array<u32>` with >= 3 elements.
+//
+// Intentional drifts from the host helper:
+//   - per-dim `limit` is hardcoded to the WebGPU spec floor (65535) instead of
+//     being read from device limits;
+//   - intermediate `size` is f32 (no f64 in core WGSL);
+//   - no error surface past the 3D tier - the driver rejects the dispatch
+//     instead.
+//
+// Consumers are unaffected by the chosen layout: shader_helper always flattens
+// workgroup_id (x, y, z) into a single linear workgroup_idx.
+constexpr const char kNormalizeDispatchGroupSizeFn[] = R"(
+fn normalize_dispatch_group_size(x: u32, y: u32, z: u32) {
+  let limit = 65535u;  // WebGPU spec maxComputeWorkgroupsPerDimension
+  if (x <= limit && y <= limit && z <= limit) {
+    indirect_buffer[0] = x;
+    indirect_buffer[1] = y;
+    indirect_buffer[2] = z;
+    return;
+  }
+  let size = f32(x) * f32(y) * f32(z);
+  let dispatch_avg_2d = u32(ceil(sqrt(size)));
+  if (dispatch_avg_2d <= limit) {
+    indirect_buffer[0] = dispatch_avg_2d;
+    indirect_buffer[1] = dispatch_avg_2d;
+    indirect_buffer[2] = 1u;
+    return;
+  }
+  let dispatch_avg_3d = u32(ceil(pow(size, 1.0 / 3.0)));
+  indirect_buffer[0] = dispatch_avg_3d;
+  indirect_buffer[1] = dispatch_avg_3d;
+  indirect_buffer[2] = dispatch_avg_3d;
+}
+)";
+
 Status SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram::GenerateShaderCode(ShaderHelper& sh) const {
   const auto& packed_qkv = sh.AddInput("packed_qkv", ShaderUsage::UseUniform);
   const auto& seqlens = sh.AddInput("seqlens", ShaderUsage::UseUniform);
@@ -28,6 +70,7 @@ Status SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram::GenerateShaderCode(Sha
 
   if (prepare_indirect_dispatch_) {
     sh.AddOutput("indirect_buffer", ShaderUsage::None);
+    sh.AdditionalImplementation() << kNormalizeDispatchGroupSizeFn;
   }
 
   return WGSL_TEMPLATE_APPLY(sh, "bert/split_packed_qkv_with_rotary_embedding_and_copykv.wgsl.template",
@@ -85,15 +128,12 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const {
     shader.MainFunctionBody() << "  let present_offset = " << present_key.IndicesToOffset("present_key_indices_t(batch, num_head_id, sequence_id, head_size_id)") << ";\n";
   }
 
-  // Add indirect dispatch logic for thread 0
   if (prepare_indirect_dispatch_) {
-    // TODO: Add NormalizeDispatchGroupSize logic here to avoid exceeding max dispatch size.
-    shader.MainFunctionBody() << "  // Prepare indirect dispatch buffer for thread 0\n"
-                              << "  if (global_idx == 0u) {\n"
+    shader.AdditionalImplementation() << kNormalizeDispatchGroupSizeFn;
+    // Match the direct-dispatch shape (batch, num_heads, num_total_seq_length_tile).
+    shader.MainFunctionBody() << "  if (global_idx == 0u) {\n"
                               << "    let num_total_seq_length_tile = (total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size;\n"
-                              << "    indirect_buffer[0] = num_total_seq_length_tile;\n"
-                              << "    indirect_buffer[1] = uniforms.num_heads;\n"
-                              << "    indirect_buffer[2] = 1u;\n"
+                              << "    normalize_dispatch_group_size(num_total_seq_length_tile, uniforms.num_heads, uniforms.copy_kv_shape_shape[0]);\n"
                               << "  }\n\n";
   }
 

onnxruntime/contrib_ops/webgpu/bert/split_packed_qkv_with_rotary_embedding_and_copykv.wgsl.template

@@ -41,12 +41,10 @@ $MAIN {
 #endif
 
 #if prepare_indirect_dispatch
-  // Prepare indirect dispatch buffer for thread 0
   if (global_idx == 0u) {
     let num_total_seq_length_tile = (total_seqlen + uniforms.tile_size - 1u) / uniforms.tile_size;
-    indirect_buffer[0] = num_total_seq_length_tile;
-    indirect_buffer[1] = uniforms.num_heads;
-    indirect_buffer[2] = 1u;
+    // Match the direct-dispatch shape (batch, num_heads, num_total_seq_length_tile).
+    normalize_dispatch_group_size(num_total_seq_length_tile, uniforms.num_heads, uniforms.query_shape[0]);
   }
 #endif