[None][perf] fuse Kimi K3 routing and MXFP8 quantization

cpp/tensorrt_llm/kernels/noAuxTcKernels.cu

@@ -20,6 +20,7 @@
 #include "tensorrt_llm/common/cudaTypeUtils.cuh"
 #include "tensorrt_llm/common/envUtils.h"
 #include "tensorrt_llm/kernels/noAuxTcKernels.h"
+#include "tensorrt_llm/kernels/quantization.cuh"
 #include <cmath>
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
@@ -47,14 +48,10 @@ static __device__ inline float sigmoid_accurate(float x)
 
 template <typename InputT, typename BiasT, typename OutputT, typename IdxT, int MaxNumExperts, bool UseGroups,
     int MaxNumTopExperts = DefaultMaxNumTopExperts, int MaxNumTopGroups = DefaultMaxNumTopGroups>
-__global__ void deepseek_v3_topk_kernel(InputT* scores, OutputT* topkValues, IdxT* topkIndices, BiasT* routingBias,
-    int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup, int64_t const topk,
-    int64_t const numExperts, int64_t const numExpertsPerGroup, double const routedScalingFactor)
+__device__ __forceinline__ void deepseek_v3_topk_block(InputT* scores, OutputT* topkValues, IdxT* topkIndices,
+    BiasT* routingBias, int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup, int64_t const topk,
+    int64_t const numExperts, int64_t const numExpertsPerGroup, double const routedScalingFactor, int64_t tokenIdx)
 {
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    cudaGridDependencySynchronize();
-#endif
-
     __shared__ float __attribute((aligned(128))) smemScoreSigmoid[MaxNumExperts];
     __shared__ float __attribute((aligned(128))) smemScoreBias[MaxNumExperts];
 
@@ -66,23 +63,18 @@ __global__ void deepseek_v3_topk_kernel(InputT* scores, OutputT* topkValues, Idx
 
     static constexpr float invalidScoreFloat = float{-INFINITY};
 
-    topkValues += blockIdx.x * topk;
-    topkIndices += blockIdx.x * topk;
+    topkValues += tokenIdx * topk;
+    topkIndices += tokenIdx * topk;
 
     if constexpr (UseGroups)
     {
         int constexpr NumWarps = MaxNumExperts / WARP_SIZE;
         __shared__ float __attribute((aligned(128))) smemGroupScores[NumWarps];
 
-        if (warpIdx >= numGroup)
-        {
-            return;
-        }
-
         auto threadExpert = warpIdx * numExpertsPerGroup + laneIdx;
         bool expertSelected = laneIdx < numExpertsPerGroup;
 
-        auto scoreIdx = int64_t{blockIdx.x} * int64_t{numExperts} + threadExpert;
+        auto scoreIdx = tokenIdx * static_cast<int64_t>(numExperts) + threadExpert;
         auto biasVal = expertSelected ? static_cast<float>(routingBias[threadExpert]) : invalidScoreFloat;
         float score = expertSelected ? static_cast<float>(scores[scoreIdx]) : invalidScoreFloat;
         auto scoreSigmoid = sigmoid_accurate(score);
@@ -149,7 +141,7 @@ __global__ void deepseek_v3_topk_kernel(InputT* scores, OutputT* topkValues, Idx
     {
         for (int e = threadIdx.x; e < numExperts; e += blockDim.x)
         {
-            auto scoreIdx = int64_t{blockIdx.x} * int64_t{numExperts} + e;
+            auto scoreIdx = tokenIdx * static_cast<int64_t>(numExperts) + e;
             auto biasVal = static_cast<float>(routingBias[e]);
             float score = static_cast<float>(scores[scoreIdx]);
             auto scoreSigmoid = sigmoid_accurate(score);
@@ -188,12 +180,101 @@ __global__ void deepseek_v3_topk_kernel(InputT* scores, OutputT* topkValues, Idx
             }
         }
     }
+}
 
+template <typename InputT, typename BiasT, typename OutputT, typename IdxT, int MaxNumExperts, bool UseGroups,
+    int MaxNumTopExperts = DefaultMaxNumTopExperts, int MaxNumTopGroups = DefaultMaxNumTopGroups>
+__global__ void deepseek_v3_topk_kernel(InputT* scores, OutputT* topkValues, IdxT* topkIndices, BiasT* routingBias,
+    int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup, int64_t const topk,
+    int64_t const numExperts, int64_t const numExpertsPerGroup, double const routedScalingFactor)
+{
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    cudaGridDependencySynchronize();
+#endif
+    if constexpr (UseGroups)
+    {
+        if (threadIdx.x / WARP_SIZE >= numGroup)
+        {
+            return;
+        }
+    }
+    deepseek_v3_topk_block<InputT, BiasT, OutputT, IdxT, MaxNumExperts, UseGroups, MaxNumTopExperts, MaxNumTopGroups>(
+        scores, topkValues, topkIndices, routingBias, numTokens, numGroup, topkGroup, topk, numExperts,
+        numExpertsPerGroup, routedScalingFactor, blockIdx.x);
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
     cudaTriggerProgrammaticLaunchCompletion();
 #endif
 }
 
+// K3 decode specialization: CTAs [0, M) route 896 experts to top-16, while
+// CTAs [M, 2M) independently quantize one 3584-wide BF16 activation row to
+// MXFP8 with linear UE8M0 group-32 scales.
+static constexpr int KimiK3NumExperts = 896;
+static constexpr int KimiK3TopK = 16;
+static constexpr int KimiK3HiddenSize = 3584;
+static constexpr int MxFp8SfVecSize = 32;
+static constexpr int KimiK3QuantThreads = KimiK3HiddenSize / CVT_ELTS_PER_THREAD;
+
+__global__ __launch_bounds__(KimiK3QuantThreads) void kimi_k3_noaux_tc_mxfp8_quant_kernel(float* scores,
+    float* routingBias, __nv_bfloat16* hiddenStates, __nv_bfloat16* topkValues, int32_t* topkIndices,
+    int64_t* quantizedHiddenStates, int32_t* hiddenStatesScale, int64_t numTokens, double routedScalingFactor)
+{
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    cudaGridDependencySynchronize();
+#endif
+
+    if (blockIdx.x < numTokens)
+    {
+        deepseek_v3_topk_block<float, float, __nv_bfloat16, int32_t, MaxSupportedExpertCount, false,
+            MaxSupportedTopExperts>(scores, topkValues, topkIndices, routingBias, numTokens, 1, 1, KimiK3TopK,
+            KimiK3NumExperts, KimiK3NumExperts, routedScalingFactor, blockIdx.x);
+    }
+    else
+    {
+        using QuantT = __nv_bfloat16;
+        using QuantPackedVec = PackedVec<QuantT>;
+        static constexpr int CvtNumThreadsPerSf = MxFp8SfVecSize / CVT_ELTS_PER_THREAD;
+        int const rowIdx = blockIdx.x - numTokens;
+        int const colIdx = threadIdx.x;
+        int const numColThreads = KimiK3HiddenSize / CVT_ELTS_PER_THREAD;
+
+        std::optional<int> optionalNumRows = numTokens;
+        auto sfOut = cvt_quant_get_sf_out_offset<uint32_t, CvtNumThreadsPerSf>(std::nullopt, rowIdx, colIdx,
+            optionalNumRows, KimiK3HiddenSize / MxFp8SfVecSize, reinterpret_cast<uint32_t*>(hiddenStatesScale),
+            QuantizationSFLayout::LINEAR);
+        int64_t const offset = static_cast<int64_t>(rowIdx) * numColThreads + colIdx;
+        QuantPackedVec inVec = reinterpret_cast<QuantPackedVec const*>(hiddenStates)[offset];
+        reinterpret_cast<uint64_t*>(quantizedHiddenStates)[offset]
+            = cvt_warp_fp16_to_mxfp8<QuantT, MxFp8SfVecSize>(inVec, sfOut);
+    }
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    __threadfence();
+    __syncthreads();
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
+void invokeKimiK3NoAuxTcMxFp8Quant(float* scores, float* bias, __nv_bfloat16* hidden_states, __nv_bfloat16* topk_values,
+    int32_t* topk_indices, int64_t* quantized_hidden_states, int32_t* hidden_states_scale, int64_t const num_tokens,
+    double const routed_scaling_factor, cudaStream_t const stream)
+{
+    cudaLaunchConfig_t config;
+    config.gridDim = 2 * num_tokens;
+    config.blockDim = KimiK3QuantThreads;
+    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;
+
+    cudaLaunchKernelEx(&config, kimi_k3_noaux_tc_mxfp8_quant_kernel, scores, bias, hidden_states, topk_values,
+        topk_indices, quantized_hidden_states, hidden_states_scale, num_tokens, routed_scaling_factor);
+    sync_check_cuda_error(stream);
+}
+
 template <typename InputT, typename BiasT, typename OutputT, typename IdxT>
 void invokeNoAuxTc(InputT* scores, BiasT* bias, OutputT* topk_values, IdxT* topk_indices, int64_t const num_tokens,
     int64_t const num_experts, int64_t const n_group, int64_t const topk_group, int64_t const topk,

cpp/tensorrt_llm/kernels/noAuxTcKernels.h

@@ -33,6 +33,10 @@ void invokeNoAuxTc(InputT* scores, BiasT* bias, OutputT* topk_values, IdxT* topk
     int64_t const num_experts, int64_t const n_group, int64_t const topk_group, int64_t const topk,
     double const routed_scaling_factor, cudaStream_t const stream = 0);
 
+void invokeKimiK3NoAuxTcMxFp8Quant(float* scores, float* bias, __nv_bfloat16* hidden_states, __nv_bfloat16* topk_values,
+    int32_t* topk_indices, int64_t* quantized_hidden_states, int32_t* hidden_states_scale, int64_t const num_tokens,
+    double const routed_scaling_factor, cudaStream_t const stream = 0);
+
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/thop/noAuxTcOp.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,10 +15,12 @@
  * limitations under the License.
  */
 
+#include "tensorrt_llm/common/cudaUtils.h"
 #include "tensorrt_llm/common/opUtils.h"
 #include "tensorrt_llm/runtime/torchUtils.h"
 
 #include "tensorrt_llm/kernels/noAuxTcKernels.h"
+#include "tensorrt_llm/thop/thUtils.h"
 
 // #include <c10/cuda/CUDAStream.h>
 // #include <cassert>
@@ -156,6 +158,50 @@ std::tuple<at::Tensor, at::Tensor> noaux_tc_op(th::Tensor const& scores, th::Ten
     return {topk_values, topk_indices};
 }
 
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> kimi_k3_noaux_tc_mxfp8_quant(
+    th::Tensor const& scores, th::Tensor const& bias, th::Tensor const& hiddenStates, double routedScalingFactor)
+{
+    constexpr int64_t numExperts = 896;
+    constexpr int64_t topK = 16;
+    constexpr int64_t hiddenSize = 3584;
+    constexpr int64_t maxNumTokens = 64;
+    constexpr int64_t sfVecSize = 32;
+
+    int const smVersion = tl::common::getSMVersion();
+    TORCH_CHECK(smVersion >= 100 && smVersion < 110, "kimi_k3_noaux_tc_mxfp8_quant requires an SM10x architecture");
+    TORCH_CHECK(scores.is_cuda() && bias.is_cuda() && hiddenStates.is_cuda(), "all inputs must be CUDA tensors");
+    TORCH_CHECK(scores.get_device() == bias.get_device() && scores.get_device() == hiddenStates.get_device(),
+        "all inputs must be on the same device");
+    TORCH_CHECK(scores.scalar_type() == torch::kFloat32 && bias.scalar_type() == torch::kFloat32,
+        "scores and bias must be float32");
+    TORCH_CHECK(hiddenStates.scalar_type() == torch::kBFloat16, "hidden_states must be bfloat16");
+    TORCH_CHECK(scores.is_contiguous() && bias.is_contiguous() && hiddenStates.is_contiguous(),
+        "all inputs must be contiguous");
+    TORCH_CHECK(scores.dim() == 2 && scores.size(1) == numExperts, "scores must have shape [M, 896]");
+    TORCH_CHECK(bias.dim() == 1 && bias.numel() == numExperts, "bias must have shape [896]");
+    TORCH_CHECK(hiddenStates.dim() == 2 && hiddenStates.size(0) == scores.size(0) && hiddenStates.size(1) == hiddenSize,
+        "hidden_states must have shape [M, 3584] with the same M as scores");
+    int64_t const numTokens = scores.size(0);
+    TORCH_CHECK(numTokens > 0 && numTokens <= maxNumTokens, "M must be in [1, 64]");
+
+    auto const device = th::Device(th::kCUDA, scores.get_device());
+    th::Tensor topkValues = th::empty({numTokens, topK}, th::dtype(torch::kBFloat16).device(device));
+    th::Tensor topkIndices = th::empty({numTokens, topK}, th::dtype(torch::kInt32).device(device));
+    th::Tensor quantizedHiddenStates
+        = th::empty({numTokens, hiddenSize}, th::dtype(torch::kFloat8_e4m3fn).device(device));
+    th::Tensor hiddenStatesScale = th::empty({numTokens, hiddenSize / sfVecSize}, th::dtype(SF_DTYPE).device(device));
+
+    auto stream = at::cuda::getCurrentCUDAStream(scores.get_device());
+    tk::invokeKimiK3NoAuxTcMxFp8Quant(reinterpret_cast<float*>(scores.mutable_data_ptr()),
+        reinterpret_cast<float*>(bias.mutable_data_ptr()),
+        reinterpret_cast<__nv_bfloat16*>(hiddenStates.mutable_data_ptr()),
+        reinterpret_cast<__nv_bfloat16*>(topkValues.mutable_data_ptr()),
+        reinterpret_cast<int32_t*>(topkIndices.mutable_data_ptr()),
+        reinterpret_cast<int64_t*>(quantizedHiddenStates.mutable_data_ptr()),
+        reinterpret_cast<int32_t*>(hiddenStatesScale.mutable_data_ptr()), numTokens, routedScalingFactor, stream);
+    return {topkIndices, topkValues, quantizedHiddenStates, hiddenStatesScale};
+}
+
 } // end namespace torch_ext
 
 TRTLLM_NAMESPACE_END
@@ -165,9 +211,13 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
     m.def(
         "noaux_tc_op(Tensor scores, Tensor bias, int n_group, int topk_group, int topk, float "
         "routed_scaling_factor) -> (Tensor, Tensor)");
+    m.def(
+        "kimi_k3_noaux_tc_mxfp8_quant(Tensor scores, Tensor bias, Tensor hidden_states, float "
+        "routed_scaling_factor) -> (Tensor, Tensor, Tensor, Tensor)");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
     m.impl("noaux_tc_op", &tensorrt_llm::torch_ext::noaux_tc_op);
+    m.impl("kimi_k3_noaux_tc_mxfp8_quant", &tensorrt_llm::torch_ext::kimi_k3_noaux_tc_mxfp8_quant);
 }

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -310,6 +310,17 @@ def _(scores, scores_with_bias, n_group, topk_group, topk,
                                 dtype=scores_with_bias.dtype), scores.new_empty(
                                     shape, dtype=torch.int32)
 
+    @torch.library.register_fake("trtllm::kimi_k3_noaux_tc_mxfp8_quant")
+    def _(scores, bias, hidden_states, routed_scaling_factor):
+        num_tokens = scores.shape[0]
+        return (
+            scores.new_empty((num_tokens, 16), dtype=torch.int32),
+            scores.new_empty((num_tokens, 16), dtype=torch.bfloat16),
+            hidden_states.new_empty((num_tokens, 3584),
+                                    dtype=torch.float8_e4m3fn),
+            hidden_states.new_empty((num_tokens, 112), dtype=torch.uint8),
+        )
+
     @torch.library.register_fake("trtllm::inplace_slice_copy")
     def _(dest, src, dim1_start, dim1_end):
         pass

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

@@ -29,7 +29,7 @@
     fp4_block_scale_fake_output_without_finalize
 from ...model_config import ModelConfig
 from ...utils import (ActivationType, ActType_TrtllmGen, AuxStreamType,
-                      Fp4QuantizedTensor)
+                      Fp4QuantizedTensor, MxFp8QuantizedTensor)
 from ..gated_mlp import GatedMLP
 from .impl_contract import MoEInputRequirement, MoERunContext, require_comm_plan
 from .interface import FORCE_SEPARATED_ROUTING, MoE, MoEWeightLoadingMode
@@ -648,6 +648,45 @@ def load_weights(self,
         self.quant_method.load_weights(self, weights, self.weight_loading_mode,
                                        **kargs)
 
+    def try_fused_kimi_route_quant(
+        self,
+        x: Union[torch.Tensor, MxFp8QuantizedTensor],
+        router_logits: torch.Tensor,
+    ) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor,
+                        torch.Tensor]]:
+        """Fuse Kimi K3 no-aux routing and MXFP8 input quantization.
+
+        This launch-overhead optimization is deliberately specialized to the
+        K3 decode shape. Returning ``None`` keeps every other model, shape,
+        architecture, and op backend on the existing unfused path.
+        """
+        if (os.environ.get("TLLM_K3_DISABLE_FUSED_ROUTE_QUANT", "0") == "1"
+                or isinstance(x, MxFp8QuantizedTensor)):
+            return None
+
+        sm_version = get_sm_version()
+        if (not 100 <= sm_version < 110 or not self.has_w4a8_mxfp4_mxfp8
+                or not isinstance(self.op_backend, TRTLLMOpBackend)
+                or not isinstance(self.routing_method,
+                                  DeepSeekV3MoeRoutingMethod)):
+            return None
+
+        routing = self.routing_method.routing_impl
+        bias = self.routing_method.e_score_correction_bias
+        if (not routing.is_fused or routing.n_group != 1
+                or routing.topk_group != 1 or routing.top_k != 16
+                or router_logits.ndim != 2 or router_logits.shape[1] != 896
+                or router_logits.dtype != torch.float32
+                or not router_logits.is_contiguous()
+                or bias.dtype != torch.float32 or not bias.is_contiguous()
+                or x.ndim != 2 or x.shape != (router_logits.shape[0], 3584)
+                or not 0 < x.shape[0] <= 64 or x.dtype != torch.bfloat16
+                or not x.is_contiguous()):
+            return None
+
+        return torch.ops.trtllm.kimi_k3_noaux_tc_mxfp8_quant(
+            router_logits, bias, x, routing.routed_scaling_factor)
+
     def quantize_input(self, x, post_quant_comm: bool = True):
         """Quantize inputs prior to post-communication (alltoall/allgather) or before MoE computation.
 
@@ -677,7 +716,7 @@ def quantize_input(self, x, post_quant_comm: bool = True):
                 x, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(
                     x, self.fc31_input_gate_dequant[0])
         elif self.has_nvfp4:
-            if isinstance(x, Fp4QuantizedTensor):
+            if isinstance(x, (Fp4QuantizedTensor, MxFp8QuantizedTensor)):
                 assert not x.is_sf_swizzled, "Fp4QuantizedTensor should not be swizzled before communication"
                 x_row = x.shape[0]
                 x, x_sf = x.fp4_tensor, x.scaling_factor