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..373a59ee6765 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,50 @@ 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; + + 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(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 +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); } 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..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 @@ -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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index f86ecf887aee..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,11 +381,29 @@ 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: - # Separated routing: ConfigurableMoE calls routing_method - token_selected_experts, token_final_scales = moe.routing_method.apply( - router_logits, input_ids - ) + if ( + supports_post_quant + 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) @@ -490,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() @@ -504,7 +521,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. @@ -537,7 +555,8 @@ def _forward_chunk_impl( 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/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) 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 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..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 @@ -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 an SM10x architecture", +) +@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)) 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,