From af9c7767c38a2c8d5e6162301fc2a347e5dbc02c Mon Sep 17 00:00:00 2001 From: Nenad Banfic Date: Mon, 27 Jul 2026 19:58:17 +0000 Subject: [PATCH] Fix MoE out of bounds for expert selection --- .../contrib_ops/cpu/moe/moe_base_cpu.h | 4 + onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc | 4 + .../cpu/moe/moe_quantization_cpu.cc | 4 + onnxruntime/test/contrib_ops/moe_test.cc | 114 ++++++++++++++++++ 4 files changed, 126 insertions(+) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_base_cpu.h b/onnxruntime/contrib_ops/cpu/moe/moe_base_cpu.h index d7522e46811e5..b0d5533c54b1a 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_base_cpu.h +++ b/onnxruntime/contrib_ops/cpu/moe/moe_base_cpu.h @@ -25,6 +25,10 @@ class MoEBaseCPU { protected: MoEBaseCPU(const OpKernelInfo& op_kernel_info) { ORT_ENFORCE(op_kernel_info.GetAttr("k", &k_).IsOK()); + // Defense-in-depth: k must be at least 1. The upper bound (k <= num_experts) + // cannot be checked here because num_experts is derived from a runtime input + // shape; it is enforced at Compute() entry in each derived operator. + ORT_ENFORCE(k_ >= 1, "MoE attribute 'k' must be >= 1, got k=", k_); std::string activation_type_str; ORT_ENFORCE(op_kernel_info.GetAttr("activation_type", &activation_type_str).IsOK()); diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc index 5db2d2e6256d2..7ed5e8fcc78a9 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc @@ -76,6 +76,10 @@ Status MoE::ComputeMoE(const OpKernelContext* context, const int64_t num_tokens = input_shape.Size() / input_shape[input_shape.NumDimensions() - 1]; const int64_t hidden_size = input_shape[input_shape.NumDimensions() - 1]; const int64_t num_experts = router_shape[1]; + + ORT_RETURN_IF_NOT(k_ <= num_experts, + "MoE attribute 'k' must be <= num_experts; got k=", k_, + ", num_experts=", num_experts); const int64_t inter_size = (fc2_shape[1] * fc2_shape[2]) / hidden_size; const bool is_swiglu = activation_type_ == ActivationType::SwiGLU; const int64_t fc1_output_size = is_swiglu ? (inter_size * 2) : inter_size; diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc index 06cccf5a283b7..fad6c5aecf512 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc @@ -954,6 +954,10 @@ Status QMoECPU::ComputeCommon(OpKernelContext* context, const ComputeInputs& const int64_t hidden_size = moe_params.hidden_size; const int64_t inter_size = moe_params.inter_size; const int64_t num_experts = moe_params.num_experts; + + ORT_RETURN_IF_NOT(k_ <= num_experts, + "QMoE attribute 'k' must be <= num_experts; got k=", k_, + ", num_experts=", num_experts); const int64_t fc1_out_features = inter_size * (swiglu_fusion_ > 0 ? 2 : 1); auto* output = context->Output(0, input_shape); diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 38cc3014aa873..1cdd22f789816 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -2584,6 +2584,120 @@ TEST(MoETest, MoECpuTest_BasicSwiGLU) { fc3_experts_weights, fc1_experts_bias, fc2_experts_bias, output_data, num_rows, num_experts, hidden_size, inter_size, "swiglu"); } + +// Regression test: MoE must reject k > num_experts. Without the guard the per-token +// sorted_logits vector (size = num_experts) is passed to std::partial_sort with +// mid = begin + k_, writing past the vector's heap allocation. +TEST(MoETest, MoECpuTest_KExceedsNumExperts) { + auto cpu_ep = DefaultCpuExecutionProvider(); + if (!cpu_ep) { + GTEST_SKIP() << "CPU execution provider not available"; + } + + constexpr int num_rows = 1; + constexpr int num_experts = 2; + constexpr int hidden_size = 4; + constexpr int inter_size = 8; + constexpr int64_t k = 3; // > num_experts: the bug trigger. + + const std::vector input = {1.0f, 2.0f, 3.0f, 4.0f}; + const std::vector router_probs = {0.6f, 0.4f}; + const std::vector fc1_experts_weights(num_experts * hidden_size * (2 * inter_size), 0.1f); + const std::vector fc2_experts_weights(num_experts * inter_size * hidden_size, 0.1f); + const std::vector dummy_output(num_rows * hidden_size, 0.0f); + + OpTester tester("MoE", 1, onnxruntime::kMSDomain); + tester.AddAttribute("k", k); + tester.AddAttribute("activation_type", "swiglu"); + tester.AddAttribute("normalize_routing_weights", int64_t{1}); + tester.AddAttribute("swiglu_fusion", int64_t{1}); + tester.AddAttribute("activation_beta", 1.0f); + + std::vector input_dims = {num_rows, hidden_size}; + std::vector router_probs_dims = {num_rows, num_experts}; + std::vector fc1_experts_weights_dims = {num_experts, hidden_size, 2 * inter_size}; + std::vector fc2_experts_weights_dims = {num_experts, inter_size, hidden_size}; + std::vector output_dims = {num_rows, hidden_size}; + + tester.AddInput("input", input_dims, input); + tester.AddInput("router_probs", router_probs_dims, router_probs); + tester.AddInput("fc1_experts_weights", fc1_experts_weights_dims, fc1_experts_weights); + tester.AddOptionalInputEdge(); // fc1_experts_bias + tester.AddInput("fc2_experts_weights", fc2_experts_weights_dims, fc2_experts_weights); + tester.AddOptionalInputEdge(); // fc2_experts_bias + tester.AddOptionalInputEdge(); // fc3_experts_weights + tester.AddOptionalInputEdge(); // fc3_experts_bias + tester.AddOutput("output", output_dims, dummy_output); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, + "'k' must be <= num_experts", + {}, nullptr, &execution_providers); +} + +// Regression test: QMoE must reject k > num_experts. Same partial_sort OOB pattern +// as MoE above, but on the quantized CPU kernel path. +TEST(MoETest, QMoETest_CPU_KExceedsNumExperts) { +#ifdef USE_MLAS + auto cpu_ep = DefaultCpuExecutionProvider(); + if (!cpu_ep) { + GTEST_SKIP() << "CPU execution provider not available"; + } + + constexpr int64_t num_rows = 1; + constexpr int64_t num_experts = 2; + constexpr int64_t hidden_size = 8; + constexpr int64_t inter_size = 8; + constexpr int64_t expert_weight_bits = 4; + constexpr int64_t pack_size = 8 / expert_weight_bits; + constexpr int64_t fc1_inter_size = 2 * inter_size; // swiglu fused + constexpr int64_t k = 3; // > num_experts: the bug trigger. + + const std::vector input = {0.1f, -0.2f, 0.3f, -0.4f, 0.5f, -0.6f, 0.7f, -0.8f}; + const std::vector router_probs = {0.5f, 0.5f}; + std::vector fc1_experts_weights(num_experts * fc1_inter_size * (hidden_size / pack_size), 0x01); + std::vector fc2_experts_weights(num_experts * hidden_size * (inter_size / pack_size), 0x10); + std::vector fc1_scales(num_experts * fc1_inter_size, 0.1f); + std::vector fc2_scales(num_experts * hidden_size, 0.05f); + std::vector dummy_output(num_rows * hidden_size, 0.0f); + + OpTester tester("QMoE", 1, onnxruntime::kMSDomain); + tester.AddAttribute("k", k); + tester.AddAttribute("activation_type", "swiglu"); + tester.AddAttribute("normalize_routing_weights", int64_t{1}); + tester.AddAttribute("expert_weight_bits", expert_weight_bits); + + std::vector input_dims = {num_rows, hidden_size}; + std::vector router_probs_dims = {num_rows, num_experts}; + std::vector fc1_experts_weights_dims = {num_experts, fc1_inter_size, hidden_size / pack_size}; + std::vector fc2_experts_weights_dims = {num_experts, hidden_size, inter_size / pack_size}; + std::vector fc1_scales_dims = {num_experts, fc1_inter_size}; + std::vector fc2_scales_dims = {num_experts, hidden_size}; + std::vector output_dims = {num_rows, hidden_size}; + + tester.AddInput("input", input_dims, ToFloat16(input)); + tester.AddInput("router_probs", router_probs_dims, ToFloat16(router_probs)); + tester.AddInput("fc1_experts_weights", fc1_experts_weights_dims, fc1_experts_weights); + tester.AddInput("fc1_scales", fc1_scales_dims, fc1_scales); + tester.AddOptionalInputEdge(); // fc1_experts_bias + tester.AddInput("fc2_experts_weights", fc2_experts_weights_dims, fc2_experts_weights); + tester.AddInput("fc2_scales", fc2_scales_dims, fc2_scales); + tester.AddOptionalInputEdge(); // fc2_experts_bias + tester.AddOptionalInputEdge(); // fc3_experts_weights + tester.AddOptionalInputEdge(); // fc3_scales + tester.AddOptionalInputEdge(); // fc3_experts_bias + tester.AddOutput("output", output_dims, ToFloat16(dummy_output)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, + "'k' must be <= num_experts", + {}, nullptr, &execution_providers); +#else + GTEST_SKIP() << "Skipping CPU QMoE test"; +#endif +} #endif } // namespace test