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 new file mode 100644 index 0000000000000..80887b845f915 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/final_mix.wgsl.template @@ -0,0 +1,24 @@ +// 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; + // 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; + 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..1214777009a8d --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/gate.wgsl.template @@ -0,0 +1,85 @@ +// 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 output_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[(row + uniforms.token_offset) * cols + local_idx]; + max_idx = local_idx; + } + 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 + 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; + } + if (local_idx == 0u) { + // softmax + var sum : f32 = 0.0; + for (var i = 0u; i < K; i++) { + sum += exp(f32(shared_vals[i])); + } + 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..d64d949b9d93d --- /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 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; + 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..a753b2d4c70a6 --- /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& /*unused*/) const { + return Status::OK(); +} + +Status MoE::ComputeInternal(ComputeContext& /*unused*/) 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..5e329dc12b5c9 --- /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..bab99fe51c88b --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/moe_base.h @@ -0,0 +1,33 @@ +// 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, +}; + +} // 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..c67cf8e37be69 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -0,0 +1,355 @@ +// 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}, + {"token_offset", 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); + // 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, 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_)); + + const auto& input_shape = hidden_state->Shape(); + + // SwiGLU validation + 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 on WebGPU."); + } + + // 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; + 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; + const int64_t N_fc1 = fc1_output_size; + const int64_t K_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; + Status status; + + Tensor* output_tensor = context.Output(0, input_shape); + 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::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)); + + // 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 + // token_idx is the global index into hidden_state + 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, 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_fc1, accuracy_level, expert_weight_bits_, context, + &fc1_outputs, expert_idx); + ORT_RETURN_IF_ERROR(status); + + // + // Step 4: apply swiglu + // + 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 WebGPU."); + } + + // + // 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_fc2, accuracy_level, expert_weight_bits_, context, + &fc2_outputs, expert_idx); + ORT_RETURN_IF_ERROR(status); + + // + // Step 6: 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, + static_cast(token_offset)}); + + 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..355f31228255f --- /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 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) +// 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..c45a49246cde0 --- /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(); 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;