diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py index 0b571a515109..b1d411b1649e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py @@ -26,6 +26,7 @@ from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, OptimizationProfile, TunableRunner, TuningConfig) from ...custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper +from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ...model_config import ModelConfig from ...utils import (ActivationType, AuxStreamType, EventType, Fp4QuantizedTensor, @@ -39,6 +40,26 @@ from .quantization import MoEWeightLoadingMode, NVFP4CuteDslFusedMoEMethod from .routing import BaseMoeRoutingMethod +# These runners are defined inside cute_dsl_custom_ops' ``if +# IS_CUTLASS_DSL_AVAILABLE:`` block, which has no else-branch, so importing them +# unconditionally would break every importer of this file -- and create_moe +# imports it eagerly under _torch.models, so that reaches all model startup +# rather than just this backend. Guard the same way the sibling custom_ops +# modules do, and leave the tuple empty when the DSL is absent: no CuteDSL +# runner can be tuned in that case, so nothing can match it. +_TILE_SIZE_CHECKED_RUNNERS: Tuple[type, ...] = () +if IS_CUTLASS_DSL_AVAILABLE: + from ...custom_ops.cute_dsl_custom_ops import ( + Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner, + Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner, + Sm100BlockScaledContiguousGroupedGemmRunner, + Sm100BlockScaledContiguousGroupedGemmSwigluFusionRunner) + _TILE_SIZE_CHECKED_RUNNERS = ( + Sm100BlockScaledContiguousGroupedGemmRunner, + Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner, + Sm100BlockScaledContiguousGroupedGemmSwigluFusionRunner, + Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner) + @dataclass class NvFp4WeightView: @@ -328,19 +349,9 @@ def runner_tactic_comb_checker( # eagerly under _torch.models, so that reaches all model startup rather # than just this backend. Reaching this line means a CuteDSL runner is # already being tuned, so the DSL is installed. - from ...custom_ops.cute_dsl_custom_ops import ( - Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner, - Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner, - Sm100BlockScaledContiguousGroupedGemmRunner, - Sm100BlockScaledContiguousGroupedGemmSwigluFusionRunner) for runner, tactic in comb: - if isinstance( - runner, - (Sm100BlockScaledContiguousGroupedGemmRunner, - Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner, - Sm100BlockScaledContiguousGroupedGemmSwigluFusionRunner, - Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner)): + if isinstance(runner, _TILE_SIZE_CHECKED_RUNNERS): mma_tiler_mn, *_ = tactic if mma_tiler_mn[0] != tile_size: return False diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 8b1745638364..7952f4fa9fee 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -3,6 +3,7 @@ from utils.util import check_accuracy from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import cute_dsl_nvfp4_grouped_gemm_ref from tensorrt_llm._torch.modules.fused_moe.quantization import interleave_linear_and_gate from tensorrt_llm._torch.utils import ( @@ -413,6 +414,7 @@ def test_moe_gelu(dtype: str, num_tokens: int, top_k: int, tile_size: int): get_sm_version() not in (100, 103), reason="This test is only supported on SM 100 and SM 103 GPUs", ) +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available") @pytest.mark.parametrize("tile_size", [128, 256]) @pytest.mark.parametrize("ep_size", [1, 8, 32]) @pytest.mark.parametrize("top_k", [1, 2, 8]) @@ -509,6 +511,7 @@ def test_nvfp4_grouped_gemm_blackwell(num_tokens: int, top_k: int, ep_size: int, get_sm_version() not in (100, 103), reason="This test is only supported on SM 100 and SM 103 GPUs", ) +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available") @pytest.mark.parametrize("tile_size", [128, 256]) @pytest.mark.parametrize("ep_size", [1, 8, 32]) @pytest.mark.parametrize("top_k", [1, 2, 8]) @@ -610,6 +613,7 @@ def test_nvfp4_grouped_gemm_finalize_blackwell( get_sm_version() not in (100, 103), reason="This test is only supported on SM 100 and SM 103 GPUs", ) +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available") @pytest.mark.parametrize("tile_size", [128, 256]) @pytest.mark.parametrize("ep_size", [1, 8, 32]) @pytest.mark.parametrize("top_k", [1, 2, 8]) @@ -735,6 +739,7 @@ def test_nvfp4_grouped_gemm_swiglu_blackwell( get_sm_version() not in (100, 103), reason="This test is only supported on SM 100 and SM 103 GPUs", ) +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available") @pytest.mark.parametrize( "activation_type", [ActivationType.Swiglu, ActivationType.Relu2],