[TRTLLM-12288][feat] Support Nemotron-H nvfp4 ckpt on Hopper

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -13,7 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import re
+from contextlib import contextmanager
 from dataclasses import replace
 from typing import TYPE_CHECKING
 
@@ -31,6 +33,7 @@
 from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.logger import logger
 from tensorrt_llm.lora_helper import LoraConfig
+from tensorrt_llm.models.modeling_utils import QuantAlgo  # noqa: E402
 
 from ..attention_backend import AttentionMetadata
 from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams
@@ -39,7 +42,11 @@
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.fused_moe import MoEWeightLoadingMode, create_moe
-from ..modules.linear import Linear, TensorParallelMode
+from ..modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE
+from ..modules.fused_moe.quantization import (NVFP4CutlassFusedMoEMethod,
+                                              W4A16NVFP4CutlassFusedMoEMethod)
+from ..modules.linear import (Linear, NVFP4LinearMethod, TensorParallelMode,
+                              W4A16NVFP4LinearMethod)
 from ..modules.mamba.mamba2_mixer import Mamba2Mixer
 from ..modules.mlp import MLP
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
@@ -431,10 +438,14 @@ def __init__(
 
         quant_mode = (model_config.quant_config.quant_mode
                       if model_config.quant_config is not None else None)
-        self.is_nvfp4 = quant_mode is not None and quant_mode.has_nvfp4()
+        # We don't use the RMSNorm+NVFP4 on SM < 100
+        _has_fp4_hw = get_sm_version() >= 100
+        self.is_nvfp4 = (quant_mode is not None and quant_mode.has_nvfp4()
+                         and _has_fp4_hw)
         # For MIXED_PRECISION models, the global quant_mode is QuantMode(0). Check per-layer
         # quant_config_dict to see if this specific layer is NVFP4-quantized.
-        if not self.is_nvfp4 and model_config.quant_config_dict is not None:
+        if (not self.is_nvfp4 and _has_fp4_hw
+                and model_config.quant_config_dict is not None):
             layer_prefix = f"model.layers.{layer_idx}."
             for key, cfg in model_config.quant_config_dict.items():
                 if key.startswith(layer_prefix) and cfg.quant_mode.has_nvfp4():
@@ -506,6 +517,11 @@ def __init__(
             )
             if fuse_allreduce_norm:
                 self.mixer.out_proj.reduce_output = False
+            # Hopper: route RMSNormGated to its bf16 Triton fallback
+            # (fused_gated_rmsnorm_quant is SM100-only).
+            if not _has_fp4_hw:
+                self.mixer.is_nvfp4 = False
+                self.mixer.norm.is_nvfp4 = False
         elif layer_type == "-":
             self.mixer = MLPLayer(
                 model_config,
@@ -754,6 +770,93 @@ def forward(
         return hidden_states
 
 
+def _force_moe_backend_for_w4a16_on_hopper(
+        model_config: NemotronHModelConfig) -> None:
+    """SM<100 + NVFP4: force ``moe_backend=CUTLASS`` (only backend with the
+    W4A16 fallback) and disable attention FP4 output fusion.
+    """
+    if get_sm_version() >= 100:
+        return
+
+    # NVFP4 may live in global quant_config OR per-layer quant_config_dict
+    # (MIXED_PRECISION ckpts).
+    qcfg = model_config.quant_config
+    has_nvfp4 = qcfg is not None and qcfg.layer_quant_mode.has_nvfp4()
+    if not has_nvfp4 and model_config.quant_config_dict is not None:
+        has_nvfp4 = any(cfg.quant_mode.has_nvfp4()
+                        for cfg in model_config.quant_config_dict.values())
+    if not has_nvfp4:
+        return
+
+    # o_proj.has_nvfp4 stays True under W4A16 -- the property reads quant_config.
+    # Use the documented env override to keep attention output in bf16.
+    if os.environ.get("TRTLLM_ENABLE_ATTENTION_NVFP4_OUTPUT") != "0":
+        logger.warning(
+            f"Nemotron-H SM{get_sm_version()}: TRTLLM_ENABLE_ATTENTION_NVFP4_OUTPUT=0"
+        )
+        os.environ["TRTLLM_ENABLE_ATTENTION_NVFP4_OUTPUT"] = "0"
+
+    if model_config.moe_backend.upper() in ('CUTLASS', 'AUTO'):
+        return
+    logger.warning(
+        f"Nemotron-H SM{get_sm_version()}: forcing moe_backend "
+        f"'{model_config.moe_backend}' -> 'CUTLASS' for W4A16 fallback")
+    model_config._frozen = False
+    model_config.moe_backend = 'CUTLASS'
+    model_config._frozen = True
+
+
+@contextmanager
+def _use_w4a16_for_nvfp4_on_hopper():
+    """SM<100 + NVFP4: swap NVFP4 quant methods -> W4A16 fallback,
+    loosen MoE SM constraint, and disable MLP's fused relu2+FP4 quant.
+    Class-level patches; model construction is single-threaded today.
+    """
+    if get_sm_version() >= 100:
+        yield
+        return
+
+    original_linear = Linear.get_quant_method
+    original_moe = CutlassFusedMoE._get_quant_method
+    original_mlp_create_weights = MLP.create_weights
+    nvfp4_entry = CutlassFusedMoE._QUANT_SUPPORT_TABLE[QuantAlgo.NVFP4]
+    original_sm_constraint = nvfp4_entry["sm_constraint"]
+
+    def _patched_linear(self, quant_config):
+        method = original_linear(self, quant_config)
+        if type(method) is NVFP4LinearMethod:
+            return W4A16NVFP4LinearMethod()
+        return method
+
+    def _patched_moe(self):
+        method = original_moe(self)
+        if type(method) is NVFP4CutlassFusedMoEMethod:
+            return W4A16NVFP4CutlassFusedMoEMethod()
+        return method
+
+    def _patched_mlp_create_weights(self):
+        # Original sets _use_fused_relu2_quant=True for NVFP4 ckpts; off here
+        # so MLP.forward emits bf16 (the SM100-only fused kernel never runs).
+        original_mlp_create_weights(self)
+        self._use_fused_relu2_quant = False
+
+    # Allow SM 90 through can_implement(); existing entries preserved.
+    constraint_type, constraint_set = original_sm_constraint
+    nvfp4_entry["sm_constraint"] = (constraint_type,
+                                    frozenset(constraint_set) | {90})
+
+    Linear.get_quant_method = _patched_linear
+    CutlassFusedMoE._get_quant_method = _patched_moe
+    MLP.create_weights = _patched_mlp_create_weights
+    try:
+        yield
+    finally:
+        Linear.get_quant_method = original_linear
+        CutlassFusedMoE._get_quant_method = original_moe
+        MLP.create_weights = original_mlp_create_weights
+        nvfp4_entry["sm_constraint"] = original_sm_constraint
+
+
 @register_auto_model("NemotronHPuzzleForCausalLM")
 @register_auto_model("NemotronHForCausalLM")
 class NemotronHForCausalLM(SpecDecOneEngineForCausalLM[NemotronHModel,
@@ -797,10 +900,12 @@ def __init__(
             }
             model_config._frozen = True
 
-        super().__init__(
-            model=NemotronHModel(model_config),
-            model_config=model_config,
-        )
+        _force_moe_backend_for_w4a16_on_hopper(model_config)
+        with _use_w4a16_for_nvfp4_on_hopper():
+            super().__init__(
+                model=NemotronHModel(model_config),
+                model_config=model_config,
+            )
         self.model_nextn = 0
         if (model_config.spec_config is not None
                 and model_config.spec_config.spec_dec_mode.is_mtp_one_model()):
@@ -832,6 +937,16 @@ def __init__(
             self.epilogue.extend(self.draft_model.mtp_layers)
             self.epilogue.append(self.spec_worker)
 
+    def __post_init__(self):
+        # PostInitCaller metaclass invokes __post_init__ AFTER __init__ returns,
+        # so our W4A16 context manager from __init__ has already exited. For
+        # MIXED_PRECISION checkpoints, ``apply_layerwise_quant_config`` rebinds
+        # per-layer ``quant_config`` to NVFP4 and then re-runs ``create_weights``
+        # (see modeling_utils.py:543). Re-enter the context manager so the
+        # patched ``_get_quant_method`` catches that second pass.
+        with _use_w4a16_for_nvfp4_on_hopper():
+            super().__post_init__()
+
     @staticmethod
     def _normalize_puzzle_config(config):
         """Set global MoE defaults from block_configs for models with per-layer MoE params."""

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

@@ -21,13 +21,12 @@
 from .quantization import UnquantizedFusedMoEMethod
 
 # isort: off
-from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethod,
-                           FP8QDQFusedMoEMethod, MoEWeightLoadingMode,
-                           NVFP4CutlassFusedMoEMethod,
-                           INT8WoqPerChannelFusedMoEMethod,
-                           W4A8MXFP4FP8CutlassFusedMoEMethod,
-                           W4A8MXFP4MXFP8CutlassFusedMoEMethod,
-                           WFP4A16FusedMoEMethod, WInt4AFP8FusedMoEMethod)
+from .quantization import (
+    DeepSeekFP8BlockScalesFusedMoEMethod, FP8QDQFusedMoEMethod,
+    MoEWeightLoadingMode, NVFP4CutlassFusedMoEMethod,
+    INT8WoqPerChannelFusedMoEMethod, W4A16NVFP4CutlassFusedMoEMethod,
+    W4A8MXFP4FP8CutlassFusedMoEMethod, W4A8MXFP4MXFP8CutlassFusedMoEMethod,
+    WFP4A16FusedMoEMethod, WInt4AFP8FusedMoEMethod)
 # isort: on
 from .routing import BaseMoeRoutingMethod
 
@@ -445,6 +444,9 @@ def quantize_input(
         """
         x_sf = None
         if self.has_any_quant:
+            # W4A16 NVFP4 path keeps activations hp; skip FP4 quant below.
+            if isinstance(self.quant_method, W4A16NVFP4CutlassFusedMoEMethod):
+                return x, None
             if self.has_fp8_qdq or self.has_w4a8_mxfp4_fp8:
                 x, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(
                     x, self.fc31_input_dequant)
@@ -599,6 +601,19 @@ def run_moe(
         Returns:
             final_hidden_states: Output tensor from MoE computation
         """
+        # W4A16 NVFP4 fallback (SM<100).
+        if isinstance(self.quant_method, W4A16NVFP4CutlassFusedMoEMethod):
+            return self._run_moe_w4a16_nvfp4(
+                x,
+                token_selected_experts,
+                token_final_scales,
+                output_dtype=output_dtype,
+                tuner_num_tokens=tuner_num_tokens,
+                tuner_top_k=tuner_top_k,
+                moe_output=moe_output,
+                enable_alltoall=enable_alltoall,
+            )
+
         # SM120 + FP8 block scales: use Triton kernel (CUTLASS TMA fails on SM120
         # for large token counts due to cuTensorMapEncodeTiled limitations).
         if self.has_deepseek_fp8_block_scales and get_sm_version() == 120:
@@ -713,6 +728,84 @@ def run_moe(
 
         return final_hidden_states
 
+    def _run_moe_w4a16_nvfp4(
+        self,
+        x: torch.Tensor,
+        token_selected_experts: torch.Tensor,
+        token_final_scales: torch.Tensor,
+        output_dtype: Optional[torch.dtype] = None,
+        tuner_num_tokens: Optional[int] = None,
+        tuner_top_k: Optional[int] = None,
+        moe_output: Optional[torch.Tensor] = None,
+        enable_alltoall: Optional[bool] = None,
+    ) -> torch.Tensor:
+        """W4A16 fallback for NVFP4 MoE on SM<100. Active-mask dequant into
+        a static [E_total, N, K] bf16 workspace, then bf16 fused_moe with the
+        original (global) token_selected_experts. CUDA-graph capturable.
+        """
+        assert isinstance(self.quant_method, W4A16NVFP4CutlassFusedMoEMethod)
+
+        if enable_alltoall is None:
+            enable_alltoall = self.enable_alltoall
+        if output_dtype is None:
+            output_dtype = x.dtype
+
+        # Same EP id convention as the FP8 path above: global ids (or
+        # ``local_n``-padded under alltoall). Clamp to local range so the
+        # active-mask scatter is in-bounds; non-local tokens collapse onto a
+        # boundary expert (1 extra dequant/rank). ``trtllm.fused_moe`` below
+        # still gets the original global ids -- it does its own remap.
+        local_n = self.expert_size_per_partition
+        if enable_alltoall:
+            local_ids = token_selected_experts.clamp(0, local_n - 1)
+        else:
+            local_ids = (token_selected_experts - self.slot_start).clamp(
+                0, local_n - 1)
+
+        w3_w1_hp, w2_hp = self.quant_method.dequant_active_experts_to_hp(
+            self, local_ids, output_dtype)
+
+        # bf16 fused_moe with empty quant_scales (matches unquantized path).
+        result = torch.ops.trtllm.fused_moe(
+            x,
+            token_selected_experts,
+            token_final_scales,
+            w3_w1_hp,
+            self.w3_w1_bias,
+            w2_hp,
+            self.w2_bias,
+            output_dtype,
+            quant_scales=[],
+            input_sf=None,
+            swizzled_input_sf=False,
+            swiglu_alpha=self.swiglu_alpha,
+            swiglu_beta=self.swiglu_beta,
+            swiglu_limit=self.swiglu_limit,
+            tp_size=self.tp_size,
+            tp_rank=self.tp_rank,
+            ep_size=self.ep_size,
+            ep_rank=self.ep_rank,
+            cluster_size=self.cluster_size,
+            cluster_rank=self.cluster_rank,
+            enable_alltoall=enable_alltoall,
+            use_deepseek_fp8_block_scale=False,
+            use_w4_group_scaling=False,
+            use_int8_woq_per_channel=False,
+            use_mxfp8_act_scaling=False,
+            min_latency_mode=False,
+            use_fused_finalize=self.use_fused_finalize,
+            tune_max_num_tokens=self.tune_max_num_tokens,
+            tuner_num_tokens=tuner_num_tokens,
+            tuner_top_k=tuner_top_k,
+            activation_type=self.activation_type,
+            unpadded_hidden_size=self.unpadded_hidden_size,
+            out_tensor=moe_output,
+            use_dynamic_fc2_scale=False,
+        )
+        if moe_output is not None:
+            return moe_output
+        return result[0]
+
     def forward_chunk(
             self,
             x: Union[torch.Tensor, Fp4QuantizedTensor],