[None][feat] Don't review Stack PR to support Qwen3.8-2.4T-A95B-FP8 best perf

ATTRIBUTIONS-Python.md

@@ -5261,7 +5261,7 @@ For more information, please refer to <http://unlicense.org>
   - `Tracker`: https://github.com/tox-dev/py-filelock/issues
 
 
-## flashinfer-python (0.6.16)
+## flashinfer-python (0.6.17.dev20260806)
 
 ### Licenses
 License: `Apache-2.0`

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;
+        // Keep an all-zero activation block finite. Without the floor, scaleOut
+        // is zero and quantizing the zero values evaluates 0 / 0, producing FP8
+        // NaNs that poison FC2 and all following layers. This matches the
+        // epsilon used by the DeepGEMM FP8 activation quantizer.
+        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();

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu

@@ -1629,10 +1629,10 @@ 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).
+    // Cooperative block kernel: fastest path for tiny batches in the large-expert tiers.
+    // It requires an elementwise preprocess (any but softmax-over-experts) and one CUDA
+    // block's worth of experts. The classic one-warp-per-token TopK is faster through the
+    // 512-expert tier, but spills registers in larger tiers (e.g. 896 experts / topK 16).
     bool const preprocessIsElementwise = data.mPreprocessType == RoutingPreprocessType::None
         || data.mPreprocessType == RoutingPreprocessType::Sigmoid
         || data.mPreprocessType == RoutingPreprocessType::SigmoidBias;
@@ -1643,7 +1643,7 @@ void run(Data const& data, void* stream)
         return env != nullptr && env[0] == '1';
     }();
     bool const useCoopBlock = !disableCoopBlock && useStaticBlock && preprocessIsElementwise
-        && dispatchedMaxExperts <= CoopBlockKernelMaxNumExperts;
+        && dispatchedMaxExperts >= CoopBlockKernelMinNumExperts && dispatchedMaxExperts <= CoopBlockKernelMaxNumExperts;
     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,9 @@ static constexpr int MaxNumTokensSingleClusterScores = NumBlocksPerCluster * Num
 static constexpr int BlockKernelMaxNumTokens = 4;
 static constexpr int DynBlockKernelMaxNumTokens = 16;
 static constexpr int DynBlockKernelMaxNumExperts = 256;
+// The classic block kernel is faster through the 512-expert tier. The cooperative
+// kernel avoids register spilling for the larger tiers.
+static constexpr int CoopBlockKernelMinNumExperts = 576;
 // Cooperative block kernel: one thread per expert, so at most 1024 experts (1 CUDA block).
 static constexpr int CoopBlockKernelMaxNumExperts = 1024;
 

cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp

@@ -212,6 +212,15 @@ 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));
+    // FC2 reads activation_output through TRTLLM-Gen's TMA-OOB path. Blackwell TMA requires at
+    // least 128 KiB of backing memory when the descriptor's addressable extent is at least
+    // 128 KiB, even when every logical access is in bounds. A gated activation is half as wide as
+    // gemm1_output, so reusing max_num_padded_tokens_gemm1 can leave its backing allocation below
+    // that contract for small decode batches. Compute the capacity from its own row width instead;
+    // this also grows the associated FP32 scale allocation from 2 KiB to the required 4 KiB minimum.
+    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 +263,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);
 

cpp/tensorrt_llm/thop/moeUtilOp.cpp

@@ -138,9 +138,15 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Te
         = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
     auto permuted_token_selected_experts_tensor
         = torch::empty({num_moe_inputs}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
-    auto permuted_data_tensor = torch::empty({num_moe_inputs, hidden_size}, input.options().requires_grad(false));
+    auto permuted_data_tensor = torch::empty({0, hidden_size}, input.options().requires_grad(false));
     auto permuted_token_final_scales_tensor
-        = torch::empty({num_moe_inputs}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
+        = torch::empty({0}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
+    if (!skip_data_expand)
+    {
+        permuted_data_tensor = torch::empty({num_moe_inputs, hidden_size}, input.options().requires_grad(false));
+        permuted_token_final_scales_tensor
+            = torch::empty({num_moe_inputs}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
+    }
     auto expert_first_token_offset_tensor = torch::empty(
         {num_experts_per_node + 1}, torch::dtype(torch::kInt64).device(torch::kCUDA).requires_grad(false));
     auto unpermuted_row_to_permuted_row_tensor = torch::empty({static_cast<int64_t>(experts_per_token * num_rows)},

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,13 +44,15 @@
 
 #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"
 
 #include <cutlass/numeric_types.h>
 
 #include <algorithm>
+#include <cmath>
 #include <cstdint>
 #include <cstring>
 #include <memory>
@@ -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<std::pair<int32_t, int32_t>> 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 and emit FP8 NaNs, which then poison FC2 and
+// all following layers. 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,50 @@ TEST_F(BlockScaleMoeActivationEquivalenceTest, ZeroScaleBlockProducesIdenticalNa
     auto const permuted = runOnce(kTileForcePermuted);
 
     auto const sfIdx = static_cast<int64_t>(zeroedRow) + static_cast<int64_t>(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<int64_t>(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<uint8_t>(legacy.bytes[idx]), static_cast<uint8_t>(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. Reusing FC1's capacity would therefore allocate only
+    // about half of the backing required by Blackwell's TMA-OOB contract.
+    constexpr int32_t maxNumPaddedTokens = 32;
+    constexpr int32_t intermediateSize = 2304;
+    constexpr int64_t minActivationBytes = 128 * 1024;
+    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<int64_t>(activationCapacity) * intermediateSize * fp8Bits / 8;
+    auto const activationBytesWithGemm1Capacity = static_cast<int64_t>(gemm1Capacity) * intermediateSize * fp8Bits / 8;
+    auto const scaleBytes = static_cast<int64_t>(activationCapacity) * (intermediateSize / kEltsPerSf) * sizeof(float);
+    auto const scaleBytesWithGemm1Capacity
+        = static_cast<int64_t>(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

cpp/tests/unit_tests/kernels/routing/routingCustomTest.cpp

@@ -295,6 +295,30 @@ 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);
+};
+
 TYPED_TEST(RoutingCustomKernelTest, BlockLevelParallelizationWithInvalidTopKInput)
 {
     auto param = RoutingKernelTestParam()

requirements.txt

@@ -58,7 +58,7 @@ ordered-set
 peft>=0.18.1,<0.19.0
 patchelf
 einops
-flashinfer-python==0.6.16
+flashinfer-python @ https://github.com/flashinfer-ai/flashinfer/releases/download/nightly-v0.6.17-20260806/flashinfer_python-0.6.17.dev20260806-py3-none-any.whl#sha256=4ed64b717bf979d268fd7249dfa354d100effe36bc9d340efac8dcebd354611d
 xgrammar==0.1.32
 llguidance==0.7.29
 jsonschema

security_scanning/pyproject.toml

@@ -55,7 +55,7 @@ dependencies = [
     "peft (>=0.18.1,<0.19.0)",
     "patchelf (>=0.19.1.0,<0.20.0.0)",
     "einops (>=0.8.2,<0.9.0)",
-    "flashinfer-python (==0.6.16)",
+    "flashinfer-python @ https://github.com/flashinfer-ai/flashinfer/releases/download/nightly-v0.6.17-20260806/flashinfer_python-0.6.17.dev20260806-py3-none-any.whl#sha256=4ed64b717bf979d268fd7249dfa354d100effe36bc9d340efac8dcebd354611d",
     "xgrammar (==0.1.32)",
     "llguidance (==0.7.29)",
     "jsonschema (>=4.26.0,<5.0.0)",

tensorrt_llm/_torch/model_config.py

@@ -177,6 +177,12 @@ class ModelConfig(Generic[TConfig]):
     max_seq_len: Optional[int] = None
 
     moe_max_num_tokens: Optional[int] = None
+    # Preserve whether the value came from the user across dataclasses.replace().
+    # DeepGEMM uses this metadata to apply its conservative default only when
+    # max_num_tokens was not explicitly configured.
+    _moe_max_num_tokens_is_default: Optional[bool] = field(default=None,
+                                                           repr=False,
+                                                           compare=False)
     moe_load_balancer: Optional[MoeLoadBalancerConfig] = None
 
     attn_backend: str = 'TRTLLM'
@@ -277,10 +283,16 @@ def get_all_reduce_strategy(strategy: str = "AUTO"):
             self.allreduce_strategy = get_all_reduce_strategy(
                 self.allreduce_strategy)
 
-        # Set default moe_max_num_tokens if not specified
-        # The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled
+        # Set default moe_max_num_tokens if not specified. The maximum number
+        # of tokens in MoE is multiplied by DP size when attention DP is
+        # enabled.
+        if self._moe_max_num_tokens_is_default is None:
+            self._moe_max_num_tokens_is_default = (self.moe_max_num_tokens
+                                                   is None)
         if self.moe_max_num_tokens is None:
             self.moe_max_num_tokens = self.max_num_tokens * self.mapping.dp_size
+        if self.moe_max_num_tokens <= 0:
+            raise ValueError("moe_max_num_tokens must be a positive integer")
 
     @property
     def torch_dtype(self) -> torch.dtype: