From ba9e27159f2ed93945ab8634127518d407a58fda Mon Sep 17 00:00:00 2001 From: Jie Chen Date: Tue, 27 Oct 2026 19:38:33 +0800 Subject: [PATCH] FlashAttention: use subgroup size 32 on Intel Xe2 and Xe3 An explicit subgroup size lets Tint use a constant lane clamp bound for subgroupShuffle, so IGC keeps direct register addressing. --- .../contrib_ops/webgpu/bert/flash_attention.cc | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 7fbd162b2e39d..cfd084d58c2d9 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -633,10 +633,11 @@ Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, 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"} && + const bool is_intel = adapter_info.vendor == std::string_view{"intel"}; + const bool is_intel_xe2 = is_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"} && + const bool is_intel_xe3 = is_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. @@ -700,6 +701,17 @@ Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, program.AddInputs({{head_sink, ProgramTensorMetadataDependency::Type}}); } program.AddOutputs({{attn_output, ProgramTensorMetadataDependency::TypeAndRank, 4}}); + + // An explicit subgroup size lets Tint use a constant lane clamp bound for subgroupShuffle, which keeps IGC from + // falling back to indirect register addressing. + constexpr uint32_t kIntelSubgroupSize = 32; + if ((is_intel_xe2 || is_intel_xe3) && context.HasFeature(wgpu::FeatureName::Subgroups) && + context.HasFeature(wgpu::FeatureName::SubgroupSizeControl) && + adapter_info.subgroupMinSize <= kIntelSubgroupSize && kIntelSubgroupSize <= adapter_info.subgroupMaxSize && + workgroup_size % kIntelSubgroupSize == 0) { + program.SetSubgroupSize(kIntelSubgroupSize); + } + program.SetDispatchGroupSize(dispatch_size) .SetWorkgroupSize(workgroup_size) .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_,