From 1775ca96b2a84621faa5b3816ed0d413bdee33f6 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Thu, 30 Jul 2026 01:11:39 -0700 Subject: [PATCH 1/5] [None][perf] fuse Kimi route and MXFP8 quantization Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/noAuxTcKernels.cu | 113 +++++++++++++++--- cpp/tensorrt_llm/kernels/noAuxTcKernels.h | 4 + cpp/tensorrt_llm/thop/noAuxTcOp.cpp | 51 +++++++- .../_torch/custom_ops/cpp_custom_ops.py | 11 ++ .../modules/fused_moe/fused_moe_trtllm_gen.py | 46 +++++++ .../_torch/modules/fused_moe/moe_scheduler.py | 33 ++++- .../modules/moe/test_kimi_k3_moe_gate.py | 32 +++++ 7 files changed, 268 insertions(+), 22 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu b/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu index 8256c4e6ca73..986e030c71c0 100644 --- a/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu +++ b/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 #include #include @@ -47,14 +48,10 @@ static __device__ inline float sigmoid_accurate(float x) template -__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(numExperts) + threadExpert; auto biasVal = expertSelected ? static_cast(routingBias[threadExpert]) : invalidScoreFloat; float score = expertSelected ? static_cast(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(numExperts) + e; auto biasVal = static_cast(routingBias[e]); float score = static_cast(scores[scoreIdx]); auto scoreSigmoid = sigmoid_accurate(score); @@ -188,12 +180,101 @@ __global__ void deepseek_v3_topk_kernel(InputT* scores, OutputT* topkValues, Idx } } } +} +template +__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( + 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(scores, topkValues, topkIndices, routingBias, numTokens, 1, 1, KimiK3TopK, + KimiK3NumExperts, KimiK3NumExperts, routedScalingFactor, blockIdx.x); + } + else + { + using QuantT = __nv_bfloat16; + using QuantPackedVec = PackedVec; + 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 optionalNumRows = numTokens; + auto sfOut = cvt_quant_get_sf_out_offset(std::nullopt, rowIdx, colIdx, + optionalNumRows, KimiK3HiddenSize / MxFp8SfVecSize, reinterpret_cast(hiddenStatesScale), + QuantizationSFLayout::LINEAR); + int64_t const offset = static_cast(rowIdx) * numColThreads + colIdx; + QuantPackedVec inVec = reinterpret_cast(hiddenStates)[offset]; + reinterpret_cast(quantizedHiddenStates)[offset] + = cvt_warp_fp16_to_mxfp8(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 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, diff --git a/cpp/tensorrt_llm/kernels/noAuxTcKernels.h b/cpp/tensorrt_llm/kernels/noAuxTcKernels.h index dfe6908723e3..2d093126fe87 100644 --- a/cpp/tensorrt_llm/kernels/noAuxTcKernels.h +++ b/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 diff --git a/cpp/tensorrt_llm/thop/noAuxTcOp.cpp b/cpp/tensorrt_llm/thop/noAuxTcOp.cpp index 4dfb20072734..1f381a327bfd 100644 --- a/cpp/tensorrt_llm/thop/noAuxTcOp.cpp +++ b/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 // #include @@ -156,6 +158,49 @@ std::tuple noaux_tc_op(th::Tensor const& scores, th::Ten return {topk_values, topk_indices}; } +std::tuple 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; + + TORCH_CHECK(tl::common::getSMVersion() == 100, "kimi_k3_noaux_tc_mxfp8_quant requires SM100"); + 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(scores.mutable_data_ptr()), + reinterpret_cast(bias.mutable_data_ptr()), + reinterpret_cast<__nv_bfloat16*>(hiddenStates.mutable_data_ptr()), + reinterpret_cast<__nv_bfloat16*>(topkValues.mutable_data_ptr()), + reinterpret_cast(topkIndices.mutable_data_ptr()), + reinterpret_cast(quantizedHiddenStates.mutable_data_ptr()), + reinterpret_cast(hiddenStatesScale.mutable_data_ptr()), numTokens, routedScalingFactor, stream); + return {topkIndices, topkValues, quantizedHiddenStates, hiddenStatesScale}; +} + } // end namespace torch_ext TRTLLM_NAMESPACE_END @@ -165,9 +210,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); } diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 1dda29ec7d8b..03c8619f3fcf 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 79c2f064e1ac..5abe22127745 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -648,6 +648,52 @@ 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) + or get_sm_version() != 100 + 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. diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index f86ecf887aee..55ae6644910e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -380,11 +380,31 @@ def _forward_chunk_impl( or moe.comm is not None or FORCE_SEPARATED_ROUTING ) + used_fused_route_quant = False if requires_separated_routing: - # Separated routing: ConfigurableMoE calls routing_method - token_selected_experts, token_final_scales = moe.routing_method.apply( - router_logits, input_ids + can_quantize_before_dispatch = ( + moe.comm is None or moe.comm.supports_post_quant_dispatch() ) + if ( + can_quantize_before_dispatch + and isinstance(moe.backend, TRTLLMGenFusedMoE) + and not moe._using_load_balancer() + and not moe.apply_router_weight_on_input + ): + fused_result = moe.backend.try_fused_kimi_route_quant(x, router_logits) + else: + fused_result = None + + if fused_result is None: + # Separated routing: ConfigurableMoE calls routing_method. + token_selected_experts, token_final_scales = moe.routing_method.apply( + router_logits, input_ids + ) + if token_final_scales is not None and isinstance(moe.backend, TRTLLMGenFusedMoE): + token_final_scales = token_final_scales.to(torch.bfloat16) + else: + token_selected_experts, token_final_scales, x, x_sf = fused_result + used_fused_route_quant = True token_selected_experts = token_selected_experts.to(torch.int32) @@ -504,7 +524,8 @@ def _forward_chunk_impl( if supports_post_quant: # Quantize -> Dispatch - x, x_sf = moe.backend.quantize_input(x) + if not used_fused_route_quant: + x, x_sf = moe.backend.quantize_input(x) # W4AFP8 + DeepEPLowLatency needs pre_quant_scale_1; other strategies # absorb the kwarg via **kwargs so unconditional passing is safe. @@ -534,10 +555,12 @@ def _forward_chunk_impl( use_dp_padding=use_dp_padding, **dispatch_kwargs, ) + assert not used_fused_route_quant x, x_sf = moe.backend.quantize_input(x, post_quant_comm=False) else: # No comm: just quantize - x, x_sf = moe.backend.quantize_input(x, post_quant_comm=False) + if not used_fused_route_quant: + x, x_sf = moe.backend.quantize_input(x, post_quant_comm=False) # ========== Step 6: MoE computation ========== # If EPLB is enabled, token_selected_slots is slot ids; otherwise expert ids. diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py index 7027a118cab3..74da41f824e2 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py @@ -125,3 +125,35 @@ def test_ineligible_configs_disable_fused_routing(cfg): # softmax scoring, grouped routing, renormalize off, and top_k == 1 all # diverge from the fused kernel's fixed contract -> eager path. assert KimiK3MoEGate(cfg)._use_fused_routing is False + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="Kimi fused route+MXFP8 quant requires SM100", +) +@pytest.mark.parametrize("num_tokens", [1, 5, 64]) +def test_fused_route_quant_matches_unfused_chain(num_tokens): + torch.manual_seed(0x5EED + num_tokens) + scores = torch.randn(num_tokens, 896, dtype=torch.float32, device="cuda") + bias = torch.randn(896, dtype=torch.float32, device="cuda") + hidden_states = torch.randn(num_tokens, 3584, dtype=torch.bfloat16, device="cuda") + routed_scaling_factor = 2.446 + + ref_scales, ref_experts = torch.ops.trtllm.noaux_tc_op( + scores, bias, 1, 1, 16, routed_scaling_factor + ) + ref_quantized, ref_quant_scales = torch.ops.trtllm.mxfp8_quantize( + hidden_states, False, alignment=256 + ) + + experts, scales, quantized, quant_scales = torch.ops.trtllm.kimi_k3_noaux_tc_mxfp8_quant( + scores, + bias, + hidden_states, + routed_scaling_factor, + ) + + assert torch.equal(experts, ref_experts) + assert torch.equal(scales.view(torch.int16), ref_scales.to(torch.bfloat16).view(torch.int16)) + assert torch.equal(quantized.view(torch.uint8), ref_quantized.view(torch.uint8)) + assert torch.equal(quant_scales, ref_quant_scales.view(num_tokens, -1)) From 4e906ef13f7cd5f696ec7d3f99b9c5ce00559611 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Thu, 30 Jul 2026 01:24:19 -0700 Subject: [PATCH 2/5] [None][perf] enable Kimi fused route quant on SM10x Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- cpp/tensorrt_llm/thop/noAuxTcOp.cpp | 3 ++- tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py | 3 ++- tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py | 2 +- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/thop/noAuxTcOp.cpp b/cpp/tensorrt_llm/thop/noAuxTcOp.cpp index 1f381a327bfd..373a59ee6765 100644 --- a/cpp/tensorrt_llm/thop/noAuxTcOp.cpp +++ b/cpp/tensorrt_llm/thop/noAuxTcOp.cpp @@ -167,7 +167,8 @@ std::tuple kimi_k3_noaux_tc_mxfp constexpr int64_t maxNumTokens = 64; constexpr int64_t sfVecSize = 32; - TORCH_CHECK(tl::common::getSMVersion() == 100, "kimi_k3_noaux_tc_mxfp8_quant requires SM100"); + 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"); diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 5abe22127745..c3024682e4a8 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -659,10 +659,11 @@ def try_fused_kimi_route_quant( K3 decode shape. Returning ``None`` keeps every other model, shape, architecture, and op backend on the existing unfused path. """ + sm_version = get_sm_version() if ( os.environ.get("TLLM_K3_DISABLE_FUSED_ROUTE_QUANT", "0") == "1" or isinstance(x, MxFp8QuantizedTensor) - or get_sm_version() != 100 + or 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) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py index 74da41f824e2..06c917c5d113 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py @@ -129,7 +129,7 @@ def test_ineligible_configs_disable_fused_routing(cfg): @pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, - reason="Kimi fused route+MXFP8 quant requires SM100", + reason="Kimi fused route+MXFP8 quant requires an SM10x architecture", ) @pytest.mark.parametrize("num_tokens", [1, 5, 64]) def test_fused_route_quant_matches_unfused_chain(num_tokens): From 953f131bc9fd9b224622221beef36299afff79ea Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Sun, 2 Aug 2026 20:38:36 -0700 Subject: [PATCH 3/5] [None][test] cover prequantized Kimi route handoff Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../modules/fused_moe/fused_moe_trtllm_gen.py | 46 ++++++++----------- .../_torch/modules/moe/test_moe_backend.py | 26 ++++++++++- 2 files changed, 43 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index c3024682e4a8..7ebe8fd2eb51 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -652,7 +652,8 @@ 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]]: + ) -> 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 @@ -660,40 +661,29 @@ def try_fused_kimi_route_quant( architecture, and op backend on the existing unfused path. """ sm_version = get_sm_version() - if ( - os.environ.get("TLLM_K3_DISABLE_FUSED_ROUTE_QUANT", "0") == "1" - or isinstance(x, MxFp8QuantizedTensor) - or 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) - ): + if (os.environ.get("TLLM_K3_DISABLE_FUSED_ROUTE_QUANT", "0") == "1" + or isinstance(x, MxFp8QuantizedTensor) + or 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() - ): + 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 - ) + 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. diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 4b8764c548a6..555ff5979a67 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -61,12 +61,15 @@ from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend, get_moe_cls from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_marlin import MarlinFusedMoE +from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoECommPlan, MoERunContext from tensorrt_llm._torch.modules.fused_moe.interface import ( MoE, MoESchedulerKind, MoEWeightLoadingMode, ) +from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE +from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoEWeightLoadingMode from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm from tensorrt_llm._torch.modules.fused_moe.quantization import ( FusedMoEMethodBase, @@ -76,8 +79,10 @@ W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, W4A16NVFP4CutlassFusedMoEMethod, ) -from tensorrt_llm._torch.utils import ActivationType, is_gated_activation +from tensorrt_llm._torch.utils import ActivationType, MxFp8QuantizedTensor, is_gated_activation from tensorrt_llm._utils import get_sm_version, mpi_rank +from tensorrt_llm._torch.utils import ActivationType, MxFp8QuantizedTensor, is_gated_activation +from tensorrt_llm._utils import mpi_rank from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -158,6 +163,25 @@ def should_skip_gptoss( return None +def test_kimi_fused_route_quant_skips_prequantized_input(monkeypatch) -> None: + """An upstream fused down projection owns quantization on this path.""" + monkeypatch.delenv("TLLM_K3_DISABLE_FUSED_ROUTE_QUANT", raising=False) + monkeypatch.setattr( + "tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen.get_sm_version", + MagicMock(side_effect=AssertionError("SM probe must be short-circuited")), + ) + backend = TRTLLMGenFusedMoE.__new__(TRTLLMGenFusedMoE) + hidden_states = MxFp8QuantizedTensor( + fp8_tensor=torch.empty(1, 3584, dtype=torch.float8_e4m3fn), + scaling_factor=torch.empty(1, 112, dtype=torch.uint8), + ) + + assert ( + backend.try_fused_kimi_route_quant(hidden_states, torch.empty(1, 896, dtype=torch.float32)) + is None + ) + + def create_test_backend( backend_type: MoeBackendType, routing_method: RenormalizeMoeRoutingMethod, From 34c37e442368234c06683f7aeee9cd517c3eeed9 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Thu, 13 Aug 2026 18:59:36 -0700 Subject: [PATCH 4/5] [None][fix] address Kimi route quant review feedback Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../_torch/modules/fused_moe/fused_moe_trtllm_gen.py | 12 +++++++----- .../_torch/modules/fused_moe/moe_scheduler.py | 10 +++------- tests/integration/test_lists/test-db/l0_b200.yml | 1 + 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 7ebe8fd2eb51..116821e6cce6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/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 @@ -660,10 +660,12 @@ def try_fused_kimi_route_quant( K3 decode shape. Returning ``None`` keeps every other model, shape, architecture, and op backend on the existing unfused path. """ - sm_version = get_sm_version() if (os.environ.get("TLLM_K3_DISABLE_FUSED_ROUTE_QUANT", "0") == "1" - or isinstance(x, MxFp8QuantizedTensor) - or not 100 <= sm_version < 110 or not self.has_w4a8_mxfp4_mxfp8 + 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)): @@ -714,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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index 55ae6644910e..82e207aa8167 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -54,6 +54,7 @@ from .communication import DeepEP, DeepEPLowLatency, NcclEP, NVLinkOneSided, NVLinkTwoSided from .communication.nvlink_two_sided_flashinfer import NVLinkTwoSidedFlashinfer from .fused_moe_cutlass import raise_moe_lora_multichunk_unsupported +from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .impl_contract import MoECommPlan, MoERunContext from .interface import FORCE_SEPARATED_ROUTING, MoESchedulerKind @@ -380,13 +381,11 @@ def _forward_chunk_impl( or moe.comm is not None or FORCE_SEPARATED_ROUTING ) + supports_post_quant = moe.comm is None or moe.comm.supports_post_quant_dispatch() used_fused_route_quant = False if requires_separated_routing: - can_quantize_before_dispatch = ( - moe.comm is None or moe.comm.supports_post_quant_dispatch() - ) if ( - can_quantize_before_dispatch + supports_post_quant and isinstance(moe.backend, TRTLLMGenFusedMoE) and not moe._using_load_balancer() and not moe.apply_router_weight_on_input @@ -510,8 +509,6 @@ def _forward_chunk_impl( # ========== Step 5: Quantization + dispatch (pre/post-quant adaptive ordering) ========== if moe.comm is not None: - supports_post_quant = moe.comm.supports_post_quant_dispatch() - # Debug: optional dummy AllReduce to break load-balancing artifacts if moe.enable_dummy_allreduce: moe.dummy_allreduce() @@ -555,7 +552,6 @@ def _forward_chunk_impl( use_dp_padding=use_dp_padding, **dispatch_kwargs, ) - assert not used_fused_route_quant x, x_sf = moe.backend.quantize_input(x, post_quant_comm=False) else: # No comm: just quantize diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 2167cb209837..5f4309969606 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -127,6 +127,7 @@ l0_b200: - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py - unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- MoE: test_moe_backend (by backend) --------------- + - unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py::test_fused_route_quant_matches_unfused_chain - unittest/_torch/modules/moe/test_megamoe_streaming_load.py - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts From c1bf03fd51937d4120375e83b230e10071448855 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Thu, 13 Aug 2026 22:59:41 -0700 Subject: [PATCH 5/5] [None][fix] add Kimi MXFP8 handoff carrier Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../modules/fused_moe/fused_moe_trtllm_gen.py | 2 +- tensorrt_llm/_torch/utils.py | 31 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 116821e6cce6..c9313a759934 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -716,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, MxFp8QuantizedTensor): + 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 diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index cb62ec99f76b..218c8fd14126 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -195,6 +195,37 @@ def shape(self): return self.fp4_tensor.shape +@dataclass +class MxFp8QuantizedTensor: + """MXFP8 activation and per-1x32 UE8M0 scaling factors.""" + + fp8_tensor: torch.Tensor + scaling_factor: torch.Tensor + is_sf_swizzled: bool = False + + @property + def shape(self): + return self.fp8_tensor.shape + + @property + def dtype(self): + return self.fp8_tensor.dtype + + def numel(self): + return self.fp8_tensor.numel() + + def split(self, split_size_or_sections, dim=0): + if dim != 0: + raise ValueError( + "MxFp8QuantizedTensor can only be split along the token dimension" + ) + fp8_chunks = self.fp8_tensor.split(split_size_or_sections, dim=dim) + sf_chunks = self.scaling_factor.split(split_size_or_sections, dim=dim) + return tuple( + MxFp8QuantizedTensor(fp8_chunk, sf_chunk, self.is_sf_swizzled) + for fp8_chunk, sf_chunk in zip(fp8_chunks, sf_chunks)) + + def compute_swizzled_sf_shape(row: int, col: int): padded_row = pad_up(row, 128) padded_col = pad_up(col, 4)