From 8dd979fe72c87171aa867de45a63454d8998196d Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Wed, 29 Jul 2026 12:41:14 +0000 Subject: [PATCH] onnxruntime: fix cuda build --- .../drop-orphaned-sharded-moe.patch | 320 ++++++++++++++++++ pkgs/by-name/on/onnxruntime/package.nix | 5 + 2 files changed, 325 insertions(+) create mode 100644 pkgs/by-name/on/onnxruntime/drop-orphaned-sharded-moe.patch diff --git a/pkgs/by-name/on/onnxruntime/drop-orphaned-sharded-moe.patch b/pkgs/by-name/on/onnxruntime/drop-orphaned-sharded-moe.patch new file mode 100644 index 0000000000000..1467d867e025f --- /dev/null +++ b/pkgs/by-name/on/onnxruntime/drop-orphaned-sharded-moe.patch @@ -0,0 +1,320 @@ +Drop the orphaned ShardedMoE CUDA contrib op. + +ShardedMoE's backend, contrib_ops/cuda/moe/ft_moe/moe_kernel.h and the +ort_fastertransformer::CutlassMoeFCRunner it wraps, was removed in the MoE GEMM +refactor (microsoft/onnxruntime#28467) which moved everything to +contrib_ops/cuda/llm/moe_gemm/. moe.cc/moe_base.h were ported to the new +onnxruntime::llm::kernels::cutlass_kernels API, but sharded_moe.{h,cc} were not, +so the file still includes the deleted header and fails to compile whenever +onnxruntime_USE_NCCL is enabled: + + contrib_ops/cuda/collective/sharded_moe.h:6:10: fatal error: + contrib_ops/cuda/moe/ft_moe/moe_kernel.h: No such file or directory + +The op has been dead (uncompilable) since that refactor. Remove it along with +its build/registration entries so NCCL builds again; all other NCCL collective +ops (AllReduce, AllGather, AllToAll, DistributedMatMul, ...) are unaffected. + +diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake +index 2c574f298f..62b47465a9 100644 +--- a/cmake/onnxruntime_providers_cuda.cmake ++++ b/cmake/onnxruntime_providers_cuda.cmake +@@ -89,8 +89,6 @@ + if (NOT onnxruntime_USE_NCCL) + list(REMOVE_ITEM onnxruntime_cuda_contrib_ops_cc_srcs + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/nccl_kernels.cc" +- "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/sharded_moe.h" +- "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/sharded_moe.cc" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/sharding_spec.cc" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/sharding.cc" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/collective/distributed_matmul.cc" +diff --git a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc +deleted file mode 100644 +index 5170c982f2..0000000000 +--- a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc ++++ /dev/null +@@ -1,218 +0,0 @@ +-// 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" +- +-using namespace onnxruntime::cuda; +-using namespace ::onnxruntime::common; +-using namespace ONNX_NAMESPACE; +- +-namespace onnxruntime { +-namespace contrib { +-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, \ +- (*KernelDefBuilder::Create()).MayInplace(0, 0).TypeConstraint("T", DataTypeImpl::GetTensorType()), \ +- ShardedMoE); +- +-REGISTER_KERNEL_TYPED(float) +-REGISTER_KERNEL_TYPED(MLFloat16) +- +-template +-ShardedMoE::ShardedMoE(const OpKernelInfo& op_kernel_info) : NcclKernel(op_kernel_info), MoEBase(op_kernel_info) { +- 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()); +-} +- +-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); +- 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); +- const Tensor* fc3_experts_bias_optional = context->Input(7); +- +- MoEParameters moe_params(tensor_shards_); +- ORT_RETURN_IF_ERROR(::onnxruntime::contrib::moe_helper::CheckInputs( +- moe_params, input, router_probs, +- fc1_experts_weights, fc1_experts_bias_optional, nullptr, nullptr, +- 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, +- 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)); +- +- 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"); +- } +- +- 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())); +- } +- +- 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()); +- } +- +- 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(); +-} +- +-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 +-} // namespace contrib +-} // namespace onnxruntime +diff --git a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h b/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h +deleted file mode 100644 +index a0f1032834..0000000000 +--- a/onnxruntime/contrib_ops/cuda/collective/sharded_moe.h ++++ /dev/null +@@ -1,37 +0,0 @@ +-// Copyright (c) Microsoft Corporation. All rights reserved. +-// Licensed under the MIT License. +- +-#pragma once +- +-#include "contrib_ops/cuda/moe/ft_moe/moe_kernel.h" +-#include "contrib_ops/cuda/moe/moe_base.h" +-#include "core/common/common.h" +-#include "nccl_kernels.h" +- +-namespace onnxruntime { +-namespace contrib { +-namespace cuda { +- +-#if defined(ORT_USE_NCCL) +- +-using namespace onnxruntime::cuda; +- +-template +-class ShardedMoE final : public NcclKernel, public MoEBase { +- public: +- explicit ShardedMoE(const OpKernelInfo& op_kernel_info); +- 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_; +-}; +- +-#endif +- +-} // namespace cuda +-} // namespace contrib +-} // namespace onnxruntime +diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +index da7ef35d25..e4b3112cd0 100644 +--- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc ++++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +@@ -250,9 +250,6 @@ class CUDA_MS_OP_CLASS_NAME(1, AllReduce); + class CUDA_MS_OP_CLASS_NAME(1, AllGather); + class CUDA_MS_OP_CLASS_NAME(1, AllToAll); + +-class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, ShardedMoE); +-class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, ShardedMoE); +- + class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, DistributedMatMul); + class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, DistributedMatMul); + +@@ -516,9 +513,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { + BuildKernelCreateInfo, + BuildKernelCreateInfo, + +- BuildKernelCreateInfo, +- BuildKernelCreateInfo, +- + BuildKernelCreateInfo, + BuildKernelCreateInfo, diff --git a/pkgs/by-name/on/onnxruntime/package.nix b/pkgs/by-name/on/onnxruntime/package.nix index b1a972d4d2206..ef4d70237ce00 100644 --- a/pkgs/by-name/on/onnxruntime/package.nix +++ b/pkgs/by-name/on/onnxruntime/package.nix @@ -131,6 +131,11 @@ effectiveStdenv.mkDerivation (finalAttrs: { # https://github.com/microsoft/onnxruntime/issues/9155 # Patch adapted from https://gitlab.alpinelinux.org/alpine/aports/-/raw/462dfe0eb4b66948fe48de44545cc22bb64fdf9f/community/onnxruntime/0001-Remove-MATH_NO_EXCEPT-macro.patch ./remove-MATH_NO_EXCEPT-macro.patch + ] + ++ lib.optionals cudaSupport [ + # Drop the orphaned ShardedMoE CUDA contrib op, whose ft_moe backend was removed upstream + # Fix submitted upstream: https://github.com/microsoft/onnxruntime/pull/31139 + ./drop-orphaned-sharded-moe.patch ]; postPatch = ''