diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index fcea8829442b..6cb73708b01c 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -82,8 +82,8 @@ if(FAST_BUILD) STATUS "FAST_BUILD enabled for kernels: using -O1 for CUDA compilation") endif() -# Marlin NVFP4: Hopper-only OBJECT library. Pinned to sm_90 so the global -# CMAKE_CUDA_ARCHITECTURES doesn't propagate. +# Marlin NVFP4: Ada (SM89, e.g. L40S) and Hopper (SM90) OBJECT library. Pinned +# to sm_89/sm_90 so the global CMAKE_CUDA_ARCHITECTURES doesn't propagate. file(GLOB_RECURSE MARLIN_SRC "marlin/*.cu" "marlin/*.cpp") if(MARLIN_SRC) add_library(marlin_src OBJECT ${MARLIN_SRC}) @@ -95,7 +95,7 @@ if(MARLIN_SRC) $ ) target_link_libraries(marlin_src PRIVATE trtllm_gen_fmha_interface) - set_cuda_architectures(marlin_src 90) + set_cuda_architectures(marlin_src 89 90) endif() add_library( diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin.cuh b/cpp/tensorrt_llm/kernels/marlin/marlin.cuh index 63864080b0f2..b04a42b66cc3 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin.cuh +++ b/cpp/tensorrt_llm/kernels/marlin/marlin.cuh @@ -34,6 +34,13 @@ #define MARLIN_NAMESPACE_NAME marlin #endif +// Ada Lovelace (SM89, e.g. L40S) and Hopper (SM90-99). +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890 && __CUDA_ARCH__ < 1000) +#define MARLIN_NVFP4_DEVICE_SUPPORTED 1 +#else +#define MARLIN_NVFP4_DEVICE_SUPPORTED 0 +#endif + namespace MARLIN_NAMESPACE_NAME { diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h index ad7d55471514..932baba29f63 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h @@ -25,6 +25,11 @@ namespace marlin_nvfp4 { +inline bool isMarlinNvfp4SmSupported(int sm) +{ + return sm >= 89 && sm < 100; +} + void dequantFp4Activations( void const* act_fp4, void const* act_sf, float const* alpha, void* act_bf16, int m, int k, cudaStream_t stream); diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu index f9cbbe3d64ab..185b6d370fb9 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu @@ -327,8 +327,8 @@ void marlinNvfp4Gemm(void const* act_bf16, void const* weight, void* output, voi bool use_fp32_reduce, cudaStream_t stream) { int const sm = tensorrt_llm::common::getSMVersion(); - TLLM_CHECK_WITH_INFO( - sm >= 90 && sm < 100, "Marlin NVFP4 GEMM is only supported on Hopper (SM 9.x); current SM = %d", sm); + TLLM_CHECK_WITH_INFO(isMarlinNvfp4SmSupported(sm), + "Marlin NVFP4 GEMM is only supported on Ada (SM89) and Hopper (SM90-99); current SM = %d", sm); int dev; cudaGetDevice(&dev); diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu index 35f28b309345..97a75344d150 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_gemm.cu @@ -253,8 +253,8 @@ void marlinNvfp4MoeGemmDispatcher(void const* A, void const* B, void* C, void* C cudaStream_t stream) { int const sm = tensorrt_llm::common::getSMVersion(); - TLLM_CHECK_WITH_INFO( - sm >= 90 && sm < 100, "Marlin NVFP4 MoE GEMM is only supported on Hopper (SM 9.x); current SM = %d", sm); + TLLM_CHECK_WITH_INFO(isMarlinNvfp4SmSupported(sm), + "Marlin NVFP4 MoE GEMM is only supported on Ada (SM89) and Hopper (SM90-99); current SM = %d", sm); int dev; cudaGetDevice(&dev); diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h index 9159e0feea61..2dc9d6390e6d 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_moe_template.h @@ -70,8 +70,8 @@ __global__ void Marlin(MARLIN_KERNEL_PARAMS); static_assert(std::is_same::value || std::is_same::value, \ "only float16 and bfloat16 is supported"); -// Empty kernel stub for non-Hopper device passes; see marlin.cuh. -#if defined(__CUDA_ARCH__) && !(__CUDA_ARCH__ >= 900 && __CUDA_ARCH__ < 1000) +// Empty kernel stub for unsupported device passes; see marlin.cuh. +#if defined(__CUDA_ARCH__) && !MARLIN_NVFP4_DEVICE_SUPPORTED template = 900 && __CUDA_ARCH__ < 1000) +// Empty kernel stub for unsupported device passes; see marlin.cuh. +#if defined(__CUDA_ARCH__) && !MARLIN_NVFP4_DEVICE_SUPPORTED template = 90 && sm < 100, "Marlin NVFP4 repack is only supported on Hopper (SM 9.x); current SM = %d", sm); + TLLM_CHECK_WITH_INFO(isMarlinNvfp4SmSupported(sm), + "Marlin NVFP4 repack is only supported on Ada (SM89) and Hopper (SM90-99); current SM = %d", sm); int blocks; int dev; diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 714c77be63ab..cd162744f72d 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -48,7 +48,7 @@ from ..utils import (ActivationType, deep_gemm_gen_tuning_buckets, fp4_scale_infer_shape, get_last_power_of_2_num_tokens_buckets, - last_positive_power_of_2) + is_nvfp4_marlin_supported_sm, last_positive_power_of_2) if IS_CUTLASS_DSL_AVAILABLE: from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ @@ -752,7 +752,7 @@ def forward( class MarlinNVFP4Runner(TunableRunner): - """Marlin-based NVFP4 GEMM for SM90 (Hopper). + """Marlin-based NVFP4 GEMM for SM89 (Ada, e.g. L40S) and SM90 (Hopper). Weights are eagerly repacked to Marlin tiled format during ``get_valid_tactics`` (before CUDA graph capture) so that ``forward()`` does @@ -761,8 +761,6 @@ class MarlinNVFP4Runner(TunableRunner): tuning_config = TuningConfig() # single tactic, no tuning - MIN_SM_VERSION = 90 - MAX_SM_VERSION = 99 # SM90-series only (Hopper) NVFP4_SCALE_VECTOR_SIZE = 16 def __init__(self, output_buffer_kind: int, output_dtype: torch.dtype): @@ -774,9 +772,7 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs) -> List[int]: if not torch.cuda.is_available(): return [] - capability = torch.cuda.get_device_capability(torch.device('cuda:0')) - sm_version = capability[0] * 10 + capability[1] - if sm_version < self.MIN_SM_VERSION or sm_version > self.MAX_SM_VERSION: + if not is_nvfp4_marlin_supported_sm(): return [] # Eagerly prepare Marlin weights so that forward() never allocates @@ -1024,7 +1020,7 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], tactics = [] act_fp4, weight, act_sf, weight_scale, alpha = inputs - # Add Marlin tactics (SM90 Hopper only) — users must opt-in explicitly + # Add Marlin tactics (SM89 Ada / SM90 Hopper) — users must opt-in explicitly # by listing "marlin" in ``allowed_backends``. if self._is_backend_allowed("marlin"): marlin_runner = MarlinNVFP4Runner(self.output_buffer_kind, @@ -1036,7 +1032,7 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], elif self._is_only_backend("marlin"): sm_version = get_sm_version() raise ValueError( - f"Marlin backend requires SM 90-99 (Hopper), but got SM " + f"Marlin backend requires SM 89-99 (Ada L40S or Hopper), but got SM " f"{sm_version}. Please add other backends to " "allowed_backends.") @@ -1153,12 +1149,12 @@ def forward( ) -> torch.Tensor: # Handle fallback tactic on cache miss if tactic == -1: - # Prefer marlin on Hopper (SM90) when explicitly allowed, cutlass + # Prefer marlin on Ada/Hopper (SM89-99) when explicitly allowed, cutlass # otherwise, falling back to whatever backend is available. assert len( self.allowed_backends) > 0, "No allowed backends available" - sm_version = get_sm_version() - if "marlin" in self.allowed_backends and 90 <= sm_version <= 99: + if ("marlin" in self.allowed_backends + and is_nvfp4_marlin_supported_sm()): tactic = ("marlin", -1) elif "cutlass" in self.allowed_backends: tactic = ("cutlass", -1) @@ -1219,7 +1215,7 @@ def nvfp4_gemm( - cuBLASLt: Heuristic-based algorithms from cuBLASLt library - CuteDSL: Blackwell-optimized persistent kernels (when available and inputs are valid) - CUDA Core: CUDA Core implementation (requires SM >= 100 and M <= 8) - - Marlin: Hopper W4A16 NVFP4 implementation (requires SM 90-99) + - Marlin: Ada/Hopper W4A16 NVFP4 implementation (requires SM 89-99) The AutoTuner profiles all available backends during the first run and caches the best choice for each input shape. Subsequent calls use the cached selection 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 77895edb07ff..df028be20263 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -150,7 +150,7 @@ Still on old path (standalone, with embedded communication): | `fused_moe_cute_dsl_b12x.py` | `CuteDslB12xFusedMoE` | SM120/SM121 | NVFP4 hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode — best perf on RTX PRO 6000 (SM120) and DGX Spark (SM121); select via the `CUTEDSL` backend path (auto-promoted when flashinfer is importable) | `EXTERNAL_COMM` | | `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` | SM90 only | W4A16 NVFP4 on 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) fall back to Cutlass in `get_moe_cls`; no dynamic EPLB | `EXTERNAL_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) fall back to Cutlass in `get_moe_cls`; 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_wide_ep.py` | `WideEPMoE` | All GPUs | Deprecating — use ConfigurableMoE instead | (legacy path) | | `fused_moe_vanilla.py` | `VanillaMoE` | All devices | Reference / debugging only | (legacy path) | @@ -203,7 +203,7 @@ Each backend's `can_implement(quant_algo, dtype_activation, swiglu_gptoss_style, | Unquantized (BF16/FP16) | Y (SM80+) | N | N | N | N | N | N | Y (SM90, BF16) | N | Y | Y | | FP8 QDQ | Y (SM89+) | N | N | N | N | N | N | Y (SM90) | N | Y | Y | | FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | Y (SM100/103) | N | N | N | N | Y | Y | -| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM90, W4A16) | Y | Y | +| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM89-SM99, W4A16) | Y | Y | | W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | N | N | | W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | N | | W4A8 MXFP4 FP8 | Y (SM100/103) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | N | diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 3761fab5a06e..94c1b821271e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -64,7 +64,7 @@ def get_moe_cls( quant_config = override_quant_config layer_prefix = f"[layer_idx={layer_idx}] " if layer_idx is not None else "" if moe_backend.upper() == "MARLIN": - # Marlin MoE is a Hopper-specific NVFP4 W4A16 backend. Layers without + # Marlin MoE is an Ada/Hopper NVFP4 W4A16 backend. Layers without # NVFP4 quantization (e.g. deliberately-unquantized MTP draft layers in # MIXED_PRECISION checkpoints) fall back to CutlassFusedMoE, matching # the CUTEDSL / DENSEGEMM / MEGAMOE_* fallback behavior below. diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py index c4728be2b140..91fd4919ed9d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py @@ -12,7 +12,7 @@ # 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. -"""Marlin-based MoE backend for NVFP4 on SM90 (Hopper). +"""Marlin-based MoE backend for NVFP4 on SM89 (Ada, e.g. L40S) and SM90 (Hopper). Uses a fused ``marlin_nvfp4_moe_gemm`` CUDA kernel that processes ALL experts in a single launch. W4A16 approach: BF16 activations + FP4 weights, @@ -25,7 +25,7 @@ import torch import torch.nn.functional as F -from tensorrt_llm._torch.utils import Fp4QuantizedTensor +from tensorrt_llm._torch.utils import Fp4QuantizedTensor, is_nvfp4_marlin_supported_sm from tensorrt_llm._utils import get_sm_version from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -44,7 +44,7 @@ def _has_fused_moe_kernel() -> bool: class MarlinFusedMoE(CutlassFusedMoE): - """MoE backend using Marlin W4A16 NVFP4 GEMM for SM90 (Hopper). + """MoE backend using Marlin W4A16 NVFP4 GEMM for SM89/SM90. Uses ``marlin_nvfp4_moe_gemm`` with BF16 activations to process all experts in a single kernel launch via sorted token dispatch. In-kernel topk_weights @@ -54,7 +54,7 @@ class MarlinFusedMoE(CutlassFusedMoE): _QUANT_SUPPORT_TABLE = { QuantAlgo.NVFP4: { - "sm_constraint": ("in", {90}), + "sm_constraint": ("in", {89, 90}), "dtypes": {torch.bfloat16}, }, } @@ -73,9 +73,9 @@ def can_implement( f"MarlinFusedMoE only supports NVFP4 (got quant_algo={quant_algo})" ) - if sm_version != 90: + if not is_nvfp4_marlin_supported_sm(sm_version): return _warn_and_return( - f"MarlinFusedMoE only supports SM90 (Hopper), got SM{sm_version}" + f"MarlinFusedMoE only supports SM89 (Ada) and SM90 (Hopper), got SM{sm_version}" ) if swiglu_gptoss_style: diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 10504cf6ecd7..eb34a624a8f2 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -2896,7 +2896,7 @@ def apply(self, module: Linear, input: torch.Tensor, class MarlinNVFP4LinearMethod(NVFP4LinearMethod): - """NVFP4 Linear method backed by the Marlin W4A16 kernel (Hopper only).""" + """NVFP4 Linear method backed by the Marlin W4A16 kernel (SM89/SM90).""" def transform_weights(self, module: Linear) -> None: from tensorrt_llm.quantization.utils import marlin_utils diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index ba0453fed790..5e3eb254461f 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -110,13 +110,20 @@ def get_model_extra_attrs(): return getattr(_model_extra_attrs, 'attrs', None) +def is_nvfp4_marlin_supported_sm(sm_version: int | None = None) -> bool: + """Return True on Ada Lovelace (SM89, e.g. L40S) and Hopper (SM90-99).""" + if sm_version is None: + sm_version = get_sm_version() + return 89 <= sm_version < 100 + + def is_nvfp4_marlin_enabled() -> bool: - is_hopper = get_sm_version() == 90 + is_supported_sm = is_nvfp4_marlin_supported_sm() has_marlin_kernel = hasattr(torch.ops.trtllm, "marlin_nvfp4_gemm") attrs = get_model_extra_attrs() is_marlin_specified = attrs is not None and "marlin" in attrs.get( 'nvfp4_gemm_allowed_backends', []) - return is_hopper and has_marlin_kernel and is_marlin_specified + return is_supported_sm and has_marlin_kernel and is_marlin_specified @contextlib.contextmanager diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 378aa74751db..4abf764b4824 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -39,7 +39,8 @@ # isort: on from tensorrt_llm.quantization import QuantAlgo -from ..conftest import (check_device_contain, get_device_count, llm_models_root, +from ..conftest import (check_device_contain, get_device_count, + get_mpi_world_size, llm_models_root, parametrize_with_ids, skip_no_hopper, skip_no_mxfp4_swizzle, skip_post_blackwell, skip_post_hopper, skip_pre_ada, skip_pre_blackwell, @@ -6882,12 +6883,15 @@ def test_fp8(self): task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) - @skip_pre_hopper + @skip_pre_ada @skip_post_hopper - @pytest.mark.skip_less_device_memory(80000) - @pytest.mark.skip_less_mpi_world_size(8) + @pytest.mark.skip_less_device_memory(45000) @parametrize_with_ids("tp_size", [1, 2, 4, 8]) def test_nvfp4_marlin_multi_gpus(self, tp_size): + if get_mpi_world_size() < tp_size: + pytest.skip( + f"Requires MPI world size >= {tp_size} for TP={tp_size}") + ep_size = tp_size with LLM( f"{llm_models_root()}/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", diff --git a/tests/integration/test_lists/test-db/l0_l40s.yml b/tests/integration/test_lists/test-db/l0_l40s.yml index bf71ba62fc8c..c70644f24754 100644 --- a/tests/integration/test_lists/test-db/l0_l40s.yml +++ b/tests/integration/test_lists/test-db/l0_l40s.yml @@ -36,6 +36,8 @@ l0_l40s: # AutoDeploy: Nemotron-Nano-V3 on Ada uses flashinfer; trtllm has no # (E4M3 input, BF16 output, paged_kv, head_dim=128, sm_89) FMHA cubin. - accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[fp8-1-attn_dp_off-flashinfer] + # PyTorch: force NVFP4 Marlin on Ada to cover the SM89-compiled Marlin kernels. + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Nano::test_nvfp4_marlin_multi_gpus[tp_size=1] # ------------- Encoder-decoder tests --------------- # Representative pre-merge set: eager smoke, CUDA-graph greedy/beam, kv-v2, # overlap scheduler, mBART, mixed batches. diff --git a/tests/unittest/_torch/modules/moe/moe_test_utils.py b/tests/unittest/_torch/modules/moe/moe_test_utils.py index 59aeb7fb8333..25c3e2cf5b43 100644 --- a/tests/unittest/_torch/modules/moe/moe_test_utils.py +++ b/tests/unittest/_torch/modules/moe/moe_test_utils.py @@ -1462,15 +1462,16 @@ def should_skip_to_accelerate_ci( ): return "[CI accel] Skip unquantized (quant=None) in CI" - # --- Rule 0a: MARLIN backend only runs NVFP4 on Hopper (SM90) --- + # --- Rule 0a: MARLIN backend only runs NVFP4 on Ada/Hopper (SM89/SM90) --- if backend_type == MoeBackendType.MARLIN: + from tensorrt_llm._torch.utils import is_nvfp4_marlin_supported_sm from tensorrt_llm._utils import get_sm_version if quant_algo != QuantAlgo.NVFP4: return f"[CI accel] MARLIN only tests NVFP4 in CI (got {quant_algo})" sm_version = get_sm_version() - if sm_version != 90: - return f"[CI accel] MARLIN only runs on Hopper (SM90) in CI (got SM{sm_version})" + if not is_nvfp4_marlin_supported_sm(sm_version): + return f"[CI accel] MARLIN only runs on Ada/Hopper (SM89-99) in CI (got SM{sm_version})" # Any e256-class model_config triggers CI Rule-1 minimal coverage: # the full dtype x seq_len x swiglu x routing matrix on e256 models diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 18888cab9c3d..0b9ee7ffc9b6 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -8,7 +8,8 @@ from tensorrt_llm._torch.autotuner import autotune from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._torch.modules.linear import Linear -from tensorrt_llm._torch.utils import model_extra_attrs +from tensorrt_llm._torch.utils import (is_nvfp4_marlin_supported_sm, + model_extra_attrs) from tensorrt_llm._utils import get_sm_version from tensorrt_llm.math_utils import pad_up from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -751,8 +752,8 @@ def test_fp4_linear_cuda_core(dtype, mnk): @pytest.mark.skipif( - get_sm_version() < 90 or get_sm_version() >= 100, - reason="Marlin NVFP4 backend runs Hopper", + not is_nvfp4_marlin_supported_sm(), + reason="Marlin NVFP4 backend runs on Ada (SM89) and Hopper (SM90-99)", ) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py b/tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py index 765fe516bd19..9a703bc3d82f 100644 --- a/tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py +++ b/tests/unittest/_torch/thop/serial/test_nvfp4_marlin_gemm.py @@ -24,7 +24,7 @@ import torch from parameterized import parameterized -from utils.util import skip_non_hopper_unittest, unittest_name_func +from utils.util import skip_neither_ada_nor_hopper_unittest, unittest_name_func import tensorrt_llm # noqa: F401 # registers torch.ops.tensorrt_llm / torch.ops.trtllm ops @@ -65,7 +65,7 @@ class TestNvfp4MarlinGemm(unittest.TestCase): list(product([1024, 2048], [1024, 2048], [1, 8, 128], [16], [1.0, 2.0], ["nvfp4", "bf16"])), name_func=unittest_name_func, ) - @skip_non_hopper_unittest + @skip_neither_ada_nor_hopper_unittest def test_nvfp4_marlin_gemm( self, input_dim, output_dim, batch_size, sf_vec_size, alpha, act_dtype ):