[Jianhui Dai (daijh)]

Description

Add FlashAttentionPrefillProgram, a simplified variant of
FlashAttentionProgram used for dense prefill on Intel GPUs.

Kernel design (and key differences from FlashAttentionProgram):

• One query row per thread.
• The full-head K tile is staged before the 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. The default kernel stages both K
  and V at full head size.
• K and V tiles use a transposed layout ([head_size_vec][kv_step]),
  so adjacent lanes hit adjacent shared-memory addresses and avoid bank
  conflicts. The default kernel uses [kv_step][head_size_vec].
• Subgroup shuffle is a static code path chosen on the host from the
  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.
• The workgroup size and KV step are chosen per adapter and shapes.
• No quantized KV cache, TurboQuant LUT, or compressed K/V decoding
  paths.

Configuration (host-selected per adapter):

• Default: 64 threads, 16 KV rows per step.
• Xe2: 128 threads / 64 KV rows for query length <= 128, otherwise
  256 threads / 32 KV rows.
• Xe3: 128 threads / 64 KV rows.
• If the K tile and V slice would exceed
  maxComputeWorkgroupStorageSize, fall back to the default
  configuration.

Integration:

• Adds CanApplyFlashAttentionPrefill and ApplyFlashAttentionPrefill.
  ApplyDensePrefill selects this kernel when the adapter is Intel and
  the KV cache is not quantized. Otherwise it uses the existing
  FlashAttentionProgram.
• Quantized KV cache (int8 / TurboQuant) is not supported by this
  kernel yet.

Tests:

• Parameterize RunFp16HighMagnitudeAttention by head size and add
  head_size 80 and 96 cases, whose head_size / 4 is not a multiple of
  the V slice width.

Performance (Phi-4-mini, Intel Panther Lake):

Metric	Baseline Prefill TPS	Optimized Prefill TPS	Improve
Prefill-128	1961.46	1945.12	-0.83%
Prefill-256	2481.54	2502.52	0.85%
Prefill-512	2577.39	2774.83	7.66%
Prefill-1024	2505.69	2751.47	9.81%
Prefill-2048	2346.94	2673.69	13.92%
Prefill-4096	1668.55	1985.91	19.02%

TODO (follow-up PRs):

• Support quantized KV cache (int8 / TurboQuant).
• Merge back into FlashAttentionProgram.

Motivation and Context

See above.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Jianhui Dai (daijh)]

Prerequisite: #32605

[Copilot]

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 · 1 · 1

Open (3)

• Accumulate online softmax state in f32 for f16 queries · New
• Avoid regressing short Intel prefill sequences · New
• Add automated coverage for the Intel-specific kernel path · New

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.

[Jianhui Dai (daijh)]

Jie Chen (@jchen10) Jiajia Qin (@qjia7) PTAL.