From 55a105a509b4e0418b54d02982f9c8c626bbbb25 Mon Sep 17 00:00:00 2001 From: Xin Guan <294044352+xguannv@users.noreply.github.com> Date: Thu, 3 Sep 2026 07:14:17 +0000 Subject: [PATCH] [None][fix] Kimi K3: admit every trtllm-gen SiTu quant format, not just MXFP4 KimiK3MoERuntime rejected any routed-expert quantization other than W4A8_MXFP4_MXFP8 on the TRTLLM backend. That was correct when it was written in #17865 -- MXFP4 was then the only fused SiTu FC1 drop -- but (Bmm_E2m1_E2m1E2m1_..._siTuGlu_*) and updated only TRTLLMGenFusedMoE._SITU_SUPPORTED_QUANT_ALGOS, leaving the model layer's copy behind. Since then the tree has contradicted itself: the backend declares NVFP4 servable and the model refuses to build it. The blast radius is wider than an explicit backend request. ModelConfig.resolve_moe_backend maps every Kimi K3 architecture to TRTLLM unconditionally, so an NVFP4 K3 checkpoint could not start under the default AUTO configuration either -- the documented workaround ("backend: CUTLASS is required, not a preference") is in examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml and is itself now stale. Nothing caught it because every existing SiTu test reaches the kernels through create_moe and never enters this guard, so the NVFP4 x TRTLLM kernel path stayed green for the whole window. Read the admitted set off the backend instead of restating it, and add the two tests that enter through the guard: one asserting the model admits everything the backend claims, one asserting formats without a fused cubin are still rejected and still name the fix. Signed-off-by: Xin Guan <294044352+xguannv@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 62 +++++++++++++------ .../moe/fused_moe/fused_moe_trtllm_gen.py | 14 ++++- .../_torch/moe/test_kimi_k3_situ_moe.py | 51 +++++++++++++++ 3 files changed, 106 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 86291f2ea770..0eb5b92b526f 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/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 diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py index b51894bce1fc..7dd1200933cc 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_trtllm_gen.py +++ b/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. 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 c1892a1c6023..90b9b41b56d1 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/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). #