diff --git a/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml index 71b07d93dd5b..e3a726a5c56a 100644 --- a/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml +++ b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml @@ -20,8 +20,13 @@ # sizing pass. Lowering the batch lowers the minimum instead of the headroom. # It costs eval wall time, not accuracy. # -# backend: CUTLASS is required, not a preference -- AUTO resolves Kimi K3 to -# TRTLLM, and trtllm-gen ships SiTu cubins for W4A8_MXFP4_MXFP8 only. +# backend: CUTLASS is a choice here, not a requirement. AUTO resolves Kimi K3 +# to TRTLLM, which serves NVFP4 SiTu since #17940 (group-16 +# Bmm_E2m1_E2m1E2m1_..._siTuGlu_* cubins) -- the model-layer guard that used to +# refuse it was removed in #18709. CUTEDSL and MEGAMOE_CUTEDSL also serve NVFP4 +# SiTu. This file pins CUTLASS so the recipe below (max_num_tokens, KV sizing) +# stays the one it was measured with; switching backends means re-checking +# those, not just this line. # # moe_config.max_num_tokens stays at the inherited value: for CUTLASS it is a # per-call chunking bound. Do NOT carry it over to MEGAMOE_* backends, where diff --git a/examples/kimi_k3/quick_start_kimi_k3.py b/examples/kimi_k3/quick_start_kimi_k3.py index c1db17bc9272..5854ae28f6f9 100644 --- a/examples/kimi_k3/quick_start_kimi_k3.py +++ b/examples/kimi_k3/quick_start_kimi_k3.py @@ -10,7 +10,7 @@ import argparse from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MambaStateConfig +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MambaStateConfig, MoeConfig SAMPLES = [ ("The capital of France is", "Paris"), @@ -26,6 +26,12 @@ def parse_arguments() -> argparse.Namespace: + """CLI for the single-node-per-rank Kimi K3 smoke run. + + The defaults describe a DEP16 deployment: ``--tp-size`` sets tensor and + expert parallelism together, since Kimi K3 runs attention-DP with + expert-parallel MoE and the two sizes are the same number. + """ parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--model", @@ -45,12 +51,24 @@ def parse_arguments() -> argparse.Namespace: "manager with KDA recurrent-state snapshots; prefix-cache hits " "skip recomputing shared prompt prefixes).", ) + parser.add_argument( + "--moe-backend", + default=None, + help="Force moe_config.backend (e.g. CUTEDSL). Default leaves AUTO, " + "which Kimi K3 resolves to TRTLLM.", + ) return parser.parse_args() def main() -> None: + """Generate the four sample prompts and report whether each hit. + + Prints the expected substring check per prompt rather than asserting, so + a run that loads and generates but answers wrongly is visible in the log + instead of collapsing into a single non-zero exit. + """ args = parse_arguments() - llm = LLM( + llm_kwargs = dict( model=args.model, tensor_parallel_size=args.tp_size, enable_attention_dp=True, @@ -72,6 +90,9 @@ def main() -> None: else MambaStateConfig(), ), ) + if args.moe_backend: + llm_kwargs["moe_config"] = MoeConfig(backend=args.moe_backend) + llm = LLM(**llm_kwargs) sampling_params = SamplingParams(max_tokens=64, temperature=0.0) prompts = [prompt for prompt, _ in SAMPLES] diff --git a/examples/kimi_k3/quick_start_kimi_k3.sbatch b/examples/kimi_k3/quick_start_kimi_k3.sbatch index 336753c5305d..46a162d3b653 100644 --- a/examples/kimi_k3/quick_start_kimi_k3.sbatch +++ b/examples/kimi_k3/quick_start_kimi_k3.sbatch @@ -27,7 +27,8 @@ set -euo pipefail usage() { - echo "Usage: sbatch $0 --model PATH --image PATH [--enable-block-reuse]" + echo "Usage: sbatch $0 --model PATH --image PATH [--enable-block-reuse]" \ + "[--moe-backend BACKEND]" } MODEL="" @@ -40,6 +41,15 @@ while [[ $# -gt 0 ]]; do FEATURE_ARGS="$FEATURE_ARGS $1" shift ;; + --moe-backend) + [[ $# -ge 2 ]] || { echo "error: --moe-backend requires a value" >&2; usage >&2; exit 2; } + FEATURE_ARGS="$FEATURE_ARGS --moe-backend $2" + shift 2 + ;; + --moe-backend=*) + FEATURE_ARGS="$FEATURE_ARGS --moe-backend ${1#*=}" + shift + ;; --model) [[ $# -ge 2 ]] || { echo "error: --model requires a value" >&2; usage >&2; exit 2; } MODEL=$2 @@ -96,6 +106,14 @@ srun --mpi=pmix \ export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} mkdir -p \"\$TRITON_CACHE_DIR\" + # DeepGEMM JIT writes compiled artifacts through a temporary file + + # rename. Sharing one cache across ranks lets two compilers race that + # rename (444447: compiler.hpp:147 runtime != nullptr in + # fp8_swap_ab_gemm). Keep it node-local and per-rank, like GSM8K. + export DG_JIT_CACHE_DIR=/tmp/deep-gemm-rank\${SLURM_PROCID:-0} + mkdir -p \"\$DG_JIT_CACHE_DIR\" + echo \"[k3] DG_JIT_CACHE_DIR=\$DG_JIT_CACHE_DIR\" + # Node-local flashinfer cache: the default (\$HOME) is shared NFS and # races across ranks during cubin download (stale file handles). export FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-rank\${SLURM_PROCID:-0} 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 08264a03c0bd..768b09104691 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -3371,13 +3371,19 @@ def __init__(self, tile_size: int, scaling_vector_size: int = 16, activation_type: ActivationType = ActivationType.Swiglu, - swiglu_limit_scalar: float = float("inf")): + swiglu_limit_scalar: float = float("inf"), + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None): """Initialize the runner. Args: - activation_type: ``ActivationType`` for the fused epilogue. Only - ``Swiglu`` (gated) and ``Relu2`` (non-gated) are supported. + activation_type: ``ActivationType`` for the fused epilogue. + ``Swiglu`` (gated), ``Relu2`` (non-gated) and ``SiTu`` + (gated) are supported. swiglu_limit_scalar: Uniform clamp limit for SwiGLU. ``+inf`` disables clamp. + situ_beta: Gate-side SiTU soft-cap. Required for -- and only + valid with -- ``ActivationType.SiTu``. + situ_linear_beta: Linear-side SiTU soft-cap, same rule. """ super().__init__() self.activation_type = validate_activation_type(activation_type) @@ -3393,6 +3399,12 @@ def __init__(self, self.tile_size = tile_size self.scaling_vector_size = scaling_vector_size self.swiglu_limit_scalar = swiglu_limit_scalar + # Trace-time constants, so they are part of the kernel identity -- + # see ``unique_id`` and the compile cache key below. Betas that are + # not keyed would let a layer silently reuse a kernel compiled for + # different soft-caps. + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta if (sm_version := get_sm_version()) not in (100, 103): raise ValueError( @@ -3405,6 +3417,14 @@ def __init__(self, ) def unique_id(self): + """Identity of the compiled kernel, for the autotuner's cache. + + Every entry here is a trace-time constant folded into the kernel, + so two runners that differ in any of them are different kernels + and must not share a tuning result. That is why the activation + soft-caps appear: ``swiglu_limit_scalar`` and the two SiTU betas + are baked in as ``const_expr``, not passed at launch. + """ return ( self.num_experts, self.top_k, @@ -3414,6 +3434,8 @@ def unique_id(self): self.scaling_vector_size, self.activation_type, self.swiglu_limit_scalar, + self.situ_beta, + self.situ_linear_beta, ) def get_valid_tactics( @@ -3616,7 +3638,8 @@ def forward(self, inputs: List, cache_key = (self.scaling_vector_size, self.tile_size, self.top_k, mma_tiler_mn, cluster_shape_mn, raster_along_m, - self.activation_type, self.swiglu_limit_scalar) + self.activation_type, self.swiglu_limit_scalar, + self.situ_beta, self.situ_linear_beta) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( @@ -3628,6 +3651,8 @@ def forward(self, inputs: List, raster_along_m=raster_along_m, activation_type=self.activation_type, swiglu_limit=self.swiglu_limit_scalar, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -3713,12 +3738,18 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( scaling_vector_size: int = 16, activation_type: int = int(ActivationType.Swiglu), swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """CuteDSL-based NVFP4 gather grouped GEMM with activation fusion. - Supports ``ActivationType.Swiglu`` (gated) and ``ActivationType.Relu2`` - (non-gated) epilogues; other ``ActivationType`` values raise an - assertion in the runner. + Supports ``ActivationType.Swiglu`` (gated), ``ActivationType.Relu2`` + (non-gated) and ``ActivationType.SiTu`` (gated) epilogues; other + ``ActivationType`` values raise an assertion in the runner. + + ``situ_beta`` / ``situ_linear_beta`` carry the two SiTU soft-caps, and + are ``None`` for every other activation. The runner rejects a mismatch + against ``activation_type`` in either direction. """ tuner = AutoTuner.get() swiglu_limit_scalar = _canonicalize_swiglu_limit_scalar( @@ -3732,7 +3763,9 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( tile_size, scaling_vector_size, activation_type=ActivationType(activation_type), - swiglu_limit_scalar=swiglu_limit_scalar) + swiglu_limit_scalar=swiglu_limit_scalar, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta) inputs = [ input, weight, input_scale, weight_scale, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, @@ -3769,7 +3802,19 @@ def _fake_single_b( scaling_vector_size: int = 16, activation_type: int = int(ActivationType.Swiglu), swiglu_limit_scalar: float = SWIGLU_LIMIT_SCALAR_DISABLED, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: + """Meta-device shapes for the FC1 output and its block scales. + + A gated activation halves the N it emits, so the interleaved + gate/up pair collapses to one value per output element; the extra + ``// 2`` on the tensor itself is NVFP4's two values per byte. + + The activation soft-caps are accepted and ignored: they change what + the kernel computes, never the shape it returns, but the fake must + still mirror the op's schema exactly. + """ m = permuted_idx_to_expanded_idx.size(0) n = weight.size(1) is_gated = is_gated_activation(ActivationType(activation_type)) 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 3c94490c3ad2..c7f87202dd1d 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 @@ -47,10 +47,15 @@ griddepcontrol_launch_dependents, griddepcontrol_wait, is_power_of_2, + sigmoid_f32, silu_f32, ) -SUPPORTED_ACTIVATION_TYPES = (ActivationType.Swiglu, ActivationType.Relu2) +SUPPORTED_ACTIVATION_TYPES = ( + ActivationType.Swiglu, + ActivationType.Relu2, + ActivationType.SiTu, +) def validate_activation_type(activation_type) -> ActivationType: @@ -72,6 +77,8 @@ def validate_activation_type(activation_type) -> ActivationType: Supported fused activations (selected at construction via ``activation_type``): - ActivationType.Swiglu: C = up * silu(gate), where up/gate come from interleaved weight matrix B - ActivationType.Relu2: C = relu(alpha * x)^2 + - ActivationType.SiTu: C = (beta*tanh(gate/beta)*sigmoid(gate)) * (linear_beta*tanh(up/linear_beta)), + gated like Swiglu; requires situ_beta / situ_linear_beta Any other ``ActivationType`` value raises an assertion at construction time. @@ -181,13 +188,16 @@ def validate_activation_type(activation_type) -> ActivationType: class BlockScaledContiguousGatherGroupedGemmKernel: """This class implements contiguous grouped matrix multiplication with gather operation and a - fused activation (selected via ``activation_type``: SwiGLU or Relu2) for FC1 layer computation. + fused activation (selected via ``activation_type``: SwiGLU, Relu2 or SiTU) for FC1 layer + computation. The computation flow: 1. GEMM: acc = alpha * (SFA * A[token_ids]) * (SFB * B) 2. Activation (selected via ``activation_type``): - ActivationType.Swiglu: C = up * silu(gate), from interleaved acc with granularity=64 - ActivationType.Relu2: C = relu(acc)^2 + - ActivationType.SiTu: C = situ_gate * situ_up, from the same interleaved acc as + Swiglu; requires ``situ_beta`` and ``situ_linear_beta`` Any other ``ActivationType`` value raises an assertion in ``__init__``. 3. Optional Quant: When c_dtype is Float4E2M1FN, generates SFC and quantizes output @@ -276,13 +286,16 @@ def __init__( raster_along_m: bool = False, activation_type: ActivationType = ActivationType.Swiglu, swiglu_limit: cutlass.Float32 = float("inf"), + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, ): """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel with gather operation and fused activation. ``activation_type`` accepts a value from ``ActivationType``; only - ``ActivationType.Swiglu`` (gated path) and ``ActivationType.Relu2`` - (non-gated path) are currently supported. + ``ActivationType.Swiglu`` (gated path), ``ActivationType.Relu2`` + (non-gated path) and ``ActivationType.SiTu`` (gated path, Kimi K3) + are currently supported. This configuration includes several key aspects: @@ -316,8 +329,15 @@ def __init__( :param topk: Number of experts selected per token (used for token ID mapping). :type topk: cutlass.Int64 :param activation_type: Fused activation. Must be ``ActivationType.Swiglu`` - (gated, default) or ``ActivationType.Relu2`` (non-gated). + (gated, default), ``ActivationType.Relu2`` (non-gated) or + ``ActivationType.SiTu`` (gated). :type activation_type: ActivationType + :param situ_beta: Gate-side SiTU constant. Required for -- and only + valid with -- ``ActivationType.SiTu``. + :type situ_beta: Optional[float] + :param situ_linear_beta: Linear-side (up) SiTU constant. Required for -- + and only valid with -- ``ActivationType.SiTu``. + :type situ_linear_beta: Optional[float] """ self.sf_vec_size = sf_vec_size @@ -407,6 +427,43 @@ def __init__( self.swiglu_limit = swiglu_limit self.has_swiglu_limit = swiglu_limit != float("inf") + # SiTU constants. They are per-model scalars (not per-expert), so they + # are folded at trace time -- which also means they belong in the + # caller's compiled-kernel cache key. + 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}." + ) + # A chained comparison rather than ``<= 0``: NaN compares false + # against every operand, so ``<= 0`` lets it through, and so does + # a positive infinity. The epilogue folds ``2/beta`` and + # ``2*beta`` at trace time, so either one is baked into the + # compiled kernel and returns quietly wrong activations instead of + # failing. ``math`` here is the MLIR dialect, not Python's, hence + # the bare comparison instead of ``isfinite``. + _INF = float("inf") + if not (0 < situ_beta < _INF and 0 < situ_linear_beta < _INF): + raise ValueError( + "SiTU beta parameters must be finite and positive, got " + f"{situ_beta} and {situ_linear_beta}." + ) + if self.has_swiglu_limit: + # Matches MegaMoE (both backends) and DeepGEMM, which reject + # activation_clamp together with SiTU. + raise ValueError( + "ActivationType.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) + def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -2373,6 +2430,9 @@ def kernel( if cutlass.const_expr(self.activation_type == ActivationType.Swiglu): acc_vec_gate = tTR_rAcc_gate.load() self._apply_swiglu_epilogue(acc_vec_up, acc_vec_gate, alpha_val, tCompute) + elif cutlass.const_expr(self.activation_type == ActivationType.SiTu): + acc_vec_gate = tTR_rAcc_gate.load() + self._apply_situ_epilogue(acc_vec_up, acc_vec_gate, alpha_val, tCompute) elif cutlass.const_expr(self.activation_type == ActivationType.Relu2): self._apply_relu2_epilogue(acc_vec_up, alpha_val, tCompute) @@ -2661,6 +2721,99 @@ def _apply_swiglu_epilogue( 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 ``kimi_k3_moe/_mlp.py::SituAndMul`` + (itself byte-identical to HF ``modeling_kimi.py``):: + + 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 + + ``up`` and ``gate`` come from the two interleaved accumulator subtiles + loaded by the caller, same as the SwiGLU epilogue. + + There is no packed tanh, so the vectorized path uses the identity + ``tanh(z) = 2 * sigmoid(2z) - 1`` (the same one ``utils.gelu_tanh_f32`` + uses) to stay on the packed f32x2 path -- calling a scalar tanh would + force the whole loop back to scalar. The reciprocals and ``2*beta`` + factors fold at trace time because both betas are ``const_expr``:: + + beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) + = 2*beta*sigmoid((2/beta)*x) - beta + """ + 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): + """``sigmoid`` on a packed f32x2 pair, as ``rcp(1 + exp2(-x*log2e))``. + + Called three times per element pair -- once for the gate's own + sigmoid and once inside each ``tanh`` -- so it is written + against the packed intrinsics rather than reused from the + scalar helpers, which would unpack and repack every call. + """ + 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 + @cute.jit def _apply_relu2_epilogue( self, @@ -3342,8 +3495,8 @@ def wrapper( """Single-B wrapper. ``l`` is the number of experts in the (sole) B tensor. ``activation_type`` - must match the one passed to ``__init__``; only ``Swiglu`` and ``Relu2`` - are supported. + must match the one passed to ``__init__``; only ``Swiglu``, ``Relu2`` + and ``SiTu`` are supported. """ is_gated = is_gated_activation(activation_type) scale_k = k // scaling_vector_size diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 38dad076d57f..076f97519b6b 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1069,6 +1069,16 @@ def __init__( layer_idx: int, aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], ): + """Build the routed experts and the shared expert for one MoE layer. + + ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: + the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields + that ``ModelConfig`` does not carry. + + ``aux_stream_dict`` is shared across every layer of the model, so the + streams reached through it are borrowed and must not be synchronized + or reassigned here. + """ super().__init__() self.layer_idx = layer_idx self.hidden_size = cfg.hidden_size @@ -1125,12 +1135,20 @@ def __init__( gate_softcap=situ_beta, linear_softcap=situ_linear_beta, ), - # A MegaMoE request that silently degraded to CUTLASS would be - # benchmarked as if it were MegaMoE, and the decline is easy to - # trigger (EP-only, own token / top-k limits). Fail in the resolver + # A request that silently degraded to CUTLASS would be benchmarked + # as if it were the backend that was asked for, and the decline is + # easy to trigger: MegaMoE has its own token / top-k limits and is + # EP-only, and CuteDSL declines on activation shape, SM version and + # the CuTe DSL dependency. Measured 2026-09-08: a CUTEDSL request + # was turned down on every one of the 92 MoE layers, on all 16 + # ranks, and still produced correct text and a zero exit -- the + # only trace was a warning line per layer. Fail in the resolver # instead, which reports the rejection trail. + # + # 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"), + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), ) self._check_trtllm_situ_quant( routed_moe_model_config.moe_backend, routed_quant_config.quant_algo @@ -1313,16 +1331,20 @@ def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: """Build a private routed-expert mapping without mutating the shared config. Default split is EP-only; see ``_select_moe_tp_ep``.""" + # Every backend here declares ``ActivationType.SiTu`` in its + # ``activation_support``; the list is not a preference order. CUTEDSL + # joined once its act-fusion kernel grew the SiTU epilogue. supported_backends = { "CUTLASS", "TRTLLM", + "CUTEDSL", "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, " - "MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " + "CUTEDSL, 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/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py index 0a956593a8be..521ff87704f3 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py @@ -604,12 +604,30 @@ class CuteDslFusedMoE(MoEImplBase): input_requirement = MoEInputRequirement(routing_scales_dtype=torch.float32) - # Kinds mirror the kernel's own SUPPORTED_ACTIVATION_TYPES in - # cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py. + # Kinds mirror the act-fusion kernels' own SUPPORTED_ACTIVATION_TYPES. + # There are two of them and ``run_moe_nvfp4`` picks between them by SM: + # cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.py + # cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py + # They agreed until SiTu, which only the Blackwell one implements. This + # attribute is the union, because resolution reads it off the class and a + # class attribute cannot see which kernel an instance will dispatch to; + # ``can_implement`` narrows it back for the SM that lacks the epilogue. + # # The clamp is a kernel-cache-key scalar and the epilogue has no # "clamp absent" branch, so an absent clamp is +inf, not None. + # + # alpha/beta are declared here rather than narrowed per instance: + # ``moe_resolution._activation_rejection`` states the invariant -- an + # instance may narrow a shape, never admit one its class refuses. + # Declaring UNSUPPORTED here and widening per instance made every K3 layer + # resolve away to CUTLASS with "CuteDslFusedMoE kernels take no activation + # alpha". Safe for the other two kinds because neither supplies the pair: + # ``SwigluActivation.constants()`` fills only ``limit`` and Relu2 fills + # nothing. ``SwigluBias`` is the kind that does, and it is not in ``kinds``. activation_support = MoEActivationSupport( - kinds=frozenset({ActivationType.Swiglu, ActivationType.Relu2}), + kinds=frozenset( + {ActivationType.Swiglu, ActivationType.Relu2, ActivationType.SiTu}), + alpha_beta=ActivationParamShape.UNIFORM_SCALAR, limit=ActivationParamShape.UNIFORM_SCALAR, limit_when_absent=float("inf"), ) @@ -734,6 +752,16 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: MoERejectReason.DEP_MISSING, "NVFP4 CuteDSL MoE on SM107 requires Rubin support in CuTe DSL" ) + # ``activation_support`` above is the union over the act-fusion + # kernels ``run_moe_nvfp4`` picks between; only one of them has a + # SiTU epilogue. Turn the layer down where the other one would be + # chosen, so it stays resolvable by another backend instead of + # raising inside the kernel at forward time. + if p.activation == "SiTu" and sm_version == 107: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "this device's CuteDSL act-fusion kernel has no SiTU " + "epilogue") # process_weights_after_loading() unswizzles the FC1 block scales, # which asserts 128-row tiles; without this gate an unaligned shard # dies mid weight load with a bare swizzle error. @@ -955,9 +983,10 @@ def run_moe_nvfp4( assert self.has_nvfp4 assert weight_view is not None if self.activation_type not in (ActivationType.Swiglu, - ActivationType.Relu2): + ActivationType.Relu2, + ActivationType.SiTu): raise NotImplementedError( - "CuteDSL NVFP4 FC1 supports only SwiGLU and Relu2; " + "CuteDSL NVFP4 FC1 supports only SwiGLU, Relu2 and SiTU; " f"got {self.activation_type.name}") output_dtype = torch.bfloat16 @@ -1113,6 +1142,15 @@ def run_moe_nvfp4_impl( gather_act_kwargs["activation_type"] = self.activation_type gather_act_kwargs["swiglu_limit_scalar"] = self.act_clamp gather_act_kwargs["activation_type"] = self.activation_type + # ``act_alpha`` / ``act_beta`` are where ``SiTuActivation.constants()`` + # lands: gate_softcap -> alpha, linear_softcap -> beta, both reduced to + # a uniform scalar by the shape this backend declares. Only forwarded + # for SiTU so every other activation keeps the op's ``None`` default -- + # passing them unconditionally would make the op signature lie about + # which kinds have soft-caps. + if self.activation_type == ActivationType.SiTu: + gather_act_kwargs["situ_beta"] = self.act_alpha + gather_act_kwargs["situ_linear_beta"] = self.act_beta x, x_sf = gather_act_op(**gather_act_kwargs) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 08feca1018fa..a7b21270f5c3 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -72,8 +72,18 @@ l0_b300: - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_drains_staging_per_expert - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streamed_experts_forward_runs - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutlass_situ_bf16_matches_reference - - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ[CUTLASS] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ[CUTEDSL] - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ[TRTLLM] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_kimi_k3_allow_list_matches_what_the_backends_declare + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_explicit_cutedsl_fails_instead_of_degrading_to_cutlass + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[nan_beta] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[nan_linear] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[inf_beta] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[inf_linear] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[zero] + - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_cutedsl_kernel_rejects_unusable_situ_betas[negative] + - unittest/_torch/moe/test_moe_backend.py::test_situ_survives_resolution_not_just_construction[cutlass] + - unittest/_torch/moe/test_moe_backend.py::test_situ_survives_resolution_not_just_construction[cutedsl] - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_experts_match_situ_reference[CUTLASS-static_1.0] - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_nvfp4_experts_match_situ_reference[TRTLLM-static_1.0] - unittest/_torch/moe/test_kimi_k3_situ_moe.py::test_megamoe_streamed_coverage_survives_per_expert_drain 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 57ecf4ae12ea..19e692a13df5 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -122,6 +122,45 @@ def test_kimi_situ_betas_must_be_positive(situ_beta, situ_linear_beta): 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")), + (0.0, 25.0), + (-2.0, 25.0), + ], + ids=["nan_beta", "nan_linear", "inf_beta", "inf_linear", "zero", "negative"], +) +def test_cutedsl_kernel_rejects_unusable_situ_betas(situ_beta, situ_linear_beta): + """NaN and infinity must be refused, not just zero and negatives. + + The betas fold into the kernel at trace time as ``2/beta`` and + ``2*beta``, so a non-finite one is compiled in and comes back as quietly + wrong activations. ``<= 0`` does not catch either: every comparison + against NaN is false, and an infinity is positive. + """ + pytest.importorskip("cutlass") + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.blockscaled_contiguous_gather_grouped_gemm_act_fusion import ( # noqa: E501 + BlockScaledContiguousGatherGroupedGemmKernel, + ) + from tensorrt_llm._torch.utils import ActivationType + + with pytest.raises(ValueError, match="finite and positive"): + BlockScaledContiguousGatherGroupedGemmKernel( + sf_vec_size=16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + vectorized_f32=True, + topk=8, + activation_type=ActivationType.SiTu, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + + def test_clear_checkpoint_fp8_pairs_releases_unconsumed_stashes(): linear = torch.nn.Linear(2, 2, bias=False) setattr(linear.weight, modeling_kimi_linear._K3_CKPT_FP8_ATTR, (torch.ones(1), torch.ones(1))) @@ -591,6 +630,96 @@ def test_kimi_k3_routed_config_logs_megamoe_capacity_override(monkeypatch): ) +def test_kimi_k3_allow_list_matches_what_the_backends_declare(): + """The K3 allow-list must be *exactly* the backends that execute SiTU. + + Both sides are derived, neither is written down here. The expected set + comes from each family's own ``activation_support``; the actual set comes + from asking the allow-list. A hand-maintained second copy of a capability + set is the defect this module exists to catch -- it is how CUTEDSL stayed + shut for weeks after its kernel grew the epilogue -- and a test that + restates the list has the same defect. + + Checking both directions matters: inclusion alone would pass while the + allow-list offered a backend that cannot serve SiTU, which resolves and + then fails at construction instead of being declined. + + ``any`` rather than ``all`` over a family: resolution walks family members + in ``IMPL_PRIORITY`` order, so a family serves SiTU when one member does. + ``CUTEDSL`` is exactly that case -- ``CuteDslB12xFusedMoE`` does not + declare it, ``CuteDslFusedMoE`` does. + """ + from tensorrt_llm._torch.moe.fused_moe.moe_resolution import BACKEND_FAMILY + from tensorrt_llm._torch.utils import ActivationType + + def admits_situ(name): + model_config = ModelConfig( + mapping=Mapping(world_size=1, rank=0, tp_size=1), + moe_backend=name, + ) + try: + KimiK3MoERuntime._routed_moe_model_config(model_config) + except ValueError: + return False + return True + + declares_situ = { + name + for name, family in BACKEND_FAMILY.items() + if any(ActivationType.SiTu in cls.activation_support.kinds for cls in family) + } + allow_listed = {name for name in BACKEND_FAMILY if admits_situ(name)} + + assert allow_listed == declares_situ, ( + "Kimi K3's routed-expert allow-list disagrees with the backends' own " + f"activation_support.\n" + f" offered but cannot serve SiTU: {sorted(allow_listed - declares_situ)}\n" + f" serves SiTU but not offered: {sorted(declares_situ - allow_listed)}" + ) + # The set is derived, so guard against it being derived as empty -- that + # would satisfy the equality above while testing nothing. + assert "CUTEDSL" in declares_situ + + +def test_explicit_cutedsl_fails_instead_of_degrading_to_cutlass(): + """An ineligible CUTEDSL request must raise, not silently pick CUTLASS. + + This is the failure this branch was built around: CUTEDSL declined every + K3 layer on all 16 ranks, CUTLASS took over, and the run produced correct + text and a zero exit while being attributed to CUTEDSL. K3 therefore + passes ``allow_backend_degradation=False`` for it; the assertion here is + that the resolver honours that rather than that K3 sets it. + + The ``allow_degradation=True`` half is the control -- without it, a + resolver that had stopped substituting altogether would pass the first + half for the wrong reason. + """ + from tensorrt_llm._torch.moe.fused_moe.moe_resolution import resolve_moe_impl + + # SM90 has no NVFP4 CuteDSL path at all, so the request is declined for a + # reason that does not depend on this branch's activation work. + model_config = ModelConfig( + mapping=Mapping(world_size=1, rank=0, tp_size=1), + moe_backend="CUTEDSL", + ) + kwargs = dict( + num_experts=8, + hidden_size=512, + intermediate_size=512, + dtype=torch.bfloat16, + ) + + degraded = resolve_moe_impl(model_config, allow_degradation=True, **kwargs) + if not degraded.degraded: + pytest.skip("CUTEDSL is eligible on this device; nothing to degrade from") + + with pytest.raises(ValueError) as excinfo: + resolve_moe_impl(model_config, allow_degradation=False, **kwargs) + # The trail, not just the refusal: without it the caller cannot tell which + # gate declined and has to re-derive it. + assert "CUTEDSL" in str(excinfo.value) + + def test_kimi_k3_routed_config_rejects_backend_without_situ_support(): model_config = ModelConfig( mapping=Mapping(world_size=1, rank=0, tp_size=1), @@ -1221,6 +1350,39 @@ 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``: +#: 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"] + + +def _skip_if_backend_unavailable(moe_backend): + """Skip a CUTEDSL parametrization when the CuTe DSL wheel is absent. + + Probed here rather than in a ``pytest.mark.skipif``, because the marker + is evaluated at collection time and importing ``cute_dsl_utils`` that + early puts the wheel's package directory on ``sys.path`` for every other + test file in the session. + """ + if moe_backend != "CUTEDSL": + return + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + if not IS_CUTLASS_DSL_AVAILABLE: + pytest.skip("CuteDSL MoE requires the CuTe DSL wheel") + # nvfp4_moe_supported admits every SM >= 100, but only the Blackwell + # act-fusion kernel carries the SiTU epilogue. Elsewhere -- SM107 included + # -- can_implement() declines SiTu, so without this the case would fail in + # backend resolution on hardware it was never claimed to support. + if get_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()}" + ) + + def _make_nvfp4_expert_bank(num_experts, intermediate, hidden, seed=907): """Random NVFP4 tensors in ``nvidia/Kimi-K3-NVFP4`` checkpoint layout.""" gen = torch.Generator().manual_seed(seed) @@ -1702,11 +1864,18 @@ def test_nvfp4_experts_match_situ_reference(moe_backend, input_scale): assert cosine > 0.95, f"cosine={cosine.item()}, rel_l2={rel_l2.item()}" -def _swiglu_reference_moe(x, router_logits, routing_method, w1, w2, w3, alpha, beta): - """Same routing/geometry as the SiTU reference, but CUTLASS's SwigluBias. +def _plain_swiglu_reference_moe(x, router_logits, routing_method, w1, w2, w3): + """Same routing/geometry as the SiTU reference, but plain SwiGLU. + + ``up * silu(gate)`` -- what the FC1 epilogue actually computes when it + does not run SiTU. This is the realistic wrong answer, so it is the one + worth measuring against: an earlier revision compared with SwigluBias + (``gate*sigmoid(gate*alpha)*(up+beta)``), which no epilogue on this path + computes, so a genuine SwiGLU fallback sat far from both references and + still satisfied a "closer to SiTU" ordering. - ``gate*sigmoid(gate*alpha)*(linear+beta)`` -- what the FC1 epilogue would - compute if the activation enum did not resolve to SiTu. + The dequant alpha is already folded into ``g`` and ``u`` here, exactly as + the kernel applies it to the accumulator before either epilogue. """ ids, weights = routing_method.apply(router_logits) out = torch.zeros_like(x, dtype=torch.float32) @@ -1716,13 +1885,13 @@ def _swiglu_reference_moe(x, router_logits, routing_method, w1, w2, w3, alpha, b e = int(ids[token, slot]) g = xf[token] @ w1[e].float().t() u = xf[token] @ w3[e].float().t() - h = g * torch.sigmoid(g * alpha) * (u + beta) + h = u * (g * torch.sigmoid(g)) out[token] += float(weights[token, slot]) * (h @ w2[e].float().t()) return out @nvfp4_moe_supported -@pytest.mark.parametrize("moe_backend", ["CUTLASS", "TRTLLM"]) +@pytest.mark.parametrize("moe_backend", _NVFP4_SITU_BACKENDS) def test_nvfp4_kernel_actually_applies_situ(moe_backend): """Which activation does the QUANTIZED kernel actually run? @@ -1734,16 +1903,23 @@ def test_nvfp4_kernel_actually_applies_situ(moe_backend): wrong in exactly the way the GSM8K collapse showed, while every shape-and-buffer check stayed green. - Run for both FP4 backends: CUTLASS resolves SiTU through the activation - enum, TRTLLM-Gen through a distinct fused-cubin family - (``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``). A silent fallback is a different - failure on each, and this comparison is tolerance-free, so it catches - both -- including the degenerate all-zero FC1 output, which scores 0 - against both references and so fails the assertion below. + 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. + + The comparison is tolerance-free, so it also catches the degenerate + all-zero FC1 output, which scores 0 against both references. Reported rather than merely asserted: which reference the kernel is closer to is the diagnosis. """ + _skip_if_backend_unavailable(moe_backend) num_experts, hidden, inter = _TP_EXPERTS, _TP_HIDDEN, _TP_INTERMEDIATE gate = _make_test_gate(num_experts=num_experts) @@ -1773,9 +1949,7 @@ def test_nvfp4_kernel_actually_applies_situ(moe_backend): situ = _situ_reference_moe( x, router_logits, gate.routing_method, w1, w2, w3, beta=4.0, linear_beta=25.0 ) - swiglu = _swiglu_reference_moe( - x, router_logits, gate.routing_method, w1, w2, w3, alpha=4.0, beta=25.0 - ) + swiglu = _plain_swiglu_reference_moe(x, router_logits, gate.routing_method, w1, w2, w3) def score(ref): cos = torch.nn.functional.cosine_similarity(actual.flatten(), ref.flatten(), dim=0) @@ -1784,15 +1958,17 @@ def score(ref): situ_cos, situ_l2 = score(situ) swiglu_cos, swiglu_l2 = score(swiglu) + # Printed, not yet asserted on: the bounds below are set from measured + # spread rather than guessed, so the numbers have to be visible first. print( - f"NVFP4[{moe_backend}] kernel vs SiTU ref: " + f"NVFP4[{moe_backend}] kernel vs SiTU ref: " f"cosine={situ_cos:.6f} rel_l2={situ_l2:.6f}\n" - f"NVFP4[{moe_backend}] kernel vs SwiGLU ref: " + f"NVFP4[{moe_backend}] kernel vs plain SwiGLU ref: " f"cosine={swiglu_cos:.6f} rel_l2={swiglu_l2:.6f}" ) assert situ_cos > swiglu_cos, ( - f"the NVFP4 {moe_backend} kernel matches a SwiGLU reference better than " - f"the SiTU one (situ={situ_cos:.6f}, swiglu={swiglu_cos:.6f}): the " + f"the NVFP4 {moe_backend} kernel matches a plain SwiGLU reference better " + f"than the SiTU one (situ={situ_cos:.6f}, swiglu={swiglu_cos:.6f}): the " f"quantized path is not applying SiTU" ) diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index f23699e3ce70..a89626f57cba 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -14,6 +14,7 @@ # limitations under the License. """MoE backend unit tests.""" +import dataclasses import importlib import itertools import logging @@ -84,6 +85,7 @@ from tensorrt_llm._torch.moe.fused_moe.interface import MoE, MoESchedulerKind, MoEWeightLoadingMode from tensorrt_llm._torch.moe.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm from tensorrt_llm._torch.moe.fused_moe.moe_resolution import ( + _reject_unsupported_activation, build_moe_deployment, impl_class_for, resolve_moe_impl, @@ -2541,6 +2543,30 @@ def test_nvfp4_fc1_row_alignment_gate( assert verdict.reject_reason is not MoERejectReason.SHAPE_UNALIGNED +@pytest.mark.parametrize( + "backend_cls", + [CutlassFusedMoE, CuteDslFusedMoE], + ids=["cutlass", "cutedsl"], +) +def test_situ_survives_resolution_not_just_construction(backend_cls): + """A SiTU layer must be admitted by the *resolver*, not only build. + + Every other SiTU test constructs a backend directly and so never consults + ``activation_support``. Resolution does, and it reads the **class** + attribute, because it judges candidates before any instance exists. A + backend that declared its alpha/beta shape per instance instead passed + every unit test and then resolved away to CUTLASS on real hardware with + "kernels take no activation alpha, which this layer's SiTu supplies" -- + silently, because Kimi K3 permits degradation for this backend. + """ + problem = dataclasses.replace( + _nvfp4_problem(2048, "SiTu"), + activation_constants=frozenset({"alpha", "beta"}), + ) + rejection = _reject_unsupported_activation(backend_cls, problem) + assert rejection is None, f"{backend_cls.__name__} refuses SiTU at resolution: {rejection}" + + def test_unresolvable_layer_error_carries_rejection_details(): """describe() prints reason codes only, so impl_class_for has to add the details -- without them a shape rejection reaches the operator as a bare