[Yang Gu (gyagp)]

Description

Optimize FP16 prefill attention on Qualcomm GPUs with explicit 128-lane subgroup support. Four adjacent subgroup lanes share a query row, reducing per-lane query/output storage. A 256-thread workgroup processes 64 rows and stages 16 keys/values at a time. Softmax and output accumulation use FP32, with normalization at the final store.

The cooperative kernel supports causal and noncausal attention, softcap, broadcast attention bias, head sinks, sliding windows, and per-batch query lengths. Softcap precedes bias addition; the window's left bound stays relative to the query even in noncausal mode. Per-batch past lengths use signed, clamped subtraction. Dispatch requires FP16, head size 128, at least 128 query tokens, unquantized KV caches, subgroup-size control supporting 128 lanes, and the subgroup_id WGSL language feature. Decode retains its existing paths.

The specialization is shared by dense-prefill callers rather than tied to Aion. Existing operator-level limits still apply: WebGPU GroupQueryAttention rejects causal=0, and WebGPU PagedAttention rejects nonzero softcap. Noncausal kernel coverage comes from MHA and PagedAttention; softcap coverage comes from GQA.

Motivation and Context

Adreno X1-85 has wide subgroups. Sharing a row across four lanes reduces register storage per lane and fills 128-lane subgroups; FP32 accumulation avoids repeated FP16 rounding while processing key tiles.

On Surface Laptop 7 / Snapdragon X Elite X1E80100 / Adreno X1-85 (driver 31.0.170.0), a fresh source-built ORT GenAI native benchmark comparison on 2026-11-02 measured Aion-1.1-Instruct base-model prefill:

Measurement	Main baseline	Revised PR
Mean prefill latency	4.04888 s	3.46377 s
Prompt throughput	252.91 TPS	295.63 TPS

This is 14.45% less latency / 16.89% more TPS. All 28 attention layers are now eligible, including the 21 sliding-window layers.

Both ORT variants use identical build settings and base c666fd413e (ORT 1.32.0), which includes merged #33560. GenAI is source-built 6b7ff3ca3efa15da219050fdb4ecf8bed67875af (0.18.0-dev). Tests use native benchmark/c/model_benchmark.exe, batch 1, exactly 1,024 prompt tokens, KV capacity 8,192, one generated token, generator reuse, D3D12, graph capture disabled, and validation disabled for timing. Four processes per variant each ran two warmups and five measured iterations, in baseline/PR/PR/baseline/PR/baseline/baseline/PR order: 20 measured samples per variant. Model loading and first-token sampling are outside the prompt-processing scope. Loaded ORT DLL paths and hashes were verified.

Model: https://huggingface.co/webai-community/ai-models/tree/c9697d1b476a9dbcf86771dfbd2abdaa8122c78a/aion/onnx-webgpu . No adapter was loaded. The prompt repeats hello 1,024 times. The local GenAI integration patch adjusts external-QNN registration and teardown output flushing; benchmark timing logic is unchanged.

Validation

• Windows ARM64 Release runtime and provider tests rebuilt from source; lintrunner and whitespace checks passed.

• 101 distinct attention tests passed on this device without skips, including cooperative-prefill dispatch and FP16 graph-capture replay.

• Tests ran with shader-cache validation enabled. Generated shaders confirm cooperative dispatch for the new feature combinations and the Aion sliding-window path.

• New provider-specific tests use 129-token FP16 queries, batch 1, two query heads, one KV head, head size 128, and KV capacity 208. Global attention and a 17-token window are checked against an independent FP64 reference and complete KV-cache contents. The capture test performs two direct graph replays with updated device-resident sequence lengths and total lengths 146/170/132, while retaining fixed buffer addresses.

• Both new tests inspect executed dispatch metadata for subgroup IDs, a 256-thread workgroup, and explicit 128-lane subgroups. Deliberately disabling cooperative dispatch makes the selection test fail; retaining stale replay lengths makes the reference comparison fail. Faults were restored before the successful 101-test run.

• The shared WebGPU GQA reference runner disables CPU fallback, including the cooperative and high-magnitude cases. The new tests also disable fallback.

• New feature tests use independent CPU calculations with FP64 accumulation. Coverage includes softcap, four bias broadcast layouts, sinks, window sizes 1/17/2048, combined features, noncausal cross-attention with shorter/longer KV lengths, and BSNH/BNSH query paths.

• Ragged PagedAttention tests exercise per-batch query lengths 129/101, past lengths 17/7, noncontiguous pages, partial tiles, causal/noncausal mode, and sliding windows through dense fallback.

• Existing PR coverage includes sequence lengths 128/129/257, batch sizes 1/2, GQA ratio 3, rotary, high-magnitude FP16 inputs, past caches, and right-padded key lengths.

• Fresh GenAI validation on the fixed 1,024-token prompt produced finite, bit-identical logits (top token 357) and bit-identical sampled KV entries from layers 0 and 27. This is a fixed-prompt smoke check; FP32 accumulation can change results on other inputs.

• MHA query/key data vary by row, head, and dimension; an independent FP64 reference exercises QK reduction and BNSH indexing.

• WebGPU-only MHA and ragged PagedAttention cooperative-prefill cases live in test/providers/webgpu/contrib_ops, with ShaderF16 capability checks retained. PagedAttention reuses a shared CPU reference runner. Both relocated cases passed in a fresh 101-test run with no skips and shader-cache validation enabled.

• Query ownership uses subgroup_id and subgroup_invocation_id, independently of workgroup-local invocation ordering. Shader dumps confirm the new mapping.

• The three affected test translation units compiled successfully in ARM64 adapter mode with ORT_USE_EP_API_ADAPTERS. A standalone capability probe linked and ran without provider libraries; both helpers return false in adapter mode, and test objects have no undefined WebGPU-context symbols. This is targeted compilation coverage, rather than a full shared-provider build.

• Performance testing is limited to this device and workload. Browser WASM and other Qualcomm adapters were not benchmarked.

[azure-pipelines]

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

[Yang Gu (gyagp)]

/run bots

[Yang Gu (gyagp)]

Jiajia Qin (@qjia7) this is an optimization specific for QC, though the overall speedup is modest. Please take a look.

[Jiajia Qin (qjia7)]

The request was received. Approve workflow runs. Thanks.

[Copilot]

🟡 Changes recommended

The WebGPU shared-library test build can fail, and the shader can mix query rows by assuming workgroup and subgroup lane ordering match; tests also lack FP16 capability checks and meaningful Q/K coverage.

4 open issues

• Derive row ownership from subgroup lane IDs · New
• Guard provider context lookup in adapter builds · New
• Skip FP16 tests when ShaderF16 is unsupported · New
• Use nonzero Q/K values to test BNSH indexing · New

What changed in this PR

Adds a Qualcomm-specific FP16 prefill attention path to ONNX Runtime’s WebGPU provider, using cooperative query processing and FP32 accumulation while retaining existing decode paths.

Changes:

• Dispatches eligible workloads to a 256-thread kernel with explicitly requested 128-lane subgroups.
• Supports attention bias, softcap, sinks, sliding windows, and per-batch sequence lengths.
• Adds regression coverage for numerical accuracy, cross-attention, query layouts, caches, padding, and ragged sequences.

File	Description
onnxruntime/​test/​contrib_ops/​paged_attention_op_test.cc	Adds ragged dense-fallback cases and sliding-window reference calculations.
onnxruntime/​test/​contrib_ops/​multihead_attention_op_test.cc	Adds noncausal cross-attention cases covering two query layouts.
onnxruntime/​test/​contrib_ops/​group_query_attention_op_test.cc	Adds capability-gated cooperative prefill and feature tests.
onnxruntime/​contrib_ops/​webgpu/​bert/​flash_attention.wgsl.template	Implements cooperative attention with FP32 accumulation.
onnxruntime/​contrib_ops/​webgpu/​bert/​flash_attention.h	Adds softcap configuration and uniform declarations.
onnxruntime/​contrib_ops/​webgpu/​bert/​flash_attention.cc	Adds dispatch eligibility, subgroup sizing, and softcap wiring.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

The tests can silently fall back to CPU, do not verify selection of the optimized WebGPU path, and omit FP16 graph-capture replay coverage with live sequence lengths.

3 open issues

• Test FP16 graph capture and replay with a 129-token query · New
• Assert tests select the 128-lane cooperative path · New
• Disable CPU fallback in WebGPU GQA reference tests · New

4 resolved since last review

• Guard provider context lookup in adapter builds
• Derive row ownership from subgroup lane IDs
• Use nonzero Q/K values to test BNSH indexing
• Skip FP16 tests when ShaderF16 is unsupported

_{Review effort: Balanced}

[Yang Gu (gyagp)]

/run bots

[Yang Gu (gyagp)]

Jiajia Qin (@qjia7) please review again

[Jiajia Qin (qjia7)]

The request was received. Approve workflow runs. Thanks.

[Copilot]

🟡 Changes recommended

FP16 query scaling can erase valid attention scores, and key MHA, ragged PagedAttention, and GQA feature tests do not verify cooperative dispatch.

2 open issues

• Verify cooperative dispatch for MHA, ragged PagedAttention, and GQA tests · New
• Scale Q after the FP32 dot product to avoid FP16 underflow · New

3 resolved since last review

• Disable CPU fallback in WebGPU GQA reference tests
• Assert tests select the 128-lane cooperative path
• Test FP16 graph capture and replay with a 129-token query

_{Review effort: Balanced}