[WebGPU] Unroll bounded NVIDIA FlashAttention reductions in WGSL

[Yang Gu (gyagp)]

Change

Unroll the inner QK/PV reductions for NVIDIA's dense shared-memory FlashAttention prefill path with head sizes up to 128. Preserve FP32 score accumulation, normalized-score conversion, accumulation order, and cache specialization. Larger heads, other vendors, paged attention, and decode retain their existing paths.

Add a high-magnitude FP16 regression with 65 query tokens and head sizes 64/128/256, covering partial tiles, expanded reductions, and the rolled fallback.

Performance

Phi-4-mini-instruct on RTX 5080 / NVIDIA 617.42. Prefill tokens/s, median of three process averages; two warmups and three repetitions per process, batch 1, 32 generated tokens. ORT uses graph capture and generator reuse. Profiling and shader logging are disabled during timing.

Matched PR A/B: Agility SDK 1.721.3-preview / D3D12 SM 6.10, subgroup matrices disabled, DP4A verified in both variants.

Prompt	ORT WebGPU D3D12 before	After	Change	ORT WebGPU Vulkan before	After	Change
512	4,745.36	5,260.95	+10.87%	4,587.78	4,581.26	-0.14%
1,024	4,753.84	5,436.72	+14.36%	4,816.16	4,811.42	-0.10%
2,048	4,101.01	5,254.00	+28.11%	4,788.19	4,776.38	-0.25%
4,096	3,218.81	4,723.09	+46.73%	5,289.50	5,233.24	-1.06%

Cross-runtime baseline context, without this change. SGM means subgroup matrix; llama.cpp uses the b11669 Vulkan release, evaluates full prompts, and disables prompt-cache reuse.

Prompt	ORT WebGPU D3D12 DP4A	ORT WebGPU Vulkan DP4A	ORT WebGPU D3D12 SGM	ORT WebGPU Vulkan SGM	llama.cpp Vulkan
512	4,747.16	4,586.13	8,243.76	7,728.76	10,780.16
1,024	4,760.50	4,801.05	8,765.53	9,785.29	10,828.61
2,048	4,109.71	4,792.29	6,684.01	8,926.87	11,859.29
4,096	3,221.59	5,238.37	4,626.31	9,087.97	11,402.90

ONNX uses Q4/block-32 linear weights plus a Q8 output head; GGUF uses mixed Q4_K_M weights. Cross-runtime comparisons include quantization, kernels, and runtime overhead. Only the matched ORT A/B isolates this patch's effect.

Accepted SM 6.8 tradeoff: the shader-target recheck measured -1.54% / -2.17% full-model prefill at 2,048 / 4,096 tokens and +6.35% attention-kernel time at 2,048. SM 6.9 improved prefill 27.78% / 46.84% at those lengths. The current NVIDIA/head-size gate is retained; this PR does not add a shader-model gate. Broader NVIDIA architecture coverage remains limited.

Validation

• 89 WGSL template tests pass; diff whitespace checks pass.
• All three native WebGPU_FP16_HighMagnitude* tests pass without skips on SM 6.8, 6.9, and 6.10. Captured shaders verify expanded head64/128, rolled head256, and FP32 accumulation.
• Full-model A/B/A checks at 512/1,024/2,048/4,096 tokens have finite, bit-identical sampled logits and matching tokens within each backend, exact A replay, and verified graph capture.
• Rebased onto latest main. The affected attention code and dependencies are unchanged, so the existing benchmark evidence is retained. Cross-backend numerical differences exist in both variants; these checks establish preservation within each backend.

[azure-pipelines]

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

[Yang Gu (gyagp)]

/run bots

[Jiajia Qin (qjia7)]

The request was received. Approve workflow runs. Thanks.