Repository navigation
[WebGPU]: Add Intel subgroup prefill flash-attention kernel - #32702
Draft
Jianhui Dai (daijh) wants to merge 6 commits into
Draft
Jianhui Dai (daijh) wants to merge 6 commits into
Jianhui Dai (daijh) wants to merge 6 commits into
Conversation
|
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/...
Jianhui Dai (daijh)
marked this pull request as draft
September 22, 2026 05:15
Jianhui Dai (daijh)
force-pushed
the
pr/flash-attn-prefill
branch
from
October 15, 2026 10:00
b2a0248 to
5e758ee
Compare
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% |
Jianhui Dai (daijh)
force-pushed
the
pr/flash-attn-prefill
branch
from
October 19, 2026 05:41
5e758ee to
1402c7b
Compare
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
Add FlashAttentionPrefillProgram and flash_attention_prefill.wgsl.template,
a simplified variant of FlashAttentionProgram used for dense prefill on Intel
GPUs. (TODO: merge back into FlashAttentionProgram.)
Kernel design:
QK^T + online-softmax stage. V is staged just before the PV matmul, in
slices of up to 16 vec4 values, so global-memory latency overlaps with
softmax compute and V needs less workgroup memory.
adjacent lanes hit adjacent shared-memory addresses and avoid bank
conflicts.
adapter's guaranteed minimum subgroup size (shuffle_chunk =
min(kv_step, subgroup_min_size)). Below 8, the kernel reads K/V directly
from shared memory, which also covers adapters without subgroups.
Configuration (host-selected per adapter):
256 threads / 32 KV rows.
fall back to the default configuration.
Integration:
ApplyDensePrefill selects this kernel when the adapter is Intel and the
KV cache is not quantized. Otherwise it uses the existing
FlashAttentionProgram.
yet.
Tests:
head_size 80 and 96 cases, whose head_size / 4 is not a multiple of the
V slice width.
Test on Phi-4-mini, EdgeLLMOnDeviceModel\2026.5.8.1, Panther Lake:
Motivation and Context
See above.