webgpu / nbitmm support for bias and weight_index

[Guenther Schmuelling (guschmue)]

add support for bias and weight_index, move subgroup_matrix_matmul_nbits to template and make program callable from other ops.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull Request Overview

Adds WebGPU support for bias and weight_index parameters to N-bit matrix multiplication operations, enabling features like stacked weights and bias addition in quantized operations.

Key changes:

• Extended matmul_nbits operations to support optional bias parameter across multiple implementations (DP4A, wide tile, subgroup matrix)
• Added weight_index uniform variable to enable weight stacking and offset computation in quantized matmul
• Refactored Apple-specific shader generation to use WGSL templates instead of inline string concatenation

Reviewed Changes

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

Show a summary per file

File	Description
matmul_4bits_test.cc	Added test case for WebGPU with bias support
subgroup_matrix_matmul_nbits_apple.wgsl.template	New template file consolidating Apple shader generation with bias support
subgroup_matrix_matmul_nbits.h	Added has_bias and weight_idx parameters to program interface
subgroup_matrix_matmul_nbits.cc	Refactored to use template system and added bias/weight_index support
matmul_nbits_zero_pt.wgsl.template	Added has_bias parameter declaration
matmul_nbits_wide_tile.wgsl.template	Implemented bias addition and weight offset calculations
matmul_nbits.h	Added has_bias parameter and exposed ApplyMatMulNBits function
matmul_nbits.cc	Removed bias constraint, added ApplyMatMulNBits function with extensive documentation
dp4a_matmul_small_m.wgsl.template	Added bias support with offset calculations
dp4a_matmul_nbits.h	Added has_bias and weight_index parameters to program interfaces
dp4a_matmul_nbits.cc	Integrated bias support across DP4A implementations
dp4a_matmul_common.wgsl.template	Added has_bias parameter declaration
dp4a_matmul.wgsl.template	Implemented bias addition with vectorized operations

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Jiajia Qin (qjia7)]

In current change, it seems that ApplyMatMulNBits calculate a with one weight index of b. For QMoE case, it only compute one expert. I remember you said you need 4 selected experts. So will ApplyMatMulNBits be called 4 times to get the up projection results? Why not directly generate the result by calling ApplyMatMulNBits once?

[Guenther Schmuelling (guschmue)]

In current change, it seems that ApplyMatMulNBits calculate a with one weight index of b. For QMoE case, it only compute one expert. I remember you said you need 4 selected experts. So will ApplyMatMulNBits be called 4 times to get the up projection results? Why not directly generate the result by calling ApplyMatMulNBits once?
Earlier on I has a QMoE implementation that does 4 experts in 1 shot but in that case I need to go token by token and the prefill would get pretty costly.
The current implementation (

onnxruntime/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc

Line 198 in 3c65aaa

for (int token_offset = 0; token_offset < moe_params.num_rows; token_offset += max_tokens) {

)
does it different. I looks at all tokens and assigns them to experts, than walks over the experts and runs them 1 by one if there are tokens assigned to them. I think this is more in line with other implementations and should be better for prefill performance.
If you have say 1000 tokens it is most likely we need to run all experts, for generation we'd run 4.
Also helps with memory usage.

[Jiajia Qin (qjia7)]

The current code changes are clean. No more questions for the nbitmm.

If you have say 1000 tokens it is most likely we need to run all experts, for generation we'd run 4.

Does that mean that for generation, we can do specific optimization to calculate the four experts in one ApplyMatMulNBits?

[Jiajia Qin (qjia7)]

The code changes in matmulnbits are clean for me. One unclear place is that whether we should treat prefill and generation different for QMoE which may impact the interface of ApplyMatMulNBits. But we can revisit this part again if needed.