[WebGPU] Pin subgroup size to 32 for subgroup-matrix MatMul/Gemm

[Jie Chen (jchen10)]

The 8x16x16 subgroup-matrix templates hard-code 32 lanes per subgroup and derive sg_index/sg_lane from it. Request the size explicitly with SetSubgroupSize(32).

[azure-pipelines]

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

[Copilot]

Pull request overview

Pins WebGPU subgroup-matrix MatMul and Gemm kernels to their required 32-lane subgroup size.

Changes:

• Requests subgroup size 32 during program setup.
• Falls back when subgroup-size control is unavailable.
• Removes resolved TODOs.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
subgroup_matrix_matmul.cc	Pins and gates MatMul subgroup size.
subgroup_matrix_gemm.cc	Pins and gates Gemm subgroup size.

No actionable issues found.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

[Hariharan Seshadri (hariharans29)]

Review — PR #32306: [WebGPU] Pin subgroup size to 32 for subgroup-matrix MatMul/Gemm

Bug — what's actually broken pre-PR

The 8x16x16 F16 subgroup-matrix templates in subgroup_matrix_matmul.cc and subgroup_matrix_gemm.cc hard-code 32 lanes per subgroup (kSubgroupMatrixSubgroupSize = 32) and derive sg_index / sg_lane arithmetic from that constant. The workgroup size is set as kSubgroupMatrixSubgroupSize * split_k, again with 32 as the multiplier.

The adapters that report the 8x16x16 F16 config (Intel Xe HPG / Arc-family) expose a 16-32 subgroup size range — i.e., WebGPU is free to pick 16 or 32 depending on shader occupancy heuristics. When the runtime picks 16, the kernel:

• Runs with a workgroup of 32 * split_k threads but only 16 * split_k actual lanes-per-subgroup, so split_k "subgroups" is under-counted at the workgroup level.
• Derives sg_index = local_invocation_id / 32 and sg_lane = local_invocation_id % 32, which no longer align to actual subgroup boundaries — the lanes intended to be one subgroup straddle two.
• Silent numerical corruption in the F16 GEMM/MatMul output.

Fix: explicitly request 32 lanes via SetSubgroupSize(32) and gate on the SubgroupSizeControl feature so the request is honored (or the kernel decline itself if the feature is unavailable).

The change — mechanics

Symmetric edits to both kernels:

1. program.SetSubgroupSize(kSubgroupMatrixSubgroupSize) added right after SetWorkgroupSize(kSubgroupMatrixSubgroupSize * split_k).

2. Selector gate extended:

   if (!IsSubgroupMatrixConfigSupported(context, /*is_fp16=*/true, config_index) ||
       !supported_subgroup_matrix_configs[config_index].Is(8, 16, 16) ||
       !context.HasFeature(wgpu::FeatureName::SubgroupSizeControl)) {
     return nullptr;
   }

3. TODO comment (// TODO: use subgroup-size-control to enforce the subgroup size is 32.) removed above the kSubgroupMatrixSubgroupSize constant. Now that the mechanism is in place, keeping the TODO would be stale.

4. Doc comment on CreateSubgroupMatrixMatMulImpl / CreateSubgroupMatrixGemmImpl extended to explain why the pin is needed:

   That config's adapters expose a 16-32 subgroup size range, so the kernel's fixed 32 lanes per subgroup must be pinned with subgroup-size control.

Concise and correct explanation, gives the next reader the specific hardware rationale without hand-waving.

Correctness notes

• Cache key: CacheHint(has_c, trans_a, trans_b, config_index_, sg_mat_count_m, sg_mat_count_n, split_k) does not include the subgroup size. That's fine — subgroup size is always 32 for this kernel by construction, and there's no configurable path. If a future PR ever lets subgroup size be selectable per-invocation, the cache key would need updating. Not this PR's problem.
• Workgroup-size interaction: SetWorkgroupSize(32 * split_k) combined with SetSubgroupSize(32) is internally consistent — the workgroup runs exactly split_k subgroups of 32 lanes each, matching the WGSL kernel's assumption. split_k is small (typically ≤ 8), so 32 × 8 = 256 stays comfortably under WebGPU's 256-thread workgroup limit. ✓
• Fallback path: CreateSubgroupMatrixMatMulImpl returning nullptr causes the caller to try the next MatMul::MatMulOptImpl in whatever priority chain the MatMul op maintains. This is exactly the pre-existing fallback flow — the PR only adds one more condition to the early bail, not a new failure mode.
• Adapter feature coverage: SubgroupSizeControl is one of the standard WebGPU features. Adapters that report the ChromiumExperimentalSubgroupMatrix feature (prerequisite for the 8x16x16 F16 config check to even fire) are effectively the same Intel Arc-family / Meteor Lake iGPU set that also expose SubgroupSizeControl, so the added guard is unlikely to disable the fast path on any real hardware that previously used it. If some adapter reports subgroup-matrix but not subgroup-size-control, the fallback is correct — that's exactly the case that would have silently miscomputed pre-PR. ✓

Prerequisite check — SetSubgroupSize on ProgramBase

Grepping onnxruntime/core/providers/webgpu/program.h in my local snapshot, I see only the three SetWorkgroupSize(x[, y[, z]]) overloads; no SetSubgroupSize. The PR touches only the two kernel .cc files (+12/-6), so ProgramBase::SetSubgroupSize and the wgpu pipeline plumbing that forwards it into the ComputePipelineDescriptor's nextInChain: WGPUComputePipelineDescriptorFullSubgroups (or the equivalent Dawn extension struct) must already exist on main. My workspace snapshot may just be behind main.

Confirm this by checking: git log --all -- onnxruntime/core/providers/webgpu/program.h for a recent commit that added a SetSubgroupSize overload. If none, this PR needs to also add the API (which would be a substantially bigger change) — but the +12/-6 diff strongly implies the API is already there and this PR is a targeted consumer of it.

Non-blocking, but worth verifying before merge — a build failure on SetSubgroupSize would be caught by CI, so if CI is green, the API is in place.

Test coverage

None added. Reasonable — this is a hardware-behavior-dependent fix (Intel adapters that pick 16 vs. 32 subgroups at runtime), and there's no portable unit test that would catch it. Existing MatMul/Gemm numerical tests already cover the F16 output; if they didn't fail on Intel Arc pre-PR, that's either (a) the runtime happened to pick 32 in the tested workloads, or (b) the numerical tolerance was loose enough to mask the corruption. Neither is testable in a portable way. If the WebGPU CI has an Intel-Arc-capable leg, running the standard MatMul suite pre- and post-PR would show the delta; if not, this fix relies on hand-validation by the author.

CI status

40/81 checks OK at last look — that's a mid-run snapshot, not a settled state. Hariharan requested a Copilot re-review 30 minutes before I looked; several CI legs are likely still queuing/running. Copilot itself already reviewed and generated no comments on both passes. Wait for CI to settle before merge — the interesting checks are the WebGPU CI legs (webgpu_build_x64_RelWithDebInfo (vcpkg, static) and (novcpkg, static)) that would catch the SetSubgroupSize API existence question above. If those are green, prerequisite is confirmed.

Style / nits

• The updated comment reads well and explicitly names the "16-32 subgroup size range" behavior. Excellent — it gives the next reader the vendor-specific fact without requiring them to consult a WebGPU spec section.
• Could optionally reference the WebGPU subgroup-matrix proposal or the specific Intel adapter class (e.g., "Xe HPG / Arc") for even more grounding. Non-blocking.
• Small consistency thing: both files now have identical selector-gate logic + identical rationale comment. Since this is a two-kernel special case, no need to factor it out. Fine as-is.

Recommendation

Approve — pending:

1. Confirming that ProgramBase::SetSubgroupSize exists on the target main branch (implicit if CI compiles). If the workspace snapshot is stale and the API is on main, no action needed.
2. Waiting for CI to settle from 40/81 — specifically the two WebGPU build legs.

The fix is minimal, correctly diagnosed, symmetrically applied to both kernels, and the fallback semantics preserve pre-existing behavior for adapters that can't honor the pin. This is precisely the right shape for a "kernel-invariant pinning" bug fix.

[Jie Chen (jchen10)]

Jiajia Qin (@qjia7) PTAL