Repository navigation
[WebGPU] Optimize DP4A for large Qualcomm subgroups - #33560
Open
Yang Gu (gyagp) wants to merge 2 commits into
Open
Yang Gu (gyagp) wants to merge 2 commits into
Yang Gu (gyagp) wants to merge 2 commits into
Conversation
Reduce each lane's four activation components before XOR shuffles on wide subgroups, keeping each 128-value quantization block independent. Use a 128x32 matrix tile with two explicitly requested 128-lane subgroups on compatible Qualcomm GPUs for 4-bit weights, block size 32, and M >= 128. A fresh source-built ORT GenAI native benchmark on Adreno X1-85 measures Aion batch-1, 1024-token prefill at 3.76782 s before and 3.63298 s after: 3.58% lower latency and 3.71% higher throughput (271.78 -> 281.86 TPS). This uses KV capacity 8192, identical build settings, D3D12, generator reuse, graph capture disabled, and one generated token. Reordered runs collect 20 measured prefills per variant with two warmups per process; loaded ORT DLL paths and hashes are verified. At capacity 1025, ten samples per variant measure 4.08% lower latency and 4.25% more TPS. Earlier GPU profiling measured quantization at 244.4 -> 82.9 ms (2.95x faster) and matrix execution at 5.586 -> 5.342 s (4.4% less time). Those profiles are from a separate session, not a breakdown of the GenAI run. Performance results are limited to this device and workload. Add regression coverage for independent quantization blocks, partial workgroups, M/N tile tails, batching, bias, and zero points. Validation: Windows ARM64 Release builds, four focused provider tests, C++ lint, and whitespace checks passed. Fresh GenAI prefill logits and sampled KV values match the baseline bit-for-bit. Earlier Aion prefill and one-token decode checks also matched at 128 and 1024 tokens.
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
Author
|
Jiajia Qin (@qjia7) The perf gain is moderate, but repeatable. Please take a look. |
Contributor
Author
|
/run bots |
Contributor
|
The request was received. Approve workflow runs. Thanks. |
Rename the row tile dimension to tile_m in host dispatch and WGSL, matching tile_n and clarifying rectangular tiles. Load each activation scale only at the start of its 128-value block. The following three 32-value K tiles reuse the workgroup scale; batch and row mapping stay fixed and both workgroup barriers are unchanged. This reduces scale fetches from 512 to 128 per full 128-row tile and 128 K values. No additional end-to-end speedup is claimed. Add a deterministic regression selecting every K tile across three activation blocks with different scales per block, row, and batch. Cover regular and Qualcomm wave128 tile shapes, M/N tails, and 4/8-bit weights. Use tile_n for the B zero-point array and correct the comment on activation quantization block size. Validation: Windows ARM64 Release build, six focused provider tests, C++ lintrunner, and git diff --check passed on Adreno X1-85.
Contributor
Author
|
/run bots |
Contributor
Author
|
Jiajia Qin (@qjia7) I have addressed your comments, please take another look. |
Contributor
|
The request was received. Approve workflow runs. Thanks. |
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
Optimize WebGPU DP4A activation quantization and matrix tiling for wide Qualcomm subgroups.
Reduce each lane's four absolute activation components locally before XOR shuffles, preserving independent 32-lane partitions for each 128-value quantization block.
Use a 128x32 output tile with two explicitly requested 128-lane subgroups on compatible Qualcomm adapters. This route requires 4-bit weights, block size 32, and M >= 128.
Reuse each activation scale across its four 32-value K tiles, retaining both workgroup barriers.
Add regression coverage for independent quantization blocks, scale reuse across K tiles, partial workgroups, matrix tile tails, batching, bias, and zero points.
Motivation and Context
Adreno X1-85 exposes 64-128-lane subgroups. Local component reduction reduces quantization shuffle traffic; the matrix tile assigns one activation row and 16 output columns to each lane.
A fresh source-built ONNX Runtime GenAI native benchmark comparison on Windows ARM64 / Adreno X1-85 (driver 31.0.170.0), Aion batch 1, and exactly 1,024 prompt tokens measured:
| KV-cache capacity | Baseline prefill | PR prefill | Baseline TPS | PR TPS | Latency reduction | TPS gain |
|---|---:|---:|---:|---:|---:|---:|
| 8,192 | 3.76782 s | 3.63298 s | 271.78 | 281.86 | 3.58% | 3.71% |
| 1,025 | 3.77155 s | 3.61767 s | 271.51 | 283.06 | 4.08% | 4.25% |
Both ORT variants were rebuilt with identical settings and used the same source-built GenAI executable, model and prompt. Runs used D3D12, generator reuse, graph capture disabled, one generated token, two warmups and five measured iterations per process. Reordered runs collected 20 measured prefills per version at capacity 8,192 and 10 at capacity 1,025. Loaded ORT DLL paths and hashes were verified. GenAI times prompt processing separately from first-token sampling; timing scopes and loops were unchanged. The local GenAI integration patch only adjusts external-QNN registration and output flushing.
Earlier GPU profiling measured activation quantization at 244.4 ms before and 82.9 ms after (approximately 3x faster), and exploratory matrix-stage timing at 5.586 s versus 5.342 s. These stage profiles came from an earlier session and are not a decomposition of the fresh GenAI results. Performance evidence is limited to this device and workload.
The performance measurements above precede the activation-scale reuse review follow-up. No additional end-to-end gain is claimed for that follow-up.
Validation
Fresh native Windows ARM64 Release build with WebGPU/D3D12 succeeded.
Six focused provider tests passed after the review follow-up:
WebGpuMatMulNBitsTest.DP4AActivationScaleReuse,MatMul2Bits.Float32_2b_Accuracy4,WebGpuMatMulNBitsTest.DP4AQualcommWave128Tiles,WebGpuMatMulNBitsTest.DP4AQuantizationLargeSubgroups,MatMulNBits.Float32_4b_Accuracy4, andMatMulNBits.Float32_4b_Accuracy4_Batch.Fresh GenAI prefill logits and sampled KV-cache values were bit-identical between the rebuilt variants at capacity 8,192.
Aion prefill and one-token decode logits were bit-identical to the original build at 128 and 1,024 input tokens.
C++ lintrunner and whitespace checks passed.
Three broader FP16 MatMulNBits cases at M=100, N=64, K=32 fail on both the original build and this branch through a different kernel. Browser WASM and other GPU vendors have not been benchmarked.