[Rithik Sharma (SharmaRithik)]

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.

[azure-pipelines]

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

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[Rithik Sharma (SharmaRithik)]

Jiajia Qin (@qjia7) could you take a look when you have time? This is the tuning follow-up I mentioned in #32937. It is NVIDIA only and was tested on Linux with Vulkan on an RTX 5080 and an RTX 4070. The second commit only regenerates the WGSL template goldens. The workflows need approval to run, thanks.

[Jiajia Qin (qjia7)]

Rithik Sharma (@SharmaRithik), are you using the f16 model with f16 weights or quantized weights? For quantized weights, I think the model will be generated with MatMulNbits instead of MatMul. And the MatMulNbits path has many optimizations for the language model prefill and decode stage. I am just curious why you have so many matmul ops in your model.

[Rithik Sharma (SharmaRithik)]

Thanks Jiajia Qin (@qjia7), I used both. The fp16 rows use the onnx-community fp16 exports, where every projection is an FP16 MatMul. The q4f16 rows use the q4f16 exports, where the projections are MatMulNBits and only lm_head is an FP16 MatMul, which is why those rows gain less. The models are:

1. Llama-3.2-1B-Instruct fp16: https://huggingface.co/onnx-community/Llama-3.2-1B-Instruct-ONNX/blob/main/onnx/model_fp16.onnx
2. Llama-3.2-1B-Instruct q4f16: https://huggingface.co/onnx-community/Llama-3.2-1B-Instruct-ONNX/blob/main/onnx/model_q4f16.onnx
3. Qwen3-0.6B fp16: https://huggingface.co/onnx-community/Qwen3-0.6B-ONNX/blob/main/onnx/model_fp16.onnx
4. Qwen3-0.6B q4f16: https://huggingface.co/onnx-community/Qwen3-0.6B-ONNX/blob/main/onnx/model_q4f16.onnx
5. gemma-3-270m-it fp16: https://huggingface.co/onnx-community/gemma-3-270m-it-ONNX/blob/main/onnx/model_fp16.onnx