onnxruntime
0ba2a859 - [WebGPU] Use shared memory tiles in subgroup matrix MatMul on NVIDIA (#33639)

Commit
7 days ago
[WebGPU] Use shared memory tiles in subgroup matrix MatMul on NVIDIA (#33639) ### Description The subgroup matrix MatMul kernel loads A and B straight from global memory with one subgroup per K split, so this PR adds a path where a workgroup of subgroups copies a K slice of A and B into shared memory and loads its subgroup matrices from there. On NVIDIA a 4x4 grid of subgroups stages K slices 64 deep and computes a 128x128 output tile for large problems and a 64x64 tile otherwise, while small problems keep the existing split K tiling because it fills the GPU better there, and each tiling is checked against the device limits with the existing tiling as the fallback. The tile sizes and the K depth were measured in ORT on an RTX 5080 and an RTX 4070, and since the path is selected for NVIDIA only, other vendors are unchanged. ### Motivation and Context Prefill speedup over main for the onnx-community models on Linux with Vulkan, P50 of 100 runs with `onnxruntime_perf_test`, averaged over five rounds: | Model, prompt tokens | RTX 5080 | RTX 4070 | |---|---:|---:| | Llama-3.2-1B-Instruct fp16, 512 | 10.6% | 25.3% | | Llama-3.2-1B-Instruct fp16, 1024 | 5.6% | 21.0% | | Qwen3-0.6B fp16, 512 | 3.1% | 10.5% | | Qwen3-0.6B fp16, 1024 | -0.6% | 13.7% | | gemma-3-270m-it fp16, 512 | 6.3% | 7.2% | | gemma-3-270m-it fp16, 1024 | 2.8% | 9.4% | | Llama-3.2-1B-Instruct q4f16, 512 | 5.7% | 11.5% | | Llama-3.2-1B-Instruct q4f16, 1024 | 5.4% | 4.9% | | Qwen3-0.6B q4f16, 512 | 3.1% | 4.5% | | Qwen3-0.6B q4f16, 1024 | 0.3% | 6.4% | | Geometric mean | 4.2% | 11.3% | The q4f16 models gain because their `lm_head` MatMul is FP16.
Author
Parents
Loading