[TRTLLM-12288][fix] add the follow-up fix of nemotron nvfp4 on hopper

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -13,7 +13,6 @@
 # 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
@@ -438,8 +437,9 @@ def __init__(
 
         quant_mode = (model_config.quant_config.quant_mode
                       if model_config.quant_config is not None else None)
-        # We don't use the RMSNorm+NVFP4 on SM < 100
-        _has_fp4_hw = get_sm_version() >= 100
+        # We don't use the RMSNorm+NVFP4 on SM < 100 or SM >= 120
+        sm = get_sm_version()
+        _has_fp4_hw = 100 <= sm < 120
         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
@@ -772,10 +772,10 @@ def forward(
 
 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.
+    """No native NVFP4 HW (SM<100 or SM>=120) + NVFP4 ckpt: force
+    ``moe_backend=CUTLASS`` (only backend with the W4A16 fallback).
     """
-    if get_sm_version() >= 100:
+    if 100 <= get_sm_version() < 120:
         return
 
     # NVFP4 may live in global quant_config OR per-layer quant_config_dict
@@ -788,14 +788,6 @@ def _force_moe_backend_for_w4a16_on_hopper(
     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(
@@ -808,17 +800,20 @@ def _force_moe_backend_for_w4a16_on_hopper(
 
 @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.
+    """No native NVFP4 HW (SM<100 or SM>=120) + NVFP4: swap NVFP4 quant
+    methods -> W4A16 fallback, loosen MoE SM constraint, disable MLP's
+    fused relu2+FP4 quant, and disable per-instance FP4 attention output.
     Class-level patches; model construction is single-threaded today.
     """
-    if get_sm_version() >= 100:
+    sm = get_sm_version()
+    if 100 <= sm < 120:
         yield
         return
 
     original_linear = Linear.get_quant_method
     original_moe = CutlassFusedMoE._get_quant_method
     original_mlp_create_weights = MLP.create_weights
+    original_attn_create_weights = Attention.create_weights
     nvfp4_entry = CutlassFusedMoE._QUANT_SUPPORT_TABLE[QuantAlgo.NVFP4]
     original_sm_constraint = nvfp4_entry["sm_constraint"]
 
@@ -840,6 +835,13 @@ def _patched_mlp_create_weights(self):
         original_mlp_create_weights(self)
         self._use_fused_relu2_quant = False
 
+    def _patched_attn_create_weights(self):
+        # Per-instance disable of FP4 attention output (kernel is SM100-only).
+        # Mirrors what `TRTLLM_ENABLE_ATTENTION_NVFP4_OUTPUT=0` does.
+        original_attn_create_weights(self)
+        if hasattr(self.attn, 'use_nvfp4_output'):
+            self.attn.use_nvfp4_output = lambda *args, **kwargs: False
+
     # Allow SM 90 through can_implement(); existing entries preserved.
     constraint_type, constraint_set = original_sm_constraint
     nvfp4_entry["sm_constraint"] = (constraint_type,
@@ -848,12 +850,14 @@ def _patched_mlp_create_weights(self):
     Linear.get_quant_method = _patched_linear
     CutlassFusedMoE._get_quant_method = _patched_moe
     MLP.create_weights = _patched_mlp_create_weights
+    Attention.create_weights = _patched_attn_create_weights
     try:
         yield
     finally:
         Linear.get_quant_method = original_linear
         CutlassFusedMoE._get_quant_method = original_moe
         MLP.create_weights = original_mlp_create_weights
+        Attention.create_weights = original_attn_create_weights
         nvfp4_entry["sm_constraint"] = original_sm_constraint
 
 

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

@@ -885,17 +885,12 @@ def _run_moe_w4a16_nvfp4(
         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.
+        # ids are GLOBAL (alltoall dispatch does not rebase them; C++ fused_moe
+        # does). Subtract slot_start to get local indices; out-of-range ids
+        # clamp to a boundary expert -- harmless 1-expert over-dequant.
         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)
+        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)

tensorrt_llm/_torch/modules/fused_moe/quantization.py

@@ -37,6 +37,8 @@
 from ...utils import (ActivationType, replace_parameter_and_save_metadata,
                       swizzle_sf, unswizzle_sf)
 from ..linear import TensorParallelMode, load_weight_shard
+from ..triton_dequant_nvfp4 import (build_active_expert_mask,
+                                    dequant_nvfp4_active_triton)
 from .interface import MoEWeightLoadingMode
 from .moe_load_balancer import advise_tensor_pageout
 
@@ -2900,18 +2902,31 @@ class W4A16NVFP4CutlassFusedMoEMethod(NVFP4CutlassFusedMoEMethod):
     ``CutlassFusedMoE.run_moe`` dispatches here and uses an active-mask Triton
     kernel (``dequant_active_experts_to_hp``) to dequant only routed experts
     into a static [E_total, N, K] workspace, then runs the bf16 ``fused_moe``.
+
+    Online EPLB is disabled: only the resident expert scales are un-swizzled
+    at load time, so a migrated expert's scale would still be in CUTLASS
+    swizzled layout and ``dequant_active_experts_to_hp`` would interpret it
+    as linear -- producing garbage outputs. EPLB on the SM<100 fallback path
+    is not a perf goal worth supporting; subclasses can opt back in by
+    also un-swizzling the shared EPLB scale pool.
     """
 
+    eplb_support_status = EplbSupportStatus.NOT_SUPPORTED
+
+    def create_weights(self, module: torch.nn.Module):
+        super().create_weights(module)
+        # Fail fast at construction time if a module with online EPLB
+        # attached lands here.
+        self._online_eplb_not_supported(module)
+
     def process_weights_after_loading(self, module: torch.nn.Module):
         super().process_weights_after_loading(module)
 
         # Scale buffer: int32-packed FP8, viewed as uint8 has shape
         # [E, pad_up(N, 128), pad_up(K/sf_vec, 4)] -- the 3D layout
-        # block_scale_interleave_reverse accepts.
+        # block_scale_interleave_reverse explicitly supports.
         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)
 
@@ -2932,9 +2947,6 @@ def dequant_active_experts_to_hp(
         MoE loader stores alpha = amax_in*amax_w/(448*6)**2 and
         input_scale = (448*6)/amax_in).
         """
-        from .triton_dequant_nvfp4 import (build_active_expert_mask,
-                                           dequant_nvfp4_active_triton)
-
         fc31_w_scale_2 = module.fc31_alpha * module.fc31_input_scale
         fc2_w_scale_2 = module.fc2_alpha * module.fc2_input_scale
 

tensorrt_llm/_torch/modules/linear.py

@@ -14,6 +14,8 @@
 
 import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils
 from tensorrt_llm._torch.custom_ops.torch_custom_ops import BufferKind
+from tensorrt_llm._torch.modules.triton_dequant_nvfp4 import \
+    dequant_nvfp4_2d_triton
 from tensorrt_llm._torch.peft.lora.layer import LoraLayer
 from tensorrt_llm._utils import is_device_integrated, mpi_disabled
 from tensorrt_llm.bindings import ipc_nvls_supported
@@ -1888,10 +1890,14 @@ def apply(self, module: Linear, input: torch.Tensor,
                 "W4A16NVFP4LinearMethod: hp input required; disable upstream "
                 "FP4 fusion (e.g. TRTLLM_ENABLE_ATTENTION_NVFP4_OUTPUT=0)")
 
-        ## FP8 input from upstream FMHA pre-quant: invert by / module.inv_input_scale.
+        # FP8 input from upstream FMHA pre-quant: invert by / inv_input_scale.
         if input.dtype == torch.float8_e4m3fn:
-            assert module.inv_input_scale is not None, \
-                "W4A16NVFP4LinearMethod: FP8 input requires static inv_input_scale"
+            assert (
+                not module.force_dynamic_quantization
+                and module.inv_input_scale is not None), (
+                    "W4A16NVFP4LinearMethod: FP8 input requires static "
+                    "inv_input_scale (force_dynamic_quantization must be False "
+                    "and the ckpt must provide the scale)")
             input = (input.to(module.dtype) / module.inv_input_scale).to(
                 module.dtype)
 
@@ -1906,8 +1912,6 @@ def apply(self, module: Linear, input: torch.Tensor,
                 "Input dtype and pre_quant_scale dtype must match")
             input = input * module.pre_quant_scale
 
-        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,

tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py → tensorrt_llm/_torch/modules/triton_dequant_nvfp4.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Active-only NVFP4 weight dequant for MoE on SM<100 (used by
 W4A16NVFP4CutlassFusedMoEMethod). Static shapes -> CUDA-graph capturable.
@@ -10,8 +10,10 @@
 """
 
 import torch
-import triton  # type: ignore[import]
-import triton.language as tl  # type: ignore[import]
+import triton
+import triton.language as tl
+
+from tensorrt_llm.quantization.utils.fp4_utils import pad_up
 
 # E2M1 codebook (signed-magnitude nibble layout). Index 0b1000 nominally
 # encodes "-0" and is treated as 0.0. Kept as a Python list so we can build
@@ -152,16 +154,42 @@ def dequant_nvfp4_active_triton(
         Tiles belonging to inactive experts are left uninitialized; they
         are never read by the downstream MoE kernel.
     """
-    assert packed_weight.dim() == 3, "packed_weight must be 3D [E, N, K/2]"
-    assert sf_vec_size == 16, "NVFP4 fixed at 16-element blocks"
-    assert block_k % sf_vec_size == 0, (
-        f"block_k={block_k} must be a multiple of sf_vec_size={sf_vec_size}"
-    )
+    if packed_weight.dim() != 3:
+        raise ValueError("packed_weight must be 3D [E, N, K/2]")
+    if sf_vec_size != 16:
+        raise ValueError("NVFP4 fixed at 16-element blocks")
+    if block_k % sf_vec_size != 0:
+        raise ValueError(f"block_k={block_k} must be a multiple of sf_vec_size={sf_vec_size}")
 
     E, N, K_packed = packed_weight.shape
     K = K_packed * 2
     device = packed_weight.device
 
+    if packed_weight.stride(-1) != 1:
+        raise ValueError("packed_weight innermost stride must be 1 (contiguous K dim)")
+    if scale_linear.dim() != 3:
+        raise ValueError(f"scale_linear must be 3D, got {tuple(scale_linear.shape)}")
+    if scale_linear.stride(-1) != 1:
+        raise ValueError("scale_linear innermost stride must be 1 (contiguous K_sf dim)")
+    if scale_linear.shape[0] != E:
+        raise ValueError(f"scale_linear E mismatch: {scale_linear.shape[0]} vs packed_weight E={E}")
+    if not (active_mask.dim() == 1 and active_mask.is_contiguous() and active_mask.shape[0] == E):
+        raise ValueError(
+            f"active_mask must be 1D contiguous of length E={E}, got "
+            f"shape={tuple(active_mask.shape)} "
+            f"contiguous={active_mask.is_contiguous()}"
+        )
+    if not (
+        weight_scale_2.dim() == 1
+        and weight_scale_2.is_contiguous()
+        and weight_scale_2.shape[0] == E
+    ):
+        raise ValueError(
+            f"weight_scale_2 must be 1D contiguous of length E={E}, got "
+            f"shape={tuple(weight_scale_2.shape)} "
+            f"contiguous={weight_scale_2.is_contiguous()}"
+        )
+
     if active_mask.dtype != torch.uint8:
         active_mask = active_mask.to(torch.uint8)
 
@@ -287,32 +315,41 @@ def dequant_nvfp4_2d_triton(
     Returns:
         ``[N, K]`` in ``target_dtype``.
     """
-    assert packed_weight.dim() == 2, "packed_weight must be 2D [N, K/2]"
-    assert sf_vec_size == 16, "NVFP4 fixed at 16-element blocks"
-    assert block_k % sf_vec_size == 0, (
-        f"block_k={block_k} must be a multiple of sf_vec_size={sf_vec_size}"
-    )
+    if packed_weight.dim() != 2:
+        raise ValueError("packed_weight must be 2D [N, K/2]")
+    if sf_vec_size != 16:
+        raise ValueError("NVFP4 fixed at 16-element blocks")
+    if block_k % sf_vec_size != 0:
+        raise ValueError(f"block_k={block_k} must be a multiple of sf_vec_size={sf_vec_size}")
 
     N, K_packed = packed_weight.shape
     K = K_packed * 2
     device = packed_weight.device
 
+    if packed_weight.stride(-1) != 1:
+        raise ValueError("packed_weight innermost stride must be 1 (contiguous K dim)")
+
     # Reshape (possibly flat) scale to its 2D [pad_rows, pad_cols] form so
     # the kernel can use ``scale.stride(0)`` directly.
     if weight_scale.dim() == 1:
-        from tensorrt_llm.quantization.utils.fp4_utils import pad_up
-
         pad_rows = pad_up(N, 128)
         pad_cols = pad_up(K // sf_vec_size, 4)
         weight_scale = weight_scale.view(pad_rows, pad_cols)
     elif weight_scale.dim() != 2:
         raise ValueError(f"weight_scale must be 1D or 2D, got shape {tuple(weight_scale.shape)}")
+    if weight_scale.stride(-1) != 1:
+        raise ValueError("weight_scale innermost stride must be 1 (contiguous K_sf dim)")
 
     out = torch.empty(N, K, dtype=target_dtype, device=device)
     e2m1_table = _get_e2m1_codebook(device)
 
-    # The kernel reads the per-tensor scale via a single pointer load; flatten
-    # to ensure a contiguous, 1-D-addressable buffer regardless of caller shape.
+    # The kernel does ``tl.load(weight_scale_2_ptr)`` -- a single scalar read.
+    # Reject multi-element buffers explicitly: any trailing values would be
+    # silently dropped and the matrix would dequantize against the wrong scale.
+    if weight_scale_2.numel() != 1:
+        raise ValueError(
+            f"weight_scale_2 must have exactly 1 element, got {weight_scale_2.numel()}"
+        )
     weight_scale_2 = weight_scale_2.reshape(-1)
 
     grid = (triton.cdiv(N, block_n), triton.cdiv(K, block_k))

tests/integration/test_lists/qa/llm_function_core.txt

@@ -681,6 +681,7 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_off-python_mamba_cache]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-cpp_mamba_cache]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_on-python_mamba_cache]
+accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_hopper_w4a16
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpu_mtp_ar
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_eplb[moe_backend=CUTLASS]
 accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_online_eplb[moe_backend=TRTLLM]