[TRTLLM-15304][fix] MoE and MTP fixes for large hybrid-attention FP8 serving

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<Type>(out[i] / scaleOut);
         }
         *reinterpret_cast<PackedIo*>(params.outPtr + static_cast<int64_t>(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
                 {

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 <cstdlib>
 
@@ -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;

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;
 

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 <cstdint>
+
+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

cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp

@@ -22,6 +22,7 @@
 #include <ATen/cuda/EmptyTensor.h>
 #include <torch/library.h>
 
+#include <algorithm>
 #include <cstdint>
 #include <cstdlib>
 #include <memory>
@@ -212,6 +213,13 @@ at::Tensor run_fp8_block_scale_moe(at::optional<at::Tensor> 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<at::Tensor> 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<at::Tensor> const& routing_logit
 
     // setup workspace
     workspace.total_num_padded_tokens = total_num_padded_tokens.data_ptr<int>();
-    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<int>();
     workspace.permuted_idx_size = total_num_padded_tokens.data_ptr<int>();
     workspace.expanded_idx_to_permuted_idx