diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index a8dae5de5b3f..c6c6fe4a292a 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 1993-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 @@ -67,7 +67,7 @@ list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "compressorKernels/.*") -# Marlin is built as its own Hopper-only OBJECT library below. +# Marlin is built as its own architecture-scoped OBJECT library below. list(FILTER SRC_CPP EXCLUDE REGEX "marlin/.*") list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") @@ -88,8 +88,8 @@ if(FAST_BUILD) STATUS "FAST_BUILD enabled for kernels: using -O1 for CUDA compilation") endif() -# 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. +# Marlin NVFP4: build Ada/Hopper kernels and the dense W4A16/repack kernels used +# on SM12x. Unsupported device passes emit empty kernel stubs. file(GLOB_RECURSE MARLIN_SRC "marlin/*.cu" "marlin/*.cpp") if(MARLIN_SRC) add_library(marlin_src OBJECT ${MARLIN_SRC}) @@ -101,7 +101,7 @@ if(MARLIN_SRC) $ ) target_link_libraries(marlin_src PRIVATE trtllm_gen_fmha_interface) - set_cuda_architectures(marlin_src 89 90) + set_cuda_architectures(marlin_src 89 90 120f) endif() add_library( diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin.cuh b/cpp/tensorrt_llm/kernels/marlin/marlin.cuh index b04a42b66cc3..8b5c40a52dcd 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin.cuh +++ b/cpp/tensorrt_llm/kernels/marlin/marlin.cuh @@ -41,6 +41,15 @@ #define MARLIN_NVFP4_DEVICE_SUPPORTED 0 #endif +// Dense Marlin also supports Blackwell GeForce (SM120/121). Keep this +// separate from MARLIN_NVFP4_DEVICE_SUPPORTED because Marlin MoE does not. +#if defined(__CUDA_ARCH__) \ + && ((__CUDA_ARCH__ >= 890 && __CUDA_ARCH__ < 1000) || (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)) +#define MARLIN_NVFP4_DENSE_DEVICE_SUPPORTED 1 +#else +#define MARLIN_NVFP4_DENSE_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 932baba29f63..13194fd76e29 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4.h @@ -30,6 +30,11 @@ inline bool isMarlinNvfp4SmSupported(int sm) return sm >= 89 && sm < 100; } +inline bool isMarlinNvfp4DenseSmSupported(int sm) +{ + return isMarlinNvfp4SmSupported(sm) || sm == 120 || sm == 121; +} + 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 185b6d370fb9..294d6c75f858 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_gemm.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -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(isMarlinNvfp4SmSupported(sm), - "Marlin NVFP4 GEMM is only supported on Ada (SM89) and Hopper (SM90-99); current SM = %d", sm); + TLLM_CHECK_WITH_INFO(isMarlinNvfp4DenseSmSupported(sm), + "Marlin NVFP4 GEMM is only supported on SM89, SM90-99, and SM120/121; current SM = %d", sm); int dev; cudaGetDevice(&dev); diff --git a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h index 1eb5b1657d88..b2c7bbaa1f85 100644 --- a/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h +++ b/cpp/tensorrt_llm/kernels/marlin/marlin_nvfp4_template.h @@ -49,8 +49,8 @@ namespace MARLIN_NAMESPACE_NAME { -// Empty kernel stub for unsupported device passes; see marlin.cuh. -#if defined(__CUDA_ARCH__) && !MARLIN_NVFP4_DEVICE_SUPPORTED +// Empty kernel stub outside the architectures supported by dense Marlin. +#if defined(__CUDA_ARCH__) && !MARLIN_NVFP4_DENSE_DEVICE_SUPPORTED template 1 or mapping.dp_size > 1: + logger.warning( + "CuteDslB12xFusedMoE does not support expert " + "parallelism or attention-DP/all-to-all; selecting " + "CutlassFusedMoE.") + return CutlassFusedMoE try: import flashinfer # noqa: F401 logger.info( @@ -101,13 +112,24 @@ def get_moe_cls( except ImportError: logger.warning( "CuteDslB12xFusedMoE eligible (SM%d + NVFP4) " - "but flashinfer is not importable; using CuteDslFusedMoE.", + "but flashinfer is not importable; using %s.", sm_version, + "CutlassFusedMoE" + if has_w4a16_nvfp4 else "CuteDslFusedMoE", ) + if has_w4a16_nvfp4: + return CutlassFusedMoE + elif has_w4a16_nvfp4: + logger.warning( + "CuteDslB12xFusedMoE requires SM120/121 for W4A16_NVFP4 " + "(got SM%d). Using CutlassFusedMoE.", + sm_version, + ) + return CutlassFusedMoE return CuteDslFusedMoE else: logger.warning( - f"{layer_prefix}CuteDslFusedMoE only supports fp8_block_scales and nvfp4. " + f"{layer_prefix}CuteDslFusedMoE only supports fp8_block_scales, nvfp4, and w4a16_nvfp4. " f"Check out details in quant_config: {quant_config}. Using CutlassFusedMoE instead." ) return CutlassFusedMoE diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py index f9cc59724ec1..1c1a553fe5c9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py @@ -44,24 +44,25 @@ class CuteDslB12xFusedMoE(CuteDslFusedMoE): - """Hybrid CUTLASS-prefill / b12x-decode NVFP4 fused-MoE backend for SM120 / SM121. + """B12x NVFP4 fused-MoE backend for SM120 / SM121. Member of the cuteDSL backend family: the decode kernel (``flashinfer.B12xMoEWrapper.run``) is JIT-compiled CuTe DSL, so the backend slots in next to :class:`CuteDslFusedMoE` (which targets SM100 / - SM103). The hybrid prefill path still routes through the C++ CUTLASS - NVFP4 GroupGEMM via explicit :class:`CutlassFusedMoE` method calls; the - parent class on the MRO does not change which kernels execute, only - where the b12x backend sits in the family. + SM103). Plain NVFP4 prefill can route through the C++ CUTLASS NVFP4 + GroupGEMM via explicit :class:`CutlassFusedMoE` method calls; the parent + class on the MRO does not change which kernels execute, only where the + b12x backend sits in the family. Composition (see ``MOE_DEVELOPER_GUIDE.md`` for the full explainer): - - **Prefill (``m >= _PREFILL_VIA_CUTLASS_THRESHOLD``)** explicitly + - **NVFP4 prefill (``m >= _PREFILL_VIA_CUTLASS_THRESHOLD``)** explicitly invokes :class:`CutlassFusedMoE` NVFP4 GroupGEMM. The b12x kernel's 12-CTA-per-token MMA pattern is suboptimal at large ``m``. - **Decode (``m < _PREFILL_VIA_CUTLASS_THRESHOLD``)** dispatches to FlashInfer's ``B12xMoEWrapper.run`` — a kernel purpose-built for ``m=1`` / small routed-row counts. + - **W4A16_NVFP4** stays on the b12x path for both prefill and decode. NVFP4 weights are loaded via :class:`NVFP4CuteDslB12xFusedMoEMethod` (an :class:`NVFP4CutlassFusedMoEMethod` subclass returned by @@ -79,8 +80,9 @@ class CuteDslB12xFusedMoE(CuteDslFusedMoE): The backend hard-rejects EP (b12x has no dispatch / combine kernel), MoE alltoall, ``Fp4QuantizedTensor`` input, ``swiglu_gptoss_style`` biased SwiGLU, and activations outside ``{Relu2, Swiglu}``. It is - selected on the ``CUTEDSL`` MoE path when SM120 / SM121 + NVFP4 + - flashinfer-importable gates pass (see ``create_moe.get_moe_cls``). + selected on the ``CUTEDSL`` MoE path when SM120 / SM121 + NVFP4 or + W4A16_NVFP4 + flashinfer-importable gates pass (see + ``create_moe.get_moe_cls``). """ # SM versions on which the FlashInfer b12x NVFP4 MoE kernel is available. @@ -105,9 +107,9 @@ def can_implement( if sm_version not in cls._SUPPORTED_SM_VERSIONS: sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) return _warn_and_return(f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}") - if quant_algo != QuantAlgo.NVFP4: + if quant_algo not in {QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4}: return _warn_and_return( - f"CuteDslB12xFusedMoE only supports NVFP4 quantization " + f"CuteDslB12xFusedMoE only supports NVFP4 or W4A16_NVFP4 quantization " f"(got quant_algo={quant_algo})" ) if dtype_activation not in {torch.float16, torch.bfloat16}: @@ -167,9 +169,12 @@ def _get_quant_method(self): def _route_to_cutlass(self, x) -> bool: """Return ``True`` iff this call should fall back to the inherited - CUTLASS path (prefill chunk). ``Fp4QuantizedTensor`` inputs always - stay on the b12x path (which rejects them) so the existing error - message is preserved.""" + CUTLASS path (NVFP4 prefill chunk). ``Fp4QuantizedTensor`` inputs + always stay on the b12x path (which rejects them) so the existing + error message is preserved.""" + quant_config = getattr(self, "quant_config", None) + if quant_config is not None and quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: + return False return isinstance(x, torch.Tensor) and x.shape[0] >= self._PREFILL_VIA_CUTLASS_THRESHOLD # ``post_load_weights`` is inherited from ``CutlassFusedMoE`` and @@ -190,12 +195,12 @@ def quantize_input( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Hybrid dispatch entrypoint for activation handling. - Prefill chunks (``x.shape[0] >= _PREFILL_VIA_CUTLASS_THRESHOLD``) take - the inherited :meth:`CutlassFusedMoE.quantize_input` path so the - downstream ``run_moe`` can call CUTLASS NVFP4 GroupGEMM. Decode - chunks pass through unchanged because b12x quantizes activations - internally (consumes a bf16 / fp16 ``x`` and produces its own scale - factors). + NVFP4 prefill chunks take the inherited + :meth:`CutlassFusedMoE.quantize_input` path so the downstream + ``run_moe`` can call CUTLASS NVFP4 GroupGEMM. Decode chunks and + W4A16_NVFP4 chunks pass through unchanged because b12x quantizes + activations internally (consumes a bf16 / fp16 ``x`` and produces its + own scale factors). """ if self._route_to_cutlass(x): return CutlassFusedMoE.quantize_input( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index cad3629f0d44..70b87aa7d488 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -813,6 +813,8 @@ def _get_quant_method(self): return FP8QDQFusedMoEMethod() elif self.quant_config.layer_quant_mode.has_fp8_block_scales(): return DeepSeekFP8BlockScalesFusedMoEMethod() + elif self.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: + return W4A16NVFP4CutlassFusedMoEMethod() elif self.quant_config.layer_quant_mode.has_nvfp4(): return NVFP4CutlassFusedMoEMethod() elif self.quant_config.layer_quant_mode.is_int4_weight_only_per_group( 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 c01a99da51d2..5dc95306f55c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_marlin.py @@ -57,6 +57,10 @@ class MarlinFusedMoE(CutlassFusedMoE): "sm_constraint": ("in", set(range(89, 100))), "dtypes": {torch.bfloat16}, }, + QuantAlgo.W4A16_NVFP4: { + "sm_constraint": ("in", set(range(89, 100))), + "dtypes": {torch.bfloat16}, + }, } @classmethod @@ -68,9 +72,9 @@ def can_implement( ) -> Tuple[bool, Optional[str]]: sm_version = get_sm_version() - if quant_algo != QuantAlgo.NVFP4: + if quant_algo not in cls._QUANT_SUPPORT_TABLE: return _warn_and_return( - f"MarlinFusedMoE only supports NVFP4 (got quant_algo={quant_algo})" + f"MarlinFusedMoE only supports NVFP4 or W4A16_NVFP4 (got quant_algo={quant_algo})" ) if not is_nvfp4_marlin_supported_sm(sm_version): diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index 1e7ecd06845f..dc082eef24c6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -1078,6 +1078,11 @@ def has_nvfp4(self): return self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( ) + @property + def has_nvfp4_activation_quantization(self): + assert self._weights_created + return self.quant_method.quantizes_nvfp4_activations + @property def has_w4a8_nvfp4_fp8(self): assert self._weights_created diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 943e629ad396..26e074669882 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -26,6 +26,7 @@ from tensorrt_llm._utils import get_sm_version, is_device_integrated, is_sm_100f from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import QuantAlgo from tensorrt_llm.quantization.functional import \ preprocess_weights_for_mixed_gemm from tensorrt_llm.quantization.utils.fp4_utils import ( @@ -236,6 +237,9 @@ class FusedMoEMethodBase(ABC): weight_alignment: int = 1 """int: Required byte alignment for MoE weight tensors.""" + quantizes_nvfp4_activations: bool = False + """Whether this method converts high-precision activations to NVFP4.""" + eplb_support_status: EplbSupportStatus = EplbSupportStatus.NOT_SUPPORTED """EplbSupportStatus: Online EPLB support status for this quantization method. @@ -2151,6 +2155,7 @@ class NVFP4FusedMoEMethod(FusedMoEMethodBase): Base class for NVFP4 fused MoE methods for all backends. """ eplb_support_status = EplbSupportStatus.SUPPORTED + quantizes_nvfp4_activations = True # Whether raw per-expert block-scale staging is an EPLB migration # target. Children that migrate derived formats and free the raw @@ -3027,6 +3032,9 @@ class NVFP4MarlinFusedMoEMethod(NVFP4CutlassFusedMoEMethod): raw ``weight_scale_2`` values. """ + # BF16 activations in, so the NVFP4FusedMoEMethod default does not hold. + quantizes_nvfp4_activations = False + # Marlin's ``transform_weights`` repacks weights into Marlin tiled format # and rebuilds the module parameters, which is incompatible with dynamic # EPLB weight migration. @@ -3132,6 +3140,8 @@ class W4A16NVFP4CutlassFusedMoEMethod(NVFP4CutlassFusedMoEMethod): into a static [E_total, N, K] workspace, then runs the bf16 ``fused_moe``. """ + quantizes_nvfp4_activations = False + def process_weights_after_loading(self, module: torch.nn.Module): super().process_weights_after_loading(module) @@ -3140,8 +3150,6 @@ def process_weights_after_loading(self, module: torch.nn.Module): # block_scale_interleave_reverse accepts. def _unswizzle_inplace(scale_param: torch.nn.Parameter): sf_view = scale_param.data.view(float4_sf_dtype) - E, pad_rows, pad_cols = (sf_view.shape[0], sf_view.shape[1], - sf_view.shape[2]) linear = torch.ops.trtllm.block_scale_interleave_reverse(sf_view) scale_param.data.view(float4_sf_dtype).copy_(linear) @@ -3359,6 +3367,10 @@ def transform_weights(self, module: torch.nn.Module) -> None: # layers across the model share one wrapper-owned output buffer. from .fused_moe_cute_dsl_b12x import _SHARED_MOE_OUTPUT_BUF + quant_config = getattr(module, "quant_config", None) + is_w4a16_nvfp4 = (quant_config is not None + and quant_config.quant_algo == QuantAlgo.W4A16_NVFP4) + num_local_experts = module.w3_w1_weight.shape[0] # Tensor shapes use the *padded* per-rank dims because TP partitions # may pad ``intermediate_size`` up to a kernel-friendly boundary. @@ -3388,16 +3400,20 @@ def transform_weights(self, module: torch.nn.Module) -> None: w2_w_scale_2 = (module.fc2_alpha * module.fc2_input_scale).to( torch.float32) - w1_sf_fp8_norm = module.w3_w1_weight_scale.view( - torch.float8_e4m3fn).float() - w2_sf_fp8_norm = module.w2_weight_scale.view( - torch.float8_e4m3fn).float() + w1_sf_fp8_src = module.w3_w1_weight_scale.view(torch.float8_e4m3fn) + w2_sf_fp8_src = module.w2_weight_scale.view(torch.float8_e4m3fn) + w1_sf_fp8_norm = w1_sf_fp8_src.float() + w2_sf_fp8_norm = w2_sf_fp8_src.float() - # Broadcast per-expert scalar over the trailing dims (E, *). - bcast1 = w1_w_scale_2.view(-1, *([1] * (w1_sf_fp8_norm.dim() - 1))) - bcast2 = w2_w_scale_2.view(-1, *([1] * (w2_sf_fp8_norm.dim() - 1))) - w1_sf_fp8 = (w1_sf_fp8_norm * bcast1).to(torch.float8_e4m3fn) - w2_sf_fp8 = (w2_sf_fp8_norm * bcast2).to(torch.float8_e4m3fn) + if is_w4a16_nvfp4: + w1_sf_fp8 = w1_sf_fp8_src + w2_sf_fp8 = w2_sf_fp8_src + else: + # Broadcast per-expert scalar over the trailing dims (E, *). + bcast1 = w1_w_scale_2.view(-1, *([1] * (w1_sf_fp8_norm.dim() - 1))) + bcast2 = w2_w_scale_2.view(-1, *([1] * (w2_sf_fp8_norm.dim() - 1))) + w1_sf_fp8 = (w1_sf_fp8_norm * bcast1).to(torch.float8_e4m3fn) + w2_sf_fp8 = (w2_sf_fp8_norm * bcast2).to(torch.float8_e4m3fn) w1_sf_b12x = convert_sf_to_mma_layout(w1_sf_fp8, m=w3w1_out_dim, @@ -3408,11 +3424,20 @@ def transform_weights(self, module: torch.nn.Module) -> None: k=w2_in_dim, num_groups=num_local_experts) - w1_alpha_b12x = ((1.0 / module.fc31_input_scale).expand( - module.num_experts).to(torch.float32).contiguous()) - w2_alpha_b12x = ((1.0 / module.fc2_input_scale).expand( - module.num_experts).to(torch.float32).contiguous()) - fc2_input_scale_b12x = (1.0 / module.fc2_input_scale).to(torch.float32) + if is_w4a16_nvfp4: + # W4A16 path: BF16/FP16 activations multiplied by FP4 weights. + # FlashInfer's W4A16 packer expects the ModelOpt scale contract: + # normalized FP8 block scales plus per-expert ``weight_global_scale``. + w1_alpha_b12x = w1_w_scale_2.to(torch.float32).contiguous() + w2_alpha_b12x = w2_w_scale_2.to(torch.float32).contiguous() + fc2_input_scale_b12x = None + else: + w1_alpha_b12x = ((1.0 / module.fc31_input_scale).expand( + module.num_experts).to(torch.float32).contiguous()) + w2_alpha_b12x = ((1.0 / module.fc2_input_scale).expand( + module.num_experts).to(torch.float32).contiguous()) + fc2_input_scale_b12x = (1.0 / module.fc2_input_scale).to( + torch.float32) # TRT-LLM packs 16 FP4 values per int64. flashinfer's internal # ``view(torch.float4_e2m1fn_x2)`` requires byte-contiguous storage @@ -3435,14 +3460,20 @@ def transform_weights(self, module: torch.nn.Module) -> None: f"{ActivationType(module.activation_type).name}; " f"supported: {supported}.") + # The model config may carry the logical intermediate size while the + # NVFP4 weight tensors are padded for kernel alignment. FlashInfer's + # CUDA-graph workspace must match the stored tensors. + b12x_intermediate_size = w2_in_dim + module.b12x_wrapper = B12xMoEWrapper( num_experts=module.num_experts, top_k=module.routing_method.experts_per_token, hidden_size=module.hidden_size, - intermediate_size=module.intermediate_size_per_partition, + intermediate_size=b12x_intermediate_size, use_cuda_graph=getattr(module, "_b12x_use_cuda_graph", False), max_num_tokens=module.moe_max_num_tokens, activation=self._ACTIVATION_MAP[module.activation_type], + quant_mode="w4a16" if is_w4a16_nvfp4 else "nvfp4", ) # Replace the wrapper's per-instance output buffer with a shared one. @@ -3463,10 +3494,11 @@ def transform_weights(self, module: torch.nn.Module) -> None: logger.info_once( f"NVFP4CuteDslB12xFusedMoEMethod active: hidden={module.hidden_size}, " - f"intermediate={module.intermediate_size_per_partition}, " + f"intermediate={b12x_intermediate_size}, " f"experts={module.num_experts}, top_k=" f"{module.routing_method.experts_per_token}, " - f"activation={self._ACTIVATION_MAP[module.activation_type]}.", + f"activation={self._ACTIVATION_MAP[module.activation_type]}, " + f"quant_mode={'w4a16' if is_w4a16_nvfp4 else 'nvfp4'}.", key="cute_dsl_b12x_moe_active", ) diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index d5787ccd0044..a9655ce45f2b 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -12,7 +12,8 @@ from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType from ..utils import Fp4QuantizedTensor -from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig +from .linear import (Linear, TensorParallelMode, WeightMode, + WeightsLoadingConfig, is_static_nvfp4_input_eligible) from .swiglu import swiglu @@ -219,7 +220,8 @@ def _can_fuse_gate_up_swiglu(self): - gate_up_proj has no bias (bias not supported in fused kernel) """ return (self.use_cute_dsl_blockscaling_mm and self.activation == F.silu - and self._is_plain_swiglu() and self.gate_up_proj.has_nvfp4 + and self._is_plain_swiglu() + and self.gate_up_proj.has_nvfp4_activation_quantization and not self.gate_up_proj.has_bias) def _can_fuse_gate_up_swiglu_fp4out(self): @@ -232,13 +234,7 @@ def _can_fuse_gate_up_swiglu_fp4out(self): """ if not self._can_fuse_gate_up_swiglu(): return False - if not self.down_proj.has_nvfp4: - return False - if self.down_proj.force_dynamic_quantization: - return False - if self.down_proj.input_scale is None: - return False - return True + return is_static_nvfp4_input_eligible(self.down_proj) def _fused_gate_up_swiglu(self, x, fp4_out=False): """Fused FC1 GEMM + SwiGLU using CuteDSL dense kernel. diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 53f91b6a4a06..9524c4db2851 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -31,7 +31,7 @@ from ..._utils import get_sm_version, is_sm_100f from ...models.modeling_utils import QuantConfig from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, - is_nvfp4_marlin_enabled, + is_nvfp4_marlin_supported_sm, replace_parameter_and_save_metadata, unswizzle_sf) @@ -99,6 +99,24 @@ def flip(cls, mode): return cls.ROW if mode == cls.COLUMN else cls.COLUMN +def quant_config_has_nvfp4_activation_quantization( + quant_config: Optional[QuantConfig]) -> bool: + """Whether the quantization algorithm converts activations to NVFP4.""" + return (quant_config is not None + and quant_config.layer_quant_mode.has_nvfp4() + and quant_config.quant_algo != QuantAlgo.W4A16_NVFP4) + + +def _uses_marlin_nvfp4_backend(module) -> bool: + """Whether a regular NVFP4 linear explicitly selects Marlin.""" + allowed_backends = getattr(module, "nvfp4_allowed_backends", ()) + return ("marlin" in allowed_backends and 89 <= get_sm_version() < 100 + and getattr(module, "dtype", None) == torch.bfloat16 + and not getattr(module, "use_fused_gemm_allreduce", False) + and hasattr(torch.ops.trtllm, "marlin_nvfp4_gemm") + and hasattr(torch.ops.trtllm, "gptq_marlin_repack")) + + def load_weight_shard( weight, tensor_parallel_size: int = 1, @@ -354,6 +372,7 @@ class LinearMethodBase(ABC): # window buffer. apply() reads this ClassVar to derive output_buffer_kind # internally; callers do not pass output_buffer_kind as a parameter. supports_nccl_symmetric_memory_window_output: ClassVar[bool] = False + quantizes_nvfp4_activations: ClassVar[bool] = False @abstractmethod def create_weights(self, module: Linear, in_features: int, @@ -400,6 +419,9 @@ def load_weights(self, def transform_weights(self, module: Linear) -> None: return None + def cache_derived_state(self, module: Linear) -> None: + return None + def post_load_weights(self, module: Linear) -> None: self.transform_weights(module) @@ -1335,6 +1357,7 @@ def transform_weights(self, module: Linear) -> None: class NVFP4LinearMethod(LinearMethodBase): supports_nccl_symmetric_memory_window_output: ClassVar[bool] = True + quantizes_nvfp4_activations: ClassVar[bool] = True # Temporary workaround which will be resolved by TRTLLM-11958 # When True, use tunable_fp4_quantize (AutoTuner selects TRTLLM vs @@ -1486,7 +1509,9 @@ def apply(self, module: Linear, input: torch.Tensor, # Use unified interface - supports CUTLASS, cuBLASLt, CuteDSL # Convert list to comma-separated string for torch.compile compatibility - allowed_backends_str = ','.join(module.nvfp4_allowed_backends) + use_marlin = _uses_marlin_nvfp4_backend(module) + allowed_backends_str = ('marlin' if use_marlin else ','.join( + module.nvfp4_allowed_backends)) output_buffer_kind = ( int(BufferKind.NCCL_WINDOW) if self.supports_nccl_symmetric_memory_window_output @@ -1498,7 +1523,7 @@ def apply(self, module: Linear, input: torch.Tensor, and module.mapping is not None else None) # Fuse bias inside the GEMM op when N is unpadded and the output is a # plain buffer; otherwise fall back to post-op `out + bias` below. - fuse_bias_in_gemm = (bias is not None + fuse_bias_in_gemm = (bias is not None and not use_marlin and output_buffer_kind == int(BufferKind.DEFAULT) and module.weight.shape[0] == module.out_features) output = torch.ops.trtllm.nvfp4_gemm( @@ -1665,12 +1690,16 @@ def _cleanup_nvfp4_tmp_attrs(self, def process_weights_after_loading_vanilla(self, module: Linear): input_scale, weight_scale_2, alpha = self._finalize_nvfp4_scales(module) - # For dynamic activation quantization, input_scale and alpha are computed at runtime - if input_scale is not None: + # For dynamic activation quantization, input_scale and alpha are computed at runtime. + # The destination checks cover weight-only subclasses (W4A16 / Marlin), which leave + # input_scale / inv_input_scale / alpha as None in create_weights: a checkpoint may + # still carry activation scales, but they are meaningless for a high-precision + # activation path, so drop them instead of copying into a missing Parameter. + if input_scale is not None and module.input_scale is not None: copy_weight(module.input_scale, input_scale) E2M1_MAX = 6.0 module.inv_input_scale.data = module.input_scale / E2M1_MAX - if alpha is not None: + if alpha is not None and module.alpha is not None: copy_weight(module.alpha, alpha) module.scalar_alpha = alpha.item() if weight_scale_2 is not None: @@ -1761,11 +1790,13 @@ def process_weights_after_loading_fused_qkv_linear(self, module: Linear): weight_scale = torch.ops.trtllm.block_scale_interleave(weight_scale) copy_weight(module.weight_scale, weight_scale) - # Finalize input_scale, weight_scale_2, alpha + # Finalize input_scale, weight_scale_2, alpha. The destination checks skip + # scales this linear method never allocated (see the weight-only note in + # process_weights_after_loading_vanilla). input_scale, weight_scale_2, alpha = self._finalize_nvfp4_scales(module) - if input_scale is not None: + if input_scale is not None and module.input_scale is not None: copy_weight(module.input_scale, input_scale) - if alpha is not None: + if alpha is not None and module.alpha is not None: copy_weight(module.alpha, alpha) module.scalar_alpha = alpha.item() if weight_scale_2 is not None: @@ -1847,11 +1878,13 @@ def process_weights_after_loading_fused_gate_up_linear( weight_scale = torch.ops.trtllm.block_scale_interleave(weight_scale) copy_weight(module.weight_scale, weight_scale) - # Finalize input_scale, weight_scale_2, alpha + # Finalize input_scale, weight_scale_2, alpha. The destination checks skip + # scales this linear method never allocated (see the weight-only note in + # process_weights_after_loading_vanilla). input_scale, weight_scale_2, alpha = self._finalize_nvfp4_scales(module) - if input_scale is not None: + if input_scale is not None and module.input_scale is not None: copy_weight(module.input_scale, input_scale) - if alpha is not None: + if alpha is not None and module.alpha is not None: copy_weight(module.alpha, alpha) module.scalar_alpha = alpha.item() if weight_scale_2 is not None: @@ -1969,70 +2002,103 @@ def transform_weights(self, module: Linear) -> None: class W4A16NVFP4LinearMethod(NVFP4LinearMethod): - """W4A16 dequant fallback for NVFP4 on SM<100. Only used by - modeling_nemotron_h. ``apply_linear_allreduce`` is inherited unchanged: - its fused path is SM>=100-gated upstream. - """ + """W4A16 NVFP4 linear using on-the-fly weight dequantization.""" + + quantizes_nvfp4_activations: ClassVar[bool] = False + + def create_weights(self, module: Linear, in_features: int, + out_features: int, bias: bool, dtype: torch.dtype): + super().create_weights(module, in_features, out_features, bias, dtype) + + # W4A16 consumes high-precision activations, so there is no calibrated + # activation scale. These must be None from create_weights onwards (not + # merely left uninitialized): eligibility checks such as + # ``is_static_nvfp4_input_eligible`` run right after create_weights, + # before any weights are loaded, and read ``input_scale is not None``. + # The parent's process_weights_after_loading_* skip scales whose + # destination is None, so checkpoint-provided values are dropped here. + module.input_scale = None + module.inv_input_scale = None + module.alpha = None def transform_weights(self, module: Linear) -> None: - # Skip parent's 32x16 weight padding (apply() accepts [N, K/2] as-is) - # and un-swizzle per-block scale once at load. + # Materialize the smaller linear scale view once for Triton dequantization. LinearMethodBase.transform_weights(self, module) + self.cache_derived_state(module) + + def cache_derived_state(self, module: Linear) -> None: pad_rows = fp4_utils.pad_up(module.out_features, 128) pad_cols = fp4_utils.pad_up( module.in_features // module.scaling_vector_size, 4) scale_swizzled = module.weight_scale.data.view( fp4_utils.float4_sf_dtype).reshape(pad_rows, pad_cols) scale_linear = torch.ops.trtllm.block_scale_interleave_reverse( - scale_swizzled) - module.weight_scale.data.view(fp4_utils.float4_sf_dtype).copy_( - scale_linear.reshape(-1)) + scale_swizzled).reshape(-1) + buffer_name = "_w4a16_weight_scale_linear" + if buffer_name in module._buffers: + module._buffers[buffer_name] = scale_linear + else: + module.register_buffer(buffer_name, scale_linear, persistent=False) - def apply(self, module: Linear, input: torch.Tensor, - bias: Optional[torch.Tensor]): + @staticmethod + def _prepare_input(module: Linear, input: torch.Tensor): if isinstance(input, (Fp4QuantizedTensor, tuple)): raise RuntimeError( - "W4A16NVFP4LinearMethod: hp input required; disable upstream " - "FP4 fusion (e.g. TRTLLM_ENABLE_ATTENTION_NVFP4_OUTPUT=0)") + "W4A16NVFP4LinearMethod requires a high-precision input; " + "disable upstream FP4 fusion") - ## FP8 input from upstream FMHA pre-quant: invert by / module.inv_input_scale. + # create_weights leaves inv_input_scale None on this path, so an FP8 + # activation has no scale to dequantize it with. Reject it up front the + # same way as a pre-quantized FP4 input. if input.dtype == torch.float8_e4m3fn: - assert module.inv_input_scale is not None, \ - "W4A16NVFP4LinearMethod: FP8 input requires static inv_input_scale" - input = (input.to(module.dtype) / module.inv_input_scale).to( - module.dtype) + raise RuntimeError( + "W4A16NVFP4LinearMethod requires a high-precision input; " + "disable upstream FP8 attention output") original_shape = None if input.dim() > 2: original_shape = input.shape input = input.reshape(-1, input.shape[-1]) - # NVFP4_AWQ pre_quant_scale (mirrors parent's _input_prepare branch). if module.pre_quant_scale is not None: - assert input.dtype == module.pre_quant_scale.dtype, ( - "Input dtype and pre_quant_scale dtype must match") + assert input.dtype == module.pre_quant_scale.dtype, \ + "Input dtype and pre_quant_scale dtype must match" input = input * module.pre_quant_scale + return input, original_shape + + @staticmethod + def _restore_output(output: torch.Tensor, original_shape, + bias: Optional[torch.Tensor]): + if original_shape is not None: + output = output.reshape(*original_shape[:-1], output.shape[-1]) + if bias is not None: + output = output + bias + return output + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor]): + input, original_shape = self._prepare_input(module, input) from tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4 import \ dequant_nvfp4_2d_triton weight_deq = dequant_nvfp4_2d_triton( module.weight.view(torch.uint8), - module.weight_scale, + module._w4a16_weight_scale_linear, module.weight_scale_2, target_dtype=module.dtype, sf_vec_size=module.scaling_vector_size, ) - if module.use_custom_cublas_mm: + if getattr(module, "use_custom_cublas_mm", False): output_buffer_kind = ( int(BufferKind.NCCL_WINDOW) if self.supports_nccl_symmetric_memory_window_output - and module.all_reduce is not None + and getattr(module, "all_reduce", None) is not None and module.all_reduce.uses_nccl_symmetric_memory_window() else int(BufferKind.DEFAULT)) - group = (module.mapping.tp_group + mapping = getattr(module, "mapping", None) + group = (mapping.tp_group if output_buffer_kind == int(BufferKind.NCCL_WINDOW) - and module.mapping is not None else None) + and mapping is not None else None) output = torch.ops.trtllm.cublas_mm( input, weight_deq.t(), @@ -2043,10 +2109,148 @@ def apply(self, module: Linear, input: torch.Tensor, ) else: output = F.linear(input, weight_deq, bias) + return self._restore_output(output, original_shape, bias=None) - if original_shape is not None: - output = output.reshape(*original_shape[:-1], output.shape[-1]) - return output + +class MarlinNVFP4LinearMethod(W4A16NVFP4LinearMethod): + """W4A16 NVFP4 linear backed by Marlin.""" + + # ``apply`` always allocates a plain output buffer (the Marlin GEMM has no + # NCCL-window output path) and ``apply_linear_allreduce`` is unsupported, so + # this must not inherit the True from NVFP4LinearMethod: Linear.forward reads + # this ClassVar to take a zero-copy branch that assumes the GEMM wrote into + # the symmetric-memory window. + supports_nccl_symmetric_memory_window_output: ClassVar[bool] = False + + def get_tp_alignment(self, + tp_mode: Optional[TensorParallelMode], + quant_config: Optional[QuantConfig] = None) -> int: + # Same 32-element alignment as the parent NVFP4 path. The Marlin kernel + # itself wants K%64 and N%128, but ``transform_weights`` pads the weight + # and scales up to those bounds (and ``apply`` slices the N padding back + # off), so the sharding constraint stays unchanged. + return 32 + + @staticmethod + def is_supported(module: Linear) -> bool: + sm_version = get_sm_version() + return ((89 <= sm_version < 100 or sm_version in (120, 121)) + and getattr(module, "dtype", None) == torch.bfloat16 + and not getattr(module, "use_fused_gemm_allreduce", False) + and hasattr(torch.ops.trtllm, "marlin_nvfp4_gemm") + and hasattr(torch.ops.trtllm, "gptq_marlin_repack")) + + @classmethod + def is_enabled(cls, module: Linear) -> bool: + """Apply the architecture policy on top of kernel capability.""" + if not cls.is_supported(module): + return False + sm_version = get_sm_version() + return (sm_version in (120, 121) + or "marlin" in getattr(module, "nvfp4_allowed_backends", ())) + + def transform_weights(self, module: Linear) -> None: + from tensorrt_llm.quantization.utils import marlin_utils + + weight = module.weight.data + weight_scale = module.weight_scale.data + size_n = module.out_features + size_k = module.in_features + group_size = module.scaling_vector_size + + assert size_k % group_size == 0, ( + f"size_k {size_k} must be divisible by group_size {group_size}") + + size_k_pad = fp4_utils.pad_up(size_k, 64) + size_n_pad = fp4_utils.pad_up(size_n, 128) + num_groups = size_k // group_size + scale_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + weight_scale.view(size_n_pad, -1)) + scale_2d = scale_unswizzled[:size_n, :num_groups] + + if size_k_pad != size_k or size_n_pad != size_n: + num_groups_pad = size_k_pad // group_size + weight = F.pad(weight, + (0, + (size_k_pad - size_k) // 2, 0, size_n_pad - size_n)) + # Marlin's S0E5M3 fast dequantization requires a non-zero scale for + # zero-weight K padding. 0x08 is the smallest normal E4M3 value. + scale_2d = F.pad(scale_2d, (0, num_groups_pad - num_groups), + value=0x08) + scale_2d = F.pad(scale_2d, (0, 0, 0, size_n_pad - size_n), value=0) + + qweight_int32 = weight.view(torch.int32).T.contiguous() + perm = torch.empty(0, dtype=torch.int32, device=weight.device) + marlin_weight = torch.ops.trtllm.gptq_marlin_repack( + b_q_weight=qweight_int32, + perm=perm, + size_k=size_k_pad, + size_n=size_n_pad, + num_bits=4, + is_a_8bit=False, + ) + + scale_2d = scale_2d.view(torch.float8_e4m3fn).T.contiguous() + marlin_scale = marlin_utils.marlin_permute_scales(scale_2d.to( + torch.half), + size_k_pad, + size_n_pad, + group_size=group_size) + marlin_scale = marlin_utils.nvfp4_marlin_process_scales(marlin_scale) + + weight_scale_2 = module.weight_scale_2.data + if (weight_scale_2.numel() == 0 + or not torch.isfinite(weight_scale_2).all() + or weight_scale_2.item() == 0): + weight_scale_2 = torch.tensor([1.0], + dtype=torch.float32, + device=weight.device) + weight_global_scale = marlin_utils.nvfp4_marlin_process_global_scale( + weight_scale_2.to(torch.bfloat16)) + + module.weight = Parameter(marlin_weight, requires_grad=False) + module.weight_scale = Parameter(marlin_scale, requires_grad=False) + module.weight_global_scale = Parameter(weight_global_scale, + requires_grad=False) + self.cache_derived_state(module) + + def cache_derived_state(self, module: Linear) -> None: + module._marlin_size_k = fp4_utils.pad_up(module.in_features, 64) + module._marlin_size_n = fp4_utils.pad_up(module.out_features, 128) + + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor]): + input, original_shape = self._prepare_input(module, input) + size_k = module.in_features + size_n = module.out_features + size_k_pad = getattr(module, "_marlin_size_k", size_k) + size_n_pad = getattr(module, "_marlin_size_n", size_n) + + input_bf16 = input.bfloat16() + if size_k_pad != size_k: + input_bf16 = F.pad(input_bf16, (0, size_k_pad - size_k)) + output = torch.ops.trtllm.marlin_nvfp4_gemm( + input_bf16, + module.weight, + scale_a=None, + scale_b=module.weight_scale, + alpha=None, + weight_global_scale=module.weight_global_scale, + bias=None, + out_dtype=module.dtype, + size_n=size_n_pad, + size_k=size_k_pad, + output_buffer_kind=int(BufferKind.DEFAULT), + ) + if size_n_pad != size_n: + output = output[..., :size_n].contiguous() + return self._restore_output(output, original_shape, bias) + + def apply_linear_allreduce(self, module: Linear, input: torch.Tensor, + bias: Optional[torch.Tensor], tp_rank: int, + tp_group: List[int]): + raise RuntimeError( + "MarlinNVFP4LinearMethod does not support apply_linear_allreduce") class W4A8NVFP4FP8LinearMethod(LinearMethodBase): @@ -2903,124 +3107,6 @@ def apply(self, module: Linear, input: torch.Tensor, return output -class MarlinNVFP4LinearMethod(NVFP4LinearMethod): - """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 - - weight = module.weight.data - weight_scale = module.weight_scale.data - size_n = module.out_features - size_k = module.in_features - group_size = module.scaling_vector_size # 16 - - assert size_k % group_size == 0, ( - f"size_k {size_k} must be divisible by group_size {group_size}") - - size_k_pad = fp4_utils.pad_up(size_k, 64) - size_n_pad = fp4_utils.pad_up(size_n, 128) - - num_groups = size_k // group_size - n_padded = size_n_pad - scale_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - weight_scale.view(n_padded, -1)) - # [size_n, num_groups] block scales; uint8 storage (reverse interleave), - # reinterpreted as E4M3 after any padding. Pad in uint8 since F.pad does - # not support float8. - scale_2d = scale_unswizzled[:size_n, :num_groups] - - if size_k_pad != size_k or size_n_pad != size_n: - num_groups_pad = size_k_pad // group_size - # weight: [N, K/2] uint8 -> [N_pad, K_pad/2] (FP4 zero == 0.0) - weight = F.pad(weight, - (0, - (size_k_pad - size_k) // 2, 0, size_n_pad - size_n)) - # scales: [N, num_groups] -> [N_pad, num_groups_pad]. - # The Marlin S0E5M3 fast-dequant is NOT zero-safe: a zero scale on - # a (zero-weight) padded K-group still corrupts that tile's output. - # Since K is the contraction dim, one bad group-scale poisons every - # output row, so padded K-groups must carry a valid non-zero fp8 - # scale -- use the smallest-normal e4m3 value (0x08), matching the - # quantizer's own zero-block scale. N-row padding is sliced off in - # ``apply`` and can stay zero. - scale_2d = F.pad(scale_2d, (0, num_groups_pad - num_groups), - value=0x08) - scale_2d = F.pad(scale_2d, (0, 0, 0, size_n_pad - size_n), value=0) - - qweight_int32 = weight.view( - torch.int32).T.contiguous() # [K_pad/4, N_pad] - perm = torch.empty(0, dtype=torch.int32, device=weight.device) - marlin_weight = torch.ops.trtllm.gptq_marlin_repack( - b_q_weight=qweight_int32, - perm=perm, - size_k=size_k_pad, - size_n=size_n_pad, - num_bits=4, - is_a_8bit=False, - ) - - scale_2d = scale_2d.view( - torch.float8_e4m3fn).T.contiguous() # [num_groups_pad, N_pad] - marlin_scale = marlin_utils.marlin_permute_scales(scale_2d.to( - torch.half), - size_k_pad, - size_n_pad, - group_size=group_size) - marlin_scale = marlin_utils.nvfp4_marlin_process_scales(marlin_scale) - - ws2 = module.weight_scale_2.data - if ws2.numel() == 0 or not torch.isfinite(ws2).all() or ws2.item() == 0: - ws2 = torch.tensor([1.0], dtype=torch.float32, device=weight.device) - weight_global_scale = marlin_utils.nvfp4_marlin_process_global_scale( - ws2.to(torch.bfloat16)) - - module.weight = Parameter(marlin_weight, requires_grad=False) - module.weight_scale = Parameter(marlin_scale, requires_grad=False) - module.weight_global_scale = Parameter(weight_global_scale, - requires_grad=False) - # Padded GEMM dims consumed by ``apply``; default to the real sizes. - module._marlin_size_k = size_k_pad - module._marlin_size_n = size_n_pad - - def apply(self, module: Linear, input: torch.Tensor, - bias: Optional[torch.Tensor]): - assert is_nvfp4_marlin_enabled() - size_k = module.in_features - size_n = module.out_features - # Set by transform_weights; equal to size_k/size_n when 64-aligned. - size_k_pad = getattr(module, "_marlin_size_k", size_k) - size_n_pad = getattr(module, "_marlin_size_n", size_n) - - x = input.bfloat16() - if size_k_pad != size_k: - x = F.pad(x, (0, size_k_pad - size_k)) - output = torch.ops.trtllm.marlin_nvfp4_gemm( - x, - module.weight, - scale_a=None, - scale_b=module.weight_scale, - alpha=None, - weight_global_scale=module.weight_global_scale, - bias=None, - out_dtype=module.dtype, - size_n=size_n_pad, - size_k=size_k_pad, - output_buffer_kind=int(BufferKind.DEFAULT), - ) - if size_n_pad != size_n: - output = output[..., :size_n].contiguous() - if bias is not None: - output = output + bias - return output - - def apply_linear_allreduce(self, module: Linear, input: torch.Tensor, - bias: Optional[torch.Tensor], tp_rank: int, - tp_group: List[int]): - raise RuntimeError( - "MarlinNVFP4LinearMethod does not support apply_linear_allreduce") - - def _mxfp8_cutlass_op_available() -> bool: """Cached check for whether the CUTLASS MXFP8xMXFP8 GEMM op is compiled in. @@ -3214,13 +3300,12 @@ def get_quant_method(quant_config: Optional[QuantConfig] = None): return FP8RowwiseLinearMethod() if quant_config.layer_quant_mode.has_fp8_block_scales(): return FP8BlockScalesLinearMethod() + if quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: + return W4A16NVFP4LinearMethod() if quant_config.layer_quant_mode.has_nvfp4(): if quant_config.quant_algo == QuantAlgo.NVFP4_ARC: return NVFP4ARCLinearMethod() - elif is_nvfp4_marlin_enabled(): - return MarlinNVFP4LinearMethod() - else: - return NVFP4LinearMethod() + return NVFP4LinearMethod() if quant_config.layer_quant_mode.has_w4a8_nvfp4_fp8(): return W4A8NVFP4FP8LinearMethod() if quant_config.layer_quant_mode.has_w4a8_mxfp4_fp8(): @@ -3275,7 +3360,9 @@ def __init__( nvfp4_allowed_backends: List of backends to consider for NVFP4 GEMM auto-selection. Default (via config): ['cutlass', 'cublaslt', 'cuda_core'] - excludes cutedsl for faster build. Add 'cutedsl' for extreme performance at the cost of longer build time. - Valid backends: 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core'. + Valid backends: 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core', 'marlin'. + Marlin is opt-in on SM89-99. W4A16 BF16 linear layers use + Marlin by default on SM120/121. Configure via nvfp4_gemm_config.allowed_backends in extra_llm_api_options.yaml. """ from ..distributed import AllReduce @@ -3389,8 +3476,8 @@ def __init__( in_features_aligned = self.in_features % 128 == 0 out_features_aligned = self.out_features % 64 == 0 tp_valid = self.tp_mode is not None and self.tp_mode == TensorParallelMode.ROW and self.tp_size > 1 - quant_valid = self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( - ) + quant_valid = quant_config_has_nvfp4_activation_quantization( + self.quant_config) device_supported = get_sm_version() >= 100 enable_gemm_allreduce_fusion_env = (os.environ.get( @@ -3416,7 +3503,21 @@ def __init__( self.create_weights() def get_quant_method(self, quant_config: Optional[QuantConfig] = None): - return get_quant_method(quant_config) + quant_method = get_quant_method(quant_config) + method_type = type(quant_method) + if method_type is W4A16NVFP4LinearMethod: + # Marlin by default on SM120/121, explicit opt-in elsewhere. + if MarlinNVFP4LinearMethod.is_enabled(self): + return MarlinNVFP4LinearMethod() + elif method_type is NVFP4LinearMethod: + # The Marlin kernel is W4A16, so an explicit opt-in on a W4A4 + # checkpoint runs the weight-only method, which pads N/K to the + # tile sizes the kernel requires. + if ("marlin" in self.nvfp4_allowed_backends + and is_nvfp4_marlin_supported_sm() + and MarlinNVFP4LinearMethod.is_supported(self)): + return MarlinNVFP4LinearMethod() + return quant_method @staticmethod def _calc_shard(total, tp_size, rank): @@ -3582,6 +3683,12 @@ def create_weights(self): if self._weights_created: return + # Mixed-precision loading may replace quant_config after __init__. + # Weight-only W4A16 must not retain the activation-quantized fused path. + if not quant_config_has_nvfp4_activation_quantization( + self.quant_config): + self.use_fused_gemm_allreduce = False + self.rebuild_tensor_metadata = {} self.quant_method = self.get_quant_method(self.quant_config) @@ -3622,6 +3729,18 @@ def has_nvfp4(self): return self.quant_config is not None and self.quant_config.layer_quant_mode.has_nvfp4( ) + @property + def has_nvfp4_activation_quantization(self): + assert self._weights_created + return self.quant_method.quantizes_nvfp4_activations + + @property + def uses_marlin_nvfp4(self): + assert self._weights_created + return (isinstance(self.quant_method, MarlinNVFP4LinearMethod) + or (type(self.quant_method) is NVFP4LinearMethod + and _uses_marlin_nvfp4_backend(self))) + @property def has_weight_only_quant(self): assert self._weights_created @@ -3777,6 +3896,7 @@ def transform_weights(self) -> None: self._weights_transformed = True def cache_derived_state(self) -> None: + self.quant_method.cache_derived_state(self) self._weights_transformed = True def post_load_weights(self) -> None: @@ -3793,16 +3913,16 @@ def is_static_nvfp4_input_eligible(linear) -> bool: """Whether `linear` consumes a static (calibrated) NVFP4 input, making it eligible to have its input-quantize folded into a producing RMSNorm. - Eligible iff the Linear has NVFP4 weights, a calibrated (static) - `input_scale`, no AWQ `pre_quant_scale`, and is not forced to dynamic - quantization. This is the single canonical definition shared by every - NVFP4-fold site (the layer-boundary / dense folds in modeling_deepseekv3.py - and the q_a_layernorm -> q_b_proj fold in attention.py's MLA) so the gate - cannot drift between them. + Eligible iff the Linear quantizes activations to NVFP4, has a calibrated + (static) `input_scale`, has no AWQ `pre_quant_scale`, and is not forced to + dynamic quantization. This is the single canonical definition shared by + every NVFP4-fold site (the layer-boundary / dense folds in + modeling_deepseekv3.py and the q_a_layernorm -> q_b_proj fold in + attention.py's MLA) so the gate cannot drift between them. """ if linear is None: return False - return (getattr(linear, "has_nvfp4", False) + return (getattr(linear, "has_nvfp4_activation_quantization", False) and not getattr(linear, "force_dynamic_quantization", False) and getattr(linear, "input_scale", None) is not None and getattr(linear, "pre_quant_scale", None) is None) diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py index b676d28ab69d..42aae9956219 100644 --- a/tensorrt_llm/_torch/modules/mlp.py +++ b/tensorrt_llm/_torch/modules/mlp.py @@ -10,7 +10,8 @@ from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType from ..utils import Fp4QuantizedTensor, gelu_tanh, relu2 -from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig +from .linear import (Linear, TensorParallelMode, WeightMode, + WeightsLoadingConfig, is_static_nvfp4_input_eligible) class MLP(nn.Module): @@ -99,13 +100,8 @@ def create_weights(self): self.up_proj.create_weights() self.down_proj.create_weights() - has_nvfp4 = hasattr(self.down_proj, - 'has_nvfp4') and self.down_proj.has_nvfp4 + has_static_nvfp4_input = is_static_nvfp4_input_eligible(self.down_proj) has_kernel = hasattr(torch.ops.trtllm, 'fused_relu2_quantize') - # NVFP4LinearMethod.create_weights always allocates input_scale as a - # Parameter, but excluded layers reset it to None at load time. Check - # for a real tensor. - has_scale = getattr(self.down_proj, 'input_scale', None) is not None is_relu2 = self.activation is relu2 # The fused relu2+fp4_quantize kernel body is guarded by # ``__CUDA_ARCH__ >= 1000`` (see fusedActivationQuant.cu). On pre-SM100 @@ -113,7 +109,7 @@ def create_weights(self): # quantize in the downstream linear layer. is_sm100_or_later = get_sm_version() >= 100 - self._use_fused_relu2_quant = (has_nvfp4 and has_kernel and has_scale + self._use_fused_relu2_quant = (has_static_nvfp4_input and has_kernel and is_relu2 and is_sm100_or_later) # Static eligibility for the fused GELU(tanh) CuteDSL epilogue (mirrors @@ -139,11 +135,14 @@ def forward( # may be downgraded to unquantized after create_weights (e.g. LTX-2 # quant-exclusion), so re-check the NVFP4 _input_prepare at runtime (a # torch.compile trace-time guard, not a per-step cost); else fall back to eager. - if self._use_fused_gelu and hasattr( - getattr(self.up_proj, "quant_method", None), "_input_prepare"): - if self._use_fused_gelu_fp4out and hasattr( - getattr(self.down_proj, "quant_method", None), - "_input_prepare"): + if (self._use_fused_gelu + and self.up_proj.has_nvfp4_activation_quantization + and hasattr(getattr(self.up_proj, "quant_method", None), + "_input_prepare")): + if (self._use_fused_gelu_fp4out + and is_static_nvfp4_input_eligible(self.down_proj) + and hasattr(getattr(self.down_proj, "quant_method", None), + "_input_prepare")): m = self._token_count(x) return self.down_proj( self._fused_gelu(x, fp4_out=m >= MLP._FP4OUT_MIN_M)) @@ -151,7 +150,10 @@ def forward( x_up = self.up_proj(x) - if self._use_fused_relu2_quant: + # Weight loading may replace the quantization method after + # create_weights(), so do not rely on the cached eligibility alone. + if (self._use_fused_relu2_quant + and is_static_nvfp4_input_eligible(self.down_proj)): x_act = self._fused_relu2_quant(x_up) else: x_act = self.activation(x_up) @@ -169,16 +171,14 @@ def _gelu_fusion_eligibility(self) -> Tuple[bool, bool]: applied in forward (quant_method can be downgraded after this). """ if (self.activation is not gelu_tanh - or get_sm_version() not in (100, 103) - or not getattr(self.up_proj, "has_nvfp4", False)): + or get_sm_version() not in (100, 103) or not getattr( + self.up_proj, "has_nvfp4_activation_quantization", False)): return False, False bf16_ok = hasattr(torch.ops.trtllm, "cute_dsl_nvfp4_dense_gemm_gelu_blackwell") fp4_ok = (bf16_ok and hasattr( torch.ops.trtllm, "cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell") - and getattr(self.down_proj, "has_nvfp4", False) - and not self.down_proj.force_dynamic_quantization - and self.down_proj.input_scale is not None) + and is_static_nvfp4_input_eligible(self.down_proj)) return bf16_ok, fp4_ok @staticmethod diff --git a/tests/integration/test_lists/test-db/l0_gb10.yml b/tests/integration/test_lists/test-db/l0_gb10.yml index efc6a64530f9..da77f6e22ace 100644 --- a/tests/integration/test_lists/test-db/l0_gb10.yml +++ b/tests/integration/test_lists/test-db/l0_gb10.yml @@ -17,6 +17,8 @@ l0_gb10: tests: # ------------- PyTorch tests --------------- - unittest/_torch/attention/test_attention_mla.py + - unittest/_torch/modules/test_w4a16_nvfp4_linear.py + - unittest/_torch/thop/parallel/test_fp4_linear.py - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index fb57eb7c1a8f..34d892ebf8fe 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -32,6 +32,7 @@ l0_h100: - unittest/_torch/modules/test_rotary_embedding.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules + - unittest/_torch/modules/test_w4a16_nvfp4_linear.py # ------------- MoE components tests --------------- - unittest/_torch/modules/test_moe_load_balancer.py - unittest/_torch/modules/test_moe_routing.py diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py new file mode 100644 index 000000000000..1a25b1e18a25 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_moe_quant.py @@ -0,0 +1,145 @@ +# 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. + +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from torch import nn + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHMOE, NemotronHMTP +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def _make_nemotron_h_moe_config( + quant_config: QuantConfig, moe_backend: str = "CUTLASS" +) -> ModelConfig: + return ModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=16, + intermediate_size=32, + mlp_bias=False, + moe_intermediate_size=64, + moe_latent_size=None, + n_group=1, + n_routed_experts=4, + n_shared_experts=0, + num_experts_per_tok=1, + routed_scaling_factor=1.0, + topk_group=1, + torch_dtype=torch.float16, + ), + moe_backend=moe_backend, + quant_config=quant_config, + ) + + +def test_nemotron_h_moe_passes_w4a16_config_through_unchanged(): + """Every MoE backend resolves W4A16_NVFP4 itself, so the layer must not + rewrite quant_algo on its way to create_moe.""" + quant_config = QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] + ) + model_config = _make_nemotron_h_moe_config(quant_config) + captured = {} + + def fake_create_moe(**kwargs): + captured.update(kwargs) + return nn.Identity() + + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.create_moe", side_effect=fake_create_moe + ): + with patch("torch.cuda.Event", side_effect=lambda: object()): + aux_stream_dict = {AuxStreamType.MoeShared: None} + NemotronHMOE(model_config=model_config, layer_idx=1, aux_stream_dict=aux_stream_dict) + + effective = captured["override_quant_config"] or captured["model_config"].quant_config + assert effective.quant_algo == QuantAlgo.W4A16_NVFP4 + assert effective.group_size == 16 + assert model_config.quant_config.quant_algo == QuantAlgo.W4A16_NVFP4 + + +def test_nemotron_h_moe_uses_mixer_expert_layer_quant_config(): + global_quant_config = QuantConfig() + layer_quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16) + model_config = _make_nemotron_h_moe_config(global_quant_config) + model_config.quant_config_dict = { + "model.layers.1.mixer.experts.0.up_proj": layer_quant_config, + } + captured = {} + + def fake_create_moe(**kwargs): + captured.update(kwargs) + return nn.Identity() + + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.create_moe", + side_effect=fake_create_moe, + ): + with patch("torch.cuda.Event", side_effect=lambda: object()): + NemotronHMOE( + model_config=model_config, + layer_idx=1, + aux_stream_dict={AuxStreamType.MoeShared: None}, + ) + + assert captured["override_quant_config"] is layer_quant_config + + +def test_nemotron_h_mtp_overrides_quant_and_inherits_moe_backend(): + quant_config = QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, group_size=16, exclude_modules=["lm_head"] + ) + model_config = ModelConfig( + pretrained_config=SimpleNamespace( + mtp_hybrid_override_pattern="*E", + torch_dtype=torch.bfloat16, + ), + moe_backend="CUTEDSL", + quant_config=quant_config, + ) + captured = [] + + def fake_decoder_layer(**kwargs): + captured.append(kwargs) + return nn.Identity() + + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.NemotronHMTPDecoderLayer", + side_effect=fake_decoder_layer, + ): + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.DeepseekV3MTPHead", + side_effect=lambda model_config: nn.Identity(), + ): + with patch( + "tensorrt_llm._torch.models.modeling_nemotron_h.get_sm_version", + return_value=121, + ): + NemotronHMTP( + model_config=model_config, + layer_idx=52, + aux_stream_dict={}, + ) + + assert len(captured) == 2 + for layer_kwargs in captured: + sublayer_model_config = layer_kwargs["model_config"] + assert sublayer_model_config.quant_config.quant_algo is None + assert sublayer_model_config.moe_backend == model_config.moe_backend + assert model_config.quant_config is quant_config diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 1feac564b176..85dc479943d8 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -23,6 +23,8 @@ SM120/SM121 hardware. """ +import sys +import types from unittest.mock import patch import pytest @@ -33,6 +35,12 @@ from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE +from tensorrt_llm._torch.modules.fused_moe.quantization import ( + NVFP4CuteDslB12xFusedMoEMethod, + NVFP4CutlassFusedMoEMethod, +) +from tensorrt_llm._torch.utils import ActivationType +from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig _FUSED_MOE_MODULE = "tensorrt_llm._torch.modules.fused_moe.fused_moe_cute_dsl_b12x" @@ -48,9 +56,10 @@ def test_can_implement_rejects_unsupported_sm(sm_version): @pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version): +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) +def test_can_implement_accepts_supported_sm(sm_version, quant_algo): with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): - ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) + ok, reason = CuteDslB12xFusedMoE.can_implement(quant_algo) assert ok assert reason is None @@ -130,17 +139,52 @@ def test_get_moe_cls_cutedsl_returns_plain_cutedsl_on_unsupported_sm(): assert cls is CuteDslFusedMoE +def test_get_moe_cls_cutedsl_returns_cutlass_for_w4a16_nvfp4_on_unsupported_sm(): + """CUTEDSL + W4A16_NVFP4 + non-SM120/121 → CutlassFusedMoE.""" + cfg = ModelConfig() + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + with patch("tensorrt_llm._utils.get_sm_version", return_value=100): + cls = get_moe_cls(cfg) + assert cls is CutlassFusedMoE + + @pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version): - """CUTEDSL + NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) +def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version, quant_algo): + """CUTEDSL + NVFP4/W4A16_NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" cfg = ModelConfig() cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + cfg.quant_config = QuantConfig(quant_algo=quant_algo) with patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version): cls = get_moe_cls(cfg) assert cls is CuteDslB12xFusedMoE +@pytest.mark.parametrize( + "mapping", + [ + Mapping(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2), + Mapping( + world_size=2, + tp_size=2, + enable_attention_dp=True, + dwdp_size=2, + dwdp_rank=0, + ), + ], +) +def test_get_moe_cls_cutedsl_falls_back_to_cutlass_for_distributed_b12x(mapping): + cfg = ModelConfig(mapping=mapping) + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + + with patch("tensorrt_llm._utils.get_sm_version", return_value=120): + cls = get_moe_cls(cfg) + + assert cls is CutlassFusedMoE + + def test_get_moe_cls_cutedsl_falls_back_to_plain_cutedsl_when_flashinfer_missing(monkeypatch): """CUTEDSL + NVFP4 + SM120/121 + flashinfer NOT importable → CuteDslFusedMoE.""" import builtins @@ -200,6 +244,114 @@ def test_dispatch_decode_shape_takes_b12x(): assert stub._route_to_cutlass(x) is False +def test_w4a16_nvfp4_prefill_quantize_input_stays_on_b12x(): + moe = object.__new__(CuteDslB12xFusedMoE) + moe.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + x = torch.empty(CuteDslB12xFusedMoE._PREFILL_VIA_CUTLASS_THRESHOLD, 1024) + + with patch.object( + CutlassFusedMoE, + "quantize_input", + side_effect=AssertionError("W4A16_NVFP4 prefill must not route through CUTLASS"), + ): + out, out_sf = CuteDslB12xFusedMoE.quantize_input(moe, x) + + assert out is x + assert out_sf is None + + +def test_w4a16_nvfp4_post_load_uses_modelopt_scale_contract(monkeypatch): + class _RoutingMethod: + experts_per_token = 4 + + class _FakeB12xWrapper: + calls = [] + + def __init__(self, **kwargs): + self._moe_output = None + self.calls.append(kwargs) + + def _convert_sf_to_mma_layout(scales, *, m, k, num_groups): + return scales + + flashinfer = types.ModuleType("flashinfer") + flashinfer.B12xMoEWrapper = _FakeB12xWrapper + cute_dsl = types.ModuleType("flashinfer.cute_dsl") + utils = types.ModuleType("flashinfer.cute_dsl.utils") + utils.convert_sf_to_mma_layout = _convert_sf_to_mma_layout + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.cute_dsl", cute_dsl) + monkeypatch.setitem(sys.modules, "flashinfer.cute_dsl.utils", utils) + + num_experts = 2 + hidden_size = 128 + logical_intermediate_size = 1856 + padded_intermediate_size = 1920 + module = torch.nn.Module() + module.num_experts = num_experts + module.hidden_size = hidden_size + module.intermediate_size_per_partition = logical_intermediate_size + module.moe_max_num_tokens = 8 + module.routing_method = _RoutingMethod() + module.activation_type = ActivationType.Swiglu + module.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + module.w3_w1_weight = torch.empty( + num_experts, + 2 * padded_intermediate_size, + hidden_size // 16, + dtype=torch.int64, + ) + module.w2_weight = torch.empty( + num_experts, + hidden_size, + padded_intermediate_size // 16, + dtype=torch.int64, + ) + w3_w1_weight_scale = torch.ones( + num_experts, + 2 * padded_intermediate_size, + hidden_size // 16, + dtype=torch.float8_e4m3fn, + ) + w2_weight_scale = torch.ones( + num_experts, + hidden_size, + padded_intermediate_size // 16, + dtype=torch.float8_e4m3fn, + ) + module.w3_w1_weight_scale = w3_w1_weight_scale.clone() + module.w2_weight_scale = w2_weight_scale.clone() + module.fc31_alpha = torch.tensor([0.25, 0.5], dtype=torch.float32) + module.fc2_alpha = torch.tensor([0.125, 0.25], dtype=torch.float32) + module.fc31_input_scale = torch.tensor(2.0, dtype=torch.float32) + module.fc2_input_scale = torch.tensor(4.0, dtype=torch.float32) + + with patch.object(NVFP4CutlassFusedMoEMethod, "transform_weights", return_value=None): + NVFP4CuteDslB12xFusedMoEMethod().transform_weights(module) + + assert _FakeB12xWrapper.calls + wrapper_kwargs = _FakeB12xWrapper.calls[0] + assert wrapper_kwargs.get("quant_mode") == "w4a16", wrapper_kwargs + assert wrapper_kwargs["intermediate_size"] == padded_intermediate_size + assert module._b12x_weights["fc2_input_scale"] is None + assert torch.equal( + module._b12x_weights["w1_weight_sf"].float(), + w3_w1_weight_scale.float(), + ) + assert torch.equal( + module._b12x_weights["w2_weight_sf"].float(), + w2_weight_scale.float(), + ) + assert torch.allclose( + module._b12x_weights["w1_alpha"], + torch.tensor([0.5, 1.0], dtype=torch.float32), + ) + assert torch.allclose( + module._b12x_weights["w2_alpha"], + torch.tensor([0.5, 1.0], dtype=torch.float32), + ) + + def test_dispatch_rejects_non_tensor(): """Non-tensor inputs (e.g. Fp4QuantizedTensor) stay on the b12x path so the existing ValueError surfaces in quantize_input.""" diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 00ec9070a70b..0a797f515628 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -65,9 +65,11 @@ from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm from tensorrt_llm._torch.modules.fused_moe.quantization import ( FusedMoEMethodBase, + NVFP4FusedMoEMethod, NVFP4MarlinFusedMoEMethod, UnquantizedFusedMoEMethod, W4A8MXFP4MXFP8MegaMoEDeepGemmMethod, + W4A16NVFP4CutlassFusedMoEMethod, ) from tensorrt_llm._torch.utils import ActivationType, is_gated_activation from tensorrt_llm._utils import get_sm_version, mpi_rank @@ -354,6 +356,11 @@ def test_configurable_moe_load_weights_invalidates_wrapper_transform_guard(): assert configurable_moe._weights_transformed is False +def test_moe_nvfp4_activation_quantization_capability(): + assert NVFP4FusedMoEMethod.quantizes_nvfp4_activations + assert not W4A16NVFP4CutlassFusedMoEMethod.quantizes_nvfp4_activations + + def test_marlin_moe_repack_is_transform_stage(): assert "transform_weights" in NVFP4MarlinFusedMoEMethod.__dict__ assert "post_load_weights" not in NVFP4MarlinFusedMoEMethod.__dict__ diff --git a/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py new file mode 100644 index 000000000000..f169105914f4 --- /dev/null +++ b/tests/unittest/_torch/modules/test_w4a16_nvfp4_linear.py @@ -0,0 +1,593 @@ +# 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. + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.modules.attention import Attention +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + Linear, + MarlinNVFP4LinearMethod, + NVFP4LinearMethod, + W4A16NVFP4LinearMethod, + quant_config_has_nvfp4_activation_quantization, +) +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def test_nvfp4_activation_quantization_excludes_w4a16(): + assert quant_config_has_nvfp4_activation_quantization(QuantConfig(quant_algo=QuantAlgo.NVFP4)) + assert not quant_config_has_nvfp4_activation_quantization( + QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + ) + + +def test_w4a16_attention_does_not_quantize_output_to_fp4(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + o_proj = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, + ) + attention = SimpleNamespace( + attn=SimpleNamespace(has_nvfp4=False), + o_proj=o_proj, + quant_config=quant_config, + has_quant_scale=True, + attn_output_gate=False, + is_marlin_enabled=False, + ) + + assert not Attention._use_quantize_output(attention) + + +def test_static_nvfp4_attention_can_quantize_output_to_fp4_on_blackwell(): + quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + with patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=100): + o_proj = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, + ) + attention = SimpleNamespace( + attn=SimpleNamespace(has_nvfp4=False), + o_proj=o_proj, + quant_config=quant_config, + has_quant_scale=True, + attn_output_gate=False, + is_marlin_enabled=False, + ) + + assert Attention._use_quantize_output(attention) + + +def test_nvfp4_attention_keeps_high_precision_output_for_hopper_marlin(): + quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=90), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + o_proj = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=quant_config, + reduce_output=False, + nvfp4_allowed_backends=["marlin"], + ) + attention = SimpleNamespace( + attn=SimpleNamespace(has_nvfp4=False), + o_proj=o_proj, + quant_config=quant_config, + has_quant_scale=o_proj.has_nvfp4_activation_quantization, + attn_output_gate=False, + is_marlin_enabled=o_proj.uses_marlin_nvfp4, + ) + + assert o_proj.uses_marlin_nvfp4 + assert isinstance(o_proj.quant_method, MarlinNVFP4LinearMethod) + assert not o_proj.has_nvfp4_activation_quantization + assert not Attention._use_quantize_output(attention) + + +@pytest.mark.parametrize( + ("allowed_backends", "expected_backends"), + [ + (["cutlass", "cublaslt", "cuda_core"], "cutlass,cublaslt,cuda_core"), + (["marlin"], "marlin"), + ], +) +def test_nvfp4_linear_forwards_allowed_backends_to_gemm(allowed_backends, expected_backends): + """The module's backend list reaches the unified GEMM op verbatim. Which + backends are eligible in the first place is covered by + ``test_nvfp4_linear_keeps_activation_quant_method``.""" + method = NVFP4LinearMethod() + input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((4, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + out_features=4, + dtype=torch.bfloat16, + nvfp4_allowed_backends=allowed_backends, + all_reduce=None, + mapping=None, + ) + act_fp4 = torch.empty((2, 16), dtype=torch.uint8) + act_sf = torch.empty((128 * 4,), dtype=torch.uint8) + alpha = torch.ones((1,), dtype=torch.float32) + captured = {} + + def fake_nvfp4_gemm(*args, **kwargs): + captured["allowed_backends"] = kwargs["allowed_backends"] + return torch.ones((2, 4), dtype=torch.bfloat16) + + with ( + patch.object(method, "_input_prepare", return_value=(act_fp4, act_sf, alpha)), + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=90), + patch("torch.ops.trtllm.nvfp4_gemm", side_effect=fake_nvfp4_gemm, create=True), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + output = method.apply(module, input_tensor, bias=None) + + assert captured["allowed_backends"] == expected_backends + assert output.shape == (2, 4) + + +@pytest.mark.parametrize( + ("sm_version", "allowed_backends", "expect_marlin"), + [ + # Honoured only on Ada/Hopper, and only when opted in. + (89, ["marlin"], True), + (90, ["marlin"], True), + (89, None, False), + (90, None, False), + (120, None, False), + (121, None, False), + # Opt-in ignored off Ada/Hopper: Marlin is not the right NVFP4 backend + # on SM120/121, where the W4A4 kernels are faster. + (120, ["marlin"], False), + (121, ["marlin"], False), + ], +) +def test_nvfp4_linear_marlin_opt_in_switches_to_weight_only_method( + sm_version, allowed_backends, expect_marlin +): + """The Marlin kernel is W4A16, so opting a W4A4 checkpoint into it converts + the layer to the weight-only method, which pads N/K for the kernel. Without + the opt-in the layer keeps NVFP4LinearMethod and its activation quantize.""" + kwargs = {} if allowed_backends is None else {"nvfp4_allowed_backends": allowed_backends} + + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=sm_version), + # is_nvfp4_marlin_supported_sm reads its own get_sm_version binding in + # utils.py, so it has to be patched alongside linear's. + patch( + "tensorrt_llm._torch.modules.linear.is_nvfp4_marlin_supported_sm", + return_value=89 <= sm_version < 100, + ), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + linear = Linear( + 32, + 32, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + reduce_output=False, + **kwargs, + ) + + # Asserted inside the patch: uses_marlin_nvfp4 re-reads the SM on every + # access rather than caching what create_weights decided. + assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) is expect_marlin + assert linear.uses_marlin_nvfp4 is expect_marlin + # Marlin consumes BF16 activations; the plain NVFP4 path quantizes them. + assert linear.has_nvfp4_activation_quantization is not expect_marlin + if not expect_marlin: + assert type(linear.quant_method) is NVFP4LinearMethod + + +def test_nvfp4_linear_hopper_marlin_applies_bias_as_post_op(): + method = NVFP4LinearMethod() + input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((4, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + out_features=4, + dtype=torch.bfloat16, + nvfp4_allowed_backends=["marlin"], + all_reduce=None, + mapping=None, + ) + act_fp4 = torch.empty((2, 16), dtype=torch.uint8) + act_sf = torch.empty((128 * 4,), dtype=torch.uint8) + alpha = torch.ones((1,), dtype=torch.float32) + captured = {} + + def fake_nvfp4_gemm(*args, **kwargs): + captured.update(kwargs) + return torch.ones((2, 4), dtype=torch.bfloat16) + + with ( + patch.object(method, "_input_prepare", return_value=(act_fp4, act_sf, alpha)), + patch( + "tensorrt_llm._torch.modules.linear.get_sm_version", + return_value=90, + ), + patch( + "torch.ops.trtllm.nvfp4_gemm", + side_effect=fake_nvfp4_gemm, + create=True, + ), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + output = method.apply(module, input_tensor, bias=bias) + + assert captured["allowed_backends"] == "marlin" + assert captured["bias"] is None + torch.testing.assert_close(output, torch.ones((2, 4), dtype=torch.bfloat16) + bias) + + +def test_w4a16_nvfp4_linear_uses_high_precision_activation_without_fp4_quantize(): + method = W4A16NVFP4LinearMethod() + input_tensor = torch.ones((2, 32), dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((4, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + _w4a16_weight_scale_linear=torch.empty((128 * 4,), dtype=torch.uint8), + weight_scale_2=torch.tensor([0.25], dtype=torch.float32), + dtype=torch.bfloat16, + out_features=4, + scaling_vector_size=16, + pre_quant_scale=None, + use_custom_cublas_mm=False, + ) + captured = {} + + def fake_dequant(weight, weight_scale, weight_scale_2, **kwargs): + captured["weight"] = weight + captured["weight_scale"] = weight_scale + captured["weight_scale_2"] = weight_scale_2 + captured.update(kwargs) + return torch.ones((4, 32), dtype=torch.bfloat16) + + def fail_fp4_quantize(*args, **kwargs): + raise AssertionError("W4A16 NVFP4 must not quantize activations") + + with ( + patch( + "tensorrt_llm._torch.modules.fused_moe.triton_dequant_nvfp4.dequant_nvfp4_2d_triton", + side_effect=fake_dequant, + ), + patch("torch.ops.trtllm.fp4_quantize", side_effect=fail_fp4_quantize, create=True), + ): + output = method.apply(module, input_tensor, bias) + + assert captured["weight"].data_ptr() == module.weight.data_ptr() + assert captured["weight_scale"] is module._w4a16_weight_scale_linear + assert captured["weight_scale_2"] is module.weight_scale_2 + assert captured["target_dtype"] is torch.bfloat16 + assert captured["sf_vec_size"] == 16 + expected = torch.tensor( + [[33.0, 34.0, 35.0, 36.0], [33.0, 34.0, 35.0, 36.0]], + dtype=torch.bfloat16, + ) + torch.testing.assert_close(output, expected) + + +@pytest.mark.parametrize( + "bad_input", + [ + pytest.param((torch.empty(1), torch.empty(1)), id="prequantized_fp4_tuple"), + pytest.param(torch.ones((2, 32), dtype=torch.float8_e4m3fn), id="fp8_activation"), + ], +) +def test_w4a16_nvfp4_linear_rejects_quantized_input(bad_input): + """W4A16 has no activation scale, so an upstream FP4/FP8 fusion must fail + loudly here rather than silently produce garbage.""" + method = W4A16NVFP4LinearMethod() + module = SimpleNamespace( + weight=torch.empty((4, 16), dtype=torch.uint8), + weight_scale=torch.empty((128 * 4,), dtype=torch.uint8), + _w4a16_weight_scale_linear=torch.empty((128 * 4,), dtype=torch.uint8), + weight_scale_2=torch.tensor([0.25], dtype=torch.float32), + dtype=torch.bfloat16, + out_features=4, + scaling_vector_size=16, + inv_input_scale=None, + pre_quant_scale=None, + use_custom_cublas_mm=False, + ) + + with pytest.raises(RuntimeError, match="high-precision input"): + method.apply(module, bad_input, bias=None) + + +def test_w4a16_nvfp4_linear_scale_cache_is_nonpersistent_buffer(): + with patch.object(MarlinNVFP4LinearMethod, "is_supported", return_value=False): + linear = Linear( + 32, + 4, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + reduce_output=False, + ) + + scale_linear = torch.arange(128 * 4, dtype=torch.int32).to(torch.uint8) + with patch( + "torch.ops.trtllm.block_scale_interleave_reverse", + return_value=scale_linear, + create=True, + ): + linear.quant_method.cache_derived_state(linear) + + assert linear._buffers["_w4a16_weight_scale_linear"].data_ptr() == scale_linear.data_ptr() + assert "_w4a16_weight_scale_linear" not in linear.state_dict() + + +def test_w4a16_nvfp4_linear_uses_marlin_op_after_weight_transform(): + method = MarlinNVFP4LinearMethod() + input_tensor = torch.ones((1, 32), dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((8, 16), dtype=torch.int32), + weight_scale=torch.empty((2, 128), dtype=torch.float8_e4m3fn), + weight_global_scale=torch.tensor([0.5], dtype=torch.bfloat16), + dtype=torch.bfloat16, + in_features=32, + out_features=3, + pre_quant_scale=None, + _marlin_size_k=32, + _marlin_size_n=32, + ) + captured = {} + + def fake_marlin_nvfp4_gemm(input_arg, weight, **kwargs): + captured["input"] = input_arg + captured["weight"] = weight + captured.update(kwargs) + return torch.ones((input_arg.shape[0], kwargs["size_n"]), dtype=kwargs["out_dtype"]) + + with patch( + "torch.ops.trtllm.marlin_nvfp4_gemm", + side_effect=fake_marlin_nvfp4_gemm, + create=True, + ): + output = method.apply(module, input_tensor, bias=None) + + assert captured["input"] is input_tensor + assert captured["weight"] is module.weight + assert captured["out_dtype"] is torch.bfloat16 + assert captured["bias"] is None + assert captured["scale_b"] is module.weight_scale + assert captured["weight_global_scale"] is module.weight_global_scale + assert output.shape == (1, 3) + + +def test_w4a16_nvfp4_linear_marlin_restores_high_rank_input_shape(): + method = MarlinNVFP4LinearMethod() + input_tensor = torch.ones((2, 9, 32), dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((8, 16), dtype=torch.int32), + weight_scale=torch.empty((2, 128), dtype=torch.float8_e4m3fn), + weight_global_scale=torch.tensor([0.5], dtype=torch.bfloat16), + dtype=torch.bfloat16, + in_features=32, + out_features=3, + pre_quant_scale=None, + _marlin_size_k=32, + _marlin_size_n=32, + ) + captured = {} + + def fake_marlin_nvfp4_gemm(input_arg, weight, **kwargs): + captured["input_shape"] = input_arg.shape + return torch.ones((input_arg.shape[0], kwargs["size_n"]), dtype=kwargs["out_dtype"]) + + with patch( + "torch.ops.trtllm.marlin_nvfp4_gemm", + side_effect=fake_marlin_nvfp4_gemm, + create=True, + ): + output = method.apply(module, input_tensor, bias=bias) + + assert captured["input_shape"] == (18, 32) + assert output.shape == (2, 9, 3) + expected = torch.tensor([2.0, 3.0, 4.0], dtype=torch.bfloat16).expand(2, 9, 3) + torch.testing.assert_close(output, expected) + + +def test_w4a16_nvfp4_linear_marlin_applies_pre_quant_scale_once(): + method = MarlinNVFP4LinearMethod() + input_tensor = torch.ones((1, 32), dtype=torch.bfloat16) + module = SimpleNamespace( + weight=torch.empty((8, 16), dtype=torch.int32), + weight_scale=torch.empty((2, 128), dtype=torch.float8_e4m3fn), + weight_global_scale=torch.tensor([0.5], dtype=torch.bfloat16), + dtype=torch.bfloat16, + in_features=32, + out_features=3, + pre_quant_scale=torch.full((32,), 2.0, dtype=torch.bfloat16), + _marlin_size_k=32, + _marlin_size_n=32, + ) + captured = {} + + def fake_marlin_nvfp4_gemm(input_arg, weight, **kwargs): + captured["input"] = input_arg + return torch.ones((input_arg.shape[0], kwargs["size_n"]), dtype=kwargs["out_dtype"]) + + with patch( + "torch.ops.trtllm.marlin_nvfp4_gemm", + side_effect=fake_marlin_nvfp4_gemm, + create=True, + ): + method.apply(module, input_tensor, bias=None) + + torch.testing.assert_close(captured["input"], input_tensor * module.pre_quant_scale) + + +@pytest.mark.parametrize( + ("dtype", "use_fused_gemm_allreduce"), + [ + (torch.float16, False), + (torch.bfloat16, True), + ], +) +def test_w4a16_nvfp4_marlin_selection_requires_supported_module(dtype, use_fused_gemm_allreduce): + module = SimpleNamespace(dtype=dtype, use_fused_gemm_allreduce=use_fused_gemm_allreduce) + + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=120), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + assert not MarlinNVFP4LinearMethod.is_supported(module) + + +@pytest.mark.parametrize( + ("sm_version", "dtype", "allowed_backends", "expect_marlin"), + [ + # Marlin is the default on SM120/121, opt-in on SM89-99, never on SM100. + (120, torch.bfloat16, None, True), + (121, torch.bfloat16, None, True), + (89, torch.bfloat16, ["marlin"], True), + (90, torch.bfloat16, ["marlin"], True), + (89, torch.bfloat16, None, False), + (90, torch.bfloat16, None, False), + (100, torch.bfloat16, None, False), + (120, torch.float16, None, False), # Marlin is bf16-only + ], +) +def test_w4a16_nvfp4_linear_method_selection(sm_version, dtype, allowed_backends, expect_marlin): + kwargs = {} if allowed_backends is None else {"nvfp4_allowed_backends": allowed_backends} + + with ( + patch("tensorrt_llm._torch.modules.linear.get_sm_version", return_value=sm_version), + patch("torch.ops.trtllm.marlin_nvfp4_gemm", create=True), + patch("torch.ops.trtllm.gptq_marlin_repack", create=True), + ): + linear = Linear( + 32, + 32, + bias=False, + dtype=dtype, + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4), + reduce_output=False, + **kwargs, + ) + + assert isinstance(linear.quant_method, W4A16NVFP4LinearMethod) + assert isinstance(linear.quant_method, MarlinNVFP4LinearMethod) is expect_marlin + + +def test_w4a16_nvfp4_post_load_preserves_checkpoint_weight_global_scale(): + method = W4A16NVFP4LinearMethod() + module = SimpleNamespace( + input_scale=None, + inv_input_scale=None, + alpha=None, + weight_scale_2=torch.empty([1], dtype=torch.float32), + tmp_nvfp4_input_scales_list=[torch.tensor(1.0, dtype=torch.float32)], + tmp_nvfp4_weight_scale_2_list=[torch.tensor(0.25, dtype=torch.float32)], + ) + + method.process_weights_after_loading_vanilla(module) + + assert module.input_scale is None + assert module.inv_input_scale is None + assert module.alpha is None + torch.testing.assert_close(module.weight_scale_2, torch.tensor([0.25], dtype=torch.float32)) + assert not hasattr(module, "tmp_nvfp4_input_scales_list") + assert not hasattr(module, "tmp_nvfp4_weight_scale_2_list") + + +def test_lm_head_uses_w4a16_nvfp4_quant_method_for_packed_lm_head(): + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16_NVFP4) + + lm_head = LMHead( + num_embeddings=3, embedding_dim=16, dtype=torch.float16, quant_config=quant_config + ) + + assert isinstance(lm_head.quant_method, W4A16NVFP4LinearMethod) + assert lm_head.weight.dtype == torch.uint8 + assert lm_head.weight.shape == (3, 8) + assert lm_head.weight_scale.shape == (128 * 4,) + assert lm_head.weight_scale_2.shape == (1,) + + +@pytest.mark.parametrize( + ("checkpoint_has_scale", "exclude_modules", "expected_width"), + [ + (True, None, 1344), + (True, ["lm_head"], 2688), + (False, None, 2688), + ], +) +def test_causal_lm_head_uses_global_w4a16_nvfp4_config( + checkpoint_has_scale, + exclude_modules, + expected_width, +): + quant_config = QuantConfig( + quant_algo=QuantAlgo.W4A16_NVFP4, + exclude_modules=exclude_modules, + ) + model_config = ModelConfig( + pretrained_config=SimpleNamespace( + torch_dtype=torch.float16, + tie_word_embeddings=False, + ), + quant_config=quant_config, + ) + + with patch.object( + DecoderModelForCausalLM, + "_checkpoint_has_lm_head_scale", + return_value=checkpoint_has_scale, + ): + causal_lm = DecoderModelForCausalLM( + torch.nn.Module(), + config=model_config, + hidden_size=2688, + vocab_size=32, + ) + + assert causal_lm.lm_head.weight.shape == (32, expected_width) + if checkpoint_has_scale and exclude_modules is None: + assert causal_lm.lm_head.quant_config is quant_config + assert isinstance(causal_lm.lm_head.quant_method, W4A16NVFP4LinearMethod) + else: + assert not causal_lm.lm_head.has_any_quant diff --git a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py index a16dd771923c..586973673f87 100644 --- a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py +++ b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py @@ -210,9 +210,16 @@ def test_mlp_fp4out_min_m_switch(): fake_qm = types.SimpleNamespace(_input_prepare=lambda *a, **k: None) mlp._use_fused_gelu = True mlp._use_fused_gelu_fp4out = True - mlp.up_proj.quant_method = fake_qm + up = torch.nn.Identity() + up.quant_method = fake_qm + up.has_nvfp4_activation_quantization = True + mlp.up_proj = up down = torch.nn.Identity() # skip the real down GEMM down.quant_method = fake_qm + down.has_nvfp4_activation_quantization = True + down.force_dynamic_quantization = False + down.input_scale = torch.ones(1) + down.pre_quant_scale = None mlp.down_proj = down seen = {} diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 0b9ee7ffc9b6..0d829aeb1b54 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + import sys import pytest @@ -8,8 +23,7 @@ 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 (is_nvfp4_marlin_supported_sm, - model_extra_attrs) +from tensorrt_llm._torch.utils import 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 @@ -752,9 +766,10 @@ def test_fp4_linear_cuda_core(dtype, mnk): @pytest.mark.skipif( - not is_nvfp4_marlin_supported_sm(), - reason="Marlin NVFP4 backend runs on Ada (SM89) and Hopper (SM90-99)", + not (89 <= get_sm_version() < 100 or get_sm_version() in (120, 121)), + reason="Dense Marlin NVFP4 runs on SM89-99 and SM120/121", ) +@pytest.mark.parametrize("quant_algo", [QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize( "mnk", @@ -776,7 +791,10 @@ def test_fp4_linear_cuda_core(dtype, mnk): (3, 176, 144), (128, 928, 1360), ]) -def test_fp4_linear_marlin(dtype, mnk): +def test_fp4_linear_marlin(quant_algo, dtype, mnk): + if quant_algo == QuantAlgo.NVFP4 and get_sm_version() in (120, 121): + pytest.skip( + "Marlin backend shouldn't be used for NVFP4 quant on SM120/121") SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk torch.manual_seed(0) @@ -798,7 +816,7 @@ def test_fp4_linear_marlin(dtype, mnk): out_features=OUTPUT_SIZE, bias=False, dtype=dtype, - quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + quant_config=QuantConfig(quant_algo=quant_algo), nvfp4_allowed_backends=['marlin'], # key ) 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 9a703bc3d82f..57ce07d0dcf3 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_neither_ada_nor_hopper_unittest, unittest_name_func +from utils.util import getSMVersion, unittest_name_func import tensorrt_llm # noqa: F401 # registers torch.ops.tensorrt_llm / torch.ops.trtllm ops @@ -65,7 +65,10 @@ 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_neither_ada_nor_hopper_unittest + @unittest.skipUnless( + 89 <= getSMVersion() < 100 or getSMVersion() in (120, 121), + "Marlin NVFP4 GEMM requires SM89-99 or SM120/121", + ) def test_nvfp4_marlin_gemm( self, input_dim, output_dim, batch_size, sf_vec_size, alpha, act_dtype ):