webgpu qmoe support - #26130
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR implements QMoE (Quantized Mixture of Experts) support for WebGPU, along with enhanced MatMulNbits functionality that supports bias and stacked weights/outputs.
- Adds QMoE operator implementation with quantized expert weights (4-bit and 8-bit support)
- Extends MatMulNbits to support bias addition and stacked weight operations via offsets
- Includes new shader templates for gate computation, SwiGLU activation, and final mixing
Reviewed Changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| webgpu_execution_provider.cc | Adds commented TopK operator declarations for future use |
| webgpu_contrib_kernels.cc | Registers QMoE kernel and comments out MoE kernel |
| matmul_nbits.h | Extends function signatures to support bias and offsets parameters |
| matmul_nbits.cc | Implements ApplyMatMulNBits function with bias/offsets support and optimized kernel selection |
| *.wgsl.template | New shader templates for QMoE gate computation, SwiGLU activation, and final mixing |
| qmoe.h/.cc | Complete QMoE operator implementation with multi-stage processing |
| moe_base.h/moe.h/.cc | Base classes and infrastructure for MoE operations |
| attention.cc | Fixes type casting issue in head sink computation |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| const Tensor* g_idx = context.Input(4); | ||
| const Tensor* bias = context.Input(5); | ||
|
|
||
| // FIXME: this is for debgguing only, will remove it later. |
There was a problem hiding this comment.
There's a typo in the comment: 'debgguing' should be 'debugging'.
| Tensor *y; | ||
| auto output_shape = helper.OutputShape(); | ||
| if (offsets != nullptr) { | ||
| // FIXME: this is for debgguing only, will remove it later. |
There was a problem hiding this comment.
There's a typo in the comment: 'debgguing' should be 'debugging'.
| // FIXME: this is for debgguing only, will remove it later. | |
| // FIXME: this is for debugging only, will remove it later. |
|
|
||
| #use guardAgainstOutOfBoundsWorkgroupSizes | ||
|
|
||
| // This implemnts: |
There was a problem hiding this comment.
There's a typo in the comment: 'implemnts' should be 'implements'.
| // This implemnts: | |
| // This implements: |
|
|
||
| Status ComputeInternal(ComputeContext& context) const override; | ||
|
|
||
| // FIXME: for debug only, will remove it later. |
There was a problem hiding this comment.
Debug code should be removed before merging to production. The DumpTensor method and related debug functionality should either be removed or properly conditionally compiled.
| class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, Softmax); | ||
| class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, Softmax); | ||
| class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, Softmax); | ||
| // class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 9, TopK); |
There was a problem hiding this comment.
What is the reason for adding the registrations for TopK but commenting them out?
There was a problem hiding this comment.
I have some code for topk but wanted to keep this PR limited to QMoe.
| topk_indices_fc1[output_base + tid].x = 0u; | ||
| topk_indices_fc1[output_base + tid].y = idx; | ||
| topk_indices_fc1[output_base + tid].z = row * K + tid; | ||
| topk_indices_fc1[output_base + tid].w = 99u; |
There was a problem hiding this comment.
It will be better to add an explanation for x/y/z/w's meanings.
| // | ||
| // Step 4: multiply fc1_activated with fc2 (gate_down) of the selected experts | ||
| // | ||
| TensorShape fc2_output_shape({k_, moe_params.num_rows, moe_params.inter_size}); |
There was a problem hiding this comment.
Why fc2_output_shape is not {k_, moe_params.num_rows, moe_params.hidden_size}? Should N_fc2 be equal to moe_params.hidden_size?
|
|
||
| class MoEProgram final : public Program<MoEProgram> { | ||
| public: | ||
| MoEProgram(TensorShape output_shape) : Program<MoEProgram>{"MoE"}, output_shape_{output_shape} {} |
| } | ||
|
|
||
| Status MoE::ComputeInternal(ComputeContext& context) const { | ||
| return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "MoE is not implemented in WebGPU"); |
|
this is using way too much memory. closing this one, new PR coming in a day. |
Still some issues with gpt-oss but UT is fine.
The support for bias and stacked weights/outputs is currently implemented only in the default shader, going to add the others this week.