diff --git a/cpp/include/tensorrt_llm/common/cudaUtils.h b/cpp/include/tensorrt_llm/common/cudaUtils.h index cd58a7abb5d9..ef1986e5c34f 100644 --- a/cpp/include/tensorrt_llm/common/cudaUtils.h +++ b/cpp/include/tensorrt_llm/common/cudaUtils.h @@ -304,10 +304,10 @@ inline int getSMVersion(bool queryRealSmArch = false) return sm; } -inline bool isSM100Family() +inline bool isBlackwellFamily() { int const sm = getSMVersion(); - return sm == 100 || sm == 103; // To be continued... + return sm == 100 || sm == 103 || sm == 120 || sm == 121; } inline int getDevice() diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp index 68e3e4d60040..28c81138d9c9 100644 --- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp @@ -52,7 +52,7 @@ FmhaDispatcher::FmhaDispatcher(MHARunnerFixedParams fixedParams) // TRTLLM-GEN only supports power of 2 head sizes. // The exception will fall back to fmha v2. // Please update fmha_v2/setup.py if you want to add more supported head sizes. - , mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 80 && fixedParams.headSize != 72) + , mUseTllmGen(tensorrt_llm::common::isBlackwellFamily() && fixedParams.headSize != 80 && fixedParams.headSize != 72) { if (mUseTllmGen) { diff --git a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp index 41aab02d3043..8d06164b9963 100644 --- a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp @@ -131,7 +131,7 @@ QKVPreprocessingParams makeQKVPreprocessingParams(XQAParams co XqaDispatcher::XqaDispatcher(XqaFixedParams fixedParams) : mFixedParams(fixedParams) , mQDataType(mFixedParams.inputDataType) - , mUseTllmGen(tensorrt_llm::common::isSM100Family()) + , mUseTllmGen(tensorrt_llm::common::isBlackwellFamily()) , mMultiProcessorCount(getMultiProcessorCount()) { if (mUseTllmGen) diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 9a7af4da49f6..dbdf12464496 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -493,7 +493,7 @@ class Runner : public RunnerBase } if (op.mIsSpecDecodingEnabled && op.mUseSpecDecoding) { - bool useTllmGen = tensorrt_llm::common::isSM100Family(); + bool useTllmGen = tensorrt_llm::common::isBlackwellFamily(); if (useTllmGen) { TORCH_CHECK(spec_decoding_tensor_params.size() == 6, diff --git a/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp index 3440a59737fc..95e399084bd2 100644 --- a/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp @@ -48,7 +48,8 @@ std::vector run_fp4_block_scale_moe_runner(torch::optional const& out_tensor = torch::nullopt) { TORCH_CHECK(dtype == btg::Dtype::E4m3 || dtype == btg::Dtype::E2m1, "dtype can only be e4m3 or e2m1."); - TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP4 block scale MOE"); + TORCH_CHECK(tensorrt_llm::common::isBlackwellFamily(), + "Blackwell family (SM100/103/120/121) is required for FP4 block scale MOE"); TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64 || tile_tokens_dim == 128 || tile_tokens_dim == 256, "tile_tokens_dim must be 8, 16, 32, 64, 128, 256"); diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index 3c13c695991c..79f9e24e7602 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -46,7 +46,8 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit MoeRunnerType& moe_runner, int64_t moeConfigIndex, std::optional const& topk_weights, std::optional const& topk_ids, std::optional const& out_tensor = std::nullopt) { - TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE"); + TORCH_CHECK(tensorrt_llm::common::isBlackwellFamily(), + "Blackwell family (SM100/103/120/121) is required for FP8 block scale MOE"); if (topk_ids.has_value() && topk_weights.has_value()) { diff --git a/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp index 092f8f013620..77285c37492d 100644 --- a/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp @@ -38,7 +38,8 @@ torch::Tensor fp8_per_tensor_scale_moe_runner(torch::optional con int64_t const tile_tokens_dim, int64_t const routing_method_type, torch::optional const& topk_weights, torch::optional const& topk_ids) { - TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE"); + TORCH_CHECK(tensorrt_llm::common::isBlackwellFamily(), + "Blackwell family (SM100/103/120/121) is required for FP8 per-tensor scale MOE"); TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64 || tile_tokens_dim == 128 || tile_tokens_dim == 192 || tile_tokens_dim == 256, "tile_tokens_dim must be 8, 16, 32, 64, 128, 256"); diff --git a/cpp/tensorrt_llm/thop/fp8Quantize.cpp b/cpp/tensorrt_llm/thop/fp8Quantize.cpp index 9e94e02950f7..5bdab0143020 100644 --- a/cpp/tensorrt_llm/thop/fp8Quantize.cpp +++ b/cpp/tensorrt_llm/thop/fp8Quantize.cpp @@ -69,7 +69,7 @@ std::tuple fp8_quantize_1x128(at::Tensor const& self, bo act_buffer, act_scale_buffer, reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), n, m, stream, use_ue8m0); // Post-process the scale tensor for sm100 gemm/moe kernel - if (tensorrt_llm::common::isSM100Family()) + if (tensorrt_llm::common::isBlackwellFamily()) { auto const num_n_blocks = (n + 127) / 128; auto const act_scal_elesize = num_n_blocks * m_padded; diff --git a/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp index 5e8331b77c3f..e3df2ff19e94 100644 --- a/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp @@ -52,7 +52,8 @@ torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::optional torch::optional const& topk_weights, torch::optional const& topk_ids, torch::optional const& out_tensor) { - TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by MXFP4 block scale MOE"); + TORCH_CHECK(tensorrt_llm::common::isBlackwellFamily(), + "Blackwell family (SM100/103/120/121) is required for MXFP4 block scale MOE"); TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64 || tile_tokens_dim == 128 || tile_tokens_dim == 256, "tile_tokens_dim must be 8, 16, 32, 64, 128, 256"); diff --git a/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp b/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp index 116013de2ffe..387eeae53f89 100644 --- a/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp +++ b/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp @@ -170,7 +170,7 @@ class PrepareCustomMaskTest : public ::testing::Test protected: static bool shouldSkip() { - return !tensorrt_llm::common::isSM100Family(); + return !tensorrt_llm::common::isBlackwellFamily(); } void SetUp() override diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index eaee69e54ebb..9c324129a2dd 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -1301,6 +1301,17 @@ def _optimization_profiles( for spec in tuning_config.dynamic_tensor_specs: assert callable(spec.gen_tuning_buckets) or isinstance(spec.gen_tuning_buckets, (list, tuple)), \ "The given dynamic dimension must provide a opt value generation function or a list of opt values" + if spec.input_idx >= len(base_profile.shapes): + logger.debug( + f"[Autotuner] Skipping DynamicTensorSpec with input_idx={spec.input_idx}: " + f"only {len(base_profile.shapes)} inputs available.") + continue + if spec.dim_idx >= len(base_profile.shapes[spec.input_idx]): + logger.debug( + f"[Autotuner] Skipping DynamicTensorSpec with dim_idx={spec.dim_idx} for " + f"input {spec.input_idx}: shape has only " + f"{len(base_profile.shapes[spec.input_idx])} dims.") + continue if self.skip_dynamic_tuning_buckets: if spec.map_to_tuning_buckets is not None: # Still include the bucketed value of the actual shape so the @@ -1358,10 +1369,21 @@ def _optimization_profiles( # Adjust the profile to satisfy the constraints for spec in tuning_config.constraint_specs: - min_value = opt_value = max_value = spec.infer_shape( - p.get_opt_shapes()) + if spec.input_idx >= len(p.shapes): + logger.debug( + f"[Autotuner] Skipping ConstraintSpec with input_idx={spec.input_idx}: " + f"only {len(p.shapes)} inputs available.") + continue if p.shapes[spec.input_idx] == [StaticDim(0)]: continue + if spec.dim_idx >= len(p.shapes[spec.input_idx]): + logger.debug( + f"[Autotuner] Skipping ConstraintSpec with dim_idx={spec.dim_idx} for " + f"input {spec.input_idx}: shape has only " + f"{len(p.shapes[spec.input_idx])} dims.") + continue + min_value = opt_value = max_value = spec.infer_shape( + p.get_opt_shapes()) p.shapes[spec.input_idx][spec.dim_idx] = DynamicDim( min_value, opt_value, max_value) generated_profiles.append(p) @@ -1395,6 +1417,21 @@ def _find_nearest_profile( base_profile = list(list(shape) for shape in shapes) for spec in dynamic_tensor_specs: + # Bounds check: skip specs that reference inputs or dimensions not present in the + # current shapes tuple. This can happen on hardware (e.g. SM121 / DGX Spark) where + # ops produce fewer or differently-shaped tensors than the specs were authored for. + if spec.input_idx >= len(base_profile): + logger.debug( + f"[Autotuner] Skipping DynamicTensorSpec with input_idx={spec.input_idx}: " + f"only {len(base_profile)} inputs available.") + continue + if spec.dim_idx >= len(base_profile[spec.input_idx]): + logger.debug( + f"[Autotuner] Skipping DynamicTensorSpec with dim_idx={spec.dim_idx} for " + f"input {spec.input_idx}: shape has only {len(base_profile[spec.input_idx])} dims." + ) + continue + # During runtime: apply map_to_tuning_buckets to map input to bucket # During tuning: no mapper, use raw bucket value if apply_map_to_tuning_buckets: @@ -1409,8 +1446,20 @@ def _find_nearest_profile( # associated dimensions dependent on other free dynamic dimensions, so assign -1 in the profile for spec in constraint_specs: + # Bounds check: same defensive guard as above for constraint specs. + if spec.input_idx >= len(base_profile): + logger.debug( + f"[Autotuner] Skipping ConstraintSpec with input_idx={spec.input_idx}: " + f"only {len(base_profile)} inputs available.") + continue if base_profile[spec.input_idx] == [0]: continue + if spec.dim_idx >= len(base_profile[spec.input_idx]): + logger.debug( + f"[Autotuner] Skipping ConstraintSpec with dim_idx={spec.dim_idx} for " + f"input {spec.input_idx}: shape has only {len(base_profile[spec.input_idx])} dims." + ) + continue base_profile[spec.input_idx][spec.dim_idx] = -1 return tuple(tuple(shape) for shape in base_profile) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index ee150d1be98a..2d5d37acfb06 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -809,13 +809,13 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], # Add CuteDSL tactics if available if self._is_backend_allowed("cutedsl"): if IS_CUTLASS_DSL_AVAILABLE: - # Check SM version first - CuteDSL NVFP4 only supports SM 100 (B200) + # Check SM version first - CuteDSL NVFP4 only supports Blackwell sm_version = get_sm_version() - if sm_version not in [100, 103]: + if sm_version not in [100, 103, 120, 121]: if self._is_only_backend("cutedsl"): # Explicitly forced CuteDSL but SM version not supported raise ValueError( - f"CuteDSL NVFP4 backend requires SM 100 (B200) or SM 103 (B300), but got SM {sm_version}. " + f"CuteDSL NVFP4 backend requires Blackwell (SM100/103/120/121), but got SM {sm_version}. " f"CuteDSL NVFP4 is not supported on this GPU architecture. " "Please add other backends to allowed_backends.") else: diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index b1a932953b06..2d6d825a7b44 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -11,11 +11,11 @@ import transformers from transformers.utils import HF_MODULES_CACHE +import tensorrt_llm._utils as _utils from tensorrt_llm import logger from tensorrt_llm._torch.pyexecutor.config_utils import ( get_qwen3_hybrid_num_attention_layers, is_nemotron_hybrid, is_qwen3_hybrid, load_pretrained_config) -from tensorrt_llm._utils import get_sm_version, torch_dtype_to_binding from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.functional import AllReduceStrategy from tensorrt_llm.llmapi.llm_args import (DeepSeekSparseAttentionConfig, @@ -251,15 +251,14 @@ def resolve_moe_backend(moe_backend: str, architecture: str) -> str: if moe_backend.upper() != "AUTO": return moe_backend + # SM100 family: use TRTLLM (trtllm-gen kernels need tcgen05.mma, not available on SM120/SM121) + if _utils.is_sm_100f(): + return "TRTLLM" + if architecture == "GptOssForCausalLM": - sm_version = get_sm_version() - # Select the best performing backend based on SM version - if 100 <= sm_version < 120: # Blackwell - return "TRTLLM" - elif 90 <= sm_version < 100: # Hopper + sm_version = _utils.get_sm_version() + if 90 <= sm_version < 100: # Hopper return "TRITON" - else: - return "CUTLASS" # Fallback to CUTLASS for other SM versions (e.g., SM120) return "CUTLASS" @@ -335,7 +334,7 @@ def load_modelopt_quant_config(quant_config_file, checkpoint_dir, def get_mxfp4_quant_algo(moe_backend, is_dynamic_quant=False): quant_algo = ModelConfig.override_quant_algo() if quant_algo is None and not is_dynamic_quant: - if get_sm_version() >= 100: + if _utils.get_sm_version() >= 100: if moe_backend == 'TRITON': return QuantAlgo.W4A8_MXFP4_FP8 else: @@ -677,7 +676,7 @@ def ceil_div(a, b): num_rnn_layers=0, num_heads=num_heads, hidden_size=hidden_size, - data_type=torch_dtype_to_binding( + data_type=_utils.torch_dtype_to_binding( self.pretrained_config.torch_dtype)) # For kv cache size calculation: set tokens_per_block 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 1273262f5f42..07f8af12b0d4 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 @@ -19,7 +19,7 @@ import torch import torch.nn.functional as F -from tensorrt_llm._utils import get_sm_version, is_sm_100f +import tensorrt_llm._utils as _utils from tensorrt_llm.models.modeling_utils import QuantAlgo from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, @@ -62,8 +62,8 @@ def cute_dsl_fp8_group_blockwise_gemm_ref( a_tmp = a.as_strided((m, k, 1), (k, 1, m * k)) b_tmp = b.permute(1, 2, 0) - # Note: we have different output scale shape for fp8_quantize_1x128, so we need to handle it differently for sm100 and other archs. - if is_sm_100f(): + # Note: we have different output scale shape for fp8_quantize_1x128, so we need to handle it differently for Blackwell and other archs. + if _utils.is_blackwell(): input_scale_tmp = a_sf.permute(1, 0).as_strided((m, w_k, 1), (1, m, m * w_k)) else: @@ -339,7 +339,7 @@ def can_implement( Check if CuteDslFusedMoE can implement the given quantization algorithm. CuteDslFusedMoE supports: - - NVFP4: SM in {100, 103} + - NVFP4: SM in {100, 103, 120, 121} Does NOT support unquantized mode. Output dtype is hardcoded to bfloat16. Does NOT support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit). @@ -356,7 +356,7 @@ def can_implement( """ from .interface import _warn_and_return - sm_version = get_sm_version() + sm_version = _utils.get_sm_version() # CuteDslFusedMoE requires at least SM90 if sm_version < 90: @@ -381,11 +381,12 @@ def can_implement( "CuteDslFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)" ) - # NVFP4 - SM in {100, 103} + # NVFP4 - SM100 family only (SM120/SM121 has scale dtype mismatch) if quant_algo == QuantAlgo.NVFP4: - if sm_version not in {100, 103}: + if not _utils.is_sm_100f(sm_version): return _warn_and_return( - f"NVFP4 requires SM100 or SM103, got SM{sm_version}") + f"CuteDSL NVFP4 requires SM100 family (SM120/SM121 excluded due to scale dtype mismatch), got SM{sm_version}" + ) return True, None return _warn_and_return( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 23354f5a5b34..8122aab97668 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -21,9 +21,9 @@ import torch from torch import nn +import tensorrt_llm._utils as _utils from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll -from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -129,7 +129,7 @@ def can_implement( """ from .interface import _warn_and_return - sm_version = get_sm_version() + sm_version = _utils.get_sm_version() # TRTLLMGenFusedMoE requires SM in {100, 103} if sm_version not in {100, 103}: @@ -203,7 +203,7 @@ def __init__( activation_type=activation_type, ) - sm_version = get_sm_version() + sm_version = _utils.get_sm_version() if sm_version >= 120: raise NotImplementedError( "TRTLLMGenFusedMoE does not support SM120 and above.") diff --git a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py index c9771996d55b..96358c82d6ac 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py +++ b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py @@ -223,11 +223,11 @@ def select_op(module: 'MoE') -> MoEOp: from .moe_op_deepgemm import DeepGemmMoEOp # Check if we should use DeepGemm op - # Blackwell has SM version 100 - is_blackwell = is_sm_100f() + # DeepGemm supports SM100/SM103 (datacenter Blackwell) only + use_deepgemm_arch = is_sm_100f() has_block_fp8 = module.has_deepseek_fp8_block_scales - if is_blackwell and has_block_fp8: + if use_deepgemm_arch and has_block_fp8: # Use DeepGemm op for Blackwell with block FP8 return DeepGemmMoEOp() else: diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index c30b7d771aa9..91fe5b66ae97 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -1099,7 +1099,7 @@ class DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm( DeepSeekFP8BlockScalesFusedMoEMethod): def _needs_e8m0_resmooth(self): - return is_sm_100f() or get_sm_version() == 120 + return is_sm_100f() or get_sm_version() in (120, 121) def post_load_weights(self, module: torch.nn.Module): if self._needs_e8m0_resmooth(): diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 47a6a88499ea..580928962507 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -796,6 +796,18 @@ def is_sm_100f(sm_version=None): return sm_version == 100 or sm_version == 103 +def is_sm_120f(sm_version=None): + if sm_version is None: + sm_version = get_sm_version() + return sm_version == 120 or sm_version == 121 + + +def is_blackwell(sm_version=None): + if sm_version is None: + sm_version = get_sm_version() + return is_sm_100f(sm_version) or is_sm_120f(sm_version) + + def print_all_stacks(): """Print stack traces for all threads""" for thread_id, frame in sys._current_frames().items(): diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index cd399acb4155..c1c2c74da172 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -1919,6 +1919,18 @@ def is_sm_100f(sm_version=None): return sm_version == 100 or sm_version == 103 +def is_sm_120f(sm_version=None): + if sm_version is None: + sm_version = get_sm_version() + return sm_version == 120 or sm_version == 121 + + +def is_blackwell(sm_version=None): + if sm_version is None: + sm_version = get_sm_version() + return is_sm_100f(sm_version) or is_sm_120f(sm_version) + + def get_gpu_device_list(): "get device list" with tempfile.TemporaryDirectory() as temp_dirname: @@ -1970,7 +1982,7 @@ def check_device_contain(keyword_list): get_sm_version() != 90, reason="This test is only supported in Hopper architecture") -skip_no_sm120 = pytest.mark.skipif(get_sm_version() != 120, +skip_no_sm120 = pytest.mark.skipif(not is_sm_120f(), reason="This test is for SM120") skip_arm = pytest.mark.skipif( diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 2e1d97326d4a..7fd21d2be46c 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -645,3 +645,21 @@ def run_moe(): with torch.inference_mode(): output = run_moe() ref_fused_moe.check_accuracy(output, ref_output) + + +class TestResolveMoeBackendSM121: + """Test that resolve_moe_backend returns CUTLASS (not TRTLLM) on SM121. + + SM120/SM121 lack tcgen05.mma instructions required by trtllm-gen kernels, + so the AUTO backend must fall back to CUTLASS. + """ + + def test_resolve_moe_backend_returns_cutlass_on_sm121(self): + from unittest.mock import patch + + with patch("tensorrt_llm._utils.get_sm_version", return_value=121): + result = ModelConfig.resolve_moe_backend("AUTO", "SomeArchitecture") + assert result == "CUTLASS", ( + f"Expected CUTLASS on SM121 but got {result}; " + "trtllm-gen kernels use tcgen05.mma which is unavailable on SM120/SM121" + ) 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 91024f5e4b77..c463e7ff5b1b 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -313,7 +313,7 @@ def test_moe_swiglu(dtype: str, num_tokens: int, top_k: int, tile_size: int): @pytest.mark.skipif( - get_sm_version() not in (100, 103), + get_sm_version() not in (100, 103, 120, 121), reason="This test is only supported on SM 100 and SM 103 GPUs", ) @pytest.mark.parametrize("tile_size", [128, 256])