Skip to content

[TRTLLM-14703][fix] 1st cherrypick [combine MR99+MR100+MR102 on latest feat (with MR100 dispatch wiring)], 2nd cherrypick [deferred-finalize scales fix, fused-epilogue batch gate, hardening] #17088

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
brnguyen2 merged 4 commits into from
Jul 31, 2026
Merged

[TRTLLM-14703][fix] 1st cherrypick [combine MR99+MR100+MR102 on latest feat (with MR100 dispatch wiring)], 2nd cherrypick [deferred-finalize scales fix, fused-epilogue batch gate, hardening] #17088

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
574265e
[None][perf] Kimi K3: combine MR99+MR100+MR102 on latest feat (with M…
litaotju Jul 30, 2026
2e4c343
[None][fix] Kimi K3: unblock !103 - deferred-finalize scales fix, fus…
moraxu Jul 30, 2026
575d7f3
[None][chore] Apply pre-commit formatting to the cherry-picked changes
moraxu Jul 30, 2026
7003910
[None][chore] Move FORCE_SEPARATED_ROUTING below the imports to fix r…
moraxu Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 151 additions & 2 deletions cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -335,6 +335,9 @@ struct MnnvlAllReduceKernelParams
uint32_t* bufferFlags;
bool waitForResults;
QuantizationSFLayout layout;
void const* expertScaleFactorPtr = nullptr;
int32_t const* expandedIdxToPermutedIdx = nullptr;
int topK = 0;
};

template <typename PackedType, typename T>
Expand All @@ -350,6 +353,42 @@ inline __device__ void sanitizeLamportPayload(PackedVec<PackedType, T>& value)
}
}

template <typename T, typename ScaleType, typename PackedType>
inline __device__ PackedVec<PackedType, T> finalizeMoeRoutes(
MnnvlAllReduceKernelParams<T> const& params, int token, int packedIdx)
{
constexpr int kEltsPerThread = sizeof(PackedType) / sizeof(T);
PackedVec<PackedType, T> accumulated;
#pragma unroll
for (int i = 0; i < kEltsPerThread; ++i)
{
accumulated.elements[i] = cuda_cast<T, float>(0.F);
}

auto const* expertScaleFactor = static_cast<ScaleType const*>(params.expertScaleFactorPtr);
for (int k = 0; k < params.topK; ++k)
{
int const expandedIdx = token * params.topK + k;
int32_t const permutedIdx = params.expandedIdxToPermutedIdx[expandedIdx];
if (permutedIdx < 0)
{
continue;
}

PackedVec<PackedType, T> routeOutput;
routeOutput.packed
= loadPacked<PackedType>(&params.shardPtr[permutedIdx * params.tokenDim + packedIdx * kEltsPerThread]);
float const routeScale = cuda_cast<float, ScaleType>(expertScaleFactor[expandedIdx]);
#pragma unroll
for (int i = 0; i < kEltsPerThread; ++i)
{
T const weightedRoute = cuda_cast<T, float>(cuda_cast<float, T>(routeOutput.elements[i]) * routeScale);
accumulated.elements[i] += weightedRoute;
}
}
return accumulated;
}

template <int Rank, uint8_t WorldSize, uint8_t LocalRank, typename T, typename PackedType, int kELTS_PER_THREAD>
inline __device__ bool pollOneshotRemoteRank(
PackedVec<PackedType, T>* remoteValues, T* stagePtrLocal, int token, int tokenDim, int packedIdx)
Expand Down Expand Up @@ -568,6 +607,7 @@ using detail::loadPacked;
using detail::blockReduceSum;
using detail::divUp;
using detail::copyF4;
using detail::finalizeMoeRoutes;
using detail::MnnvlAllReduceKernelParams;
using detail::sanitizeLamportPayload;
using detail::accumulateLamportRanksChunked;
Expand All @@ -576,7 +616,7 @@ using detail::writeEpilogueOutput;

template <uint8_t WorldSize, typename T,
ar_fusion::AllReduceFusionPattern Pattern = ar_fusion::AllReduceFusionPattern::kAllReduce,
typename PackedType = float4>
bool FuseMoeFinalize = false, typename ScaleType = T, typename PackedType = float4>
__global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllReduceKernelParams<T> params)
{
constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T);
Expand Down Expand Up @@ -606,7 +646,14 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed
PackedVec<PackedType, T> val;
if (inBounds)
{
val.packed = loadPacked<PackedType>(&params.shardPtr[threadOffset]);
if constexpr (FuseMoeFinalize)
{
val = finalizeMoeRoutes<T, ScaleType, PackedType>(params, token, packedIdx);
}
else
{
val.packed = loadPacked<PackedType>(&params.shardPtr[threadOffset]);
}
sanitizeLamportPayload(val);

reinterpret_cast<PackedType*>(
Expand Down Expand Up @@ -865,6 +912,108 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params)
}
}

void oneshotMoeFinalizeAllreduceRMSNormOp(MoeFinalizeAllReduceRMSNormParams const& params)
{
static int const kSMVersion = tensorrt_llm::common::getSMVersion();
TLLM_CHECK_WITH_INFO(kSMVersion >= 90, "[MNNVL MoE Finalize AllReduce RMSNorm] requires SM 90 or newer.");
TLLM_CHECK_WITH_INFO(params.input != nullptr, "MoE FC2 output must be provided.");
TLLM_CHECK_WITH_INFO(params.output != nullptr, "RMSNorm output must be provided.");
TLLM_CHECK_WITH_INFO(params.gamma != nullptr, "RMSNorm gamma must be provided.");
TLLM_CHECK_WITH_INFO(params.expertScaleFactor != nullptr, "MoE route weights must be provided.");
TLLM_CHECK_WITH_INFO(
params.expandedIdxToPermutedIdx != nullptr, "MoE expanded-to-permuted indices must be provided.");
TLLM_CHECK_WITH_INFO(params.topK > 0, "MoE top-k must be positive.");

int const numTokens = params.numTokens;
int const tokenDim = params.tokenDim;
int const eltsPerThread = sizeof(float4) / getDTypeSize(params.dType);
auto [blockSize, clusterSize, loadsPerThread] = adjustGridConfig<true>(numTokens, tokenDim, eltsPerThread);
dim3 grid(numTokens, clusterSize, 1);

TLLM_LOG_DEBUG(
"[MNNVL MoE Finalize AllReduce RMSNorm] Dispatch: grid size: (%d, %d, 1), block_size: %d, "
"cluster_size: %d, loads_per_thread: %d, top_k: %d",
numTokens, clusterSize, blockSize, clusterSize, loadsPerThread, params.topK);

TLLM_CHECK_WITH_INFO(blockSize <= 1024 && loadsPerThread == 1,
"Hidden Dimension %d exceeds the maximum supported hidden dimension (%d)", tokenDim, 1024 * 8 * eltsPerThread);

cudaLaunchAttribute attrs[2];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0;
attrs[1].id = cudaLaunchAttributeClusterDimension;
attrs[1].val.clusterDim.x = 1;
attrs[1].val.clusterDim.y = clusterSize;
attrs[1].val.clusterDim.z = 1;

cudaLaunchConfig_t config{
.gridDim = grid,
.blockDim = blockSize,
.dynamicSmemBytes = 0,
.stream = params.stream,
.attrs = attrs,
.numAttrs = 2U,
};

#define LAUNCH_MNNVL_MOE_FINALIZE_KERNEL(WORLD_SIZE, T, SCALE_TYPE) \
TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, \
&oneshotAllreduceFusionKernel<WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARRMSNorm, true, SCALE_TYPE>, \
kernelParams));

auto dispatchImpl = [&](auto* typePtr, auto* scaleTypePtr) -> bool
{
using T = std::remove_pointer_t<decltype(typePtr)>;
using ScaleType = std::remove_pointer_t<decltype(scaleTypePtr)>;
T** ucPtrs = reinterpret_cast<T**>(params.bufferPtrsDev);
T* mcPtr = reinterpret_cast<T*>(params.multicastPtr);
T* output = reinterpret_cast<T*>(params.output);
T const* input = reinterpret_cast<T const*>(params.input);
T const* gamma = reinterpret_cast<T const*>(params.gamma);
MnnvlAllReduceKernelParams<T> kernelParams{output, nullptr, input, nullptr, gamma, ucPtrs,
reinterpret_cast<T*>(params.bufferPtrLocal), mcPtr, nullptr, nullptr, nullptr, numTokens, tokenDim,
params.nRanks, params.rank, static_cast<float>(params.epsilon), params.bufferFlags, false, params.layout,
params.expertScaleFactor, params.expandedIdxToPermutedIdx, params.topK};

switch (params.nRanks)
{
case 2: LAUNCH_MNNVL_MOE_FINALIZE_KERNEL(2, T, ScaleType); return true;
case 4: LAUNCH_MNNVL_MOE_FINALIZE_KERNEL(4, T, ScaleType); return true;
case 8: LAUNCH_MNNVL_MOE_FINALIZE_KERNEL(8, T, ScaleType); return true;
case 16: LAUNCH_MNNVL_MOE_FINALIZE_KERNEL(16, T, ScaleType); return true;
}
return false;
};

#undef LAUNCH_MNNVL_MOE_FINALIZE_KERNEL

bool launched = false;
if (params.dType == nvinfer1::DataType::kBF16)
{
if (params.scaleDType == nvinfer1::DataType::kFLOAT)
{
launched = dispatchImpl((__nv_bfloat16*) nullptr, (float*) nullptr);
}
else if (params.scaleDType == nvinfer1::DataType::kBF16)
{
launched = dispatchImpl((__nv_bfloat16*) nullptr, (__nv_bfloat16*) nullptr);
}
}
else if (params.dType == nvinfer1::DataType::kHALF)
{
if (params.scaleDType == nvinfer1::DataType::kFLOAT)
{
launched = dispatchImpl((__nv_half*) nullptr, (float*) nullptr);
}
else if (params.scaleDType == nvinfer1::DataType::kHALF)
{
launched = dispatchImpl((__nv_half*) nullptr, (__nv_half*) nullptr);
}
}
TLLM_CHECK_WITH_INFO(launched,
"Failed to dispatch MNNVL MoE Finalize AllReduce RMSNorm kernel for nranks=%d, dtype=%d, scale_dtype=%d.",
params.nRanks, static_cast<int>(params.dType), static_cast<int>(params.scaleDType));
}

enum MNNVLTwoShotStage : uint8_t
{
SCATTER = 0,
Expand Down
16 changes: 16 additions & 0 deletions cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,8 +73,24 @@ struct AllReduceFusionParams
//! @}
};

/**
* \brief Parameters for fusing MoE finalize, MNNVL AllReduce, and RMSNorm.
*
* The local MoE FC2 output remains in permuted-route layout. The kernel
* applies the per-route weights, reduces the local top-k routes, performs
* MNNVL AllReduce, and normalizes the global latent output.
*/
struct MoeFinalizeAllReduceRMSNormParams : public AllReduceFusionParams
{
int topK; //!< Number of selected routes per token
nvinfer1::DataType scaleDType; //!< Data type of expertScaleFactor
void const* expertScaleFactor; //!< Per-token route weights [numTokens, topK]
int32_t const* expandedIdxToPermutedIdx; //!< Route-to-permuted-row map [numTokens, topK]
};

void oneshotAllreduceFusionOp(AllReduceFusionParams const& params);
void twoshotAllreduceFusionOp(AllReduceFusionParams const& params);
void oneshotMoeFinalizeAllreduceRMSNormOp(MoeFinalizeAllReduceRMSNormParams const& params);
} // namespace kernels::mnnvl

TRTLLM_NAMESPACE_END
Expand Down
Loading
Loading