Fix MoE security issue when the attribute k is larger than the number of experts

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<int64_t>("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<std::string>("activation_type", &activation_type_str).IsOK());

onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc

@@ -76,6 +76,10 @@ Status MoE<T>::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;

onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc

@@ -954,6 +954,10 @@ Status QMoECPU<T>::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);

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<float> input = {1.0f, 2.0f, 3.0f, 4.0f};
+  const std::vector<float> router_probs = {0.6f, 0.4f};
+  const std::vector<float> fc1_experts_weights(num_experts * hidden_size * (2 * inter_size), 0.1f);
+  const std::vector<float> fc2_experts_weights(num_experts * inter_size * hidden_size, 0.1f);
+  const std::vector<float> dummy_output(num_rows * hidden_size, 0.0f);
+
+  OpTester tester("MoE", 1, onnxruntime::kMSDomain);
+  tester.AddAttribute<int64_t>("k", k);
+  tester.AddAttribute<std::string>("activation_type", "swiglu");
+  tester.AddAttribute<int64_t>("normalize_routing_weights", int64_t{1});
+  tester.AddAttribute<int64_t>("swiglu_fusion", int64_t{1});
+  tester.AddAttribute<float>("activation_beta", 1.0f);
+
+  std::vector<int64_t> input_dims = {num_rows, hidden_size};
+  std::vector<int64_t> router_probs_dims = {num_rows, num_experts};
+  std::vector<int64_t> fc1_experts_weights_dims = {num_experts, hidden_size, 2 * inter_size};
+  std::vector<int64_t> fc2_experts_weights_dims = {num_experts, inter_size, hidden_size};
+  std::vector<int64_t> output_dims = {num_rows, hidden_size};
+
+  tester.AddInput<float>("input", input_dims, input);
+  tester.AddInput<float>("router_probs", router_probs_dims, router_probs);
+  tester.AddInput<float>("fc1_experts_weights", fc1_experts_weights_dims, fc1_experts_weights);
+  tester.AddOptionalInputEdge<float>();  // fc1_experts_bias
+  tester.AddInput<float>("fc2_experts_weights", fc2_experts_weights_dims, fc2_experts_weights);
+  tester.AddOptionalInputEdge<float>();  // fc2_experts_bias
+  tester.AddOptionalInputEdge<float>();  // fc3_experts_weights
+  tester.AddOptionalInputEdge<float>();  // fc3_experts_bias
+  tester.AddOutput<float>("output", output_dims, dummy_output);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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<float> input = {0.1f, -0.2f, 0.3f, -0.4f, 0.5f, -0.6f, 0.7f, -0.8f};
+  const std::vector<float> router_probs = {0.5f, 0.5f};
+  std::vector<uint8_t> fc1_experts_weights(num_experts * fc1_inter_size * (hidden_size / pack_size), 0x01);
+  std::vector<uint8_t> fc2_experts_weights(num_experts * hidden_size * (inter_size / pack_size), 0x10);
+  std::vector<float> fc1_scales(num_experts * fc1_inter_size, 0.1f);
+  std::vector<float> fc2_scales(num_experts * hidden_size, 0.05f);
+  std::vector<float> dummy_output(num_rows * hidden_size, 0.0f);
+
+  OpTester tester("QMoE", 1, onnxruntime::kMSDomain);
+  tester.AddAttribute<int64_t>("k", k);
+  tester.AddAttribute<std::string>("activation_type", "swiglu");
+  tester.AddAttribute<int64_t>("normalize_routing_weights", int64_t{1});
+  tester.AddAttribute<int64_t>("expert_weight_bits", expert_weight_bits);
+
+  std::vector<int64_t> input_dims = {num_rows, hidden_size};
+  std::vector<int64_t> router_probs_dims = {num_rows, num_experts};
+  std::vector<int64_t> fc1_experts_weights_dims = {num_experts, fc1_inter_size, hidden_size / pack_size};
+  std::vector<int64_t> fc2_experts_weights_dims = {num_experts, hidden_size, inter_size / pack_size};
+  std::vector<int64_t> fc1_scales_dims = {num_experts, fc1_inter_size};
+  std::vector<int64_t> fc2_scales_dims = {num_experts, hidden_size};
+  std::vector<int64_t> output_dims = {num_rows, hidden_size};
+
+  tester.AddInput<MLFloat16>("input", input_dims, ToFloat16(input));
+  tester.AddInput<MLFloat16>("router_probs", router_probs_dims, ToFloat16(router_probs));
+  tester.AddInput<uint8_t>("fc1_experts_weights", fc1_experts_weights_dims, fc1_experts_weights);
+  tester.AddInput<float>("fc1_scales", fc1_scales_dims, fc1_scales);
+  tester.AddOptionalInputEdge<MLFloat16>();  // fc1_experts_bias
+  tester.AddInput<uint8_t>("fc2_experts_weights", fc2_experts_weights_dims, fc2_experts_weights);
+  tester.AddInput<float>("fc2_scales", fc2_scales_dims, fc2_scales);
+  tester.AddOptionalInputEdge<MLFloat16>();  // fc2_experts_bias
+  tester.AddOptionalInputEdge<uint8_t>();    // fc3_experts_weights
+  tester.AddOptionalInputEdge<float>();      // fc3_scales
+  tester.AddOptionalInputEdge<MLFloat16>();  // fc3_experts_bias
+  tester.AddOutput<MLFloat16>("output", output_dims, ToFloat16(dummy_output));
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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