From dd56dc24a5a3bbc5218f8ac0a5a565107f193278 Mon Sep 17 00:00:00 2001
From: Xin Guan <294044352+xguannv@users.noreply.github.com>
Date: Thu, 10 Sep 2026 02:02:07 +0000
Subject: [PATCH 1/8] [None][feat] CuteDSL Blackwell act-fusion: add the SiTU
 epilogue

Kimi K3's routed experts use SiTU, which CUTLASS, TRTLLM-Gen and both
MegaMoE backends already implement. The CuteDSL grouped-GEMM act-fusion
kernel did not, so its SUPPORTED_ACTIVATION_TYPES stopped at
(Swiglu, Relu2) and a K3 layer could never reach it.

    situ_gate = beta        * tanh(g / beta) * sigmoid(g)
    situ_up   = linear_beta * tanh(u / linear_beta)

is the same expression the MegaMoE CuteDSL kernel and the CUTLASS
SiTuAdaptor evaluate, so the backends stay numerically comparable.

Two constants, not one: the branches are soft-capped independently and
both must be positive because the kernel divides by them. They are
per-model scalars, so they fold at trace time and belong in the
compiled-kernel cache key.

There is no packed tanh intrinsic, so the vectorized path uses
tanh(z) = 2*sigmoid(2z) - 1 -- the identity utils.gelu_tanh_f32 already
uses -- to stay on the packed f32x2 path.

A SwiGLU clamp is rejected alongside SiTU rather than silently ignored,
matching what both MegaMoE backends and DeepGEMM do.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>
---
 ...ntiguous_gather_grouped_gemm_act_fusion.py | 148 +++++++++++++++++-
 1 file changed, 142 insertions(+), 6 deletions(-)

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..f76fc6feebfa 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.
 
@@ -276,13 +283,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 +326,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 +424,35 @@ 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}."
+                )
+            if situ_beta <= 0 or situ_linear_beta <= 0:
+                raise ValueError(
+                    "SiTU beta parameters must be 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 +2419,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 +2710,93 @@ 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)
+
+            # sigmoid(x) = rcp(1 + exp2(-x * log2e)), shared by both cores.
+            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
+
     @cute.jit
     def _apply_relu2_epilogue(
         self,
@@ -3342,8 +3478,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

From 1e2755afe4480f659eae819590f7c258c67c563c Mon Sep 17 00:00:00 2001
From: Xin Guan <294044352+xguannv@users.noreply.github.com>
Date: Thu, 10 Sep 2026 02:03:36 +0000
Subject: [PATCH 2/8] [None][feat] CuteDSL MoE: carry the SiTU soft-caps to the
 act-fusion op

The Blackwell act-fusion kernel now implements SiTU, but nothing could
reach it: the op in between had no parameters for the soft-caps.

Both betas are trace-time constants -- the kernel folds them rather than
reading them from memory -- so they are added to three places, not one:
the runner constructor, unique_id, and the compile cache key. Missing
either of the latter two would let a layer silently reuse a kernel
compiled for different soft-caps, a wrong-numbers bug with no error
attached to it.

The op boundary carries SITU_BETA_DISABLED = -1.0 rather than None. Zero
is not usable as the neutral value because the epilogue divides by both
betas, and a negative soft-cap is already impossible (SiTuActivation
rejects it at construction), so the negative range is free to reserve.
_canonicalize_situ_beta maps the sentinel back to None at the runner
boundary, mirroring _canonicalize_swiglu_limit_scalar directly above it.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  | 58 ++++++++++++++++---
 1 file changed, 50 insertions(+), 8 deletions(-)

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..c55fc8310ebf 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -79,6 +79,20 @@ def _canonicalize_swiglu_limit_scalar(swiglu_limit_scalar: float) -> float:
     return float("inf") if swiglu_limit_scalar < 0 else swiglu_limit_scalar
 
 
+#: Sentinel for "this layer is not SiTU". A torch custom op schema cannot carry
+#: ``Optional[float]`` here the way a Python signature can, and 0.0 is not
+#: available as the neutral value: the SiTU epilogue divides by both betas, so
+#: zero is a division by zero rather than a no-op. A negative value is
+#: impossible for a real soft-cap -- ``SiTuActivation`` rejects it at
+#: construction -- which makes it safe to reserve.
+SITU_BETA_DISABLED = -1.0
+
+
+def _canonicalize_situ_beta(situ_beta: float) -> Optional[float]:
+    """Map the op-boundary sentinel back to ``None`` for the kernel."""
+    return None if situ_beta is None or situ_beta <= 0 else float(situ_beta)
+
+
 def _get_cute_dsl_swap_ab_candidates(
     m: int,
     output_aligned: bool,
@@ -3371,13 +3385,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 +3413,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(
@@ -3414,6 +3440,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 +3644,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 +3657,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 +3744,19 @@ 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: float = SITU_BETA_DISABLED,
+        situ_linear_beta: float = SITU_BETA_DISABLED,
     ) -> 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.
+        They default to ``SITU_BETA_DISABLED`` rather than ``None`` because the
+        op schema takes plain floats; the runner maps the sentinel back to
+        ``None`` and then rejects a mismatch against ``activation_type``.
         """
         tuner = AutoTuner.get()
         swiglu_limit_scalar = _canonicalize_swiglu_limit_scalar(
@@ -3732,7 +3770,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=_canonicalize_situ_beta(situ_beta),
+            situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta))
         inputs = [
             input, weight, input_scale, weight_scale, alpha,
             tile_idx_to_group_idx, tile_idx_to_mn_limit,
@@ -3769,6 +3809,8 @@ 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: float = SITU_BETA_DISABLED,
+        situ_linear_beta: float = SITU_BETA_DISABLED,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         m = permuted_idx_to_expanded_idx.size(0)
         n = weight.size(1)

From b0f0b53b799ec16a5d0266e7997ba06a1531bf8c Mon Sep 17 00:00:00 2001
From: Xin Guan <294044352+xguannv@users.noreply.github.com>
Date: Thu, 10 Sep 2026 02:32:27 +0000
Subject: [PATCH 3/8] [None][feat] Kimi K3: unlock the CUTEDSL MoE backend for
 NVFP4

With the Blackwell act-fusion kernel carrying the SiTU epilogue and the
op carrying the soft-caps, what remains is to stop refusing the
activation and to hand the constants over.

CuteDslFusedMoE declares the alpha/beta pair on the *class*. Resolution
reads the class attribute -- it has to, since it judges candidates before
any instance exists -- and moe_resolution._activation_rejection states
the invariant outright: an instance may narrow a shape, never admit one
its class refuses. An earlier revision of this branch had it backwards,
declaring UNSUPPORTED on the class and widening per instance, and every
K3 layer was turned down on all 16 ranks with "CuteDslFusedMoE kernels
take no activation alpha, which this layer's SiTu supplies". Kimi K3
permitted degradation for this backend, so the run produced correct text
and exited zero while running CUTLASS.

Declaring the pair on the class is safe for the other two kinds this
backend executes: SwigluActivation.constants fills only limit and Relu2
fills nothing. SwigluBias is the kind that fills alpha/beta, and it is
not in kinds.

run_moe_nvfp4 admits SiTu and forwards act_alpha / act_beta as the two
betas -- that is where SiTuActivation.constants() lands, reduced to
uniform scalars by the declared shape. They are forwarded only for SiTU
so every other kind keeps hitting the op default. The other three
activation gates in the file are deliberately untouched: the unquantized
BF16 method interleaves FC1 weights for a kernel that fuses SwiGLU by
name, the locality-domain half-GEMM has no SiTU parameters on its op, and
the FP8 block-scale path evaluates SwiGLU in Python.

SiTU is turned down on SM107. run_moe_nvfp4 dispatches to the Rubin
act-fusion kernel there, whose SUPPORTED_ACTIVATION_TYPES is still
(Swiglu, Relu2); reaching it would trip an assert inside the kernel
instead of resolving to another backend.

CUTEDSL also joins the list of backends whose K3 request must not degrade
silently. That list already held both MegaMoE backends for the reason
above; CuteDSL declines for more causes than they do (activation shape,
SM version, the CuTe DSL dependency), and the failure described above is
what an unnoticed decline looks like.

Also correct the NVFP4 eval recipe's comment, which still told readers
CUTLASS was required because trtllm-gen served MXFP4 only. That stopped
being true in #17940 and the guard repeating it was removed in #18709.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>
---
 .../eval_extra_llm_options_nvfp4_dep16.yaml   |  9 +++-
 .../_torch/models/modeling_kimi_linear.py     | 22 +++++++--
 .../moe/fused_moe/fused_moe_cute_dsl.py       | 48 +++++++++++++++++--
 3 files changed, 67 insertions(+), 12 deletions(-)

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/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 38dad076d57f..957ff1e3613a 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -1125,12 +1125,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 +1321,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..d99be994de31 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 hitting the op's sentinel
+        # 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)
 

From ceb1ba033838f7be3be9f3b62123396942e2420a Mon Sep 17 00:00:00 2001
From: Xin Guan <294044352+xguannv@users.noreply.github.com>
Date: Thu, 10 Sep 2026 02:32:41 +0000
Subject: [PATCH 4/8] [None][test] cover CUTEDSL SiTU at the kernel and at the
 resolver

Three additions, each aimed at a way this backend can fail quietly.

test_nvfp4_kernel_actually_applies_situ gains CUTEDSL. The other two
backends fail loudly if SiTU goes missing -- CUTLASS has no matching
activation enum, TRTLLM-Gen no matching cubin. CuteDSL does not: its
soft-caps are trace-time scalars folded into a JIT-compiled epilogue, so
dropping them compiles a SwiGLU kernel and returns plausible numbers.
Scoring the output against a SiTU reference and a SwiGLU reference, with
no tolerance, is the only way to tell those apart.

test_situ_survives_resolution_not_just_construction calls
_activation_rejection directly. Every other SiTU test constructs a
backend and so never consults activation_support; resolution does, and it
reads the class attribute. A per-instance declaration passed the entire
unit suite and then resolved away to CUTLASS on hardware. CUTLASS is
parametrized alongside so the next backend to grow SiTU inherits it.

test_kimi_k3_allow_list_matches_what_the_backends_declare asserts the
model-layer allow-list against each backend's own activation_support
rather than a literal list. A hand-maintained second copy of a capability
set is what #18709 had to fix; this keeps a new one from forming.

The CUTEDSL parameter probes the CuTe DSL wheel inside the test body
rather than in a skipif: importing cute_dsl_utils at module scope pulls
in a package that appends its own directory to sys.path, which this
repository's magic_import hooks reject at session level.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>
---
 .../_torch/moe/test_kimi_k3_situ_moe.py       | 73 +++++++++++++++++--
 tests/unittest/_torch/moe/test_moe_backend.py | 26 +++++++
 2 files changed, 92 insertions(+), 7 deletions(-)

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..413c42e5d555 100644
--- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py
+++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py
@@ -591,6 +591,41 @@ 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.
+
+    Asserting agreement with each backend's own ``activation_support``
+    rather than against a literal list: a hand-maintained second copy of a
+    capability set is the defect this module already exists to catch, and
+    it is how CUTEDSL stayed shut after its kernels grew the epilogue.
+    """
+    from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE
+    from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE
+    from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE
+    from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_cute_dsl import MegaMoECuteDsl
+    from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_deepgemm import MegaMoEDeepGemm
+    from tensorrt_llm._torch.utils import ActivationType
+
+    declares_situ = {
+        "CUTLASS": CutlassFusedMoE,
+        "TRTLLM": TRTLLMGenFusedMoE,
+        "CUTEDSL": CuteDslFusedMoE,
+        "MEGAMOE_CUTEDSL": MegaMoECuteDsl,
+        "MEGAMOE_DEEPGEMM": MegaMoEDeepGemm,
+    }
+    for name, cls in declares_situ.items():
+        assert ActivationType.SiTu in cls.activation_support.kinds, (
+            f"{name} lost SiTU from activation_support; the K3 allow-list "
+            "still offers it, so a layer would resolve here and then fail."
+        )
+        model_config = ModelConfig(
+            mapping=Mapping(world_size=1, rank=0, tp_size=1),
+            moe_backend=name,
+        )
+        # Must not raise: this is the gate that kept CUTEDSL out.
+        KimiK3MoERuntime._routed_moe_model_config(model_config)
+
+
 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 +1256,23 @@ 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):
+    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")
+
+
 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)
@@ -1722,7 +1774,7 @@ def _swiglu_reference_moe(x, router_logits, routing_method, w1, w2, w3, alpha, b
 
 
 @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 +1786,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)
 
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

From 8a0795b93eea485b30c5496974e2aacd91362aa4 Mon Sep 17 00:00:00 2001
From: Xin Guan <294044352+xguannv@users.noreply.github.com>
Date: Thu, 10 Sep 2026 02:32:54 +0000
Subject: [PATCH 5/8] [None][chore] Kimi K3 quickstart: let it select a MoE
 backend

The shipped example took --model and --image only, so there was no way to
exercise a backend other than the resolver's default -- including the one
this series opens. Add --moe-backend to both the Python entry point and
the sbatch wrapper that forwards to it.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>
---
 examples/kimi_k3/quick_start_kimi_k3.py     | 13 +++++++++++--
 examples/kimi_k3/quick_start_kimi_k3.sbatch | 17 +++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/examples/kimi_k3/quick_start_kimi_k3.py b/examples/kimi_k3/quick_start_kimi_k3.py
index c1db17bc9272..520638cf6e62 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"),
@@ -45,12 +45,18 @@ 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:
     args = parse_arguments()
-    llm = LLM(
+    llm_kwargs = dict(
         model=args.model,
         tensor_parallel_size=args.tp_size,
         enable_attention_dp=True,
@@ -72,6 +78,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..7308d4362e58 100644
--- a/examples/kimi_k3/quick_start_kimi_k3.sbatch
+++ b/examples/kimi_k3/quick_start_kimi_k3.sbatch
@@ -40,6 +40,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 +105,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}

From a89c8a9176285d6a84b746bf477df5da849fe7ac Mon Sep 17 00:00:00 2001
From: Xin Guan <294044352+xguannv@users.noreply.github.com>
Date: Thu, 10 Sep 2026 08:52:35 +0000
Subject: [PATCH 6/8] [None][chore] Docstring the functions this branch touches

Docstring coverage over the diff was 72%, under the 80% gate. Seven
functions were missing one; each now says the thing a reader could not
recover from the signature.

The two that carry real information:

unique_id() lists trace-time constants, which is why the activation
soft-caps belong in it -- they are folded into the compiled kernel as
const_expr, so two runners differing in a beta are different kernels and
must not share a tuning result. Omitting them would silently serve one
layer's kernel to a layer with different soft-caps.

_skip_if_backend_unavailable() probes at call time rather than through
pytest.mark.skipif, because the marker is evaluated during collection and
importing cute_dsl_utils that early puts the CuTe DSL wheel's package
directory on sys.path for every other test file in the session.

No behaviour change.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>
---
 examples/kimi_k3/quick_start_kimi_k3.py        | 12 ++++++++++++
 .../_torch/custom_ops/cute_dsl_custom_ops.py   | 18 ++++++++++++++++++
 ...ontiguous_gather_grouped_gemm_act_fusion.py |  8 +++++++-
 .../_torch/models/modeling_kimi_linear.py      | 10 ++++++++++
 .../_torch/moe/test_kimi_k3_situ_moe.py        |  7 +++++++
 5 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/examples/kimi_k3/quick_start_kimi_k3.py b/examples/kimi_k3/quick_start_kimi_k3.py
index 520638cf6e62..5854ae28f6f9 100644
--- a/examples/kimi_k3/quick_start_kimi_k3.py
+++ b/examples/kimi_k3/quick_start_kimi_k3.py
@@ -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",
@@ -55,6 +61,12 @@ def parse_arguments() -> argparse.Namespace:
 
 
 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_kwargs = dict(
         model=args.model,
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 c55fc8310ebf..77f0b1373700 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -3431,6 +3431,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,
@@ -3812,6 +3820,16 @@ def _fake_single_b(
         situ_beta: float = SITU_BETA_DISABLED,
         situ_linear_beta: float = SITU_BETA_DISABLED,
     ) -> 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 f76fc6feebfa..677cf6d15e0c 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
@@ -2752,8 +2752,14 @@ def _apply_situ_epilogue(
             two_lbeta = cutlass.Float32(2.0 * linear_beta)
             neg_lbeta = cutlass.Float32(-linear_beta)
 
-            # sigmoid(x) = rcp(1 + exp2(-x * log2e)), shared by both cores.
             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),
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 957ff1e3613a..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
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 413c42e5d555..1cc87e62ea44 100644
--- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py
+++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py
@@ -1265,6 +1265,13 @@ def test_tp16_nvfp4_padded_loaders_preserve_rank_ownership():
 
 
 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

From a860be6a28803cb1f28488fc756c424c6423d370 Mon Sep 17 00:00:00 2001
From: Xin Guan <294044352+xguannv@users.noreply.github.com>
Date: Thu, 10 Sep 2026 08:58:19 +0000
Subject: [PATCH 7/8] [None][chore] Address review: SiTU beta validation and
 its coverage

Five review findings, all verified against the code first.

_canonicalize_situ_beta mapped the whole `<= 0` range to None, not just
the sentinel. On a SwiGLU layer that turned situ_beta=0.0 into "no
soft-caps supplied" and accepted it, when the kernel's own check would
have refused the combination. Only SITU_BETA_DISABLED disables now;
everything else is forwarded so the kernel sees it.

The kernel's `situ_beta <= 0` test let NaN and positive infinity through
-- every comparison against NaN is false, and an infinity is positive.
The epilogue folds 2/beta and 2*beta at trace time, so either one is
compiled in and returns quietly wrong activations rather than failing.
Replaced with `0 < beta < inf`. Note this file's `math` is the MLIR
dialect, not Python's, so isfinite() is not available here.

Three new tests: the sentinel is the only disabling value; the kernel
refuses NaN, infinity, zero and negative betas; and an ineligible
explicit CUTEDSL request raises rather than degrading to CUTLASS, with
the allow_degradation=True case as the control so the assertion cannot
pass because substitution stopped working altogether.

_skip_if_backend_unavailable() checked only for the CuTe DSL wheel, but
nvfp4_moe_supported admits every SM >= 100 and only the Blackwell
act-fusion kernel carries the SiTU epilogue, so on other architectures
the CUTEDSL case failed in resolution instead of skipping.

Also the launcher's usage() text, which omitted --moe-backend while the
parser accepted it -- and printed that incomplete text on the error path
-- and the kernel class docstring, which still said SwiGLU or Relu2 only.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>
---
 examples/kimi_k3/quick_start_kimi_k3.sbatch   |   3 +-
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  13 ++-
 ...ntiguous_gather_grouped_gemm_act_fusion.py |  17 ++-
 .../_torch/moe/test_kimi_k3_situ_moe.py       | 108 ++++++++++++++++++
 4 files changed, 135 insertions(+), 6 deletions(-)

diff --git a/examples/kimi_k3/quick_start_kimi_k3.sbatch b/examples/kimi_k3/quick_start_kimi_k3.sbatch
index 7308d4362e58..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=""
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 77f0b1373700..275feb48be09 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -89,8 +89,17 @@ def _canonicalize_swiglu_limit_scalar(swiglu_limit_scalar: float) -> float:
 
 
 def _canonicalize_situ_beta(situ_beta: float) -> Optional[float]:
-    """Map the op-boundary sentinel back to ``None`` for the kernel."""
-    return None if situ_beta is None or situ_beta <= 0 else float(situ_beta)
+    """Map the op-boundary sentinel back to ``None`` for the kernel.
+
+    Only the sentinel becomes ``None``. Every other value is forwarded so the
+    kernel's own validation sees it: mapping the whole ``<= 0`` range here
+    would turn ``situ_beta=0.0`` on a SwiGLU layer into "no soft-caps
+    supplied", silently accepting an argument that combination has no meaning
+    for, instead of raising.
+    """
+    if situ_beta is None or situ_beta == SITU_BETA_DISABLED:
+        return None
+    return float(situ_beta)
 
 
 def _get_cute_dsl_swap_ab_candidates(
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 677cf6d15e0c..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
@@ -188,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
 
@@ -433,9 +436,17 @@ def __init__(
                     "ActivationType.SiTu requires both situ_beta and "
                     f"situ_linear_beta, got {situ_beta} and {situ_linear_beta}."
                 )
-            if situ_beta <= 0 or situ_linear_beta <= 0:
+            # 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 positive, got "
+                    "SiTU beta parameters must be finite and positive, got "
                     f"{situ_beta} and {situ_linear_beta}."
                 )
             if self.has_swiglu_limit:
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 1cc87e62ea44..7869682fa0ff 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,66 @@ def test_kimi_situ_betas_must_be_positive(situ_beta, situ_linear_beta):
         modeling_kimi_linear._resolve_kimi_situ_betas(cfg)
 
 
+def test_situ_beta_sentinel_is_the_only_value_that_disables():
+    """Only ``SITU_BETA_DISABLED`` means "this layer is not SiTU".
+
+    The op signature cannot take ``Optional[float]``, so the absence of a
+    soft-cap travels as a sentinel. Canonicalizing the whole non-positive
+    range instead would make ``situ_beta=0.0`` on a SwiGLU layer read as
+    "none supplied" and be accepted, when it should reach the kernel and be
+    refused there.
+    """
+    pytest.importorskip("cutlass")
+    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import (
+        SITU_BETA_DISABLED,
+        _canonicalize_situ_beta,
+    )
+
+    assert _canonicalize_situ_beta(SITU_BETA_DISABLED) is None
+    assert _canonicalize_situ_beta(None) is None
+    for forwarded in (0.0, -2.0, 1.0, 4.0, 25.0):
+        assert _canonicalize_situ_beta(forwarded) == forwarded
+
+
+@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)))
@@ -626,6 +686,45 @@ def test_kimi_k3_allow_list_matches_what_the_backends_declare():
         KimiK3MoERuntime._routed_moe_model_config(model_config)
 
 
+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),
@@ -1278,6 +1377,15 @@ def _skip_if_backend_unavailable(moe_backend):
 
     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):

From 0c6740ce8f693afb6743a02116c7ac649b69d6a2 Mon Sep 17 00:00:00 2001
From: Xin Guan <294044352+xguannv@users.noreply.github.com>
Date: Wed, 16 Sep 2026 08:23:45 +0000
Subject: [PATCH 8/8] [None][chore] Review round 2: drop the sentinel, fix the
 test's control

Four findings from review, all verified against the code first.

The op schema does take Optional[float]. My earlier commit message claimed
it cannot, and that claim was the whole reason SITU_BETA_DISABLED existed;
trtllm::kda_mtp_decode in this repo has carried `scale: Optional[float] =
None` all along. The sentinel, its canonicalization helper and the test
that covered the helper are gone -- None now means "not a SiTU layer" the
way it does everywhere else.

test_nvfp4_kernel_actually_applies_situ compared against SwigluBias,
`gate*sigmoid(gate*alpha)*(up+beta)`. No epilogue on this path computes
that. What a CuteDSL layer computes when it does not run SiTU is
`up * silu(gate)`, so a genuine SwiGLU fallback landed far from both
references and still satisfied `situ_cos > swiglu_cos` -- the test could
not fail the way it claimed to. The control is now the realistic wrong
answer. The absolute bounds the review also asked for are deliberately
not in this commit: cosine is scale-invariant, so the assertion that
catches a mis-scaled SiTU has to be on rel_l2, and its threshold should
come from the measured spread rather than a guess. Both scores are
printed; the bounds land once there is a number to set them from.

test_kimi_k3_allow_list_matches_what_the_backends_declare said "exactly"
and tested inclusion, walking a dict of backends written out by hand --
a second copy of the capability set, which is the defect this module
exists to catch. Both sides are derived now, from BACKEND_FAMILY and from
asking the allow-list, and compared as sets so an offered-but-incapable
backend fails too. `any` rather than `all` over a family because
resolution walks members in IMPL_PRIORITY order: CUTEDSL qualifies
through CuteDslFusedMoE while CuteDslB12xFusedMoE does not declare SiTu.

None of the five tests this PR adds were in any CI list, including
test_nvfp4_kernel_actually_applies_situ[CUTEDSL] -- the one the PR is
for. All are listed in l0_b300.yml now. The [CUTLASS] arm is replaced
rather than joined, per review: K3 rarely runs CUTLASS, and the
parametrization keeps the case for local use.

Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  42 ++-----
 .../moe/fused_moe/fused_moe_cute_dsl.py       |   6 +-
 .../test_lists/test-db/l0_b300.yml            |  12 +-
 .../_torch/moe/test_kimi_k3_situ_moe.py       | 118 +++++++++---------
 4 files changed, 83 insertions(+), 95 deletions(-)

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 275feb48be09..768b09104691 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -79,29 +79,6 @@ def _canonicalize_swiglu_limit_scalar(swiglu_limit_scalar: float) -> float:
     return float("inf") if swiglu_limit_scalar < 0 else swiglu_limit_scalar
 
 
-#: Sentinel for "this layer is not SiTU". A torch custom op schema cannot carry
-#: ``Optional[float]`` here the way a Python signature can, and 0.0 is not
-#: available as the neutral value: the SiTU epilogue divides by both betas, so
-#: zero is a division by zero rather than a no-op. A negative value is
-#: impossible for a real soft-cap -- ``SiTuActivation`` rejects it at
-#: construction -- which makes it safe to reserve.
-SITU_BETA_DISABLED = -1.0
-
-
-def _canonicalize_situ_beta(situ_beta: float) -> Optional[float]:
-    """Map the op-boundary sentinel back to ``None`` for the kernel.
-
-    Only the sentinel becomes ``None``. Every other value is forwarded so the
-    kernel's own validation sees it: mapping the whole ``<= 0`` range here
-    would turn ``situ_beta=0.0`` on a SwiGLU layer into "no soft-caps
-    supplied", silently accepting an argument that combination has no meaning
-    for, instead of raising.
-    """
-    if situ_beta is None or situ_beta == SITU_BETA_DISABLED:
-        return None
-    return float(situ_beta)
-
-
 def _get_cute_dsl_swap_ab_candidates(
     m: int,
     output_aligned: bool,
@@ -3761,8 +3738,8 @@ 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: float = SITU_BETA_DISABLED,
-        situ_linear_beta: float = SITU_BETA_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.
 
@@ -3770,10 +3747,9 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell(
         (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.
-        They default to ``SITU_BETA_DISABLED`` rather than ``None`` because the
-        op schema takes plain floats; the runner maps the sentinel back to
-        ``None`` and then rejects a mismatch against ``activation_type``.
+        ``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(
@@ -3788,8 +3764,8 @@ def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell(
             scaling_vector_size,
             activation_type=ActivationType(activation_type),
             swiglu_limit_scalar=swiglu_limit_scalar,
-            situ_beta=_canonicalize_situ_beta(situ_beta),
-            situ_linear_beta=_canonicalize_situ_beta(situ_linear_beta))
+            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,
@@ -3826,8 +3802,8 @@ 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: float = SITU_BETA_DISABLED,
-        situ_linear_beta: float = SITU_BETA_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.
 
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 d99be994de31..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
@@ -1145,9 +1145,9 @@ def run_moe_nvfp4_impl(
         # ``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 hitting the op's sentinel
-        # default -- passing them unconditionally would make the op signature
-        # lie about which kinds have soft-caps.
+        # 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
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 7869682fa0ff..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,27 +122,6 @@ def test_kimi_situ_betas_must_be_positive(situ_beta, situ_linear_beta):
         modeling_kimi_linear._resolve_kimi_situ_betas(cfg)
 
 
-def test_situ_beta_sentinel_is_the_only_value_that_disables():
-    """Only ``SITU_BETA_DISABLED`` means "this layer is not SiTU".
-
-    The op signature cannot take ``Optional[float]``, so the absence of a
-    soft-cap travels as a sentinel. Canonicalizing the whole non-positive
-    range instead would make ``situ_beta=0.0`` on a SwiGLU layer read as
-    "none supplied" and be accepted, when it should reach the kernel and be
-    refused there.
-    """
-    pytest.importorskip("cutlass")
-    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import (
-        SITU_BETA_DISABLED,
-        _canonicalize_situ_beta,
-    )
-
-    assert _canonicalize_situ_beta(SITU_BETA_DISABLED) is None
-    assert _canonicalize_situ_beta(None) is None
-    for forwarded in (0.0, -2.0, 1.0, 4.0, 25.0):
-        assert _canonicalize_situ_beta(forwarded) == forwarded
-
-
 @pytest.mark.parametrize(
     "situ_beta,situ_linear_beta",
     [
@@ -652,38 +631,54 @@ 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.
-
-    Asserting agreement with each backend's own ``activation_support``
-    rather than against a literal list: a hand-maintained second copy of a
-    capability set is the defect this module already exists to catch, and
-    it is how CUTEDSL stayed shut after its kernels grew the epilogue.
+    """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.fused_moe_cute_dsl import CuteDslFusedMoE
-    from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE
-    from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE
-    from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_cute_dsl import MegaMoECuteDsl
-    from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_deepgemm import MegaMoEDeepGemm
+    from tensorrt_llm._torch.moe.fused_moe.moe_resolution import BACKEND_FAMILY
     from tensorrt_llm._torch.utils import ActivationType
 
-    declares_situ = {
-        "CUTLASS": CutlassFusedMoE,
-        "TRTLLM": TRTLLMGenFusedMoE,
-        "CUTEDSL": CuteDslFusedMoE,
-        "MEGAMOE_CUTEDSL": MegaMoECuteDsl,
-        "MEGAMOE_DEEPGEMM": MegaMoEDeepGemm,
-    }
-    for name, cls in declares_situ.items():
-        assert ActivationType.SiTu in cls.activation_support.kinds, (
-            f"{name} lost SiTU from activation_support; the K3 allow-list "
-            "still offers it, so a layer would resolve here and then fail."
-        )
+    def admits_situ(name):
         model_config = ModelConfig(
             mapping=Mapping(world_size=1, rank=0, tp_size=1),
             moe_backend=name,
         )
-        # Must not raise: this is the gate that kept CUTEDSL out.
-        KimiK3MoERuntime._routed_moe_model_config(model_config)
+        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():
@@ -1869,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.
 
-    ``gate*sigmoid(gate*alpha)*(linear+beta)`` -- what the FC1 epilogue would
-    compute if the activation enum did not resolve to SiTu.
+    ``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.
+
+    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)
@@ -1883,7 +1885,7 @@ 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
 
@@ -1947,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)
@@ -1958,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"
     )