diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 48e7d214a612..aa27f740576e 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -808,13 +808,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 39a7289fee60..6a4c99189261 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -250,12 +250,12 @@ def resolve_moe_backend(moe_backend: str, architecture: str) -> str: if architecture == "GptOssForCausalLM": sm_version = get_sm_version() # Select the best performing backend based on SM version - if 100 <= sm_version < 120: # Blackwell + if 100 <= sm_version < 120 or sm_version in (120, 121): # Blackwell return "TRTLLM" elif 90 <= sm_version < 100: # Hopper return "TRITON" else: - return "CUTLASS" # Fallback to CUTLASS for other SM versions (e.g., SM120) + return "CUTLASS" # Fallback for other SM versions return "CUTLASS" 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..693332ffa1ff 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 +from tensorrt_llm._utils import get_sm_version, is_blackwell, is_sm_100f 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 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). @@ -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 - SM in {100, 103, 120, 121} (Blackwell family) if quant_algo == QuantAlgo.NVFP4: - if sm_version not in {100, 103}: + if sm_version not in {100, 103, 120, 121}: return _warn_and_return( - f"NVFP4 requires SM100 or SM103, got SM{sm_version}") + f"NVFP4 requires Blackwell (SM100/103/120/121), 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 956542fd6ff1..7b37569efd85 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 @@ -60,7 +60,7 @@ class TRTLLMGenFusedMoE(MoE): aux_stream_dict (Optional[Dict[AuxStreamType, torch.cuda.Stream]]): Auxiliary CUDA streams for overlapping. MoE torch custom op: - Only support min-latency mode now (SM100 Blackwell only). + Only support min-latency mode now (Blackwell: SM100/103/120/121). Quant: fp8 block scales quant and nvfp4 quant and w4a16_mxfp4 quant FusedMoE Op: routing(topK, etc.) + scatter + gemm1 + swiglu + gemm2 + finalize MoeRoute @@ -105,7 +105,7 @@ def can_implement( """ Check if TRTLLMGenFusedMoE can implement the given quantization algorithm. - TRTLLMGenFusedMoE only supports SM in {100, 103} and the following quantizations: + TRTLLMGenFusedMoE only supports SM in {100, 103, 120, 121} and the following quantizations: - NVFP4 - FP8_BLOCK_SCALES - W4A8_NVFP4_FP8 @@ -129,10 +129,10 @@ def can_implement( sm_version = get_sm_version() - # TRTLLMGenFusedMoE requires SM in {100, 103} - if sm_version not in {100, 103}: + # TRTLLMGenFusedMoE requires SM in {100, 103, 120, 121} (Blackwell family) + if sm_version not in {100, 103, 120, 121}: return _warn_and_return( - f"TRTLLMGenFusedMoE requires SM100 or SM103, got SM{sm_version}" + f"TRTLLMGenFusedMoE requires Blackwell (SM100/103/120/121), got SM{sm_version}" ) # Check dtype_activation: only bfloat16 is supported @@ -201,11 +201,6 @@ def __init__( activation_type=activation_type, ) - sm_version = get_sm_version() - if sm_version >= 120: - raise NotImplementedError( - "TRTLLMGenFusedMoE does not support SM120 and above.") - assert not self.smart_router, "Smart router is not supported in TRTLLMGenFusedMoE." # Note: Load balancer initialization is handled by base class _init_load_balancer() diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index c53b96516d19..7c33c77df6e7 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -781,6 +781,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..82948213b539 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: