diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index be0c68577dda..286746bc92f5 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -18,7 +18,7 @@ This guide uses Slurm and the `trtllm-llmapi-launch` multi-node launcher. The co | TEP8 | attention-TP, EP8 | 213 GB | GB300-class per-GPU memory | | TEP16 | attention-TP, EP16 | 115 GB | validated on GB200 (`SM100`) | - DEP16 replicates the BF16 non-expert weights on every rank (114 GB per rank) on top of the MXFP4 routed experts at 16-way expert parallelism (90 GB per rank). TEP16 shards those non-expert weights instead, which is what brings it within `SM100` per-GPU memory; TEP8 does not fit because its 8-way expert share alone is 181 GB per rank. On B200 (`SM100`), Kimi K3 is functionally supported at the kernel and module level and covered by unit tests in CI. Other GPU architectures are not supported. + DEP16 replicates the BF16 non-expert weights on every rank (114 GB per rank) on top of the MXFP4 routed experts at 16-way expert parallelism (90 GB per rank). TEP16 shards those non-expert weights instead, which is what brings it within `SM100` per-GPU memory; TEP8 does not fit because its 8-way expert share alone is 181 GB per rank. On B200 (`SM100`), Kimi K3 is functionally supported at the kernel and module level and covered by unit tests in CI. Full-model deployments on other GPU architectures are not covered by this guide. * Multi-node launcher: Slurm with the pyxis/enroot container plugin (or an equivalent MPI launcher) to start one rank per GPU across the nodes. * High-speed inter-node interconnect (e.g., NVLink/InfiniBand) for the expert-parallel traffic. * Shared filesystem visible to all nodes for the repository, the model weights, and the configuration file. @@ -34,7 +34,8 @@ The checkpoint and the configuration file must live on a shared filesystem visib ## Feature Support Notes -* **Blackwell only.** NVIDIA Blackwell GPUs are supported. The performance results in this guide were validated on NVIDIA GB300 NVL GPUs. Kimi K3 kernels and modules are also functional on B200 (`SM100`) and covered by unit tests in CI, and the TEP16 deployment is validated end-to-end on GB200 (`SM100`); DEP16 and TEP8 require GB300-class per-GPU memory (see Prerequisites). Support for other GPU architectures may be added in a future release. +* **Full-model deployments in this guide use Blackwell.** The performance results were validated on NVIDIA GB300 NVL GPUs. Kimi K3 kernels and modules are also functional on B200 (`SM100`) and covered by unit tests in CI, and the TEP16 deployment is validated end-to-end on GB200 (`SM100`); DEP16 and TEP8 require GB300-class per-GPU memory (see Prerequisites). +* **Rubin FC12 MoE kernel support.** On SM107, `CUTEDSL_FC12` supports SwiGLU and Kimi K3 SiTU for NVFP4 routed experts. This kernel/backend support does not validate the full Kimi K3 deployment on Rubin; the configurations and results below are for Blackwell. * **High-throughput and low-latency deployments are provided.** DEP16 (`enable_attention_dp: true`, `moe_expert_parallel_size: 16`) is the high-throughput deployment. TEP16 (`enable_attention_dp: false`, `moe_expert_parallel_size: 16`) is the low-latency deployment. An 8-GPU deployment, TEP8 (`enable_attention_dp: false`, `moe_expert_parallel_size: 8`), is also provided. Select the deployment and concurrency appropriate for your workload. * **CUDA graphs and the overlap scheduler are enabled.** The performance-sweep recipes set `disable_overlap_scheduler: false` and enable CUDA graphs. DEP16 additionally sets `cuda_graph_config.enable_padding: true`. * **Chunked prefill is supported and enabled** (`enable_chunked_prefill: true`), so prompts longer than `max_num_tokens` are scheduled across multiple steps. diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 5d5b34b9da3d..7889ae631ceb 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -17943,7 +17943,7 @@ def _( # builds that do not ship cutlass.utils.rubin_helpers (SM107 only). if IS_CUTLASS_DSL_RUBIN_AVAILABLE: # ---------------------------------------------------------------- - # Rubin NVFP4 Fused FC12 (FC1 gather+SwiGLU + FC2 finalize in ONE kernel) + # Rubin NVFP4 Fused FC12 (FC1 gather+gated-act + FC2 finalize in ONE kernel) # ---------------------------------------------------------------- # Compat shim: the delivered fused kernel accesses ``cutlass.memory.*`` # (SmemAllocator / TmemAllocator / get_smem_capacity_in_bytes) and @@ -18036,16 +18036,20 @@ class Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner( kernel_cache = dict() tuning_config_cache = dict() - def __init__(self, - num_experts: int, - top_k: int, - num_local_experts: int, - local_expert_offset: int, - tile_size: int, - scaling_vector_size: int = 16, - swiglu_limit: float = float("inf"), - ep_size: int = 1, - enable_alltoall: bool = False): + def __init__( + self, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + swiglu_limit: float = float("inf"), + ep_size: int = 1, + enable_alltoall: bool = False, + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None): super().__init__() self.num_experts = num_experts self.top_k = top_k @@ -18054,6 +18058,9 @@ def __init__(self, self.tile_size = tile_size self.scaling_vector_size = scaling_vector_size self.swiglu_limit = swiglu_limit + self.activation_type = ActivationType(activation_type) + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta # Used only by the in-op output memset (moved here so the memset # is the fused kernel's immediate stream predecessor). self.ep_size = ep_size @@ -18078,6 +18085,9 @@ def unique_id(self): self.tile_size, self.scaling_vector_size, self.swiglu_limit, + int(self.activation_type), + self.situ_beta, + self.situ_linear_beta, ) def get_valid_tactics( @@ -18375,7 +18385,8 @@ def forward(self, inputs: List[torch.Tensor], cache_key = (self.scaling_vector_size, self.tile_size, self.top_k, mma_tiler, mma_inst_shape, cluster_shape_mn, max_active_clusters, - self.swiglu_limit) + self.swiglu_limit, int(self.activation_type), + self.situ_beta, self.situ_linear_beta) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( self.scaling_vector_size, @@ -18387,6 +18398,9 @@ def forward(self, inputs: List[torch.Tensor], use_pdl=True, swiglu_limit=self.swiglu_limit, scheduler="l2_atomic", + activation_type=self.activation_type, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) @cute.jit @@ -18622,6 +18636,9 @@ def _run_nvfp4_fc12_fused_rubin( ep_size: int, enable_alltoall: bool, tuner_key: str, + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ) -> torch.Tensor: tuner = AutoTuner.get() runner = Sm107BlockScaledContiguousGroupedGemmFusedFc12Runner( @@ -18634,6 +18651,9 @@ def _run_nvfp4_fc12_fused_rubin( swiglu_limit=swiglu_limit, ep_size=ep_size, enable_alltoall=enable_alltoall, + activation_type=activation_type, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, ) # Input order matches Fc12FusedInputsHelper (FC1 prefix 0..9 mirrors # GatherGroupedGemmInputsHelper; FC2 tensors 10..14; expanded_idx 15 @@ -18670,7 +18690,9 @@ def _run_nvfp4_fc12_fused_rubin( "SymInt local_expert_offset, SymInt tile_size, float swiglu_limit, " "SymInt ep_size, bool enable_alltoall, " "SymInt scaling_vector_size=16, " - "str? precomputed_tactic=None) -> ()", + "str? precomputed_tactic=None, " + f"SymInt activation_type={int(ActivationType.Swiglu)}, " + "float? situ_beta=None, float? situ_linear_beta=None) -> ()", device_types="cuda") def cute_dsl_nvfp4_fc12_fused_rubin( input: torch.Tensor, @@ -18699,6 +18721,9 @@ def cute_dsl_nvfp4_fc12_fused_rubin( enable_alltoall: bool, scaling_vector_size: int = 16, precomputed_tactic: Optional[str] = None, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ) -> None: # In-place: finalize scatter-adds into ``output`` (mutates_args); # the op returns nothing so it does not alias its own input. @@ -18711,7 +18736,8 @@ def cute_dsl_nvfp4_fc12_fused_rubin( local_expert_offset, tile_size, scaling_vector_size, swiglu_limit, precomputed_tactic, expanded_idx_to_permuted_idx, ep_size, enable_alltoall, - "trtllm::cute_dsl_nvfp4_fc12_fused_rubin") + "trtllm::cute_dsl_nvfp4_fc12_fused_rubin", + ActivationType(activation_type), situ_beta, situ_linear_beta) @torch.library.register_fake("trtllm::cute_dsl_nvfp4_fc12_fused_rubin") def _( @@ -18741,5 +18767,8 @@ def _( enable_alltoall: bool, scaling_vector_size: int = 16, precomputed_tactic: Optional[str] = None, + activation_type: int = int(ActivationType.Swiglu), + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ) -> None: return None diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py index 8c9d85e68f46..f06f092507c4 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py @@ -2857,6 +2857,9 @@ def _apply_situ_epilogue( beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) = 2*beta*sigmoid((2/beta)*x) - beta + + Keep the packed SiTU algebra in sync with the Rubin fused FC12 + ``rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py`` epilogue. """ beta = self.situ_beta linear_beta = self.situ_linear_beta diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py index 7681fbde79bc..1e3408801c0f 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_fused_fc12.py @@ -29,6 +29,8 @@ from dataclasses import dataclass from typing import NamedTuple, Optional, Tuple, Type, Union +from ....utils import ActivationType + import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute @@ -1076,6 +1078,22 @@ def silu_f32( return a * sigmoid_f32(a, fastmath=fastmath) +_FC12_ACTIVATION_TYPES = ( + ActivationType.Swiglu, + ActivationType.SiTu, +) + + +def validate_fc12_activation_type(activation_type) -> ActivationType: + """Fused FC12 is gated-only: SwiGLU or SiTU. Relu2 would keep full N.""" + activation_type = ActivationType(int(activation_type)) + if activation_type not in _FC12_ACTIVATION_TYPES: + raise ValueError( + f"Fused FC12 supports SwiGLU and SiTU only, got {activation_type}" + ) + return activation_type + + class S2TCopyBundle(NamedTuple): """Bundle of tiled copy and partitioned tensors for smem-to-tmem copies.""" @@ -1127,7 +1145,7 @@ def sm100_tcgen05_st_32x32b_x4( Compute: fc1_acc = fc1_alpha * (SFA * A[token_ids]) * (SFB * B) - C = up * silu(gate) # SwiGLU on interleaved acc + C = gated act on interleaved acc (SwiGLU or SiTU; both halve N) fc2_acc = fc2_alpha * (SFC * C) * (FC2_SFB * FC2_B) + optional NVFP4 quantization (generates SFC) when c_dtype == Float4E2M1FN. @@ -1169,8 +1187,9 @@ class Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel: - A/SFA load: four gather warps issue CpAsync128.CG into a shared fc1_a_pipeline. SFA is consumed directly by a Cp128x128b UTCCP in the MMA warp using the 128dp_Unique layout. - - SwiGLU epilogue: C = up * silu(gate), where up/gate come from the - gate16/up16 FC1 weight layout → output N is halved. + - Gated epilogue: SwiGLU (``up * silu(gate)``) or SiTU (Kimi K3); + gate/up are decoded from the gate16/up16 FC1 weight layout → + output N is halved. Selected by ``activation_type`` at trace time. - Optional NVFP4 quant: when c_dtype == Float4E2M1FN, the epilogue also generates SFC and quantizes the output. @@ -1189,6 +1208,11 @@ class Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel: :param vectorized_f32: Use vectorized f32x2 ops in epilogue. :param topk: Experts selected per token. :param swiglu_limit: DS-v4 SwiGLU clamp limit; ``+inf`` disables clamp. + Mutually exclusive with ``ActivationType.SiTu``. + :param activation_type: ``Swiglu`` (default) or ``SiTu``. Folded at trace + time; a different value compiles a different kernel. + :param situ_beta: Gate-side SiTU constant; required for ``SiTu`` only. + :param situ_linear_beta: Linear-side SiTU constant; required for ``SiTu`` only. :param scheduler: Persistent work-ID scheduler: ``static`` or ``l2_atomic``. FC1 B and SFB must already use gate16/up16 order (tile N 128/256). FC1 output remains in logical channel order; FC2 weight layout is unchanged. @@ -1205,6 +1229,9 @@ def __init__( use_pdl: bool = True, swiglu_limit: cutlass.Float32 = float("inf"), scheduler: str = "static", + activation_type: ActivationType = ActivationType.Swiglu, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ): self.sf_vec_size = sf_vec_size self.topk = topk @@ -1226,6 +1253,31 @@ def __init__( swiglu_limit = float("inf") self.swiglu_limit = swiglu_limit self.has_swiglu_limit = swiglu_limit != float("inf") + self.activation_type = validate_fc12_activation_type(activation_type) + if self.activation_type == ActivationType.SiTu: + if situ_beta is None or situ_linear_beta is None: + raise ValueError( + "ActivationType.SiTu requires both situ_beta and " + f"situ_linear_beta, got {situ_beta} and {situ_linear_beta}." + ) + if not (0 < situ_beta < float("inf") and 0 < situ_linear_beta < float("inf")): + raise ValueError( + f"SiTU betas must be finite and positive, got {situ_beta} and {situ_linear_beta}." + ) + if self.has_swiglu_limit: + raise ValueError( + "Fused FC12 SiTU does not support a SwiGLU clamp; " + "drop swiglu_limit for SiTU checkpoints." + ) + elif situ_beta is not None or situ_linear_beta is not None: + raise ValueError( + "situ_beta / situ_linear_beta require " + f"ActivationType.SiTu, got {self.activation_type.name}." + ) + self.situ_beta = None if situ_beta is None else float(situ_beta) + self.situ_linear_beta = ( + None if situ_linear_beta is None else float(situ_linear_beta) + ) self.use_2cta_instrs = mma_inst_shape[0] == 256 self.cta_group = ( @@ -1968,7 +2020,7 @@ def __call__( This method performs FC1 layer computation: 1. GEMM: acc = fc1_alpha * (SFA * A[token_ids]) * (SFB * B) - 2. SwiGLU: C = up * silu(gate), decoded from gate16/up16 accumulator groups + 2. Gated act: SwiGLU or SiTU, decoded from gate16/up16 accumulator groups (N/2) 3. Optional Quant: When c_dtype is Float4E2M1FN, generates SFC and quantizes output Data loading: @@ -1988,7 +2040,7 @@ def __call__( shared memory when use_2cta_instrs is True - TMA warp: Load B and SFB with multicast - MMA warp: Perform matrix multiply-accumulate - - Epilogue warps: Apply SwiGLU activation, optional quantization, and store results + - Epilogue warps: Apply gated activation, optional quantization, and store results :param fc1_a: FC1 input A (MxKx1), gathered with the shared route mapping :type fc1_a: cute.Tensor @@ -2615,6 +2667,176 @@ class SharedStorageCpasync2cta: ) return + @cute.jit + def _apply_swiglu_epilogue( + self, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + alpha_val, + tCompute: cute.Tensor, + ): + """SwiGLU: ``tCompute[i] = (alpha * up[i]) * silu(alpha * gate[i])``. + + Optional DS-v4 clamp via ``has_swiglu_limit`` is folded at trace time. + """ + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + acc_vec_up_alpha = cute.arch.mul_packed_f32x2( + (acc_vec_up[i], acc_vec_up[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + ) + acc_vec_gate_alpha = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + ) + if cutlass.const_expr(self.has_swiglu_limit): + acc_vec_gate_alpha = ( + fmin(acc_vec_gate_alpha[0], self.swiglu_limit), + fmin(acc_vec_gate_alpha[1], self.swiglu_limit), + ) + acc_vec_up_alpha = ( + fclip_xorsign(acc_vec_up_alpha[0], self.swiglu_limit), + fclip_xorsign(acc_vec_up_alpha[1], self.swiglu_limit), + ) + tCompute_log2e = cute.arch.mul_packed_f32x2( + (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), + (-LOG2_E, -LOG2_E), + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(tCompute_log2e[0], fastmath=True), + cute.math.exp2(tCompute_log2e[1], fastmath=True), + ), + (1.0, 1.0), + ) + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_up_alpha[0], acc_vec_up_alpha[1]), + ) + else: + for i in cutlass.range_constexpr(cute.size(acc_vec_up.shape)): + acc_vec_up_alpha = acc_vec_up[i] * cutlass.Float32(alpha_val) + acc_vec_gate_alpha = acc_vec_gate[i] * cutlass.Float32(alpha_val) + if cutlass.const_expr(self.has_swiglu_limit): + acc_vec_gate_alpha = fmin(acc_vec_gate_alpha, self.swiglu_limit) + acc_vec_up_alpha = fclip_xorsign( + acc_vec_up_alpha, self.swiglu_limit + ) + tCompute[i] = acc_vec_up_alpha * silu_f32( + acc_vec_gate_alpha, fastmath=True + ) + + @cute.jit + def _apply_situ_epilogue( + self, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + alpha_val, + tCompute: cute.Tensor, + ): + """SiTU (Kimi K3), matching ``SituAndMul`` / the two-op Rubin kernel:: + + g = alpha * gate, u = alpha * up + situ_gate = beta * tanh(g / beta) * sigmoid(g) + situ_up = linear_beta * tanh(u / linear_beta) + tCompute = situ_gate * situ_up + + Packed path uses ``tanh(z) = 2 * sigmoid(2z) - 1``. Betas fold at + trace time; a different pair compiles a different kernel. + + Keep the packed SiTU algebra in sync with the Blackwell two-op + ``blockscaled_contiguous_gather_grouped_gemm_act_fusion.py`` epilogue. + """ + beta = self.situ_beta + linear_beta = self.situ_linear_beta + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + neg_log2e_pair = (-LOG2_E, -LOG2_E) + one_pair = (cutlass.Float32(1.0), cutlass.Float32(1.0)) + + inv_2beta = cutlass.Float32(2.0 / beta) + two_beta = cutlass.Float32(2.0 * beta) + neg_beta = cutlass.Float32(-beta) + inv_2lbeta = cutlass.Float32(2.0 / linear_beta) + two_lbeta = cutlass.Float32(2.0 * linear_beta) + neg_lbeta = cutlass.Float32(-linear_beta) + + def _sigmoid(p0, p1): + neg = cute.arch.mul_packed_f32x2((p0, p1), neg_log2e_pair) + e = ( + cute.math.exp2(neg[0], fastmath=True), + cute.math.exp2(neg[1], fastmath=True), + ) + d = cute.arch.add_packed_f32x2(e, one_pair) + return (cute.arch.rcp_approx(d[0]), cute.arch.rcp_approx(d[1])) + + alpha_pair = (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)) + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + g = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), alpha_pair + ) + u = cute.arch.mul_packed_f32x2( + (acc_vec_up[i], acc_vec_up[i + 1]), alpha_pair + ) + + sigmoid_g = _sigmoid(g[0], g[1]) + + gs = _sigmoid( + *cute.arch.mul_packed_f32x2(g, (inv_2beta, inv_2beta)) + ) + tanh_g = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(gs, (two_beta, two_beta)), + (neg_beta, neg_beta), + ) + + us = _sigmoid( + *cute.arch.mul_packed_f32x2(u, (inv_2lbeta, inv_2lbeta)) + ) + tanh_u = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(us, (two_lbeta, two_lbeta)), + (neg_lbeta, neg_lbeta), + ) + + situ_gate = cute.arch.mul_packed_f32x2(tanh_g, sigmoid_g) + out_pair = cute.arch.mul_packed_f32x2(situ_gate, tanh_u) + tCompute[i] = out_pair[0] + tCompute[i + 1] = out_pair[1] + else: + inv_2beta = cutlass.Float32(2.0 / beta) + two_beta = cutlass.Float32(2.0 * beta) + beta_f32 = cutlass.Float32(beta) + inv_2lbeta = cutlass.Float32(2.0 / linear_beta) + two_lbeta = cutlass.Float32(2.0 * linear_beta) + lbeta_f32 = cutlass.Float32(linear_beta) + for i in cutlass.range_constexpr(cute.size(acc_vec_up.shape)): + g = acc_vec_gate[i] * cutlass.Float32(alpha_val) + u = acc_vec_up[i] * cutlass.Float32(alpha_val) + tanh_g = ( + two_beta * sigmoid_f32(g * inv_2beta, fastmath=True) + - beta_f32 + ) + tanh_u = ( + two_lbeta * sigmoid_f32(u * inv_2lbeta, fastmath=True) + - lbeta_f32 + ) + tCompute[i] = (tanh_g * sigmoid_f32(g, fastmath=True)) * tanh_u + # GPU device kernel @cute.kernel def kernel( @@ -4792,7 +5014,8 @@ def kernel( acc_pipeline.consumer_wait(acc_consumer_state) # Each N64 output subtile consumes two consecutive N64 - # accumulator chunks, each [gate16, up16, gate16, up16]. + # accumulator chunks, each [gate16, up16, gate16, up16], + # decoded into gate/up for SwiGLU or SiTU. # tFC1TR_tAcc: (T2R, T2R_M, T2R_N, EPI_M, EPI_N, # STAGE), sliced on STAGE. bSG_gC: ((ATOM_V, REST_V), # EPI_M, EPI_N, loopM, loopN, loopL). @@ -4841,103 +5064,22 @@ def kernel( acc_vec_up = tFC1TR_rAcc_up.load() acc_vec_gate = tFC1TR_rAcc_gate.load() - # SwiGLU: output = up * silu(gate), silu(x) = x * sigmoid(x). + # Gate/up are decoded from gate16/up16 groups; + # C has N/2 cols. SwiGLU vs SiTU is a const_expr + # specialization (trace-time). tCompute = cute.make_rmem_tensor( acc_vec_gate.shape, self.acc_dtype ) - if cutlass.const_expr(self.vectorized_f32): - # SwiGLU Packed Version: uses f32x2 packed operations for better performance - # Computes: output = (alpha * up) * silu(alpha * gate) - # where silu(x) = x * sigmoid(x) = x / (1 + exp(-x)) - LOG2_E = cutlass.Float32(1.4426950408889634) - for i in cutlass.range_constexpr( - 0, cute.size(tFC1TR_rAcc_up), 2 - ): - acc_vec_up_alpha = cute.arch.mul_packed_f32x2( - (acc_vec_up[i], acc_vec_up[i + 1]), - ( - cutlass.Float32(alpha_val), - cutlass.Float32(alpha_val), - ), - ) - acc_vec_gate_alpha = cute.arch.mul_packed_f32x2( - (acc_vec_gate[i], acc_vec_gate[i + 1]), - ( - cutlass.Float32(alpha_val), - cutlass.Float32(alpha_val), - ), - ) - if cutlass.const_expr(self.has_swiglu_limit): - acc_vec_gate_alpha = ( - fmin( - acc_vec_gate_alpha[0], self.swiglu_limit - ), - fmin( - acc_vec_gate_alpha[1], self.swiglu_limit - ), - ) - acc_vec_up_alpha = ( - fclip_xorsign( - acc_vec_up_alpha[0], self.swiglu_limit - ), - fclip_xorsign( - acc_vec_up_alpha[1], self.swiglu_limit - ), - ) - tCompute_log2e = cute.arch.mul_packed_f32x2( - (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), - (-LOG2_E, -LOG2_E), - ) - ( - tCompute[i], - tCompute[i + 1], - ) = cute.arch.add_packed_f32x2( - ( - cute.math.exp2( - tCompute_log2e[0], fastmath=True - ), - cute.math.exp2( - tCompute_log2e[1], fastmath=True - ), - ), - (1.0, 1.0), - ) - tCompute[i] = cute.arch.rcp_approx(tCompute[i]) - tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) - ( - tCompute[i], - tCompute[i + 1], - ) = cute.arch.mul_packed_f32x2( - (tCompute[i], tCompute[i + 1]), - (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), - ) - ( - tCompute[i], - tCompute[i + 1], - ) = cute.arch.mul_packed_f32x2( - (tCompute[i], tCompute[i + 1]), - (acc_vec_up_alpha[0], acc_vec_up_alpha[1]), - ) + if cutlass.const_expr( + self.activation_type == ActivationType.SiTu + ): + self._apply_situ_epilogue( + acc_vec_up, acc_vec_gate, alpha_val, tCompute + ) else: - # SwiGLU Unpacked Version: scalar operations - # Computes: output = (alpha * up) * silu(alpha * gate) - for i in cutlass.range_constexpr(cute.size(tFC1TR_rAcc_up)): - acc_vec_up_alpha = acc_vec_up[i] * cutlass.Float32( - alpha_val - ) - acc_vec_gate_alpha = acc_vec_gate[i] * cutlass.Float32( - alpha_val - ) - if cutlass.const_expr(self.has_swiglu_limit): - acc_vec_gate_alpha = fmin( - acc_vec_gate_alpha, self.swiglu_limit - ) - acc_vec_up_alpha = fclip_xorsign( - acc_vec_up_alpha, self.swiglu_limit - ) - tCompute[i] = acc_vec_up_alpha * silu_f32( - acc_vec_gate_alpha, fastmath=True - ) + self._apply_swiglu_epilogue( + acc_vec_up, acc_vec_gate, alpha_val, tCompute + ) if cutlass.const_expr(self.fc1_generate_sfc): # Float4E2M1FN quantization: per-vector absmax → diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index c726cb6c290f..fedceda55f4a 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -260,9 +260,9 @@ def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " "None means an identity linear branch that the fused kernels cannot represent." ) - if situ_beta <= 0 or situ_linear_beta <= 0: + if not (0 < situ_beta < float("inf") and 0 < situ_linear_beta < float("inf")): raise ValueError( - f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." + f"Kimi K3 SiTu betas must be finite and positive; got {situ_beta} and {situ_linear_beta}." ) return float(situ_beta), float(situ_linear_beta) @@ -1190,7 +1190,12 @@ def __init__( # CUTLASS is absent on purpose: it is the fallback target, so # "degraded to CUTLASS" is not a thing that can happen to it. allow_backend_degradation=routed_moe_model_config.moe_backend - not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), + not in ( + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + "CUTEDSL", + "CUTEDSL_FC12", + ), ) self._check_trtllm_situ_quant( routed_moe_model_config.moe_backend, routed_quant_config.quant_algo @@ -1403,13 +1408,15 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: "CUTLASS", "TRTLLM", "CUTEDSL", + "CUTEDSL_FC12", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", } if model_config.moe_backend not in supported_backends: raise ValueError( "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " - "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " + "CUTEDSL, CUTEDSL_FC12, MEGAMOE_DEEPGEMM, and " + "MEGAMOE_CUTEDSL backends; " f"got {model_config.moe_backend!r}." ) if model_config.moe_load_balancer is not None: diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 99891bd3d3a4..6862c99ac176 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -227,7 +227,7 @@ Still on old path (standalone, with embedded communication): | `fused_moe_densegemm.py` | `TrtllmCutedslDenseGemmNvfp4Impl` (aliased as `DenseGEMMFusedMoE`) | SM100/SM103 | NVFP4 min-latency; CuTe DSL dense GEMM packs all experts into one matrix (vs Cutlass per-expert scatter), efficient for small token counts | `EXTERNAL_COMM` | | `fused_moe_cute_dsl.py` | `CuteDslFusedMoE` | SM100/SM103/SM107 | High throughput NVFP4, generally faster than Cutlass | `EXTERNAL_COMM` | | `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xFusedMoE` | SM120/SM121 | NVFP4 hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode — best perf on RTX PRO 6000 (SM120) and DGX Spark (SM121); select via the `CUTEDSL` backend path (it heads that family's candidate list, so it wins on SM120/121 when flashinfer is present and yields to `CuteDslFusedMoE` otherwise); single-GPU-shaped topology only — it rejects both `ep_size > 1` and attention-DP, because it has no dispatch/combine kernel and has never been exercised behind a DP allgather | `EXTERNAL_COMM` | -| `fused_moe_cute_dsl_fc12.py` | `TrtllmCutedslFusedFc12Nvfp4Impl` (`trtllm.cutedsl.fused_fc12.nvfp4`, aliased as `CuteDslFc12FusedMoE`) | SM107 (Rubin) | NVFP4 fused FC1+FC2: a `MoEImplBase` leaf sharing `CuteDslFusedMoE`'s NVFP4 weight layout and outer autotune runner, driving the single persistent `cute_dsl_nvfp4_fc12_fused_rubin` kernel (keeps the FC1->FC2 intermediate on-chip, removing the global round-trip and one launch); reachable through the `CUTEDSL_FC12` backend family or a pinned `impl_id`; requires CuTe DSL Rubin support (`MoEDep.CUTEDSL_RUBIN` in the collected environment); routing tile 128 (1-CTA) or 256 (2-CTA, cluster (2,1)), chosen by the autotuner per shape; non-uGPU | `EXTERNAL_COMM` | +| `fused_moe_cute_dsl_fc12.py` | `TrtllmCutedslFusedFc12Nvfp4Impl` (`trtllm.cutedsl.fused_fc12.nvfp4`, aliased as `CuteDslFc12FusedMoE`) | SM107 (Rubin) | NVFP4 fused FC1+FC2 with SwiGLU or SiTU (uniform scalar softcaps): a `MoEImplBase` leaf sharing `CuteDslFusedMoE`'s NVFP4 weight layout and outer autotune runner, driving the single persistent `cute_dsl_nvfp4_fc12_fused_rubin` kernel (keeps the FC1->FC2 intermediate on-chip, removing the global round-trip and one launch); reachable through the `CUTEDSL_FC12` backend family or a pinned `impl_id`; requires CuTe DSL Rubin support (`MoEDep.CUTEDSL_RUBIN` in the collected environment); routing tile 128 (1-CTA) or 256 (2-CTA, cluster (2,1)), chosen by the autotuner per shape; non-uGPU | `EXTERNAL_COMM` | | `mega_moe/mega_moe_deepgemm.py` | `DeepgemmCudaW4a8Mxfp4Mxfp8Impl` (aliased as `MegaMoEDeepGemm`) | SM100/SM103/SM107 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | | `mega_moe/mega_moe_cute_dsl.py` | `TrtllmCutedslMegaMoeNvfp4Impl` (aliased as `MegaMoECuteDsl`) | SM100/SM103/SM107 | NVFP4 fused dispatch+FC1+act+FC2+combine; internally selects one of its per-architecture kernel variants. Requires CUDA 13 Cutlass DSL and a symmetric-memory provider. Uses uniform activation constants and applies routing weights before FC2 quantization for DeepSeek-V4, after FC2 for other models. | `FUSED_COMM` | | `marlin/` (+ the `fused_moe_marlin.py` alias module) | the abstract `MarlinFusedMoEBase` and its two leaves, `MarlinCudaNvfp4Impl` and `MarlinCudaW4a16Nvfp4Impl` | SM89-SM99 | W4A16 NVFP4 on Ada/Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) degrade to Cutlass in `resolve_moe_impl`, recorded in the layer's `MoEResolutionReport`; no dynamic EPLB | `EXTERNAL_COMM` | diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py index 5969ad27d6c3..54a010b6654e 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py @@ -15,7 +15,7 @@ """TrtllmCutedslFusedFc12Nvfp4Impl: ``trtllm.cutedsl.fused_fc12.nvfp4``. FC1+FC2-fused CuteDSL NVFP4 MoE backend for Rubin (SM107). One persistent kernel -(``trtllm::cute_dsl_nvfp4_fc12_fused_rubin``) does gather + FC1 GEMM + SwiGLU + +(``trtllm::cute_dsl_nvfp4_fc12_fused_rubin``) does gather + FC1 GEMM + SwiGLU/SiTU + requant + FC2 GEMM + finalize (scatter-add into ``moe_output``), so the FC1->FC2 intermediate never round-trips through global memory. @@ -78,6 +78,25 @@ class CuteDslFc12FusedMoENvfp4Runner(CuteDslFusedMoENvfp4Runner): def _tile_sizes(): return [128, 256] + def forward( + self, + inputs: List[torch.Tensor], + tactic: Optional[int], + do_preparation: bool = False, + ) -> torch.Tensor: + """Prime both fused tiles before outer CUDA-graph profiling. + + The inner fused op may need to compile or tune, which cannot happen + inside the outer profiler's CUDA graph. The parent's preparation also + passes a two-op memset argument that FC12 does not accept. This path + runs for SwiGLU as well as SiTU when the outer tuner prepares a shape. + """ + if do_preparation: + for tile_size in self._tile_sizes(): + super().forward(inputs, tactic=tile_size) + return inputs[4] + return super().forward(inputs, tactic=tactic) + @register_moe_impl class TrtllmCutedslFusedFc12Nvfp4Impl(MoEImplBase): @@ -106,7 +125,7 @@ class TrtllmCutedslFusedFc12Nvfp4Impl(MoEImplBase): supports_apply_router_weight_on_input=True, ), input_requirement=MoEInputRequirement(routing_scales_dtype=torch.float32), - doc="Rubin (SM107) NVFP4 fused FC1 + SwiGLU + FC2 + finalize persistent " + doc="Rubin (SM107) NVFP4 fused FC1 + gated activation + FC2 + finalize persistent " "CuTe DSL grouped GEMM.", ) # Taken off the descriptor rather than restated: the scheduler reads these @@ -116,11 +135,11 @@ class TrtllmCutedslFusedFc12Nvfp4Impl(MoEImplBase): capabilities = descriptor.capabilities input_requirement = descriptor.input_requirement - # The fused FC12 kernel has no activation selector: FC1 weights are the - # gate/up pair and the epilogue is SwiGLU (+ clamp). Relu2 is not gated and - # is not supported, so the resolver turns such layers down here. + # Both activations share the gate/up geometry. SiTU softcaps are baked + # into the epilogue, so they must be uniform across the local experts. activation_support = MoEActivationSupport( - kinds=frozenset({ActivationType.Swiglu}), + kinds=frozenset({ActivationType.Swiglu, ActivationType.SiTu}), + alpha_beta=ActivationParamShape.UNIFORM_SCALAR, limit=ActivationParamShape.UNIFORM_SCALAR, limit_when_absent=float("inf"), ) @@ -380,6 +399,7 @@ def run_moe_nvfp4( local_expert_offset=weight_view.slot_start, enable_finalize_fusion=self.use_fused_finalize, enable_alltoall=enable_alltoall, + workload_identity=(int(self.activation_type), self.act_alpha, self.act_beta), ) inputs = [x, token_selected_experts, token_final_scales, x_sf, moe_output, weight_view] _, best_tactic = tuner.choose_one( @@ -436,7 +456,7 @@ def run_moe_nvfp4_impl( tile_tokens_dim=tile_size, ) - # One fused op: gather + FC1 GEMM + SwiGLU + requant + FC2 GEMM + + # One fused op: gather + FC1 GEMM + gated activation + requant + FC2 GEMM + # finalize (scatter-add into moe_output = a2a combine workspace). # fc1_alpha/fc2_alpha map 1:1 to the two-op path's per-expert global # scales (the kernel takes split alphas). The three atomic counters are @@ -471,6 +491,9 @@ def run_moe_nvfp4_impl( ep_size=self.mapping.moe_ep_size, enable_alltoall=enable_alltoall, scaling_vector_size=16, + activation_type=int(self.activation_type), + situ_beta=self.act_alpha, + situ_linear_beta=self.act_beta, ) return moe_output diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 1b7d163cbe9e..8c8b90b240c2 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -210,7 +210,11 @@ l0_b200: - unittest/_torch/moe/test_moe_backend.py::test_megamoe_streaming_reload_resets_slot_claims - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale - - unittest/_torch/moe/test_moe_backend.py::test_megamoe_bakes_situ_softcaps_as_uniform_scalars + - unittest/_torch/moe/test_moe_backend.py::test_codegen_baked_situ_softcaps_are_uniform_scalars + - unittest/_torch/moe/test_moe_backend.py::test_fc12_rejects_nonuniform_situ_softcaps + - unittest/_torch/moe/test_moe_backend.py::test_fc12_swiglu_keeps_clamp_without_situ_constants + - unittest/_torch/moe/test_moe_backend.py::test_fc12_outer_tuning_separates_activation_and_softcaps + - unittest/_torch/moe/test_moe_backend.py::test_fc12_preparation_primes_fused_tiles_without_two_op_memset_knob - unittest/_torch/moe/test_moe_backend.py::test_megamoe_plain_swiglu_carries_no_constants - unittest/_torch/moe/test_moe_backend.py::test_create_moe_forwards_situ_activation_as_one_carrier - unittest/_torch/moe/test_moe_backend.py::test_create_moe_backend_rejects_apply_router_weight_on_input_by_declaration diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index 384474fd1448..8fc218dba977 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -110,7 +110,16 @@ def test_kimi_situ_betas_require_linear_beta(): @pytest.mark.parametrize( "situ_beta,situ_linear_beta", - [(0.0, 25.0), (4.0, 0.0), (-1.0, 25.0), (4.0, -1.0)], + [ + (0.0, 25.0), + (4.0, 0.0), + (-1.0, 25.0), + (4.0, -1.0), + (float("nan"), 25.0), + (4.0, float("nan")), + (float("inf"), 25.0), + (4.0, float("inf")), + ], ) def test_kimi_situ_betas_must_be_positive(situ_beta, situ_linear_beta): cfg = SimpleNamespace( @@ -118,10 +127,42 @@ def test_kimi_situ_betas_must_be_positive(situ_beta, situ_linear_beta): activation_situ_linear_beta=situ_linear_beta, ) - with pytest.raises(ValueError, match="must be positive"): + with pytest.raises(ValueError, match="must be finite and positive"): modeling_kimi_linear._resolve_kimi_situ_betas(cfg) +@pytest.mark.parametrize( + "situ_beta,situ_linear_beta", + [ + (float("nan"), 25.0), + (4.0, float("nan")), + (float("inf"), 25.0), + (4.0, float("inf")), + ], + ids=["nan_gate", "nan_linear", "inf_gate", "inf_linear"], +) +def test_fc12_kernel_rejects_nonfinite_situ_betas(situ_beta, situ_linear_beta): + pytest.importorskip("cutlass") + pytest.importorskip("cutlass.utils.rubin_helpers") + from tensorrt_llm._torch.cute_dsl_kernels.rubin.moe.rubin_contiguous_grouped_blockscaled_gemm_fused_fc12 import ( # noqa: E501 + Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel, + ) + from tensorrt_llm._torch.utils import ActivationType + + with pytest.raises(ValueError, match="finite and positive"): + Sm107BlockScaledContiguousGroupedGemmFusedFc12Kernel( + sf_vec_size=16, + mma_inst_shape=(128, 128, 256), + mma_tiler=(128, 128, 256), + cluster_shape_mn=(1, 1), + vectorized_f32=True, + topk=8, + activation_type=ActivationType.SiTu, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + + @pytest.mark.parametrize( "situ_beta,situ_linear_beta", [ @@ -555,7 +596,17 @@ def test_kimi_k3_moe_split_selection() -> None: assert KimiK3MoERuntime._select_moe_tp_ep(tep) == (4, 2) -@pytest.mark.parametrize("backend", ["CUTLASS", "TRTLLM", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"]) +@pytest.mark.parametrize( + "backend", + [ + "CUTLASS", + "TRTLLM", + "CUTEDSL", + "CUTEDSL_FC12", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + ], +) def test_kimi_k3_routed_config_preserves_explicit_backend(backend): model_config = ModelConfig( mapping=Mapping(world_size=1, rank=0, tp_size=1), @@ -673,7 +724,8 @@ def admits_situ(name): assert "CUTEDSL" in declares_situ -def test_explicit_cutedsl_fails_instead_of_degrading_to_cutlass(monkeypatch): +@pytest.mark.parametrize("backend", ["CUTEDSL", "CUTEDSL_FC12"]) +def test_explicit_cutedsl_fails_instead_of_degrading_to_cutlass(monkeypatch, backend): """K3 must propagate strict backend selection through create_moe.""" from transformers.configuration_utils import PretrainedConfig @@ -688,7 +740,7 @@ def test_explicit_cutedsl_fails_instead_of_degrading_to_cutlass(monkeypatch): from tensorrt_llm.models.modeling_utils import QuantConfig # Keep the real resolver and K3 caller; only make eligibility deterministic. - for backend_cls in BACKEND_FAMILY["CUTEDSL"]: + for backend_cls in BACKEND_FAMILY[backend]: monkeypatch.setattr( backend_cls, "can_implement", @@ -707,7 +759,7 @@ def test_explicit_cutedsl_fails_instead_of_degrading_to_cutlass(monkeypatch): model_config = ModelConfig( pretrained_config=pretrained_config, mapping=Mapping(world_size=1, rank=0, tp_size=1), - moe_backend="CUTEDSL", + moe_backend=backend, quant_config_dict={"layers.0.mlp.experts": quant_config}, ) cfg = _K3Config(routed_expert_hidden_size=512, latent_moe_use_norm=True) @@ -726,8 +778,8 @@ def test_explicit_cutedsl_fails_instead_of_degrading_to_cutlass(monkeypatch): assert report.degraded assert impl_class_for(report) is CutlassFusedMoE - # Removing CUTEDSL from K3's no-degradation list must fail this assertion. - with pytest.raises(ValueError, match="CUTEDSL.*degradation disallowed") as excinfo: + # Removing either backend from K3's no-degradation list must fail this assertion. + with pytest.raises(ValueError, match=rf"{backend}.*degradation disallowed") as excinfo: KimiK3MoERuntime(model_config, cfg, layer_idx=0, aux_stream_dict={}) assert "dep_missing" in str(excinfo.value) @@ -1383,12 +1435,12 @@ def test_tp16_nvfp4_padded_loaders_preserve_rank_ownership(): assert torch.count_nonzero(w2_sf_dst[:, 12:].float()) == 0 -#: The FP4 backends that serve SiTU. CUTEDSL additionally needs the CuTe DSL -#: wheel, which is checked inside the test rather than in a ``skipif``: +#: The FP4 backends that serve SiTU. The CuteDSL backends additionally need +#: the matching CuTe DSL wheel, checked inside the test rather than a ``skipif``: #: importing ``cute_dsl_utils`` pulls in the DSL package, which appends its own #: directory to ``sys.path``, and this repository fails the whole pytest #: session when a test file does that at collection time. -_NVFP4_SITU_BACKENDS = ["CUTLASS", "TRTLLM", "CUTEDSL"] +_NVFP4_SITU_BACKENDS = ["CUTLASS", "TRTLLM", "CUTEDSL", "CUTEDSL_FC12"] def _skip_if_backend_unavailable(moe_backend): @@ -1399,19 +1451,29 @@ def _skip_if_backend_unavailable(moe_backend): early puts the wheel's package directory on ``sys.path`` for every other test file in the session. """ - if moe_backend != "CUTEDSL": + if moe_backend not in ("CUTEDSL", "CUTEDSL_FC12"): return - from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE, + ) if not IS_CUTLASS_DSL_AVAILABLE: pytest.skip("CuteDSL MoE requires the CuTe DSL wheel") + sm_version = get_sm_version() + if moe_backend == "CUTEDSL_FC12": + if sm_version != 107: + pytest.skip(f"FC12 SiTU MoE requires Rubin (SM107), got SM{sm_version}") + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("FC12 SiTU MoE requires CuTe DSL Rubin support") + return # nvfp4_moe_supported admits every SM >= 100, but this integration enables # SiTU only on Blackwell. Match can_implement() rather than exercising an # end-to-end path that has not been enabled on SM107. - if get_sm_version() not in (100, 103): + if sm_version not in (100, 103): pytest.skip( f"CuteDSL SiTU MoE needs the Blackwell act-fusion kernel " - f"(SM100/SM103), got SM{get_sm_version()}" + f"(SM100/SM103), got SM{sm_version}" ) @@ -1453,13 +1515,13 @@ def block_scales(*shape): def _make_nvfp4_moe( gate, num_experts=_TP_EXPERTS, moe_backend="CUTLASS", gate_softcap=4.0, linear_softcap=25.0 ): - """NVFP4 + SiTU routed MoE on either FP4 backend. + """NVFP4 + SiTU routed MoE on a supported FP4 backend. - Both take the same ``SiTuActivation`` carrier; CUTLASS serves it as an - ``ActivationType`` its kernels branch on, TRTLLM-Gen with the fused - ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*`` FC1 cubins (group-16 block scales). - ``_make_routed_moe`` already mirrors both of KimiK3MoERuntime's branches, - so the backend is the only variable. + All take the same ``SiTuActivation`` carrier. CUTLASS branches on the + ``ActivationType`` in its kernels, TRTLLM-Gen selects the fused + ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*`` FC1 cubins, Blackwell CuteDSL folds + the soft-caps into a JIT epilogue, and FC12 specializes its Rubin fused + epilogue. """ from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -2005,12 +2067,11 @@ def test_nvfp4_kernel_actually_applies_situ(moe_backend): Run for every FP4 backend, because each reaches SiTU by a different mechanism and so fails differently: CUTLASS resolves it through the activation enum, TRTLLM-Gen through a distinct fused-cubin family - (``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``), and CuteDSL through soft-caps - folded into a JIT-compiled epilogue at trace time. The CuteDSL case is - the one this test exists for: its betas travel as trace-time scalars - keyed into the kernel cache, so dropping them does not raise -- it - compiles a SwiGLU kernel and returns plausible numbers. An internal - branch shipped exactly that defect on a sibling backend for weeks. + (``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``), Blackwell CuteDSL through + soft-caps folded into a JIT-compiled epilogue at trace time, and FC12 + through its Rubin fused epilogue. The CuteDSL cases are why this test + exists: dropping their trace-time activation parameters does not raise -- + it compiles a SwiGLU kernel and returns plausible numbers. The degenerate all-zero FC1 output is caught as well: it scores 0 against both references. @@ -2058,10 +2119,33 @@ def test_nvfp4_kernel_actually_applies_situ(moe_backend): bank = [_quantize_expert_to_nvfp4(w1[e], w2[e], w3[e], act_scale) for e in range(num_experts)] moe = _make_nvfp4_moe(gate, num_experts=num_experts, moe_backend=moe_backend) + if moe_backend == "CUTEDSL_FC12": + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_fc12 import ( + TrtllmCutedslFusedFc12Nvfp4Impl, + ) + + assert type(moe.backend) is TrtllmCutedslFusedFc12Nvfp4Impl _load_nvfp4_bank_for(moe, bank, moe_backend) router_logits = gate.compute_logits(x) actual = moe.forward(x, router_logits, all_rank_num_tokens=None).float() + assert torch.isfinite(actual).all() + + if moe_backend == "CUTEDSL_FC12": + # Warm up on a side stream before capture; compilation and tuning + # must finish before CUDA Graph records the fused op. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + moe.forward(x, router_logits, all_rank_num_tokens=None) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = moe.forward(x, router_logits, all_rank_num_tokens=None) + for _ in range(2): + graph.replay() + torch.testing.assert_close(captured.float(), actual, rtol=0.02, atol=0.02) # Both references read the weights BACK OUT of the checkpoint tensors, so # they see the same 4-bit values the kernel was loaded with. Against the diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index adbe7eaf2f0b..91b8db3d62bd 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -73,6 +73,10 @@ CuteDslFusedMoENvfp4Runner, ) from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl_fc12 import ( + CuteDslFc12FusedMoENvfp4Runner, + TrtllmCutedslFusedFc12Nvfp4Impl, +) from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_marlin import MarlinFusedMoE @@ -1890,14 +1894,15 @@ def test_megamoe_cutedsl_cache_derived_state_survives_the_read_only_reader_walk( wrapper.cache_derived_state.assert_not_called() -def test_megamoe_bakes_situ_softcaps_as_uniform_scalars(): - # MegaMoE declares UNIFORM_SCALAR for alpha/beta because the kernels bake - # them at codegen time, so a per-expert tensor is reduced here. +@pytest.mark.parametrize("impl", [MegaMoEDeepGemm, TrtllmCutedslFusedFc12Nvfp4Impl]) +def test_codegen_baked_situ_softcaps_are_uniform_scalars(impl): + # These kernels bake alpha/beta at codegen time, so the activation adapter + # accepts only one value shared by every local expert. params = materialize_activation_params( SiTuActivation(gate_softcap=torch.full((8,), 4.0), linear_softcap=25.0), - MegaMoEDeepGemm.activation_support, + impl.activation_support, num_local_experts=8, - owner="MegaMoEDeepGemm", + owner=impl.__name__, ) assert params.activation_type is ActivationType.SiTu @@ -1905,6 +1910,82 @@ def test_megamoe_bakes_situ_softcaps_as_uniform_scalars(): assert params.beta == 25.0 +@pytest.mark.parametrize("nonuniform_softcap", ["gate", "linear"]) +def test_fc12_rejects_nonuniform_situ_softcaps(nonuniform_softcap: str) -> None: + gate_softcap = torch.tensor([4.0, 5.0]) if nonuniform_softcap == "gate" else 4.0 + linear_softcap = torch.tensor([25.0, 26.0]) if nonuniform_softcap == "linear" else 25.0 + with pytest.raises(ValueError, match="uniform"): + materialize_activation_params( + SiTuActivation(gate_softcap=gate_softcap, linear_softcap=linear_softcap), + TrtllmCutedslFusedFc12Nvfp4Impl.activation_support, + num_local_experts=2, + owner="FC12", + ) + + +@pytest.mark.parametrize("clamp", [None, 7.0]) +def test_fc12_swiglu_keeps_clamp_without_situ_constants(clamp: Optional[float]) -> None: + params = materialize_activation_params( + SwigluActivation(clamp=clamp), + TrtllmCutedslFusedFc12Nvfp4Impl.activation_support, + num_local_experts=2, + owner="FC12", + ) + assert params.activation_type is ActivationType.Swiglu + assert (params.alpha, params.beta) == (None, None) + assert params.clamp == (float("inf") if clamp is None else clamp) + + +def test_fc12_outer_tuning_separates_activation_and_softcaps() -> None: + identities = [ + (int(ActivationType.Swiglu), None, None), + (int(ActivationType.SiTu), 4.0, 25.0), + (int(ActivationType.SiTu), 5.0, 25.0), + (int(ActivationType.SiTu), 4.0, 30.0), + ] + keys = [ + CuteDslFc12FusedMoENvfp4Runner( + forward_impl=MagicMock(), + num_experts=8, + top_k=2, + num_local_experts=8, + local_expert_offset=0, + workload_identity=identity, + ).unique_id() + for identity in identities + ] + assert len(set(keys)) == len(identities) + + +def test_fc12_preparation_primes_fused_tiles_without_two_op_memset_knob() -> None: + def fused_forward( + *inputs, + enable_alltoall, + tile_size, + recv_expert_count, + deep_ep_expert_capacity, + use_count_native_expert_metadata, + ): + assert not enable_alltoall + assert recv_expert_count is None and deep_ep_expert_capacity is None + assert not use_count_native_expert_metadata + tiles.append(tile_size) + return inputs[4] + + tiles = [] + runner = CuteDslFc12FusedMoENvfp4Runner( + forward_impl=fused_forward, + num_experts=8, + top_k=2, + num_local_experts=8, + local_expert_offset=0, + workload_identity=(int(ActivationType.SiTu), 4.0, 25.0), + ) + inputs = [object() for _ in range(6)] + assert runner.forward(inputs, tactic=None, do_preparation=True) is inputs[4] + assert tiles == [128, 256] + + def test_megamoe_plain_swiglu_carries_no_constants(): params = materialize_activation_params( SwigluActivation(), diff --git a/tests/unittest/_torch/moe/test_moe_impl.py b/tests/unittest/_torch/moe/test_moe_impl.py index b751cb0b5306..cf260b910f79 100644 --- a/tests/unittest/_torch/moe/test_moe_impl.py +++ b/tests/unittest/_torch/moe/test_moe_impl.py @@ -23,6 +23,7 @@ or numerical parity across backends, belong in ``test_moe_backend.py``. """ +from dataclasses import replace from unittest.mock import MagicMock import pytest @@ -988,6 +989,23 @@ def test_fc12_identity_round_trips_through_registry(): assert impl.input_requirement is impl.descriptor.input_requirement +@pytest.mark.parametrize("activation,accepted", [("SiTu", True), ("Relu2", False)]) +def test_fc12_pinned_activation_support(activation: str, accepted: bool) -> None: + problem = replace( + _fc12_problem(), + activation=activation, + activation_constants=frozenset({"alpha", "beta"}) if activation == "SiTu" else frozenset(), + ) + report = resolve_moe_impl( + ModelConfig(), problem=problem, deployment=_fc12_deployment(), impl_id=_FC12_IMPL_ID + ) + if accepted: + assert impl_class_for(report) is TrtllmCutedslFusedFc12Nvfp4Impl + else: + assert report.winner is None + assert report.rejected[0].reason is MoERejectReason.ACTIVATION_UNSUPPORTED + + def test_pinned_fc12_identity_fails_hard_where_the_backend_literal_degrades(): """On Rubin both tracks land on FC12; off Rubin the literal degrades and the pin does not.""" config = ModelConfig()