diff --git a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md index f2765e1ac57a..63bb456a37db 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -159,7 +159,7 @@ Still on old path (standalone, with embedded communication): | `mega_moe/mega_moe_deepgemm.py` | `MegaMoEDeepGemm` | SM100/SM103 | W4A8_MXFP4_MXFP8 via DeepGEMM `fp8_fp4_mega_moe` fused dispatch+GEMM+act+GEMM+combine kernel; requires `hidden_size % 512 == 0` | `FUSED_COMM` | | `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider (hard gate); threads per-expert `fc31_alpha`/`fc2_alpha`/`fc1_norm_const` through the kernel ABI and supports SwiGLU clamp via `swiglu_limit`; default deepgemm graph (topk score folded before fc1-out quant, host `combine_output.sum(dim=1)`) | `FUSED_COMM` | | `fused_moe_marlin.py` | `MarlinFusedMoE` | SM89-SM99 | W4A16 NVFP4 on Ada/Hopper (BF16 activations + FP4 weights, fused single-launch `marlin_nvfp4_moe_gemm` kernel); supports attention-DP + EP via external comm (scheduler precomputes routing; dispatch payload is plain BF16, no activation scales); non-NVFP4 layers (e.g. unquantized MTP draft layers) degrade to Cutlass in `resolve_moe_impl`, recorded in the layer's `MoEResolutionReport`; no dynamic EPLB | `EXTERNAL_COMM` | -| `fused_moe_triton.py` | `TritonFusedMoE` | SM90 only | GPT-OSS on Hopper (requires `swiglu_gptoss_style=True`) | (legacy path) | +| `fused_moe_triton.py` | `TritonFusedMoE` | SM90 only | GPT-OSS and plain-SwiGLU MXFP4 on Hopper (SwiGLU family only) | (legacy path) | | `fused_moe_vanilla.py` | `VanillaMoE` | All devices | Reference / debugging only | (legacy path) | ### Communication (`fused_moe/communication/`) @@ -374,6 +374,14 @@ every specialized backend — `CuteDslFusedMoE`, `CuteDslB12xFusedMoE`, `DeepGemmFusedMoE`, `DenseGEMMFusedMoE`, `MarlinFusedMoE` — while `TRTLLMGenFusedMoE` accepts only the algorithms in its `_GPTOSS_SUPPORTED_ALGOS`. +`TritonFusedMoE` gates the activation *family* (`Swiglu`, `SwigluBias`), not +`swiglu_gptoss_style`: each quant method's `apply` branches on `beta == 1.0`, +serving gpt-oss through the fused Triton activation and plain SwiGLU through +`swiglu_torch` at the `alpha=1.0` / `beta=0.0` defaults. Gating it on +`swiglu_gptoss_style` instead degraded explicit `TRITON` requests for +plain-SwiGLU MXFP4 checkpoints to Cutlass, which cannot load the `weight_scale` +scale spelling those checkpoints ship (nvbugs/6660905). + Cutlass gates gpt-oss / MiniMax SwiGLU on unquantized, NVFP4, and the MXFP4 family (`CutlassFusedMoE._GPTOSS_SUPPORTED_ALGOS` = `None`, `NVFP4`, `W4A16_MXFP4`, `W4A8_MXFP4_FP8`, `W4A8_MXFP4_MXFP8`). The CUDA kernel is not diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py index e8c90eaabf9f..b3654edda6fa 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py @@ -38,6 +38,7 @@ from tensorrt_llm.models.modeling_utils import QuantAlgo from ...model_config import ModelConfig +from ...utils import ActivationType from ..linear import TensorParallelMode, load_weight_shard from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, MoERejectReason) @@ -1550,7 +1551,7 @@ class TritonFusedMoE(MoE): @classmethod def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: - """Triton MoE: SM90 only, and only the gpt-oss style swiglu. + """Triton MoE: SM90 only, SwiGLU family only. Supports unquantized BF16, FP8 per-tensor QDQ, W4A8_MXFP4_FP8 and W4A16_MXFP4. @@ -1569,11 +1570,18 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: MoERejectReason.EPLB_UNSUPPORTED, "TritonFusedMoE does not implement the EPLB slot hooks") - # Require gpt-oss SwiGLU; abstain when the style is unknown. - if p.swiglu_gptoss_style is False: + # Gate the activation family, NOT ``swiglu_gptoss_style``: every + # ``apply`` branches on ``beta == 1.0``, so gpt-oss takes the fused + # activation and the plain-SwiGLU alpha=1/beta=0 defaults take + # ``swiglu_torch``. Rejecting plain SwiGLU degraded an explicit TRITON + # request to Cutlass, which then died loading the MXFP4 scales this + # checkpoint spells ``weight_scale`` (nvbugs/6660905). + if p.activation_type not in (ActivationType.Swiglu, + ActivationType.SwigluBias): return _reject( MoERejectReason.ACTIVATION_UNSUPPORTED, - "TritonFusedMoE only supports swiglu_gptoss_style=True") + "TritonFusedMoE fuses Swiglu and SwigluBias only, got " + f"{p.activation}") if d.smart_router: return _reject( diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index c7b31356f516..f80abdeb5303 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -234,7 +234,6 @@ full:H20/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[t full:H20/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=True] SKIP (https://nvbugs/6546909) full:H20/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=TRTLLM-torch_compile=True] SKIP (https://nvbugs/6580087) full:H20/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6546909) -full:H20/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a16_mxfp4[latency-TRITON] SKIP (https://nvbugs/6660905) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[bf16] SKIP (https://nvbugs/6618649) diff --git a/tests/unittest/_torch/modules/moe/test_triton_moe_resolution.py b/tests/unittest/_torch/modules/moe/test_triton_moe_resolution.py new file mode 100644 index 000000000000..46ad467be906 --- /dev/null +++ b/tests/unittest/_torch/modules/moe/test_triton_moe_resolution.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Which activations ``TritonFusedMoE`` claims during MoE resolution. + +Triton gates the activation *family*, not ``swiglu_gptoss_style`` — see +``TritonFusedMoE.can_implement`` for why (nvbugs/6660905). + +Selection-time only: these run without a cubin and need no GPU. Kernel-level +MXFP4 numerics live in ``_torch/modules/test_fused_moe.py`` +(``test_fused_moe_triton_mxfp4``, which covers bias on and off). +""" + +import pytest +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.fused_moe import RenormalizeMoeRoutingMethod +from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import TritonFusedMoE +from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoEEnvironment, MoERejectReason +from tensorrt_llm._torch.modules.fused_moe.impl_environment import override_moe_environment +from tensorrt_llm._torch.modules.fused_moe.moe_resolution import ( + impl_class_for, + infer_swiglu_gptoss_style, + resolve_moe_impl, +) +from tensorrt_llm._torch.utils import ActivationType +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def _resolve_triton(*, bias=False, swiglu_alpha=None, swiglu_beta=None, activation_type=None): + """Resolve a W4A16_MXFP4 TRITON request the way ``create_moe`` does. + + ``swiglu_gptoss_style`` goes through ``infer_swiglu_gptoss_style`` rather + than being hard-coded, because ``resolve_moe_impl`` takes it as a parameter: + passing a literal would leave the activation gate unreached and every + assertion below would hold vacuously. + """ + cfg = ModelConfig() + cfg.moe_backend = "TRITON" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_MXFP4) + # SM90 is Triton's only window, so activation stays the lone variable. + with override_moe_environment(MoEEnvironment(sm=90)): + return resolve_moe_impl( + cfg, + dtype=torch.bfloat16, + routing=RenormalizeMoeRoutingMethod(top_k=4), + swiglu_gptoss_style=infer_swiglu_gptoss_style( + bias=bias, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + activation_type=activation_type, + ), + bias=bias, + activation_type=activation_type, + ) + + +def test_triton_serves_plain_swiglu_mxfp4(): + """Qwen3-30B-A3B W4A16_MXFP4: no bias, no alpha/beta -> Triton, not Cutlass.""" + report = _resolve_triton() + # The regression degraded to Cutlass here and only failed later, at load. + assert impl_class_for(report) is TritonFusedMoE + assert report.selected_by == "pinned" + assert not report.degraded + + +def test_triton_still_serves_gptoss_swiglu(): + """The gpt-oss package keeps resolving to Triton (the pre-existing case).""" + report = _resolve_triton( + bias=True, + swiglu_alpha=torch.tensor([1.702]), + swiglu_beta=torch.tensor([1.0]), + ) + assert impl_class_for(report) is TritonFusedMoE + assert not report.degraded + + +@pytest.mark.parametrize( + "activation_type", + [ + pytest.param(ActivationType.Geglu, id="geglu"), + pytest.param(ActivationType.Relu2, id="relu2"), + ], +) +def test_triton_degrades_on_non_swiglu_activation(activation_type): + """A non-SwiGLU activation has no Triton path and must still degrade.""" + report = _resolve_triton(activation_type=activation_type) + assert impl_class_for(report) is CutlassFusedMoE + assert report.degraded + assert report.degraded_from.reason is MoERejectReason.ACTIVATION_UNSUPPORTED