diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc index f09c6b4a593f3..19bf6ba45e133 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.cc @@ -290,12 +290,20 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { (fc1_zero_points == nullptr && fc2_zero_points == nullptr && fc3_zero_points == nullptr), "WebGPU QMoE does not support explicit zero points on the optimized single-token path."); - if (fc1_expert_weight_bits_ != expert_weight_bits_ || - fc2_expert_weight_bits_ != expert_weight_bits_ || - fc3_expert_weight_bits_ != expert_weight_bits_) { - return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "Mixed-width QMoE execution is not yet implemented on WebGPU."); - } + // MatMulNBits addresses packed rows and blocks as u32 words; partial words would drop weights. + const auto has_word_aligned_rows = [&](int64_t k, int64_t n, int64_t bits) { + const int64_t packed_block = block_size_ != 0 && block_size_ != k * n ? block_size_ : k; + return (k * bits) % 32 == 0 && (packed_block * bits) % 32 == 0; + }; + ORT_RETURN_IF_NOT(has_word_aligned_rows(moe_params.hidden_size, + is_fused_swiglu ? 2 * moe_params.inter_size : moe_params.inter_size, + fc1_expert_weight_bits_), + "WebGPU QMoE FC1 packed rows and blocks must be 32-bit aligned."); + ORT_RETURN_IF_NOT(has_word_aligned_rows(moe_params.inter_size, moe_params.hidden_size, fc2_expert_weight_bits_), + "WebGPU QMoE FC2 packed rows and blocks must be 32-bit aligned."); + ORT_RETURN_IF_NOT(fc3_experts_weights_optional == nullptr || + has_word_aligned_rows(moe_params.hidden_size, moe_params.inter_size, fc3_expert_weight_bits_), + "WebGPU QMoE FC3 packed rows and blocks must be 32-bit aligned."); const auto& input_shape = hidden_state->Shape(); @@ -363,7 +371,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { TensorShape fc1_output_shape({static_cast(k), fc1_output_size}); Tensor fc1_outputs = context.CreateGPUTensor(dtype, fc1_output_shape); status = ApplyMatMulNBits(hidden_state, fc1_experts_weights, fc1_scales, fc1_zero_points, fc1_experts_bias_optional, - K_fc1, N_fc1, block_size_fc1, accuracy_level, expert_weight_bits_, context, + K_fc1, N_fc1, block_size_fc1, accuracy_level, fc1_expert_weight_bits_, context, &fc1_outputs, 0, &indirect_experts, /*override_M=*/k); ORT_RETURN_IF_ERROR(status); @@ -376,7 +384,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { ORT_RETURN_IF_ERROR(ApplyMatMulNBits(hidden_state, fc3_experts_weights_optional, fc3_scales_optional, fc3_zero_points, fc3_experts_bias_optional, K_fc1, moe_params.inter_size, block_size_fc3, accuracy_level, - expert_weight_bits_, context, &*fc3_outputs, 0, &indirect_experts, + fc3_expert_weight_bits_, context, &*fc3_outputs, 0, &indirect_experts, /*override_M=*/k)); } MoEActivationProgram activation{activation_type_, swiglu_fusion, fc3_outputs.has_value()}; @@ -398,7 +406,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { TensorShape fc2_output_shape({static_cast(k), N_fc2}); Tensor fc2_outputs = context.CreateGPUTensor(dtype, fc2_output_shape); status = ApplyMatMulNBits(&fc1_activated, fc2_experts_weights, fc2_scales, fc2_zero_points, fc2_experts_bias_optional, - K_fc2, N_fc2, block_size_fc2, accuracy_level, expert_weight_bits_, context, + K_fc2, N_fc2, block_size_fc2, accuracy_level, fc2_expert_weight_bits_, context, &fc2_outputs, 0, &indirect_experts, /*override_M=*/0); ORT_RETURN_IF_ERROR(status); @@ -515,7 +523,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { // Step 3: matmul the hidden_state with fc1 (gate_up) of the selected experts // status = ApplyMatMulNBits(&expert_hidden, fc1_experts_weights, fc1_scales, fc1_zero_points, fc1_experts_bias_optional, - K_fc1, N_fc1, block_size_fc1, accuracy_level, expert_weight_bits_, context, + K_fc1, N_fc1, block_size_fc1, accuracy_level, fc1_expert_weight_bits_, context, &fc1_outputs, expert_idx); ORT_RETURN_IF_ERROR(status); @@ -528,7 +536,7 @@ Status QMoE::ComputeInternal(ComputeContext& context) const { ORT_RETURN_IF_ERROR(ApplyMatMulNBits(&expert_hidden, fc3_experts_weights_optional, fc3_scales_optional, fc3_zero_points, fc3_experts_bias_optional, K_fc1, moe_params.inter_size, block_size_fc3, accuracy_level, - expert_weight_bits_, context, &*fc3_outputs, expert_idx)); + fc3_expert_weight_bits_, context, &*fc3_outputs, expert_idx)); } MoEActivationProgram activation{activation_type_, swiglu_fusion, fc3_outputs.has_value()}; activation.AddInputs({{&fc1_outputs, ProgramTensorMetadataDependency::Type}}); @@ -548,7 +556,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, fc2_zero_points, fc2_experts_bias_optional, - K_fc2, N_fc2, block_size_fc2, accuracy_level, expert_weight_bits_, context, + K_fc2, N_fc2, block_size_fc2, accuracy_level, fc2_expert_weight_bits_, context, &fc2_outputs, expert_idx); ORT_RETURN_IF_ERROR(status); diff --git a/onnxruntime/contrib_ops/webgpu/moe/qmoe.h b/onnxruntime/contrib_ops/webgpu/moe/qmoe.h index 8981b39b98b0e..d3e4de36adeba 100755 --- a/onnxruntime/contrib_ops/webgpu/moe/qmoe.h +++ b/onnxruntime/contrib_ops/webgpu/moe/qmoe.h @@ -20,8 +20,8 @@ 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_); + ORT_ENFORCE(expert_weight_bits_ == 2 || expert_weight_bits_ == 4 || expert_weight_bits_ == 8, + "expert_weight_bits must be 2, 4, or 8, but got ", expert_weight_bits_); fc1_expert_weight_bits_ = info.GetAttrOrDefault("fc1_expert_weight_bits", expert_weight_bits_); fc2_expert_weight_bits_ = info.GetAttrOrDefault("fc2_expert_weight_bits", expert_weight_bits_); fc3_expert_weight_bits_ = info.GetAttrOrDefault("fc3_expert_weight_bits", expert_weight_bits_); diff --git a/onnxruntime/test/providers/webgpu/contrib_ops/moe_test.cc b/onnxruntime/test/providers/webgpu/contrib_ops/moe_test.cc index 7119309e3e6ce..7960e1bf6c84e 100644 --- a/onnxruntime/test/providers/webgpu/contrib_ops/moe_test.cc +++ b/onnxruntime/test/providers/webgpu/contrib_ops/moe_test.cc @@ -697,53 +697,160 @@ TEST(MoETest, QMoETest_WebGPU_ActivationsAndSwiGLUFusion) { run_case("swiglu", 2); } -TEST(MoETest, QMoETest_MixedWidthContract_WebGPU) { +static void RunQMoEWebGpuIntegerWidths(int num_rows, int num_experts, int k, + int base_bits, int fc1_bits, int fc2_bits, int fc3_bits = 0) { GET_WEBGPU_EP_OR_SKIP(webgpu_ep); - constexpr int64_t num_rows = 1; - constexpr int64_t num_experts = 1; - constexpr int64_t hidden_size = 8; - constexpr int64_t inter_size = 8; - constexpr int64_t fc1_bits = 2; - constexpr int64_t fc2_bits = 4; - constexpr int64_t fc1_pack_size = 8 / fc1_bits; - constexpr int64_t fc2_pack_size = 8 / fc2_bits; - - const std::vector input_dims = {num_rows, hidden_size}; - const std::vector router_probs_dims = {num_rows, num_experts}; - const std::vector fc1_weights_dims = {num_experts, inter_size, hidden_size / fc1_pack_size}; - const std::vector fc2_weights_dims = {num_experts, hidden_size, inter_size / fc2_pack_size}; - const std::vector fc1_scales_dims = {num_experts, inter_size}; - const std::vector fc2_scales_dims = {num_experts, hidden_size}; + constexpr int size = 64; + const bool has_fc3 = fc3_bits != 0; + const std::vector fc1_weights(num_experts * size * size * fc1_bits / 8, + fc1_bits == 2 ? 0xFF : 0x99); + const std::vector fc2_weights(num_experts * size * size * fc2_bits / 8, + fc2_bits == 2 ? 0xFF : 0x99); + const std::vector fc3_weights(has_fc3 ? num_experts * size * size * fc3_bits / 8 : 0, + fc3_bits == 2 ? 0xFF : 0x99); + std::vector input(num_rows * size), router_probs(num_rows * num_experts, 0.0f); + std::vector fc1_scales(num_experts * size), fc2_scales(num_experts * size); + std::vector fc3_scales(has_fc3 ? num_experts * size : 0); + std::vector expected(num_rows * size); + for (int expert = 0; expert < num_experts; ++expert) { + std::fill_n(fc1_scales.begin() + expert * size, size, expert == 0 ? 0.5f : 1.0f); + std::fill_n(fc2_scales.begin() + expert * size, size, expert == 0 ? 0.25f : 0.5f); + if (has_fc3) { + std::fill_n(fc3_scales.begin() + expert * size, size, expert == 0 ? 0.125f : 0.25f); + } + } + for (int row = 0; row < num_rows; ++row) { + const float value = 0.25f * (row + 1); + std::fill_n(input.begin() + row * size, size, value); + float result = 0.0f; + for (int expert = 0; expert < num_experts; ++expert) { + if (k == 1 && expert != row % num_experts) { + continue; + } + router_probs[row * num_experts + expert] = 10.0f; + const float fc1 = MLFloat16(size * value * fc1_scales[expert * size]).ToFloat(); + float activated = fc1; + if (has_fc3) { + const float fc3 = MLFloat16(size * value * fc3_scales[expert * size]).ToFloat(); + activated = MLFloat16(fc1 / (1.0f + std::exp(-fc1)) * fc3).ToFloat(); + } + result += MLFloat16(size * activated * fc2_scales[expert * size]).ToFloat() / k; + } + std::fill_n(expected.begin() + row * size, size, result); + } OpTester tester("QMoE", 1, onnxruntime::kMSDomain); - tester.AddAttribute("k", 1); - tester.AddAttribute("activation_type", "identity"); - tester.AddAttribute("expert_weight_bits", fc2_bits); - tester.AddAttribute("fc1_expert_weight_bits", fc1_bits); - tester.AddInput("input", input_dims, std::vector(num_rows * hidden_size)); - tester.AddInput("router_probs", router_probs_dims, - std::vector(num_rows * num_experts)); - tester.AddInput("fc1_experts_weights", fc1_weights_dims, - std::vector(static_cast(fc1_weights_dims[1] * fc1_weights_dims[2]))); - tester.AddInput("fc1_scales", fc1_scales_dims, - std::vector(num_experts * inter_size, 1.0f)); + tester.AddAttribute("k", k); + tester.AddAttribute("activation_type", has_fc3 ? "swiglu" : "identity"); + tester.AddAttribute("normalize_routing_weights", 1); + tester.AddAttribute("expert_weight_bits", base_bits); + if (fc1_bits != base_bits) tester.AddAttribute("fc1_expert_weight_bits", fc1_bits); + if (fc2_bits != base_bits) tester.AddAttribute("fc2_expert_weight_bits", fc2_bits); + if (has_fc3 && fc3_bits != base_bits) tester.AddAttribute("fc3_expert_weight_bits", fc3_bits); + tester.AddInput("input", {num_rows, size}, ToFloat16(input)); + tester.AddInput("router_probs", {num_rows, num_experts}, ToFloat16(router_probs)); + tester.AddInput("fc1_experts_weights", {num_experts, size, size * fc1_bits / 8}, fc1_weights); + tester.AddInput("fc1_scales", {num_experts, size}, ToFloat16(fc1_scales)); tester.AddOptionalInputEdge(); - tester.AddInput("fc2_experts_weights", fc2_weights_dims, - std::vector(static_cast(fc2_weights_dims[1] * fc2_weights_dims[2]))); - tester.AddInput("fc2_scales", fc2_scales_dims, - std::vector(num_experts * hidden_size, 1.0f)); + tester.AddInput("fc2_experts_weights", {num_experts, size, size * fc2_bits / 8}, fc2_weights); + tester.AddInput("fc2_scales", {num_experts, size}, ToFloat16(fc2_scales)); tester.AddOptionalInputEdge(); - tester.AddOptionalInputEdge(); - tester.AddOptionalInputEdge(); + if (has_fc3) { + tester.AddInput("fc3_experts_weights", {num_experts, size, size * fc3_bits / 8}, fc3_weights); + tester.AddInput("fc3_scales", {num_experts, size}, ToFloat16(fc3_scales)); + } else { + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + } tester.AddOptionalInputEdge(); - tester.AddOutput("output", input_dims, std::vector(num_rows * hidden_size)); + tester.AddOutput("output", {num_rows, size}, ToFloat16(expected)); + tester.SetOutputTolerance(0.02f); + + RunWebGpuOnly(tester, std::move(webgpu_ep)); +} + +TEST(MoETest, QMoETest_WebGPU_MixedWidth_SingleToken) { + RunQMoEWebGpuIntegerWidths(1, 1, 1, 4, 2, 4); +} + +TEST(MoETest, QMoETest_WebGPU_Int2InheritedWidths) { + RunQMoEWebGpuIntegerWidths(1, 1, 1, 2, 2, 2); +} - std::vector> execution_providers; - execution_providers.push_back(std::move(webgpu_ep)); - tester.Run(OpTester::ExpectResult::kExpectFailure, - "Mixed-width QMoE execution is not yet implemented on WebGPU.", - {}, nullptr, &execution_providers); +TEST(MoETest, QMoETest_WebGPU_MixedWidth_TwoExperts) { + RunQMoEWebGpuIntegerWidths(1, 2, 2, 4, 2, 4); +} + +TEST(MoETest, QMoETest_WebGPU_MixedWidth_MultiToken) { + RunQMoEWebGpuIntegerWidths(2, 2, 1, 4, 2, 4); +} + +TEST(MoETest, QMoETest_WebGPU_MixedWidth_SeparateFc3) { + RunQMoEWebGpuIntegerWidths(1, 2, 2, 4, 4, 2, 2); +} + +TEST(MoETest, QMoETest_WebGPU_Int2RejectUnalignedRows) { + const auto run_case = [](int hidden_size, int inter_size) { + GET_WEBGPU_EP_OR_SKIP(webgpu_ep); + OpTester tester("QMoE", 1, onnxruntime::kMSDomain); + tester.AddAttribute("k", 1); + tester.AddAttribute("activation_type", "identity"); + tester.AddAttribute("expert_weight_bits", 2); + tester.AddInput("input", {1, hidden_size}, + ToFloat16(std::vector(hidden_size, 0.25f))); + tester.AddInput("router_probs", {1, 1}, ToFloat16({1.0f})); + tester.AddInput("fc1_experts_weights", {1, inter_size, hidden_size / 4}, + std::vector(inter_size * hidden_size / 4, 0xFF)); + tester.AddInput("fc1_scales", {1, inter_size}, + ToFloat16(std::vector(inter_size, 0.5f))); + tester.AddOptionalInputEdge(); + tester.AddInput("fc2_experts_weights", {1, hidden_size, inter_size / 4}, + std::vector(hidden_size * inter_size / 4, 0xFF)); + tester.AddInput("fc2_scales", {1, hidden_size}, + ToFloat16(std::vector(hidden_size, 0.25f))); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + tester.AddOutput("output", {1, hidden_size}, + ToFloat16(std::vector(hidden_size, 0.0f))); + RunWebGpuOnly(tester, std::move(webgpu_ep), OpTester::ExpectResult::kExpectFailure, "32-bit aligned"); + }; + run_case(12, 64); + run_case(64, 12); +} + +TEST(MoETest, QMoETest_WebGPU_Int2RejectUnalignedFc3) { + GET_WEBGPU_EP_OR_SKIP(webgpu_ep); + constexpr int hidden_size = 8; + constexpr int inter_size = 64; + OpTester tester("QMoE", 1, onnxruntime::kMSDomain); + tester.AddAttribute("k", 1); + tester.AddAttribute("activation_type", "swiglu"); + tester.AddAttribute("expert_weight_bits", 4); + tester.AddAttribute("fc3_expert_weight_bits", 2); + tester.AddInput("input", {1, hidden_size}, + ToFloat16(std::vector(hidden_size, 0.25f))); + tester.AddInput("router_probs", {1, 1}, ToFloat16({1.0f})); + tester.AddInput("fc1_experts_weights", {1, inter_size, hidden_size / 2}, + std::vector(inter_size * hidden_size / 2, 0x99)); + tester.AddInput("fc1_scales", {1, inter_size}, + ToFloat16(std::vector(inter_size, 0.5f))); + tester.AddOptionalInputEdge(); + tester.AddInput("fc2_experts_weights", {1, hidden_size, inter_size / 2}, + std::vector(hidden_size * inter_size / 2, 0x99)); + tester.AddInput("fc2_scales", {1, hidden_size}, + ToFloat16(std::vector(hidden_size, 0.25f))); + tester.AddOptionalInputEdge(); + tester.AddInput("fc3_experts_weights", {1, inter_size, hidden_size / 4}, + std::vector(inter_size * hidden_size / 4, 0xFF)); + tester.AddInput("fc3_scales", {1, inter_size}, + ToFloat16(std::vector(inter_size, 0.125f))); + tester.AddOptionalInputEdge(); + tester.AddOutput("output", {1, hidden_size}, + ToFloat16(std::vector(hidden_size, 0.0f))); + RunWebGpuOnly(tester, std::move(webgpu_ep), OpTester::ExpectResult::kExpectFailure, "FC3 packed rows"); } #undef GET_WEBGPU_EP_OR_SKIP