diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu index 43a1a6846cdc..a8130ad8f80b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu @@ -265,12 +265,14 @@ constexpr int DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA = 128; // and strides over the row space. This visits the per-expert tile padding that // the expanded-space kernel skips (~4% extra rows at 32 local experts); those // rows are dropped by the finalize kernel. The arithmetic below deliberately -// preserves the legacy kernel's 0/0 -> NaN behavior for an all-zero block. +// matches activationDeepSeekKernel bit for bit, including its finite all-zero +// block handling. constexpr int kDsActWarpSize = 32; constexpr int kDsActEltsPerSf = 128; constexpr int kDsActEltsPerThread = kDsActEltsPerSf / kDsActWarpSize; constexpr int kDsActWarpsPerCta = 4; constexpr int kDsActPermutedNumThreadsPerCta = kDsActWarpSize * kDsActWarpsPerCta; +constexpr float kDsActAmaxEpsilon = 1.0e-10F; constexpr bool shouldUsePermutedActivation(int innerDim, int numTokens, int topK, int numExperts, int tileTokensDim) { @@ -352,7 +354,11 @@ __global__ void activationDeepSeekPermutedKernel(KernelParams params) aMax = fmaxf(aMax, __shfl_xor_sync(0xffffffffu, aMax, offset)); } - float const scaleOut = aMax / kE4m3MaxVal; + // Floor aMax so an all-zero block stays finite: without it scaleOut is + // zero and quantizing evaluates 0 / 0, which is undefined and writes FP8 + // NaNs into that row. Same epsilon as the DeepGEMM FP8 activation + // quantizer (fp8_utils.py). + float const scaleOut = fmaxf(aMax, kDsActAmaxEpsilon) / kE4m3MaxVal; if (lane == 0) { @@ -367,7 +373,7 @@ __global__ void activationDeepSeekPermutedKernel(KernelParams params) // Divide; do NOT hoist a reciprocal. `x / s` and `x * (1/s)` round // differently, and an equivalence run showed that single ulp flip a // greedy-decoded token. This must match activationDeepSeekKernel - // bit for bit, including 0/0 -> NaN on an all-zero scale block. + // bit for bit. outElts[i] = static_cast(out[i] / scaleOut); } *reinterpret_cast(params.outPtr + static_cast(permutedIdx) * outputDim + hiddenBase) @@ -504,10 +510,11 @@ __global__ void activationDeepSeekKernel(KernelParams params) { continue; } - s_scaleOutArr[tokenInCtaIdx] = aMaxArr[tokenInCtaIdx] / E4m3MaxVal; + float const scaleOut = fmaxf(aMaxArr[tokenInCtaIdx], kDsActAmaxEpsilon) / E4m3MaxVal; + s_scaleOutArr[tokenInCtaIdx] = scaleOut; int const scaleOut_idx = permutedIdxArr[tokenInCtaIdx] + totalNumPaddedTokens * (hiddenIdx / 128); - params.outDqSfsPtr[scaleOut_idx] = aMaxArr[tokenInCtaIdx] / E4m3MaxVal; + params.outDqSfsPtr[scaleOut_idx] = scaleOut; } } __syncthreads(); @@ -1101,9 +1108,15 @@ __global__ void finalizeDeepSeekKernel(KernelParams params) { if (params.outDqSfsPtr) { - s_scaleOut = aMax / E4m3MaxVal; + // Same all-zero-block hazard as the activation kernels: without the floor + // an all-zero accumulator makes the division below evaluate 0 / 0. This + // branch is unreachable today because every thop entry point passes + // args.output_scale = nullptr, so nothing observable changes; the floor is + // here so the first caller to wire up outDqSfsPtr does not inherit it. + float const scaleOut = fmaxf(aMax, activation::kDsActAmaxEpsilon) / E4m3MaxVal; + s_scaleOut = scaleOut; int const scaleOut_idx = tokenIdx + hiddenIdx / 128 * params.numTokens; - params.outDqSfsPtr[scaleOut_idx] = aMax / E4m3MaxVal; + params.outDqSfsPtr[scaleOut_idx] = scaleOut; } else { diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu index 6a0ba120d4cc..e8f682ca47ab 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu @@ -27,6 +27,7 @@ // 7. routingIndicesOffsetsKernel — prefix-scan + permutation (defined in RoutingKernel.cuh) #include "RoutingCustomPolicy.cuh" +#include "RoutingCustomSelection.h" #include @@ -1594,6 +1595,46 @@ void launchOffsetsKernel(Data const& data, int numBlocksOffsets, uint32_t numThr // //////////////////////////////////////////////////////////////////////////////////////////////////// +bool prefersCoopBlockKernel(RoutingPreprocessType preprocessType, RoutingPostprocessType postprocessType, + int32_t numTokens, int32_t dispatchedMaxExperts, int32_t minNumExpertsForCoopOverride) +{ + // The cooperative block kernel is the fastest path for tiny batches. It needs an + // elementwise preprocess (anything but softmax-over-experts) and one CUDA block's + // worth of experts, since it runs one thread per expert. + bool const useStaticBlock = numTokens <= BlockKernelMaxNumTokens; + bool const preprocessIsElementwise = preprocessType == RoutingPreprocessType::None + || preprocessType == RoutingPreprocessType::Sigmoid || preprocessType == RoutingPreprocessType::SigmoidBias; + + // The lower tier bound applies to the Renormalize policy only, which is the one that + // was measured. With no per-expert preprocess the classic one-warp-per-token TopK is + // faster through the 512-expert tier. Policies that do preprocess per expert push the + // classic kernel into register spilling long before that -- at E512/topK 22 SigmoidBias + // it needs 64 registers and a 176-byte stack against 32 registers and no stack for the + // cooperative kernel -- so they keep using the cooperative kernel across the whole tier + // range. The None + None fallback policy is left alone for the same reason: it is + // unmeasured, and no routing method in runner.cu selects it today. + // + // The bound is one tier lower at a single token. Measured across GB300 (SM103) and + // B200 (SM100) with the same launcher harness, the classic kernel wins every tier up to + // 512 from two tokens up, but at one token the two parts disagree at the 512 tier and + // both prefer the cooperative kernel at 576. + // + // The bound is the only part of this predicate that rests on measurement, and the + // measurement is SM100-family only, so it is the part a deployment may need to undo + // without a rebuild. minNumExpertsForCoopOverride carries + // TLLM_ROUTING_COOP_BLOCK_MIN_EXPERTS in from the caller: 0 restores the parent + // selection, a value above every tier forces the classic kernel. + bool const isRenormalize + = preprocessType == RoutingPreprocessType::None && postprocessType == RoutingPostprocessType::Softmax; + int32_t const minNumExpertsForCoop = minNumExpertsForCoopOverride >= 0 + ? minNumExpertsForCoopOverride + : (numTokens == 1 ? CoopBlockKernelSingleTokenMinNumExperts : CoopBlockKernelMinNumExperts); + bool const meetsMinNumExperts = !isRenormalize || dispatchedMaxExperts >= minNumExpertsForCoop; + + return useStaticBlock && preprocessIsElementwise && meetsMinNumExperts + && dispatchedMaxExperts <= CoopBlockKernelMaxNumExperts; +} + void run(Data const& data, void* stream) { TLLM_CHECK_WITH_INFO(data.mPtrTopKPacked != nullptr || data.mPtrScores != nullptr || data.mPtrTopKIds != nullptr, @@ -1629,21 +1670,23 @@ void run(Data const& data, void* stream) bool const useStaticBlock = data.mNumTokens <= BlockKernelMaxNumTokens; int32_t const dispatchedMaxExperts = queryDispatchedMaxExperts(data); - // Cooperative block kernel: fastest path for tiny batches. Requires an elementwise - // preprocess (any but softmax-over-experts) and one CUDA block's worth of experts. - // Critical for large expert counts, where the classic one-warp-per-token TopK spills - // registers under the 1024-thread launch bounds (e.g. 896 experts / topK 16 at decode). - bool const preprocessIsElementwise = data.mPreprocessType == RoutingPreprocessType::None - || data.mPreprocessType == RoutingPreprocessType::Sigmoid - || data.mPreprocessType == RoutingPreprocessType::SigmoidBias; // Escape hatch for A/B validation and emergency fallback to the classic block kernel. static bool const disableCoopBlock = [] { char const* env = std::getenv("TLLM_ROUTING_DISABLE_COOP_BLOCK"); return env != nullptr && env[0] == '1'; }(); - bool const useCoopBlock = !disableCoopBlock && useStaticBlock && preprocessIsElementwise - && dispatchedMaxExperts <= CoopBlockKernelMaxNumExperts; + // The opposite direction: move the Renormalize lower tier bound instead of disabling + // the cooperative kernel outright. 0 restores the parent selection for every tier. + // Both are read once into a function-static, so they must be set before the first call. + static int32_t const coopBlockMinNumExpertsOverride = [] + { + char const* env = std::getenv("TLLM_ROUTING_COOP_BLOCK_MIN_EXPERTS"); + return env != nullptr ? std::atoi(env) : -1; + }(); + bool const useCoopBlock = !disableCoopBlock + && prefersCoopBlockKernel(data.mPreprocessType, data.mPostprocessType, data.mNumTokens, dispatchedMaxExperts, + coopBlockMinNumExpertsOverride); bool const useDynBlock = !useStaticBlock && data.mNumTokens <= DynBlockKernelMaxNumTokens && dispatchedMaxExperts <= DynBlockKernelMaxNumExperts; bool const useSingleBlock = useStaticBlock || useDynBlock; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh index ad0163cf9294..86caaf63afd0 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomPolicy.cuh @@ -391,6 +391,17 @@ static constexpr int MaxNumTokensSingleClusterScores = NumBlocksPerCluster * Num static constexpr int BlockKernelMaxNumTokens = 4; static constexpr int DynBlockKernelMaxNumTokens = 16; static constexpr int DynBlockKernelMaxNumExperts = 256; +// For the Renormalize policy (None + Softmax) the classic block kernel is faster through +// the 512-expert tier, so the cooperative kernel is only preferred from this tier up. +// Every other policy is excluded from this bound; see prefersCoopBlockKernel(). +// TLLM_ROUTING_COOP_BLOCK_MIN_EXPERTS overrides both bounds below at runtime. +static constexpr int CoopBlockKernelMinNumExperts = 576; +// At a single token the classic kernel gives up its advantage one tier earlier. The +// cooperative kernel runs one thread per expert, so it scales with the expert count even +// when there is only one token to route, while the classic kernel has one warp of work in +// total. Both measured parts agree the cooperative kernel wins the 576 tier at one token, +// and they disagree at 512, so 512 stays cooperative there. See prefersCoopBlockKernel(). +static constexpr int CoopBlockKernelSingleTokenMinNumExperts = 512; // Cooperative block kernel: one thread per expert, so at most 1024 experts (1 CUDA block). static constexpr int CoopBlockKernelMaxNumExperts = 1024; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomSelection.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomSelection.h new file mode 100644 index 000000000000..79db2a48c88a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomSelection.h @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "RoutingKernel.h" + +#include + +namespace moe::dev::routing::routingCustom +{ + +//! Whether the cooperative block kernel is the preferred launcher for this shape. +//! +//! Split out of run() so the selection table can be unit tested without launching a +//! kernel. This is host-only and holds no state: both escape hatches, +//! TLLM_ROUTING_DISABLE_COOP_BLOCK and TLLM_ROUTING_COOP_BLOCK_MIN_EXPERTS, are read by +//! the caller and applied here only through arguments. +//! +//! \param preprocessType routing preprocess applied before top-k. +//! \param postprocessType routing postprocess applied to the top-k scores. Paired with +//! preprocessType it identifies the policy, exactly as dispatchRoutingPolicy() does. +//! \param numTokens number of routing tokens in this launch. +//! \param dispatchedMaxExperts compile-time tier from queryDispatchedMaxExperts(), which +//! is not the model's raw expert count. +//! \param minNumExpertsForCoopOverride replaces the built-in Renormalize lower tier bound +//! when non-negative. 0 restores the parent behaviour of always preferring the +//! cooperative kernel; a value above every tier forces the classic kernel. It has +//! no effect on any other policy, which is not subject to the bound. +bool prefersCoopBlockKernel(RoutingPreprocessType preprocessType, RoutingPostprocessType postprocessType, + int32_t numTokens, int32_t dispatchedMaxExperts, int32_t minNumExpertsForCoopOverride = -1); + +} // namespace moe::dev::routing::routingCustom diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index c3007343bf68..e5debeda9e6f 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -22,6 +22,7 @@ #include #include +#include #include #include #include @@ -212,6 +213,13 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit int32_t max_num_padded_tokens_gemm1 = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( max_num_padded_tokens, 2 * args.intermediate_size, btg::dtypeGetNumBits(args.mDtypeElt)); + // maybeGetMinTokenCount pads a buffer up to the 128 KiB floor using the row width it is + // handed, so its result is only valid for that width. A gated activation is half as wide as + // gemm1_output, so reusing max_num_padded_tokens_gemm1 delivers half the intended floor on + // small decode batches. Derive the capacity from the activation's own row width instead. + int32_t max_num_padded_tokens_activation + = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( + max_num_padded_tokens, args.intermediate_size, btg::dtypeGetNumBits(args.mDtypeElt)); int32_t max_num_padded_tokens_gemm2 = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( max_num_padded_tokens, args.hidden_size, btg::dtypeGetNumBits(args.mDtypeOut)); @@ -254,10 +262,11 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit 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_gemm1}, at::ScalarType::Float, routing_device, std::nullopt); - at::Tensor activation_output = at::detail::empty_cuda( - {max_num_padded_tokens_gemm1, 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_gemm1}, at::ScalarType::Float, routing_device, std::nullopt); + at::Tensor activation_output = at::detail::empty_cuda({max_num_padded_tokens_activation, 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_activation}, at::ScalarType::Float, + routing_device, std::nullopt); at::Tensor gemm2_output = at::detail::empty_cuda( {max_num_padded_tokens_gemm2, args.hidden_size}, at::ScalarType::BFloat16, routing_device, std::nullopt); @@ -340,7 +349,11 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit // setup workspace workspace.total_num_padded_tokens = total_num_padded_tokens.data_ptr(); - workspace.total_max_padded_tokens = std::max(max_num_padded_tokens_gemm1, max_num_padded_tokens_gemm2); + // The activation is the narrowest of the three buffers, so its 128 KiB floor needs the most + // rows and it is now the largest capacity of the three on small batches. Include it here or + // this descriptor under-counts the workspace it claims to describe. + workspace.total_max_padded_tokens + = std::max({max_num_padded_tokens_gemm1, max_num_padded_tokens_activation, max_num_padded_tokens_gemm2}); workspace.routing_expert_indexes = expert_indexes.data_ptr(); workspace.permuted_idx_size = total_num_padded_tokens.data_ptr(); workspace.expanded_idx_to_permuted_idx diff --git a/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu index 4a56dcd8f558..9616a13d3995 100644 --- a/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu +++ b/cpp/tests/unit_tests/kernels/blockScaleMoeActivationTest.cu @@ -22,12 +22,11 @@ // * `activationDeepSeekPermutedKernel` - grids directly over the permuted row // space with one warp per (row, 128-element scale block), // via `shouldUsePermutedActivation()`. Both must produce *identical bits* for -// every row that carries a real token: DevKernel.cu documents that the permuted +// every row that carries a real token. DevKernel.cu documents that the permuted // kernel must not hoist a reciprocal out of `out / scaleOut`, because `x / s` // and `x * (1/s)` round differently and one ulp was enough to flip a // greedy-decoded token. An `isClose`-style comparison would not catch that -// regression, so everything below compares raw bit patterns (which also makes -// the NaN cases comparable). +// regression, so everything below compares raw bit patterns. // // Note on coverage: fp8 e4m3 carries three mantissa bits, so most 1-ulp fp32 // differences vanish when the result is rounded back down to fp8 -- only values @@ -45,6 +44,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h" +#include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/iBuffer.h" @@ -52,6 +52,7 @@ #include #include +#include #include #include #include @@ -212,8 +213,8 @@ protected: // Allocates device buffers, fills the inputs deterministically and uploads // them. `zeroedRowBlock`, when set, forces one (row, scale block) pair of - // the input to all-zero so both kernels take the aMax == 0 -> 0/0 -> NaN - // path on exactly the same element. + // the input to all-zero so both kernels exercise the finite aMax floor on + // exactly the same element. void setUp(ActivationEquivParam const& param, PermutedLayout const& layout, std::optional> zeroedRowBlock = std::nullopt) { @@ -401,12 +402,11 @@ INSTANTIATE_TEST_SUITE_P(BlockScaleMoeActivation, BlockScaleMoeActivationEquival //////////////////////////////////////////////////////////////////////////////////////////////////// -// An all-zero scale block yields aMax == 0, so the quantization does 0 / 0. The -// resulting NaN encoding is unspecified, but both kernels evaluate the same -// expression and must therefore land on the same bits -- which is exactly what -// would break if one of them replaced the division with a multiply by the -// reciprocal. -TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesIdenticalNaNs) +// An all-zero scale block must remain finite. A zero dequantization scale would +// make quantization evaluate 0 / 0, which is undefined and writes FP8 NaNs into +// that row. Both kernels floor aMax with the same epsilon and must emit +// identical zero bytes and a finite, positive scale. +TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesFiniteZeros) { ActivationEquivParam const param{"zero_block", /*numTokens=*/64, /*topK=*/4, /*numExperts=*/32, /*numLocalExperts=*/8, /*intermediateSize=*/256, /*paddingTile=*/8, /*hasSwigluLimit=*/false, @@ -422,17 +422,53 @@ TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesIdenticalNa auto const permuted = runOnce(kTileForcePermuted); auto const sfIdx = static_cast(zeroedRow) + static_cast(mTotalRows) * zeroedBlock; - EXPECT_EQ(floatBits(legacy.scales[sfIdx]), 0U) << "an all-zero block must give scaleOut == +0"; + EXPECT_TRUE(std::isfinite(legacy.scales[sfIdx])); + EXPECT_GT(legacy.scales[sfIdx], 0.F); EXPECT_EQ(floatBits(legacy.scales[sfIdx]), floatBits(permuted.scales[sfIdx])); for (int32_t elt = 0; elt < kEltsPerSf; ++elt) { auto const idx = static_cast(zeroedRow) * mOutputDim + zeroedBlock * kEltsPerSf + elt; + EXPECT_EQ(legacy.bytes[idx], toFp8Byte(0.F)) << "zero block emitted non-zero FP8 at element " << elt; ASSERT_EQ(static_cast(legacy.bytes[idx]), static_cast(permuted.bytes[idx])) - << "0/0 encoding differs at element " << elt; + << "zero-block encoding differs at element " << elt; } } //////////////////////////////////////////////////////////////////////////////////////////////////// +TEST(BlockScaleMoeActivationBackingTest, PadsActivationUsingItsOwnRowWidth) +{ + // A single-token Qwen-style decode can have only 32 padded rows. FC1 writes + // 2 * intermediateSize elements per row, while the gated activation read by + // FC2 is half as wide, so reusing FC1's capacity delivers only half of + // maybeGetMinTokenCount's 128 KiB floor. This is host-side arithmetic: it + // pins the sizing invariant, it does not observe the actual allocation. + constexpr int32_t maxNumPaddedTokens = 32; + constexpr int32_t intermediateSize = 2304; + constexpr int64_t minActivationBytes = 128 * 1024; + // Not an independent requirement: the FP32 scales are indexed by the same + // token capacity, so they scale with the activation and land just past 4 KiB. + constexpr int64_t minScaleBytes = 4 * 1024; + auto const fp8Bits = tg::dtypeGetNumBits(tg::Dtype::E4m3); + + auto const gemm1Capacity = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( + maxNumPaddedTokens, 2 * intermediateSize, fp8Bits); + auto const activationCapacity = tensorrt_llm::kernels::trtllmGenFp8BlockScaleMoe::Routing::maybeGetMinTokenCount( + maxNumPaddedTokens, intermediateSize, fp8Bits); + + auto const activationBytes = static_cast(activationCapacity) * intermediateSize * fp8Bits / 8; + auto const activationBytesWithGemm1Capacity = static_cast(gemm1Capacity) * intermediateSize * fp8Bits / 8; + auto const scaleBytes = static_cast(activationCapacity) * (intermediateSize / kEltsPerSf) * sizeof(float); + auto const scaleBytesWithGemm1Capacity + = static_cast(gemm1Capacity) * (intermediateSize / kEltsPerSf) * sizeof(float); + + EXPECT_GE(activationBytes, minActivationBytes); + EXPECT_GE(scaleBytes, minScaleBytes); + EXPECT_LT(activationBytesWithGemm1Capacity, minActivationBytes); + EXPECT_LT(scaleBytesWithGemm1Capacity, minScaleBytes); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + } // namespace tensorrt_llm::tests::kernels::blockscalemoe diff --git a/cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp b/cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp index 4d1120420486..ac8042600bc3 100644 --- a/cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp +++ b/cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustomSelection.h" #include "tests/unit_tests/kernels/routing/routingTest.h" #include @@ -295,6 +296,186 @@ TYPED_TEST(RoutingCustomKernelTest, BlockLevelParallelizationWithExpertParalleli this->runTest(param); }; +TYPED_TEST(RoutingCustomKernelTest, BlockLevelClassicBoundaryE512K16) +{ + auto param = RoutingKernelTestParam() + .withRoutingMethod(RoutingMethodType::Renormalize) + .withNumTokens(4) + .withNumExperts(512) + .withTopK(16) + .withTileTokensDim(256) + .build(); + this->runTest(param); +}; + +TYPED_TEST(RoutingCustomKernelTest, BlockLevelCooperativeBoundaryE576K8) +{ + auto param = RoutingKernelTestParam() + .withRoutingMethod(RoutingMethodType::Renormalize) + .withNumTokens(4) + .withNumExperts(576) + .withTopK(8) + .withTileTokensDim(256) + .build(); + this->runTest(param); +}; + +// Tier 512 is the only tier whose launcher changes with the token count: one token keeps +// the cooperative kernel, two and up take the classic one. These two straddle that flip. +// Tier 576 has no such flip -- it is cooperative across the whole one-to-four range -- so +// it is covered once, above, at four tokens. +TYPED_TEST(RoutingCustomKernelTest, BlockLevelCooperativeBoundaryE512K16SingleToken) +{ + auto param = RoutingKernelTestParam() + .withRoutingMethod(RoutingMethodType::Renormalize) + .withNumTokens(1) + .withNumExperts(512) + .withTopK(16) + .withTileTokensDim(256) + .build(); + this->runTest(param); +}; + +TYPED_TEST(RoutingCustomKernelTest, BlockLevelClassicBoundaryE512K16TwoTokens) +{ + auto param = RoutingKernelTestParam() + .withRoutingMethod(RoutingMethodType::Renormalize) + .withNumTokens(2) + .withNumExperts(512) + .withTopK(16) + .withTileTokensDim(256) + .build(); + this->runTest(param); +}; + +// SigmoidBias keeps the cooperative kernel at every tier, including the ones Renormalize +// gives up. Four tokens is already covered by DeepSeekNoGroupBlockLevel and +// MiniMax2BlockLevel; one token is not covered anywhere else in this file, and it is where +// the cooperative kernel has the least work to amortise its one-thread-per-expert shape. +TYPED_TEST(RoutingCustomKernelTest, BlockLevelSigmoidBiasBoundaryE512K16SingleToken) +{ + auto param = RoutingKernelTestParam() + .withRoutingMethod(RoutingMethodType::MiniMax2) + .withNumTokens(1) + .withNumExperts(512) + .withTopK(16) + .withTileTokensDim(256) + .withRoutedScalingFactor(2.5f) + .build(); + this->runTest(param); +}; + +// The boundary tests above validate outputs, which both launchers produce identically, +// so they cannot detect a change in which launcher is selected. This one pins the selection +// table itself. Tiers are the compile-time values queryDispatchedMaxExperts() returns, not +// raw expert counts. +TEST(RoutingCustomSelectionTest, CoopBlockKernelSelectionTable) +{ + using moe::dev::routing::RoutingPostprocessType; + using moe::dev::routing::RoutingPreprocessType; + using moe::dev::routing::routingCustom::prefersCoopBlockKernel; + + auto const renormalize = [](int32_t numTokens, int32_t tier) + { return prefersCoopBlockKernel(RoutingPreprocessType::None, RoutingPostprocessType::Softmax, numTokens, tier); }; + + // The tiers below are exactly the ones each policy's PolicyTraits::Pairs can produce, + // so every cell here is a combination that queryDispatchedMaxExperts() can actually + // return. Renormalize: 128, 160, 256, 512, 576, 2048. + // + // Classic through the 512 tier, cooperative from 576 up. + for (int32_t numTokens : {2, 3, 4}) + { + for (int32_t tier : {128, 160, 256, 512}) + { + EXPECT_FALSE(renormalize(numTokens, tier)) << "tier " << tier << ", " << numTokens << " tokens"; + } + EXPECT_TRUE(renormalize(numTokens, 576)) << numTokens << " tokens"; + } + + // At a single token the boundary moves down one tier: 512 stays cooperative. + EXPECT_FALSE(renormalize(1, 128)); + EXPECT_FALSE(renormalize(1, 160)); + EXPECT_FALSE(renormalize(1, 256)); + EXPECT_TRUE(renormalize(1, 512)); + EXPECT_TRUE(renormalize(1, 576)); + + // Per-expert preprocessing spills registers in the classic kernel well below 512 + // experts, so the lower bound must not apply to it. SigmoidBias tiers: 128, 256, 384, + // 512, 1024. Sigmoid has a single tier. + for (int32_t tier : {128, 256, 384, 512, 1024}) + { + EXPECT_TRUE(prefersCoopBlockKernel( + RoutingPreprocessType::SigmoidBias, RoutingPostprocessType::ScaledSumNormalize, 4, tier)) + << "tier " << tier; + } + EXPECT_TRUE(prefersCoopBlockKernel(RoutingPreprocessType::Sigmoid, RoutingPostprocessType::SumNormalize, 4, 128)); + + // The None + None fallback shares a preprocess with Renormalize but was never measured, + // so it keeps the cooperative kernel that the parent selected. + EXPECT_TRUE(prefersCoopBlockKernel(RoutingPreprocessType::None, RoutingPostprocessType::None, 4, 128)); + + // Above one CUDA block of experts the cooperative kernel cannot run at all. + EXPECT_FALSE(renormalize(1, 2048)); + EXPECT_FALSE(prefersCoopBlockKernel( + RoutingPreprocessType::SigmoidBias, RoutingPostprocessType::ScaledSumNormalize, 1, 2048)); + + // Softmax-over-experts is not elementwise, so it never uses the cooperative kernel. + EXPECT_FALSE(prefersCoopBlockKernel(RoutingPreprocessType::Softmax, RoutingPostprocessType::None, 1, 576)); + + // The whole path is gated on the static-block token limit. + EXPECT_TRUE( + prefersCoopBlockKernel(RoutingPreprocessType::SigmoidBias, RoutingPostprocessType::ScaledSumNormalize, 4, 512)); + EXPECT_FALSE( + prefersCoopBlockKernel(RoutingPreprocessType::SigmoidBias, RoutingPostprocessType::ScaledSumNormalize, 5, 512)); +}; + +// TLLM_ROUTING_COOP_BLOCK_MIN_EXPERTS reaches the predicate as an argument, so the two +// directions of the escape hatch are asserted here rather than through the environment. +TEST(RoutingCustomSelectionTest, CoopBlockKernelMinNumExpertsOverride) +{ + using moe::dev::routing::RoutingPostprocessType; + using moe::dev::routing::RoutingPreprocessType; + using moe::dev::routing::routingCustom::prefersCoopBlockKernel; + + auto const renormalize = [](int32_t numTokens, int32_t tier, int32_t override) + { + return prefersCoopBlockKernel( + RoutingPreprocessType::None, RoutingPostprocessType::Softmax, numTokens, tier, override); + }; + + // 0 restores the parent selection: every Renormalize tier the cooperative kernel can + // run at all takes it again, at every token count in the static-block range. + for (int32_t numTokens : {1, 2, 3, 4}) + { + for (int32_t tier : {128, 160, 256, 512, 576}) + { + EXPECT_TRUE(renormalize(numTokens, tier, 0)) << "tier " << tier << ", " << numTokens << " tokens"; + } + } + + // A bound above every tier forces the classic kernel across the range. + for (int32_t numTokens : {1, 2, 3, 4}) + { + for (int32_t tier : {128, 160, 256, 512, 576}) + { + EXPECT_FALSE(renormalize(numTokens, tier, 4096)) << "tier " << tier << ", " << numTokens << " tokens"; + } + } + + // A negative value means unset and leaves the built-in bounds in place. + EXPECT_FALSE(renormalize(4, 512, -1)); + EXPECT_TRUE(renormalize(1, 512, -1)); + + // The hard constraints are not negotiable: the override cannot put more than one CUDA + // block of experts, or more than four tokens, on the cooperative kernel. + EXPECT_FALSE(renormalize(1, 2048, 0)); + EXPECT_FALSE(renormalize(5, 512, 0)); + + // Policies outside the bound are untouched in both directions. + EXPECT_TRUE(prefersCoopBlockKernel( + RoutingPreprocessType::SigmoidBias, RoutingPostprocessType::ScaledSumNormalize, 4, 512, 4096)); +}; + TYPED_TEST(RoutingCustomKernelTest, BlockLevelParallelizationWithInvalidTopKInput) { auto param = RoutingKernelTestParam() diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index aa2626bf1da7..919e1e6e18ce 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -881,44 +881,60 @@ def forward( spec_metadata: Optional[SpecMetadata] = None, **kwargs, ) -> torch.Tensor: - del all_rank_num_tokens - - def norm_embeds(): - return self.pre_fc_norm_embedding(embed_tokens(input_ids)) - - def norm_hidden(): - return self.pre_fc_norm_hidden(hidden_states) - - inputs_embeds, hidden_states = maybe_execute_in_parallel( - norm_embeds, - norm_hidden, - self.event_dict[EventType.Main], - self.event_dict[EventType.MoeShared], - self.aux_stream, - disable_on_compile=True, - ) - hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1) - - tp_size = self.model_config.mapping.tp_size - tp_rank = self.model_config.mapping.tp_rank - if tp_size > 1 and not self.model_config.mapping.enable_attention_dp: - hidden_states = torch.chunk(hidden_states, tp_size, dim=-1)[tp_rank] + # Install this draft step's attention-DP token distribution: the draft + # loop passes it as a kwarg and leaves attn_metadata alone, matching + # Eagle3DraftModel.forward. The MoE below reads the per-rank counts off + # attn_metadata, and an MTP Eagle draft runs one token per sequence + # after step 0, so a stale target value mismatches collective sizes. + previous_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + + try: + + def norm_embeds(): + return self.pre_fc_norm_embedding(embed_tokens(input_ids)) + + def norm_hidden(): + return self.pre_fc_norm_hidden(hidden_states) + + inputs_embeds, hidden_states = maybe_execute_in_parallel( + norm_embeds, + norm_hidden, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + disable_on_compile=True, + ) + hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1) - hidden_states = self.fc(hidden_states) + tp_size = self.model_config.mapping.tp_size + tp_rank = self.model_config.mapping.tp_rank + if tp_size > 1 and not self.model_config.mapping.enable_attention_dp: + hidden_states = torch.chunk(hidden_states, tp_size, + dim=-1)[tp_rank] - hidden_states, residual = super().forward( - position_ids=position_ids, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - residual=None, - spec_metadata=spec_metadata, - **kwargs, - ) - hidden_states, _ = self.shared_head.norm(hidden_states, residual) - if spec_metadata is not None: - spec_metadata.maybe_capture_hidden_states(0, hidden_states, None) + hidden_states = self.fc(hidden_states) - return hidden_states + hidden_states, residual = super().forward( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=None, + spec_metadata=spec_metadata, + **kwargs, + ) + hidden_states, _ = self.shared_head.norm(hidden_states, residual) + if spec_metadata is not None: + spec_metadata.maybe_capture_hidden_states( + 0, hidden_states, None) + + return hidden_states + finally: + # Shared with the target forward, so restore on the exception path + # too; otherwise a failed draft step corrupts every later forward. + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous_all_rank_num_tokens ALL_DECODER_LAYER_TYPES = { diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py index 4ab178c75239..0afe4697cf25 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -145,11 +145,17 @@ def create_strategy( if mapping.moe_tp_size != 1: return AllGatherReduceScatter(mapping) - # Check if forced method is specified via environment variable - force_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD", communication_method) + # A forced method comes either from the environment, which wins, or from the + # model-selected argument. Keep the source with the value so the log below can + # name the one the reader can actually go and change. + env_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD") + if env_method is not None: + force_method, force_source = env_method, "TRTLLM_FORCE_COMM_METHOD" + else: + force_method, force_source = communication_method, "communication_method" if force_method is not None: - return CommunicationFactory._create_forced_method( + strategy = CommunicationFactory._create_forced_method( force_method, model_config, num_experts, @@ -161,6 +167,11 @@ def create_strategy( use_flashinfer, hidden_size=hidden_size, ) + logger.info( + f"Selected communication strategy: {strategy.__class__.__name__} " + f"({force_source}={force_method})" + ) + return strategy # Auto-selection: Try strategies in priority order using try-catch # Priority: NVLinkOneSided > NVLinkTwoSided > NcclEP > DeepEP > DeepEPLowLatency > AllGather diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index f86ecf887aee..698fb5b7594a 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -638,7 +638,13 @@ def _forward_multiple_chunks( ) # ========== Empty-chunk substitution (DP only) ========== - chunked_used = torch.ones(num_chunks, dtype=torch.bool) + # Host-only bookkeeping, so keep it in Python state. A tensor here + # costs an allocation per call plus a Tensor.__bool__ dispatch per + # chunk read below, on a path that runs once per MoE layer per step. + # It is also a latent hazard: the tensor only lands on the host + # because no device is requested, and a CUDA one would turn each read + # into a device-to-host sync that is illegal under CUDA Graph capture. + chunked_used = [True] * num_chunks if moe.use_dp: # The split heuristic guarantees chunk 0 has >= 1 token, so it can # stand in for any empty chunk on this rank. Without substitution, diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index e085d0c390d2..057d0ae50a3a 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -37,6 +37,7 @@ l0_a10: # test list either). - unittest/_torch/models/checkpoints - unittest/_torch/models/test_qwen3_next_moe_quant.py + - unittest/_torch/models/test_qwen3_next_eager_fusion.py - unittest/_torch/weight_sharing - unittest/inputs/test_multimodal_input_processor.py - unittest/others/test_cache_transceiver_precheck_config.py diff --git a/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py b/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py index e0a73d9c1f05..f00cb74fc320 100644 --- a/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py +++ b/tests/unittest/_torch/models/test_qwen3_next_eager_fusion.py @@ -16,16 +16,20 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import pytest import torch from torch import nn from tensorrt_llm._torch.distributed import AllReduceFusionOp from tensorrt_llm._torch.models.modeling_qwen3_next import ( Qwen3NextForCausalLM, + Qwen3NextFullAttentionDecoderLayer, Qwen3NextLinearDecoderLayer, + Qwen3NextMTP, _eager_fusion_enabled, ) from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.utils import EventType def _new_causal_lm() -> Qwen3NextForCausalLM: @@ -34,6 +38,76 @@ def _new_causal_lm() -> Qwen3NextForCausalLM: return model +def _new_mtp() -> Qwen3NextMTP: + model = Qwen3NextMTP.__new__(Qwen3NextMTP) + nn.Module.__init__(model) + model.pre_fc_norm_embedding = nn.Identity() + model.pre_fc_norm_hidden = nn.Identity() + model.fc = nn.Identity() + model.shared_head = nn.Module() + model.shared_head.norm = MagicMock( + side_effect=lambda hidden_states, residual: (hidden_states, None) + ) + model.event_dict = {EventType.Main: None, EventType.MoeShared: None} + model.aux_stream = None + model.model_config = SimpleNamespace( + mapping=SimpleNamespace(tp_size=1, tp_rank=0, enable_attention_dp=True) + ) + return model + + +@torch.no_grad() +def test_mtp_forward_uses_and_restores_draft_rank_token_counts(monkeypatch) -> None: + model = _new_mtp() + target_rank_tokens = [1024, 1024] + draft_rank_tokens = [1, 1] + attn_metadata = SimpleNamespace(all_rank_num_tokens=target_rank_tokens) + + def decoder_forward(self, **kwargs): + assert kwargs["attn_metadata"].all_rank_num_tokens is draft_rank_tokens + return kwargs["hidden_states"], None + + monkeypatch.setattr(Qwen3NextFullAttentionDecoderLayer, "forward", decoder_forward) + + hidden_states = model( + input_ids=torch.tensor([0]), + position_ids=torch.tensor([0]), + hidden_states=torch.ones(1, 2), + embed_tokens=nn.Embedding(1, 2), + attn_metadata=attn_metadata, + all_rank_num_tokens=draft_rank_tokens, + ) + + assert hidden_states.shape == (1, 4) + assert attn_metadata.all_rank_num_tokens is target_rank_tokens + + +@torch.no_grad() +def test_mtp_forward_restores_rank_token_counts_after_failure(monkeypatch) -> None: + model = _new_mtp() + target_rank_tokens = [1024, 1024] + draft_rank_tokens = [1, 1] + attn_metadata = SimpleNamespace(all_rank_num_tokens=target_rank_tokens) + + def decoder_forward(self, **kwargs): + assert kwargs["attn_metadata"].all_rank_num_tokens is draft_rank_tokens + raise RuntimeError("draft forward failed") + + monkeypatch.setattr(Qwen3NextFullAttentionDecoderLayer, "forward", decoder_forward) + + with pytest.raises(RuntimeError, match="draft forward failed"): + model( + input_ids=torch.tensor([0]), + position_ids=torch.tensor([0]), + hidden_states=torch.ones(1, 2), + embed_tokens=nn.Embedding(1, 2), + attn_metadata=attn_metadata, + all_rank_num_tokens=draft_rank_tokens, + ) + + assert attn_metadata.all_rank_num_tokens is target_rank_tokens + + @torch.no_grad() def test_setup_aliases_does_not_read_meta_weights() -> None: model = _new_causal_lm()