peihu-nv · Oct 31, 2026 · Oct 31, 2026 · Oct 31, 2026 · Oct 31, 2026 · Oct 31, 2026
@@ -91,6 +91,9 @@ list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*")
 # architecture-scoped OBJECT library below.
 list(FILTER SRC_CU EXCLUDE REGEX
      "minimaxM3/(fusedDecodeAttn|attnMega|attnC1)\\.cu$")
+# The MiniMax-M3 sm_107 decode kernels are built as their own sm_107-only OBJECT
+# library below.
+list(FILTER SRC_CU EXCLUDE REGEX "minimaxM3Sm107(Tp[248])?/.*")
 
 if(NOT ENABLE_MULTI_DEVICE)
   list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$")
@@ -141,11 +144,38 @@ set_property(TARGET minimax_m3_attn_src PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_property(TARGET minimax_m3_attn_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
 set_cuda_architectures(minimax_m3_attn_src 100 103)
 
+# MiniMax-M3 sm_107 decode kernels (minimaxM3Sm107Tp{2,4,8}/): written for
+# sm_107 only (PDL griddepcontrol, sm_10x-only PTX), so they are compiled for
+# the sm_10x entries of this build's architectures only (the same targets, e.g.
+# the sm_100f family, that kernels_src compiles its sm_10x code for), not for
+# the older ones. A build without an sm_10x architecture compiles them for
+# sm_107a: their torch ops link against the host launchers, and only the sm_107
+# targets launch them.
+file(GLOB MINIMAX_M3_SM107_SRC minimaxM3Sm107/*.cu minimaxM3Sm107Tp2/*.cu
+     minimaxM3Sm107Tp4/*.cu minimaxM3Sm107Tp8/*.cu)
+add_library(minimax_m3_sm107_src OBJECT ${MINIMAX_M3_SM107_SRC})
+set_property(TARGET minimax_m3_sm107_src PROPERTY POSITION_INDEPENDENT_CODE ON)
+set_property(TARGET minimax_m3_sm107_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS
+                                                  ON)
+target_include_directories(
+  minimax_m3_sm107_src
+  PRIVATE
+    $<TARGET_PROPERTY:${INTERNAL_CUTLASS_KERNELS_TARGET},INTERFACE_INCLUDE_DIRECTORIES>
+)
+target_link_libraries(minimax_m3_sm107_src PRIVATE trtllm_gen_fmha_interface)
+set(MINIMAX_M3_SM107_ARCHS ${CMAKE_CUDA_ARCHITECTURES})
+list(FILTER MINIMAX_M3_SM107_ARCHS INCLUDE REGEX "^10[0-9]")
+if(NOT MINIMAX_M3_SM107_ARCHS)
+  set(MINIMAX_M3_SM107_ARCHS 107a-real)
+endif()
+set_property(TARGET minimax_m3_sm107_src PROPERTY CUDA_ARCHITECTURES
+                                                  ${MINIMAX_M3_SM107_ARCHS})
+
 add_library(
   kernels_src STATIC
   ${SRC_CPP} ${SRC_CU} ${MOE_KERNELS_SRC}
   $<$<TARGET_EXISTS:marlin_src>:$<TARGET_OBJECTS:marlin_src>>
-  $<TARGET_OBJECTS:minimax_m3_attn_src>)
+  $<TARGET_OBJECTS:minimax_m3_attn_src> $<TARGET_OBJECTS:minimax_m3_sm107_src>)
 set_property(TARGET kernels_src PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_property(TARGET kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
 target_include_directories(

@@ -520,36 +520,13 @@ void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, i
     launchGemv<1>(x, ldx, m, k, w, e, out, qOut, sfOut, sfScale, pdl, stream);
 }
 
-void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out,
-    uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream)
+int minimaxM3RouterTcCluster(int32_t k)
 {
-    int64_t const M = m, K = k;
-    TLLM_CHECK_WITH_INFO(
-        M >= 1 && M <= kTcMaxM && K % (kTcSlices * 64) == 0, "unsupported router shape M=%d K=%d", m, k);
-    int const ks = K / kTcSlices;
-    auto launch_tc = [&](auto kernel, int cl) -> cudaError_t
-    {
-        int const kc = (kTcSlices / cl) * ks;
-        size_t const smem = static_cast<size_t>(kTcExperts) * (kc + 4) * 4 + static_cast<size_t>(M) * (kc + 8) * 2;
-        cudaLaunchConfig_t cfg = {};
-        cfg.gridDim = dim3(8 * cl);
-        cfg.blockDim = dim3(256);
-        cfg.dynamicSmemBytes = smem;
-        cfg.stream = stream;
-        cudaLaunchAttribute attr[2];
-        attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-        attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
-        attr[1].id = cudaLaunchAttributeClusterDimension;
-        attr[1].val.clusterDim.x = cl;
-        attr[1].val.clusterDim.y = 1;
-        attr[1].val.clusterDim.z = 1;
-        cfg.attrs = attr;
-        cfg.numAttrs = 2;
-        return cudaLaunchKernelEx(
-            &cfg, kernel, x, ldx, w, out, static_cast<int>(M), static_cast<int>(K), qOut, sfOut, sfScale);
-    };
-    // 16-CTA clusters (128 CTAs, one K slice each) when the device co-schedules all 8, else 8-CTA clusters.
-    static int const cluster = [&]()
+    // 16-CTA clusters (128 CTAs, one K slice each) when the device co-schedules all 8, else 8-CTA clusters; decided
+    // once per process (device-static, for the first K seen: MiniMax-M3 has one hidden size) and configures both
+    // kernels.
+    int const ks = k / kTcSlices;
+    static int const cluster = [ks]()
     {
         int device = 0, max_optin = 0;
         cudaGetDevice(&device);
@@ -592,6 +569,38 @@ void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int
         }
         return clusters >= 8 ? 16 : 8;
     }();
+    return cluster;
+}
+
+void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out,
+    uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream)
+{
+    int64_t const M = m, K = k;
+    TLLM_CHECK_WITH_INFO(
+        M >= 1 && M <= kTcMaxM && K % (kTcSlices * 64) == 0, "unsupported router shape M=%d K=%d", m, k);
+    int const ks = K / kTcSlices;
+    auto launch_tc = [&](auto kernel, int cl) -> cudaError_t
+    {
+        int const kc = (kTcSlices / cl) * ks;
+        size_t const smem = static_cast<size_t>(kTcExperts) * (kc + 4) * 4 + static_cast<size_t>(M) * (kc + 8) * 2;
+        cudaLaunchConfig_t cfg = {};
+        cfg.gridDim = dim3(8 * cl);
+        cfg.blockDim = dim3(256);
+        cfg.dynamicSmemBytes = smem;
+        cfg.stream = stream;
+        cudaLaunchAttribute attr[2];
+        attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+        attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+        attr[1].id = cudaLaunchAttributeClusterDimension;
+        attr[1].val.clusterDim.x = cl;
+        attr[1].val.clusterDim.y = 1;
+        attr[1].val.clusterDim.z = 1;
+        cfg.attrs = attr;
+        cfg.numAttrs = 2;
+        return cudaLaunchKernelEx(
+            &cfg, kernel, x, ldx, w, out, static_cast<int>(M), static_cast<int>(K), qOut, sfOut, sfScale);
+    };
+    int const cluster = minimaxM3RouterTcCluster(k);
     cudaError_t const err
         = cluster == 16 ? launch_tc(m3_router_tc_kernel<16>, 16) : launch_tc(m3_router_tc_kernel<8>, 8);
     if (err != cudaSuccess)

@@ -41,6 +41,10 @@ void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, i
 void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out,
     uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream);
 
+//! The cluster size `invokeMinimaxM3RouterTc` launches with (16 or 8): decided once per process from the device's
+//! co-scheduling (device-static) and configures both kernels. SM107 callers use the router only when this is 16.
+int minimaxM3RouterTcCluster(int32_t k);
+
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END
@@ -43,6 +43,9 @@ namespace specTail
 namespace
 {
 constexpr int kVec = 8; // bf16 per 16-byte access
+// draft_input loads its constant norm weight row before the grid dependency wait (the same values, bit for bit).
+// Kill: false (load the row after the wait, with the normalization).
+constexpr bool kDraftInputPrewaitWeight = true;
 
 __device__ __forceinline__ float block_sum(float v, float* red)
 {
@@ -72,6 +75,20 @@ __global__ void __launch_bounds__(kThreads)
     __shared__ float red[kThreads / 32];
     int const row = blockIdx.x;
     bool const is_hidden = blockIdx.y == 1;
+    // The norm weight row is constant (no grid of the step writes it): load it before the grid dependency wait, so
+    // after the wait only the dependent loads (the id / hidden row) remain ahead of the reduction. The same values and
+    // arithmetic.
+    __nv_bfloat16 const* const w = is_hidden ? w_hidden : w_embed;
+    uint4 wv[kIters];
+    if constexpr (kDraftInputPrewaitWeight)
+    {
+#pragma unroll
+        for (int it = 0; it < kIters; ++it)
+        {
+            int const c = (it * kThreads + threadIdx.x) * kVec;
+            wv[it] = c < H ? *reinterpret_cast<uint4 const*>(w + c) : make_uint4(0u, 0u, 0u, 0u);
+        }
+    }
 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
     asm volatile("griddepcontrol.wait;" ::: "memory");
 #endif
@@ -112,16 +129,18 @@ __global__ void __launch_bounds__(kThreads)
     asm volatile("griddepcontrol.launch_dependents;");
 #endif
     float const rs = rsqrtf(block_sum(ss, red) / H + (is_hidden ? eps_hidden : eps_embed));
-    __nv_bfloat16 const* w = is_hidden ? w_hidden : w_embed;
     __nv_bfloat16* dst = out + static_cast<long long>(row) * 2 * H + (is_hidden ? H : 0);
 #pragma unroll
     for (int it = 0; it < kIters; ++it)
     {
         int const c = (it * kThreads + threadIdx.x) * kVec;
         if (c < H)
         {
-            uint4 const wr = *reinterpret_cast<uint4 const*>(w + c);
-            __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr);
+            if constexpr (!kDraftInputPrewaitWeight)
+            {
+                wv[it] = *reinterpret_cast<uint4 const*>(w + c);
+            }
+            __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wv[it]);
             uint4 o;
             __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o);
 #pragma unroll

@@ -0,0 +1,189 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * 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.
+ */
+
+// The sm_107 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's
+// GPT-NeoX path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin
+// negated, dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale
+// E4M3 store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query
+// saturates the same way (448, not NaN).
+
+#include "tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.h"
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+namespace minimaxM3Sm107
+{
+namespace drafterRopeFp8Kv
+{
+namespace
+{
+
+constexpr int kElemsPerThread = kHeadDim / 32;
+constexpr int kBlockSize = 256;
+
+template <int kRopeMode>
+__device__ __forceinline__ float rotate(float c, float x, float s, float pair)
+{
+    if constexpr (kRopeMode == 0)
+    {
+        return __fadd_rn(__fmul_rn(c, x), __fmul_rn(s, pair));
+    }
+    else if constexpr (kRopeMode == 1)
+    {
+        return __fmaf_rn(c, x, __fmul_rn(s, pair));
+    }
+    else
+    {
+        return __fmaf_rn(s, pair, __fmul_rn(c, x));
+    }
+}
+
+__device__ __forceinline__ void storeE4m3Pairs(__nv_fp8_e4m3* dst, __nv_bfloat16 const (&v)[kElemsPerThread])
+{
+#pragma unroll
+    for (int i = 0; i < kElemsPerThread; i += 2)
+    {
+        __nv_bfloat162 pair;
+        pair.x = v[i];
+        pair.y = v[i + 1];
+        __nv_fp8x2_e4m3 const packed(pair);
+        reinterpret_cast<__nv_fp8x2_storage_t*>(dst)[i / 2] = packed.__x;
+    }
+}
+
+template <int kRopeMode>
+__global__ void __launch_bounds__(kBlockSize) drafterRopeFp8KvKernel(Params const p)
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+    asm volatile("griddepcontrol.launch_dependents;");
+#endif
+    int const lane = threadIdx.x % 32;
+    int const warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
+    int const slots = p.numHeadsQ + p.numHeadsKv;
+    int const token = warp / slots;
+    if (token >= p.numTokens)
+    {
+        return;
+    }
+    int const slot = warp % slots;
+    bool const isQ = slot < p.numHeadsQ;
+    int const head = isQ ? slot : slot - p.numHeadsQ;
+    int const request = token / p.qLen;
+    int const row = token % p.qLen;
+    int const cacheLen = p.seqLens[request];
+    int const kvPos = cacheLen - p.qLen + row;
+    int const ropePos
+        = p.specPositionOffsets != nullptr ? p.specPositionOffsets[request * p.qLen + row] + cacheLen - p.qLen : kvPos;
+    int const rowWidth = (p.numHeadsQ + 2 * p.numHeadsKv) * kHeadDim;
+    __nv_bfloat16 const* src
+        = p.qkv + static_cast<long long>(token) * rowWidth + (isQ ? head * kHeadDim : (p.numHeadsQ + head) * kHeadDim);
+    int const half = p.rotaryDim / 2;
+    float2 const* coef = p.cosSin + static_cast<long long>(ropePos) * half;
+
+    __nv_bfloat16 rotated[kElemsPerThread];
+#pragma unroll
+    for (int i = 0; i < kElemsPerThread; ++i)
+    {
+        int const d = lane * kElemsPerThread + i;
+        bool const firstHalf = d < half;
+        float const x = __bfloat162float(src[d]);
+        float const pair = __bfloat162float(src[firstHalf ? d + half : d - half]);
+        float2 const cs = d < p.rotaryDim ? coef[firstHalf ? d : d - half] : make_float2(1.0f, 0.0f);
+        float const s = firstHalf ? -cs.y : cs.y;
+        rotated[i] = __float2bfloat16(rotate<kRopeMode>(cs.x, x, s, pair));
+    }
+
+    if (isQ)
+    {
+        storeE4m3Pairs(
+            p.qOut + (static_cast<long long>(token) * p.numHeadsQ + head) * kHeadDim + lane * kElemsPerThread, rotated);
+        return;
+    }
+
+    int const block = kvPos / p.tokensPerBlock;
+    long long const inPage = (static_cast<long long>(head) * p.tokensPerBlock + kvPos % p.tokensPerBlock) * kHeadDim
+        + lane * kElemsPerThread;
+    long long const pageElems = static_cast<long long>(p.numHeadsKv) * p.tokensPerBlock * kHeadDim;
+    int const* table = p.blockTable + request * p.blockTableRequestStride;
+    long long const kPage = table[block];
+    long long const vPage = table[p.blockTableKvStride + block];
+    storeE4m3Pairs(p.pool + kPage * pageElems + inPage, rotated);
+
+    __nv_bfloat16 const* vSrc = p.qkv + static_cast<long long>(token) * rowWidth
+        + (p.numHeadsQ + p.numHeadsKv + head) * kHeadDim + lane * kElemsPerThread;
+    __nv_bfloat16 v[kElemsPerThread];
+#pragma unroll
+    for (int i = 0; i < kElemsPerThread; ++i)
+    {
+        v[i] = vSrc[i];
+    }
+    storeE4m3Pairs(p.pool + vPage * pageElems + inPage, v);
+}
+
+//! With PDL and maxSharedCarveout, the maximum shared-memory carveout lets the full-context decode's large-smem CTAs
+//! (the dependent this kernel releases right after its wait) co-reside on the SMs this small kernel occupies. Set once
+//! per instantiation.
+template <int kRopeMode>
+cudaError_t launchWithMaxSharedCarveout(
+    cudaLaunchConfig_t const& config, Params const& params, bool pdl, bool maxSharedCarveout)
+{
+    auto const kernel = drafterRopeFp8KvKernel<kRopeMode>;
+    if (pdl && maxSharedCarveout)
+    {
+        static cudaError_t const carveout = cudaFuncSetAttribute(
+            kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared);
+        if (carveout != cudaSuccess)
+        {
+            return carveout;
+        }
+    }
+    return cudaLaunchKernelEx(&config, kernel, params);
+}
+
+} // namespace
+
+cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream)
+{
+    int const warps = params.numTokens * (params.numHeadsQ + params.numHeadsKv);
+    int const warpsPerBlock = kBlockSize / 32;
+    cudaLaunchConfig_t config = {};
+    config.gridDim = dim3((warps + warpsPerBlock - 1) / warpsPerBlock);
+    config.blockDim = dim3(kBlockSize);
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    cudaLaunchAttribute attribute[1];
+    attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0;
+    config.attrs = attribute;
+    config.numAttrs = 1;
+    switch (params.ropeMode)
+    {
+    case 0: return launchWithMaxSharedCarveout<0>(config, params, pdl, maxSharedCarveout);
+    case 1: return launchWithMaxSharedCarveout<1>(config, params, pdl, maxSharedCarveout);
+    case 2: return launchWithMaxSharedCarveout<2>(config, params, pdl, maxSharedCarveout);
+    default: return cudaErrorInvalidValue;
+    }
+}
+
+} // namespace drafterRopeFp8Kv
+} // namespace minimaxM3Sm107
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END