[Jie Chen (jchen10)]

A trailing M-tile is not required to be a multiple of kTileM. Writing it straight to output with subgroupMatrixStore can run past M.
Bounds-check the tail tile instead of writing it unconditionally. The has_bias path already staged results through a workgroup scratch array and added bias with a per-element bounds check, so it is left as-is. For the no-bias path, avoid paying for that workgroup scratch on every dispatch: the tail tile is written via the same branch-free fast path into a small padded [tile_size_a, N] auxiliary buffer instead, and a follow-up
SubgroupMatrixMatMulNBitsTailCopyProgram dispatch crops the valid rows into the real output.

[azure-pipelines]

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

[Jie Chen (jchen10)]

Jianhui Dai (@daijh) PTAL

[Copilot]

🟡 Changes recommended

One or more issues must be addressed before approval.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Fixes OOB stores for non-aligned trailing M tiles in the WebGPU 8x16x16 MatMulNBits kernel.

Changes:

• Adds a padded tail buffer for no-bias partial tiles.
• Adds a follow-up shader to copy valid rows into the output.
• Preserves the existing bias path and updates shader/cache wiring.

File summaries

File	Description
onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.h	Updated as part of this pull request.
onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc	Updated as part of this pull request.
onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_tail_copy.wgsl.template	Updated as part of this pull request.
onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template	Updated as part of this pull request.

Review details

Suppressed comments (1)

onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits_8x16x16.wgsl.template:258

• This raw subgroupMatrixStore path is not safe when tail_buffer is larger than the device's maxStorageBufferBindingSize. ComputeContext transparently splits such outputs into tail_output, tail_output1, etc., but these stores always address only tail_output; the follow-up copy can therefore read missing/old values for columns in later segments. This is possible even when y fits because the padded buffer has 64 rows while the supported range allows a partial M as small as 4. Keep this buffer within one binding or use a segmented-safe/fallback write path for oversized tails.

        } else {
            let sg_mat_offset_c = global_base_b + sg_idx * kSgMatM * uniforms.N;
            subgroupMatrixStore<row_major>(&tail_output, sg_mat_offset_c, sg_mat_c0, uniforms.N);
            subgroupMatrixStore<row_major>(&tail_output, sg_mat_offset_c + kSgMatN, sg_mat_c1, uniforms.N);
            subgroupMatrixStore<row_major>(&tail_output, sg_mat_offset_c + 2 * kSgMatN, sg_mat_c2, uniforms.N);
            subgroupMatrixStore<row_major>(&tail_output, sg_mat_offset_c + 3 * kSgMatN, sg_mat_c3, uniforms.N);

• Files reviewed: 4/4 changed files
• Comments generated: 1
• Review effort level: Lite (auto)

Note

Copilot is running an experiment and ran this review at Lite.

---

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

[Jianhui Dai (daijh)]

Can we drop the N % 64 == 0 restriction in CanApplySubgroupMatrixMatMulNBits(), with this fixing?

[Jie Chen (jchen10)]

Can we drop the N % 64 == 0 restriction in CanApplySubgroupMatrixMatMulNBits(), with this fixing?

Not yet. If there is an actual need to lift the N % 64 == 0 restriction, most likely the same idea of this PR could be applied.

[Jianhui Dai (daijh)]

It could be great to remove the restriction, as it allows us to apply the subgroup matrix feature more broadly.

[Jie Chen (jchen10)]

It could be great to remove the restriction, as it allows us to apply the subgroup matrix feature more broadly.

Let's address this with a separate PR. Currently N % 64 == 0 is a common restriction, not just specific to Intel.

[Jianhui Dai (daijh)]

It could be great to remove the restriction, as it allows us to apply the subgroup matrix feature more broadly.

Let's address this with a separate PR. Currently N % 64 == 0 is a common restriction, not just specific to Intel.

Thanks. LGTM.

[Jie Chen (jchen10)]

Jiajia Qin (@qjia7) Hariharan Seshadri (@hariharans29) PTAL