webgpu: Refactor SubgroupMatrixMatMulNBits to vendor-agnostic config … (#28109)
…+ add NVIDIA 16x16x16
Refactor subgroup matrix MatMulNBits support from vendor-specific
(Apple/Intel) to a vendor-agnostic config-based approach. Any GPU
reporting a matching subgroup matrix config from Dawn is now
automatically supported.
Key changes:
- Replace vendor-specific config table with
SupportedSubgroupMatrixConfig struct containing {componentType,
resultComponentType, M, N, K, subgroupMinSize, subgroupMaxSize,
needsPrepack}. No architecture or backendType required.
- Remove vendor_ member from SubgroupMatrixMatMulNBitsProgram. Shader
selection is now driven by config dimensions (8x8x8, 8x16x16, 16x16x16).
- Remove vendor gate in matmul_nbits.cc call site.
- Rename shader templates: _apple -> _8x8x8, _intel -> _8x16x16.
- Add new 16x16x16 shader template for NVIDIA Blackwell (RTX 5080).
- 4 subgroups x 32 lanes = 128 threads per workgroup
- 64x64 tile with 16x16 subgroup matrices
- Bounds-checked output via scratch buffer for partial M tiles
- Fix prepack shader OOB reads: add scalar fallback with zero-fill for
partial blocks where M is not a multiple of kSgMatM.
- Prioritize larger configs (16x16x16 > 8x16x16 > 8x8x8) when multiple
match.
Verified on NVIDIA RTX 5080 (Blackwell, Vulkan backend):
- Correctness: model-qa.py with phi4-graph-prune produces identical
output to D3D12 baseline
- Prefill (phi4, l=1024):
phi4-graph-prune | D3D12 DP4A | Vulkan DP4A | Vulkan TC (16x16x16) |
Vulkan TC (16x16x16_128)
-- | -- | -- | -- | --
Prefill (tps) | 3,134 | 6,389 | 7,089 | 10,744
- NVIDIA reports ChromiumExperimentalSubgroupMatrix with F16/F16
16x16x16 config
### Description
<!-- Describe your changes. -->
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->