From 6dd6e705bab94afbbbf7269087a1c4ec40305597 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:40:31 +0000 Subject: [PATCH 1/3] Initial plan From 9e4add39940ad7f92e6bc40027ee60892d71644d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:57:08 +0000 Subject: [PATCH 2/3] Add CPU MoE intermediate-buffer reuse regression tests Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com> --- onnxruntime/test/contrib_ops/moe_test.cc | 103 +++++++++++++++++++++++ 1 file changed, 103 insertions(+) diff --git a/onnxruntime/test/contrib_ops/moe_test.cc b/onnxruntime/test/contrib_ops/moe_test.cc index 4cf6a6c4ccfba..04d93d136aa92 100644 --- a/onnxruntime/test/contrib_ops/moe_test.cc +++ b/onnxruntime/test/contrib_ops/moe_test.cc @@ -27,6 +27,7 @@ #include "test/common/tensor_op_test_utils.h" #include "test/common/cuda_op_test_utils.h" #include "test/providers/provider_test_utils.h" +#include "test/unittest_util/graph_transform_test_builder.h" #include "test/util/include/inference_session_wrapper.h" #include "test/util/include/test/test_environment.h" #ifdef USE_CUDA @@ -150,6 +151,108 @@ static void RunMoETest(const std::vector& input, const std::vector } } +template +static void RunMoEIntermediateInputOutputTest() { + const std::vector shape{3, 2}; + const std::vector tokens{1.0f, 2.0f, -1.0f, -2.0f, 0.5f, -1.5f}; + std::vector fc1(12, T(0.0f)); + std::vector fc2; + for (int expert = 0; expert < 3; ++expert) { + fc1[expert * 4] = T(static_cast(expert + 1)); + fc1[expert * 4 + 3] = T(1.0f); + fc2.push_back(T(static_cast(expert + 1))); + fc2.push_back(T(static_cast(-expert - 1))); + } + + for (const std::string input_source : {"input", "add", "reshape(add)"}) { + for (bool expose_moe : {false, true}) { + for (bool normalize : {false, true}) { + SCOPED_TRACE(MakeString("input_source=", input_source, ", expose_moe=", expose_moe, + ", normalize=", normalize)); + Model model("MoEIntermediateInputOutput", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kOnnxDomain, 20}, {kMSDomain, 1}}, {}, + DefaultLoggingManager().DefaultLogger()); + ModelTestBuilder builder(model.MainGraph()); + std::vector input; + for (float value : tokens) { + input.push_back(T(input_source == "input" ? value : value * 0.5f)); + } + auto* moe_input = builder.MakeInput(shape, input); + if (input_source != "input") { + auto* sum = builder.MakeIntermediate(); + builder.AddNode("Add", {moe_input, moe_input}, {sum}); + moe_input = sum; + if (input_source == "reshape(add)") { + moe_input = builder.MakeIntermediate(); + builder.AddNode("Reshape", {sum, builder.MakeInitializer({2}, {-1, 2})}, {moe_input}); + } + } + auto* logits = builder.MakeInput({3, 3}, {T(0.0f), T(1.0f), T(2.0f), + T(0.0f), T(1.0f), T(2.0f), + T(0.0f), T(1.0f), T(2.0f)}); + auto* y = builder.MakeOutput(); + auto* moe_output = expose_moe ? builder.MakeOutput() : builder.MakeIntermediate(); + auto& moe = builder.AddNode("MoE", + {moe_input, logits, builder.MakeInitializer({3, 2, 2}, fc1), + builder.MakeOptionalTensor(), builder.MakeInitializer({3, 2, 1}, fc2)}, + {moe_output}, kMSDomain); + moe.AddAttribute("k", int64_t{2}); + moe.AddAttribute("activation_type", std::string("swiglu")); + moe.AddAttribute("swiglu_fusion", int64_t{1}); + moe.AddAttribute("activation_alpha", 1.0f); + moe.AddAttribute("activation_beta", 0.0f); + moe.AddAttribute("normalize_routing_weights", static_cast(normalize)); + builder.AddNode("Relu", {moe_output}, {y}); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + const std::string model_data = model.ToProto().SerializeAsString(); + + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + options.intra_op_param.thread_pool_size = 1; + InferenceSession session(options, GetEnvironment()); + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultCpuExecutionProvider())); + ASSERT_STATUS_OK(session.Load(model_data.data(), static_cast(model_data.size()))); + ASSERT_STATUS_OK(session.Initialize()); + + // The two largest logits select experts 1 and 2, with interleaved gate/up FC1 rows. + std::vector expected(tokens.size(), 0.0f); + const float denominator = (normalize ? 0.0f : 1.0f) + std::exp(1.0f) + std::exp(2.0f); + for (size_t row = 0; row < 3; ++row) { + for (int expert = 1; expert < 3; ++expert) { + const float gate = static_cast(expert + 1) * tokens[2 * row]; + expected[2 * row] += std::exp(static_cast(expert)) / denominator * + static_cast(expert + 1) * gate / (1.0f + std::exp(-gate)) * tokens[2 * row + 1]; + } + expected[2 * row + 1] = -expected[2 * row]; + } + for (int run = 0; run < 2; ++run) { + std::vector fetches; + ASSERT_STATUS_OK(session.Run(RunOptions(), builder.feeds_, builder.output_names_, &fetches)); + ASSERT_EQ(fetches.size(), expose_moe ? 2U : 1U); + for (size_t output = 0; output < fetches.size(); ++output) { + const auto& tensor = fetches[output].Get(); + ASSERT_EQ(tensor.Shape(), TensorShape(shape)); + const T* actual = tensor.Data(); + for (size_t i = 0; i < expected.size(); ++i) { + const float value = output == 0 ? std::max(expected[i], 0.0f) : expected[i]; + EXPECT_NEAR(static_cast(actual[i]), value, (std::is_same_v ? 1e-5f : 0.02f)); + } + } + } + } + } + } +} + +TEST(MoETest, IntermediateInputOutputFloat32) { + RunMoEIntermediateInputOutputTest(); +} + +TEST(MoETest, IntermediateInputOutputFloat16) { + RunMoEIntermediateInputOutputTest(); +} + // TODO(wy): Add python parity tests that can serve as examples. Need cutlass upgrade to build cutlass extensions to // add weights preprocesser to onnxruntime_pybind_quant.cc template From ef0ddf86496b150c555a8e1d00df902d3168efa5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:18:03 +0000 Subject: [PATCH 3/3] Preserve CPU MoE input when output reuses its buffer Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com> --- onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc b/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc index 223f90f56e871..59e47667c46fd 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc +++ b/onnxruntime/contrib_ops/cpu/moe/moe_cpu.cc @@ -103,7 +103,8 @@ Status MoE::ComputeMoE(const OpKernelContext* context, const T* input_data_to_use = input_data; IAllocatorUniquePtr input_data_copy_ptr; - if (normalize_routing_weights_) { + // Preserve the tokens before clearing an output that reuses the input buffer. + if (input_data == output_data) { input_data_copy_ptr = IAllocator::MakeUniquePtr(allocator, static_cast(num_tokens * hidden_size)); T* input_data_copy = input_data_copy_ptr.get(); std::copy(input_data, input_data + (num_tokens * hidden_size), input_data_copy);