[Jianhui Dai (daijh)]

Description

Broadcast each lane's A vecs across the subgroup via subgroupShuffle instead of having every lane re-read the full A tile from workgroup memory. Each lane loads its own row once and shuffles it to the other rows that need it. The workgroup tile is stored as [KAVecSizeForBlock32/2][kTileM][2] pairs so the two vecs consumed per dequantized weight column are read in one contiguous access.

tile_m now adapts to the element size (16 for f32, 32 for f16, since f16 packs twice as many elements per register), and the compute-stage reduction strategy is chosen from the adapter's guaranteed minimum subgroup size:

• subgroup_min_size >= tile_m: single full-tile subgroupShuffle band
• tile_m % subgroup_min_size == 0 (subgroup_min_size > 0): chunked subgroupShuffle bands
• otherwise (including when Subgroups isn't supported): direct reduction, no shuffle

The actual runtime subgroup size can exceed the adapter's guaranteed minimum, so a raw sg_id can index a_data_tile out of bounds in both shuffle branches. Cap it via capped_sg_id = min(sg_id, kTileM - 1u) before the main loop and use it in place of sg_id in both branches.

Reduces redundant workgroup-memory reads for the A operand in the MatMulNBits wide-tile prefill kernel, and lets it run correctly with the best available strategy across adapters with varying or no subgroup support.

max_length = 8K	Prefill Length	Default Prefill TPS	Opt TileM-16 Prefill TPS	Opt TileM-32 Prefill TPS	Improvement
Panther Lake	128	221.01	242.49	243.02	110%
Panther Lake	1024	594.20	644.74	696.70	117%
Panther Lake	4096	596.47	667.17	727.84	122%
Alder Lake	128	46.76	86.13	113.09	242%
Alder Lake	1024	44.72	94.43	129.06	289%

[1] https://huggingface.co/onnx-community/Phi-4-mini-instruct-ONNX

Motivation and Context

See above.

[azure-pipelines]

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

[Jianhui Dai (daijh)]

Add some performance data for next step.

[Jianhui Dai (daijh)]

Update the description w/ perf number.
Jiajia Qin (@qjia7) Hariharan Seshadri (@hariharans29) PTAL

[Copilot]

Pull request overview

This PR optimizes the WebGPU MatMulNBits wide-tile (Block32 prefill) shader by reducing redundant workgroup-memory reads of the A tile using subgroup shuffles when supported, while providing a safe fallback path when subgroups are unavailable or undesirable (e.g., NVIDIA).

Changes:

• Added subgroup_min_size specialization to select between full-tile shuffle, chunked shuffle, or direct reduction in the wide-tile WGSL template.
• Changed wide-tile A-tile storage layout to pack K vectors as contiguous pairs for more efficient reads during dequantized dot-products.
• Made tile_m adaptive to element size (f32: 16, f16: 32) and updated program cache hinting accordingly.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h	Extends wide-tile program specialization parameters to include subgroup_min_size.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc	Selects tile_m based on A datatype and passes adapter subgroup minimum size into shader specialization + cache key.
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_wide_tile.wgsl.template	Implements subgroup-shuffle broadcast strategy (with safe sg_id capping) and reorganizes A-tile layout for paired reads.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Jianhui Dai (daijh)]

urllib.error.URLError: <urlopen error [Errno 8] nodename nor servname provided, or not known>
socket.gaierror: [Errno 8] nodename nor servname provided, or not known

The MacOS CI Pipeline fails because of network connectivity issue.
Could you please restart the CI?

[Jianhui Dai (daijh)]

urllib.error.URLError: <urlopen error [Errno 8] nodename nor servname provided, or not known>
socket.gaierror: [Errno 8] nodename nor servname provided, or not known

The MacOS CI Pipeline fails because of network connectivity issue. Could you please restart the CI?

Fail again as network connectivity issue. Rebase to try.

[Jianhui Dai (daijh)]

Hariharan Seshadri (@hariharans29) PTAL. The GPU Doc Gen CI failure is unrelated to these changes.