webgpu qmoe

[Guenther Schmuelling (guschmue)]

qmoe implementation for webgpu ep

[Copilot]

Pull Request Overview

This PR implements the QMoE (Quantized Mixture of Experts) operator for the WebGPU execution provider, adding support for 4-bit and 8-bit quantized expert weights. The implementation includes gate computation with top-k selection, hidden state gathering, quantized matrix multiplication using existing MatMulNBits infrastructure, SwiGLU activation, and final output mixing.

Key Changes

• Added QMoE kernel registration to webgpu_contrib_kernels.cc (MoE kernel itself remains commented out/unimplemented)
• Implemented QMoE operator with support for 4-bit and 8-bit quantized expert weights
• Added multiple WGSL shader templates for gate computation, hidden state gathering, zero initialization, SwiGLU activation, and final mixing

Reviewed Changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
webgpu_contrib_kernels.cc	Registers QMoE kernel while keeping MoE commented out
qmoe.h	Defines QMoE class extending MoE base class with quantization parameters
qmoe.cc	Implements QMoE compute logic with gate, gather, matmul, and activation steps
moe_base.h	Defines base enums for activation types, quantization types, and parallelization modes
moe.h	Defines base MoE class with common activation and routing parameters
moe.cc	Stub implementation of base MoE class (not yet implemented)
gate.wgsl.template	WGSL shader for gate computation using bubble sort for top-k selection
hidden_state_gather.wgsl.template	WGSL shader for gathering hidden states assigned to each expert
zero_tensor.wgsl.template	WGSL shader for zero-initializing output tensor
swiglu.wgsl.template	WGSL shader implementing SwiGLU activation function
final_mix.wgsl.template	WGSL shader for mixing expert outputs with router weights

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

You can commit the suggested changes from lintrunner.

[Guenther Schmuelling (guschmue)]

• some fixes around chunking tokens into chunks of max_tokens
• merge main
• reflect some changes around zero_points that came in with main

testing looks good on nvidia, apple gpu and intel integrated gpu.

[Jiajia Qin (qjia7)]

Guenther Schmuelling (@guschmue), two additional points may need your attention in follow-ups aside from the existing TODOs in your PR:

Firstly, please refactor all your shaders using input.getByOffset/output.setByOffset. This approach helps avoid errors such as tensor size exceeding the maxStorageBufferBindingSize.

Secondly, consider optimizing the generation phase when moe_params.num_rows is 1. Your current solution is suitable for the prefill phase where moe_params.num_rows is greater than 1. I think your previous solution #26130 might be more effective for generation.