From 67a974e0bbb9d2d297a90cb435f9c22e26a8f4b4 Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Mon, 3 Nov 2025 18:29:35 -0800 Subject: [PATCH 01/15] webgpu qmoe --- .../webgpu/moe/final_mix.wgsl.template | 22 ++ .../contrib_ops/webgpu/moe/gate.wgsl.template | 87 +++++ .../moe/hidden_state_gather.wgsl.template | 28 ++ onnxruntime/contrib_ops/webgpu/moe/moe.cc | 38 ++ onnxruntime/contrib_ops/webgpu/moe/moe.h | 73 ++++ onnxruntime/contrib_ops/webgpu/moe/moe_base.h | 40 ++ onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 345 ++++++++++++++++++ onnxruntime/contrib_ops/webgpu/moe/qmoe.h | 37 ++ .../webgpu/moe/swiglu.wgsl.template | 25 ++ .../webgpu/moe/zero_tensor.wgsl.template | 10 + .../webgpu/webgpu_contrib_kernels.cc | 6 +- 11 files changed, 710 insertions(+), 1 deletion(-) create mode 100644 onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template create mode 100644 onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template create mode 100644 onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template create mode 100755 onnxruntime/contrib_ops/webgpu/moe/moe.cc create mode 100755 onnxruntime/contrib_ops/webgpu/moe/moe.h create mode 100755 onnxruntime/contrib_ops/webgpu/moe/moe_base.h create mode 100755 onnxruntime/contrib_ops/webgpu/moe/qmoe.cc create mode 100755 onnxruntime/contrib_ops/webgpu/moe/qmoe.h create mode 100644 onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template create mode 100644 onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template diff --git a/onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template new file mode 100644 index 0000000000000..1da1469df32f7 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// in: fc2_outputs [used_by, inter_size] +// in: router_values [num_tokens, num_experts] +// in: expert_tokens [used_by], mapping token idx to original token index +// out: output +// uniform: used_by, hidden_size, num_experts, expert_idx + +$MAIN { + let token_idx = expert_tokens[workgroup_idx]; + let step = uniforms.hidden_size / workgroup_size_x; + let wg_offset = local_idx * step; + let router_value_offset = token_idx * uniforms.num_experts + uniforms.expert_idx; + let router_value = router_values[router_value_offset]; + let fc2_outputs_offset = workgroup_idx * uniforms.hidden_size + wg_offset; + let output_offset = token_idx * uniforms.hidden_size + wg_offset; + for (var i = 0u; i < step; i++) { + let weight = fc2_outputs[fc2_outputs_offset + i]; + output[output_offset + i] += router_value * weight; + } +} diff --git a/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template new file mode 100644 index 0000000000000..d2af00158a925 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// +// MOE gate shader +// +// called with expert as local_idx and token_idx as workgroup_idx +// in: router_values [num_tokens, num_experts], per expert float we multiply final results with +// out: gate_counts [num_experts], number of tokens assigned to each expert +// out: gate_hidden [num_experts, num_tokens], token_idx assigned to each expert +// uniform: rows(num_tokens), cols(num_experts), token_offset +// + +#param is_fp16 +#param k + +const K: u32 = k; +#if is_fp16 +const MAX_FLOAT: f16 = 65504.0; +#else +const MAX_FLOAT: f32 = 3.402823466e+38; +#endif + +var shared_vals: array; +var shared_idxs: array; + +$MAIN { + let row = workgroup_idx; + if (row >= uniforms.rows) { + return; + } + let cols = uniforms.cols; + let base = row * cols; + + var max_val: hidden_state_element_t = -MAX_FLOAT; + var max_idx: u32 = 0u; + + if (global_idx < cols) { + atomicStore(&tokencount_for_expert[global_idx], 0u); + } + if (local_idx < cols) { + max_val = hidden_state[base + local_idx + uniforms.token_offset]; + max_idx = local_idx; + } + shared_vals[local_idx] = max_val; + shared_idxs[local_idx] = max_idx; + workgroupBarrier(); + + // K is small, use a simple bubble sort + for (var i = 0u; i < workgroup_size_x - 1u; i++) { + for (var j = 0u; j < workgroup_size_x - 1u - i; j++) { + if (local_idx == j && local_idx < cols && (local_idx + 1u) < cols) { + // Compare adjacent elements and swap if needed (descending order) + if (shared_vals[local_idx] < shared_vals[local_idx + 1u]) { + let temp_val = shared_vals[local_idx]; + let temp_idx = shared_idxs[local_idx]; + shared_vals[local_idx] = shared_vals[local_idx + 1u]; + shared_idxs[local_idx] = shared_idxs[local_idx + 1u]; + shared_vals[local_idx + 1u] = temp_val; + shared_idxs[local_idx + 1u] = temp_idx; + } + } + workgroupBarrier(); + } + } + if (local_idx < K) { + // found the top K experts for token, write to output + let expert_idx = shared_idxs[local_idx]; + let expert_base = expert_idx * uniforms.rows; + let target_idx = atomicAdd(&tokencount_for_expert[expert_idx], 1u); + hiddenstate_for_expert[expert_base + target_idx] = row + uniforms.token_offset; + } + workgroupBarrier(); + if (local_idx == 0u) { + // softmax + var sum : f32 = 0.0; + for (var i = 0u; i < K; i++) { + sum += exp(f32(shared_vals[i])); + } + // rows(num_tokens), cols(num_experts) + let output_base = row * uniforms.cols; + for (var i = 0u; i < K; i++) { + let expert_idx = shared_idxs[i]; + topk_values[output_base + expert_idx] = topk_values_value_t(exp(f32(shared_vals[i])) / sum); + } + } +} // MAIN diff --git a/onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template new file mode 100644 index 0000000000000..e07e0dc5edbc4 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// local_idx is used to copy hidden state row, workgroup_idx is token_idx +// workgroup_idx is the token index in this expert + +// in: hiddenstate_for_expert [num_experts, num_tokens] +// in: hidden_state(vec4) +// out: new_hidden_state(vec4) [used_by, hidden_size] +// out: expert_tokens [used_by] +// uniform: expert_idx, num_experts, num_tokens, hidden_size(vec4) + +$MAIN { + let expert_base = uniforms.expert_idx * uniforms.num_tokens; + let token_idx = hiddenstate_for_expert[expert_base + workgroup_idx]; + tokens[workgroup_idx] = token_idx; + + // copy hiden state for this token + let step = (uniforms.hidden_size + workgroup_size_x - 1) / workgroup_size_x; + let wg_offset = local_idx * step; + let src_offset = token_idx * uniforms.hidden_size + wg_offset; + let dst_offset = workgroup_idx * uniforms.hidden_size + wg_offset; + + for (var i = 0u; i < step; i++) { + let src = hidden_state[src_offset + i]; + new_hidden_state[dst_offset + i] = src; + } +} // MAIN diff --git a/onnxruntime/contrib_ops/webgpu/moe/moe.cc b/onnxruntime/contrib_ops/webgpu/moe/moe.cc new file mode 100755 index 0000000000000..4df2a7924c1fc --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/moe.cc @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_utils.h" +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "contrib_ops/webgpu/moe/moe_base.h" +#include "contrib_ops/webgpu/moe/moe.h" +#include "contrib_ops/cpu/moe/moe_helper.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +Status MoEProgram::GenerateShaderCode(ShaderHelper& shader) const { + return Status::OK(); +} + +Status MoE::ComputeInternal(ComputeContext& context) const { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "MoE is not implemented in WebGPU"); +} + +ONNX_OPERATOR_KERNEL_EX( + MoE, + kMSDomain, + 1, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()), + MoE); + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/moe/moe.h b/onnxruntime/contrib_ops/webgpu/moe/moe.h new file mode 100755 index 0000000000000..5c34b13fe77b0 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/moe.h @@ -0,0 +1,73 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/program.h" +#include "core/providers/webgpu/webgpu_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +class MoEProgram final : public Program { + public: + MoEProgram(TensorShape output_shape) : Program{"MoE"}, output_shape_{output_shape} {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"output_size", ProgramUniformVariableDataType::Uint32}); + + private: + TensorShape output_shape_; +}; + +class MoE : public WebGpuKernel { + public: + MoE(const OpKernelInfo& info) : WebGpuKernel(info) { + activation_alpha_ = static_cast(info.GetAttrOrDefault("activation_alpha", 1.0)); + activation_beta_ = static_cast(info.GetAttrOrDefault("activation_beta", 1.0)); + swiglu_fusion_ = static_cast(info.GetAttrOrDefault("swiglu_fusion", 0)); + swiglu_limit_ = info.GetAttrOrDefault("swiglu_limit", 0); + k_ = static_cast(info.GetAttrOrDefault("k", 4)); + normalize_routing_weights_ = info.GetAttrOrDefault("normalize_routing_weights", 0) == 1; + use_sparse_mixer_ = info.GetAttrOrDefault("use_sparse_mixer", 0) == 1; + std::string activation_type = info.GetAttrOrDefault("activation_type", "relu"); + if (activation_type == "relu") { + activation_type_ = MoEActivationType::Relu; + } else if (activation_type == "gelu") { + activation_type_ = MoEActivationType::Gelu; + } else if (activation_type == "silu") { + activation_type_ = MoEActivationType::Silu; + } else if (activation_type == "identity") { + activation_type_ = MoEActivationType::Identity; + } else if (activation_type == "swiglu") { + activation_type_ = MoEActivationType::SwiGLU; + } else { + ORT_THROW("Unsupported MoE activation type: ", activation_type); + } + + // for now webgpu only implements a subset of MoE features + // ORT_ENFORCE(normalize_routing_weights_ == 0, "normalize_routing_weights not supported"); + ORT_ENFORCE(use_sparse_mixer_ == 0, "use_sparse_mixer not supported"); + } + + Status ComputeInternal(ComputeContext& context) const override; + +protected: + int k_; + bool normalize_routing_weights_; + bool use_sparse_mixer_; + MoEActivationType activation_type_; + int swiglu_fusion_; + float swiglu_limit_; + float activation_alpha_; + float activation_beta_; +}; + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/moe/moe_base.h b/onnxruntime/contrib_ops/webgpu/moe/moe_base.h new file mode 100755 index 0000000000000..8bb036d34ca33 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/moe_base.h @@ -0,0 +1,40 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/program.h" +#include "core/providers/webgpu/webgpu_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +enum class MoEActivationType { + Relu, + Gelu, + Silu, + Identity, + SwiGLU, + +}; + +enum class MoEQuantType { + None = 0, + UINT4 = 1, + UINT8 = 2, +}; + +enum class MoEParallelType { + None = 0, + EP = 1, + TP = 2, + EPAndTP = 3, +}; + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc new file mode 100755 index 0000000000000..1185afeccb80b --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -0,0 +1,345 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_utils.h" +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "contrib_ops/webgpu/webgpu_contrib_kernels.h" +#include "contrib_ops/webgpu/moe/qmoe.h" +#include "contrib_ops/cpu/moe/moe_helper.h" +#include "contrib_ops/webgpu/quantization/matmul_nbits.h" +#include "core/providers/webgpu/math/gemm_packed.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +class GateProgram final : public Program { + public: + GateProgram(int k, bool is_fp16) : Program{"QmoeGate"}, k_{k}, is_fp16_{is_fp16} {}; + + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("hidden_state", ShaderUsage::UseElementTypeAlias); + shader.AddOutput("topk_values"); + shader.AddOutput("hiddenstate_for_expert"); + shader.AddOutput("tokencount_for_expert"); + + return WGSL_TEMPLATE_APPLY(shader, "moe/gate.wgsl.template", + WGSL_TEMPLATE_PARAMETER(is_fp16, is_fp16_), + WGSL_TEMPLATE_PARAMETER(k, k_)); + }; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"rows", ProgramUniformVariableDataType::Uint32}, + {"cols", ProgramUniformVariableDataType::Uint32}, + {"token_offset", ProgramUniformVariableDataType::Uint32}); + + private: + int k_; + bool is_fp16_; +}; + +class HiddenStateGatherProgram final : public Program { + public: + HiddenStateGatherProgram() : Program{"QmoeHiddenStateGather"} {}; + + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("hiddenstate_for_expert", ShaderUsage::UseElementTypeAlias); + shader.AddInput("hidden_state", ShaderUsage::UseElementTypeAlias); + shader.AddOutput("new_hidden_state"); + shader.AddOutput("tokens"); + + return WGSL_TEMPLATE_APPLY(shader, "moe/hidden_state_gather.wgsl.template"); + }; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"expert_idx", ProgramUniformVariableDataType::Uint32}, + {"num_experts", ProgramUniformVariableDataType::Uint32}, + {"num_tokens", ProgramUniformVariableDataType::Uint32}, + {"hidden_size", ProgramUniformVariableDataType::Uint32}); + + private: +}; + +class ZeroTensorProgram final : public Program { + public: + ZeroTensorProgram() : Program{"QmoeZeroTensor"} {}; + + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddOutput("tensor", ShaderUsage::UseElementTypeAlias); + return WGSL_TEMPLATE_APPLY(shader, "moe/zero_tensor.wgsl.template"); + }; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"size", ProgramUniformVariableDataType::Uint32}); + + private: +}; + +class SwigLuProgram final : public Program { + public: + SwigLuProgram() : Program{"SwigLu"} { + }; + + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("input", ShaderUsage::UseElementTypeAlias); + shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); + + return WGSL_TEMPLATE_APPLY(shader, "moe/swiglu.wgsl.template"); + }; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"rows", ProgramUniformVariableDataType::Uint32}, + {"cols", ProgramUniformVariableDataType::Uint32}, + {"alpha", ProgramUniformVariableDataType::Float32}, + {"beta", ProgramUniformVariableDataType::Float32}, + {"swiglu_limit", ProgramUniformVariableDataType::Float32}); + + private: +}; + +class QMoEFinalMixProgram final : public Program { + public: + QMoEFinalMixProgram() : Program{"QMoEFinalMix"} {} + + Status GenerateShaderCode(ShaderHelper& shader) const override { + shader.AddInput("fc2_outputs", ShaderUsage::UseElementTypeAlias); + shader.AddInput("router_values", ShaderUsage::UseElementTypeAlias); + shader.AddInput("expert_tokens", ShaderUsage::UseElementTypeAlias); + shader.AddOutput("output", ShaderUsage::UseElementTypeAlias); + + return WGSL_TEMPLATE_APPLY(shader, "moe/final_mix.wgsl.template"); + } + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES( + {"used_by", ProgramUniformVariableDataType::Uint32}, + {"hidden_size", ProgramUniformVariableDataType::Uint32}, + {"num_experts", ProgramUniformVariableDataType::Uint32}, + {"expert_idx", ProgramUniformVariableDataType::Uint32}); + + private: +}; + + +Status QMoE::ComputeInternal(ComputeContext& context) const { + const Tensor* hidden_state = context.Input(0); + const Tensor* router_logits = context.Input(1); + // fc1 is gate_up_proj + const Tensor* fc1_experts_weights = context.Input(2); + const Tensor* fc1_scales = context.Input(3); + const Tensor* fc1_experts_bias_optional = context.Input(4); + // fc2 is gate_down_proj + const Tensor* fc2_experts_weights = context.Input(5); + const Tensor* fc2_scales = context.Input(6); + const Tensor* fc2_experts_bias_optional = context.Input(7); + const Tensor* fc3_experts_weights_optional = context.Input(8); + const Tensor* fc3_scales_optional = context.Input(9); + const Tensor* fc3_experts_bias_optional = context.Input(10); + + MoEParameters moe_params; + + ORT_RETURN_IF_ERROR(::onnxruntime::contrib::moe_helper::CheckInputs( + moe_params, hidden_state, router_logits, + fc1_experts_weights, fc1_experts_bias_optional, fc1_scales, + fc2_experts_weights, fc2_experts_bias_optional, fc2_scales, + fc3_experts_weights_optional, fc3_experts_bias_optional, fc3_scales_optional, + expert_weight_bits_ == 4 ? 2 : 1, + activation_type_ == MoEActivationType::SwiGLU, block_size_)); + + const auto& input_shape = hidden_state->Shape(); + + // SwiGLU validation + bool is_swiglu = (activation_type_ == MoEActivationType::SwiGLU); + if (is_swiglu && fc3_experts_weights_optional != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "SwiGLU activation is not supported with fc3. Gate weights should be concatenated with FC1 weights."); + } + if (!is_swiglu && fc3_experts_weights_optional != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "FC3 gating is not yet implemented for non-SwiGLU activations on CPU."); + } + + const int max_tokens = 256; // TODO: maybe 512 ? + const uint32_t num_experts = static_cast(moe_params.num_experts); + const uint32_t hidden_size = static_cast(moe_params.hidden_size); + const int64_t fc1_output_size = is_swiglu && swiglu_fusion_ > 0 ? 2 * moe_params.inter_size : moe_params.inter_size; + const bool is_fp16 = hidden_state->DataType() == DataTypeImpl::GetType(); + const auto dtype = is_fp16 ? DataTypeImpl::GetType() : DataTypeImpl::GetType(); + const auto dtype_uint32 = DataTypeImpl::GetType(); + + const int64_t K_fc1 = moe_params.hidden_size; // left_shape[left_num_dims - 1] + const int64_t N_fc1 = fc1_output_size; // right_shape[right_num_dims - 1] + const int64_t K_fc2 = moe_params.inter_size; // left_shape[left_num_dims - 1] + const int64_t N_fc2 = moe_params.inter_size; // right_shape[right_num_dims - 1] + const int64_t accuracy_level = 4; + const int64_t block_size = (block_size_ != 0) ? block_size_ : fc1_experts_weights->Shape()[2]; + Status status; + + Tensor* output_tensor = context.Output(0, input_shape); + const int total_output_size = static_cast(input_shape.Size()) / 4; + + // we are accumulating expert results into output_tensor, need to initialize to zero + ZeroTensorProgram zero; + zero + .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, 4}) + .SetDispatchGroupSize((total_output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) + .AddUniformVariables({static_cast(total_output_size)}); + ORT_RETURN_IF_ERROR(context.RunProgram(zero)); + + // process tokens in chunks of max_tokens to put some cap on memory usage + for (int token_offset = 0; token_offset < moe_params.num_rows; token_offset += max_tokens) { + // + // Step 1: run the gate to get router indices and values + // + int num_tokens = static_cast(moe_params.num_rows) - token_offset; + if (num_tokens > max_tokens) { + num_tokens = max_tokens; + } + TensorShape gate_value_shape({num_tokens, num_experts}); // use max_tokens ? + TensorShape gate_hidden_shape({num_experts, num_tokens}); // use max_tokens ? + TensorShape gate_count_shape({num_experts}); + + // router_values: per expert float we multiply final results with + Tensor router_values = context.CreateGPUTensor(dtype, gate_value_shape); + // gate_counts: number of tokens assigned to each expert + Tensor gate_counts = context.CreateGPUTensor(dtype_uint32, gate_count_shape); + // gate_hidden: token_idx assigned to each expert + Tensor gate_hidden = context.CreateGPUTensor(dtype_uint32, gate_hidden_shape); + + GateProgram gate{k_, is_fp16}; + gate + .AddInputs({{router_logits, ProgramTensorMetadataDependency::Type}}) + .AddOutput({&router_values, ProgramTensorMetadataDependency::None}) + .AddOutput({&gate_hidden, ProgramTensorMetadataDependency::None}) + .AddOutput({&gate_counts, ProgramTensorMetadataDependency::None, ProgramOutput::Atomic}) + .SetWorkgroupSize(num_experts) + .SetDispatchGroupSize(static_cast(num_tokens)) + .AddUniformVariables({static_cast(num_tokens), + num_experts, + static_cast(token_offset)}) + .CacheHint(k_, is_fp16 ? "fp16" : "fp32"); + + ORT_RETURN_IF_ERROR(context.RunProgram(gate)); + + Tensor gate_counts_cpu = context.CreateCPUTensor(dtype_uint32, gate_count_shape); + ORT_RETURN_IF_ERROR(Info().GetDataTransferManager().CopyTensor(gate_counts, gate_counts_cpu)); + + for (uint32_t expert_idx = 0; expert_idx < num_experts; expert_idx++) { + uint32_t used_by = *(gate_counts_cpu.Data() + expert_idx); + if (used_by <= 0) { + continue; + } + + // + // Step 2: for each expert, gather the hidden_state rows assigned to it + // FIXME: use vec4 + // + TensorShape expert_hidden_shape({used_by, moe_params.hidden_size}); + // expert_hidden: hidden states assigned to this expert + Tensor expert_hidden = context.CreateGPUTensor(dtype, expert_hidden_shape); + TensorShape expert_tokens_shape({used_by}); + // expert_tokens: token_idx that match expert_hidden rows + Tensor expert_tokens = context.CreateGPUTensor(dtype_uint32, expert_tokens_shape); + HiddenStateGatherProgram gather; + gather + .AddInputs({{&gate_hidden, ProgramTensorMetadataDependency::Type}}) + .AddInputs({{hidden_state, ProgramTensorMetadataDependency::Type, 1}}) + .AddOutput({&expert_hidden, ProgramTensorMetadataDependency::None, 1}) + .AddOutput({&expert_tokens, ProgramTensorMetadataDependency::None}) + .SetDispatchGroupSize(used_by) + .AddUniformVariables({expert_idx, + num_experts, + static_cast(num_tokens), + hidden_size}); + ORT_RETURN_IF_ERROR(context.RunProgram(gather)); + + TensorShape fc1_output_shape({used_by, fc1_output_size}); + Tensor fc1_outputs = context.CreateGPUTensor(dtype, fc1_output_shape); + TensorShape fc1_activated_shape({used_by, moe_params.inter_size}); + Tensor fc1_activated = context.CreateGPUTensor(dtype, fc1_activated_shape); + TensorShape fc2_output_shape({used_by, moe_params.inter_size}); + Tensor fc2_outputs = context.CreateGPUTensor(dtype, fc2_output_shape); + + // + // Step 3: matmul the hidden_state with fc1 (gate_up) of the selected experts + // + status = ApplyMatMulNBits(&expert_hidden, fc1_experts_weights, fc1_scales, nullptr, fc1_experts_bias_optional, + K_fc1, N_fc1, block_size, accuracy_level, expert_weight_bits_, context, + &fc1_outputs, expert_idx); + ORT_RETURN_IF_ERROR(status); + + // + // Step 4: apply swigly + // + if (is_swiglu) { + SwigLuProgram swiglu; + swiglu + .AddInputs({{&fc1_outputs, ProgramTensorMetadataDependency::Type, 2}}) + .AddOutput({&fc1_activated, ProgramTensorMetadataDependency::None}) + .SetWorkgroupSize(128) + .SetDispatchGroupSize(((used_by * static_cast(moe_params.inter_size)) + 127) / 128) + .AddUniformVariables({static_cast(used_by), + static_cast(moe_params.inter_size), + activation_alpha_, + activation_beta_, + swiglu_limit_}); + ORT_RETURN_IF_ERROR(context.RunProgram(swiglu)); + } else { + ORT_THROW("only swiglu is supported for now."); + } + + // + // Step 5: multiply fc1_activated with fc2 (gate_down) of the selected experts + // + status = ApplyMatMulNBits(&fc1_activated, fc2_experts_weights, fc2_scales, nullptr, fc2_experts_bias_optional, + K_fc2, N_fc2, block_size, accuracy_level, expert_weight_bits_, context, + &fc2_outputs, expert_idx); + ORT_RETURN_IF_ERROR(status); + + // + // Step 5: multiply fc2_outputs with router_values and accumulate + // + QMoEFinalMixProgram final_mix; + final_mix + .AddInputs({{&fc2_outputs, ProgramTensorMetadataDependency::Type}}) + .AddInputs({{&router_values, ProgramTensorMetadataDependency::Type}}) + .AddInputs({{&expert_tokens, ProgramTensorMetadataDependency::Type}}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::None}) + .SetDispatchGroupSize(used_by) + .AddUniformVariables({used_by, + hidden_size, + num_experts, + expert_idx}); + + ORT_RETURN_IF_ERROR(context.RunProgram(final_mix)); + } + } + + return Status::OK(); +} + +namespace { +const std::vector& QMoET1Constraint() { + static std::vector types{ + DataTypeImpl::GetTensorType()}; + return types; +} +} // namespace + +ONNX_OPERATOR_KERNEL_EX( + QMoE, + kMSDomain, + 1, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedFloatTypes()) + .TypeConstraint("T1", QMoET1Constraint()) + .TypeConstraint("T2", WebGpuSupportedFloatTypes()), + QMoE); + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.h b/onnxruntime/contrib_ops/webgpu/moe/qmoe.h new file mode 100755 index 0000000000000..2b398e514c44d --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.h @@ -0,0 +1,37 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/program.h" +#include "core/providers/webgpu/webgpu_kernel.h" +#include "contrib_ops/webgpu/moe/moe_base.h" +#include "contrib_ops/webgpu/moe/moe.h" +#include "core/providers/webgpu/math/matmul.h" + +namespace onnxruntime { +namespace contrib { +namespace webgpu { + +using namespace onnxruntime::webgpu; +using onnxruntime::webgpu::ComputeContext; + +class QMoE final : public MoE { + public: + QMoE(const OpKernelInfo& info) : MoE(info) { + ORT_ENFORCE(info.GetAttr("expert_weight_bits", &expert_weight_bits_).IsOK()); + ORT_ENFORCE(expert_weight_bits_ == 8 || expert_weight_bits_ == 4, + "expert_weight_bits must be 4 or 8, but got ", expert_weight_bits_); + block_size_ = static_cast(info.GetAttrOrDefault("block_size", 0)); + } + + Status ComputeInternal(ComputeContext& context) const override; + + private: + int64_t expert_weight_bits_; + int64_t block_size_; +}; + +} // namespace webgpu +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template new file mode 100644 index 0000000000000..6658c8120bb5e --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#use guardAgainstOutOfBoundsWorkgroupSizes + +// This implemnts: +// gate, up = gate_up[..., 0::2], gate_up[..., 1::2] +// gate = gate.clamp(min=None, max=self.limit) +// up = up.clamp(min=-self.limit, max=self.limit) +// glu = gate * torch.sigmoid(gate * self.alpha) +// gated_output = (up + 1) * glu + +$MAIN { + let total = uniforms.rows * uniforms.cols; + guardAgainstOutOfBoundsWorkgroupSizes(total); + + let row = global_idx / uniforms.cols; + let col = global_idx % uniforms.cols; + let base = row * uniforms.cols; + let gate_up = vec2(input[base + col]); + let gate_val = min(gate_up.x, uniforms.swiglu_limit); + let up_val = clamp(gate_up.y, -uniforms.swiglu_limit, uniforms.swiglu_limit); + let glu = gate_val * 1.0f / (1.0f + exp(-uniforms.alpha * gate_val)); + output[global_idx] = output_element_t(glu * (up_val + uniforms.beta)); +} // MAIN diff --git a/onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template new file mode 100644 index 0000000000000..16fca3195eccb --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template @@ -0,0 +1,10 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + + +$MAIN { + if (global_idx > uniforms.size) { + return; + }; + tensor[global_idx] = vec4(0.0); +} // MAIN diff --git a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc index fe0bc5dee92ff..e3573534f94b9 100644 --- a/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc @@ -28,6 +28,8 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, SkipLayerNormalization); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, SimplifiedLayerNormalization); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, SkipSimplifiedLayerNormalization); +// class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, MoE); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSDomain, 1, QMoE); template <> KernelCreateInfo BuildKernelCreateInfo() { @@ -53,7 +55,9 @@ Status RegisterWebGpuContribKernels(KernelRegistry& kernel_registry, bool enable // LayerNormalization used to be a contrib op that (incorrectly) used kOnnxDomain so we need to version it BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo}; + BuildKernelCreateInfo, + // BuildKernelCreateInfo, + BuildKernelCreateInfo}; for (auto& function_table_entry : function_table) { KernelCreateInfo info = function_table_entry(); From 6afc761c2bd07421cb33b1f8689e8f7bb646adc8 Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 4 Nov 2025 08:41:07 -0800 Subject: [PATCH 02/15] Update onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template index 6658c8120bb5e..355f31228255f 100644 --- a/onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/moe/swiglu.wgsl.template @@ -3,7 +3,7 @@ #use guardAgainstOutOfBoundsWorkgroupSizes -// This implemnts: +// This implements: // gate, up = gate_up[..., 0::2], gate_up[..., 1::2] // gate = gate.clamp(min=None, max=self.limit) // up = up.clamp(min=-self.limit, max=self.limit) From 61dfdbcbb21a3d4edd3397ad9c6ca150d02e36bf Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 4 Nov 2025 08:41:23 -0800 Subject: [PATCH 03/15] Update onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template index 16fca3195eccb..c45a49246cde0 100644 --- a/onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/moe/zero_tensor.wgsl.template @@ -3,7 +3,7 @@ $MAIN { - if (global_idx > uniforms.size) { + if (global_idx >= uniforms.size) { return; }; tensor[global_idx] = vec4(0.0); From f64e588bf6b8664b517a21b1f1fbe8505ee54f2d Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 4 Nov 2025 08:42:02 -0800 Subject: [PATCH 04/15] Update onnxruntime/contrib_ops/webgpu/moe/qmoe.cc Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index 1185afeccb80b..a255de11a3edb 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -272,7 +272,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { ORT_RETURN_IF_ERROR(status); // - // Step 4: apply swigly + // Step 4: apply swiglu // if (is_swiglu) { SwigLuProgram swiglu; From 48f9386613f57d1bd21ecdcc5a448fb4638545c5 Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 4 Nov 2025 08:42:45 -0800 Subject: [PATCH 05/15] Update onnxruntime/contrib_ops/webgpu/moe/qmoe.cc Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index a255de11a3edb..4baf0f04c4a43 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -159,7 +159,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { } if (!is_swiglu && fc3_experts_weights_optional != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "FC3 gating is not yet implemented for non-SwiGLU activations on CPU."); + "FC3 gating is not yet implemented for non-SwiGLU activations on WebGPU."); } const int max_tokens = 256; // TODO: maybe 512 ? From 32ee5bc3e1d311a3e565e7d897424d95f40a0dc2 Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 4 Nov 2025 08:43:04 -0800 Subject: [PATCH 06/15] Update onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template index e07e0dc5edbc4..d64d949b9d93d 100644 --- a/onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/moe/hidden_state_gather.wgsl.template @@ -15,7 +15,7 @@ $MAIN { let token_idx = hiddenstate_for_expert[expert_base + workgroup_idx]; tokens[workgroup_idx] = token_idx; - // copy hiden state for this token + // copy hidden state for this token let step = (uniforms.hidden_size + workgroup_size_x - 1) / workgroup_size_x; let wg_offset = local_idx * step; let src_offset = token_idx * uniforms.hidden_size + wg_offset; From 21d8cc937e532099ca02cb9bd47d56920b519484 Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 4 Nov 2025 08:43:53 -0800 Subject: [PATCH 07/15] Update onnxruntime/contrib_ops/webgpu/moe/qmoe.cc Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index 4baf0f04c4a43..f81a0a10f09a6 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -300,7 +300,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { ORT_RETURN_IF_ERROR(status); // - // Step 5: multiply fc2_outputs with router_values and accumulate + // Step 6: multiply fc2_outputs with router_values and accumulate // QMoEFinalMixProgram final_mix; final_mix From 9152e1fdaea411852ce0000508f2aad9880fcdb7 Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 4 Nov 2025 08:45:03 -0800 Subject: [PATCH 08/15] Update onnxruntime/contrib_ops/webgpu/moe/qmoe.cc Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index f81a0a10f09a6..849da4f43bbb4 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -162,7 +162,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { "FC3 gating is not yet implemented for non-SwiGLU activations on WebGPU."); } - const int max_tokens = 256; // TODO: maybe 512 ? + const int max_tokens = 256; // TODO: maybe 512 ? const uint32_t num_experts = static_cast(moe_params.num_experts); const uint32_t hidden_size = static_cast(moe_params.hidden_size); const int64_t fc1_output_size = is_swiglu && swiglu_fusion_ > 0 ? 2 * moe_params.inter_size : moe_params.inter_size; From 9f055eca044354e025c21ea8a2c71147cc2b15f2 Mon Sep 17 00:00:00 2001 From: gs Date: Tue, 4 Nov 2025 09:45:35 -0800 Subject: [PATCH 09/15] fix lint errors --- onnxruntime/contrib_ops/webgpu/moe/moe.cc | 4 ++-- onnxruntime/contrib_ops/webgpu/moe/moe.h | 2 +- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 1 - 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/moe.cc b/onnxruntime/contrib_ops/webgpu/moe/moe.cc index 4df2a7924c1fc..a753b2d4c70a6 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/moe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/moe.cc @@ -16,11 +16,11 @@ namespace webgpu { using namespace onnxruntime::webgpu; using onnxruntime::webgpu::ComputeContext; -Status MoEProgram::GenerateShaderCode(ShaderHelper& shader) const { +Status MoEProgram::GenerateShaderCode(ShaderHelper& /*unused*/) const { return Status::OK(); } -Status MoE::ComputeInternal(ComputeContext& context) const { +Status MoE::ComputeInternal(ComputeContext& /*unused*/) const { return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "MoE is not implemented in WebGPU"); } diff --git a/onnxruntime/contrib_ops/webgpu/moe/moe.h b/onnxruntime/contrib_ops/webgpu/moe/moe.h index 5c34b13fe77b0..5e329dc12b5c9 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/moe.h +++ b/onnxruntime/contrib_ops/webgpu/moe/moe.h @@ -57,7 +57,7 @@ class MoE : public WebGpuKernel { Status ComputeInternal(ComputeContext& context) const override; -protected: + protected: int k_; bool normalize_routing_weights_; bool use_sparse_mixer_; diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index 849da4f43bbb4..0863d1bc60e20 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -123,7 +123,6 @@ class QMoEFinalMixProgram final : public Program { private: }; - Status QMoE::ComputeInternal(ComputeContext& context) const { const Tensor* hidden_state = context.Input(0); const Tensor* router_logits = context.Input(1); From 336cafad47505175adcb58580ccdf077aef9d35d Mon Sep 17 00:00:00 2001 From: gs Date: Mon, 10 Nov 2025 09:10:33 -0800 Subject: [PATCH 10/15] address review feedback --- onnxruntime/contrib_ops/webgpu/moe/moe_base.h | 7 ------- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 19 ++++++++----------- 2 files changed, 8 insertions(+), 18 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/moe_base.h b/onnxruntime/contrib_ops/webgpu/moe/moe_base.h index 8bb036d34ca33..bab99fe51c88b 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/moe_base.h +++ b/onnxruntime/contrib_ops/webgpu/moe/moe_base.h @@ -28,13 +28,6 @@ enum class MoEQuantType { UINT8 = 2, }; -enum class MoEParallelType { - None = 0, - EP = 1, - TP = 2, - EPAndTP = 3, -}; - } // namespace webgpu } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index 0863d1bc60e20..80386bdae8e3f 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -152,11 +152,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { // SwiGLU validation bool is_swiglu = (activation_type_ == MoEActivationType::SwiGLU); - if (is_swiglu && fc3_experts_weights_optional != nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "SwiGLU activation is not supported with fc3. Gate weights should be concatenated with FC1 weights."); - } - if (!is_swiglu && fc3_experts_weights_optional != nullptr) { + if (fc3_experts_weights_optional != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "FC3 gating is not yet implemented for non-SwiGLU activations on WebGPU."); } @@ -174,16 +170,17 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { const int64_t K_fc2 = moe_params.inter_size; // left_shape[left_num_dims - 1] const int64_t N_fc2 = moe_params.inter_size; // right_shape[right_num_dims - 1] const int64_t accuracy_level = 4; - const int64_t block_size = (block_size_ != 0) ? block_size_ : fc1_experts_weights->Shape()[2]; + const int64_t block_size_fc1 = (block_size_ != 0) ? block_size_ : K_fc1; + const int64_t block_size_fc2 = (block_size_ != 0) ? block_size_ : K_fc2; Status status; Tensor* output_tensor = context.Output(0, input_shape); - const int total_output_size = static_cast(input_shape.Size()) / 4; + const int total_output_size = (static_cast(input_shape.Size()) + 3) / 4; // we are accumulating expert results into output_tensor, need to initialize to zero ZeroTensorProgram zero; zero - .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, 4}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}) .SetDispatchGroupSize((total_output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .AddUniformVariables({static_cast(total_output_size)}); ORT_RETURN_IF_ERROR(context.RunProgram(zero)); @@ -259,14 +256,14 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { Tensor fc1_outputs = context.CreateGPUTensor(dtype, fc1_output_shape); TensorShape fc1_activated_shape({used_by, moe_params.inter_size}); Tensor fc1_activated = context.CreateGPUTensor(dtype, fc1_activated_shape); - TensorShape fc2_output_shape({used_by, moe_params.inter_size}); + TensorShape fc2_output_shape({used_by, N_fc2}); Tensor fc2_outputs = context.CreateGPUTensor(dtype, fc2_output_shape); // // Step 3: matmul the hidden_state with fc1 (gate_up) of the selected experts // status = ApplyMatMulNBits(&expert_hidden, fc1_experts_weights, fc1_scales, nullptr, fc1_experts_bias_optional, - K_fc1, N_fc1, block_size, accuracy_level, expert_weight_bits_, context, + K_fc1, N_fc1, block_size_fc1, accuracy_level, expert_weight_bits_, context, &fc1_outputs, expert_idx); ORT_RETURN_IF_ERROR(status); @@ -294,7 +291,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { // Step 5: multiply fc1_activated with fc2 (gate_down) of the selected experts // status = ApplyMatMulNBits(&fc1_activated, fc2_experts_weights, fc2_scales, nullptr, fc2_experts_bias_optional, - K_fc2, N_fc2, block_size, accuracy_level, expert_weight_bits_, context, + K_fc2, N_fc2, block_size_fc2, accuracy_level, expert_weight_bits_, context, &fc2_outputs, expert_idx); ORT_RETURN_IF_ERROR(status); From 75fc052230e0953cbe6ce9b82c21e8fcbac8c967 Mon Sep 17 00:00:00 2001 From: gs Date: Mon, 10 Nov 2025 10:54:54 -0800 Subject: [PATCH 11/15] address review feedback --- onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template index d2af00158a925..a2cadfc938666 100644 --- a/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template @@ -30,7 +30,6 @@ $MAIN { return; } let cols = uniforms.cols; - let base = row * cols; var max_val: hidden_state_element_t = -MAX_FLOAT; var max_idx: u32 = 0u; @@ -39,7 +38,7 @@ $MAIN { atomicStore(&tokencount_for_expert[global_idx], 0u); } if (local_idx < cols) { - max_val = hidden_state[base + local_idx + uniforms.token_offset]; + max_val = hidden_state[(row + uniforms.token_offset) * cols + local_idx]; max_idx = local_idx; } shared_vals[local_idx] = max_val; From 55fea84259832c80bd5a32e3e1b06523d2fc8aae Mon Sep 17 00:00:00 2001 From: gs Date: Tue, 11 Nov 2025 12:44:23 -0800 Subject: [PATCH 12/15] fix chunking of tokens into max_token blocks --- .../contrib_ops/webgpu/bert/attention.cc | 2 +- .../webgpu/moe/final_mix.wgsl.template | 4 +++- .../contrib_ops/webgpu/moe/gate.wgsl.template | 4 ++-- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 21 ++++++++++++------- 4 files changed, 19 insertions(+), 12 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/attention.cc b/onnxruntime/contrib_ops/webgpu/bert/attention.cc index ca20845a8184d..382f67a0f2041 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/attention.cc @@ -284,7 +284,7 @@ Status InPlaceSoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { if (has_head_sink_) { // Handle head sink - shader.MainFunctionBody() << "let sink_value: f32 = head_sink[head_idx];\n" + shader.MainFunctionBody() << "let sink_value: f32 = f32(head_sink[head_idx]);\n" << "var max_value = sink_value;\n"; } else if (use_smooth_softmax_) { shader.MainFunctionBody() << "var max_value: f32 = 0.0;\n"; diff --git a/onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template index 1da1469df32f7..80887b845f915 100644 --- a/onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template @@ -11,7 +11,9 @@ $MAIN { let token_idx = expert_tokens[workgroup_idx]; let step = uniforms.hidden_size / workgroup_size_x; let wg_offset = local_idx * step; - let router_value_offset = token_idx * uniforms.num_experts + uniforms.expert_idx; + // token_idx is the offset into hidden state while fc2_outputs is for the chunk and + // we need to substract uniforms.token_offset + let router_value_offset = (token_idx - uniforms.token_offset) * uniforms.num_experts + uniforms.expert_idx; let router_value = router_values[router_value_offset]; let fc2_outputs_offset = workgroup_idx * uniforms.hidden_size + wg_offset; let output_offset = token_idx * uniforms.hidden_size + wg_offset; diff --git a/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template index a2cadfc938666..22d965a70e877 100644 --- a/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template @@ -30,6 +30,7 @@ $MAIN { return; } let cols = uniforms.cols; + let output_base = row * cols; var max_val: hidden_state_element_t = -MAX_FLOAT; var max_idx: u32 = 0u; @@ -69,6 +70,7 @@ $MAIN { let target_idx = atomicAdd(&tokencount_for_expert[expert_idx], 1u); hiddenstate_for_expert[expert_base + target_idx] = row + uniforms.token_offset; } + topk_values[output_base + local_idx] = topk_values_value_t(0); workgroupBarrier(); if (local_idx == 0u) { // softmax @@ -76,8 +78,6 @@ $MAIN { for (var i = 0u; i < K; i++) { sum += exp(f32(shared_vals[i])); } - // rows(num_tokens), cols(num_experts) - let output_base = row * uniforms.cols; for (var i = 0u; i < K; i++) { let expert_idx = shared_idxs[i]; topk_values[output_base + expert_idx] = topk_values_value_t(exp(f32(shared_vals[i])) / sum); diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index 80386bdae8e3f..5b4a2ebf30d9f 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -118,7 +118,8 @@ class QMoEFinalMixProgram final : public Program { {"used_by", ProgramUniformVariableDataType::Uint32}, {"hidden_size", ProgramUniformVariableDataType::Uint32}, {"num_experts", ProgramUniformVariableDataType::Uint32}, - {"expert_idx", ProgramUniformVariableDataType::Uint32}); + {"expert_idx", ProgramUniformVariableDataType::Uint32}, + {"token_offset", ProgramUniformVariableDataType::Uint32}); private: }; @@ -157,7 +158,9 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { "FC3 gating is not yet implemented for non-SwiGLU activations on WebGPU."); } - const int max_tokens = 256; // TODO: maybe 512 ? + // process tokens in chunks of max_tokens to put some cap on memory usage + const int max_tokens = 512; + const uint32_t num_experts = static_cast(moe_params.num_experts); const uint32_t hidden_size = static_cast(moe_params.hidden_size); const int64_t fc1_output_size = is_swiglu && swiglu_fusion_ > 0 ? 2 * moe_params.inter_size : moe_params.inter_size; @@ -165,10 +168,10 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { const auto dtype = is_fp16 ? DataTypeImpl::GetType() : DataTypeImpl::GetType(); const auto dtype_uint32 = DataTypeImpl::GetType(); - const int64_t K_fc1 = moe_params.hidden_size; // left_shape[left_num_dims - 1] - const int64_t N_fc1 = fc1_output_size; // right_shape[right_num_dims - 1] - const int64_t K_fc2 = moe_params.inter_size; // left_shape[left_num_dims - 1] - const int64_t N_fc2 = moe_params.inter_size; // right_shape[right_num_dims - 1] + const int64_t K_fc1 = moe_params.hidden_size; + const int64_t N_fc1 = fc1_output_size; + const int64_t K_fc2 = moe_params.inter_size; + const int64_t N_fc2 = moe_params.inter_size; const int64_t accuracy_level = 4; const int64_t block_size_fc1 = (block_size_ != 0) ? block_size_ : K_fc1; const int64_t block_size_fc2 = (block_size_ != 0) ? block_size_ : K_fc2; @@ -180,7 +183,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { // we are accumulating expert results into output_tensor, need to initialize to zero ZeroTensorProgram zero; zero - .AddOutput({output_tensor, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, 4}) .SetDispatchGroupSize((total_output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .AddUniformVariables({static_cast(total_output_size)}); ORT_RETURN_IF_ERROR(context.RunProgram(zero)); @@ -203,6 +206,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { // gate_counts: number of tokens assigned to each expert Tensor gate_counts = context.CreateGPUTensor(dtype_uint32, gate_count_shape); // gate_hidden: token_idx assigned to each expert + // token_idx is the global index into hidden_state Tensor gate_hidden = context.CreateGPUTensor(dtype_uint32, gate_hidden_shape); GateProgram gate{k_, is_fp16}; @@ -308,7 +312,8 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { .AddUniformVariables({used_by, hidden_size, num_experts, - expert_idx}); + expert_idx, + static_cast(token_offset)}); ORT_RETURN_IF_ERROR(context.RunProgram(final_mix)); } From 9aac0543971337a7debfa6f6c8cf0511276def14 Mon Sep 17 00:00:00 2001 From: gs Date: Tue, 11 Nov 2025 13:18:00 -0800 Subject: [PATCH 13/15] reflect changes in main --- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index 5b4a2ebf30d9f..7f92925489123 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -138,14 +138,23 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { const Tensor* fc3_experts_weights_optional = context.Input(8); const Tensor* fc3_scales_optional = context.Input(9); const Tensor* fc3_experts_bias_optional = context.Input(10); + // zero points, not supported yet + const Tensor* fc1_zero_points = context.Input(11); + const Tensor* fc2_zero_points = context.Input(12); + const Tensor* fc3_zero_points = context.Input(13); MoEParameters moe_params; + if (fc1_zero_points || fc2_zero_points || fc3_zero_points) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "zero_points for QMoE are not yet supported on WebGPU."); + } + ORT_RETURN_IF_ERROR(::onnxruntime::contrib::moe_helper::CheckInputs( moe_params, hidden_state, router_logits, - fc1_experts_weights, fc1_experts_bias_optional, fc1_scales, - fc2_experts_weights, fc2_experts_bias_optional, fc2_scales, - fc3_experts_weights_optional, fc3_experts_bias_optional, fc3_scales_optional, + fc1_experts_weights, fc1_experts_bias_optional, fc1_scales, fc1_zero_points, + fc2_experts_weights, fc2_experts_bias_optional, fc2_scales, fc2_zero_points, + fc3_experts_weights_optional, fc3_experts_bias_optional, fc3_scales_optional, fc3_zero_points, expert_weight_bits_ == 4 ? 2 : 1, activation_type_ == MoEActivationType::SwiGLU, block_size_)); @@ -153,7 +162,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { // SwiGLU validation bool is_swiglu = (activation_type_ == MoEActivationType::SwiGLU); - if (fc3_experts_weights_optional != nullptr) { + if (fc3_experts_weights_optional) { return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "FC3 gating is not yet implemented for non-SwiGLU activations on WebGPU."); } From d6fd933a125e49416c25deea5c55a49bb787fab2 Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 11 Nov 2025 17:00:43 -0800 Subject: [PATCH 14/15] review feedback --- onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template | 3 +-- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 8 ++++---- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template index 22d965a70e877..1214777009a8d 100644 --- a/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template @@ -44,6 +44,7 @@ $MAIN { } shared_vals[local_idx] = max_val; shared_idxs[local_idx] = max_idx; + topk_values[output_base + local_idx] = topk_values_value_t(0); workgroupBarrier(); // K is small, use a simple bubble sort @@ -70,8 +71,6 @@ $MAIN { let target_idx = atomicAdd(&tokencount_for_expert[expert_idx], 1u); hiddenstate_for_expert[expert_base + target_idx] = row + uniforms.token_offset; } - topk_values[output_base + local_idx] = topk_values_value_t(0); - workgroupBarrier(); if (local_idx == 0u) { // softmax var sum : f32 = 0.0; diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index 7f92925489123..d37e25c6b27c6 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -164,7 +164,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { bool is_swiglu = (activation_type_ == MoEActivationType::SwiGLU); if (fc3_experts_weights_optional) { return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "FC3 gating is not yet implemented for non-SwiGLU activations on WebGPU."); + "FC3 gating is not yet implemented on WebGPU."); } // process tokens in chunks of max_tokens to put some cap on memory usage @@ -180,7 +180,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { const int64_t K_fc1 = moe_params.hidden_size; const int64_t N_fc1 = fc1_output_size; const int64_t K_fc2 = moe_params.inter_size; - const int64_t N_fc2 = moe_params.inter_size; + const int64_t N_fc2 = moe_params.hidden_size; const int64_t accuracy_level = 4; const int64_t block_size_fc1 = (block_size_ != 0) ? block_size_ : K_fc1; const int64_t block_size_fc2 = (block_size_ != 0) ? block_size_ : K_fc2; @@ -192,7 +192,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { // we are accumulating expert results into output_tensor, need to initialize to zero ZeroTensorProgram zero; zero - .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, 4}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}) .SetDispatchGroupSize((total_output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .AddUniformVariables({static_cast(total_output_size)}); ORT_RETURN_IF_ERROR(context.RunProgram(zero)); @@ -297,7 +297,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { swiglu_limit_}); ORT_RETURN_IF_ERROR(context.RunProgram(swiglu)); } else { - ORT_THROW("only swiglu is supported for now."); + ORT_THROW("only swiglu is supported for WebGPU."); } // From 114d22809a29199aa111fecec6cc797fde285e05 Mon Sep 17 00:00:00 2001 From: Guenther Schmuelling Date: Tue, 11 Nov 2025 18:48:22 -0800 Subject: [PATCH 15/15] fix build --- onnxruntime/contrib_ops/webgpu/moe/qmoe.cc | 2 +- onnxruntime/core/providers/webgpu/program.cc | 10 ++++++++++ onnxruntime/core/providers/webgpu/program.h | 3 +++ 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index d37e25c6b27c6..c67cf8e37be69 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -192,7 +192,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { // we are accumulating expert results into output_tensor, need to initialize to zero ZeroTensorProgram zero; zero - .AddOutput({output_tensor, ProgramTensorMetadataDependency::Type, ProgramInput::Flatten, 4}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::Type, ProgramOutput::Flatten, 4}) .SetDispatchGroupSize((total_output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .AddUniformVariables({static_cast(total_output_size)}); ORT_RETURN_IF_ERROR(context.RunProgram(zero)); diff --git a/onnxruntime/core/providers/webgpu/program.cc b/onnxruntime/core/providers/webgpu/program.cc index 2c1b70222a5f6..9c0f1e85b3021 100644 --- a/onnxruntime/core/providers/webgpu/program.cc +++ b/onnxruntime/core/providers/webgpu/program.cc @@ -305,6 +305,16 @@ ProgramOutput::ProgramOutput(Tensor* tensor, ProgramTensorMetadataDependency dep use_override_shape{false}, override_shape{} {} +ProgramOutput::ProgramOutput(Tensor* tensor, ProgramTensorMetadataDependency dependency, ProgramOutput::FlattenTag, int component) + : tensor{tensor}, + dependency{dependency}, + var_type{ToProgramVariableDataType(tensor->GetElementType(), component)}, + is_atomic{false}, + use_override_shape{true}, + override_shape{} { + override_shape = {(tensor->Shape().Size() + component - 1) / component}; +} + ProgramOutput::ProgramOutput(Tensor* tensor, ProgramTensorMetadataDependency dependency, const TensorShape& override_shape, int component) : tensor{tensor}, dependency{dependency}, diff --git a/onnxruntime/core/providers/webgpu/program.h b/onnxruntime/core/providers/webgpu/program.h index 80f6d831d0909..d23211bdff674 100644 --- a/onnxruntime/core/providers/webgpu/program.h +++ b/onnxruntime/core/providers/webgpu/program.h @@ -235,14 +235,17 @@ struct ProgramInput { struct ProgramOutput { private: struct AtomicTag {}; + struct FlattenTag {}; public: constexpr static const AtomicTag Atomic{}; + constexpr static const FlattenTag Flatten{}; ProgramOutput(Tensor* tensor); ProgramOutput(Tensor* tensor, ProgramTensorMetadataDependency dependency, int component = 1); ProgramOutput(Tensor* tensor, ProgramTensorMetadataDependency dependency, AtomicTag); ProgramOutput(Tensor* tensor, ProgramTensorMetadataDependency dependency, const TensorShape& override_shape, int component); + ProgramOutput(Tensor* tensor, ProgramTensorMetadataDependency dependency, FlattenTag, int component = 1); Tensor* tensor; ProgramTensorMetadataDependency dependency;