[None][fix] Kimi K3: admit every trtllm-gen SiTu quant format, not just MXFP4

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -118,7 +118,7 @@
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
 from ..modules.situ import SituAndMul
-from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, create_moe
+from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, TRTLLMGenFusedMoE, create_moe
 from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
 from ..utils import AuxStreamType
 from .modeling_speculative import SpecDecOneEngineForCausalLM
@@ -1133,25 +1133,9 @@ def __init__(
             allow_backend_degradation=routed_moe_model_config.moe_backend
             not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"),
         )
-        # trtllm-gen ships SiTu cubins for exactly one dtype combination
-        # (``Bmm_MxE4m3_MxE2m1MxE4m3`` = MXFP8 act x MXFP4 weight) and has no
-        # standalone SiTu activation kernel to fall back on, so a non-MXFP4
-        # routed-expert format cannot be served here. Fail with the fix rather
-        # than with a cubin lookup error deep inside the runner. Checked against
-        # the resolved backend, not the K3 architecture branch, because the
-        # generic FP8_BLOCK_SCALES fallback in ``resolve_moe_backend`` can also
-        # land on TRTLLM.
-        if (
-            routed_moe_model_config.moe_backend == "TRTLLM"
-            and routed_quant_config.quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8
-        ):
-            raise ValueError(
-                f"Kimi K3 routed experts are quantized as "
-                f"{routed_quant_config.quant_algo}, which the TRTLLM "
-                "(trtllm-gen) MoE backend cannot serve: its SiTu activation "
-                "exists only for W4A8_MXFP4_MXFP8. Set moe_config.backend to "
-                "CUTLASS or MEGAMOE_CUTEDSL."
-            )
+        self._check_trtllm_situ_quant(
+            routed_moe_model_config.moe_backend, routed_quant_config.quant_algo
+        )
 
         self.routed_experts = create_moe(**routed_moe_kwargs)
         if not isinstance(self.routed_experts, ConfigurableMoE):
@@ -1288,6 +1272,44 @@ def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> Q
         )
         return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO)
 
+    @staticmethod
+    def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None:
+        """Reject a routed-expert format trtllm-gen has no fused SiTu cubin for.
+
+        trtllm-gen has fused SiTu FC1 cubins for two input formats and no
+        standalone SiTu activation kernel, so anything else has to die here
+        rather than in a cubin lookup deep inside the runner. Checked against
+        the resolved backend, not the K3 architecture branch, because the
+        generic FP8_BLOCK_SCALES fallback in ``resolve_moe_backend`` can also
+        land on TRTLLM.
+
+        The admitted set is read off the backend rather than restated here,
+        because restating it is what broke. This guard was written in #17865
+        when MXFP4 was the only fused SiTu drop; #17940 then added the NVFP4
+        (group-16 ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``) cubins and updated
+        ``TRTLLMGenFusedMoE``'s set without touching this copy. For the week
+        in between, an NVFP4 K3 checkpoint could not start at all -- and not
+        only when TRTLLM was asked for by name, because
+        ``ModelConfig.resolve_moe_backend`` sends every K3 architecture to
+        TRTLLM, so the default AUTO configuration hit this raise too. The unit
+        tests did not catch it: they call ``create_moe`` directly and never
+        reach this guard, so the kernel path stayed green while the model path
+        was closed.
+
+        A staticmethod, not an inline block, so that the invariant is
+        reachable from a test without constructing the whole runtime.
+        """
+        situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos()
+        if moe_backend != "TRTLLM" or quant_algo in situ_supported:
+            return
+        supported = ", ".join(sorted(algo.name for algo in situ_supported))
+        raise ValueError(
+            f"Kimi K3 routed experts are quantized as {quant_algo}, which the "
+            "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins "
+            f"exist only for {supported}. Set moe_config.backend to CUTLASS "
+            "or MEGAMOE_CUTEDSL."
+        )
+
     @staticmethod
     def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig:
         """Build a private routed-expert mapping without mutating the shared

tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py

@@ -15,7 +15,7 @@
 
 import os
 from dataclasses import dataclass, replace
-from typing import Dict, List, Optional, Union
+from typing import Dict, FrozenSet, List, Optional, Union
 
 import torch
 from torch import nn
@@ -188,6 +188,18 @@ def resolve_activation_support(self) -> MoEActivationSupport:
             return replace(support, limit=ActivationParamShape.UNIFORM_SCALAR)
         return support
 
+    @classmethod
+    def situ_supported_quant_algos(cls) -> FrozenSet[QuantAlgo]:
+        """The quant algos this backend has a fused SiTu FC1 cubin for.
+
+        Public because a model that hands SiTu to this backend has to decide,
+        before construction, whether the handoff can work at all -- and the
+        only alternative to asking is restating the set. Restating it is what
+        broke once already: ``modeling_kimi_linear`` carried its own copy,
+        the NVFP4 cubins landed here, and the copy stayed at MXFP4-only.
+        """
+        return frozenset(cls._SITU_SUPPORTED_QUANT_ALGOS)
+
     @classmethod
     def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility:
         """TRTLLM-Gen kernels: the SM100 (Blackwell) family, bfloat16 activations.

tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py

@@ -610,6 +610,57 @@ def test_kimi_k3_moe_auto_backend_defaults_to_trtllm(architecture):
     assert ModelConfig.resolve_moe_backend("AUTO", architecture) == "TRTLLM"
 
 
+# ---------------------------------------------------------------------------
+# The K3 model layer and TRTLLMGenFusedMoE both have to know which routed-expert
+# formats trtllm-gen has a fused SiTu FC1 cubin for. They disagreed once: the
+# model's copy was written when MXFP4 was the only drop (#17865) and #17940 then
+# shipped the NVFP4 cubins and updated only the backend, so for a week an NVFP4
+# K3 checkpoint raised at construction on a path that the kernels supported.
+#
+# It survived because every other SiTu test calls ``create_moe`` directly and so
+# never reaches the model-layer guard -- the kernel path was green throughout.
+# These two tests enter through the guard instead.
+# ---------------------------------------------------------------------------
+
+
+def test_kimi_k3_trtllm_situ_admits_every_backend_supported_quant():
+    """The model must not narrow what the backend says it can serve.
+
+    Asserting agreement rather than a literal set is the point: a new fused
+    SiTu cubin family should require no edit here, and removing one should
+    fail loudly rather than leave a stale allow-list behind.
+    """
+    from tensorrt_llm._torch.moe.fused_moe import TRTLLMGenFusedMoE
+
+    supported = TRTLLMGenFusedMoE.situ_supported_quant_algos()
+    assert QuantAlgo.NVFP4 in supported, (
+        "trtllm-gen has shipped group-16 Bmm_E2m1_E2m1E2m1_..._siTuGlu_* cubins "
+        "since #17940; if this fails the backend regressed, not the model."
+    )
+    for algo in supported:
+        KimiK3MoERuntime._check_trtllm_situ_quant("TRTLLM", algo)
+
+
+@pytest.mark.parametrize("quant_algo", [QuantAlgo.FP8_BLOCK_SCALES, QuantAlgo.W4A16_MXFP4, None])
+def test_kimi_k3_trtllm_situ_rejects_quant_without_fused_cubin(quant_algo):
+    """...and must still reject the formats that have no fused SiTu cubin.
+
+    ``resolve_moe_backend`` sends every K3 architecture to TRTLLM, including
+    the generic FP8_BLOCK_SCALES fallback, so this rejection is reachable
+    without anyone asking for TRTLLM by name. It has to name the fix.
+    """
+    from tensorrt_llm._torch.moe.fused_moe import TRTLLMGenFusedMoE
+
+    assert quant_algo not in TRTLLMGenFusedMoE.situ_supported_quant_algos()
+    with pytest.raises(ValueError, match="fused SiTu cubins exist only for"):
+        KimiK3MoERuntime._check_trtllm_situ_quant("TRTLLM", quant_algo)
+
+    # Any other backend owns its own SiTu translation and is not this guard's
+    # business -- gating it here is how CUTLASS would get blocked by a
+    # trtllm-gen cubin inventory.
+    KimiK3MoERuntime._check_trtllm_situ_quant("CUTLASS", quant_algo)
+
+
 # ---------------------------------------------------------------------------
 # MoE tensor-parallel shard parity (ConfigurableMoE / TRTLLM-Gen, GPU).
 #