[None][feat] Kimi K3 optimizations on Rubin

cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu

@@ -18,6 +18,7 @@
 #include "tensorrt_llm/common/envUtils.h"
 
 #include "tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.h"
+#include <cstdint>
 using namespace tensorrt_llm::common;
 
 TRTLLM_NAMESPACE_BEGIN
@@ -194,6 +195,216 @@ void invokeRouterGemm(float* output, T const* mat_a, T const* mat_b, cudaStream_
         &config, router_gemm_kernel<T, kBlockSize, VPT, kNumTokens, kNumExperts, kHiddenDim>, output, mat_a, mat_b));
 }
 
+// D += A(16x16, row) * B(16x8, col), bf16 inputs, fp32 accumulate. Fragment layout (PTX ISA, m16n8k16):
+//   a[0] = A[g][2q..2q+1]  a[1] = A[g+8][2q..2q+1]  a[2] = A[g][2q+8..2q+9]  a[3] = A[g+8][2q+8..2q+9]
+//   b[0] = B[2q..2q+1][g]  b[1] = B[2q+8..2q+9][g]   c[0..1] = D[g][2q..2q+1]  c[2..3] = D[g+8][2q..2q+1]
+// with g = lane / 4, q = lane % 4.
+__device__ __forceinline__ void mma_bf16_m16n8k16(float (&c)[4], uint32_t const (&a)[4], uint32_t const (&b)[2])
+{
+    asm volatile(
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
+        : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
+}
+
+// Tensor-core variant for large expert counts (Kimi K3: 896 experts). One warp computes a 16-token x 8-expert
+// tile over a K-slice with mma.sync; the kNumWarps warps of a CTA split K and are reduced through shared memory
+// in warp order (deterministic). Each lane loads 16 bytes (8 consecutive k) per row/column per 32-wide k-chunk;
+// the k order inside a chunk is permuted identically for A and B (lane group q owns k = q*8..q*8+7, fed to the
+// two MMAs as k-pairs {0,1},{2,3} and {4,5},{6,7}), which is legal because the dot product is order-invariant.
+// Motivation (B200, M=8): router_gemm_kernel executes ~4M warp-instructions (M x 896 x 7168 FMAs + bf16->fp32
+// converts) and is issue-bound (~26% issue utilisation, time ~linear in M); one MMA replaces 64 warp-FFMAs.
+template <int kNumTokens, int kNumExperts, int kHiddenDim, int kExpertsPerCta, int kBlockSize, int kLatentDim = 0>
+__device__ __forceinline__ void routerGemmMma(float* out, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b,
+    __nv_bfloat16* projection, __nv_bfloat16 const* downWeight)
+{
+    static_assert(kNumTokens >= 1 && kNumTokens <= 16, "one m16 tile in the token dimension");
+    static_assert(kExpertsPerCta % 8 == 0, "expert columns per CTA must be whole n8 tiles");
+    static_assert(kLatentDim == 0 || (kNumExperts % kExpertsPerCta == 0 && kLatentDim % kExpertsPerCta == 0),
+        "gate and latent output boundaries must align to CTA tiles");
+    constexpr int kNumWarps = kBlockSize / 32;
+    constexpr int kTilesN = kExpertsPerCta / 8;
+    constexpr int kChunk = 32; // k elements per lane load (8 per lane group) = two m16n8k16 MMAs
+    static_assert(kHiddenDim % (kNumWarps * kChunk) == 0, "K must split evenly over warps and 32-wide chunks");
+    constexpr int kChunksPerWarp = kHiddenDim / (kNumWarps * kChunk);
+    constexpr int kBatch = kChunksPerWarp < 8 ? kChunksPerWarp : 8; // chunks whose loads are kept in flight
+    static_assert(kChunksPerWarp % kBatch == 0, "chunks per warp must be a multiple of the load batch");
+    constexpr bool kTwoRows = kNumTokens > 8;
+
+    int const warpId = threadIdx.x / 32;
+    int const lane = threadIdx.x % 32;
+    int const g = lane / 4; // token row (A) / expert column (B) owned by this lane in the fragments
+    int const q = lane % 4; // 8-wide k sub-chunk owned by this lane
+    int const n_base = blockIdx.x * kExpertsPerCta;
+    int const k_warp = warpId * (kChunksPerWarp * kChunk);
+
+    __shared__ float sm_red[kNumWarps][kTilesN][32][4];
+
+    float acc[kTilesN][4] = {};
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    cudaGridDependencySynchronize();
+#endif
+
+    __nv_bfloat16 const* a_row0 = mat_a + g * kHiddenDim;
+    __nv_bfloat16 const* a_row1 = mat_a + (g + 8) * kHiddenDim;
+    __nv_bfloat16 const* weight = mat_b;
+    int weightColumn = n_base;
+    if constexpr (kLatentDim > 0)
+    {
+        if (n_base >= kNumExperts)
+        {
+            weight = downWeight;
+            weightColumn -= kNumExperts;
+        }
+    }
+    bool const row0_valid = g < kNumTokens;
+    bool const row1_valid = kTwoRows && (g + 8 < kNumTokens);
+    uint4 const zero4 = make_uint4(0u, 0u, 0u, 0u);
+
+#pragma unroll
+    for (int batch = 0; batch < kChunksPerWarp; batch += kBatch)
+    {
+        uint4 a0[kBatch];
+        uint4 a1[kBatch];
+        uint4 b[kBatch][kTilesN];
+#pragma unroll
+        for (int c = 0; c < kBatch; c++)
+        {
+            int const k = k_warp + (batch + c) * kChunk + q * 8;
+            a0[c] = row0_valid ? *reinterpret_cast<uint4 const*>(a_row0 + k) : zero4;
+            a1[c] = row1_valid ? *reinterpret_cast<uint4 const*>(a_row1 + k) : zero4;
+#pragma unroll
+            for (int t = 0; t < kTilesN; t++)
+            {
+                b[c][t] = *reinterpret_cast<uint4 const*>(
+                    weight + static_cast<int64_t>(weightColumn + t * 8 + g) * kHiddenDim + k);
+            }
+        }
+#pragma unroll
+        for (int c = 0; c < kBatch; c++)
+        {
+            uint32_t const afrag0[4] = {a0[c].x, a1[c].x, a0[c].y, a1[c].y};
+            uint32_t const afrag1[4] = {a0[c].z, a1[c].z, a0[c].w, a1[c].w};
+#pragma unroll
+            for (int t = 0; t < kTilesN; t++)
+            {
+                uint32_t const bfrag0[2] = {b[c][t].x, b[c][t].y};
+                uint32_t const bfrag1[2] = {b[c][t].z, b[c][t].w};
+                mma_bf16_m16n8k16(acc[t], afrag0, bfrag0);
+                mma_bf16_m16n8k16(acc[t], afrag1, bfrag1);
+            }
+        }
+    }
+
+#pragma unroll
+    for (int t = 0; t < kTilesN; t++)
+    {
+#pragma unroll
+        for (int i = 0; i < 4; i++)
+        {
+            sm_red[warpId][t][lane][i] = acc[t][i];
+        }
+    }
+    __syncthreads();
+
+    // One thread per (tile, lane slot, fragment element): sum the warp partials in warp order and scatter.
+    constexpr int kNumOutputs = kTilesN * 32 * 4;
+    static_assert(kNumOutputs <= kBlockSize, "need one thread per output element");
+    if (threadIdx.x < kNumOutputs)
+    {
+        int const t = threadIdx.x / 128;
+        int const l = (threadIdx.x % 128) / 4;
+        int const i = threadIdx.x % 4;
+        float sum = 0.0f;
+#pragma unroll
+        for (int w = 0; w < kNumWarps; w++)
+        {
+            sum += sm_red[w][t][l][i];
+        }
+        int const row = (l / 4) + (i >= 2 ? 8 : 0);
+        int const col = n_base + t * 8 + (l % 4) * 2 + (i & 1);
+        if (row < kNumTokens)
+        {
+            if constexpr (kLatentDim > 0)
+            {
+                if (n_base >= kNumExperts)
+                {
+                    projection[row * kLatentDim + col - kNumExperts] = __float2bfloat16_rn(sum);
+                }
+                else
+                {
+                    out[row * kNumExperts + col] = sum;
+                }
+            }
+            else
+            {
+                out[row * kNumExperts + col] = sum;
+            }
+        }
+    }
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
+template <int kNumTokens, int kNumExperts, int kHiddenDim, int kExpertsPerCta, int kBlockSize>
+__global__ __launch_bounds__(kBlockSize, 1) void router_gemm_mma_kernel(
+    float* out, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b)
+{
+    routerGemmMma<kNumTokens, kNumExperts, kHiddenDim, kExpertsPerCta, kBlockSize>(out, mat_a, mat_b, nullptr, nullptr);
+}
+
+template <int kNumTokens>
+__global__ __launch_bounds__(448, 1) void router_latent_gemm_mma_kernel(float* logits, __nv_bfloat16* projection,
+    __nv_bfloat16 const* input, __nv_bfloat16 const* gateWeight, __nv_bfloat16 const* downWeight)
+{
+    routerGemmMma<kNumTokens, 896, 7168, 8, 448, 3584>(logits, input, gateWeight, projection, downWeight);
+}
+
+template <int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeRouterGemmMma(float* output, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b, cudaStream_t stream)
+{
+    // 8 expert columns per CTA and 14 warps splitting K: for 896 x 7168 that is 112 CTAs (one wave on B200)
+    // each doing 16 k-chunks per warp. Measured B200 (graph replay, L2-warm) vs the cuBLAS fallback:
+    // M=1 2.79 vs 7.36 us, M=8 3.77 vs 7.62 us, M=16 4.70 vs 6.64 us.
+    constexpr int kExpertsPerCta = 8;
+    constexpr int kBlockSize = 448;
+    static_assert(kNumExperts % kExpertsPerCta == 0, "kNumExperts must be a multiple of kExpertsPerCta");
+    cudaLaunchConfig_t config;
+    config.gridDim = kNumExperts / kExpertsPerCta;
+    config.blockDim = kBlockSize;
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    cudaLaunchAttribute attrs[1];
+    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL();
+    config.numAttrs = 1;
+    config.attrs = attrs;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config,
+        router_gemm_mma_kernel<kNumTokens, kNumExperts, kHiddenDim, kExpertsPerCta, kBlockSize>, output, mat_a, mat_b));
+}
+
+template <int kNumTokens>
+void invokeRouterLatentGemmMma(float* logits, __nv_bfloat16* projection, __nv_bfloat16 const* input,
+    __nv_bfloat16 const* gateWeight, __nv_bfloat16 const* downWeight, cudaStream_t stream)
+{
+    constexpr int kColumns = 896 + 3584;
+    constexpr int kColumnsPerCta = 8;
+    constexpr int kBlockSize = 448;
+    cudaLaunchConfig_t config{};
+    config.gridDim = kColumns / kColumnsPerCta;
+    config.blockDim = kBlockSize;
+    config.stream = stream;
+    cudaLaunchAttribute attr{};
+    attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attr.val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL();
+    config.numAttrs = 1;
+    config.attrs = &attr;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(
+        &config, router_latent_gemm_mma_kernel<kNumTokens>, logits, projection, input, gateWeight, downWeight));
+}
+
 template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 1, 256, 7168>(
     float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
 
@@ -340,6 +551,49 @@ template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__n
 template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 4096>(
     float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
 
+// num_experts=896, hidden_dim=7168 instantiations (Kimi K3), tensor-core variant.
+#define INSTANTIATE_ROUTER_GEMM_MMA_K3(kNumTokens)                                                                     \
+    template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemmMma<kNumTokens, 896, 7168>(            \
+        float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+INSTANTIATE_ROUTER_GEMM_MMA_K3(1)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(2)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(3)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(4)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(5)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(6)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(7)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(8)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(9)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(10)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(11)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(12)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(13)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(14)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(15)
+INSTANTIATE_ROUTER_GEMM_MMA_K3(16)
+#undef INSTANTIATE_ROUTER_GEMM_MMA_K3
+
+#define INSTANTIATE_ROUTER_LATENT_GEMM_MMA(kNumTokens)                                                                 \
+    template void invokeRouterLatentGemmMma<kNumTokens>(                                                               \
+        float*, __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(1)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(2)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(3)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(4)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(5)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(6)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(7)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(8)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(9)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(10)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(11)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(12)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(13)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(14)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(15)
+INSTANTIATE_ROUTER_LATENT_GEMM_MMA(16)
+#undef INSTANTIATE_ROUTER_LATENT_GEMM_MMA
+
 } // namespace kernels::dsv3MinLatencyKernels
 
 TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2019-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.
@@ -19,6 +19,7 @@
 #include "tensorrt_llm/common/config.h"
 #include "tensorrt_llm/common/cudaUtils.h"
 #include <assert.h>
+#include <cuda_bf16.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
@@ -30,6 +31,16 @@ namespace kernels::dsv3MinLatencyKernels
 template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
 void invokeRouterGemm(float* output, T const* mat_a, T const* mat_b, cudaStream_t stream);
 
+// Tensor-core variant (mma.sync m16n8k16, bf16 in / fp32 out) for large expert counts such as Kimi K3's 896,
+// where the scalar FFMA kernel above is instruction-bound. kNumTokens <= 16, kNumExperts % 8 == 0.
+template <int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeRouterGemmMma(float* output, __nv_bfloat16 const* mat_a, __nv_bfloat16 const* mat_b, cudaStream_t stream);
+
+// BF16 input [M,7168], gate [896,7168], down [3584,7168]; FP32 logits and BF16 projection, M <= 16.
+template <int kNumTokens>
+void invokeRouterLatentGemmMma(float* logits, __nv_bfloat16* projection, __nv_bfloat16 const* input,
+    __nv_bfloat16 const* gateWeight, __nv_bfloat16 const* downWeight, cudaStream_t stream);
+
 } // namespace kernels::dsv3MinLatencyKernels
 
 TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomKernels.cuh

@@ -840,8 +840,13 @@ void launchCoopBlockKernel(Data const& data, uint32_t numThreadsHist, void* stre
 //
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+// launchDynBlockKernel launches min(max(tokens*32, MaxExperts), 1024) threads, i.e. 1024 for the
+// 1024-expert tier now reachable from the post-topK path. Without a bound the 1024-expert
+// instantiation compiled to 89 regs/thread -> 89*1024 > 64K -> cudaLaunchKernelEx "too many
+// resources requested for launch". The <=576-expert tiers already use <=62 regs, so the bound only
+// affects the large tiers.
 template <typename KernelParams>
-__global__ void routingIndicesDynBlockKernel(KernelParams params)
+__global__ void __launch_bounds__(1024) routingIndicesDynBlockKernel(KernelParams params)
 {
     using OutputT = typename KernelParams::OutputT;
     using InputT = typename KernelParams::InputT;
@@ -1734,6 +1739,20 @@ void launchClusterKernelBlockDim1024(Data const& data, void* stream)
 #if defined(TRTLLM_ROUTING_CUSTOM_ENTRY)
 void launchClusterKernel(Data const& data, void* stream)
 {
+    // Post-topK (mPtrScores == nullptr) only permutes: the cluster kernel is thread-per-expanded-index
+    // there, so the 512-thread variant is not bound by the warp-per-token limit below (capacity
+    // NumBlocksPerCluster * 512 = 4096 tokens). Measured on B200 (896 experts/top-16, E=1024 tier),
+    // same GPU: 512 threads (2 experts/thread) 4.3/4.6/4.8/5.1/5.4/6.1 us at 17/32/64/224/512/1024
+    // tokens vs 256 threads (4 experts/thread) 5.2/5.7/6.2 us at 16/32/64 and 1024 threads
+    // (1 expert/thread) ~5.9 us at 224-512. At 4096 tokens the 512-thread variant degrades to ~9.8 us
+    // (16 expanded indices per thread), so beyond 1024 tokens keep the 1024-thread launch.
+    static constexpr int PostTopKClusterBlockDim512MaxNumTokens = 1024;
+    if (data.mPtrScores == nullptr && data.mNumTokens <= PostTopKClusterBlockDim512MaxNumTokens)
+    {
+        launchClusterKernelBlockDim512(data, stream);
+        return;
+    }
+
     // Each warp owns one token, so the reduced-thread cluster variants have lower token capacity.
     // Use them only where the requested token count fits; otherwise keep the original 1024-thread launch.
     if (data.mNumTokens <= MaxNumTokensClusterScores256)

cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingCustomPolicy.cuh

@@ -587,6 +587,8 @@ struct PolicyTraits<NoOpPreprocess, SoftmaxPostprocess>
             Tier<512, 16>,       // Various 512-expert models with high topK
             Tier<512, 22>,       // Nemotron Super V3 (512 experts, topK=22)
             Tier<576, 8>,        // Customized model with 576 experts
+            Tier<1024, 32>,      // Kimi K3 post-topK permutation (896 experts, topK 16): without this tier the
+                                 // >128-token cluster launch and the dyn-block dispatch fell to <2048,32>
             Tier<2048, 32>       // Large-expert fallback
             >;
 };

cpp/tensorrt_llm/kernels/moe/trtllmGen/routing/RoutingFromTopKIds.cu

@@ -74,8 +74,18 @@ void runPostTopKPipeline(DataType const& data, uint32_t /*numThreadsHist*/, void
     static int const smMajor = tensorrt_llm::common::getSMVersion() / 10;
     bool const useStaticBlock = data.mNumTokens <= routingCustom::BlockKernelMaxNumTokens;
     int32_t const dispatchedMaxExperts = routingCustom::queryDispatchedMaxExperts(customData);
-    bool const useDynBlock = !useStaticBlock && data.mNumTokens <= routingCustom::DynBlockKernelMaxNumTokens
-        && dispatchedMaxExperts <= routingCustom::DynBlockKernelMaxNumExperts;
+    // Post-topK only permutes (no in-kernel TopK), so the single-CTA dyn-block kernel is worth using
+    // for larger expert tiers than in the scores path: the 8-CTA cluster kernel pays a DSMEM gather +
+    // cluster-barrier latency floor regardless of token count. Measured on B200 (896 experts/top-16,
+    // E=1024 tier): dyn-block 3.1/3.3/4.2 us at 5/8/16 tokens vs the 256-thread cluster kernel
+    // 4.7/4.7/5.2 us; but the dyn-block per-expert loops run over all tokens (~0.11 us/token), while
+    // the post-topK 512-thread cluster variant (see launchClusterKernel) is ~4.3 us at 17 tokens and
+    // nearly flat -> crossover ~12 tokens (dyn-block 4.7 us vs cluster 4.3 us already at 16-17).
+    // smem = tokens * MaxNumExperts * 3 B (int8 k-idx + int16 offset) = 36 KB at 12 x 1024.
+    static constexpr int PostTopKDynBlockMaxNumTokens = 12;
+    static constexpr int PostTopKDynBlockMaxNumExperts = 1024;
+    bool const useDynBlock = !useStaticBlock && data.mNumTokens <= PostTopKDynBlockMaxNumTokens
+        && dispatchedMaxExperts <= PostTopKDynBlockMaxNumExperts;
 
     // runPostTopKPipeline only handles pre-computed topK (mPtrTopKIds or mPtrTopKPacked),
     // never raw scores. The cluster kernel's routingPermutation uses thread-per-expanded-index