[WebGPU]: Add Intel subgroup prefill flash-attention kernel - #32702
Open
Jianhui Dai (daijh) wants to merge 1 commit into
Open
Jianhui Dai (daijh) wants to merge 1 commit into
Jianhui Dai (daijh) wants to merge 1 commit into
Conversation
Introduce FlashAttentionPrefillSimpleProgram and its template (flash_attention_prefill_simple.wgsl.template). This is a simplified version of the FlashAttentionProgram kernel. (TODO: Merge back into the original FlashAttentionProgram) Key implementation details: - Latency Hiding: Only K is loaded before the QK^T + online-softmax stage. V is fetched just before the PV matmul, allowing global-memory latency to overlap with softmax compute. - Avoids Bank Conflicts: K and V tiles are stored in transposed layout ([head_size_vec][kv_step]), ensuring the subgroup column is the inner/ fastest dimension. Adjacent lanes hit adjacent shared-memory addresses. - Static Shuffle Codepath: Unconditionally uses a subgroup-shuffle path based on the host-configured minimum subgroup size. This avoids dynamic runtime SubgroupSize detection inside the WGSL kernel. - Configurations: The host determines kv_step (8 or 16) based on the adapter's guaranteed minimum subgroup size (needs minimum of 8). Workgroup size is 128 for causal and 256 for non-causal prefill. - Unsupported Features: Does not support quantized KV cache (int8/ TurboQuant) or the shared-memory fallback path (to be supported). Integration & Selection: - Adds CanApplyFlashAttentionPrefillSimple and ApplyFlashAttentionPrefillSimple, mirroring the MatMulNBits pattern. - ApplyFlashAttention selects this kernel only if 'is_intel' is true, subgroup_min_size >= 8, and the KV cache is unquantized. Otherwise, it falls back to the existing FlashAttentionProgram. **Test on Phi-4-mini, EdgeLLMOnDeviceModel\2026.5.8.1, Panther Lake:** | Metric | Baseline Prefill TPS | Optimized Prefill TPS | Improve | |---|---|---|---| | Prefill-128 | 1289.89 | 1264.72 | -1.95% | | Prefill-256 | 1647.14 | 1631.21 | -0.97% | | Prefill-512 | 1915.64 | 2012.29 | 5.05% | | Prefill-1024 | 1834.62 | 1935.33 | 5.49% | | Prefill-2048 | 1745.07 | 1947.23 | 11.58% | | Prefill-4096 | 1508.13 | 1812.86 | 20.21% |
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
Author
|
Prerequisite: #32605 |
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The kernel has out-of-bounds subgroup accesses, incorrect long-context fp16 accumulation, and no automated Intel execution coverage.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 1
Open (3)
What changed in this PR
Adds an Intel-optimized WebGPU flash-attention prefill path using subgroup shuffles and deferred V loading.
Changes:
- Adds the simplified WGSL prefill kernel.
- Adds host-side configuration and Intel-specific dispatch.
- Retains the existing kernel as fallback.
| File | Description |
|---|---|
flash_attention.h |
Declares the new program and helpers. |
flash_attention.cc |
Configures and selects the Intel kernel. |
flash_attention_prefill_simple.wgsl.template |
Implements tiled attention and online softmax. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+154
to
+158
| var previous_max : q_element_t = sink_value; | ||
| var previous_denom : q_element_t = 1; | ||
| #else | ||
| var previous_max : q_element_t = min_value; | ||
| var previous_denom : q_element_t = 0; |
Contributor
Author
There was a problem hiding this comment.
This exactly mirrors existing flash_attention.wgsl.template
|
|
||
| ORT_RETURN_IF_ERROR(context.RunProgram(program)); | ||
| // TODO: Merge back into `FlashAttentionProgram`. | ||
| if (CanApplyFlashAttentionPrefillSimple(is_intel, subgroup_min_size, kv_cache_quantization_enabled)) { |
|
|
||
| ORT_RETURN_IF_ERROR(context.RunProgram(program)); | ||
| // TODO: Merge back into `FlashAttentionProgram`. | ||
| if (CanApplyFlashAttentionPrefillSimple(is_intel, subgroup_min_size, kv_cache_quantization_enabled)) { |
Contributor
Author
There was a problem hiding this comment.
Verified locally w/ Intel Panther Lake.
Eventually, it will be merged back on flash_attention.wgsl.template, if no regression for Apple/QUALCOMM/NVDIA/...
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Introduce FlashAttentionPrefillSimpleProgram and its template (flash_attention_prefill_simple.wgsl.template). This is a simplified version of the FlashAttentionProgram kernel. (TODO: Merge back into the original FlashAttentionProgram)
Key implementation details:
Integration & Selection:
Test on Phi-4-mini, EdgeLLMOnDeviceModel\2026.5.8.1, Panther Lake:
Motivation and Context
See above.