[None][fix] Fix Qwen3.5 MoE fallback, GDN alignment, FP8 activation, and draft KV cache

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-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.
@@ -22,6 +22,9 @@
 #include <cutlass/cutlass.h>
 #include <cutlass/numeric_types.h>
 
+#include <algorithm>
+#include <cstdint>
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 // Helper function for array conversion
@@ -231,6 +234,145 @@ struct KernelTraits<1>
 
 constexpr int DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA = 128;
 
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Permuted-space SwiGLU for the DeepSeek-FP8 separate-activation path.
+//
+// `activationDeepSeekKernel` below grids over the *expanded* index space
+// (numTokens x topK) and discovers work by loading expandedIdxToPermutedIdx,
+// skipping entries that map to -1. Under expert parallelism only 1/ep_size of
+// those entries are local, so most of the launched CTAs do no memory work at
+// all -- yet they still run the unconditional cub::BlockReduce. At a large
+// context and a high expert-parallel degree the launched CTA count exceeds the
+// permuted rows of real work by the ep_size factor, and the achieved bandwidth
+// is a small fraction of what the row count alone would need.
+//
+// Every memory access in that kernel is addressed by (permutedIdx, hiddenIdx)
+// only -- the expanded index exists purely to find the work. So grid directly
+// over the permuted rows instead and the indirection, the -1 slots and the
+// ep_size-fold CTA inflation all disappear together.
+//
+// Layout: one warp owns exactly one (permutedRow, 128-element scale block).
+// 32 lanes x 4 elements = 128 = one scale block, so the amax reduction is a
+// single warp shuffle instead of a shared-memory block reduce, and each lane
+// moves 4 bytes per load instead of 1.
+//
+// totalNumPaddedTokens is only known on the device, so the grid is persistent
+// 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.
+constexpr int kDsActWarpSize = 32;
+constexpr int kDsActEltsPerSf = 128;
+constexpr int kDsActEltsPerThread = kDsActEltsPerSf / kDsActWarpSize;
+constexpr int kDsActWarpsPerCta = 4;
+constexpr int kDsActPermutedNumThreadsPerCta = kDsActWarpSize * kDsActWarpsPerCta;
+
+constexpr bool shouldUsePermutedActivation(int innerDim, int numTokens, int topK, int numExperts, int tileTokensDim)
+{
+    int const outputDim = innerDim / 2;
+    bool const layoutEligible = outputDim >= kDsActEltsPerSf && outputDim % kDsActEltsPerSf == 0 && innerDim % 8 == 0;
+    int64_t const realRowsPerExpert = numExperts > 0 ? static_cast<int64_t>(numTokens) * topK / numExperts : 0;
+    bool const paddingAmortised = tileTokensDim > 0 && realRowsPerExpert >= tileTokensDim;
+    return layoutEligible && paddingAmortised;
+}
+
+template <typename KernelParams>
+__global__ void activationDeepSeekPermutedKernel(KernelParams params)
+{
+    using Type = typename KernelParams::Type;
+    using PackedIo = uint32_t; // kDsActEltsPerThread x 8-bit elements
+
+    static_assert(kDsActEltsPerThread == 4, "PackedIo assumes 4 elements per thread");
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    if constexpr (KernelParams::UsePdl)
+    {
+        cudaTriggerProgrammaticLaunchCompletion();
+        cudaGridDependencySynchronize();
+    }
+#endif
+
+    float constexpr kE4m3MaxVal{448.F};
+
+    int const totalNumPaddedTokens = params.totalNumPaddedTokens[0];
+    int const outputDim = params.innerDim / 2;
+    int const numSfBlocks = outputDim / kDsActEltsPerSf;
+
+    bool const hasSwigluLimit = params.hasSwigluLimit;
+    float const swigluLimit = params.swigluLimit;
+
+    int const lane = threadIdx.x % kDsActWarpSize;
+    int const warpInCta = threadIdx.x / kDsActWarpSize;
+
+    int64_t const numTasks = static_cast<int64_t>(totalNumPaddedTokens) * numSfBlocks;
+    int64_t const taskStride = static_cast<int64_t>(gridDim.x) * kDsActWarpsPerCta;
+
+    for (int64_t task = static_cast<int64_t>(blockIdx.x) * kDsActWarpsPerCta + warpInCta; task < numTasks;
+         task += taskStride)
+    {
+        int const permutedIdx = static_cast<int>(task / numSfBlocks);
+        int const sfBlock = static_cast<int>(task % numSfBlocks);
+        int const hiddenBase = sfBlock * kDsActEltsPerSf + lane * kDsActEltsPerThread;
+
+        // Both scales are uniform across the warp: one per (row, scale block).
+        float const scale1 = params.inDqSfsPtr[permutedIdx + totalNumPaddedTokens * sfBlock];
+        float const scale2 = params.inDqSfsPtr[permutedIdx + totalNumPaddedTokens * (sfBlock + numSfBlocks)];
+
+        int64_t const baseIdx = static_cast<int64_t>(permutedIdx) * params.innerDim + hiddenBase;
+        PackedIo const packed1 = *reinterpret_cast<PackedIo const*>(params.inPtr + baseIdx);
+        PackedIo const packed2 = *reinterpret_cast<PackedIo const*>(params.inPtr + baseIdx + outputDim);
+
+        Type const* elts1 = reinterpret_cast<Type const*>(&packed1);
+        Type const* elts2 = reinterpret_cast<Type const*>(&packed2);
+
+        float out[kDsActEltsPerThread];
+        float aMax = 0.F;
+#pragma unroll
+        for (int i = 0; i < kDsActEltsPerThread; ++i)
+        {
+            float x1 = scale1 * static_cast<float>(elts1[i]); // up (linear)
+            float x2 = scale2 * static_cast<float>(elts2[i]); // gate (silu input)
+            if (hasSwigluLimit)
+            {
+                x2 = fminf(x2, swigluLimit);
+                x1 = fmaxf(fminf(x1, swigluLimit), -swigluLimit);
+            }
+            out[i] = silu(x2) * x1;
+            aMax = fmaxf(aMax, fabsf(out[i]));
+        }
+
+#pragma unroll
+        for (int offset = kDsActWarpSize / 2; offset > 0; offset >>= 1)
+        {
+            aMax = fmaxf(aMax, __shfl_xor_sync(0xffffffffu, aMax, offset));
+        }
+
+        float const scaleOut = aMax / kE4m3MaxVal;
+
+        if (lane == 0)
+        {
+            params.outDqSfsPtr[permutedIdx + totalNumPaddedTokens * sfBlock] = scaleOut;
+        }
+
+        PackedIo packedOut;
+        Type* outElts = reinterpret_cast<Type*>(&packedOut);
+#pragma unroll
+        for (int i = 0; i < kDsActEltsPerThread; ++i)
+        {
+            // 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.
+            outElts[i] = static_cast<Type>(out[i] / scaleOut);
+        }
+        *reinterpret_cast<PackedIo*>(params.outPtr + static_cast<int64_t>(permutedIdx) * outputDim + hiddenBase)
+            = packedOut;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
 template <typename KernelParams>
 __global__ void activationDeepSeekKernel(KernelParams params)
 {
@@ -435,8 +577,45 @@ void run(Data const& data, void* stream)
 
         const dim3 grid(gridSizeX, gridSizeY, data.topK);
 
-        LAUNCH_ACTIVATION(
-            data, activationDeepSeekKernel, numTokensPerCta, grid, DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA, 0, stream);
+        // The two kernels sweep different spaces, and which one is cheaper flips
+        // with batch size.
+        //
+        // The expanded-space kernel visits numTokens x topK slots and skips the
+        // ~(1 - 1/ep_size) of them that are not local, so it never touches the
+        // per-expert tile padding. The permuted-space kernel sweeps
+        // [0, totalNumPaddedTokens), which *is* padded: each local expert
+        // contributes up to tileTokensDim-1 rows of padding that carry no real
+        // tokens but cost a full row of load/compute/store.
+        //
+        // At prefill that padding is noise next to the real rows, and the
+        // permuted sweep wins by the ep_size factor. At decode the ratio
+        // inverts: a single token leaves well under one real row per expert
+        // against the same padding, so the permuted kernel does almost nothing
+        // but padding. Getting this wrong costs more on every decode step than
+        // the prefill win is worth over a full generation.
+        //
+        // So gate on real work per expert. tileTokensDim is exactly the padding
+        // granularity, which makes it the natural threshold: below it, an
+        // expert's real rows do not even fill the tile that must be swept for it.
+        if (shouldUsePermutedActivation(data.innerDim, data.numTokens, data.topK, data.numExperts, data.tileTokensDim))
+        {
+            int64_t const maxTasks = static_cast<int64_t>(data.numTokens) * data.topK * (outputDim / kDsActEltsPerSf);
+            int64_t const ctasForAllTasks = (maxTasks + kDsActWarpsPerCta - 1) / kDsActWarpsPerCta;
+            // Persistent grid: totalNumPaddedTokens is a device-side value, so the
+            // host can only bound it. Cap at a few waves and let the grid stride
+            // absorb the difference rather than launching the (numTokens x topK)
+            // worst case that the expanded-space kernel pays unconditionally.
+            int const numCtas = static_cast<int>(std::min<int64_t>(ctasForAllTasks, int64_t{numSms} * 32));
+            dim3 const permutedGrid(std::max(numCtas, 1), 1, 1);
+
+            LAUNCH_ACTIVATION(
+                data, activationDeepSeekPermutedKernel, 1, permutedGrid, kDsActPermutedNumThreadsPerCta, 0, stream);
+        }
+        else
+        {
+            LAUNCH_ACTIVATION(data, activationDeepSeekKernel, numTokensPerCta, grid,
+                DEEP_SEEK_ACTIVATION_NUM_THREADS_PER_CTA, 0, stream);
+        }
     }
     else
     {

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-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.
@@ -259,6 +259,14 @@ struct Data
     int32_t topK;
     int32_t* expandedIdxToPermutedIdx;
 
+    // Used only to pick between the two activation kernels; see the dispatch
+    // note in DevKernel.cu. The permuted-space kernel sweeps
+    // totalNumPaddedTokens, which carries up to one partial tile of padding per
+    // local expert, so it only pays off when there is enough real work per
+    // expert to amortise that sweep. Both are host-side values.
+    int32_t numExperts{0};
+    int32_t tileTokensDim{0};
+
     int32_t const* totalNumPaddedTokens;
 
     // Optional swiglu clamp limit (fp32, uniform across experts on the FP8

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-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.
@@ -622,6 +622,7 @@ Runner::Runner(
     : mPermuteGemm1(PermuteGemm1::Runner(dtypeAct, dtypeWeights, useDeepSeekFp8, tileTokensDim, actType))
     , mGemm2(Gemm2::Runner(dtypeAct, dtypeWeights, btg::Dtype::Bfloat16, useDeepSeekFp8, tileTokensDim))
     , mActType(actType)
+    , mTileTokensDim(tileTokensDim)
 {
     auto const& gemm1PassingIndices = mPermuteGemm1.getPassingConfigIndices();
     auto const& gemm2PassingIndices = mGemm2.getPassingConfigIndices();
@@ -670,6 +671,8 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace
     activationData.topK = args.top_k;
     activationData.numTokens = args.num_tokens;
     activationData.expandedIdxToPermutedIdx = workspace.expanded_idx_to_permuted_idx;
+    activationData.numExperts = args.num_experts;
+    activationData.tileTokensDim = mTileTokensDim;
     // For DeepSeek FP8 the activation runs as a separate kernel rather than
     // fused into the FC1 GEMM cubin; forward the scalar swiglu_limit so it
     // can honor swiglu_limit (uniform across experts; see DevKernel.h note).

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-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.
@@ -418,6 +418,9 @@ class Runner
     PermuteGemm1::Runner mPermuteGemm1;
     Gemm2::Runner mGemm2;
     ActType mActType;
+    // Kept so setOpsData can tell the activation launcher how much per-expert
+    // tile padding the permuted row space carries.
+    int32_t mTileTokensDim;
 
     // This will be the cartesian product of the passing configs for gemm1 and gemm2
     // This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2 separately

cpp/tests/unit_tests/kernels/CMakeLists.txt

@@ -100,6 +100,8 @@ set(ROUTING_KERNEL_TEST_SRC
 add_gtest(routingKernelsTest "${ROUTING_KERNEL_TEST_SRC}")
 target_link_libraries(routingKernelsTest PRIVATE Python3::Python)
 
+add_gtest(blockScaleMoeActivationTest blockScaleMoeActivationTest.cu)
+
 add_gtest(moeLoadBalanceKernelTest moeLoadBalanceKernelTest.cpp)
 
 if(USING_OSS_CUTLASS_MOE_GEMM)