[WebGPU]: Add Intel subgroup prefill flash-attention kernel

onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <string>
+
 #include "contrib_ops/cpu/bert/multihead_attention_helper.h"
 #include "contrib_ops/webgpu/bert/kv_cache_block_quant_int8.h"
 #include "contrib_ops/webgpu/bert/flash_attention.h"
@@ -10,6 +12,7 @@
 #include "contrib_ops/webgpu/webgpu_contrib_kernels.h"
 
 #include "core/providers/webgpu/webgpu_supported_types.h"
+#include "core/providers/webgpu/webgpu_utils.h"
 
 using namespace onnxruntime::webgpu;
 using namespace ::onnxruntime::common;
@@ -538,6 +541,149 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont
   return context.RunProgram(program);
 }
 
+Status FlashAttentionPrefillProgram::GenerateShaderCode(ShaderHelper& shader) const {
+  shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
+  shader.AddInput("present_key", ShaderUsage::UseUniform);
+  shader.AddInput("present_value", ShaderUsage::UseUniform);
+  if (has_attention_bias_) {
+    shader.AddInput("attention_bias", ShaderUsage::UseUniform);
+  }
+  if (use_seqlen_k_) {
+    shader.AddInput("seqlens_k", ShaderUsage::None);
+  }
+  if (use_seqlens_q_) {
+    shader.AddInput("seqlens_q", ShaderUsage::None);
+  }
+  if (has_head_sink_) {
+    shader.AddInput("head_sink", ShaderUsage::UseUniform);
+  }
+  shader.AddOutput("output", ShaderUsage::UseUniform);
+
+  return WGSL_TEMPLATE_APPLY(shader, "bert/flash_attention_prefill.wgsl.template",
+                             WGSL_TEMPLATE_PARAMETER(has_attention_bias, has_attention_bias_),
+                             WGSL_TEMPLATE_PARAMETER(has_head_sink, has_head_sink_),
+                             WGSL_TEMPLATE_PARAMETER(has_local_window, has_local_window_),
+                             WGSL_TEMPLATE_PARAMETER(is_unidirectional, is_unidirectional_),
+                             WGSL_TEMPLATE_PARAMETER(kv_step_param, kv_step_),
+                             WGSL_TEMPLATE_PARAMETER(q_BNSH, q_BNSH_),
+                             WGSL_TEMPLATE_PARAMETER(qk_f32, true),
+                             WGSL_TEMPLATE_PARAMETER(qkv_head_size, qkv_head_size_),
+                             WGSL_TEMPLATE_PARAMETER(qkv_num_heads, qkv_num_heads_),
+                             WGSL_TEMPLATE_PARAMETER(shuffle_chunk_param, shuffle_chunk_),
+                             WGSL_TEMPLATE_PARAMETER(use_seqlen_k, use_seqlen_k_),
+                             WGSL_TEMPLATE_PARAMETER(use_seqlens_q, use_seqlens_q_));
+}
+
+bool CanApplyFlashAttentionPrefill(bool is_intel,
+                                   bool kv_cache_quantization_enabled) {
+  return is_intel && !kv_cache_quantization_enabled;
+}
+
+Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context,
+                                  const Tensor* Q,
+                                  const Tensor* present_key,
+                                  const Tensor* present_value,
+                                  const Tensor* attention_bias,
+                                  const Tensor* seqlen_k,
+                                  Tensor* attn_output,
+                                  const WebgpuAttentionParameters& parameters,
+                                  bool has_attention_bias,
+                                  bool use_seqlen_k,
+                                  bool has_local_window,
+                                  int local_window_size,
+                                  uint32_t present_sequence_length,
+                                  uint32_t subgroup_min_size,
+                                  float alpha,
+                                  uint32_t attn_bias_dim0,
+                                  uint32_t attn_bias_dim1,
+                                  uint32_t attn_bias_dim3,
+                                  bool use_seqlens_q,
+                                  const Tensor* seqlens_q,
+                                  bool q_BNSH,
+                                  bool has_head_sink,
+                                  const Tensor* head_sink) {
+  const auto& adapter_info = context.AdapterInfo();
+  const bool is_intel_xe2 = adapter_info.vendor == std::string_view{"intel"} &&
+                            (adapter_info.architecture == std::string_view{"xe-2lpg"} ||
+                             adapter_info.architecture == std::string_view{"xe-2hpg"});
+  const bool is_intel_xe3 = adapter_info.vendor == std::string_view{"intel"} &&
+                            (adapter_info.architecture == std::string_view{"xe-3lpg"} ||
+                             adapter_info.architecture == std::string_view{"xe-3lpg-xs"});
+  // Default: 64 threads, 16 KV rows per step.
+  uint32_t workgroup_size = 64u;
+  uint32_t kv_step = 16u;
+  // Xe2: query length <= 128 uses 128 threads and 64 KV rows per step;
+  //      longer queries use 256 threads and 32 KV rows per step.
+  // Xe3: 128 threads, 64 KV rows per step for all query lengths.
+  if (is_intel_xe3 || (is_intel_xe2 && parameters.sequence_length_ <= 128)) {
+    workgroup_size = 128u;
+    kv_step = 64u;
+  } else if (is_intel_xe2) {
+    workgroup_size = 256u;
+    kv_step = 32u;
+  }
+  // The K tile (full head) and V slice (up to 16 vec4) are staged in workgroup memory for kv_step rows.
+  // If that exceeds the device limit, fall back to the default config (64 threads, 16 KV rows per step).
+  const uint64_t head_size_vec = static_cast<uint64_t>(parameters.head_size_) / 4;
+  const uint64_t vec4_bytes = (Q->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ? 2u : 4u) * 4u;
+  if ((head_size_vec + std::min<uint64_t>(head_size_vec, 16)) * kv_step * vec4_bytes >
+      context.DeviceLimits().maxComputeWorkgroupStorageSize) {
+    workgroup_size = 64u;
+    kv_step = 16u;
+  }
+  // Values below 8 select direct shared-memory reads, including without subgroups.
+  const uint32_t shuffle_chunk = std::min(kv_step, subgroup_min_size);
+  const uint32_t num_seq_tile = CeilDiv(static_cast<uint32_t>(parameters.sequence_length_), workgroup_size);
+  const uint32_t dispatch_size =
+      parameters.batch_size_ * parameters.num_heads_ * num_seq_tile;
+
+  FlashAttentionPrefillProgram program{"FlashAttentionPrefill",
+                                       has_attention_bias,
+                                       parameters.is_unidirectional_,
+                                       parameters.head_size_,
+                                       parameters.num_heads_,
+                                       use_seqlen_k,
+                                       kv_step,
+                                       shuffle_chunk,
+                                       has_local_window,
+                                       use_seqlens_q,
+                                       q_BNSH,
+                                       has_head_sink};
+  program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, 4},
+                     {present_key, ProgramTensorMetadataDependency::TypeAndRank, 4},
+                     {present_value, ProgramTensorMetadataDependency::TypeAndRank, 4}});
+  if (has_attention_bias) {
+    program.AddInputs({{attention_bias, ProgramTensorMetadataDependency::TypeAndRank}});
+  }
+  if (use_seqlen_k) {
+    program.AddInputs({{seqlen_k, ProgramTensorMetadataDependency::None}});
+  }
+  if (use_seqlens_q) {
+    program.AddInputs({{seqlens_q, ProgramTensorMetadataDependency::None}});
+  }
+  if (has_head_sink) {
+    program.AddInputs({{head_sink, ProgramTensorMetadataDependency::Type}});
+  }
+  program.AddOutputs({{attn_output, ProgramTensorMetadataDependency::TypeAndRank, 4}});
+  program.SetDispatchGroupSize(dispatch_size)
+      .SetWorkgroupSize(workgroup_size)
+      .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_,
+                 use_seqlen_k, workgroup_size, kv_step, shuffle_chunk,
+                 has_local_window, use_seqlens_q, q_BNSH, has_head_sink)
+      .AddUniformVariables({{static_cast<uint32_t>(parameters.sequence_length_)},
+                            {static_cast<uint32_t>(parameters.total_sequence_length_)},
+                            {static_cast<uint32_t>(present_sequence_length)},
+                            {static_cast<uint32_t>(parameters.batch_size_)},
+                            {static_cast<uint32_t>(parameters.n_reps)},
+                            {alpha},
+                            {num_seq_tile},
+                            {attn_bias_dim0},
+                            {attn_bias_dim1},
+                            {attn_bias_dim3},
+                            {static_cast<uint32_t>(has_local_window ? local_window_size : 0)}});
+  return context.RunProgram(program);
+}
+
 Status FlashAttentionDecodeQKVProgram::GenerateShaderCode(ShaderHelper& shader) const {
   const auto& q = shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
   const auto& present_key = shader.AddInput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
@@ -1047,6 +1193,17 @@ Status ApplyDensePrefill(const Tensor* Q, Tensor* present_key, Tensor* present_v
     attn_bias_dim3 = static_cast<uint32_t>(bias_shape[3]);
   }
 
+  const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"};
+  const uint32_t subgroup_min_size = has_subgroups ? context.AdapterInfo().subgroupMinSize : 0u;
+  // TODO(daijh): Merge back into `FlashAttentionProgram`.
+  if (CanApplyFlashAttentionPrefill(is_intel, kv_cache_quantization_enabled)) {
+    return ApplyFlashAttentionPrefill(
+        context, Q, present_key, present_value, attention_bias, seqlen_k, attn_output,
+        parameters, has_attention_bias, use_seqlen_k, has_local_window, local_window_size,
+        present_sequence_length, subgroup_min_size, alpha, attn_bias_dim0, attn_bias_dim1, attn_bias_dim3,
+        use_seqlens_q, seqlens_q, q_BNSH, has_head_sink, head_sink);
+  }
+
   program.SetDispatchGroupSize(parameters.batch_size_ * parameters.num_heads_ * num_seq_tile)
       .SetWorkgroupSize(prefill_tile_size)
       .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_,

onnxruntime/contrib_ops/webgpu/bert/flash_attention.h

@@ -205,6 +205,88 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont
                                        uint32_t max_num_blocks_per_seq,
                                        const Tensor* cumulative_seqlens_q = nullptr);
 
+class FlashAttentionPrefillProgram final : public Program<FlashAttentionPrefillProgram> {
+ public:
+  FlashAttentionPrefillProgram(const std::string& kernel_name,
+                               bool has_attention_bias,
+                               bool is_unidirectional,
+                               int qkv_head_size,
+                               int qkv_num_heads,
+                               bool use_seqlen_k,
+                               uint32_t kv_step,
+                               uint32_t shuffle_chunk,
+                               bool has_local_window,
+                               bool use_seqlens_q,
+                               bool q_BNSH,
+                               bool has_head_sink)
+      : Program{kernel_name},
+        has_attention_bias_(has_attention_bias),
+        is_unidirectional_(is_unidirectional),
+        qkv_head_size_(qkv_head_size),
+        qkv_num_heads_(qkv_num_heads),
+        use_seqlen_k_(use_seqlen_k),
+        kv_step_(kv_step),
+        shuffle_chunk_(shuffle_chunk),
+        has_local_window_(has_local_window),
+        use_seqlens_q_(use_seqlens_q),
+        q_BNSH_(q_BNSH),
+        has_head_sink_(has_head_sink) {}
+
+  Status GenerateShaderCode(ShaderHelper& sh) const override;
+
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"new_sequence_length", ProgramUniformVariableDataType::Uint32},
+                                          {"total_sequence_length", ProgramUniformVariableDataType::Uint32},
+                                          {"present_sequence_length", ProgramUniformVariableDataType::Uint32},
+                                          {"batch_size", ProgramUniformVariableDataType::Uint32},
+                                          {"n_reps", ProgramUniformVariableDataType::Uint32},
+                                          {"alpha", ProgramUniformVariableDataType::Float32},
+                                          {"num_seq_tile", ProgramUniformVariableDataType::Uint32},
+                                          {"attn_bias_dim0", ProgramUniformVariableDataType::Uint32},
+                                          {"attn_bias_dim1", ProgramUniformVariableDataType::Uint32},
+                                          {"attn_bias_dim3", ProgramUniformVariableDataType::Uint32},
+                                          {"local_window_size", ProgramUniformVariableDataType::Uint32});
+
+ private:
+  bool has_attention_bias_;
+  bool is_unidirectional_;
+  int qkv_head_size_;
+  int qkv_num_heads_;
+  bool use_seqlen_k_;
+  uint32_t kv_step_;
+  uint32_t shuffle_chunk_;
+  bool has_local_window_;
+  bool use_seqlens_q_;
+  bool q_BNSH_;
+  bool has_head_sink_;
+};
+
+bool CanApplyFlashAttentionPrefill(bool is_intel,
+                                   bool kv_cache_quantization_enabled);
+
+Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context,
+                                  const Tensor* Q,
+                                  const Tensor* present_key,
+                                  const Tensor* present_value,
+                                  const Tensor* attention_bias,
+                                  const Tensor* seqlen_k,
+                                  Tensor* attn_output,
+                                  const WebgpuAttentionParameters& parameters,
+                                  bool has_attention_bias,
+                                  bool use_seqlen_k,
+                                  bool has_local_window,
+                                  int local_window_size,
+                                  uint32_t present_sequence_length,
+                                  uint32_t subgroup_min_size,
+                                  float alpha,
+                                  uint32_t attn_bias_dim0,
+                                  uint32_t attn_bias_dim1,
+                                  uint32_t attn_bias_dim3,
+                                  bool use_seqlens_q,
+                                  const Tensor* seqlens_q,
+                                  bool q_BNSH,
+                                  bool has_head_sink,
+                                  const Tensor* head_sink);
+
 class FlashAttentionDecodeQKVProgram final : public Program<FlashAttentionDecodeQKVProgram> {
  public:
   FlashAttentionDecodeQKVProgram(const std::string& kernel_name,