From 0881a65acc3a3962b2c135f4b632ab31b03fc127 Mon Sep 17 00:00:00 2001 From: Christina Zhang <83400082+ChristinaZ@users.noreply.github.com> Date: Sat, 26 Jul 2025 15:06:29 +0000 Subject: [PATCH] Add topk_ids and topk_weight as inputs for moe trtllgen backend; enable post quant allgather Signed-off-by: Christina Zhang <83400082+ChristinaZ@users.noreply.github.com> --- .../kernels/customMoeRoutingKernels.cu | 54 ++- .../kernels/customMoeRoutingKernels.h | 2 +- .../blockScaleMoe/RoutingDeepSeek.cu | 238 +++++----- .../blockScaleMoe/RoutingKernel.cuh | 101 +++-- .../blockScaleMoe/RoutingKernel.h | 39 +- .../blockScaleMoe/RoutingKernelTopK.cuh | 2 +- .../blockScaleMoe/RoutingLlama4.cu | 140 ++++-- .../blockScaleMoe/RoutingRenormalize.cu | 74 ++- .../trtllmGenKernels/blockScaleMoe/runner.cu | 33 +- .../trtllmGenKernels/blockScaleMoe/runner.h | 11 +- cpp/tensorrt_llm/thop/customMoeRoutingOp.cpp | 88 ++-- cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp | 138 ++++-- cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp | 171 ++++--- .../thop/fp8PerTensorScaleMoe.cpp | 153 +++++-- cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp | 150 +++++-- .../kernels/routing/routingDeepSeekTest.cpp | 117 +++-- .../kernels/routing/routingLlama4Test.cpp | 71 ++- .../routing/routingRenormalizeTest.cpp | 52 ++- .../kernels/routing/routingTest.cpp | 42 +- .../unit_tests/kernels/routing/routingTest.h | 35 +- .../_torch/custom_ops/cpp_custom_ops.py | 10 +- .../custom_ops/trtllm_gen_custom_ops.py | 421 +++++++++++------- .../_torch/models/modeling_deepseekv3.py | 12 +- .../_torch/models/modeling_gpt_oss.py | 11 +- .../_torch/models/modeling_qwen3_moe.py | 23 +- .../modules/fused_moe/fused_moe_trtllm_gen.py | 124 +++++- .../_torch/modules/fused_moe/routing.py | 29 +- .../unittest/_torch/thop/parallel/test_moe.py | 128 ++++-- 28 files changed, 1697 insertions(+), 772 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu index eb3b958eb2d8..c7fa6d4e08ed 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu @@ -38,17 +38,18 @@ static constexpr int WARPS_PER_BLOCK = BLOCK_SIZE / WARP_SIZE; //////////////////////////////////////////////////////////////////////////////////////////////////// -template -__device__ T calcSoftmax(cg::thread_block_tile const& warp, T score, int32_t laneIdx, int32_t NumTopExperts) +template +__device__ DataType calcSoftmax( + cg::thread_block_tile const& warp, DataType score, int32_t laneIdx, int32_t NumTopExperts) { - T maxScore = T{-INFINITY}; + float maxScore = -INFINITY; if (laneIdx < NumTopExperts) { - maxScore = score >= maxScore ? score : maxScore; + maxScore = float(score) >= maxScore ? float(score) : maxScore; } - maxScore = cg::reduce(warp, maxScore, cg::greater()); + maxScore = cg::reduce(warp, maxScore, cg::greater()); - float sumScore{0.f}; + float sumScore = 0.f; float newScore; // Get the summation of scores for each token if (laneIdx < NumTopExperts) @@ -61,7 +62,7 @@ __device__ T calcSoftmax(cg::thread_block_tile const& warp, T score, if (laneIdx < NumTopExperts) { - score = static_cast(newScore / sumScore); + score = static_cast(newScore / sumScore); } return score; @@ -70,31 +71,35 @@ __device__ T calcSoftmax(cg::thread_block_tile const& warp, T score, template __device__ void calcSoftmax(cg::thread_block_tile const& warp, DataType (&scores)[VecSize]) { - DataType maxScore = DataType{-INFINITY}; - DataType sumScore = DataType{0.f}; - + // Compute in float to support half/bfloat16 inputs safely. + float maxScore = -INFINITY; + float sumScore = 0.f; // Get the max score for each token #pragma unroll for (int i = 0; i < VecSize; ++i) { - maxScore = scores[i] >= maxScore ? scores[i] : maxScore; + float si = static_cast(scores[i]); + maxScore = si >= maxScore ? si : maxScore; } - maxScore = cg::reduce(warp, maxScore, cg::greater()); + maxScore = cg::reduce(warp, maxScore, cg::greater()); // Get the summation of scores for each token #pragma unroll for (int i = 0; i < VecSize; ++i) { - scores[i] = static_cast(exp(scores[i] - maxScore)); - sumScore += scores[i]; + float si = static_cast(scores[i]); + float e = expf(si - maxScore); + scores[i] = static_cast(e); + sumScore += e; } - sumScore = cg::reduce(warp, sumScore, cg::plus()); + sumScore = cg::reduce(warp, sumScore, cg::plus()); // Normalize the scores #pragma unroll for (int i = 0; i < VecSize; ++i) { - scores[i] = static_cast(scores[i] / sumScore); + float si = static_cast(scores[i]) / sumScore; + scores[i] = static_cast(si); } } @@ -205,7 +210,7 @@ int nextPowerOfTwo(int num) break; template -void invokeRenormMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens, +void invokeCustomMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens, int64_t const numExperts, int64_t const topK, cudaStream_t const stream) { @@ -249,20 +254,25 @@ void invokeRenormMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* top } #define INSTANTIATE_RENORM_MOE_ROUTING(InputT, OutputT, IdxT, DoSoftmaxBeforeTopK) \ - template void invokeRenormMoeRouting(InputT * routerLogits, \ + template void invokeCustomMoeRouting(InputT * routerLogits, \ OutputT * topkValues, IdxT * topkIndices, int64_t const numTokens, int64_t const numExperts, \ int64_t const topK, cudaStream_t const stream); INSTANTIATE_RENORM_MOE_ROUTING(float, float, int32_t, false); INSTANTIATE_RENORM_MOE_ROUTING(half, float, int32_t, false); -#ifdef ENABLE_BF16 -INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, float, int32_t, false); -#endif - INSTANTIATE_RENORM_MOE_ROUTING(float, float, int32_t, true); INSTANTIATE_RENORM_MOE_ROUTING(half, float, int32_t, true); + #ifdef ENABLE_BF16 +INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, float, int32_t, false); +INSTANTIATE_RENORM_MOE_ROUTING(float, __nv_bfloat16, int32_t, false); +INSTANTIATE_RENORM_MOE_ROUTING(half, __nv_bfloat16, int32_t, false); +INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, int32_t, false); + INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, float, int32_t, true); +INSTANTIATE_RENORM_MOE_ROUTING(float, __nv_bfloat16, int32_t, true); +INSTANTIATE_RENORM_MOE_ROUTING(half, __nv_bfloat16, int32_t, true); +INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, int32_t, true); #endif } // namespace tensorrt_llm::kernels diff --git a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h index cfe0ae8f15e5..500889c0e5de 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h @@ -24,6 +24,6 @@ namespace tensorrt_llm::kernels { template -void invokeRenormMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens, +void invokeCustomMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens, int64_t const numExperts, int64_t const topK, cudaStream_t const stream); } // namespace tensorrt_llm::kernels diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingDeepSeek.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingDeepSeek.cu index c81a5daa4a6c..c7491e78582a 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingDeepSeek.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingDeepSeek.cu @@ -94,125 +94,129 @@ __global__ void routingMainKernel(KernelParams params) } #endif - // get our assigned thread score; each warp represents one expert group - float score = expertSelected ? static_cast(params.mPtrScores[scoreIdx]) : invalidScoreFloat; - // get the sigmoid score - // note that for invalid values, we simply use a negative value: - // sigmoig scores are always strictly positive - auto scoreSigmoid = sigmoid_accurate(score); - // write the sigmoid score to shared for later use - if (expertSelected) + if (params.mPtrScores != nullptr) { - smemScoreSigmoid[threadExpert] = scoreSigmoid; - } - // get the score with bias - // note that with invalid values, because sigmoid is < 1 and bias is -1, - // we must get a negative value, which is smaller than any valid value - auto scoreBias = float{scoreSigmoid + float{biasVal}}; - - if (expertSelected) - { - smemScoreBias[threadExpert] = scoreBias; - } - - // registers for top group score reduction - float topExpGroupScores[NumTopGroupScores]; - [[maybe_unused]] int32_t topExpGroupIdx[NumTopGroupScores]; - float topGroups[MaxNumTopGroups]; // bound of params.mNumLimitedGroups - int32_t topGroupIdx[MaxNumTopGroups]; - float expertScoreGroup[MaxNumTopGroups]; - int32_t expertIdxGroup[MaxNumTopGroups]; - float topScores[MaxNumTopExperts]; // bound of params.mTopK - int32_t topExperts[MaxNumTopExperts]; - - if constexpr (KernelParams::UseGroups) - { - topk::reduceTopK(warp, topExpGroupScores, topExpGroupIdx, scoreBias, threadExpert, - /* minValue */ invalidScoreFloat); + // get our assigned thread score; each warp represents one expert group + float score = expertSelected ? static_cast(params.mPtrScores[scoreIdx]) : invalidScoreFloat; + // get the sigmoid score + // note that for invalid values, we simply use a negative value: + // sigmoig scores are always strictly positive + auto scoreSigmoid = sigmoid_accurate(score); + // write the sigmoid score to shared for later use + if (expertSelected) + { + smemScoreSigmoid[threadExpert] = scoreSigmoid; + } + // get the score with bias + // note that with invalid values, because sigmoid is < 1 and bias is -1, + // we must get a negative value, which is smaller than any valid value + auto scoreBias = float{scoreSigmoid + float{biasVal}}; - // get the final group score and write it to shared - if (cute::elect_one_sync()) + if (expertSelected) { - auto groupScore = topExpGroupScores[0] + topExpGroupScores[1]; - smemGroupScores[warpIdx] = groupScore; + smemScoreBias[threadExpert] = scoreBias; } - } - // make group scores available to all warps - __syncthreads(); + // registers for top group score reduction + float topExpGroupScores[NumTopGroupScores]; + [[maybe_unused]] int32_t topExpGroupIdx[NumTopGroupScores]; + float topGroups[MaxNumTopGroups]; // bound of params.mNumLimitedGroups + int32_t topGroupIdx[MaxNumTopGroups]; + float expertScoreGroup[MaxNumTopGroups]; + int32_t expertIdxGroup[MaxNumTopGroups]; + float topScores[MaxNumTopExperts]; // bound of params.mTopK + int32_t topExperts[MaxNumTopExperts]; - auto localExpertExtent = params.mNumLocalExperts << params.mLocalExpertsStrideLog2; - if (warpIdx == 0) - { - // a single warp performs the selection of top groups, and goes on to select the final experts if constexpr (KernelParams::UseGroups) { - float groupScore = laneIdx < params.mNumExpertGroups ? smemGroupScores[laneIdx] : invalidScoreFloat; - - topk::reduceTopK(warp, topGroups, topGroupIdx, groupScore, laneIdx, + topk::reduceTopK(warp, topExpGroupScores, topExpGroupIdx, scoreBias, threadExpert, /* minValue */ invalidScoreFloat); - // final expert selection: get relevant indexes and scores from shared - -#pragma unroll - for (int ii = 0; ii < MaxNumTopGroups; ++ii) - { // bound of params.mNumLimitedGroups - auto groupIdx = topGroupIdx[ii]; - expertIdxGroup[ii] = groupIdx * params.mNumExpertsPerGroup + laneIdx; - // note: expertSelected implies laneIdx < params.mNumExpertsPerGroup. - // we have params.mNumExpertsPerGroup == params.mNumExperts / params.mNumExpertGroups, - // thus groupIdx <= params.mNumExpertGroups - 1 => - // groupIdx * params.mNumExpertsPerGroup <= params.mNumExperts - params.mNumExpertsPerGroup - // => expertIdxGroup[ii] < params.mNumExperts <= NumThreads, - // so the access is safe here - expertScoreGroup[ii] = groupIdx < params.mNumExpertGroups && expertSelected - ? smemScoreBias[expertIdxGroup[ii]] - : invalidScoreFloat; + // get the final group score and write it to shared + if (cute::elect_one_sync()) + { + auto groupScore = topExpGroupScores[0] + topExpGroupScores[1]; + smemGroupScores[warpIdx] = groupScore; } } - else + + // make group scores available to all warps + __syncthreads(); + + auto localExpertExtent = params.mNumLocalExperts << params.mLocalExpertsStrideLog2; + if (warpIdx == 0) { - // without groups, each thread just takes `MaxNumTopGroups` experts + // a single warp performs the selection of top groups, and goes on to select the final experts + if constexpr (KernelParams::UseGroups) + { + float groupScore = laneIdx < params.mNumExpertGroups ? smemGroupScores[laneIdx] : invalidScoreFloat; + + topk::reduceTopK(warp, topGroups, topGroupIdx, groupScore, laneIdx, + /* minValue */ invalidScoreFloat); + + // final expert selection: get relevant indexes and scores from shared #pragma unroll - for (int ii = 0; ii < MaxNumTopGroups; ++ii) - { - auto expertIdx = ii * WarpSize + laneIdx; - expertIdxGroup[ii] = expertIdx; - expertScoreGroup[ii] = expertIdx < params.mNumExperts ? smemScoreBias[expertIdx] : invalidScoreFloat; + for (int ii = 0; ii < MaxNumTopGroups; ++ii) + { // bound of params.mNumLimitedGroups + auto groupIdx = topGroupIdx[ii]; + expertIdxGroup[ii] = groupIdx * params.mNumExpertsPerGroup + laneIdx; + // note: expertSelected implies laneIdx < params.mNumExpertsPerGroup. + // we have params.mNumExpertsPerGroup == params.mNumExperts / params.mNumExpertGroups, + // thus groupIdx <= params.mNumExpertGroups - 1 => + // groupIdx * params.mNumExpertsPerGroup <= params.mNumExperts - params.mNumExpertsPerGroup + // => expertIdxGroup[ii] < params.mNumExperts <= NumThreads, + // so the access is safe here + expertScoreGroup[ii] = groupIdx < params.mNumExpertGroups && expertSelected + ? smemScoreBias[expertIdxGroup[ii]] + : invalidScoreFloat; + } } - } - - topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, expertIdxGroup, - /* minValue */ invalidScoreFloat, params.mTopK); + else + { + // without groups, each thread just takes `MaxNumTopGroups` experts - // determine our lane's expert index and write to output - int32_t expertIdx = 0; #pragma unroll - for (int ii = 0; ii < params.mTopK; ++ii) - { // bound of params.mTopK - expertIdx = laneIdx == ii ? topExperts[ii] : expertIdx; - } - // determine whether our expert is local to this GPU - auto localExpertIdx = expertIdx - params.mLocalExpertsStartIdx; - auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent - && (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; + for (int ii = 0; ii < MaxNumTopGroups; ++ii) + { + auto expertIdx = ii * WarpSize + laneIdx; + expertIdxGroup[ii] = expertIdx; + expertScoreGroup[ii] + = expertIdx < params.mNumExperts ? smemScoreBias[expertIdx] : invalidScoreFloat; + } + } - float scoreNorm = laneIdx < params.mTopK ? smemScoreSigmoid[expertIdx] : 0.F; - auto redNorm = cg::reduce(warp, scoreNorm, cg::plus{}); - auto finalScore = OutputT{scoreNorm * params.mRouteScale / redNorm}; + topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, expertIdxGroup, + /* minValue */ invalidScoreFloat, params.mTopK); - // write expert idx out already - auto idxTopK = blockIdx.x * params.mTopK + laneIdx; - if (laneIdx < params.mTopK && params.mPtrExpertIdx != nullptr) - { - PackedScoreIdx packedScore{static_cast(finalScore), static_cast(expertIdx)}; - params.mPtrExpertIdx[idxTopK] = packedScore; - } + // determine our lane's expert index and write to output + int32_t expertIdx = 0; +#pragma unroll + for (int ii = 0; ii < params.mTopK; ++ii) + { // bound of params.mTopK + expertIdx = laneIdx == ii ? topExperts[ii] : expertIdx; + } + // determine whether our expert is local to this GPU + auto localExpertIdx = expertIdx - params.mLocalExpertsStartIdx; + auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent + && (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; + + float scoreNorm = laneIdx < params.mTopK ? smemScoreSigmoid[expertIdx] : 0.F; + auto redNorm = cg::reduce(warp, scoreNorm, cg::plus{}); + auto finalScore = OutputT{scoreNorm * params.mRouteScale / redNorm}; + + // write expert idx out already + auto idxTopK = blockIdx.x * params.mTopK + laneIdx; + if (laneIdx < params.mTopK && params.mPtrTopKPacked != nullptr) + { + PackedScoreIdx packedScore{static_cast(finalScore), static_cast(expertIdx)}; + params.mPtrTopKPacked[idxTopK] = packedScore; + } - if (laneIdx < params.mTopK && params.mPtrExpertWeights != nullptr) - { - params.mPtrExpertWeights[idxTopK] = finalScore; + if (laneIdx < params.mTopK && params.mPtrTopKWeights != nullptr && params.mPtrTopKIds == nullptr) + { + params.mPtrTopKWeights[idxTopK] = finalScore; + } } } } @@ -296,7 +300,8 @@ __global__ void __launch_bounds__(NumThreads) routingIndicesCoopKernel(KernelPar // Define a lambda to avoid code duplication in both branches. auto loopBody = [&](int ii, int expandedIdx) { - int32_t expertIdx = params.mPtrExpertIdx[expandedIdx].idx; + int32_t expertIdx + = params.mPtrTopKIds != nullptr ? params.mPtrTopKIds[expandedIdx] : params.mPtrTopKPacked[expandedIdx].idx; expertIndexes[ii] = expertIdx; // check whether this expert is local to our GPU at all and ignore if not auto localExpertIdx = expertIdx - params.mLocalExpertsStartIdx; @@ -447,12 +452,17 @@ __global__ void routingIndicesCoopKernel(KernelParams params) void run(Data& data, void* stream) { - TLLM_CHECK_WITH_INFO( - data.mPtrExpertIdx != nullptr || data.mPtrPermutedIdxSize != nullptr || data.mPtrExpertWeights != nullptr, - "Routing kernel requires at least one output parameter"); + TLLM_CHECK_WITH_INFO(data.mPtrTopKPacked != nullptr || data.mPtrScores != nullptr || data.mPtrTopKIds != nullptr, + "Routing kernel requires at least one input parameter"); + if (data.mPtrTopKIds != nullptr) + { + TLLM_CHECK_WITH_INFO(data.mPtrTopKWeights != nullptr, + "When mPtrTopKIds is provided, mPtrTopKWeights must also be provided for DeepSeek routing."); + } if (data.mPtrExpandedIdxToPermutedIdx != nullptr || data.mPtrPermutedIdxToTokenIdx != nullptr) - TLLM_CHECK_WITH_INFO(data.mPtrExpertIdx != nullptr && data.mPtrPermutedIdxSize, - "If permuted index is required, `mPtrExpertIdx` is also required"); + TLLM_CHECK_WITH_INFO( + (data.mPtrTopKPacked != nullptr || data.mPtrTopKIds != nullptr) && data.mPtrPermutedIdxSize, + "If permuted index is required, `mPtrTopKPacked` or `mPtrTopKIds` is also required"); TLLM_CHECK_WITH_INFO(!data.mUseRoutingSoftmax, "Routing with softmax not implemented yet"); TLLM_CHECK_WITH_INFO(data.mNumLimitedGroups <= MaxNumTopGroups, "Routing kernel expects <= %d top groups, got %d", MaxNumTopGroups, data.mNumLimitedGroups); @@ -519,10 +529,22 @@ void run(Data& data, void* stream) // Maximum number of tokens supported by the kernel using a cooperative launch. int const maxTokensCoop = (numBlocksCoop * NumThreads * 64) / data.mTopK; - LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, - /*coopLaunch=*/false, routingMainKernel, numBlocks, NumThreads, - /*smemSize=*/0, // No dynamic smem - stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true); + + if (data.mPtrTopKIds == nullptr) + { + LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, + /*coopLaunch=*/false, routingMainKernel, numBlocks, NumThreads, + /*smemSize=*/0, // No dynamic smem + stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true); + } + else + { + // Reset the global histograms. + LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, false, routingInitExpertCounts, + (2 * data.mNumExperts - 1) / NumThreadsHist + 1, NumThreadsHist, + /*smemSize=*/0, // No dynamic smem + stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/false); + } if (data.mPtrPermutedIdxSize != nullptr) { diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.cuh b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.cuh index 92d020fd19a4..29d05e59068b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.cuh +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.cuh @@ -118,30 +118,35 @@ __device__ void initArr(int startIdx, int numElts, int stride, DataType* arr, Da template __device__ void calcSoftmax(cg::thread_block_tile const& warp, DataType (&scores)[VecSize]) { - DataType maxScore = DataType{-INFINITY}; - DataType sumScore = DataType{0.f}; - + // Compute in float to support half/bfloat16 inputs safely. + float maxScore = -INFINITY; + float sumScore = 0.f; // Get the max score for each token +#pragma unroll for (int i = 0; i < VecSize; ++i) { - maxScore = scores[i] >= maxScore ? scores[i] : maxScore; + float si = static_cast(scores[i]); + maxScore = si >= maxScore ? si : maxScore; } - maxScore = cg::reduce(warp, maxScore, cg::greater()); + maxScore = cg::reduce(warp, maxScore, cg::greater()); // Get the summation of scores for each token #pragma unroll for (int i = 0; i < VecSize; ++i) { - scores[i] = static_cast(exp(scores[i] - maxScore)); - sumScore += scores[i]; + float si = static_cast(scores[i]); + float e = expf(si - maxScore); + scores[i] = static_cast(e); + sumScore += e; } - sumScore = cg::reduce(warp, sumScore, cg::plus()); + sumScore = cg::reduce(warp, sumScore, cg::plus()); // Normalize the scores #pragma unroll for (int i = 0; i < VecSize; ++i) { - scores[i] = static_cast(scores[i] / sumScore); + float si = static_cast(scores[i]) / sumScore; + scores[i] = static_cast(si); } } @@ -232,8 +237,16 @@ __device__ void routingPermutation(KernelParams params, PackedScoreIdx TypePacked scoreIdx; if constexpr (LoadExpertIdxFromGlobal) { - scoreIdx = TypePacked{static_cast(params.mPtrExpertIdx[expandedIdx].score), - static_cast(params.mPtrExpertIdx[expandedIdx].idx)}; + if (params.mPtrTopKIds != nullptr) + { + scoreIdx = TypePacked{static_cast(params.mPtrTopKWeights[expandedIdx]), + static_cast(params.mPtrTopKIds[expandedIdx])}; + } + else + { + scoreIdx = TypePacked{static_cast(params.mPtrTopKPacked[expandedIdx].score), + static_cast(params.mPtrTopKPacked[expandedIdx].idx)}; + } } else { @@ -248,9 +261,9 @@ __device__ void routingPermutation(KernelParams params, PackedScoreIdx auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent && (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; expertOffsets[ii] = isLocalExpert ? atomicAdd(smemExpertCount + scoreIdx.idx, 1) : 0; - if (params.mPtrExpertWeights != nullptr) + if (params.mPtrTopKWeights != nullptr && params.mPtrTopKIds == nullptr) { - params.mPtrExpertWeights[expandedIdx] = OutputT{scoreIdx.score}; + params.mPtrTopKWeights[expandedIdx] = OutputT{scoreIdx.score}; } }; @@ -459,19 +472,29 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesHistogramKernel( // Define a lambda to avoid code duplication in branches. auto loopBody = [&](int expandedIdx) { - PackedScoreIdx scoreIdx = params.mPtrExpertIdx[expandedIdx]; + PackedScoreIdx scoreIdx; + int idx; + if (params.mPtrTopKIds != nullptr) + { + idx = params.mPtrTopKIds[expandedIdx]; + } + else + { + // If params.mPtrTopKIds != nullptr, we don't need to store the weights + if (params.mPtrTopKWeights != nullptr) + { + scoreIdx = params.mPtrTopKPacked[expandedIdx]; + idx = scoreIdx.idx; + params.mPtrTopKWeights[expandedIdx] = static_cast(scoreIdx.score); + } + } // check whether this expert is local to our GPU at all and ignore if not - auto localExpertIdx = scoreIdx.idx - params.mLocalExpertsStartIdx; + auto localExpertIdx = idx - params.mLocalExpertsStartIdx; auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent && (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; if (isLocalExpert) { - atomicAdd(&smemExpertCount[scoreIdx.idx], 1); - } - - if (params.mPtrExpertWeights != nullptr) - { - params.mPtrExpertWeights[expandedIdx] = static_cast(scoreIdx.score); + atomicAdd(&smemExpertCount[idx], 1); } }; @@ -625,13 +648,13 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesOffsetsKernel(Ke // Define a lambda to avoid code duplication in branches. auto loopBody = [&](int ii, int expandedIdx) { - PackedScoreIdx scoreIdx = params.mPtrExpertIdx[expandedIdx]; - expertIndexes[ii] = scoreIdx.idx; + expertIndexes[ii] + = params.mPtrTopKIds ? params.mPtrTopKIds[expandedIdx] : params.mPtrTopKPacked[expandedIdx].idx; // check whether this expert is local to our GPU at all and ignore if not - auto localExpertIdx = scoreIdx.idx - params.mLocalExpertsStartIdx; + auto localExpertIdx = expertIndexes[ii] - params.mLocalExpertsStartIdx; auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent && (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; - expertOffsets[ii] = isLocalExpert ? atomicAdd(smemExpertCount + scoreIdx.idx, 1) : 0; + expertOffsets[ii] = isLocalExpert ? atomicAdd(smemExpertCount + expertIndexes[ii], 1) : 0; }; // For all tiles but the last, all indices are in bounds. @@ -767,5 +790,33 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesOffsetsKernel(Ke #endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) } +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void __launch_bounds__(NumThreadsHist) routingInitExpertCounts(KernelParams params) +{ + // initialize the mPtrExpertCounts + int32_t expertCountsNum = 2 * params.mNumExperts; + int32_t globalThreadIdx = blockIdx.x * NumThreadsHist + threadIdx.x; + int32_t globalThreadStride = gridDim.x * NumThreadsHist; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // Wait on primary grid. + if constexpr (KernelParams::UsePdl) + { + cudaGridDependencySynchronize(); + } +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + + initArr(globalThreadIdx, expertCountsNum, globalThreadStride, params.mPtrExpertCounts, 0); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // Wait on primary grid. + if constexpr (KernelParams::UsePdl) + { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) +} } // namespace routing } // namespace moe::dev diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.h index d003e7473bc6..dfa537ac99d1 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.h @@ -58,17 +58,26 @@ struct DataBase // Note: this array (mPtrPermutedIdxToTokenIdx) is uninitialized // Any out-of-bounds values are undefined. int32_t* mPtrPermutedIdxToTokenIdx{nullptr}; + + // optional: if `nullptr`, it is not filled + // dim: [mNumTokens, mTopK] + // When mPtrTopKIds is provided, mPtrTopKWeights must be also provided as inputs. + // Otherwise, mPtrTopKWeights is the output scores of the topK experts. + void* mPtrTopKWeights{nullptr}; // optional: if `nullptr`, it is not filled // dim: [mNumTokens, mTopK] - void* mPtrExpertWeights{nullptr}; + // mPtrTopKIds[i] is the index of the expert for the i-th token in the top-k experts + // Together with mPtrTopKWeights, they form the top-k experts for each token + int32_t* mPtrTopKIds{nullptr}; + // optional: if `nullptr`, scores are used directly as input. // If it is given, it must represent a packed value s.t. the most significant // 16/32 bits represent the score without sigmoid activation and // the least significant 16 bits represent the index of the chosen expert (unsigned). // note: this is required if the number of tokens is large. // dim: [mNumTokens, mTopK] - void* mPtrExpertIdx{nullptr}; - // optional: if `nullptr`, `mPtrExpertIdx` must be provided. + void* mPtrTopKPacked{nullptr}; + // optional: if `nullptr`, `mPtrTopKPacked` must be provided. // If it is given, it represents the scores without sigmoid activation for // each token and expert. // note: if it is provided, we always re-compute the top1 scores @@ -111,7 +120,8 @@ struct KernelParamsBase int32_t* mPtrCtaIdxXyToBatchIdx = nullptr; int32_t* mPtrCtaIdxXyToMnLimit = nullptr; int32_t* mPtrNumNonExitingCtas = nullptr; - OutputT* mPtrExpertWeights = nullptr; + OutputT* mPtrTopKWeights = nullptr; + int32_t* mPtrTopKIds = nullptr; InputT const* mPtrScores = nullptr; // Public scalar members @@ -134,7 +144,8 @@ struct KernelParamsBase mPtrCtaIdxXyToBatchIdx = data.mPtrCtaIdxXyToBatchIdx; mPtrCtaIdxXyToMnLimit = data.mPtrCtaIdxXyToMnLimit; mPtrNumNonExitingCtas = data.mPtrNumNonExitingCtas; - mPtrExpertWeights = static_cast(data.mPtrExpertWeights); + mPtrTopKWeights = static_cast(data.mPtrTopKWeights); + mPtrTopKIds = static_cast(data.mPtrTopKIds); mPtrScores = (InputT const*) data.mPtrScores; mNumTokens = data.mNumTokens; @@ -176,10 +187,10 @@ struct KernelParams : public KernelParamsBase static constexpr bool UseGroups = UseGroups_; - PackedScoreIdx* mPtrExpertIdx = nullptr; + PackedScoreIdx* mPtrTopKPacked = nullptr; - // OutputT* mPtrExpertWeightsFull = nullptr; - // Note: this variable(mPtrExpertWeightsFull) might need to be added back for the low-latency kernels for MoE in + // OutputT* mPtrTopKWeightsFull = nullptr; + // Note: this variable(mPtrTopKWeightsFull) might need to be added back for the low-latency kernels for MoE in // tllm-gen in the future OutputT const* mPtrRoutingBias = nullptr; @@ -196,9 +207,9 @@ struct KernelParams : public KernelParamsBase KernelParams params; params.setBaseParams(data); - params.mPtrExpertIdx = (PackedScoreIdx*) data.mPtrExpertIdx; + params.mPtrTopKPacked = (PackedScoreIdx*) data.mPtrTopKPacked; - // params.mPtrExpertWeightsFull = static_cast(data.mPtrExpertWeightsFull); + // params.mPtrTopKWeightsFull = static_cast(data.mPtrTopKWeightsFull); params.mPtrRoutingBias = static_cast(data.mPtrRoutingBias); params.mNumExpertGroups = data.mNumExpertGroups; @@ -233,7 +244,7 @@ struct KernelParams : public KernelParamsBase using InputT = InputT_; using OutputT = OutputT_; - PackedScoreIdx* mPtrExpertIdx = nullptr; + PackedScoreIdx* mPtrTopKPacked = nullptr; int32_t mTopK; @@ -242,7 +253,7 @@ struct KernelParams : public KernelParamsBase KernelParams params; params.setBaseParams(data); - params.mPtrExpertIdx = (PackedScoreIdx*) data.mPtrExpertIdx; + params.mPtrTopKPacked = (PackedScoreIdx*) data.mPtrTopKPacked; params.mTopK = data.mTopK; return params; } @@ -276,7 +287,7 @@ struct KernelParams : public KernelParamsBase static constexpr bool DoSoftmaxBeforeTopK = DoSoftmaxBeforeTopK_; - PackedScoreIdx* mPtrExpertIdx = nullptr; + PackedScoreIdx* mPtrTopKPacked = nullptr; int32_t mTopK = 0; @@ -287,7 +298,7 @@ struct KernelParams : public KernelParamsBase KernelParams params; params.setBaseParams(data); - params.mPtrExpertIdx = (PackedScoreIdx*) data.mPtrExpertIdx; + params.mPtrTopKPacked = (PackedScoreIdx*) data.mPtrTopKPacked; params.mNormTopkProb = data.mNormTopkProb; params.mTopK = data.mTopK; return params; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernelTopK.cuh b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernelTopK.cuh index e6d04c39757a..fa7c6e250662 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernelTopK.cuh +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernelTopK.cuh @@ -56,7 +56,7 @@ struct TopKRedType { auto valueBits = cub::Traits::TwiddleIn(reinterpret_cast::UnsignedBits&>(val)); - TypeCmp compactTmp = reinterpret_cast(valueBits); + TypeCmp compactTmp = valueBits; compactTmp = (compactTmp << moveBits) | (0xFFFF & (maxIdx - idx)); // Use 65535 minus idx to give higher priority to elements with smaller indices. return compactTmp; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingLlama4.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingLlama4.cu index 5c3989203908..51e7b72eff5f 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingLlama4.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingLlama4.cu @@ -44,7 +44,7 @@ __forceinline__ __device__ void routingTopKExperts(cg::thread_block_tile; // Non-vectorized loading: directly access ptrScores with expertIdx @@ -111,7 +111,7 @@ __global__ void __launch_bounds__(WarpSize) routingIndicesWarpKernel(KernelParam } #endif - if (params.mPtrScores != nullptr) + if (params.mPtrScores != nullptr && params.mPtrTopKIds == nullptr) { // if we use `mPtrScores` as input, we need to perform the top-1 reduction // for each token, we load the scores then use `reduceTopK` for this. @@ -134,33 +134,49 @@ __global__ void __launch_bounds__(WarpSize) routingIndicesWarpKernel(KernelParam smemExpertTokenCountFull[tokenIdx][warpMaxExpertIdx[0] / ExpertsPerThread] = expertTokenCount; // we also compute the final score here and write it out if required auto finalScore = OutputT{sigmoid_accurate(float{warpMaxScore[0]})}; - if (params.mPtrExpertWeights != nullptr) + if (params.mPtrTopKWeights != nullptr) { - params.mPtrExpertWeights[tokenIdx] = finalScore; + params.mPtrTopKWeights[tokenIdx] = finalScore; } } } } else { - // if we do not have `mPtrScores` as input, we expect that `mPtrExpertWeights` - // contains the top-1 packed score and index already. - // Each thread represents a token here, and we extract the relevant score - // The assumption is that the #tokens is limited by warp-size + // if we do not have `mPtrScores` as input, we expect that `params.mPtrTopKPacked` or `params.mPtrTopKIds` and + // `params.mPtrTopKWeights` contains the top-1 packed score and index already. Each thread represents a token + // here, and we extract the relevant score The assumption is that the #tokens is limited by warp-size static_assert(WarpKernelMaxNumTokens <= WarpSize); - TypePacked scoreIdx = threadIdx.x < params.mNumTokens ? params.mPtrExpertIdx[threadIdx.x] : TypePacked{}; + TypePacked scoreIdx = TypePacked{}; + if (params.mPtrTopKIds != nullptr) + { + if (threadIdx.x < params.mNumTokens) + { + scoreIdx = TypePacked{static_cast(params.mPtrTopKWeights[threadIdx.x]), + static_cast(params.mPtrTopKIds[threadIdx.x])}; + } + } + else + { + if (threadIdx.x < params.mNumTokens) + { + scoreIdx = TypePacked{static_cast(params.mPtrTopKPacked[threadIdx.x].score), + static_cast(params.mPtrTopKPacked[threadIdx.x].idx)}; + if (params.mPtrTopKWeights != nullptr) + { + // we also compute the final score here and write it out if required + auto finalScore = OutputT{sigmoid_accurate(float{scoreIdx.score})}; + params.mPtrTopKWeights[threadIdx.x] = finalScore; + } + } + } + int32_t expertTokenCount = 0; setBits(expertTokenCount, 1, scoreIdx.idx % ExpertsPerThread); if (threadIdx.x < params.mNumTokens) { smemExpertTokenCountFull[threadIdx.x][scoreIdx.idx / ExpertsPerThread] = expertTokenCount; } - // we also compute the final score here and write it out if required - auto finalScore = OutputT{sigmoid_accurate(float{scoreIdx.score})}; - if (params.mPtrExpertWeights != nullptr && threadIdx.x < params.mNumTokens) - { - params.mPtrExpertWeights[threadIdx.x] = finalScore; - } } // make the full table available to all threads @@ -319,7 +335,16 @@ __global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(Nu cudaGridDependencySynchronize(); } - if (params.mPtrScores != nullptr) + if (params.mPtrTopKIds != nullptr) + { + if (validToken) + { + TypePacked packedScore{static_cast(params.mPtrTopKWeights[warpTokenIdx]), + static_cast(params.mPtrTopKIds[warpTokenIdx])}; + smemPackedScoreIdx[warpIdx] = packedScore; + } + } + else if (params.mPtrScores != nullptr) { // in this case, each warp represents a token // we then exchange all token max scores, s.t. afterwards, each thread @@ -338,13 +363,29 @@ __global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(Nu smemPackedScoreIdx[warpIdx] = packedScore; } } - // make packed scores available to all threads in cluster - __cluster_barrier_arrive(); - __cluster_barrier_wait(); + } + else + { + if (validToken) + { + smemPackedScoreIdx[warpIdx] = params.mPtrTopKPacked[warpTokenIdx]; + } } - routingPermutation(params, smemPackedScoreIdx, warpIdx, clusterBlockRank); + // make packed scores available to all threads in cluster + __cluster_barrier_arrive(); + __cluster_barrier_wait(); + + if (params.mPtrTopKIds != nullptr || params.mPtrScores != nullptr) + { + routingPermutation(params, smemPackedScoreIdx, warpIdx, clusterBlockRank); + } + else + { + routingPermutation(params, smemPackedScoreIdx, warpIdx, clusterBlockRank); + } } #else __global__ void routingIndicesClusterKernel(KernelParams params) @@ -376,8 +417,8 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesHistogramScoresK // initialize the mPtrExpertCounts int32_t expertCountsNum = 2 * params.mNumExperts; - int32_t globalThreadIdx = blockIdx.x * NumThreads + threadIdx.x; - int32_t globalThreadStride = gridDim.x * NumThreads; + int32_t globalThreadIdx = blockIdx.x * NumThreadsHist + threadIdx.x; + int32_t globalThreadStride = gridDim.x * NumThreadsHist; initArr(globalThreadIdx, expertCountsNum, globalThreadStride, params.mPtrExpertCounts, 0); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -397,14 +438,33 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesHistogramScoresK int32_t warpMaxExpertIdx[MaxNumTopExperts]; InputT warpMaxScore[MaxNumTopExperts]; - routingTopKExperts( - warp, warpMaxScore, warpMaxExpertIdx, laneIdx, params.mNumExperts, params.mPtrScores + scoreOffset); + if (params.mPtrTopKIds != nullptr) + { + if (laneIdx < MaxNumTopExperts) + { + warpMaxExpertIdx[laneIdx] = params.mPtrTopKIds[tokenIdx]; + warpMaxScore[laneIdx] = static_cast(params.mPtrTopKWeights[tokenIdx]); + } + } + else if (params.mPtrScores != nullptr) + { + routingTopKExperts( + warp, warpMaxScore, warpMaxExpertIdx, laneIdx, params.mNumExperts, params.mPtrScores + scoreOffset); + } + else + { + if (laneIdx < MaxNumTopExperts) + { + warpMaxExpertIdx[laneIdx] = params.mPtrTopKPacked[tokenIdx].idx; + warpMaxScore[laneIdx] = params.mPtrTopKPacked[tokenIdx].score; + } + } if (cute::elect_one_sync()) { auto finalScore = OutputT{sigmoid_accurate(float{warpMaxScore[0]})}; TypePacked packedScore{finalScore, static_cast(warpMaxExpertIdx[0])}; - params.mPtrExpertIdx[tokenIdx] = packedScore; + params.mPtrTopKPacked[tokenIdx] = packedScore; } } } @@ -413,15 +473,20 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesHistogramScoresK void run(Data const& data, void* stream) { - TLLM_CHECK_WITH_INFO(data.mPtrExpertIdx != nullptr || data.mPtrScores != nullptr, + TLLM_CHECK_WITH_INFO(data.mPtrTopKPacked != nullptr || data.mPtrScores != nullptr || data.mPtrTopKIds != nullptr, "Routing kernel requires at least one input parameter"); + if (data.mPtrTopKIds != nullptr) + { + TLLM_CHECK_WITH_INFO(data.mPtrTopKWeights != nullptr, + "When mPtrTopKIds is provided, mPtrTopKWeights must also be provided for Llama4 routing."); + } TLLM_CHECK_WITH_INFO(data.mPtrPermutedIdxSize != nullptr && data.mPtrCtaIdxXyToBatchIdx != nullptr && data.mPtrCtaIdxXyToMnLimit != nullptr && data.mPtrNumNonExitingCtas != nullptr, "Llama4 routing kernel expects permuted idx and grouped Gemm launch config buffers"); TLLM_CHECK_WITH_INFO(data.mTopK <= MaxNumTopExperts, "Routing kernel expects topK experts <= %d, got %d", MaxNumTopExperts, data.mTopK); - TLLM_CHECK_WITH_INFO(data.mNumExperts <= MaxNumExperts, - "Routing kernel expects #experts %d to be at most max #experts %d", data.mNumExperts, MaxNumExperts); + TLLM_CHECK_WITH_INFO(data.mNumExperts <= MaxNumExperts, "Routing kernel expects #experts %d to be no more than %d", + data.mNumExperts, MaxNumExperts); static_assert(MaxNumExperts <= NumThreads, "#experts must be bounded by #threads"); static_assert(MaxNumExperts <= NumThreadsHist, "#experts must be bounded by #threads"); TLLM_CHECK_WITH_INFO( @@ -430,12 +495,13 @@ void run(Data const& data, void* stream) bool const useSingleWarp = (data.mPtrScores == nullptr && data.mNumTokens <= WarpKernelMaxNumTokens) || data.mNumTokens < WarpKernelMaxNumTokens; - bool const useSingleCluster - = data.mNumTokens <= (data.mPtrScores != nullptr ? MaxNumTokensSingleClusterScores : MaxNumTokensSingleCluster); + bool const useSingleCluster = data.mNumTokens <= ((data.mPtrScores != nullptr || data.mPtrTopKIds != nullptr) + ? MaxNumTokensSingleClusterScores + : MaxNumTokensSingleCluster); if (!useSingleCluster) { - TLLM_CHECK_WITH_INFO( - data.mPtrExpertIdx != nullptr, "When #tokens is large, `mPtrExpertIdx` is a required input."); + TLLM_CHECK_WITH_INFO((data.mPtrTopKPacked != nullptr || data.mPtrTopKIds != nullptr), + "When #tokens is large, `mPtrTopKPacked` or `mPtrTopKIds` is a required input."); TLLM_CHECK_WITH_INFO( data.mPtrExpertCounts != nullptr, "When #tokens is large, `mPtrExpertCounts` is a required input."); } @@ -469,7 +535,7 @@ void run(Data const& data, void* stream) int const numBlocksOffsets = std::min((expandedIdxSize + offsetEltsPerBlock - 1) / offsetEltsPerBlock, maxNumBlocks); - if (data.mPtrScores != nullptr) + if (data.mPtrScores != nullptr && data.mPtrTopKIds == nullptr) { LAUNCH_ROUTING(data, /*coopLaunch=*/false, routingIndicesHistogramScoresKernel, maxNumBlocks, NumThreadsHist, @@ -479,8 +545,10 @@ void run(Data const& data, void* stream) else { // Reset the global histograms. - TLLM_CUDA_CHECK(cudaMemsetAsync(data.mPtrExpertCounts, 0, - static_cast(2 * NumThreads) * sizeof(int32_t), (cudaStream_t) stream)); + LAUNCH_ROUTING(data, false, routingInitExpertCounts, (2 * data.mNumExperts - 1) / NumThreadsHist + 1, + NumThreadsHist, + /*smemSize=*/0, // No dynamic smem + stream); } LAUNCH_ROUTING(data, /*coopLaunch=*/false, routingIndicesHistogramKernel, numBlocksHistogram, NumThreadsHist, diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu index f6364e0cc9a1..b5a6efe6d710 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingRenormalize.cu @@ -97,7 +97,7 @@ __global__ void __launch_bounds__(NumThreadsSingleBlock) routingIndicesBlockKern __shared__ int8_t __attribute((aligned(128))) smemOffset[totalExpertCounts]; __shared__ int8_t __attribute((aligned(128))) smemKIdx[totalExpertCounts]; - using Scan = cub::BlockScan; + using Scan = cub::BlockScan; __shared__ typename Scan::TempStorage tempStorage; auto block = cg::this_thread_block(); @@ -118,7 +118,18 @@ __global__ void __launch_bounds__(NumThreadsSingleBlock) routingIndicesBlockKern } #endif - if (params.mPtrScores != nullptr) + if (params.mPtrTopKIds != nullptr) + { + if (validToken) + { + if (laneIdx < params.mTopK) + { + int offset = warpIdx * MaxNumExperts + params.mPtrTopKIds[warpIdx * params.mTopK + laneIdx]; + smemKIdx[offset] = static_cast(laneIdx); + } + } + } + else if (params.mPtrScores != nullptr) { // in this case, each warp represents a token BaseType score[VecSize]; @@ -138,9 +149,9 @@ __global__ void __launch_bounds__(NumThreadsSingleBlock) routingIndicesBlockKern { int offset = warpIdx * MaxNumExperts + warpTopKExpertIdx[laneIdx]; smemKIdx[offset] = static_cast(laneIdx); - if (params.mPtrExpertWeights != nullptr) + if (params.mPtrTopKWeights != nullptr) { - params.mPtrExpertWeights[warpIdx * params.mTopK + laneIdx] = OutputT{warpTopKScore[laneIdx]}; + params.mPtrTopKWeights[warpIdx * params.mTopK + laneIdx] = OutputT{warpTopKScore[laneIdx]}; } } } // end if (validToken) @@ -283,14 +294,22 @@ __global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(Nu = TypePacked{warpTopKScore[laneIdx], static_cast(warpTopKExpertIdx[laneIdx])}; } } // end if (validToken) - - // make packed scores available to all threads in cluster - __cluster_barrier_arrive(); - __cluster_barrier_wait(); } - routingPermutation(params, smemPackedScoreIdx, warpIdx, clusterBlockRank); + // make packed scores available to all threads in cluster + __cluster_barrier_arrive(); + __cluster_barrier_wait(); + + if (params.mPtrScores != nullptr) + { + routingPermutation(params, smemPackedScoreIdx, warpIdx, clusterBlockRank); + } + else + { + routingPermutation(params, smemPackedScoreIdx, warpIdx, clusterBlockRank); + } } #else __global__ void __launch_bounds__(NumThreads) routingIndicesClusterKernel(KernelParams /* params */) @@ -330,8 +349,8 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesHistogramScoresK // initialize the mPtrExpertCounts int32_t expertCountsNum = 2 * params.mNumExperts; - int32_t globalThreadIdx = blockIdx.x * NumThreads + threadIdx.x; - int32_t globalThreadStride = gridDim.x * NumThreads; + int32_t globalThreadIdx = blockIdx.x * NumThreadsHist + threadIdx.x; + int32_t globalThreadStride = gridDim.x * NumThreadsHist; initArr(globalThreadIdx, expertCountsNum, globalThreadStride, params.mPtrExpertCounts, 0); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -360,7 +379,7 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesHistogramScoresK { PackedScoreIdx packedScore{ static_cast(warpTopKScore[laneIdx]), static_cast(warpTopKExpertIdx[laneIdx])}; - params.mPtrExpertIdx[tokenIdx * params.mTopK + laneIdx] = packedScore; + params.mPtrTopKPacked[tokenIdx * params.mTopK + laneIdx] = packedScore; } } } @@ -369,15 +388,20 @@ __global__ void __launch_bounds__(NumThreadsHist) routingIndicesHistogramScoresK void run(Data const& data, void* stream) { - TLLM_CHECK_WITH_INFO(data.mPtrExpertIdx != nullptr || data.mPtrScores != nullptr, + TLLM_CHECK_WITH_INFO(data.mPtrTopKPacked != nullptr || data.mPtrScores != nullptr || data.mPtrTopKIds != nullptr, "Routing kernel requires at least one input parameter"); + if (data.mPtrTopKIds != nullptr) + { + TLLM_CHECK_WITH_INFO(data.mPtrTopKWeights != nullptr, + "When mPtrTopKIds is provided, mPtrTopKWeights must also be provided for Renormalize routing."); + } TLLM_CHECK_WITH_INFO(data.mPtrPermutedIdxSize != nullptr && data.mPtrCtaIdxXyToBatchIdx != nullptr && data.mPtrCtaIdxXyToMnLimit != nullptr && data.mPtrNumNonExitingCtas != nullptr, "Llama4 routing kernel expects permuted idx and grouped Gemm launch config buffers"); TLLM_CHECK_WITH_INFO(data.mTopK <= MaxNumTopExperts, "Routing kernel expects topK experts <= %d, got %d", MaxNumTopExperts, data.mTopK); - TLLM_CHECK_WITH_INFO(data.mNumExperts <= MaxNumExperts, - "Routing kernel expects #experts %d to be at most max #experts %d", data.mNumExperts, MaxNumExperts); + TLLM_CHECK_WITH_INFO(data.mNumExperts <= MaxNumExperts, "Routing kernel expects #experts %d to be no more than %d", + data.mNumExperts, MaxNumExperts); static_assert(MaxNumExperts <= NumThreads, "#experts must be bounded by #threads"); static_assert(MaxNumExperts <= NumThreadsHist, "#experts must be bounded by #threads"); TLLM_CHECK_WITH_INFO( @@ -385,13 +409,15 @@ void run(Data const& data, void* stream) TLLM_CHECK_WITH_INFO(data.mPaddingLog2 < 8, "Routing kernel expects padding log2 < 8, got %d", data.mPaddingLog2); bool const useSingleBlock = data.mNumTokens <= BlockKernelMaxNumTokens; - bool const useSingleCluster - = data.mNumTokens <= (data.mPtrScores != nullptr ? MaxNumTokensSingleClusterScores : MaxNumTokensSingleCluster); + + bool const useSingleCluster = data.mNumTokens <= ((data.mPtrScores != nullptr || data.mPtrTopKIds != nullptr) + ? MaxNumTokensSingleClusterScores + : MaxNumTokensSingleCluster); if (!useSingleCluster && !useSingleBlock) { - TLLM_CHECK_WITH_INFO( - data.mPtrExpertIdx != nullptr, "When #tokens is large, `mPtrExpertIdx` is a required input."); + TLLM_CHECK_WITH_INFO((data.mPtrTopKPacked != nullptr || data.mPtrTopKIds != nullptr), + "When #tokens is large, `mPtrTopKPacked` or `mPtrTopKIds` is a required input."); TLLM_CHECK_WITH_INFO( data.mPtrExpertCounts != nullptr, "When #tokens is large, `mPtrExpertCounts` is a required input."); } @@ -425,7 +451,7 @@ void run(Data const& data, void* stream) int const numBlocksOffsets = std::min((expandedIdxSize + offsetEltsPerBlock - 1) / offsetEltsPerBlock, maxNumBlocks); - if (data.mPtrScores != nullptr) + if (data.mPtrScores != nullptr && data.mPtrTopKIds == nullptr) { LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, false, routingIndicesHistogramScoresKernel, maxNumBlocks, NumThreadsHist, @@ -435,8 +461,10 @@ void run(Data const& data, void* stream) else { // Reset the global histograms. - TLLM_CUDA_CHECK(cudaMemsetAsync(data.mPtrExpertCounts, 0, - static_cast(2 * NumThreads) * sizeof(int32_t), (cudaStream_t) stream)); + LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, false, routingInitExpertCounts, + (2 * data.mNumExperts - 1) / NumThreadsHist + 1, NumThreadsHist, + /*smemSize=*/0, // No dynamic smem + stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false); } LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, false, routingIndicesHistogramKernel, numBlocksHistogram, NumThreadsHist, /*smemSize=*/0, // No dynamic smem diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu index 1d267a295551..f8ef50d31600 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu @@ -60,9 +60,9 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 int32_t nGroup, int32_t topkGroup, int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, int32_t* routingExpertIndexes, int32_t* expertCountHistogram, int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx, - void* expertWeights, int32_t* numTokensPerExpert, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, - int32_t* numNonExitingCtas, btg::Dtype dtypeElt, bool useRoutingScalesOnInput, bool useDeepSeekFp8, - RoutingMethodType routingMethodType, cudaStream_t stream) + void* expertWeights, int32_t* expertIds, int32_t* numTokensPerExpert, int32_t* ctaIdxXyToBatchIdx, + int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, btg::Dtype dtypeElt, bool useRoutingScalesOnInput, + bool useDeepSeekFp8, RoutingMethodType routingMethodType, cudaStream_t stream) { if (routingMethodType == RoutingMethodType::DeepSeekV3) { @@ -73,12 +73,12 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 routingData.mUsePdl = true; // output: - routingData.mPtrExpertIdx = routingExpertIndexes; + routingData.mPtrTopKPacked = routingExpertIndexes; routingData.mPtrExpertCounts = expertCountHistogram; routingData.mPtrPermutedIdxSize = permutedIdxSize; routingData.mPtrExpandedIdxToPermutedIdx = expandedIdxToPermutedIdx; routingData.mPtrPermutedIdxToTokenIdx = permutedIdxToTokenIdx; - routingData.mPtrExpertWeights = expertWeights; + routingData.mPtrTopKWeights = expertWeights; routingData.mPtrCtaIdxXyToBatchIdx = ctaIdxXyToBatchIdx; routingData.mPtrCtaIdxXyToMnLimit = ctaIdxXyToMnLimit; @@ -86,7 +86,9 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 // input: routingData.mPtrRoutingBias = routingBias; - routingData.mPtrScores = reinterpret_cast(routingLogits); + // Pass-through raw pointer; kernels will cast to the proper InputT based on routing method + routingData.mPtrScores = expertIds == nullptr ? routingLogits : nullptr; + routingData.mPtrTopKIds = expertIds; routingData.mNumTokens = numTokens; routingData.mNumExperts = numExperts; routingData.mNumExpertGroups = nGroup; @@ -112,12 +114,12 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 routingData.mUsePdl = true; // output: - routingData.mPtrExpertIdx = routingExpertIndexes; + routingData.mPtrTopKPacked = routingExpertIndexes; routingData.mPtrExpertCounts = expertCountHistogram; routingData.mPtrPermutedIdxSize = permutedIdxSize; routingData.mPtrExpandedIdxToPermutedIdx = expandedIdxToPermutedIdx; routingData.mPtrPermutedIdxToTokenIdx = permutedIdxToTokenIdx; - routingData.mPtrExpertWeights = expertWeights; + routingData.mPtrTopKWeights = expertWeights; routingData.mPtrCtaIdxXyToBatchIdx = ctaIdxXyToBatchIdx; routingData.mPtrCtaIdxXyToMnLimit = ctaIdxXyToMnLimit; @@ -127,7 +129,10 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 // input: // routingData.mPtrRoutingWeights = args.mRoutingWeights; // routing weights (don't need if not using gemm) // routingData.mPtrRoutingBias = routingBias; - routingData.mPtrScores = routingLogits; + + // Pass-through raw pointer; kernels will cast to the proper InputT based on routing method + routingData.mPtrScores = expertIds == nullptr ? routingLogits : nullptr; + routingData.mPtrTopKIds = expertIds; // routingData.mPtrIn = args.mInputActs; routingData.mNumTokens = numTokens; // routingData.mHiddenDim = args.mHiddenDim; @@ -158,18 +163,18 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3 routingData.mDoSoftmaxBeforeTopK = routingMethodType == RoutingMethodType::RenormalizeNaive; routingData.mNormTopkProb = routingMethodType == RoutingMethodType::RenormalizeNaive; - routingData.mPtrScores = routingLogits; - + // Pass-through raw pointer; kernels will cast to the proper InputT based on routing method + routingData.mPtrScores = expertIds == nullptr ? routingLogits : nullptr; // // Outputs // - routingData.mPtrExpertIdx = routingExpertIndexes; + routingData.mPtrTopKPacked = routingExpertIndexes; routingData.mPtrExpertCounts = expertCountHistogram; routingData.mPtrPermutedIdxSize = permutedIdxSize; routingData.mPtrExpandedIdxToPermutedIdx = expandedIdxToPermutedIdx; routingData.mPtrPermutedIdxToTokenIdx = permutedIdxToTokenIdx; - routingData.mPtrExpertWeights = expertWeights; - + routingData.mPtrTopKWeights = expertWeights; + routingData.mPtrTopKIds = expertIds; // // Grouped Gemm Launch Config Buffers // diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h index a1c7113aa2bb..94c7e4d22c56 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h @@ -106,9 +106,10 @@ class Runner int32_t nGroups, int32_t topkGroups, int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor, int32_t* routingExpertIndexes, int32_t* expertCountHistogram, int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, - int32_t* permutedIdxToTokenIdx, void* expertWeights, int32_t* numTokensPerExpert, int32_t* ctaIdxXyToBatchIdx, - int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, batchedGemm::trtllm::gen::Dtype dtypeElt, - bool useRoutingScalesOnInput, bool useDeepSeekFp8, RoutingMethodType routingMethodType, cudaStream_t stream); + int32_t* permutedIdxToTokenIdx, void* expertWeights, int32_t* expertIds, int32_t* numTokensPerExpert, + int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, + batchedGemm::trtllm::gen::Dtype dtypeElt, bool useRoutingScalesOnInput, bool useDeepSeekFp8, + RoutingMethodType routingMethodType, cudaStream_t stream); private: int32_t mTileTokensDim; @@ -198,6 +199,10 @@ struct MoERunnerArgs // and [num_tokens, hidden_size/16] in float for e2m1 void* hidden_states_scale = nullptr; + // Optional inputs: + void* topk_weights = nullptr; // [num_tokens, top_k] with quantized weights + int32_t* topk_ids = nullptr; // [num_tokens, top_k] with expert ids in int32_t + // Gemm input: void* gemm1_weights = nullptr; void* gemm1_weights_scale = nullptr; diff --git a/cpp/tensorrt_llm/thop/customMoeRoutingOp.cpp b/cpp/tensorrt_llm/thop/customMoeRoutingOp.cpp index 814fdf87c3bd..81fb4acf9cff 100644 --- a/cpp/tensorrt_llm/thop/customMoeRoutingOp.cpp +++ b/cpp/tensorrt_llm/thop/customMoeRoutingOp.cpp @@ -24,9 +24,9 @@ namespace tk = tensorrt_llm::kernels; namespace torch_ext { - template -std::tuple custom_moe_routing_op(th::Tensor const& router_logits, int64_t topk) +std::tuple custom_moe_routing_op( + th::Tensor const& router_logits, int64_t topk, c10::optional output_dtype) { auto data_type = router_logits.scalar_type(); auto input_size = router_logits.sizes(); @@ -36,33 +36,69 @@ std::tuple custom_moe_routing_op(th::Tensor const& route TORCH_CHECK(topk <= 8, "topk should be smaller than or equal to 8 for now"); //@todo: remove this restriction later TORCH_CHECK(num_experts <= 128, "expert number should be smaller than or equal to 128 for now"); - th::Tensor topk_values = th::empty({num_tokens, topk}, th::dtype(torch::kFloat32).device(torch::kCUDA)); - th::Tensor topk_indices = th::empty({num_tokens, topk}, th::dtype(torch::kInt32).device(torch::kCUDA)); + // Determine output data type + at::ScalarType topk_values_dtype = output_dtype.value_or(torch::kFloat32); + TORCH_CHECK(topk_values_dtype == torch::kFloat32 || topk_values_dtype == torch::kBFloat16, + "output_dtype must be float32 or bfloat16"); + + auto opts = router_logits.options(); + th::Tensor topk_values = th::empty({num_tokens, topk}, opts.dtype(topk_values_dtype)); + th::Tensor topk_indices = th::empty({num_tokens, topk}, opts.dtype(torch::kInt32)); auto stream = at::cuda::getCurrentCUDAStream(router_logits.get_device()); switch (data_type) { case torch::kFloat32: - // Handle Float32 - tk::invokeRenormMoeRouting( - reinterpret_cast(router_logits.mutable_data_ptr()), - reinterpret_cast(topk_values.mutable_data_ptr()), - reinterpret_cast(topk_indices.mutable_data_ptr()), num_tokens, num_experts, topk, stream); + // Handle Float32 input + if (topk_values_dtype == torch::kFloat32) + { + tk::invokeCustomMoeRouting( + reinterpret_cast(router_logits.mutable_data_ptr()), + reinterpret_cast(topk_values.mutable_data_ptr()), topk_indices.data_ptr(), num_tokens, + num_experts, topk, stream); + } + else + { // bfloat16 output + tk::invokeCustomMoeRouting( + reinterpret_cast(router_logits.mutable_data_ptr()), + reinterpret_cast<__nv_bfloat16*>(topk_values.mutable_data_ptr()), topk_indices.data_ptr(), + num_tokens, num_experts, topk, stream); + } break; case torch::kBFloat16: - // Handle BFloat16 - tk::invokeRenormMoeRouting<__nv_bfloat16, float, int32_t, DoSoftmaxBeforeTopK>( - reinterpret_cast<__nv_bfloat16*>(router_logits.mutable_data_ptr()), - reinterpret_cast(topk_values.mutable_data_ptr()), - reinterpret_cast(topk_indices.mutable_data_ptr()), num_tokens, num_experts, topk, stream); + // Handle BFloat16 input + if (topk_values_dtype == torch::kFloat32) + { + tk::invokeCustomMoeRouting<__nv_bfloat16, float, int32_t, DoSoftmaxBeforeTopK>( + reinterpret_cast<__nv_bfloat16*>(router_logits.mutable_data_ptr()), + reinterpret_cast(topk_values.mutable_data_ptr()), topk_indices.data_ptr(), num_tokens, + num_experts, topk, stream); + } + else + { // bfloat16 output + tk::invokeCustomMoeRouting<__nv_bfloat16, __nv_bfloat16, int32_t, DoSoftmaxBeforeTopK>( + reinterpret_cast<__nv_bfloat16*>(router_logits.mutable_data_ptr()), + reinterpret_cast<__nv_bfloat16*>(topk_values.mutable_data_ptr()), topk_indices.data_ptr(), + num_tokens, num_experts, topk, stream); + } break; case torch::kHalf: - // Handle Half - tk::invokeRenormMoeRouting( - reinterpret_cast(router_logits.mutable_data_ptr()), - reinterpret_cast(topk_values.mutable_data_ptr()), - reinterpret_cast(topk_indices.mutable_data_ptr()), num_tokens, num_experts, topk, stream); + // Handle Half input + if (topk_values_dtype == torch::kFloat32) + { + tk::invokeCustomMoeRouting( + reinterpret_cast(router_logits.mutable_data_ptr()), + reinterpret_cast(topk_values.mutable_data_ptr()), topk_indices.data_ptr(), num_tokens, + num_experts, topk, stream); + } + else + { // bfloat16 output + tk::invokeCustomMoeRouting( + reinterpret_cast(router_logits.mutable_data_ptr()), + reinterpret_cast<__nv_bfloat16*>(topk_values.mutable_data_ptr()), topk_indices.data_ptr(), + num_tokens, num_experts, topk, stream); + } break; default: // Handle other data types @@ -72,21 +108,23 @@ std::tuple custom_moe_routing_op(th::Tensor const& route return {topk_indices, topk_values}; } -std::tuple renorm_moe_routing_op(th::Tensor const& router_logits, int64_t topk) +std::tuple renorm_moe_routing_op( + th::Tensor const& router_logits, int64_t topk, c10::optional output_dtype) { - return custom_moe_routing_op(router_logits, topk); + return custom_moe_routing_op(router_logits, topk, output_dtype); } -std::tuple default_moe_routing_op(th::Tensor const& router_logits, int64_t topk) +std::tuple default_moe_routing_op( + th::Tensor const& router_logits, int64_t topk, c10::optional output_dtype) { - return custom_moe_routing_op(router_logits, topk); + return custom_moe_routing_op(router_logits, topk, output_dtype); } } // namespace torch_ext TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( - "renorm_moe_routing_op(Tensor router_logits, SymInt topk" + "renorm_moe_routing_op(Tensor router_logits, SymInt topk, ScalarType? output_dtype=None" ") -> (Tensor, Tensor)"); } @@ -98,7 +136,7 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( - "default_moe_routing_op(Tensor router_logits, SymInt topk" + "default_moe_routing_op(Tensor router_logits, SymInt topk, ScalarType? output_dtype=None" ") -> (Tensor, Tensor)"); } diff --git a/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp index 24bb2f8d534c..0c57fe07d59f 100644 --- a/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp @@ -28,7 +28,7 @@ namespace btg = batchedGemm::trtllm::gen; using tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::RoutingMethodType; using MoeRunnerType = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::Runner; -std::vector run_fp4_block_scale_moe_runner(torch::Tensor const& routing_logits, +std::vector run_fp4_block_scale_moe_runner(torch::optional const& routing_logits, torch::optional const& routing_bias, torch::Tensor const& hidden_states, torch::Tensor const& hidden_states_scale, torch::Tensor const& gemm1_weights, torch::Tensor const& gemm1_weights_scale, torch::Tensor const& gemm2_weights, @@ -37,23 +37,62 @@ std::vector run_fp4_block_scale_moe_runner(torch::Tensor const& r int64_t const num_experts, int64_t const top_k, std::optional const n_group, std::optional const topk_group, int64_t const intermediate_size, int64_t const local_expert_offset, int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t const tile_tokens_dim, - int64_t const routing_method_type, bool const do_finalize, MoeRunnerType& moe_runner, int64_t const moeConfigIndex) + int64_t const routing_method_type, bool const do_finalize, MoeRunnerType& moe_runner, int64_t const moeConfigIndex, + torch::optional const& topk_weights, torch::optional const& topk_ids) { TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP4 block scale MOE"); TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64, "tile_tokens_dim must be 8, 16, 32, 64"); - if (static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3) + + if (topk_ids.has_value() && topk_weights.has_value()) + { + TORCH_CHECK(topk_ids.value().scalar_type() == at::ScalarType::Int, "topk_ids must be int"); + TORCH_CHECK(topk_weights.value().scalar_type() == at::ScalarType::BFloat16, "topk_weights must be bfloat16."); + TORCH_CHECK(topk_ids.value().dim() == 2, "topk_ids must be 2D."); + TORCH_CHECK(topk_ids.value().sizes()[0] == hidden_states.sizes()[0], + "topk_ids and hidden_states must have the same number of tokens."); + TORCH_CHECK(topk_ids.value().sizes()[1] == top_k, "topk_ids dim1 must match top_k."); + TORCH_CHECK(topk_weights.value().dim() == 2, "topk_weights must be 2D."); + TORCH_CHECK(topk_weights.value().sizes()[0] == hidden_states.sizes()[0], + "topk_weights and hidden_states must have the same number of tokens."); + TORCH_CHECK(topk_weights.value().sizes()[1] == top_k, "topk_weights dim1 must match top_k."); + } + else if (routing_logits.has_value()) + { + if (static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3) + { + TORCH_CHECK(routing_logits.value().scalar_type() == at::ScalarType::Float, "routing_logits must be float"); + } + else + { + TORCH_CHECK( + routing_logits.value().scalar_type() == at::ScalarType::BFloat16, "routing_logits must be bfloat16"); + } + TORCH_CHECK(routing_logits.value().dim() == 2, "routing_logits must be 2D."); + TORCH_CHECK(routing_logits.value().sizes()[1] == num_experts, "routing_logits has incorrect shape."); + } + else + { + TORCH_CHECK(false, "routing_logits or (topk_ids and topk_weights) must be provided."); + } + + if (topk_ids.has_value() && topk_weights.has_value() && routing_logits.has_value()) { - TORCH_CHECK(routing_logits.scalar_type() == at::ScalarType::Float, "routing_logits must be float"); + TLLM_LOG_WARNING( + "When logits and (topk_ids and topk_weights) are both provided, we only use (topk_ids and topk_weights)."); + } + + if (topk_ids.has_value()) + { + TORCH_CHECK(topk_ids.value().sizes()[0] == hidden_states.sizes()[0], + "topk_ids and hidden_states must have the same number of tokens."); } else { - TORCH_CHECK(routing_logits.scalar_type() == at::ScalarType::BFloat16, "routing_logits must be bfloat16"); + TORCH_CHECK(routing_logits.value().sizes()[0] == hidden_states.sizes()[0], + "routing_logits and hidden_states must have the same number of tokens."); } - TORCH_CHECK(routing_logits.dim() == 2, "routing_logits must be 2D."); - TORCH_CHECK(routing_logits.sizes()[0] == hidden_states.sizes()[0], - "routing_logits and hidden_states must have the same number of tokens."); - TORCH_CHECK(routing_logits.sizes()[1] == num_experts, "routing_logits has incorrect shape."); + if (routing_bias.has_value()) { TORCH_CHECK(routing_bias.value().scalar_type() == at::ScalarType::BFloat16, "routing_bias must be bfloat16."); @@ -90,6 +129,13 @@ std::vector run_fp4_block_scale_moe_runner(torch::Tensor const& r TORCH_CHECK(num_experts > top_k, "num_experts must be greater than top_k"); TORCH_CHECK(num_experts <= 256, "num_experts must be less than or equal to 256"); + // If both routing inputs are provided, they must be on the same device + if (routing_logits.has_value() && topk_ids.has_value()) + { + TORCH_CHECK( + routing_logits->device() == topk_ids->device(), "routing_logits and topk_ids must be on the same device"); + } + tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoERunnerArgs args; tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoEWorkspace workspace; @@ -99,10 +145,14 @@ std::vector run_fp4_block_scale_moe_runner(torch::Tensor const& r = routing_bias.has_value() ? routing_bias.value().scalar_type() : at::ScalarType::BFloat16; args.mDtypeElt = btg::Dtype::E2m1; args.mDtypeExpW = routing_bias_dtype == at::ScalarType::Float ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; - args.routing_logits = routing_logits.data_ptr(); + args.routing_logits = routing_logits.has_value() ? routing_logits.value().data_ptr() : nullptr; args.routing_bias = routing_bias.has_value() ? routing_bias.value().data_ptr() : nullptr; args.hidden_states = hidden_states.data_ptr(); args.hidden_states_scale = hidden_states_scale.data_ptr(); + + args.topk_weights = topk_weights.has_value() ? topk_weights.value().data_ptr() : nullptr; + args.topk_ids = topk_ids.has_value() ? static_cast(topk_ids.value().data_ptr()) : nullptr; + args.gemm1_weights = gemm1_weights.data_ptr(); args.gemm1_weights_scale = gemm1_weights_scale.data_ptr(); args.gemm2_weights = gemm2_weights.data_ptr(); @@ -120,61 +170,71 @@ std::vector run_fp4_block_scale_moe_runner(torch::Tensor const& r args.intermediate_size = intermediate_size; // allocate workspace for routing kernel + if (routing_logits.has_value() && topk_ids.has_value()) + { + TORCH_CHECK(routing_logits.value().device() == topk_ids.value().device(), + "routing_logits and topk_ids must be on the same device"); + } + auto routing_device = routing_logits.has_value() ? routing_logits.value().device() : topk_ids.value().device(); at::Tensor num_tokens_per_expert - = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_device, std::nullopt); int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxPermutedPaddedCount( args.num_tokens, top_k, num_experts, tile_tokens_dim); at::Tensor total_num_padded_tokens - = at::empty({}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int)); - at::Tensor expanded_idx_to_permuted_idx = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); + at::Tensor expanded_idx_to_permuted_idx + = at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor permuted_idx_to_token_idx - = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_logits.device(), std::nullopt); - at::Tensor expert_weights = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, routing_bias_dtype, routing_logits.device(), std::nullopt); - at::Tensor expert_indexes = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_device, std::nullopt); + at::Tensor expert_weights + = at::detail::empty_cuda({args.num_tokens, args.top_k}, routing_bias_dtype, routing_device, std::nullopt); + at::Tensor expert_indexes + = at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); int64_t const size_of_expert_count_histogram = std::max(num_experts * 2, int64_t(256 * 2)); - at::Tensor expert_count_histogram = at::detail::empty_cuda({size_of_expert_count_histogram}, - at::ScalarType::Int, // 256 is the max number of threads per block and max number of experts - routing_logits.device(), std::nullopt); + at::Tensor expert_count_histogram + = at::detail::empty_cuda({size_of_expert_count_histogram}, at::ScalarType::Int, routing_device, std::nullopt); // allocate workspace for activation/gemm/finalize kernels - at::Tensor gemm1_output = at::detail::empty_cuda({max_num_padded_tokens, intermediate_size / 2}, - at::ScalarType::Float8_e4m3fn, hidden_states.device(), std::nullopt); + at::Tensor gemm1_output = at::detail::empty_cuda( + {max_num_padded_tokens, intermediate_size / 2}, at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); int64_t sf_size = tensorrt_llm::computeSwizzledLayoutSFSize(max_num_padded_tokens, intermediate_size / 16); at::Tensor gemm1_output_scale - = at::detail::empty_cuda({sf_size}, at::ScalarType::Float8_e4m3fn, hidden_states.device(), std::nullopt); + = at::detail::empty_cuda({sf_size}, at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); at::Tensor gemm2_output = at::detail::empty_cuda( - {max_num_padded_tokens, args.hidden_size}, at::ScalarType::BFloat16, hidden_states.device(), std::nullopt); + {max_num_padded_tokens, args.hidden_size}, at::ScalarType::BFloat16, routing_device, std::nullopt); int32_t max_num_ctas = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxNumCtasInBatchDim( args.num_tokens, args.top_k, args.num_experts, tile_tokens_dim); at::Tensor cta_idx_xy_to_batch_idx - = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor cta_idx_xy_to_mn_limit - = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor num_non_exiting_ctas - = at::empty({}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int)); + = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); + + // Set the optional pointer to the expert weights and expert ids + void* expert_weights_ptr = args.topk_weights ? args.topk_weights : expert_weights.data_ptr(); // // TopK routing // tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::Runner routing_runner(tile_tokens_dim); - auto const& stream = at::cuda::getCurrentCUDAStream(routing_logits.get_device()); + auto const& stream = at::cuda::getCurrentCUDAStream( + routing_logits.has_value() ? routing_logits.value().get_device() : topk_ids.value().get_device()); routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), nullptr, /*permuted_idx_to_expanded_idx.data_ptr(),*/ - permuted_idx_to_token_idx.data_ptr(), expert_weights.data_ptr(), num_tokens_per_expert.data_ptr(), - cta_idx_xy_to_batch_idx.data_ptr(), cta_idx_xy_to_mn_limit.data_ptr(), - num_non_exiting_ctas.data_ptr(), args.mDtypeElt, false /* use_routing_scales_on_input */, - false /* use_deep_seek_fp8 */, static_cast(routing_method_type), stream); + permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, + num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), + cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, + false /* use_routing_scales_on_input */, false /* use_deep_seek_fp8 */, + static_cast(routing_method_type), stream); // // FC13 (gemm1) + FC2 (gemm2) @@ -247,7 +307,8 @@ std::vector run_fp4_block_scale_moe_runner(torch::Tensor const& r workspace.expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx.data_ptr(); // Needed by permute/finalize kernels workspace.permuted_idx_to_token_idx = permuted_idx_to_token_idx.data_ptr(); // Needed by permuteGemm1 kernel - workspace.expert_weights = expert_weights.data_ptr(); // Consumed by finalize kernel + + workspace.expert_weights = expert_weights_ptr; // Consumed by finalize kernel workspace.cta_idx_xy_to_batch_idx = cta_idx_xy_to_batch_idx.data_ptr(); workspace.cta_idx_xy_to_mn_limit = cta_idx_xy_to_mn_limit.data_ptr(); @@ -297,7 +358,7 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder mRunner = std::make_unique(mDtypeElt, mUseDeepSeekFp8, mTileTokensDim); } - [[nodiscard]] std::vector run(torch::Tensor const& routing_logits, + [[nodiscard]] std::vector run(torch::optional const& routing_logits, torch::optional const& routing_bias, torch::Tensor const& hidden_states, torch::Tensor const& hidden_states_scale, torch::Tensor const& gemm1_weights, torch::Tensor const& gemm1_weights_scale, torch::Tensor const& gemm2_weights, @@ -306,7 +367,8 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder int64_t const num_experts, int64_t const top_k, std::optional const n_group, std::optional const topk_group, int64_t const intermediate_size, int64_t const local_expert_offset, int64_t const local_num_experts, std::optional const routed_scaling_factor, - int64_t const routing_method_type, bool const do_finalize, int64_t moeConfigIndex) + int64_t const routing_method_type, bool const do_finalize, int64_t moeConfigIndex, + torch::optional const& topk_weights, torch::optional const& topk_ids) { // Autotuner has requested a default or 'fallback' config index @@ -325,7 +387,7 @@ class FP4BlockScaleMoeRunner : public torch::CustomClassHolder gemm1_weights, gemm1_weights_scale, gemm2_weights, gemm2_weights_scale, output1_scales_scalar, output1_scales_gate_scalar, output2_scales_scalar, num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, mTileTokensDim, - routing_method_type, do_finalize, *mRunner, moeConfigIndex); + routing_method_type, do_finalize, *mRunner, moeConfigIndex, topk_weights, topk_ids); } [[nodiscard]] std::vector getValidConfigs( diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index f1c8c560b6ab..b6f720c83464 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -30,20 +30,66 @@ namespace btg = batchedGemm::trtllm::gen; using tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::RoutingMethodType; using MoeRunnerType = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::Runner; -at::Tensor run_fp8_block_scale_moe(at::Tensor const& routing_logits, std::optional const& routing_bias, - at::Tensor const& hidden_states, at::Tensor const& hidden_states_scale, at::Tensor const& gemm1_weights, - at::Tensor const& gemm1_weights_scale, at::Tensor const& gemm2_weights, at::Tensor const& gemm2_weights_scale, - int64_t const num_experts, int64_t const top_k, std::optional const n_group, - std::optional const topk_group, int64_t const intermediate_size, int64_t const local_expert_offset, - int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t const tile_tokens_dim, - int64_t const routing_method_type, MoeRunnerType& moe_runner, int64_t moeConfigIndex) +at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logits, + std::optional const& routing_bias, at::Tensor const& hidden_states, + at::Tensor const& hidden_states_scale, at::Tensor const& gemm1_weights, at::Tensor const& gemm1_weights_scale, + at::Tensor const& gemm2_weights, at::Tensor const& gemm2_weights_scale, int64_t const num_experts, + int64_t const top_k, std::optional const n_group, std::optional const topk_group, + int64_t const intermediate_size, int64_t const local_expert_offset, int64_t const local_num_experts, + std::optional const routed_scaling_factor, int64_t const tile_tokens_dim, int64_t const routing_method_type, + MoeRunnerType& moe_runner, int64_t moeConfigIndex, std::optional const& topk_weights, + std::optional const& topk_ids) { TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE"); - TORCH_CHECK(routing_logits.scalar_type() == at::ScalarType::Float, "routing_logits must be float."); - TORCH_CHECK(routing_logits.dim() == 2, "routing_logits must be 2D."); - TORCH_CHECK(routing_logits.sizes()[0] == hidden_states.sizes()[0], - "routing_logits and hidden_states must have the same number of tokens."); - TORCH_CHECK(routing_logits.sizes()[1] == num_experts, "routing_logits dim1 must match num_experts."); + + if (topk_ids.has_value() && topk_weights.has_value()) + { + TORCH_CHECK(topk_ids.value().scalar_type() == at::ScalarType::Int, "topk_ids must be int"); + TORCH_CHECK(topk_weights.value().scalar_type() == at::ScalarType::BFloat16, "topk_weights must be bfloat16."); + TORCH_CHECK(topk_ids.value().dim() == 2, "topk_ids must be 2D."); + TORCH_CHECK(topk_ids.value().sizes()[0] == hidden_states.sizes()[0], + "topk_ids and hidden_states must have the same number of tokens."); + TORCH_CHECK(topk_ids.value().sizes()[1] == top_k, "topk_ids dim1 must match top_k."); + TORCH_CHECK(topk_weights.value().dim() == 2, "topk_weights must be 2D."); + TORCH_CHECK(topk_weights.value().sizes()[0] == hidden_states.sizes()[0], + "topk_weights and hidden_states must have the same number of tokens."); + TORCH_CHECK(topk_weights.value().sizes()[1] == top_k, "topk_weights dim1 must match top_k."); + } + else if (routing_logits.has_value()) + { + if (static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3) + { + TORCH_CHECK(routing_logits.value().scalar_type() == at::ScalarType::Float, "routing_logits must be float"); + } + else + { + TORCH_CHECK( + routing_logits.value().scalar_type() == at::ScalarType::BFloat16, "routing_logits must be bfloat16"); + } + TORCH_CHECK(routing_logits.value().dim() == 2, "routing_logits must be 2D."); + TORCH_CHECK(routing_logits.value().sizes()[1] == num_experts, "routing_logits dim1 must match num_experts."); + } + else + { + TORCH_CHECK(false, "routing_logits or (topk_ids and topk_weights) must be provided."); + } + + if (topk_ids.has_value() && topk_weights.has_value() && routing_logits.has_value()) + { + TLLM_LOG_WARNING( + "When logits and (topk_ids and topk_weights) are both provided, we only use (topk_ids and topk_weights)."); + } + + if (topk_ids.has_value()) + { + TORCH_CHECK(topk_ids.value().sizes()[0] == hidden_states.sizes()[0], + "topk_ids and hidden_states must have the same number of tokens."); + } + else + { + TORCH_CHECK(routing_logits.value().sizes()[0] == hidden_states.sizes()[0], + "routing_logits and hidden_states must have the same number of tokens."); + } if (routing_bias.has_value()) { @@ -79,6 +125,13 @@ at::Tensor run_fp8_block_scale_moe(at::Tensor const& routing_logits, std::option TORCH_CHECK(num_experts % 4 == 0, "Routing kernel expects that num_experts must be divisible by 4"); TORCH_CHECK(num_experts > top_k, "num_experts must be greater than top_k"); + // If both routing inputs are provided, they must be on the same device + if (routing_logits.has_value() && topk_ids.has_value()) + { + TORCH_CHECK( + routing_logits->device() == topk_ids->device(), "routing_logits and topk_ids must be on the same device"); + } + tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoERunnerArgs args; tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoEWorkspace workspace; @@ -89,8 +142,12 @@ at::Tensor run_fp8_block_scale_moe(at::Tensor const& routing_logits, std::option = routing_bias.has_value() ? routing_bias.value().scalar_type() : at::ScalarType::BFloat16; args.mDtypeExpW = routing_bias_dtype == at::ScalarType::Float ? btg::Dtype::Fp32 : btg::Dtype::Bfloat16; - args.routing_logits = routing_logits.data_ptr(); + args.routing_logits = routing_logits.has_value() ? routing_logits.value().data_ptr() : nullptr; args.routing_bias = routing_bias.has_value() ? routing_bias.value().data_ptr() : nullptr; + + args.topk_weights = topk_weights.has_value() ? topk_weights.value().data_ptr() : nullptr; + args.topk_ids = topk_ids.has_value() ? static_cast(topk_ids.value().data_ptr()) : nullptr; + args.hidden_states = hidden_states.data_ptr(); args.hidden_states_scale = hidden_states_scale.data_ptr(); args.gemm1_weights = gemm1_weights.data_ptr(); @@ -110,55 +167,64 @@ at::Tensor run_fp8_block_scale_moe(at::Tensor const& routing_logits, std::option args.mUseDeepSeekFp8 = true; // allocate workspace for routing kernel + if (routing_logits.has_value() && topk_ids.has_value()) + { + TORCH_CHECK(routing_logits.value().device() == topk_ids.value().device(), + "routing_logits and topk_ids must be on the same device"); + } + auto routing_device = routing_logits.has_value() ? routing_logits.value().device() : topk_ids.value().device(); at::Tensor num_tokens_per_expert - = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_device, std::nullopt); int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxPermutedPaddedCount( args.num_tokens, top_k, num_experts, tile_tokens_dim); at::Tensor total_num_padded_tokens - = at::empty({}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int)); - at::Tensor expanded_idx_to_permuted_idx = at::detail::empty_cuda( - {args.num_tokens * args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); + at::Tensor expanded_idx_to_permuted_idx + = at::detail::empty_cuda({args.num_tokens * args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor permuted_idx_to_token_idx - = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_logits.device(), std::nullopt); - at::Tensor expert_weights = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, routing_bias_dtype, routing_logits.device(), std::nullopt); - at::Tensor expert_indexes = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_device, std::nullopt); + at::Tensor expert_weights + = at::detail::empty_cuda({args.num_tokens, args.top_k}, routing_bias_dtype, routing_device, std::nullopt); + at::Tensor expert_indexes + = at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); int64_t const size_of_expert_count_histogram = std::max(num_experts * 2, int64_t(256 * 2)); - at::Tensor expert_count_histogram = at::detail::empty_cuda({size_of_expert_count_histogram}, - at::ScalarType::Int, // 256 is the max number of threads per block and max number of experts - routing_logits.device(), std::nullopt); + at::Tensor expert_count_histogram + = at::detail::empty_cuda({size_of_expert_count_histogram}, at::ScalarType::Int, routing_device, std::nullopt); // allocate workspace for activation/gemm/finalize kernels - at::Tensor gemm1_output = at::detail::empty_cuda({max_num_padded_tokens, 2 * intermediate_size}, - at::ScalarType::Float8_e4m3fn, hidden_states.device(), std::nullopt); - at::Tensor gemm1_output_scale = at::detail::empty_cuda({2 * intermediate_size / 128, max_num_padded_tokens}, - at::ScalarType::Float, hidden_states.device(), std::nullopt); - at::Tensor activation_output = at::detail::empty_cuda({max_num_padded_tokens, intermediate_size}, - at::ScalarType::Float8_e4m3fn, hidden_states.device(), std::nullopt); + at::Tensor gemm1_output = at::detail::empty_cuda( + {max_num_padded_tokens, 2 * intermediate_size}, at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); + at::Tensor gemm1_output_scale = at::detail::empty_cuda( + {2 * intermediate_size / 128, max_num_padded_tokens}, at::ScalarType::Float, routing_device, std::nullopt); + at::Tensor activation_output = at::detail::empty_cuda( + {max_num_padded_tokens, intermediate_size}, at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); at::Tensor activation_output_scale = at::detail::empty_cuda( - {intermediate_size / 128, max_num_padded_tokens}, at::ScalarType::Float, hidden_states.device(), std::nullopt); + {intermediate_size / 128, max_num_padded_tokens}, at::ScalarType::Float, routing_device, std::nullopt); at::Tensor gemm2_output = at::detail::empty_cuda( - {max_num_padded_tokens, args.hidden_size}, at::ScalarType::BFloat16, hidden_states.device(), std::nullopt); + {max_num_padded_tokens, args.hidden_size}, at::ScalarType::BFloat16, routing_device, std::nullopt); int32_t max_num_ctas = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxNumCtasInBatchDim( args.num_tokens, args.top_k, args.num_experts, tile_tokens_dim); at::Tensor cta_idx_xy_to_batch_idx - = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor cta_idx_xy_to_mn_limit - = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor num_non_exiting_ctas - = at::empty({}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int)); + = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); + + // Set the optional pointer to the expert weights and expert ids + void* expert_weights_ptr = args.topk_weights ? args.topk_weights : expert_weights.data_ptr(); tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::Runner routing_runner(tile_tokens_dim); - auto const& stream = at::cuda::getCurrentCUDAStream(routing_logits.get_device()); - routing_runner.run(routing_logits.data_ptr(), args.routing_bias, args.num_tokens, args.num_experts, - args.top_k, args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, - args.routed_scaling_factor, expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), - total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), - nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, permuted_idx_to_token_idx.data_ptr(), - expert_weights.data_ptr(), num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), + auto const& stream = at::cuda::getCurrentCUDAStream( + routing_logits.has_value() ? routing_logits.value().get_device() : topk_ids.value().get_device()); + routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, + args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, + expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), + expanded_idx_to_permuted_idx.data_ptr(), nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, + permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, + num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, false, true, static_cast(routing_method_type), stream); @@ -206,7 +272,7 @@ at::Tensor run_fp8_block_scale_moe(at::Tensor const& routing_logits, std::option workspace.expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx.data_ptr(); // Needed by activation/finalize kernels workspace.permuted_idx_to_token_idx = permuted_idx_to_token_idx.data_ptr(); // Needed by permuteGemm1 kernel - workspace.expert_weights = expert_weights.data_ptr(); // Consumed by finalize kernel + workspace.expert_weights = expert_weights_ptr; // Consumed by finalize kernel workspace.cta_idx_xy_to_batch_idx = cta_idx_xy_to_batch_idx.data_ptr(); workspace.cta_idx_xy_to_mn_limit = cta_idx_xy_to_mn_limit.data_ptr(); @@ -255,13 +321,14 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder return mRunner->getValidConfigIndices(topK, hiddenSize, intermediateSize, numLocalExperts, numTokens); } - [[nodiscard]] at::Tensor run(at::Tensor const& routing_logits, std::optional const& routing_bias, - at::Tensor const& hidden_states, at::Tensor const& hidden_states_scale, at::Tensor const& gemm1_weights, - at::Tensor const& gemm1_weights_scale, at::Tensor const& gemm2_weights, at::Tensor const& gemm2_weights_scale, - int64_t num_experts, int64_t top_k, std::optional const n_group, - std::optional const topk_group, int64_t const intermediate_size, int64_t const local_expert_offset, - int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t routing_method_type, - int64_t moeConfigIndex) + [[nodiscard]] at::Tensor run(at::optional const& routing_logits, + std::optional const& routing_bias, at::Tensor const& hidden_states, + at::Tensor const& hidden_states_scale, at::Tensor const& gemm1_weights, at::Tensor const& gemm1_weights_scale, + at::Tensor const& gemm2_weights, at::Tensor const& gemm2_weights_scale, int64_t num_experts, int64_t top_k, + std::optional const n_group, std::optional const topk_group, int64_t const intermediate_size, + int64_t const local_expert_offset, int64_t const local_num_experts, + std::optional const routed_scaling_factor, int64_t routing_method_type, int64_t moeConfigIndex, + std::optional const& topk_weights, std::optional const& topk_ids) { // Autotuner has requested a default or 'fallback' config index @@ -277,7 +344,7 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder return run_fp8_block_scale_moe(routing_logits, routing_bias, hidden_states, hidden_states_scale, gemm1_weights, gemm1_weights_scale, gemm2_weights, gemm2_weights_scale, num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, mTileTokensDim, - routing_method_type, *mRunner, moeConfigIndex); + routing_method_type, *mRunner, moeConfigIndex, topk_weights, topk_ids); } private: diff --git a/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp index b6cdbff0e9ab..f99cf73cff87 100644 --- a/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp @@ -25,7 +25,7 @@ namespace torch_ext namespace btg = batchedGemm::trtllm::gen; using tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::RoutingMethodType; -torch::Tensor fp8_per_tensor_scale_moe_runner(torch::Tensor const& routing_logits, +torch::Tensor fp8_per_tensor_scale_moe_runner(torch::optional const& routing_logits, torch::optional const& routing_bias, torch::Tensor const& hidden_states, torch::Tensor const& gemm1_weights, torch::Tensor const& output1_scales_scalar, torch::Tensor const& output1_scales_gate_scalar, torch::Tensor const& gemm2_weights, @@ -33,20 +33,67 @@ torch::Tensor fp8_per_tensor_scale_moe_runner(torch::Tensor const& routing_logit std::optional const n_group, std::optional const topk_group, int64_t const intermediate_size, int64_t const local_expert_offset, int64_t const local_num_experts, std::optional const routed_scaling_factor, bool const use_routing_scales_on_input, - int64_t const tile_tokens_dim, int64_t const routing_method_type) + int64_t const tile_tokens_dim, int64_t const routing_method_type, + torch::optional const& topk_weights, torch::optional const& topk_ids) { - TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE"); - if (use_routing_scales_on_input) + if (topk_ids.has_value() && topk_weights.has_value()) + { + TORCH_CHECK(topk_ids.value().scalar_type() == at::ScalarType::Int, "topk_ids must be int"); + TORCH_CHECK(topk_weights.value().scalar_type() == at::ScalarType::BFloat16, "topk_weights must be bfloat16."); + TORCH_CHECK(topk_ids.value().dim() == 2, "topk_ids must be 2D."); + TORCH_CHECK(topk_ids.value().sizes()[0] == hidden_states.sizes()[0], + "topk_ids and hidden_states must have the same number of tokens."); + TORCH_CHECK(topk_ids.value().sizes()[1] == top_k, "topk_ids dim1 must match top_k."); + TORCH_CHECK(topk_weights.value().dim() == 2, "topk_weights must be 2D."); + TORCH_CHECK(topk_weights.value().sizes()[0] == hidden_states.sizes()[0], + "topk_weights and hidden_states must have the same number of tokens."); + TORCH_CHECK(topk_weights.value().sizes()[1] == top_k, "topk_weights dim1 must match top_k."); + } + else if (routing_logits.has_value()) + { + if (use_routing_scales_on_input) + { + TORCH_CHECK( + routing_logits.value().scalar_type() == at::ScalarType::BFloat16, "routing_logits must be bfloat16."); + } + else + { + if (static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3) + { + TORCH_CHECK( + routing_logits.value().scalar_type() == at::ScalarType::Float, "routing_logits must be float"); + } + else + { + TORCH_CHECK(routing_logits.value().scalar_type() == at::ScalarType::BFloat16, + "routing_logits must be bfloat16"); + } + } + TORCH_CHECK(routing_logits.value().dim() == 2, "routing_logits must be 2D."); + TORCH_CHECK(routing_logits.value().sizes()[1] == num_experts, "routing_logits has incorrect shape."); + } + else + { + TORCH_CHECK(false, "routing_logits or (topk_ids and topk_weights) must be provided."); + } + + if (topk_ids.has_value() && topk_weights.has_value() && routing_logits.has_value()) + { + TLLM_LOG_WARNING( + "When logits and (topk_ids and topk_weights) are both provided, we only use (topk_ids and topk_weights)."); + } + + if (topk_ids.has_value()) { - TORCH_CHECK(routing_logits.scalar_type() == at::ScalarType::BFloat16, "routing_logits must be bfloat16."); + TORCH_CHECK(topk_ids.value().sizes()[0] == hidden_states.sizes()[0], + "topk_ids and hidden_states must have the same number of tokens."); } else { - TORCH_CHECK(routing_logits.scalar_type() == at::ScalarType::Float, "routing_logits must be float."); + TORCH_CHECK(routing_logits.value().sizes()[0] == hidden_states.sizes()[0], + "routing_logits and hidden_states must have the same number of tokens."); } - TORCH_CHECK(routing_logits.dim() == 2, "routing_logits must be 2D."); - TORCH_CHECK(routing_logits.sizes()[1] == num_experts, "routing_logits has incorrect shape."); if (routing_bias.has_value()) { @@ -82,13 +129,23 @@ torch::Tensor fp8_per_tensor_scale_moe_runner(torch::Tensor const& routing_logit TORCH_CHECK(num_experts % 4 == 0, "Routing kernel expects that num_experts must be divisible by 4"); TORCH_CHECK(num_experts > top_k, "num_experts must be greater than top_k"); + // If both routing inputs are provided, they must be on the same device + if (routing_logits.has_value() && topk_ids.has_value()) + { + TORCH_CHECK( + routing_logits->device() == topk_ids->device(), "routing_logits and topk_ids must be on the same device"); + } + tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoERunnerArgs args; tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoEWorkspace workspace; // setup args args.mDtypeElt = btg::Dtype::E4m3; - args.routing_logits = routing_logits.data_ptr(); + args.routing_logits = routing_logits.has_value() ? routing_logits.value().data_ptr() : nullptr; args.routing_bias = routing_bias.has_value() ? routing_bias.value().data_ptr() : nullptr; + + args.topk_weights = topk_weights.has_value() ? topk_weights.value().data_ptr() : nullptr; + args.topk_ids = topk_ids.has_value() ? static_cast(topk_ids.value().data_ptr()) : nullptr; args.hidden_states = hidden_states.data_ptr(); args.gemm1_weights = gemm1_weights.data_ptr(); args.output1_scales_scalar = output1_scales_scalar.data_ptr(); @@ -108,56 +165,68 @@ torch::Tensor fp8_per_tensor_scale_moe_runner(torch::Tensor const& routing_logit args.mUseRoutingScalesOnInput = use_routing_scales_on_input; // allocate workspace for routing kernel + if (routing_logits.has_value() && topk_ids.has_value()) + { + TORCH_CHECK(routing_logits.value().device() == topk_ids.value().device(), + "routing_logits and topk_ids must be on the same device"); + } + auto routing_device = routing_logits.has_value() ? routing_logits.value().device() : topk_ids.value().device(); at::Tensor num_tokens_per_expert - = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_device, std::nullopt); int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxPermutedPaddedCount( args.num_tokens, top_k, num_experts, tile_tokens_dim); at::Tensor total_num_padded_tokens - = at::empty({}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int)); - at::Tensor expanded_idx_to_permuted_idx = at::detail::empty_cuda( - {args.num_tokens * args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); + at::Tensor expanded_idx_to_permuted_idx + = at::detail::empty_cuda({args.num_tokens * args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor permuted_idx_to_token_idx - = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_logits.device(), std::nullopt); - at::Tensor expert_weights = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, at::ScalarType::BFloat16, routing_logits.device(), std::nullopt); - at::Tensor expert_indexes = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt); - at::Tensor expert_count_histogram = at::detail::empty_cuda({2 * 256}, - at::ScalarType::Int, // 256 is the max number of threads per block and max number of experts - routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_device, std::nullopt); + at::Tensor expert_weights + = at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::BFloat16, routing_device, std::nullopt); + at::Tensor expert_indexes + = at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); + + int64_t const size_of_expert_count_histogram = std::max(num_experts * 2, int64_t(256 * 2)); + at::Tensor expert_count_histogram + = at::detail::empty_cuda({size_of_expert_count_histogram}, at::ScalarType::Int, routing_device, std::nullopt); // allocate workspace for activation/gemm/finalize kernels - at::Tensor gemm1_output = at::detail::empty_cuda({max_num_padded_tokens, 2 * intermediate_size}, - at::ScalarType::Float8_e4m3fn, hidden_states.device(), std::nullopt); - at::Tensor gemm1_output_scale = at::detail::empty_cuda({2 * intermediate_size / 128, max_num_padded_tokens}, - at::ScalarType::Float, hidden_states.device(), std::nullopt); - at::Tensor activation_output = at::detail::empty_cuda({max_num_padded_tokens, intermediate_size}, - at::ScalarType::Float8_e4m3fn, hidden_states.device(), std::nullopt); + at::Tensor gemm1_output = at::detail::empty_cuda( + {max_num_padded_tokens, 2 * intermediate_size}, at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); + at::Tensor gemm1_output_scale = at::detail::empty_cuda( + {2 * intermediate_size / 128, max_num_padded_tokens}, at::ScalarType::Float, routing_device, std::nullopt); + at::Tensor activation_output = at::detail::empty_cuda( + {max_num_padded_tokens, intermediate_size}, at::ScalarType::Float8_e4m3fn, routing_device, std::nullopt); at::Tensor activation_output_scale = at::detail::empty_cuda( - {intermediate_size / 128, max_num_padded_tokens}, at::ScalarType::Float, hidden_states.device(), std::nullopt); + {intermediate_size / 128, max_num_padded_tokens}, at::ScalarType::Float, routing_device, std::nullopt); at::Tensor gemm2_output = at::detail::empty_cuda( - {max_num_padded_tokens, args.hidden_size}, at::ScalarType::BFloat16, hidden_states.device(), std::nullopt); + {max_num_padded_tokens, args.hidden_size}, at::ScalarType::BFloat16, routing_device, std::nullopt); int32_t max_num_ctas = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxNumCtasInBatchDim( args.num_tokens, args.top_k, args.num_experts, tile_tokens_dim); at::Tensor cta_idx_xy_to_batch_idx - = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor cta_idx_xy_to_mn_limit - = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor num_non_exiting_ctas - = at::empty({}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int)); + = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); + + // Set the optional pointer to the expert weights and expert ids + void* expert_weights_ptr = args.topk_weights ? args.topk_weights : expert_weights.data_ptr(); tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::Runner routing_runner(tile_tokens_dim); - auto const& stream = at::cuda::getCurrentCUDAStream(routing_logits.get_device()); - routing_runner.run(routing_logits.data_ptr(), args.routing_bias, args.num_tokens, args.num_experts, args.top_k, + auto const& stream = at::cuda::getCurrentCUDAStream( + routing_logits.has_value() ? routing_logits.value().get_device() : topk_ids.value().get_device()); + routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), nullptr /*permuted_idx_to_expanded_idx.data_ptr()*/, - permuted_idx_to_token_idx.data_ptr(), expert_weights.data_ptr(), num_tokens_per_expert.data_ptr(), - cta_idx_xy_to_batch_idx.data_ptr(), cta_idx_xy_to_mn_limit.data_ptr(), - num_non_exiting_ctas.data_ptr(), args.mDtypeElt, use_routing_scales_on_input, - false /* use_deep_seek_fp8 */, static_cast(routing_method_type), stream); + permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, + num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), + cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, + use_routing_scales_on_input, false /* use_deep_seek_fp8 */, static_cast(routing_method_type), + stream); // MoE kernel except routing TORCH_CHECK(hidden_states.scalar_type() == at::ScalarType::Float8_e4m3fn, "hidden_states must be fp8."); @@ -200,7 +269,7 @@ torch::Tensor fp8_per_tensor_scale_moe_runner(torch::Tensor const& routing_logit workspace.expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx.data_ptr(); // Needed by activation/finalize kernels workspace.permuted_idx_to_token_idx = permuted_idx_to_token_idx.data_ptr(); // Needed by permuteGemm1 kernel - workspace.expert_weights = expert_weights.data_ptr(); // Consumed by finalize kernel + workspace.expert_weights = expert_weights_ptr; // Consumed by finalize kernel workspace.cta_idx_xy_to_batch_idx = cta_idx_xy_to_batch_idx.data_ptr(); workspace.cta_idx_xy_to_mn_limit = cta_idx_xy_to_mn_limit.data_ptr(); @@ -241,7 +310,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( "fp8_per_tensor_scale_moe_runner(" - "Tensor routing_logits," + "Tensor? routing_logits," "Tensor? routing_bias," "Tensor hidden_states," "Tensor gemm1_weights," @@ -259,7 +328,9 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "float? routed_scaling_factor," "bool use_routing_scales_on_input," "int tile_tokens_dim," - "int routing_method_type) -> Tensor"); + "int routing_method_type," + "Tensor? topk_weights," + "Tensor? topk_ids) -> Tensor"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp index 99cdacd2ffac..9bfbce952d66 100644 --- a/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp @@ -31,7 +31,7 @@ namespace btg = batchedGemm::trtllm::gen; using tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::RoutingMethodType; using MoeRunnerType = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::Runner; -torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::Tensor const& routing_logits, +torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::optional const& routing_logits, torch::optional const& routing_bias, torch::Tensor const& hidden_states, std::optional const& hidden_states_scale, torch::Tensor const& gemm1_weights, torch::Tensor const& gemm1_weights_scale, std::optional const& gemm1_bias, @@ -44,18 +44,63 @@ torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::Tensor const& routing_l std::optional const n_group, std::optional const topk_group, int64_t const intermediate_size, std::optional const hidden_size_output, int64_t const local_expert_offset, int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t const tile_tokens_dim, int64_t const routing_method_type, - btg::Dtype const dtype, MoeRunnerType& moe_runner, int64_t moeConfigIndex) + btg::Dtype const dtype, MoeRunnerType& moe_runner, int64_t moeConfigIndex, + torch::optional const& topk_weights, torch::optional const& topk_ids) { TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by MXFP4 block scale MOE"); TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64, "tile_tokens_dim must be 8, 16, 32, 64"); - TORCH_CHECK(routing_logits.scalar_type() == at::ScalarType::Float - || routing_logits.scalar_type() == at::ScalarType::BFloat16, - "routing_logits must be float or bfloat16."); - TORCH_CHECK(routing_logits.dim() == 2, "routing_logits must be 2D."); - TORCH_CHECK( - routing_logits.sizes()[0] == hidden_states.sizes()[0], "routing_logits dim0 must match hidden_states dim0."); - TORCH_CHECK(routing_logits.sizes()[1] == num_experts, "routing_logits dim1 must match num_experts."); + + if (topk_ids.has_value() && topk_weights.has_value()) + { + TORCH_CHECK(topk_ids.value().scalar_type() == at::ScalarType::Int, "topk_ids must be int."); + TORCH_CHECK(topk_weights.value().scalar_type() == at::ScalarType::BFloat16, "topk_weights must be bfloat."); + TORCH_CHECK(topk_ids.value().dim() == 2, "topk_ids must be 2D."); + TORCH_CHECK(topk_weights.value().dim() == 2, "topk_weights must be 2D."); + TORCH_CHECK( + topk_ids.value().sizes()[0] == hidden_states.sizes()[0], "topk_ids dim0 must match hidden_states dim0."); + TORCH_CHECK(topk_ids.value().sizes()[1] == top_k, "topk_ids dim1 must match top_k."); + TORCH_CHECK(topk_weights.value().sizes()[0] == hidden_states.sizes()[0], + "topk_weights dim0 must match hidden_states dim0."); + TORCH_CHECK(topk_weights.value().sizes()[1] == top_k, "topk_weights dim1 must match top_k."); + } + else if (routing_logits.has_value()) + { + if (static_cast(routing_method_type) == RoutingMethodType::DeepSeekV3) + { + TORCH_CHECK(routing_logits.value().scalar_type() == at::ScalarType::Float, "routing_logits must be float"); + } + else + { + TORCH_CHECK( + routing_logits.value().scalar_type() == at::ScalarType::BFloat16, "routing_logits must be bfloat16"); + } + + TORCH_CHECK(routing_logits.value().dim() == 2, "routing_logits must be 2D."); + TORCH_CHECK(routing_logits.value().sizes()[1] == num_experts, "routing_logits dim1 must match num_experts."); + } + else + { + TORCH_CHECK(false, "routing_logits or (topk_ids and topk_weights) must be provided."); + } + + if (topk_ids.has_value() && topk_weights.has_value() && routing_logits.has_value()) + { + TLLM_LOG_WARNING( + "When logits and (topk_ids and topk_weights) are both provided, we only use (topk_ids and topk_weights)."); + } + + if (topk_ids.has_value()) + { + TORCH_CHECK(topk_ids.value().sizes()[0] == hidden_states.sizes()[0], + "topk_ids and hidden_states must have the same number of tokens."); + } + else + { + TORCH_CHECK(routing_logits.value().sizes()[0] == hidden_states.sizes()[0], + "routing_logits and hidden_states must have the same number of tokens."); + } + if (routing_bias.has_value()) { TORCH_CHECK(routing_bias.value().scalar_type() == at::ScalarType::BFloat16, "routing_bias must be bfloat16."); @@ -87,15 +132,26 @@ torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::Tensor const& routing_l TORCH_CHECK(num_experts % 4 == 0, "Routing kernel expects that num_experts must be divisible by 4"); TORCH_CHECK(num_experts > top_k, "num_experts must be greater than top_k"); + // If both routing inputs are provided, they must be on the same device + if (routing_logits.has_value() && topk_ids.has_value()) + { + TORCH_CHECK( + routing_logits->device() == topk_ids->device(), "routing_logits and topk_ids must be on the same device"); + } + tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoERunnerArgs args; tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::MoE::MoEWorkspace workspace; // setup args args.mDtypeElt = dtype; - args.routing_logits = routing_logits.data_ptr(); + args.routing_logits = routing_logits.has_value() ? routing_logits.value().data_ptr() : nullptr; args.routing_bias = routing_bias.has_value() ? routing_bias.value().data_ptr() : nullptr; args.hidden_states = hidden_states.data_ptr(); args.hidden_states_scale = hidden_states_scale.has_value() ? hidden_states_scale.value().data_ptr() : nullptr; + + args.topk_weights = topk_weights.has_value() ? topk_weights.value().data_ptr() : nullptr; + args.topk_ids = topk_ids.has_value() ? static_cast(topk_ids.value().data_ptr()) : nullptr; + args.gemm1_weights = gemm1_weights.data_ptr(); args.gemm1_weights_scale = gemm1_weights_scale.data_ptr(); args.gemm2_weights = gemm2_weights.data_ptr(); @@ -124,52 +180,61 @@ torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::Tensor const& routing_l args.intermediate_size = intermediate_size; // allocate workspace for routing kernel + if (routing_logits.has_value() && topk_ids.has_value()) + { + TORCH_CHECK(routing_logits.value().device() == topk_ids.value().device(), + "routing_logits and topk_ids must be on the same device"); + } + auto routing_device = routing_logits.has_value() ? routing_logits.value().device() : topk_ids.value().device(); at::Tensor num_tokens_per_expert - = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({num_experts}, at::ScalarType::Int, routing_device, std::nullopt); int32_t max_num_padded_tokens = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxPermutedPaddedCount( args.num_tokens, top_k, num_experts, tile_tokens_dim); at::Tensor total_num_padded_tokens - = at::empty({}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int)); - at::Tensor expanded_idx_to_permuted_idx = at::detail::empty_cuda( - {args.num_tokens * args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); + at::Tensor expanded_idx_to_permuted_idx + = at::detail::empty_cuda({args.num_tokens * args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor permuted_idx_to_token_idx - = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_logits.device(), std::nullopt); - at::Tensor expert_weights = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, at::ScalarType::BFloat16, routing_logits.device(), std::nullopt); - at::Tensor expert_indexes = at::detail::empty_cuda( - {args.num_tokens, args.top_k}, at::ScalarType::Int, routing_logits.device(), std::nullopt); - at::Tensor expert_count_histogram = at::detail::empty_cuda({2 * 256}, - at::ScalarType::Int, // 256 is the max number of threads per block and max number of experts - routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_padded_tokens}, at::ScalarType::Int, routing_device, std::nullopt); + at::Tensor expert_weights + = at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::BFloat16, routing_device, std::nullopt); + at::Tensor expert_indexes + = at::detail::empty_cuda({args.num_tokens, args.top_k}, at::ScalarType::Int, routing_device, std::nullopt); + + int64_t const size_of_expert_count_histogram = std::max(num_experts * 2, int64_t(256 * 2)); + at::Tensor expert_count_histogram + = at::detail::empty_cuda({size_of_expert_count_histogram}, at::ScalarType::Int, routing_device, std::nullopt); + // Set the optional pointer to the expert weights and expert ids + void* expert_weights_ptr = args.topk_weights ? args.topk_weights : expert_weights.data_ptr(); int32_t const sf_block_size = 32; // allocate workspace for activation/gemm/finalize kernels auto const gemm1_output_type = dtype == btg::Dtype::Bfloat16 ? at::ScalarType::BFloat16 : at::ScalarType::Float8_e4m3fn; at::Tensor gemm1_output = at::detail::empty_cuda( - {max_num_padded_tokens, intermediate_size}, gemm1_output_type, hidden_states.device(), std::nullopt); + {max_num_padded_tokens, intermediate_size}, gemm1_output_type, routing_device, std::nullopt); std::optional gemm1_output_scale; if (dtype == btg::Dtype::MxE4m3) { int64_t sf_size = tensorrt_llm::computeSwizzledLayoutSFSize(max_num_padded_tokens, intermediate_size / sf_block_size); - gemm1_output_scale = at::detail::empty_cuda({sf_size}, SF_DTYPE, hidden_states.device(), std::nullopt); + gemm1_output_scale = at::detail::empty_cuda({sf_size}, SF_DTYPE, routing_device, std::nullopt); } at::Tensor gemm2_output = at::detail::empty_cuda( - {max_num_padded_tokens, args.hidden_size}, at::ScalarType::BFloat16, hidden_states.device(), std::nullopt); + {max_num_padded_tokens, args.hidden_size}, at::ScalarType::BFloat16, routing_device, std::nullopt); int32_t max_num_ctas = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::getMaxNumCtasInBatchDim( args.num_tokens, args.top_k, args.num_experts, tile_tokens_dim); at::Tensor cta_idx_xy_to_batch_idx - = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor cta_idx_xy_to_mn_limit - = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_logits.device(), std::nullopt); + = at::detail::empty_cuda({max_num_ctas}, at::ScalarType::Int, routing_device, std::nullopt); at::Tensor num_non_exiting_ctas - = at::empty({}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int)); + = at::empty({}, at::TensorOptions().device(routing_device).dtype(at::ScalarType::Int)); // FIXME: check shape TORCH_CHECK(dtype == btg::Dtype::MxE4m3 || dtype == btg::Dtype::Bfloat16 || dtype == btg::Dtype::E4m3, @@ -189,15 +254,17 @@ torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::Tensor const& routing_l // tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::Runner routing_runner(tile_tokens_dim); - auto const& stream = at::cuda::getCurrentCUDAStream(routing_logits.get_device()); + auto const& stream = at::cuda::getCurrentCUDAStream( + routing_logits.has_value() ? routing_logits.value().get_device() : topk_ids.value().get_device()); routing_runner.run(args.routing_logits, args.routing_bias, args.num_tokens, args.num_experts, args.top_k, args.n_group, args.topk_group, args.local_expert_offset, args.local_num_experts, args.routed_scaling_factor, expert_indexes.data_ptr(), expert_count_histogram.data_ptr(), total_num_padded_tokens.data_ptr(), expanded_idx_to_permuted_idx.data_ptr(), nullptr, /*permuted_idx_to_expanded_idx.data_ptr(),*/ - permuted_idx_to_token_idx.data_ptr(), expert_weights.data_ptr(), num_tokens_per_expert.data_ptr(), - cta_idx_xy_to_batch_idx.data_ptr(), cta_idx_xy_to_mn_limit.data_ptr(), - num_non_exiting_ctas.data_ptr(), args.mDtypeElt, false /* use_routing_scales_on_input */, - false /* use_deep_seek_fp8 */, static_cast(routing_method_type), stream); + permuted_idx_to_token_idx.data_ptr(), expert_weights_ptr, args.topk_ids, + num_tokens_per_expert.data_ptr(), cta_idx_xy_to_batch_idx.data_ptr(), + cta_idx_xy_to_mn_limit.data_ptr(), num_non_exiting_ctas.data_ptr(), args.mDtypeElt, + false /* use_routing_scales_on_input */, false /* use_deep_seek_fp8 */, + static_cast(routing_method_type), stream); // // FC13 (gemm1) + FC2 (gemm2) @@ -341,7 +408,7 @@ torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::Tensor const& routing_l workspace.expanded_idx_to_permuted_idx = expanded_idx_to_permuted_idx.data_ptr(); // Needed by permute/finalize kernels workspace.permuted_idx_to_token_idx = permuted_idx_to_token_idx.data_ptr(); // Needed by permuteGemm1 kernel - workspace.expert_weights = expert_weights.data_ptr(); // Consumed by finalize kernel + workspace.expert_weights = expert_weights_ptr; // Consumed by finalize kernel workspace.cta_idx_xy_to_batch_idx = cta_idx_xy_to_batch_idx.data_ptr(); workspace.cta_idx_xy_to_mn_limit = cta_idx_xy_to_mn_limit.data_ptr(); @@ -390,7 +457,7 @@ class Bf16MxE2m1BlockScaleMoeRunner : public torch::CustomClassHolder } // BF16 run does not use hidden_states_scale - [[nodiscard]] torch::Tensor run(torch::Tensor const& routing_logits, + [[nodiscard]] torch::Tensor run(torch::optional const& routing_logits, std::optional const& routing_bias, torch::Tensor const& hidden_states, torch::Tensor const& gemm1_weights, torch::Tensor const& gemm1_weights_scale, std::optional const& gemm1_bias, std::optional const& gemm1_alpha, @@ -399,7 +466,9 @@ class Bf16MxE2m1BlockScaleMoeRunner : public torch::CustomClassHolder std::optional const& gemm2_bias, int64_t num_experts, int64_t top_k, std::optional const n_group, std::optional const topk_group, int64_t intermediate_size, int64_t local_expert_offset, int64_t local_num_experts, std::optional routed_scaling_factor, - int64_t routing_method_type, int64_t moeConfigIndex) + int64_t routing_method_type, int64_t moeConfigIndex, torch::optional const& topk_weights, + torch::optional const& topk_ids) + { // Autotuner has requested a default or 'fallback' config index if (moeConfigIndex == -1) @@ -415,7 +484,7 @@ class Bf16MxE2m1BlockScaleMoeRunner : public torch::CustomClassHolder gemm1_weights, gemm1_weights_scale, gemm1_bias, gemm1_alpha, gemm1_beta, gemm1_clamp_limit, gemm2_weights, gemm2_weights_scale, gemm2_bias, std::nullopt, std::nullopt, std::nullopt, num_experts, top_k, n_group, topk_group, intermediate_size, std::nullopt, local_expert_offset, local_num_experts, routed_scaling_factor, - mTileTokensDim, routing_method_type, mDtypeAct, *mRunner, moeConfigIndex); + mTileTokensDim, routing_method_type, mDtypeAct, *mRunner, moeConfigIndex, topk_weights, topk_ids); } private: @@ -447,7 +516,7 @@ class MxE4m3MxE2m1BlockScaleMoeRunner : public torch::CustomClassHolder return mRunner->getValidConfigIndices(topK, hiddenSize, intermediateSize, numLocalExperts, numTokens); } - [[nodiscard]] torch::Tensor run(torch::Tensor const& routing_logits, + [[nodiscard]] torch::Tensor run(torch::optional const& routing_logits, std::optional const& routing_bias, torch::Tensor const& hidden_states, std::optional const& hidden_states_scale, torch::Tensor const& gemm1_weights, torch::Tensor const& gemm1_weights_scale, std::optional const& gemm1_bias, @@ -459,7 +528,8 @@ class MxE4m3MxE2m1BlockScaleMoeRunner : public torch::CustomClassHolder std::optional const& output2_scale_scalar, int64_t num_experts, int64_t top_k, std::optional const n_group, std::optional const topk_group, int64_t intermediate_size, std::optional const hidden_size_output, int64_t local_expert_offset, int64_t local_num_experts, - std::optional routed_scaling_factor, int64_t routing_method_type, int64_t moeConfigIndex) + std::optional routed_scaling_factor, int64_t routing_method_type, int64_t moeConfigIndex, + torch::optional const& topk_weights, torch::optional const& topk_ids) { // Autotuner has requested a default or 'fallback' config index if (moeConfigIndex == -1) @@ -476,7 +546,7 @@ class MxE4m3MxE2m1BlockScaleMoeRunner : public torch::CustomClassHolder gemm2_weights_scale, gemm2_bias, output1_scale_scalar, output1_scale_gate_scalar, output2_scale_scalar, num_experts, top_k, n_group, topk_group, intermediate_size, hidden_size_output, local_expert_offset, local_num_experts, routed_scaling_factor, mTileTokensDim, routing_method_type, mDtypeAct, *mRunner, - moeConfigIndex); + moeConfigIndex, topk_weights, topk_ids); } private: diff --git a/cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp b/cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp index d4f927688111..9df5a3a9e45c 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp +++ b/cpp/tests/unit_tests/kernels/routing/routingDeepSeekTest.cpp @@ -73,40 +73,49 @@ class RoutingDeepSeekKernelTest : public RoutingKernelTest expWeightsIdx[ie] = si; } - // Calculate the group score - int32_t expertsPerGroup = param.numExperts / param.nGroup; - PackedFloat groupScoresCandidate[param.nGroup][2]; - PackedFloat groupScores[param.nGroup]; - for (int ig = 0; ig < param.nGroup; ++ig) + PackedFloat finalTopkExperts[param.topK]; + if (param.nGroup != 0) { - std::partial_sort_copy(expWeightsIdx + ig * expertsPerGroup, expWeightsIdx + (ig + 1) * expertsPerGroup, - groupScoresCandidate[ig], groupScoresCandidate[ig] + 2, comp); - PackedFloat si{ - static_cast(groupScoresCandidate[ig][0].score + groupScoresCandidate[ig][1].score), - static_cast(ig)}; - groupScores[ig] = si; - } + // Calculate the group score + int32_t expertsPerGroup = param.numExperts / param.nGroup; + PackedFloat groupScoresCandidate[param.nGroup][2]; + PackedFloat groupScores[param.nGroup]; + for (int ig = 0; ig < param.nGroup; ++ig) + { + std::partial_sort_copy(expWeightsIdx + ig * expertsPerGroup, + expWeightsIdx + (ig + 1) * expertsPerGroup, groupScoresCandidate[ig], + groupScoresCandidate[ig] + 2, comp); + PackedFloat si{ + static_cast(groupScoresCandidate[ig][0].score + groupScoresCandidate[ig][1].score), + static_cast(ig)}; + groupScores[ig] = si; + } - // Get the topkGroup group score - PackedFloat topGroupScores[param.topkGroup]; - std::partial_sort_copy( - groupScores, groupScores + param.nGroup, topGroupScores, topGroupScores + param.topkGroup, comp); + // Get the topkGroup group score + PackedFloat topGroupScores[param.topkGroup]; + std::partial_sort_copy( + groupScores, groupScores + param.nGroup, topGroupScores, topGroupScores + param.topkGroup, comp); - // Prepare the data for the final topk experts selection - PackedFloat topkExpertsCandidate[param.topkGroup * expertsPerGroup]; - PackedFloat finalTopkExperts[param.topK]; + // Prepare the data for the final topk experts selection + PackedFloat topkExpertsCandidate[param.topkGroup * expertsPerGroup]; - for (int ig = 0; ig < param.topkGroup; ++ig) - { - for (int ie = 0; ie < expertsPerGroup; ++ie) + for (int ig = 0; ig < param.topkGroup; ++ig) { - topkExpertsCandidate[ig * expertsPerGroup + ie] - = expWeightsIdx[topGroupScores[ig].idx * expertsPerGroup + ie]; + for (int ie = 0; ie < expertsPerGroup; ++ie) + { + topkExpertsCandidate[ig * expertsPerGroup + ie] + = expWeightsIdx[topGroupScores[ig].idx * expertsPerGroup + ie]; + } } - } - std::partial_sort_copy(topkExpertsCandidate, topkExpertsCandidate + param.topkGroup * expertsPerGroup, - finalTopkExperts, finalTopkExperts + param.topK, comp); + std::partial_sort_copy(topkExpertsCandidate, topkExpertsCandidate + param.topkGroup * expertsPerGroup, + finalTopkExperts, finalTopkExperts + param.topK, comp); + } + else + { + std::partial_sort_copy(expWeightsIdx, expWeightsIdx + param.numExperts, finalTopkExperts, + finalTopkExperts + param.topK, comp); + } // Normalize the score float sumScore = 0.0f; @@ -120,16 +129,24 @@ class RoutingDeepSeekKernelTest : public RoutingKernelTest finalTopkExperts[ie].score = finalScore; } - // Convert back to io_dtype and store the topk expert results in hostData.mPtrExpertIdx + // Convert back to io_dtype and store the topk expert results in hostData.mPtrTopKPacked for (int ie = 0; ie < param.topK; ++ie) { - if (param.getExpWeights) + if (param.useTopKAsInput) { - bufferCast(*this->mPtrExpertWeightsHost)[it * param.topK + ie] + bufferCast(*this->mPtrTopKIdsHost)[it * param.topK + ie] + = static_cast(finalTopkExperts[ie].idx); + bufferCast(*this->mPtrTopKWeightsHost)[it * param.topK + ie] = static_cast(finalTopkExperts[ie].score); } + else if (param.getExpWeights) + { + bufferCast(*this->mPtrTopKWeightsHost)[it * param.topK + ie] + = static_cast(finalTopkExperts[ie].score); + } + PackedType si{static_cast(finalTopkExperts[ie].score), finalTopkExperts[ie].idx}; - reinterpret_cast(bufferCast(*this->mPtrExpertIdxHost))[it * param.topK + ie] = si; + reinterpret_cast(bufferCast(*this->mPtrTopKPackedHost))[it * param.topK + ie] = si; } } } @@ -163,13 +180,24 @@ class RoutingDeepSeekKernelTest : public RoutingKernelTest { RoutingKernelTest::setCommonParams(param, routingData); routingData.mDtypeExpW = btg::Dtype::Bfloat16; - routingData.mPtrScores = bufferCast(*this->mPtrScoresDevice); + routingData.mPtrRoutingBias = bufferCast(*this->mPtrRoutingBiasDevice); routingData.mNumExpertGroups = param.nGroup; routingData.mNumLimitedGroups = param.topkGroup; routingData.mRouteScale = param.routedScalingFactor; routingData.mUseRoutingSoftmax = false; + + if (param.useTopKAsInput) + { + routingData.mPtrTopKIds = bufferCast(*this->mPtrTopKIdsDevice); + routingData.mPtrScores = nullptr; + } + else + { + routingData.mPtrTopKIds = nullptr; + routingData.mPtrScores = bufferCast(*this->mPtrScoresDevice); + } } void callTestedFunction( @@ -189,7 +217,18 @@ TYPED_TEST(RoutingDeepSeekKernelTest, ClusterLevelParallelization) /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, + /*nGroup*/ 8, /*topkGroup*/ 4, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); + this->runTest(param); +}; + +TYPED_TEST(RoutingDeepSeekKernelTest, ClusterLevelParallelizationWithTopKAsInput) +{ + RoutingKernelTestParam param(RoutingMethodType::DeepSeekV3, /*numTokens=*/1024, // 10 + /*numExperts=*/128, /*topK=*/8, + /*expertParallelization=*/1, /*expertParallelizationId=*/0, + /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*nGroup*/ 8, /*topkGroup*/ 4, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -200,7 +239,7 @@ TYPED_TEST(RoutingDeepSeekKernelTest, ClusterLevelParallelizationWithExpertParal /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/2, /*expertParallelizationId=*/1, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 8, /*topkGroup*/ 4, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -211,7 +250,7 @@ TYPED_TEST(RoutingDeepSeekKernelTest, CooperativeLevelParallelization) /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 8, /*topkGroup*/ 4, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 10); this->runTest(param); }; @@ -222,7 +261,7 @@ TYPED_TEST(RoutingDeepSeekKernelTest, CooperativeLevelParallelization) // /*numExperts=*/128, /*topK=*/8, // /*expertParallelization=*/1, /*expertParallelizationId=*/0, // /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, -// /*usePdl=*/true, /*getExpWeights=*/true, +// /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, // /*nGroup*/ 8, /*topkGroup*/ 4, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 10); // this->runTest(param); // }; @@ -233,7 +272,7 @@ TYPED_TEST(RoutingDeepSeekKernelTest, ClusterLevelParallelizationTop2) /*numExperts=*/128, /*topK=*/2, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 8, /*topkGroup*/ 4, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -244,7 +283,7 @@ TYPED_TEST(RoutingDeepSeekKernelTest, ClusterLevelParallelizationWithExpertParal /*numExperts=*/128, /*topK=*/2, /*expertParallelization=*/2, /*expertParallelizationId=*/1, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 8, /*topkGroup*/ 4, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -255,7 +294,7 @@ TYPED_TEST(RoutingDeepSeekKernelTest, CooperativeLevelParallelizationTop2) /*numExperts=*/128, /*topK=*/2, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKExpertsAsInput=*/false, /*nGroup*/ 8, /*topkGroup*/ 4, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 10); this->runTest(param); }; diff --git a/cpp/tests/unit_tests/kernels/routing/routingLlama4Test.cpp b/cpp/tests/unit_tests/kernels/routing/routingLlama4Test.cpp index 998fe45f79f3..977881d649ba 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingLlama4Test.cpp +++ b/cpp/tests/unit_tests/kernels/routing/routingLlama4Test.cpp @@ -70,15 +70,21 @@ class RoutingLlama4KernelTest : public RoutingKernelTest expIdx[ie].score = static_cast(finalScore); } - // convert back to io_dtype and store the topk expert results in hostData.mPtrExpertIdx + // convert back to io_dtype and store the topk expert results in hostData.mPtrTopKPacked for (int ie = 0; ie < param.topK; ++ie) { PackedType si{static_cast(expIdx[ie].score), expIdx[ie].idx}; - reinterpret_cast(bufferCast(*this->mPtrExpertIdxHost))[it * param.topK + ie] = si; - if (param.getExpWeights) + reinterpret_cast(bufferCast(*this->mPtrTopKPackedHost))[it * param.topK + ie] = si; + + if (param.useTopKAsInput) + { + bufferCast(*this->mPtrTopKIdsHost)[it * param.topK + ie] + = static_cast(expIdx[ie].idx); + bufferCast(*this->mPtrTopKWeightsHost)[it * param.topK + ie] = static_cast(expIdx[ie].score); + } + else if (param.getExpWeights) { - bufferCast(*this->mPtrExpertWeightsHost)[it * param.topK + ie] - = static_cast(expIdx[ie].score); + bufferCast(*this->mPtrTopKWeightsHost)[it * param.topK + ie] = static_cast(expIdx[ie].score); } } } @@ -97,8 +103,18 @@ class RoutingLlama4KernelTest : public RoutingKernelTest { RoutingKernelTest::setCommonParams(param, routingData); routingData.mDtypeExpW = btg::Dtype::Bfloat16; - routingData.mPtrScores = bufferCast(*this->mPtrScoresDevice); - routingData.mPtrExpertIdx = reinterpret_cast(bufferCast(*this->mPtrExpertIdxDevice)); + + routingData.mPtrTopKPacked = reinterpret_cast(bufferCast(*this->mPtrTopKPackedDevice)); + if (param.useTopKAsInput) + { + routingData.mPtrTopKIds = bufferCast(*this->mPtrTopKIdsDevice); + routingData.mPtrScores = nullptr; + } + else + { + routingData.mPtrTopKIds = nullptr; + routingData.mPtrScores = bufferCast(*this->mPtrScoresDevice); + } } void callTestedFunction( @@ -118,7 +134,8 @@ TYPED_TEST(RoutingLlama4KernelTest, WarpLevelParallelization) /*numExperts=*/128, /*topK=*/1, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 0.0f, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, + /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 0.0f, /*requiredComputeCapability*/ 8); this->runTest(param); }; @@ -129,7 +146,7 @@ TYPED_TEST(RoutingLlama4KernelTest, ClusterLevelParallelization) /*numExperts=*/128, /*topK=*/1, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -140,7 +157,41 @@ TYPED_TEST(RoutingLlama4KernelTest, DeviceLevelParallelization) /*numExperts=*/128, /*topK=*/1, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, + /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 8); + this->runTest(param); +}; + +TYPED_TEST(RoutingLlama4KernelTest, WarpLevelParallelizationTopKAsInput) +{ + RoutingKernelTestParam param(RoutingMethodType::Llama4, /*numTokens=*/3, + /*numExperts=*/128, /*topK=*/1, + /*expertParallelization=*/1, /*expertParallelizationId=*/0, + /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, + /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 0.0f, + /*requiredComputeCapability*/ 8); + this->runTest(param); +}; + +TYPED_TEST(RoutingLlama4KernelTest, ClusterLevelParallelizationTopKAsInput) +{ + RoutingKernelTestParam param(RoutingMethodType::Llama4, /*numTokens=*/10, + /*numExperts=*/128, /*topK=*/1, + /*expertParallelization=*/1, /*expertParallelizationId=*/0, + /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, + /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); + this->runTest(param); +}; + +TYPED_TEST(RoutingLlama4KernelTest, DeviceLevelParallelizationTopKAsInput) +{ + RoutingKernelTestParam param(RoutingMethodType::Llama4, /*numTokens=*/300, + /*numExperts=*/128, /*topK=*/1, + /*expertParallelization=*/1, /*expertParallelizationId=*/0, + /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 8); this->runTest(param); }; diff --git a/cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp b/cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp index 67cdd77ddc85..0a6187ffb4f2 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp +++ b/cpp/tests/unit_tests/kernels/routing/routingRenormalizeTest.cpp @@ -124,15 +124,20 @@ class RoutingRenormalizeKernelTest : public RoutingKernelTest } } - // convert back to io_dtype and store the topk expert results in hostData.mPtrExpertIdx + // convert back to io_dtype and store the topk expert results in hostData.mPtrTopKPacked for (int ie = 0; ie < param.topK; ++ie) { PackedType si{static_cast(expIdx[ie].score), expIdx[ie].idx}; - reinterpret_cast(bufferCast(*this->mPtrExpertIdxHost))[it * param.topK + ie] = si; - if (param.getExpWeights) + reinterpret_cast(bufferCast(*this->mPtrTopKPackedHost))[it * param.topK + ie] = si; + if (param.useTopKAsInput) { - bufferCast(*this->mPtrExpertWeightsHost)[it * param.topK + ie] - = static_cast(expIdx[ie].score); + bufferCast(*this->mPtrTopKIdsHost)[it * param.topK + ie] + = static_cast(expIdx[ie].idx); + bufferCast(*this->mPtrTopKWeightsHost)[it * param.topK + ie] = static_cast(expIdx[ie].score); + } + else if (param.getExpWeights) + { + bufferCast(*this->mPtrTopKWeightsHost)[it * param.topK + ie] = static_cast(expIdx[ie].score); } } } @@ -164,7 +169,16 @@ class RoutingRenormalizeKernelTest : public RoutingKernelTest routingData.mDoSoftmaxBeforeTopK = param.routingMethod == RoutingMethodType::RenormalizeNaive; routingData.mNormTopkProb = param.routingMethod == RoutingMethodType::RenormalizeNaive; - routingData.mPtrScores = bufferCast(*this->mPtrScoresDevice); + if (param.useTopKAsInput) + { + routingData.mPtrTopKIds = bufferCast(*this->mPtrTopKIdsDevice); + routingData.mPtrScores = nullptr; + } + else + { + routingData.mPtrTopKIds = nullptr; + routingData.mPtrScores = bufferCast(*this->mPtrScoresDevice); + } } void callTestedFunction( @@ -184,18 +198,18 @@ TYPED_TEST(RoutingRenormalizeKernelTest, BlockLevelParallelization) /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; TYPED_TEST(RoutingRenormalizeKernelTest, BlockLevelParallelizationWithExpertParallelization) { - RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/14, + RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/4, /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/2, /*expertParallelizationId=*/1, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -206,7 +220,7 @@ TYPED_TEST(RoutingRenormalizeKernelTest, ClusterLevelParallelization) /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -217,7 +231,7 @@ TYPED_TEST(RoutingRenormalizeKernelTest, ClusterLevelParallelizationWithExpertPa /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/2, /*expertParallelizationId=*/1, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -228,7 +242,7 @@ TYPED_TEST(RoutingRenormalizeKernelTest, ClusterLevelParallelizationWithRenormal /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -239,29 +253,29 @@ TYPED_TEST(RoutingRenormalizeKernelTest, DeviceLevelParallelization) /*numExperts=*/128, /*topK=*/8, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/false, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 8); this->runTest(param); }; TYPED_TEST(RoutingRenormalizeKernelTest, ClusterLevelParallelizationTop4) { - RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/10, + RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/200, /*numExperts=*/128, /*topK=*/4, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; TYPED_TEST(RoutingRenormalizeKernelTest, ClusterLevelParallelizationWithExpertParallelizationTop4) { - RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/100, + RoutingKernelTestParam param(RoutingMethodType::Renormalize, /*numTokens=*/256, /*numExperts=*/128, /*topK=*/4, /*expertParallelization=*/2, /*expertParallelizationId=*/1, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -272,7 +286,7 @@ TYPED_TEST(RoutingRenormalizeKernelTest, ClusterLevelParallelizationWithRenormal /*numExperts=*/128, /*topK=*/4, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 9); this->runTest(param); }; @@ -283,7 +297,7 @@ TYPED_TEST(RoutingRenormalizeKernelTest, DeviceLevelParallelizationTop4) /*numExperts=*/128, /*topK=*/4, /*expertParallelization=*/1, /*expertParallelizationId=*/0, /*paddingLog2=*/3, /*localExpertsStrideLog2=*/0, - /*usePdl=*/true, /*getExpWeights=*/true, + /*usePdl=*/true, /*getExpWeights=*/true, /*useTopKAsInput=*/true, /*nGroup*/ 0, /*topkGroup*/ 0, /*routedScalingFactor*/ 1.0f, /*requiredComputeCapability*/ 8); this->runTest(param); }; diff --git a/cpp/tests/unit_tests/kernels/routing/routingTest.cpp b/cpp/tests/unit_tests/kernels/routing/routingTest.cpp index f2b08a2bf7d5..b12b72663e21 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingTest.cpp +++ b/cpp/tests/unit_tests/kernels/routing/routingTest.cpp @@ -46,6 +46,7 @@ void RoutingKernelTest::allocateBuffers(RoutingKernelTestParam const& param) auto const usePdl = param.usePdl; auto const doSoftmaxBeforeTopK = param.doSoftmaxBeforeTopK; auto const normTopkProb = param.normTopkProb; + auto const useTopKAsInput = param.useTopKAsInput; int64_t countsSize = 2 * numExperts; if (param.routingMethod == RoutingMethodType::DeepSeekV3) @@ -72,8 +73,20 @@ void RoutingKernelTest::allocateBuffers(RoutingKernelTestParam const& param) = mBufferManager->gpu(ITensor::makeShape({permIdxToTokenIdxSize}), nvinfer1::DataType::kINT32); int64_t expWeightsSize = numTokens * topK; - mPtrExpertWeightsHost = mBufferManager->pinned(ITensor::makeShape({expWeightsSize}), TRTDataType::value); - mPtrExpertWeightsDevice = mBufferManager->gpu(ITensor::makeShape({expWeightsSize}), TRTDataType::value); + mPtrTopKWeightsHost = mBufferManager->pinned(ITensor::makeShape({expWeightsSize}), TRTDataType::value); + mPtrTopKWeightsDevice = mBufferManager->gpu(ITensor::makeShape({expWeightsSize}), TRTDataType::value); + + if (useTopKAsInput) + { + int64_t topKIdsSize = numTokens * topK; + mPtrTopKIdsHost = mBufferManager->pinned(ITensor::makeShape({topKIdsSize}), nvinfer1::DataType::kINT32); + mPtrTopKIdsDevice = mBufferManager->gpu(ITensor::makeShape({topKIdsSize}), nvinfer1::DataType::kINT32); + } + else + { + mPtrTopKIdsHost = nullptr; + mPtrTopKIdsDevice = nullptr; + } int64_t ctaIdxSize = numTokens * topK; mPtrCtaIdxXyToBatchIdxHost = mBufferManager->pinned(ITensor::makeShape({ctaIdxSize}), nvinfer1::DataType::kINT32); @@ -89,8 +102,8 @@ void RoutingKernelTest::allocateBuffers(RoutingKernelTestParam const& param) = mBufferManager->gpu(ITensor::makeShape({numNonExitingCtasSize}), nvinfer1::DataType::kINT32); int64_t idxSize = numTokens * topK * sizeof(PackedType); - mPtrExpertIdxHost = mBufferManager->pinned(ITensor::makeShape({idxSize}), nvinfer1::DataType::kINT8); - mPtrExpertIdxDevice = mBufferManager->gpu(ITensor::makeShape({idxSize}), nvinfer1::DataType::kINT8); + mPtrTopKPackedHost = mBufferManager->pinned(ITensor::makeShape({idxSize}), nvinfer1::DataType::kINT8); + mPtrTopKPackedDevice = mBufferManager->gpu(ITensor::makeShape({idxSize}), nvinfer1::DataType::kINT8); mCurandStatesDevice = mBufferManager->gpu(ITensor::makeShape({numTokens, sizeof(curandState_t)}), nvinfer1::DataType::kINT8); } @@ -109,7 +122,7 @@ void RoutingKernelTest::computePermutation(RoutingKernelTestParam const& para { int32_t* expertCountsHostPtr = bufferCast(*this->mPtrExpertCountsHost); - PackedType* expIdxHostPtr = reinterpret_cast(bufferCast(*this->mPtrExpertIdxHost)); + PackedType* expIdxHostPtr = reinterpret_cast(bufferCast(*this->mPtrTopKPackedHost)); auto tokenToExpertHost = mBufferManager->pinned(ITensor::makeShape({param.numTokens * param.topK}), nvinfer1::DataType::kINT32); @@ -247,7 +260,7 @@ void RoutingKernelTest::verifyExpertRoutingIndices(RoutingKernelTestParam con mStream->synchronize(); int32_t* expIdxToPermHostptr = bufferCast(*mPtrExpandedIdxToPermutedIdxHost); - PackedType* expIdxHostPtr = reinterpret_cast(bufferCast(*mPtrExpertIdxHost)); + PackedType* expIdxHostPtr = reinterpret_cast(bufferCast(*mPtrTopKPackedHost)); for (int ie = 0; ie < param.numExperts; ++ie) { @@ -283,7 +296,7 @@ void RoutingKernelTest::verifyExpertRoutingIndices(RoutingKernelTestParam con template void RoutingKernelTest::verifyResult(RoutingKernelTestParam const& param) { - auto const expertWeightsHost = mBufferManager->copyFrom(*mPtrExpertWeightsDevice, MemoryType::kCPU); + auto const expertWeightsHost = mBufferManager->copyFrom(*mPtrTopKWeightsDevice, MemoryType::kCPU); auto const expertCountsHost = mBufferManager->copyFrom(*mPtrExpertCountsDevice, MemoryType::kCPU); auto const permutedIdxSizeHost = mBufferManager->copyFrom(*mPtrPermutedIdxSizeDevice, MemoryType::kCPU); auto const numNonExitingCtasHost = mBufferManager->copyFrom(*mPtrNumNonExitingCtasDevice, MemoryType::kCPU); @@ -301,7 +314,7 @@ void RoutingKernelTest::verifyResult(RoutingKernelTestParam const& param) if (param.getExpWeights) { - EXPECT_EQ(isClose(bufferCast(*mPtrExpertWeightsHost), expertWeightsPtr, param.numTokens * param.topK, + EXPECT_EQ(isClose(bufferCast(*mPtrTopKWeightsHost), expertWeightsPtr, param.numTokens * param.topK, "expert weights"), true); } @@ -344,6 +357,15 @@ void RoutingKernelTest::runTest(RoutingKernelTestParam const& param) // Setup buffers setupBuffers(param); + // Call host function + callHostFunction(param); + if (param.useTopKAsInput) + { + // Set the topk_ids as input + mBufferManager->copy(*mPtrTopKIdsHost, *mPtrTopKIdsDevice); + mBufferManager->copy(*mPtrTopKWeightsHost, *mPtrTopKWeightsDevice); + mStream->synchronize(); + } // Retrieve the workspace size of the routing kernel. auto const workspaceSize = getDeviceWorkspaceSize(param); TensorPtr workspaceDevice @@ -351,10 +373,6 @@ void RoutingKernelTest::runTest(RoutingKernelTestParam const& param) // Call tested function routing callTestedFunction(param, workspaceDevice); - - // Call host function - callHostFunction(param); - // Verify results verifyResult(param); } diff --git a/cpp/tests/unit_tests/kernels/routing/routingTest.h b/cpp/tests/unit_tests/kernels/routing/routingTest.h index 5f437054d12a..1f37f21ff765 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingTest.h +++ b/cpp/tests/unit_tests/kernels/routing/routingTest.h @@ -220,6 +220,10 @@ struct RoutingKernelTestParam bool getExpWeights{true}; int requiredComputeCapability{9}; + + // Check the input parameters + bool useTopKAsInput{false}; + // Special for renormalize routing method bool doSoftmaxBeforeTopK{false}; bool normTopkProb{true}; @@ -243,8 +247,8 @@ struct RoutingKernelTestParam // Constructor with all parameters RoutingKernelTestParam(RoutingMethodType routingMethod, int32_t numTokens, int32_t numExperts, uint32_t topK, int32_t expertParallelization = 1, int32_t expertParallelizationId = 0, int32_t paddingLog2 = 3, - int32_t localExpertsStrideLog2 = 0, bool usePdl = true, bool getExpWeights = true, int32_t nGroup = 1, - int32_t topkGroup = 1, float routedScalingFactor = 1.0f, int requiredComputeCapability = 9) + int32_t localExpertsStrideLog2 = 0, bool usePdl = true, bool getExpWeights = true, bool useTopKAsInput = false, + int32_t nGroup = 1, int32_t topkGroup = 1, float routedScalingFactor = 1.0f, int requiredComputeCapability = 9) : routingMethod(routingMethod) , numTokens(numTokens) , numExperts(numExperts) @@ -253,6 +257,7 @@ struct RoutingKernelTestParam , localExpertsStrideLog2(localExpertsStrideLog2) , usePdl(usePdl) , getExpWeights(getExpWeights) + , useTopKAsInput(useTopKAsInput) , nGroup(nGroup) , topkGroup(topkGroup) , routedScalingFactor(routedScalingFactor) @@ -393,8 +398,8 @@ class RoutingKernelTest : public testing::Test routingData.mPtrPermutedIdxSize = bufferCast(*mPtrPermutedIdxSizeDevice); routingData.mPtrExpandedIdxToPermutedIdx = bufferCast(*mPtrExpandedIdxToPermutedIdxDevice); routingData.mPtrPermutedIdxToTokenIdx = bufferCast(*mPtrPermutedIdxToTokenIdxDevice); - routingData.mPtrExpertWeights = bufferCast(*mPtrExpertWeightsDevice); - routingData.mPtrExpertIdx = reinterpret_cast(bufferCast(*mPtrExpertIdxDevice)); + routingData.mPtrTopKWeights = bufferCast(*mPtrTopKWeightsDevice); + routingData.mPtrTopKPacked = reinterpret_cast(bufferCast(*mPtrTopKPackedDevice)); // Set grouped gemm launch config buffers routingData.mPtrCtaIdxXyToBatchIdx = bufferCast(*mPtrCtaIdxXyToBatchIdxDevice); @@ -413,7 +418,7 @@ class RoutingKernelTest : public testing::Test struct cudaDeviceProp mDeviceProp; - // optional: if `nullptr`, `mPtrExpertIdx` must be provided. + // optional: if `nullptr`, `mPtrTopKPacked` must be provided. // If it is given, it represents the scores without sigmoid activation for // each token and expert. // note: if it is provided, we always re-compute the top1 scores @@ -426,8 +431,19 @@ class RoutingKernelTest : public testing::Test // the least significant 16 bits represent the index of the chosen expert (unsigned). // note: this is required if the number of tokens is large. // dim: [mNumTokens, mTopK] - TensorPtr mPtrExpertIdxHost; - TensorPtr mPtrExpertIdxDevice; + TensorPtr mPtrTopKPackedHost; + TensorPtr mPtrTopKPackedDevice; + + // optional: Add another input format. + // dim: [mNumTokens, mTopK] + TensorPtr mPtrTopKIdsHost; + TensorPtr mPtrTopKIdsDevice; + + // optional: if `nullptr`, it is not filled + // dim: [mNumTokens, mTopK] + // Note: this might be reused as input when we take topk_ids as input. + TensorPtr mPtrTopKWeightsHost; + TensorPtr mPtrTopKWeightsDevice; // note: at least one of the optional outputs below must be provided // optional: only used as an intermediate buffer when the number of tokens is large. @@ -445,10 +461,7 @@ class RoutingKernelTest : public testing::Test // dim: [mNumTokens * mTopK + (mNumExperts << mPaddingLog2) - mNumExperts] TensorPtr mPtrPermutedIdxToTokenIdxHost; TensorPtr mPtrPermutedIdxToTokenIdxDevice; - // optional: if `nullptr`, it is not filled - // dim: [mNumTokens, mTopK] - TensorPtr mPtrExpertWeightsHost; - TensorPtr mPtrExpertWeightsDevice; + // // Grouped Gemm Launch Config Buffers // diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 07ed20be3e50..284b0f4e7b08 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -488,17 +488,19 @@ def _(input, residual, norm_weight, expanded_idx_to_permuted_idx, ] @torch.library.register_fake("trtllm::renorm_moe_routing_op") - def _(router_logits, topk): + def _(router_logits, topk, output_dtype: torch.dtype = None): num_tokens = router_logits.shape[0] sz = (num_tokens, topk) + output_dtype = output_dtype or torch.float32 return router_logits.new_empty( sz, dtype=torch.int32), router_logits.new_empty(sz, - dtype=torch.float32) + dtype=output_dtype) @torch.library.register_fake("trtllm::default_moe_routing_op") - def _(router_logits, topk): + def _(router_logits, topk, output_dtype: torch.dtype = None): num_tokens = router_logits.shape[0] sz = (num_tokens, topk) + output_dtype = output_dtype or torch.float32 return router_logits.new_empty( sz, dtype=torch.int32), router_logits.new_empty(sz, - dtype=torch.float32) + dtype=output_dtype) diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 4f90c3fbb710..65196525276a 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -34,7 +34,7 @@ def calculate_tile_tokens_dim( @dataclass(frozen=True) class FP4BlockScaleMoEInputs: - routing_logits: torch.Tensor + routing_logits: Optional[torch.Tensor] routing_bias: Optional[torch.Tensor] hidden_states: torch.Tensor hidden_states_scale: torch.Tensor @@ -45,6 +45,8 @@ class FP4BlockScaleMoEInputs: output1_scale_scalar: torch.Tensor output1_scale_gate_scalar: torch.Tensor output2_scale_scalar: torch.Tensor + topk_weights: Optional[torch.Tensor] = None + topk_ids: Optional[torch.Tensor] = None class FP4BlockScaleMoERunner(TunableRunner): @@ -114,9 +116,10 @@ def forward( tactic: int = -1, ) -> torch.Tensor: - kernel_runner = self.get_runner(inputs[0].shape[0]) args = FP4BlockScaleMoEInputs(*inputs) + kernel_runner = self.get_runner(args.hidden_states.shape[0]) + return kernel_runner.run_moe( args.routing_logits, args.routing_bias, args.hidden_states, args.hidden_states_scale, args.gemm1_weights, @@ -126,7 +129,8 @@ def forward( self.num_experts, self.top_k, self.n_group, self.topk_group, self.intermediate_size, self.local_expert_offset, self.local_num_experts, self.routed_scaling_factor, - self.routing_method_type, self.do_finalize, tactic) + self.routing_method_type, self.do_finalize, tactic, + args.topk_weights, args.topk_ids) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs) -> List[int]: @@ -227,28 +231,29 @@ def get_tuning_config(cls) -> TuningConfig: @torch.library.custom_op("trtllm::fp4_block_scale_moe_runner", mutates_args=()) def fp4_block_scale_moe_runner( - routing_logits: torch.Tensor, - routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - output1_scale_scalar: torch.Tensor, - output1_scale_gate_scalar: torch.Tensor, - output2_scale_scalar: torch.Tensor, - num_experts: int, - top_k: int, - n_group: Optional[int], - topk_group: Optional[int], - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: Optional[float], - routing_method_type: int, - do_finalize: bool, -) -> List[torch.Tensor]: + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + output1_scale_scalar: torch.Tensor, + output1_scale_gate_scalar: torch.Tensor, + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + do_finalize: bool, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None) -> List[torch.Tensor]: tuner = AutoTuner.get() @@ -265,8 +270,17 @@ def fp4_block_scale_moe_runner( do_finalize, ) - inputs = [ - routing_logits, + # Use dummy routing logits for autotuner + if routing_logits is None: + routing_logits_for_tuner = torch.randn(hidden_states.shape[0], + num_experts, + dtype=torch.bfloat16, + device=hidden_states.device) + else: + routing_logits_for_tuner = routing_logits + + input_tensors_for_tuner = [ + routing_logits_for_tuner, routing_bias, hidden_states, hidden_states_scale, @@ -283,10 +297,13 @@ def fp4_block_scale_moe_runner( "trtllm::fp4_block_scale_moe_runner", [kernel_runner], kernel_runner.tuning_config, - inputs, + input_tensors_for_tuner, ) - return kernel_runner(inputs, tactic=best_tactic) + input_tensors = input_tensors_for_tuner + [topk_weights, topk_ids] + input_tensors[ + 0] = routing_logits # replace dummy routing logits with actual routing logits + return kernel_runner(input_tensors, tactic=best_tactic) def fp4_block_scale_fake_output_without_finalize( @@ -315,29 +332,29 @@ def fp4_block_scale_fake_output_without_finalize( @fp4_block_scale_moe_runner.register_fake -def _( - routing_logits, - routing_bias, - hidden_states, - hidden_states_scale, - gemm1_weights, - gemm1_weights_scale, - gemm2_weights, - gemm2_weights_scale, - output1_scale_scalar, - output1_scale_gate_scalar, - output2_scale_scalar, - num_experts, - top_k, - n_group, - topk_group, - intermediate_size, - local_expert_offset, - local_num_experts, - routed_scaling_factor, - routing_method_type, - do_finalize, -) -> List[torch.Tensor]: +def _(routing_logits, + routing_bias, + hidden_states, + hidden_states_scale, + gemm1_weights, + gemm1_weights_scale, + gemm2_weights, + gemm2_weights_scale, + output1_scale_scalar, + output1_scale_gate_scalar, + output2_scale_scalar, + num_experts, + top_k, + n_group, + topk_group, + intermediate_size, + local_expert_offset, + local_num_experts, + routed_scaling_factor, + routing_method_type, + do_finalize, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None) -> List[torch.Tensor]: if do_finalize: num_tokens = hidden_states.shape[0] hidden_size = hidden_states.shape[1] * 2 @@ -357,7 +374,7 @@ def _( @dataclass(frozen=True) class FP8BlockScaleMoEInputs: - routing_logits: torch.Tensor + routing_logits: Optional[torch.Tensor] routing_bias: torch.Tensor hidden_states: torch.Tensor hidden_states_scale: torch.Tensor @@ -365,6 +382,8 @@ class FP8BlockScaleMoEInputs: gemm1_weights_scale: torch.Tensor gemm2_weights: torch.Tensor gemm2_weights_scale: torch.Tensor + topk_weights: Optional[torch.Tensor] = None + topk_ids: Optional[torch.Tensor] = None class FP8BlockScaleMoERunner(TunableRunner): @@ -442,7 +461,7 @@ def forward( args = FP8BlockScaleMoEInputs(*inputs) - kernel_runner = self.get_runner(inputs[0].shape[0]) + kernel_runner = self.get_runner(args.hidden_states.shape[0]) return kernel_runner.run_moe( args.routing_logits, args.routing_bias, args.hidden_states, @@ -451,7 +470,8 @@ def forward( args.gemm2_weights_scale, self.num_experts, self.top_k, self.n_group, self.topk_group, self.intermediate_size, self.local_expert_offset, self.local_num_experts, - self.routed_scaling_factor, self.routing_method_type, tactic) + self.routed_scaling_factor, self.routing_method_type, tactic, + args.topk_weights, args.topk_ids) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs) -> List[int]: @@ -532,24 +552,25 @@ def get_tuning_config(cls) -> TuningConfig: @torch.library.custom_op("trtllm::fp8_block_scale_moe_runner", mutates_args=()) def fp8_block_scale_moe_runner( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - num_experts: int, - top_k: int, - n_group: int, - topk_group: int, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: float, - routing_method_type: int, -) -> torch.Tensor: + routing_logits: Optional[torch.Tensor], + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: float, + routing_method_type: int, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None) -> torch.Tensor: tuner = AutoTuner.get() @@ -565,8 +586,17 @@ def fp8_block_scale_moe_runner( routing_method_type, ) - inputs = [ - routing_logits, + # Use dummy routing logits for autotuner + if routing_logits is None: + routing_logits_for_tuner = torch.randn(hidden_states.shape[0], + num_experts, + dtype=torch.bfloat16, + device=hidden_states.device) + else: + routing_logits_for_tuner = routing_logits + + input_tensors_for_tuner = [ + routing_logits_for_tuner, routing_bias, hidden_states, hidden_states_scale, @@ -580,32 +610,35 @@ def fp8_block_scale_moe_runner( "trtllm::fp8_block_scale_moe_runner", [kernel_runner], kernel_runner.tuning_config, - inputs, + input_tensors_for_tuner, ) - return kernel_runner(inputs, tactic=best_tactic) + input_tensors = input_tensors_for_tuner + [topk_weights, topk_ids] + input_tensors[ + 0] = routing_logits # replace dummy routing logits with actual routing logits + return kernel_runner(input_tensors, tactic=best_tactic) @fp8_block_scale_moe_runner.register_fake -def _( - routing_logits: torch.Tensor, - routing_bias: torch.Tensor, - hidden_states: torch.Tensor, - hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, - gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, - num_experts: int, - top_k: int, - n_group: int, - topk_group: int, - intermediate_size: int, - local_expert_offset: int, - local_num_experts: int, - routed_scaling_factor: float, - routing_method_type: int, -) -> torch.Tensor: +def _(routing_logits: torch.Tensor, + routing_bias: torch.Tensor, + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: float, + routing_method_type: int, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None) -> torch.Tensor: num_tokens = hidden_states.shape[0] hidden_size = hidden_states.shape[1] * 2 @@ -615,7 +648,7 @@ def _( @dataclass(frozen=True) class MxE4m3MxE2m1BlockScaleMoEInputs: - routing_logits: torch.Tensor + routing_logits: Optional[torch.Tensor] routing_bias: Optional[torch.Tensor] hidden_states: torch.Tensor hidden_states_scale: torch.Tensor @@ -628,6 +661,8 @@ class MxE4m3MxE2m1BlockScaleMoEInputs: gemm2_weights: torch.Tensor gemm2_weights_scale: torch.Tensor gemm2_bias: Optional[torch.Tensor] + topk_weights: Optional[torch.Tensor] = None + topk_ids: Optional[torch.Tensor] = None class MxE4m3MxE2m1BlockScaleMoERunner(TunableRunner): @@ -708,8 +743,7 @@ def forward( args = MxE4m3MxE2m1BlockScaleMoEInputs(*inputs) - kernel_runner = self.get_runner(inputs[0].shape[0]) - + kernel_runner = self.get_runner(args.hidden_states.shape[0]) return kernel_runner.run_moe( args.routing_logits, args.routing_bias, args.hidden_states, args.hidden_states_scale, args.gemm1_weights, @@ -719,7 +753,8 @@ def forward( self.num_experts, self.top_k, self.n_group, self.topk_group, self.intermediate_size, self.hidden_size_output, self.local_expert_offset, self.local_num_experts, - self.routed_scaling_factor, self.routing_method_type, tactic) + self.routed_scaling_factor, self.routing_method_type, tactic, + args.topk_weights, args.topk_ids) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs) -> List[int]: @@ -810,19 +845,33 @@ def get_tuning_config(cls) -> TuningConfig: @torch.library.custom_op("trtllm::mxe4m3_mxe2m1_block_scale_moe_runner", mutates_args=()) def mxe4m3_mxe2m1_block_scale_moe_runner( - routing_logits: torch.Tensor, routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, hidden_states_scale: torch.Tensor, - gemm1_weights: torch.Tensor, gemm1_weights_scale: torch.Tensor, - gemm1_bias: Optional[torch.Tensor], gemm1_alpha: Optional[torch.Tensor], + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + hidden_states_scale: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], gemm1_beta: Optional[torch.Tensor], - gemm1_clamp_limit: Optional[torch.Tensor], gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, gemm2_bias: Optional[torch.Tensor], - num_experts: int, top_k: int, n_group: Optional[int], - topk_group: Optional[int], intermediate_size: int, - hidden_size_output: int, local_expert_offset: int, - local_num_experts: int, routed_scaling_factor: Optional[float], - routing_method_type: int, act_type: int, - imbalance_factor: float) -> torch.Tensor: + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + hidden_size_output: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + act_type: int, + imbalance_factor: float, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None) -> torch.Tensor: tuner = AutoTuner.get() @@ -831,8 +880,17 @@ def mxe4m3_mxe2m1_block_scale_moe_runner( hidden_size_output, local_expert_offset, local_num_experts, routed_scaling_factor, routing_method_type, act_type, imbalance_factor) - input_tensors = [ - routing_logits, + # Use dummy routing logits for autotuner + if routing_logits is None: + routing_logits_for_tuner = torch.randn(hidden_states.shape[0], + num_experts, + dtype=torch.bfloat16, + device=hidden_states.device) + else: + routing_logits_for_tuner = routing_logits + + input_tensors_for_tuner = [ + routing_logits_for_tuner, routing_bias, hidden_states, hidden_states_scale, @@ -851,15 +909,18 @@ def mxe4m3_mxe2m1_block_scale_moe_runner( "trtllm::mxe4m3_mxe2m1_block_scale_moe_runner", [kernel_runner], kernel_runner.tuning_config, - input_tensors, + input_tensors_for_tuner, ) + input_tensors = input_tensors_for_tuner + [topk_weights, topk_ids] + input_tensors[ + 0] = routing_logits # replace dummy routing logits with actual routing logits return kernel_runner(input_tensors, tactic=best_tactic) @dataclass(frozen=True) class E4m3MxE2m1BlockScaleMoEInputs: - routing_logits: torch.Tensor + routing_logits: Optional[torch.Tensor] routing_bias: Optional[torch.Tensor] hidden_states: torch.Tensor gemm1_weights: torch.Tensor @@ -874,6 +935,8 @@ class E4m3MxE2m1BlockScaleMoEInputs: output1_scale_scalar: torch.Tensor output1_scale_gate_scalar: torch.Tensor output2_scale_scalar: torch.Tensor + topk_weights: Optional[torch.Tensor] = None + topk_ids: Optional[torch.Tensor] = None class E4m3MxE2m1BlockScaleMoERunner(TunableRunner): @@ -949,8 +1012,8 @@ def forward( ) -> torch.Tensor: args = E4m3MxE2m1BlockScaleMoEInputs(*inputs) - kernel_runner = self.get_runner(inputs[0].shape[0]) + kernel_runner = self.get_runner(args.hidden_states.shape[0]) return kernel_runner.run_moe( args.routing_logits, args.routing_bias, args.hidden_states, None, args.gemm1_weights, args.gemm1_weights_scale, args.gemm1_bias, @@ -960,7 +1023,8 @@ def forward( args.output2_scale_scalar, self.num_experts, self.top_k, self.n_group, self.topk_group, self.intermediate_size, None, self.local_expert_offset, self.local_num_experts, - self.routed_scaling_factor, self.routing_method_type, tactic) + self.routed_scaling_factor, self.routing_method_type, tactic, + args.topk_weights, args.topk_ids) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs) -> List[int]: @@ -1031,20 +1095,34 @@ def get_tuning_config(cls) -> TuningConfig: @torch.library.custom_op("trtllm::e4m3_mxe2m1_block_scale_moe_runner", mutates_args=()) def e4m3_mxe2m1_block_scale_moe_runner( - routing_logits: torch.Tensor, routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, gemm1_bias: Optional[torch.Tensor], - gemm1_alpha: Optional[torch.Tensor], gemm1_beta: Optional[torch.Tensor], - gemm1_clamp_limit: Optional[torch.Tensor], gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, gemm2_bias: Optional[torch.Tensor], + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], output1_scale_scalar: torch.Tensor, output1_scale_gate_scalar: torch.Tensor, - output2_scale_scalar: torch.Tensor, num_experts: int, top_k: int, - n_group: Optional[int], topk_group: Optional[int], - intermediate_size: int, local_expert_offset: int, - local_num_experts: int, routed_scaling_factor: Optional[float], - routing_method_type: int, act_type: int, - imbalance_factor: float) -> torch.Tensor: + output2_scale_scalar: torch.Tensor, + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + act_type: int, + imbalance_factor: float, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None) -> torch.Tensor: tuner = AutoTuner.get() @@ -1062,8 +1140,17 @@ def e4m3_mxe2m1_block_scale_moe_runner( imbalance_factor, ) - input_tensors = [ - routing_logits, + # Use dummy routing logits for autotuner + if routing_logits is None: + routing_logits_for_tuner = torch.randn(hidden_states.shape[0], + num_experts, + dtype=torch.bfloat16, + device=hidden_states.device) + else: + routing_logits_for_tuner = routing_logits + + input_tensors_for_tuner = [ + routing_logits_for_tuner, routing_bias, hidden_states, gemm1_weights, @@ -1084,15 +1171,19 @@ def e4m3_mxe2m1_block_scale_moe_runner( "trtllm::e4m3_mxe2m1_block_scale_moe_runner", [kernel_runner], kernel_runner.tuning_config, - input_tensors, + input_tensors_for_tuner, ) + # Add topk tensors for final execution + input_tensors = input_tensors_for_tuner + [topk_weights, topk_ids] + input_tensors[ + 0] = routing_logits # replace dummy routing logits with actual routing logits return kernel_runner(input_tensors, tactic=best_tactic) @dataclass(frozen=True) class Bf16MxE2m1BlockScaleMoEInputs: - routing_logits: torch.Tensor + routing_logits: Optional[torch.Tensor] routing_bias: Optional[torch.Tensor] hidden_states: torch.Tensor gemm1_weights: torch.Tensor @@ -1104,6 +1195,8 @@ class Bf16MxE2m1BlockScaleMoEInputs: gemm2_weights: torch.Tensor gemm2_weights_scale: torch.Tensor gemm2_bias: Optional[torch.Tensor] + topk_weights: Optional[torch.Tensor] = None + topk_ids: Optional[torch.Tensor] = None class Bf16MxE2m1BlockScaleMoERunner(TunableRunner): @@ -1179,8 +1272,8 @@ def forward( ) -> torch.Tensor: args = Bf16MxE2m1BlockScaleMoEInputs(*inputs) - kernel_runner = self.get_runner(inputs[0].shape[0]) + kernel_runner = self.get_runner(args.hidden_states.shape[0]) return kernel_runner.run_moe( args.routing_logits, args.routing_bias, args.hidden_states, args.gemm1_weights, args.gemm1_weights_scale, args.gemm1_bias, @@ -1189,7 +1282,7 @@ def forward( self.num_experts, self.top_k, self.n_group, self.topk_group, self.intermediate_size, self.local_expert_offset, self.local_num_experts, self.routed_scaling_factor, - self.routing_method_type, tactic) + self.routing_method_type, tactic, args.topk_weights, args.topk_ids) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs) -> List[int]: @@ -1260,17 +1353,31 @@ def get_tuning_config(cls) -> TuningConfig: @torch.library.custom_op("trtllm::bf16_mxe2m1_block_scale_moe_runner", mutates_args=()) def bf16_mxe2m1_block_scale_moe_runner( - routing_logits: torch.Tensor, routing_bias: Optional[torch.Tensor], - hidden_states: torch.Tensor, gemm1_weights: torch.Tensor, - gemm1_weights_scale: torch.Tensor, gemm1_bias: Optional[torch.Tensor], - gemm1_alpha: Optional[torch.Tensor], gemm1_beta: Optional[torch.Tensor], - gemm1_clamp_limit: Optional[torch.Tensor], gemm2_weights: torch.Tensor, - gemm2_weights_scale: torch.Tensor, gemm2_bias: Optional[torch.Tensor], - num_experts: int, top_k: int, n_group: Optional[int], - topk_group: Optional[int], intermediate_size: int, - local_expert_offset: int, local_num_experts: int, - routed_scaling_factor: Optional[float], routing_method_type: int, - act_type: int, imbalance_factor: float) -> torch.Tensor: + routing_logits: Optional[torch.Tensor], + routing_bias: Optional[torch.Tensor], + hidden_states: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm1_bias: Optional[torch.Tensor], + gemm1_alpha: Optional[torch.Tensor], + gemm1_beta: Optional[torch.Tensor], + gemm1_clamp_limit: Optional[torch.Tensor], + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm2_bias: Optional[torch.Tensor], + num_experts: int, + top_k: int, + n_group: Optional[int], + topk_group: Optional[int], + intermediate_size: int, + local_expert_offset: int, + local_num_experts: int, + routed_scaling_factor: Optional[float], + routing_method_type: int, + act_type: int, + imbalance_factor: float, + topk_weights: Optional[torch.Tensor] = None, + topk_ids: Optional[torch.Tensor] = None) -> torch.Tensor: tuner = AutoTuner.get() @@ -1288,8 +1395,17 @@ def bf16_mxe2m1_block_scale_moe_runner( imbalance_factor, ) - input_tensors = [ - routing_logits, + # Use dummy routing logits for autotuner + if routing_logits is None: + routing_logits_for_tuner = torch.randn(hidden_states.shape[0], + num_experts, + dtype=torch.bfloat16, + device=hidden_states.device) + else: + routing_logits_for_tuner = routing_logits + + input_tensors_for_tuner = [ + routing_logits_for_tuner, routing_bias, hidden_states, gemm1_weights, @@ -1303,11 +1419,16 @@ def bf16_mxe2m1_block_scale_moe_runner( gemm2_bias, ] + # Choose best tactic using autotuner _, best_tactic = tuner.choose_one( "trtllm::bf16_mxe2m1_block_scale_moe_runner", [kernel_runner], kernel_runner.tuning_config, - input_tensors, + input_tensors_for_tuner, ) + # Add topk tensors for final execution + input_tensors = input_tensors_for_tuner + [topk_weights, topk_ids] + input_tensors[ + 0] = routing_logits # replace dummy routing logits with actual routing logits return kernel_runner(input_tensors, tactic=best_tactic) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 2dd93b366275..ac10ab2ae91b 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -52,14 +52,13 @@ from ..attention_backend import AttentionMetadata from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, - MoEAllReduce, MoEAllReduceParams, allgather) + MoEAllReduce, MoEAllReduceParams) from ..model_config import ModelConfig from ..modules.attention import MLA from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import (DeepSeekV3MoeRoutingMethod, - MoEWeightLoadingMode, TRTLLMGenFusedMoE, - create_moe) + MoEWeightLoadingMode, create_moe) from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode, WeightsLoadingConfig from ..modules.multi_stream_utils import maybe_execute_in_parallel @@ -854,13 +853,6 @@ def compute_routed_output(self, hidden_states, hidden_states_fp4, all_rank_num_tokens, do_finalize): # max-throughput use_dp_padding = False - if self.use_dp and self.mapping.tp_size > 1: - if isinstance(self.experts, TRTLLMGenFusedMoE): - hidden_states = allgather(hidden_states, - self.mapping, - dim=0, - sizes=all_rank_num_tokens) - router_logits = self.gate(hidden_states) routed_output = self.experts( diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index d0e5a1f0848d..8f464d914360 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py @@ -23,7 +23,7 @@ # isort: off from ..modules.fused_moe import (MoE, MoEWeightLoadingMode, RenormalizeMoeRoutingMethod, TritonFusedMoE, - TRTLLMGenFusedMoE, create_moe) + create_moe) from ..modules.fused_moe.routing import (get_cached_perfect_router_logits, precompute_common_perfect_router_logits ) @@ -151,7 +151,10 @@ def __init__( ) self.routing_method = RenormalizeMoeRoutingMethod( - top_k=pretrained_config.num_experts_per_tok) + top_k=pretrained_config.num_experts_per_tok, + output_dtype=torch.bfloat16 + if config.moe_backend.upper() == "TRTLLM" else torch.float32) + self.swiglu_alpha = torch.tensor( [1.702] * (self.num_experts // config.mapping.moe_ep_size), dtype=torch.float32).cuda() @@ -260,7 +263,7 @@ def forward_attn_dp( all_rank_num_tokens = attn_metadata.all_rank_num_tokens if self.mapping.tp_size > 1 and all_rank_num_tokens is not None: - if (isinstance(self.experts, (TRTLLMGenFusedMoE, TritonFusedMoE))): + if (isinstance(self.experts, (TritonFusedMoE))): t = allgather(t, self.mapping, dim=0, sizes=all_rank_num_tokens) g = self.gate(t) @@ -273,7 +276,7 @@ def forward_attn_dp( g = self._create_ideal_expert_load_balanced_logits( num_tokens=num_tokens, num_experts=num_experts, device=x.device) - # Let CutlassFusedMoE handle allgather internally + # Let CutlassFusedMoE and TRTLLMGenFusedMoE handle allgather internally # Pass the normalized tensor (t) as input to experts, not x expert_output = self.experts(x=t, router_logits=g, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py index 221478edf8ee..8b6a7b957441 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py @@ -10,15 +10,14 @@ from ..attention_backend import AttentionMetadata from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, - MoEAllReduce, MoEAllReduceParams, allgather) + MoEAllReduce, MoEAllReduceParams) from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import (BaseMoeRoutingMethod, RenormalizeMoeRoutingMethod, RenormalizeNaiveMoeRoutingMethod, - RoutingMethodType, TRTLLMGenFusedMoE, - create_moe) + RoutingMethodType, create_moe) from ..modules.linear import TensorParallelMode from ..modules.rms_norm import RMSNorm from ..speculative import SpecMetadata @@ -46,6 +45,7 @@ def __init__( dtype=dtype), requires_grad=False) self.routing_method_type = routing_method_type + self.moe_backend = moe_backend # FIXME: out_dtype=float32 does not work # self.out_dtype = torch.float32 if moe_backend == "TRTLLM" else dtype self.out_dtype = dtype @@ -65,9 +65,15 @@ def load_weights(self, weights: List[Dict]): @property def routing_method(self) -> BaseMoeRoutingMethod: if self.routing_method_type == RoutingMethodType.RenormalizeNaive: - return RenormalizeNaiveMoeRoutingMethod(top_k=self.top_k) + return RenormalizeNaiveMoeRoutingMethod( + top_k=self.top_k, + output_dtype=torch.bfloat16 + if self.moe_backend.upper() == "TRTLLM" else torch.float32) elif self.routing_method_type == RoutingMethodType.Renormalize: - return RenormalizeMoeRoutingMethod(top_k=self.top_k) + return RenormalizeMoeRoutingMethod( + top_k=self.top_k, + output_dtype=torch.bfloat16 + if self.moe_backend.upper() == "TRTLLM" else torch.float32) else: raise ValueError( f"Unsupported routing method: {self.routing_method_type}") @@ -131,13 +137,6 @@ def forward( if not do_finalize: assert not self.enable_attention_dp - if self.enable_attention_dp and self.mapping.tp_size > 1: - if isinstance(self.experts, TRTLLMGenFusedMoE): - hidden_states = allgather(hidden_states, - self.mapping, - dim=0, - sizes=all_rank_num_tokens) - router_logits = self.gate(hidden_states) final_hidden_states = self.experts( hidden_states, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 47a2fffb4ba6..dbda1e4cb0f1 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -7,8 +7,9 @@ from ...custom_ops.trtllm_gen_custom_ops import \ fp4_block_scale_fake_output_without_finalize +from ...distributed import allgather from ...model_config import ModelConfig -from ...utils import Fp4QuantizedTensor +from ...utils import Fp4QuantizedTensor, ceil_div from .interface import MoE, MoEWeightLoadingMode from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethod, NVFP4TRTLLMGenFusedMoEMethod, @@ -195,6 +196,66 @@ def forward_impl( topk_group = None routed_scaling_factor = None + # Don't support post-quant allgather for fp8 block scale and has_w4a16_mxfp4 for now. + is_post_quant_allgather_supported = self.has_nvfp4 or self.has_w4a8_mxfp4_fp8 or self.has_w4a8_mxfp4_mxfp8 + run_post_quant_allgather = self.use_dp and self.parallel_size > 1 and is_post_quant_allgather_supported + + x_sf = None + token_selected_experts = None + token_final_scales = None + x_row = x.shape[0] + x_col = x.shape[1] + if run_post_quant_allgather: + # apply routing + token_selected_experts, token_final_scales = self.routing_method.apply( + router_logits) + token_final_scales = token_final_scales.to(torch.bfloat16) + assert token_final_scales.dtype == torch.bfloat16 + assert token_selected_experts.dtype == torch.int32 + # quantize inputs + if self.has_w4a8_mxfp4_fp8: + x, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + x, self.fc31_input_dequant[0]) + # Update x_row and x_col to the padded shape + x_row, x_col = x.shape[0], x.shape[1] + elif self.has_nvfp4: + if isinstance(x, Fp4QuantizedTensor): + assert not x.is_sf_swizzled, "Fp4QuantizedTensor should not be swizzled before communication" + x_row = x.shape[0] + # note: we use uint8 to store 2 fp4 values + x_col = x.shape[1] * 2 + x, x_sf = x.fp4_tensor, x.scaling_factor + else: + x_row = x.shape[0] + x_col = x.shape[1] + x, x_sf = torch.ops.trtllm.fp4_quantize( + x, self.fc31_input_scale, self.scaling_vector_size, + False, False) + elif self.has_w4a8_mxfp4_mxfp8: + x, x_sf = torch.ops.trtllm.mxfp8_quantize( + x, False, alignment=self.quant_method.weight_alignment) + # Update x_row and x_col to the padded shape + x_row, x_col = x.shape[0], x.shape[1] + else: + raise ValueError( + f"unsupported quantization mode with run_post_quant_allgather: {self.quant_config.quant_mode}" + ) + + #allgather for attention DP + if x_sf is not None: + x_sf = x_sf.view(x_row, ceil_div(x_col, + self.scaling_vector_size)) + assert len( + x_sf.shape + ) == 2, "The hidden states scaling factor should be 2D tensor before allgather" + x, x_sf, token_selected_experts, token_final_scales = allgather( + [x, x_sf, token_selected_experts, token_final_scales], + self.mapping, + dim=0, + sizes=None if use_dp_padding else all_rank_num_tokens) + if x_sf is not None: + x_sf = x_sf.flatten() + # TODO: since routing kernel is integrated into moe_runner for fp8, # here we just route the I/Os for moe_runner if self.has_deepseek_fp8_block_scales: @@ -224,18 +285,22 @@ def forward_impl( elif self.has_nvfp4: scale_factor_use_ue8m0 = False is_scale_factor_swizzled = False # use linear layout here - hidden_states_fp4, hidden_states_scale_linear_fp4 = ( - torch.ops.trtllm.fp4_quantize( - x, - self.fc31_input_scale, - self.scaling_vector_size, - scale_factor_use_ue8m0, - is_scale_factor_swizzled, - )) + + if not run_post_quant_allgather: + hidden_states_fp4, hidden_states_scale_linear_fp4 = ( + torch.ops.trtllm.fp4_quantize( + x, + self.fc31_input_scale, + self.scaling_vector_size, + scale_factor_use_ue8m0, + is_scale_factor_swizzled, + )) + else: + hidden_states_fp4, hidden_states_scale_linear_fp4 = x, x_sf outputs = torch.ops.trtllm.fp4_block_scale_moe_runner( - router_logits, - routing_bias, + router_logits if not run_post_quant_allgather else None, + routing_bias if not run_post_quant_allgather else None, hidden_states_fp4, hidden_states_scale_linear_fp4.view(torch.float8_e4m3fn), self.w3_w1_weight, @@ -256,6 +321,8 @@ def forward_impl( routed_scaling_factor, self.routing_method.routing_method_type, do_finalize=do_finalize, + topk_ids=token_selected_experts, + topk_weights=token_final_scales, ) if not do_finalize: @@ -295,20 +362,25 @@ def forward_impl( self.routing_method.routing_method_type, 0, # act_type 1.3, # imbalance_factor + token_final_scales, + token_selected_experts, ) final_hidden_states = final_hidden_states[:, :self. hidden_size].contiguous() elif self.has_w4a8_mxfp4_fp8: pad_size = self.w3_w1_weight.shape[-1] * 2 - x.shape[-1] - x = torch.nn.functional.pad(x, (0, pad_size)) - x, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( - x, self.fc31_input_gate_dequant[0]) + if not run_post_quant_allgather: + x = torch.nn.functional.pad(x, (0, pad_size)) + x, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + x, self.fc31_input_gate_dequant[0]) + else: + x = x intermediate_size_per_partition_padded = self.w3_w1_weight.shape[ -2] // 2 final_hidden_states = torch.ops.trtllm.e4m3_mxe2m1_block_scale_moe_runner( - router_logits, - routing_bias, + router_logits if not run_post_quant_allgather else None, + routing_bias if not run_post_quant_allgather else None, x, self.w3_w1_weight, self.w3_w1_weight_scale, @@ -334,19 +406,25 @@ def forward_impl( self.routing_method.routing_method_type, 0, # act_type 1.3, # imbalance_factor + token_final_scales, + token_selected_experts, ) final_hidden_states = final_hidden_states[:, :self. hidden_size].contiguous() elif self.has_w4a8_mxfp4_mxfp8: - # TRTLLM-Gen uses linear SF layout for the mxfp8 input. - mxfp8_x, sf = torch.ops.trtllm.mxfp8_quantize( - x, False, alignment=self.quant_method.weight_alignment) + if not run_post_quant_allgather: + # TRTLLM-Gen uses linear SF layout for the mxfp8 input. + mxfp8_x, sf = torch.ops.trtllm.mxfp8_quantize( + x, False, alignment=self.quant_method.weight_alignment) + else: + mxfp8_x, sf = x, x_sf + intermediate_size_per_partition_padded = self.w3_w1_weight.shape[ -2] // 2 final_hidden_states = torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner( - router_logits, - routing_bias, + router_logits if not run_post_quant_allgather else None, + routing_bias if not run_post_quant_allgather else None, mxfp8_x, sf, self.w3_w1_weight, @@ -371,10 +449,12 @@ def forward_impl( self.routing_method.routing_method_type, 0, # act_type 1.3, # imbalance_factor + token_final_scales, + token_selected_experts, ) else: raise NotImplementedError( - "TRTLLMGenFusedMoE only supports fp8_block_scaling, nvfp4, w4a16_mxfp4 and w4a8_mxfp4_fp8 dtypes." + "TRTLLMGenFusedMoE only supports fp8_block_scaling, nvfp4, w4a16_mxfp4, w4a8_mxfp4_mxfp8 and w4a8_mxfp4_fp8 dtypes." ) final_hidden_states = self.reducescatter_or_allreduce( diff --git a/tensorrt_llm/_torch/modules/fused_moe/routing.py b/tensorrt_llm/_torch/modules/fused_moe/routing.py index 635091c7ad5e..886d3a97b7d2 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/routing.py +++ b/tensorrt_llm/_torch/modules/fused_moe/routing.py @@ -183,15 +183,19 @@ def routing_method_type(self) -> RoutingMethodType: class DefaultMoeRoutingMethod(BaseMoeRoutingMethod): - def __init__(self, top_k: int, force_enable_pytorch_op: bool = False): + def __init__(self, + top_k: int, + output_dtype: torch.dtype = torch.float32, + force_enable_pytorch_op: bool = False): super().__init__() self.top_k = top_k self.force_enable_pytorch_op = force_enable_pytorch_op + self.output_dtype = output_dtype def apply_pytorch( self, router_logits: torch.Tensor) -> (torch.Tensor, torch.Tensor): topk_values, topk_indices = torch.topk(torch.nn.functional.softmax( - router_logits.float(), dim=-1), + router_logits.to(self.output_dtype), dim=-1), k=self.top_k, dim=-1) return topk_indices.to(torch.int32), topk_values @@ -203,7 +207,7 @@ def apply(self, return self.apply_pytorch(router_logits) else: return torch.ops.trtllm.default_moe_routing_op( - router_logits, self.top_k) + router_logits, self.top_k, self.output_dtype) @property def routing_method_type(self): @@ -228,11 +232,13 @@ class RenormalizeMoeRoutingMethod(BaseMoeRoutingMethod): def __init__( self, top_k: int, + output_dtype: torch.dtype = torch.float32, force_enable_pytorch_op: bool = False, ): super().__init__() self.top_k = top_k self.force_enable_pytorch_op = force_enable_pytorch_op + self.output_dtype = output_dtype def apply_pytorch( self, router_logits: torch.Tensor) -> (torch.Tensor, torch.Tensor): @@ -240,7 +246,7 @@ def apply_pytorch( k=self.top_k, dim=-1) return topk_indices.to(torch.int32), torch.nn.functional.softmax( - topk_values.float(), dim=-1) + topk_values.to(self.output_dtype), dim=-1) def apply(self, router_logits: torch.Tensor) -> (torch.Tensor, torch.Tensor): @@ -249,7 +255,7 @@ def apply(self, return self.apply_pytorch(router_logits) else: return torch.ops.trtllm.renorm_moe_routing_op( - router_logits, self.top_k) + router_logits, self.top_k, self.output_dtype) @property def routing_method_type(self): @@ -258,16 +264,18 @@ def routing_method_type(self): class Llama4RenormalizeMoeRoutingMethod(BaseMoeRoutingMethod): - def __init__(self, top_k: int): + def __init__(self, top_k: int, output_dtype: torch.dtype = torch.float32): super().__init__() self.top_k = top_k + self.output_dtype = output_dtype def apply(self, router_logits: torch.Tensor) -> (torch.Tensor, torch.Tensor): topk_values, topk_indices = torch.topk(router_logits, k=self.top_k, dim=-1) - return topk_indices.to(torch.int32), torch.sigmoid(topk_values.float()) + return topk_indices.to(torch.int32), torch.sigmoid( + topk_values.float()).to(self.output_dtype) @property def routing_method_type(self): @@ -396,15 +404,16 @@ def apply(self, class RenormalizeNaiveMoeRoutingMethod(RenormalizeMoeRoutingMethod): - def __init__(self, top_k: int): - super().__init__() + def __init__(self, top_k: int, output_dtype: torch.dtype = torch.float32): + super().__init__(top_k, output_dtype) self.top_k = top_k + self.output_dtype = output_dtype def apply(self, router_logits: torch.Tensor) -> (torch.Tensor, torch.Tensor): #x = topk(softmax()); x /= x.sum() is mathematically equivalent to softmax(topk) topk_indices, topk_values = self.apply_pytorch(router_logits) - return topk_indices.to(torch.int32), topk_values + return topk_indices.to(torch.int32), topk_values.to(self.output_dtype) @property def routing_method_type(self) -> RoutingMethodType: diff --git a/tests/unittest/_torch/thop/parallel/test_moe.py b/tests/unittest/_torch/thop/parallel/test_moe.py index 4612d31477a7..5462de2cc7a9 100644 --- a/tests/unittest/_torch/thop/parallel/test_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_moe.py @@ -819,7 +819,7 @@ class TestMoeFP8: the default tactic selection works. This reduces unnecessary test runs for CI """ - @pytest.mark.parametrize("num_tokens", [16, 64, 4096]) + @pytest.mark.parametrize("num_tokens", [16, 64, 1024]) @pytest.mark.parametrize("expert_info", [(32, 8, 4, 8), (72, 1, 1, 6), (256, 8, 4, 8)]) @pytest.mark.parametrize("hidden_size", [512]) @@ -832,26 +832,31 @@ def test_autotune(self, num_tokens: int, expert_info: Tuple[int, int, int, expert_info, hidden_size, intermediate_size, - use_autotune=True) + use_autotune=True, + use_topk_as_input=False) @pytest.mark.parametrize("num_tokens", [16]) @pytest.mark.parametrize("expert_info", [(32, 8, 4, 8)]) @pytest.mark.parametrize("hidden_size", [512]) @pytest.mark.parametrize("intermediate_size", [512]) + @pytest.mark.parametrize("use_topk_as_input", [False, True], + ids=["use_score_as_input", "use_topk_as_input"]) def test_no_autotune(self, num_tokens: int, expert_info: Tuple[int, int, int, int], - hidden_size: int, intermediate_size: int): + hidden_size: int, intermediate_size: int, + use_topk_as_input: bool): self.run_moe_fp8_test(num_tokens, expert_info, hidden_size, intermediate_size, - use_autotune=False) + use_autotune=False, + use_topk_as_input=use_topk_as_input) def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, int, int], hidden_size: int, intermediate_size: int, - use_autotune: bool): + use_autotune: bool, use_topk_as_input: bool): torch.random.manual_seed(0) # @@ -906,12 +911,21 @@ def run_moe_fp8_test(self, num_tokens: int, expert_info: Tuple[int, int, None, scores, gemm1_weights, gemm1_scales, None, gemm2_weights, gemm2_scales, None, permute_info, False) + if use_topk_as_input: + topk_ids = permute_info["topKIndices"].to(torch.int32) + topk_weights = permute_info["topKLogits"].to(torch.bfloat16) + expert_logits = None + else: + topk_ids = None + topk_weights = None + with autotune(use_autotune): output = torch.ops.trtllm.fp8_block_scale_moe_runner( expert_logits, routing_bias, hidden_states, hidden_states_scale, gemm1_weights, gemm1_scales, gemm2_weights, gemm2_scales, num_experts, top_k, n_groups, top_k_groups, intermediate_size, - 0, num_experts, routed_scaling, routing_method_type) + 0, num_experts, routed_scaling, routing_method_type, + topk_weights, topk_ids) output_dequant_actual = output.to(torch.float) # @@ -938,7 +952,7 @@ class TestMoeFp4: the default tactic selection works. This reduces unnecessary test runs for CI """ - @pytest.mark.parametrize("num_tokens", [1, 4096]) + @pytest.mark.parametrize("num_tokens", [1, 1024]) @pytest.mark.parametrize("hidden_size", [1024]) @pytest.mark.parametrize("intermediate_size", [1024, 768, 384, 192]) @pytest.mark.parametrize( @@ -1013,9 +1027,10 @@ def test_autotune(self, num_tokens, hidden_size, intermediate_size, hidden_size, intermediate_size, routing_info, - use_autotune=True) + use_autotune=True, + use_topk_as_input=False) - @pytest.mark.parametrize("num_tokens", [1]) + @pytest.mark.parametrize("num_tokens", [1, 150]) @pytest.mark.parametrize("hidden_size", [1024]) @pytest.mark.parametrize("intermediate_size", [1024]) @pytest.mark.parametrize( @@ -1033,20 +1048,35 @@ def test_autotune(self, num_tokens, hidden_size, intermediate_size, "routing_method_type": RoutingMethodType.DeepSeekV3 }, id="RoutingDSv3"), + pytest.param( + { + "num_experts": 128, + "top_k": 4, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize + }, + id="RoutingRenormalize_topk_4"), ], ) + @pytest.mark.parametrize("use_topk_as_input", [False, True], + ids=["use_score_as_input", "use_topk_as_input"]) def test_no_autotune(self, num_tokens, hidden_size, intermediate_size, - routing_info): + routing_info, use_topk_as_input): self.run_moe_fp4_test(num_tokens, hidden_size, intermediate_size, routing_info, - use_autotune=False) + use_autotune=False, + use_topk_as_input=use_topk_as_input) def run_moe_fp4_test(self, num_tokens: int, hidden_size: int, intermediate_size: int, routing_info: dict, - use_autotune: bool) -> None: + use_autotune: bool, use_topk_as_input: bool) -> None: torch.random.manual_seed(0) @@ -1074,6 +1104,12 @@ def run_moe_fp4_test(self, num_tokens: int, hidden_size: int, assert top_k <= 8 assert num_experts % 4 == 0 + if use_topk_as_input: + if routing_method_type != RoutingMethodType.DeepSeekV3 or num_tokens != 150 or use_autotune: + pytest.skip( + "use_topk_as_input is tested only with routing_method_type=DeepSeekV3 and num_tokens=150 and use_autotune=False" + ) + if are_groups_valid(top_k_groups, n_groups): assert top_k_groups <= 4 assert num_experts > n_groups @@ -1248,6 +1284,14 @@ def run_moe_fp4_test(self, num_tokens: int, hidden_size: int, scale_c_fc2 = (1.0 / args_dequant.c_global_sf) * ( 1.0 / args.gemm2_scales_global) + if use_topk_as_input: + topk_ids = permute_info["topKIndices"].to(torch.int32) + topk_weights = permute_info["topKLogits"].to(torch.bfloat16) + expert_logits = None + else: + topk_ids = None + topk_weights = None + with autotune(use_autotune): output = torch.ops.trtllm.fp4_block_scale_moe_runner( expert_logits, @@ -1270,7 +1314,9 @@ def run_moe_fp4_test(self, num_tokens: int, hidden_size: int, num_experts, routed_scaling, routing_method_type, - do_finalize=True) + do_finalize=True, + topk_ids=topk_ids, + topk_weights=topk_weights) output_dequant_actual = output[0].to(torch.float) @@ -1286,14 +1332,19 @@ def run_moe_fp4_test(self, num_tokens: int, hidden_size: int, reason="The kernel only supports Blackwell. Current SM is %d." % getSMVersion(), ) -@pytest.mark.parametrize("num_tokens", [1, 4, 384, 4096]) +@pytest.mark.parametrize("num_tokens", [1, 1024]) @pytest.mark.parametrize("expert_info", [(128, 0, 0, 1, True)]) @pytest.mark.parametrize("hidden_size", [2048]) @pytest.mark.parametrize("intermediate_size", [2048]) +@pytest.mark.parametrize("use_topk_as_input", [False, True], + ids=["use_score_as_input", "use_topk_as_input"]) def test_moe_fp8_per_tensor_scale(num_tokens, expert_info, hidden_size, - intermediate_size): + intermediate_size, use_topk_as_input): torch.random.manual_seed(0) + if use_topk_as_input: + if num_tokens != 1: + pytest.skip("use_topk_as_input is tested only with num_tokens=1") # # Data Generation # @@ -1396,13 +1447,21 @@ def test_moe_fp8_per_tensor_scale(num_tokens, expert_info, hidden_size, scale_c_fc2 = (1.0 / args_dequant.c_global_sf) * (1.0 / args.gemm2_scales_global) + # set topk_weights and topk_ids + topk_weights = None + topk_ids = None + if use_topk_as_input: + topk_ids = permute_info["topKIndices"].to(torch.int32).cuda() + topk_weights = permute_info["topKLogits"].to(torch.bfloat16).cuda() + output = torch.ops.trtllm.fp8_per_tensor_scale_moe_runner( - expert_logits.to(torch.bfloat16) if use_routing_scales_on_input else - expert_logits, routing_bias, hidden_states_quant, - gemm1_weights_fp8_shuffled, scale_c_fc1, scale_gate_fc1, - gemm2_weights_fp8_shuffled, scale_c_fc2, num_experts, top_k, n_groups, - top_k_groups, intermediate_size, 0, num_experts, routed_scaling, - use_routing_scales_on_input, tile_tokens_dim, routing_method_type) + expert_logits.to(torch.bfloat16) + if use_routing_scales_on_input else expert_logits, routing_bias, + hidden_states_quant, gemm1_weights_fp8_shuffled, scale_c_fc1, + scale_gate_fc1, gemm2_weights_fp8_shuffled, scale_c_fc2, num_experts, + top_k, n_groups, top_k_groups, intermediate_size, 0, num_experts, + routed_scaling, use_routing_scales_on_input, tile_tokens_dim, + routing_method_type, topk_weights, topk_ids) output_dequant_actual = output.to(torch.float) @@ -1418,7 +1477,7 @@ def test_moe_fp8_per_tensor_scale(num_tokens, expert_info, hidden_size, reason="The kernel only supports Blackwell. Current SM is %d." % getSMVersion(), ) -@pytest.mark.parametrize("num_tokens", [1, 256, 1024]) +@pytest.mark.parametrize("num_tokens", [1, 1024]) @pytest.mark.parametrize("hidden_size", [512]) @pytest.mark.parametrize("intermediate_size", [512]) @pytest.mark.parametrize( @@ -1452,9 +1511,11 @@ def test_moe_fp8_per_tensor_scale(num_tokens, expert_info, hidden_size, @pytest.mark.parametrize("act_type_str", ["SwiGlu", "GatedSilu"]) @pytest.mark.parametrize("use_autotune", [True, False], ids=["autotune", "no_autotune"]) +@pytest.mark.parametrize("use_topk_as_input", [False, True], + ids=["use_score_as_input", "use_topk_as_input"]) def test_moe_mxe2m1_weights(num_tokens, hidden_size, intermediate_size, routing_info, dtype_activation, act_type_str, - use_autotune): + use_autotune, use_topk_as_input): torch.random.manual_seed(0) # @@ -1493,6 +1554,11 @@ def test_moe_mxe2m1_weights(num_tokens, hidden_size, intermediate_size, pytest.skip( "1024 tokens is tested only with topk=4, mxfp8, SwiGlu, and autotune" ) + if use_topk_as_input: + if dtype_activation != "mxfp8" or top_k != 4 or act_type_str != "SwiGlu" or not use_autotune or num_tokens != 1: + pytest.skip( + "use_topk_as_input is tested only with mxfp8, topk=4, SwiGlu, autotune, and num_tokens=1" + ) assert top_k <= num_experts assert top_k <= 8 @@ -1591,6 +1657,14 @@ def test_moe_mxe2m1_weights(num_tokens, hidden_size, intermediate_size, else: raise ValueError("Invalid routing method type") + if use_topk_as_input: + topk_ids = permute_info["topKIndices"].to(torch.int32) + topk_weights = permute_info["topKLogits"] + expert_logits = None + else: + topk_ids = None + topk_weights = None + input_hidden_states = None input_hidden_global_scale = None if dtype_activation == "mxfp8": @@ -1747,7 +1821,8 @@ def test_moe_mxe2m1_weights(num_tokens, hidden_size, intermediate_size, gemm2_scales_mxe2m1_shuffled.cuda(), gemm2_bias_shuffled.cuda(), num_experts, top_k, n_groups, top_k_groups, intermediate_size, unpadded_hidden_size, 0, num_experts, routed_scaling, - routing_method_type, act_type.value, 1.0) + routing_method_type, act_type.value, 1.0, topk_weights, + topk_ids) elif dtype_activation == "bf16": output = torch.ops.trtllm.bf16_mxe2m1_block_scale_moe_runner( expert_logits, routing_bias, @@ -1759,7 +1834,7 @@ def test_moe_mxe2m1_weights(num_tokens, hidden_size, intermediate_size, gemm2_scales_mxe2m1_shuffled.cuda(), gemm2_bias_shuffled.cuda(), num_experts, top_k, n_groups, top_k_groups, intermediate_size, 0, num_experts, routed_scaling, routing_method_type, - act_type.value, 1.0) + act_type.value, 1.0, topk_weights, topk_ids) elif dtype_activation == "fp8": output = torch.ops.trtllm.e4m3_mxe2m1_block_scale_moe_runner( expert_logits, routing_bias, @@ -1771,7 +1846,8 @@ def test_moe_mxe2m1_weights(num_tokens, hidden_size, intermediate_size, gemm2_scales_mxe2m1_shuffled.cuda(), gemm2_bias_shuffled.cuda(), scale_c_fc1, scale_gate_fc1, scale_c_fc2, num_experts, top_k, n_groups, top_k_groups, intermediate_size, 0, num_experts, - routed_scaling, routing_method_type, act_type.value, 1.0) + routed_scaling, routing_method_type, act_type.value, 1.0, + topk_weights, topk_ids) else: raise ValueError("Invalid dtype_activation")