webgpu qmoe

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";

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;
+    }
+}

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<workgroup> shared_vals: array<hidden_state_element_t, workgroup_size_x>;
+var<workgroup> shared_idxs: array<u32, workgroup_size_x>;
+
+$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

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

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

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<MoEProgram> {
+ public:
+  MoEProgram(TensorShape output_shape) : Program<MoEProgram>{"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<float>(info.GetAttrOrDefault<float>("activation_alpha", 1.0));
+    activation_beta_ = static_cast<float>(info.GetAttrOrDefault<float>("activation_beta", 1.0));
+    swiglu_fusion_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("swiglu_fusion", 0));
+    swiglu_limit_ = info.GetAttrOrDefault<float>("swiglu_limit", 0);
+    k_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("k", 4));
+    normalize_routing_weights_ = info.GetAttrOrDefault<int64_t>("normalize_routing_weights", 0) == 1;
+    use_sparse_mixer_ = info.GetAttrOrDefault<int64_t>("use_sparse_mixer", 0) == 1;
+    std::string activation_type = info.GetAttrOrDefault<std::string>("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

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