From 3912854e2b132957933743251a0cd8a59efb675f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 30 Jul 2026 19:04:09 +0000 Subject: [PATCH 1/2] fix sharded moe --- onnxruntime/contrib_ops/cpu/moe/moe_helper.h | 58 ++--- .../cuda/collective/sharded_moe.cc | 205 +++++------------- .../contrib_ops/cuda/collective/sharded_moe.h | 10 +- onnxruntime/contrib_ops/cuda/moe/moe.cc | 183 +++++++++++----- onnxruntime/contrib_ops/cuda/moe/moe_base.h | 21 ++ .../sharded_moe/test_sharded_moe.py | 5 +- 6 files changed, 246 insertions(+), 236 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/moe/moe_helper.h b/onnxruntime/contrib_ops/cpu/moe/moe_helper.h index d769ef94bca17..43dd9ef141e13 100644 --- a/onnxruntime/contrib_ops/cpu/moe/moe_helper.h +++ b/onnxruntime/contrib_ops/cpu/moe/moe_helper.h @@ -131,22 +131,26 @@ Status CheckInputs(MoEParameters& parameters, const int64_t fc1_inter_size = is_fused_swiglu ? (inter_size + inter_size) : inter_size; const int64_t zp_pack_size = pack_size; // Zero points packing (1 for 8-bit, 2 for 4-bit) + // Weights, and every tensor indexed by the weights' expert dimension, only cover the experts that + // are local to this rank. They are identical to num_experts unless expert parallelism is used + // (ShardedMoE), where each rank only owns a slice of the experts. if (legacy_shape) { // legacy shape does not match column major memory layout. This is for backward compatibility. - CHECK_SHAPE(fc1_experts_weights_shape, "fc1_experts_weights", num_experts, hidden_size, fc1_inter_size / pack_size); - CHECK_SHAPE(fc2_experts_weights_shape, "fc2_experts_weights", num_experts, inter_size, hidden_size / pack_size); - CHECK_SHAPE(fc3_experts_weights_shape, "fc3_experts_weights", num_experts, hidden_size, inter_size / pack_size); + CHECK_SHAPE(fc1_experts_weights_shape, "fc1_experts_weights", local_num_experts, hidden_size, fc1_inter_size / pack_size); + CHECK_SHAPE(fc2_experts_weights_shape, "fc2_experts_weights", local_num_experts, inter_size, hidden_size / pack_size); + CHECK_SHAPE(fc3_experts_weights_shape, "fc3_experts_weights", local_num_experts, hidden_size, inter_size / pack_size); } else { - CHECK_SHAPE(fc1_experts_weights_shape, "fc1_experts_weights", num_experts, fc1_inter_size, hidden_size / pack_size); - CHECK_SHAPE(fc2_experts_weights_shape, "fc2_experts_weights", num_experts, hidden_size, inter_size / pack_size); - CHECK_SHAPE(fc3_experts_weights_shape, "fc3_experts_weights", num_experts, inter_size, hidden_size / pack_size); + CHECK_SHAPE(fc1_experts_weights_shape, "fc1_experts_weights", local_num_experts, fc1_inter_size, hidden_size / pack_size); + CHECK_SHAPE(fc2_experts_weights_shape, "fc2_experts_weights", local_num_experts, hidden_size, inter_size / pack_size); + CHECK_SHAPE(fc3_experts_weights_shape, "fc3_experts_weights", local_num_experts, inter_size, hidden_size / pack_size); } CHECK_TENSOR_SHAPE(router_probs, num_rows, num_experts); - CHECK_TENSOR_SHAPE(fc1_experts_bias, num_experts, fc1_inter_size); + CHECK_TENSOR_SHAPE(fc1_experts_bias, local_num_experts, fc1_inter_size); + // fc2_experts_bias is applied after the results of all ranks are combined, so it covers all experts. CHECK_TENSOR_SHAPE(fc2_experts_bias, num_experts, hidden_size); - CHECK_TENSOR_SHAPE(fc3_experts_bias, num_experts, inter_size); + CHECK_TENSOR_SHAPE(fc3_experts_bias, local_num_experts, inter_size); // Validate scale tensors: Handle both row-wise and block-wise quantization flexibly // First, detect the actual quantization method from the tensor shapes @@ -166,29 +170,29 @@ Status CheckInputs(MoEParameters& parameters, const int64_t fc2_blocks_per_row = (inter_size + block_size - 1) / block_size; const int64_t fc3_blocks_per_row = (hidden_size + block_size - 1) / block_size; - CHECK_TENSOR_SHAPE(fc1_experts_scales, num_experts, fc1_inter_size, fc1_blocks_per_row); - CHECK_TENSOR_SHAPE(fc2_experts_scales, num_experts, hidden_size, fc2_blocks_per_row); - CHECK_TENSOR_SHAPE(fc3_experts_scales, num_experts, inter_size, fc3_blocks_per_row); + CHECK_TENSOR_SHAPE(fc1_experts_scales, local_num_experts, fc1_inter_size, fc1_blocks_per_row); + CHECK_TENSOR_SHAPE(fc2_experts_scales, local_num_experts, hidden_size, fc2_blocks_per_row); + CHECK_TENSOR_SHAPE(fc3_experts_scales, local_num_experts, inter_size, fc3_blocks_per_row); // Validate zero-point tensors (block-wise) const int64_t fc1_zp_blocks = (fc1_blocks_per_row + zp_pack_size - 1) / zp_pack_size; const int64_t fc2_zp_blocks = (fc2_blocks_per_row + zp_pack_size - 1) / zp_pack_size; const int64_t fc3_zp_blocks = (fc3_blocks_per_row + zp_pack_size - 1) / zp_pack_size; - CHECK_TENSOR_SHAPE(fc1_zero_points, num_experts, fc1_inter_size, fc1_zp_blocks); - CHECK_TENSOR_SHAPE(fc2_zero_points, num_experts, hidden_size, fc2_zp_blocks); - CHECK_TENSOR_SHAPE(fc3_zero_points, num_experts, inter_size, fc3_zp_blocks); + CHECK_TENSOR_SHAPE(fc1_zero_points, local_num_experts, fc1_inter_size, fc1_zp_blocks); + CHECK_TENSOR_SHAPE(fc2_zero_points, local_num_experts, hidden_size, fc2_zp_blocks); + CHECK_TENSOR_SHAPE(fc3_zero_points, local_num_experts, inter_size, fc3_zp_blocks); } else { // Row-wise quantization: 2D scale tensors or 3D with last dimension = 1 // Handle both {num_experts, features} and {num_experts, features, 1} shapes if (fc1_experts_scales != nullptr) { const auto& fc1_scales_dims = fc1_experts_scales->Shape().GetDims(); if (fc1_scales_dims.size() == 2) { - CHECK_TENSOR_SHAPE(fc1_experts_scales, num_experts, fc1_inter_size); - CHECK_TENSOR_SHAPE(fc1_zero_points, num_experts, (fc1_inter_size + zp_pack_size - 1) / zp_pack_size); + CHECK_TENSOR_SHAPE(fc1_experts_scales, local_num_experts, fc1_inter_size); + CHECK_TENSOR_SHAPE(fc1_zero_points, local_num_experts, (fc1_inter_size + zp_pack_size - 1) / zp_pack_size); } else if (fc1_scales_dims.size() == 3) { - CHECK_TENSOR_SHAPE(fc1_experts_scales, num_experts, fc1_inter_size, 1); - CHECK_TENSOR_SHAPE(fc1_zero_points, num_experts, (fc1_inter_size + zp_pack_size - 1) / zp_pack_size); + CHECK_TENSOR_SHAPE(fc1_experts_scales, local_num_experts, fc1_inter_size, 1); + CHECK_TENSOR_SHAPE(fc1_zero_points, local_num_experts, (fc1_inter_size + zp_pack_size - 1) / zp_pack_size); } else { ORT_THROW("fc1_experts_scales must be 2D or 3D tensor"); } @@ -197,11 +201,11 @@ Status CheckInputs(MoEParameters& parameters, if (fc2_experts_scales != nullptr) { const auto& fc2_scales_dims = fc2_experts_scales->Shape().GetDims(); if (fc2_scales_dims.size() == 2) { - CHECK_TENSOR_SHAPE(fc2_experts_scales, num_experts, hidden_size); - CHECK_TENSOR_SHAPE(fc2_zero_points, num_experts, (hidden_size + zp_pack_size - 1) / zp_pack_size); + CHECK_TENSOR_SHAPE(fc2_experts_scales, local_num_experts, hidden_size); + CHECK_TENSOR_SHAPE(fc2_zero_points, local_num_experts, (hidden_size + zp_pack_size - 1) / zp_pack_size); } else if (fc2_scales_dims.size() == 3) { - CHECK_TENSOR_SHAPE(fc2_experts_scales, num_experts, hidden_size, 1); - CHECK_TENSOR_SHAPE(fc2_zero_points, num_experts, (hidden_size + zp_pack_size - 1) / zp_pack_size); + CHECK_TENSOR_SHAPE(fc2_experts_scales, local_num_experts, hidden_size, 1); + CHECK_TENSOR_SHAPE(fc2_zero_points, local_num_experts, (hidden_size + zp_pack_size - 1) / zp_pack_size); } else { ORT_THROW("fc2_experts_scales must be 2D or 3D tensor"); } @@ -210,11 +214,11 @@ Status CheckInputs(MoEParameters& parameters, if (fc3_experts_scales != nullptr) { const auto& fc3_scales_dims = fc3_experts_scales->Shape().GetDims(); if (fc3_scales_dims.size() == 2) { - CHECK_TENSOR_SHAPE(fc3_experts_scales, num_experts, inter_size); - CHECK_TENSOR_SHAPE(fc3_zero_points, num_experts, (inter_size + zp_pack_size - 1) / zp_pack_size); + CHECK_TENSOR_SHAPE(fc3_experts_scales, local_num_experts, inter_size); + CHECK_TENSOR_SHAPE(fc3_zero_points, local_num_experts, (inter_size + zp_pack_size - 1) / zp_pack_size); } else if (fc3_scales_dims.size() == 3) { - CHECK_TENSOR_SHAPE(fc3_experts_scales, num_experts, inter_size, 1); - CHECK_TENSOR_SHAPE(fc3_zero_points, num_experts, (inter_size + zp_pack_size - 1) / zp_pack_size); + CHECK_TENSOR_SHAPE(fc3_experts_scales, local_num_experts, inter_size, 1); + CHECK_TENSOR_SHAPE(fc3_zero_points, local_num_experts, (inter_size + zp_pack_size - 1) / zp_pack_size); } else { ORT_THROW("fc3_experts_scales must be 2D or 3D tensor"); } @@ -241,6 +245,8 @@ Status CheckInputs(MoEParameters& parameters, parameters.parallel_type = MoEParallelType::TP; } } else if (num_experts > local_num_experts) { + ORT_RETURN_IF(local_num_experts <= 0 || num_experts % local_num_experts != 0, + "num_experts (", num_experts, ") must be divisible by local_num_experts (", local_num_experts, ")"); if (parameters.tensor_shards == 1) { parameters.parallel_type = MoEParallelType::EP; } else { diff --git a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc index 5170c982f248d..de055b4efaf9c 100644 --- a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc +++ b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc @@ -1,10 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include #include -#include "core/common/safeint.h" #include "core/providers/cuda/cuda_common.h" #include "contrib_ops/cuda/bert/transformer_cuda_common.h" #include "sharded_moe.h" @@ -19,18 +17,6 @@ namespace cuda { #if defined(ORT_USE_NCCL) -#define CHECK_CUDA(res) \ - if (res != cudaSuccess) { \ - cuda_result = res; \ - return; \ - } - -#define CHECK_NCCL(res) \ - if (res != ncclSuccess) { \ - nccl_result = res; \ - return; \ - } - #define REGISTER_KERNEL_TYPED(T) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ ShardedMoE, kMSDomain, 1, T, kCudaExecutionProvider, \ @@ -41,26 +27,17 @@ REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(MLFloat16) template -ShardedMoE::ShardedMoE(const OpKernelInfo& op_kernel_info) : NcclKernel(op_kernel_info), MoEBase(op_kernel_info) { +ShardedMoE::ShardedMoE(const OpKernelInfo& op_kernel_info) + : NcclKernel(op_kernel_info), MoEBase(op_kernel_info, GetDeviceProp()) { ORT_ENFORCE(op_kernel_info.GetAttr("tensor_shards", &tensor_shards_).IsOK()); ORT_ENFORCE(op_kernel_info.GetAttr("local_experts_start_index", &local_experts_start_index_).IsOK()); - rank_to_experts_start_index_.resize(nccl_->Size()); - - auto allocator = op_kernel_info.GetAllocator(OrtMemTypeDefault); - ORT_ENFORCE(SynchronizeExpertsStartIndex(allocator) == Status::OK()); + ORT_ENFORCE(tensor_shards_ >= 1, "tensor_shards must be positive, got ", tensor_shards_); + ORT_ENFORCE(local_experts_start_index_ >= 0, + "local_experts_start_index must not be negative, got ", local_experts_start_index_); } template Status ShardedMoE::ComputeInternal(OpKernelContext* context) const { - typedef typename ToCudaType::MappedType CudaT; - auto stream = context->GetComputeStream(); - - auto& device_prop = GetDeviceProp(); - const int sm = device_prop.major * 10 + device_prop.minor; - - AllocatorPtr allocator; - ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); - const Tensor* input = context->Input(0); const Tensor* router_probs = context->Input(1); const Tensor* fc1_experts_weights = context->Input(2); @@ -77,140 +54,66 @@ Status ShardedMoE::ComputeInternal(OpKernelContext* context) const { fc2_experts_weights, fc2_experts_bias_optional, nullptr, nullptr, fc3_experts_weights_optional, fc3_experts_bias_optional, nullptr, nullptr, 1, // no quantization so pack size is 1 - activation_type_ == ort_fastertransformer::ActivationType::SwiGLU, + IsFusedSwiglu(fc3_experts_weights_optional != nullptr), 0)); // no block-wise quantization for sharded MoE - - ORT_RETURN_IF_NOT(moe_params.num_experts % nccl_->Size() == 0, "num_experts should be divisible by world_size"); - - ort_fastertransformer::CutlassMoeFCRunner moe_runner(sm, - activation_type_, - fc3_experts_weights_optional != nullptr, - normalize_routing_weights_, - use_sparse_mixer_); - - size_t ws_size = moe_runner.getWorkspaceSize( - static_cast(moe_params.num_rows), static_cast(moe_params.hidden_size), - static_cast(moe_params.inter_size), static_cast(moe_params.num_experts), static_cast(k_)); - - size_t fc2_output_size = k_ * moe_params.num_rows * moe_params.hidden_size * sizeof(CudaT); - size_t expert_scales_size = k_ * moe_params.num_rows * sizeof(CudaT); - size_t expanded_source_row_to_expanded_dest_row_size = k_ * moe_params.num_rows * sizeof(int); - size_t expert_for_source_row_size = k_ * moe_params.num_rows * sizeof(int); - - // TODO: allocate one buffer and reuse it. - IAllocatorUniquePtr work_space = IAllocator::MakeUniquePtr(allocator, ws_size, false, stream); - IAllocatorUniquePtr fc2_output = IAllocator::MakeUniquePtr(allocator, fc2_output_size, false, stream); - IAllocatorUniquePtr fc2_output_bc = IAllocator::MakeUniquePtr(allocator, fc2_output_size, false, stream); - IAllocatorUniquePtr expert_scales = - IAllocator::MakeUniquePtr(allocator, expert_scales_size, false, stream); - IAllocatorUniquePtr expanded_source_row_to_expanded_dest_row = - IAllocator::MakeUniquePtr(allocator, expanded_source_row_to_expanded_dest_row_size, false, stream); - IAllocatorUniquePtr expert_for_source_row = - IAllocator::MakeUniquePtr(allocator, expert_for_source_row_size, false, stream); - - const CudaT* fc_scales_ptr = nullptr; - - moe_runner.run_moe_fc( - reinterpret_cast(input->template Data()), - reinterpret_cast(router_probs->template Data()), - reinterpret_cast(fc1_experts_weights->template Data()), std::move(fc_scales_ptr), - fc1_experts_bias_optional == nullptr - ? nullptr - : reinterpret_cast(fc1_experts_bias_optional->template Data()), - activation_type_, - fc3_experts_weights_optional == nullptr - ? nullptr - : reinterpret_cast(fc3_experts_weights_optional->template Data()), - std::move(fc_scales_ptr), - fc3_experts_bias_optional == nullptr - ? nullptr - : reinterpret_cast(fc3_experts_bias_optional->template Data()), - reinterpret_cast(fc2_experts_weights->template Data()), std::move(fc_scales_ptr), - static_cast(moe_params.num_rows), static_cast(moe_params.hidden_size), - static_cast(moe_params.inter_size), static_cast(moe_params.num_experts), - static_cast(moe_params.local_num_experts), static_cast(local_experts_start_index_), - static_cast(k_), reinterpret_cast(work_space.get()), reinterpret_cast(fc2_output.get()), - reinterpret_cast(expert_scales.get()), - reinterpret_cast(expanded_source_row_to_expanded_dest_row.get()), - reinterpret_cast(expert_for_source_row.get()), Stream(context)); + ORT_RETURN_IF_NOT(k_ > 0 && k_ <= moe_params.num_experts, + "ShardedMoE requires 0 < k <= num_experts, got k=", k_, + " and num_experts=", moe_params.num_experts); + + const int world_size = nccl_->Size(); + + // The kernel computes the experts (expert parallelism) or the slice of the intermediate dimension + // (tensor parallelism) that this rank owns, so every rank only produces a partial result that has + // to be combined with an all-reduce afterwards. + onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig parallelism_config{}; + switch (moe_params.parallel_type) { + case MoEParallelType::None: + break; + case MoEParallelType::EP: { + ORT_RETURN_IF_NOT(moe_params.local_num_experts * world_size == moe_params.num_experts, + "num_experts (", moe_params.num_experts, ") should be local_num_experts (", + moe_params.local_num_experts, ") times world_size (", world_size, ")"); + ORT_RETURN_IF_NOT(local_experts_start_index_ % moe_params.local_num_experts == 0, + "local_experts_start_index (", local_experts_start_index_, + ") should be a multiple of local_num_experts (", moe_params.local_num_experts, ")"); + const int ep_rank = static_cast(local_experts_start_index_ / moe_params.local_num_experts); + ORT_RETURN_IF_NOT(ep_rank < world_size, + "local_experts_start_index (", local_experts_start_index_, ") is out of range"); + parallelism_config = onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig( + /*tp_size*/ 1, /*tp_rank*/ 0, /*ep_size*/ world_size, ep_rank); + break; + } + case MoEParallelType::TP: { + ORT_RETURN_IF_NOT(moe_params.tensor_shards == world_size, + "tensor_shards (", moe_params.tensor_shards, ") should be equal to world_size (", + world_size, ")"); + parallelism_config = onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig( + /*tp_size*/ world_size, /*tp_rank*/ nccl_->Rank(), /*ep_size*/ 1, /*ep_rank*/ 0); + break; + } + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Expert and Tensor Parallelism is not supported yet"); + } Tensor* output = context->Output(0, input->Shape()); if (moe_params.parallel_type == MoEParallelType::None) { - fc2_output_bc = std::move(fc2_output); - } - - if (moe_params.parallel_type == MoEParallelType::EPAndTP) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Expert and Tensor Parallelism is not supported yet"); + return RunMoe(context, moe_params, parallelism_config, mGemmProfiler, mGemmProfilerMutex, + output->MutableDataRaw()); } - if (moe_params.parallel_type == MoEParallelType::TP) { - ORT_ENFORCE(moe_params.tensor_shards == nccl_->Size()); - - ORT_RETURN_IF_ERROR(FuncCustomAllReduce(nccl_, - Stream(context), - fc2_output.get(), - fc2_output_bc.get(), - static_cast(fc2_output_size / sizeof(CudaT)), - input->DataType(), - collective::IPCMemoryResourcePack::GetGlobalInstance())); - } + AllocatorPtr allocator; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); - if (moe_params.parallel_type == MoEParallelType::EP) { - size_t stride_count = moe_params.hidden_size; - size_t stride_bytes = stride_count * sizeof(CudaT); - int64_t total_past_rows = 0; - int64_t total_covered_rows = 0; - - NCCL_RETURN_IF_ERROR(ncclGroupStart()); - for (int rank = 0; rank < nccl_->Size(); ++rank) { - int64_t experts_start_index = rank_to_experts_start_index_[rank]; - moe_runner.get_total_rows_info(experts_start_index, moe_params.local_num_experts, total_past_rows, - total_covered_rows); - const char* src = reinterpret_cast(fc2_output.get()) + total_past_rows * stride_bytes; - char* dst = reinterpret_cast(fc2_output_bc.get()) + total_past_rows * stride_bytes; - NCCL_RETURN_IF_ERROR(ncclBroadcast(src, dst, total_covered_rows * stride_count, - GetNcclDataType(input->DataType()), rank, nccl_->Comm(), Stream(context))); - } - NCCL_RETURN_IF_ERROR(ncclGroupEnd()); - } + auto partial_output = Tensor::Create(input->DataType(), input->Shape(), allocator); + ORT_RETURN_IF_ERROR(RunMoe(context, moe_params, parallelism_config, mGemmProfiler, mGemmProfilerMutex, + partial_output->MutableDataRaw())); - ort_fastertransformer::finalize_moe_routing_kernelLauncher( - reinterpret_cast(fc2_output_bc.get()), reinterpret_cast(output->template MutableData()), - fc2_experts_bias_optional == nullptr - ? nullptr - : reinterpret_cast(fc2_experts_bias_optional->template Data()), - reinterpret_cast(expert_scales.get()), - reinterpret_cast(expanded_source_row_to_expanded_dest_row.get()), - reinterpret_cast(expert_for_source_row.get()), static_cast(moe_params.num_rows), - static_cast(moe_params.hidden_size), static_cast(k_), Stream(context)); - - return Status::OK(); + // Tokens that are not routed to a local expert contribute zero, and the bias of the second GEMM is + // only applied by rank 0, so summing the partial results of all ranks gives the final output. + return FuncAllReduce(nccl_->Comm(), Stream(context), partial_output.get(), output); } -template -Status ShardedMoE::SynchronizeExpertsStartIndex(AllocatorPtr& allocator) const { - using IndexType = int64_t; - size_t IndexTypeSize = sizeof(IndexType); - - IAllocatorUniquePtr experts_start_index_d = - IAllocator::MakeUniquePtr(allocator, 1, false); - IAllocatorUniquePtr rank_to_experts_start_index_d = - IAllocator::MakeUniquePtr(allocator, nccl_->Size(), false); - - CUDA_RETURN_IF_ERROR(cudaMemcpy(experts_start_index_d.get(), &local_experts_start_index_, IndexTypeSize, - cudaMemcpyHostToDevice)); - NCCL_RETURN_IF_ERROR(ncclAllGather(reinterpret_cast(experts_start_index_d.get()), - reinterpret_cast(rank_to_experts_start_index_d.get()), 1, - GetNcclDataType(DataTypeImpl::GetType()), nccl_->Comm(), - nullptr)); - - CUDA_RETURN_IF_ERROR(cudaMemcpy(const_cast(rank_to_experts_start_index_.data()), - rank_to_experts_start_index_d.get(), nccl_->Size() * IndexTypeSize, - cudaMemcpyDeviceToHost)); - - return Status::OK(); -} #endif } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h index a0f10328349e6..d0bd0fd2851ff 100644 --- a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h +++ b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h @@ -3,11 +3,13 @@ #pragma once -#include "contrib_ops/cuda/moe/ft_moe/moe_kernel.h" #include "contrib_ops/cuda/moe/moe_base.h" +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.h" #include "core/common/common.h" #include "nccl_kernels.h" +#include + namespace onnxruntime { namespace contrib { namespace cuda { @@ -23,11 +25,11 @@ class ShardedMoE final : public NcclKernel, public MoEBase { Status ComputeInternal(OpKernelContext* ctx) const override; private: - Status SynchronizeExpertsStartIndex(AllocatorPtr& alloc) const; - int64_t local_experts_start_index_; int64_t tensor_shards_; - InlinedVector rank_to_experts_start_index_; + + mutable onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler mGemmProfiler; + mutable std::mutex mGemmProfilerMutex; }; #endif diff --git a/onnxruntime/contrib_ops/cuda/moe/moe.cc b/onnxruntime/contrib_ops/cuda/moe/moe.cc index 1aa06a83efcde..cd74bc4c77b2b 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe.cc @@ -39,6 +39,27 @@ REGISTER_KERNEL_TYPED(float) REGISTER_KERNEL_TYPED(MLFloat16) REGISTER_KERNEL_TYPED(BFloat16) +int MoEBase::GetSwigluFusion(bool has_fc3) const { + using onnxruntime::llm::kernels::cutlass_kernels::ActivationType; + + // Backward compatibility: the published gpt-oss-20b model (and any model exported by ORT < 1.27) + // hard-coded the interleaved SwiGLU fusion layout and did not emit a swiglu_fusion attribute, so it + // falls back to the default of 0 ("not fused"). When the activation is SwiGLU, swiglu_fusion is 0, + // and there is no separate FC3 weight, the gate and value projections are actually pre-fused into FC1 + // (interleaved layout). Treat this as swiglu_fusion == 1 so those legacy models keep working. + if (activation_type_ == ActivationType::Swiglu && swiglu_fusion_ == 0 && !has_fc3) { + LogSwigluFusionRemapOnce(); + return 1; + } + + return swiglu_fusion_; +} + +bool MoEBase::IsFusedSwiglu(bool has_fc3) const { + using onnxruntime::llm::kernels::cutlass_kernels::ActivationType; + return activation_type_ == ActivationType::Swiglu && !has_fc3 && GetSwigluFusion(has_fc3) != 0; +} + template MoE::MoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoEBase(op_kernel_info, GetDeviceProp()) { } @@ -54,24 +75,6 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { const Tensor* fc3_experts_weights_optional = context->Input(6); const Tensor* fc3_experts_bias_optional = context->Input(7); - using onnxruntime::llm::kernels::cutlass_kernels::ActivationType; - - // Backward compatibility: the published gpt-oss-20b model (and any model exported by ORT < 1.27) - // hard-coded the interleaved SwiGLU fusion layout and did not emit a swiglu_fusion attribute, so it - // falls back to the default of 0 ("not fused"). When the activation is SwiGLU, swiglu_fusion is 0, - // and there is no separate FC3 weight, the gate and value projections are actually pre-fused into FC1 - // (interleaved layout). Treat this as swiglu_fusion == 1 so those legacy models keep working. - int swiglu_fusion = swiglu_fusion_; - if (activation_type_ == ActivationType::Swiglu && swiglu_fusion == 0 && - fc3_experts_weights_optional == nullptr) { - swiglu_fusion = 1; - LogSwigluFusionRemapOnce(); - } - - bool is_fused_swiglu = (activation_type_ == ActivationType::Swiglu) && - (swiglu_fusion != 0) && - (fc3_experts_weights_optional == nullptr); - MoEParameters moe_params; ORT_RETURN_IF_ERROR(::onnxruntime::contrib::moe_helper::CheckInputs( moe_params, input, router_probs, @@ -79,19 +82,47 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { fc2_experts_weights, fc2_experts_bias_optional, nullptr, nullptr, fc3_experts_weights_optional, fc3_experts_bias_optional, nullptr, nullptr, 1, // no quantization so pack size is 1 - is_fused_swiglu, + IsFusedSwiglu(fc3_experts_weights_optional != nullptr), 0)); // no block-wise quantization for regular MoE ORT_RETURN_IF_NOT(k_ > 0 && k_ <= moe_params.num_experts, "MoE requires 0 < k <= num_experts, got k=", k_, " and num_experts=", moe_params.num_experts); + Tensor* output = context->Output(0, input->Shape()); + + return RunMoe(context, moe_params, + onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig{}, + mGemmProfiler, mGemmProfilerMutex, output->MutableDataRaw()); +} + +template +Status MoEBase::RunMoe(OpKernelContext* context, + const MoEParameters& moe_params, + onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig parallelism_config, + onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler& gemm_profiler, + std::mutex& gemm_profiler_mutex, + void* output_data) const { + const Tensor* input = context->Input(0); + const Tensor* router_probs = context->Input(1); + const Tensor* fc1_experts_weights = context->Input(2); + const Tensor* fc1_experts_bias_optional = context->Input(3); + const Tensor* fc2_experts_weights = context->Input(4); + const Tensor* fc2_experts_bias_optional = context->Input(5); + const Tensor* fc3_experts_weights_optional = context->Input(6); + + using onnxruntime::llm::kernels::cutlass_kernels::ActivationType; + + const int swiglu_fusion = GetSwigluFusion(fc3_experts_weights_optional != nullptr); + using CudaT = typename OrtToCudaType::type; - void* stream_obj = GetComputeStream(context); - cudaStream_t stream = Stream(context); + onnxruntime::Stream* stream_obj = context->GetComputeStream(); + cudaStream_t stream = stream_obj == nullptr ? nullptr : static_cast(stream_obj->GetHandle()); + + AllocatorPtr allocator; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); - auto& device_prop = GetDeviceProp(); - int sm = device_prop.major * 10 + device_prop.minor; + int sm = sm_; // SM90 TMA WS kernels only support f16/bf16, not float32. // Force SM80 path for float32 to use legacy kernels. @@ -115,13 +146,18 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { " for SM", sm, ", got ", moe_params.inter_size); } - using onnxruntime::llm::kernels::cutlass_kernels::ActivationType; ActivationType kernel_activation_type = activation_type_; - if (activation_type_ == ActivationType::Silu && fc3_experts_weights_optional != nullptr) { - // Mixtral case: SiLU activation with separate FC3. - // Kernel supports SwiGLU which is Linear * SiLU(Gate). - // We map Mixtral to SwiGLU by packing weights as [FC3, FC1] (Linear, Gate). - kernel_activation_type = ActivationType::Swiglu; + if (fc3_experts_weights_optional != nullptr) { + // A separate FC3 means a gated activation: output = Act(FC1 * x) * (FC3 * x). The kernel expresses + // that as SwiGLU/GeGLU and expects FC1 and FC3 packed into a single [E, 2 * inter_size, hidden_size] + // buffer, so the activation has to be mapped to its gated variant. Otherwise the kernel would use + // inter_size as the per-expert stride of a buffer that has 2 * inter_size rows per expert and would + // read the weights of the wrong expert. + if (activation_type_ == ActivationType::Silu) { + kernel_activation_type = ActivationType::Swiglu; + } else if (activation_type_ == ActivationType::Gelu) { + kernel_activation_type = ActivationType::Geglu; + } } onnxruntime::llm::kernels::cutlass_kernels::CutlassMoeFCRunner moe_runner(sm, @@ -130,7 +166,6 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { use_sparse_mixer_); constexpr bool use_awq = false; - onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig parallelism_config{}; if (onnxruntime::llm::common::getEnvForceDeterministicMOE()) { auto tactics = moe_runner.getTactics(); @@ -138,11 +173,9 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { moe_runner.setTactic(tactics[0], tactics[0]); } } else { - std::lock_guard profiler_lock(mGemmProfilerMutex); - AllocatorPtr allocator; - ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); - mGemmProfiler.setAllocator(std::move(allocator)); - mGemmProfiler.setProfilerParams(static_cast(moe_params.num_experts), static_cast(this->k_), + std::lock_guard profiler_lock(gemm_profiler_mutex); + gemm_profiler.setAllocator(allocator); + gemm_profiler.setProfilerParams(static_cast(moe_params.num_experts), static_cast(this->k_), static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size), static_cast(this->block_size_), kernel_activation_type, false, true, parallelism_config, sm); @@ -171,18 +204,18 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { // (small M) and prefill (large M) each profile and select their own best tile shape. GemmDims dims(static_cast(moe_params.num_rows), static_cast(moe_params.num_rows), static_cast(moe_params.inter_size), static_cast(moe_params.hidden_size)); - mGemmProfiler.profileTactics(&moe_runner, dims, id1, stream); + gemm_profiler.profileTactics(&moe_runner, dims, id1, stream); } - auto config1 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), id1); + auto config1 = gemm_profiler.getBestConfig(static_cast(moe_params.num_rows), id1); // GEMM 2 MoeGemmId id2(static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size), dtype, MoeGemmId::GemmType::Gemm2); if (!stream_is_capturing) { GemmDims dims(static_cast(moe_params.num_rows), static_cast(moe_params.num_rows), static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size)); - mGemmProfiler.profileTactics(&moe_runner, dims, id2, stream); + gemm_profiler.profileTactics(&moe_runner, dims, id2, stream); } - auto config2 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), id2); + auto config2 = gemm_profiler.getBestConfig(static_cast(moe_params.num_rows), id2); // Capture-safe fallback: if profiling was skipped (graph capture) and no tuned config was // cached from a prior non-capturing run, use the runner's default tactic instead of leaving @@ -215,7 +248,7 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { size_t permutation_bytes = expanded_rows * sizeof(int); size_t total_scratch_bytes = SafeInt(ws_size) + scales_bytes + indices_bytes + permutation_bytes; - auto work_space = GetScratchBuffer(total_scratch_bytes, stream_obj); + auto work_space = IAllocator::MakeUniquePtr(allocator, total_scratch_bytes, false, stream_obj); char* workspace_ptr = reinterpret_cast(work_space.get()); float* expert_scales = reinterpret_cast(workspace_ptr + ws_size); int* expert_indices = reinterpret_cast(workspace_ptr + ws_size + scales_bytes); @@ -273,8 +306,6 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { } } - Tensor* output = context->Output(0, input->Shape()); - onnxruntime::llm::kernels::cutlass_kernels::QuantParams quant_params{}; // ============================================================================= @@ -287,7 +318,11 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { // Calculate buffer sizes size_t fc1_block_size = static_cast(moe_params.inter_size) * static_cast(moe_params.hidden_size); - int E = static_cast(moe_params.num_experts); + + // The weights only contain the experts owned by this rank (all of them unless expert parallelism + // is used), and the kernel indexes them relative to the first local expert. + int E = static_cast(moe_params.local_num_experts); + int64_t const start_expert = moe_params.local_num_experts * parallelism_config.ep_rank; // FC1 Handling const CudaT* fc1_input_ptr = reinterpret_cast(fc1_experts_weights->DataRaw()); @@ -306,17 +341,16 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { // Each expert has 2*I*H elements = 2 * fc1_block_size const CudaT* fc3_input_ptr = reinterpret_cast(fc3_experts_weights_optional->DataRaw()); size_t fc1_total_size = E * 2 * fc1_block_size * sizeof(CudaT); - fc1_processed_buffer = GetScratchBuffer(fc1_total_size, stream_obj); + fc1_processed_buffer = IAllocator::MakeUniquePtr(allocator, fc1_total_size, false, stream_obj); CudaT* fc1_fc3_processed_ptr = reinterpret_cast(fc1_processed_buffer.get()); fc1_processed_ptr = fc1_fc3_processed_ptr; for (int e = 0; e < E; ++e) { - // Horizontally stack [FC3 | FC1] within each expert's block to match SwiGLU convention - // Kernel computes: Linear(1st half) * SiLU(Gate(2nd half)) - // Mixtral wants: FC3 * SiLU(FC1) - // So: 1st half = FC3 (Linear), 2nd half = FC1 (Gate) - CudaT* dest_fc1 = fc1_fc3_processed_ptr + e * 2 * fc1_block_size; // First half of expert e (Gate/FC1) - CudaT* dest_fc3 = fc1_fc3_processed_ptr + e * 2 * fc1_block_size + fc1_block_size; // Second half of expert e (Linear/FC3) + // Horizontally stack [FC1 | FC3] within each expert's block: the gated activation kernel reads the + // gate from the first inter_size rows and the linear part from the next inter_size rows, and + // computes Act(Gate) * Linear, which is what FC1/FC3 mean here. + CudaT* dest_fc1 = fc1_fc3_processed_ptr + e * 2 * fc1_block_size; // Gate (FC1) + CudaT* dest_fc3 = fc1_fc3_processed_ptr + e * 2 * fc1_block_size + fc1_block_size; // Linear (FC3) // Copy [I, H] directly CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(dest_fc1, fc1_input_ptr + e * fc1_block_size, fc1_block_size * sizeof(CudaT), cudaMemcpyDeviceToDevice, stream)); @@ -324,27 +358,55 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { } } + // FC1 bias handling + const CudaT* fc1_bias_ptr = fc1_experts_bias_optional == nullptr + ? nullptr + : reinterpret_cast(fc1_experts_bias_optional->Data()); + IAllocatorUniquePtr fc1_bias_processed_buffer; + if (fc3_experts_weights_optional != nullptr && fc1_bias_ptr != nullptr) { + // The gated GEMM produces 2 * inter_size columns per expert and indexes the bias with that same + // stride, so the [E, inter_size] bias of FC1 has to be padded to [E, 2 * inter_size]. Without the + // padding the kernel would read the bias of the wrong expert. FC3 has no bias here, so the second + // (linear) half is zero. + size_t const bias_row_bytes = static_cast(moe_params.inter_size) * sizeof(CudaT); + size_t const total_bias_bytes = static_cast(E) * 2 * bias_row_bytes; + fc1_bias_processed_buffer = IAllocator::MakeUniquePtr(allocator, total_bias_bytes, false, stream_obj); + CudaT* padded_bias = reinterpret_cast(fc1_bias_processed_buffer.get()); + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(padded_bias, 0, total_bias_bytes, stream)); + CUDA_RETURN_IF_ERROR(cudaMemcpy2DAsync(padded_bias, 2 * bias_row_bytes, fc1_bias_ptr, bias_row_bytes, + bias_row_bytes, static_cast(E), cudaMemcpyDeviceToDevice, stream)); + fc1_bias_ptr = padded_bias; + } + // FC2 Handling const CudaT* fc2_input_ptr = reinterpret_cast(fc2_experts_weights->DataRaw()); // Layout matches kernel expectation [H, I]. Use directly. const CudaT* fc2_processed_ptr = fc2_input_ptr; + // fc2_experts_bias covers all experts (it is applied after the ranks are combined), but the kernel + // indexes it by local expert, so skip the experts that belong to the previous ranks. + const CudaT* fc2_bias_ptr = nullptr; + if (fc2_experts_bias_optional != nullptr) { + fc2_bias_ptr = reinterpret_cast(fc2_experts_bias_optional->Data()) + + start_expert * moe_params.hidden_size; + } + moe_runner.runMoe( - reinterpret_cast(input->template Data()), + reinterpret_cast(input->Data()), nullptr, // input_sf expert_indices, // token_selected_experts expert_scales, // token_final_scales fc1_processed_ptr, - fc1_experts_bias_optional == nullptr ? nullptr : reinterpret_cast(fc1_experts_bias_optional->template Data()), + fc1_bias_ptr, kernel_activation_type, fc2_processed_ptr, - fc2_experts_bias_optional == nullptr ? nullptr : reinterpret_cast(fc2_experts_bias_optional->template Data()), + fc2_bias_ptr, quant_params, static_cast(moe_params.num_rows), static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size), static_cast(moe_params.num_experts), static_cast(k_), workspace_ptr, - reinterpret_cast(output->template MutableData()), + output_data, unpermuted_row_to_permuted_row, parallelism_config, [&]() { @@ -361,6 +423,19 @@ Status MoE::ComputeInternal(OpKernelContext* context) const { return Status::OK(); } +template Status MoEBase::RunMoe(OpKernelContext*, const MoEParameters&, + onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig, + onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler&, + std::mutex&, void*) const; +template Status MoEBase::RunMoe(OpKernelContext*, const MoEParameters&, + onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig, + onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler&, + std::mutex&, void*) const; +template Status MoEBase::RunMoe(OpKernelContext*, const MoEParameters&, + onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig, + onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler&, + std::mutex&, void*) const; + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_base.h b/onnxruntime/contrib_ops/cuda/moe/moe_base.h index 4964259fd8e90..1f061e528546b 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_base.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_base.h @@ -12,10 +12,12 @@ #include "core/common/common.h" #include "core/framework/op_kernel.h" #include "contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.h" +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.h" #include "contrib_ops/cpu/moe/moe_helper.h" #include "core/providers/cuda/cuda_common.h" #include "contrib_ops/cuda/llm/moe_gemm/common.h" #include +#include #ifdef __GNUC__ #pragma GCC diagnostic pop @@ -76,6 +78,25 @@ class MoEBase { sm_ = device_prop.major * 10 + device_prop.minor; } + // Returns the effective swiglu_fusion. Models exported before the swiglu_fusion attribute existed + // (like the published gpt-oss-20b) hard-coded the interleaved fused layout but leave the attribute + // at its default of 0, so treat SwiGLU without a separate FC3 weight as interleaved fusion. + int GetSwigluFusion(bool has_fc3) const; + + // True when FC1 holds the pre-fused gate and value projections (shape [E, 2 * inter_size, K]). + bool IsFusedSwiglu(bool has_fc3) const; + + // Shared implementation of the MoE computation used by both the MoE and the ShardedMoE kernels. + // The result (which is only a partial result when parallelism_config spans multiple ranks, so the + // caller has to combine the ranks with an all-reduce) is written to output_data. + template + Status RunMoe(OpKernelContext* context, + const MoEParameters& moe_params, + onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig parallelism_config, + onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler& gemm_profiler, + std::mutex& gemm_profiler_mutex, + void* output_data) const; + bool normalize_routing_weights_; bool use_sparse_mixer_; int64_t k_; diff --git a/onnxruntime/test/python/transformers/sharded_moe/test_sharded_moe.py b/onnxruntime/test/python/transformers/sharded_moe/test_sharded_moe.py index 2d29135726839..2c3ad07bb6c55 100644 --- a/onnxruntime/test/python/transformers/sharded_moe/test_sharded_moe.py +++ b/onnxruntime/test/python/transformers/sharded_moe/test_sharded_moe.py @@ -40,7 +40,10 @@ def print_out(*args): ORT_DTYPE = TensorProto.FLOAT16 NP_TYPE = np.float16 if ORT_DTYPE == TensorProto.FLOAT16 else np.float32 THRESHOLD_TP = 3e-2 -THRESHOLD_EP = 1e-6 +# Expert parallelism is not bit-exact anymore: the sharded and the unsharded session run the same +# CUTLASS MoE kernels with a different number of experts per rank, so they may pick a different GEMM +# tactic and accumulate in a different order. The difference stays within a couple of fp16 ulps. +THRESHOLD_EP = 3e-2 def create_moe_onnx_graph( From 0ad62e3135981e632ee5d7835fa13a1bdc640252 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 30 Jul 2026 20:37:44 +0000 Subject: [PATCH 2/2] fix plugin ep build --- onnxruntime/contrib_ops/cuda/moe/moe.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/moe/moe.cc b/onnxruntime/contrib_ops/cuda/moe/moe.cc index cd74bc4c77b2b..a1eb0effc4a94 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe.cc @@ -116,8 +116,8 @@ Status MoEBase::RunMoe(OpKernelContext* context, using CudaT = typename OrtToCudaType::type; - onnxruntime::Stream* stream_obj = context->GetComputeStream(); - cudaStream_t stream = stream_obj == nullptr ? nullptr : static_cast(stream_obj->GetHandle()); + void* stream_obj = GetComputeStream(context); + cudaStream_t stream = Stream(context); AllocatorPtr allocator; ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator));