[TRTLLM-12288][feat] revert all the changes of the PR#14775

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -13,9 +13,7 @@
 # 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
 
@@ -33,7 +31,6 @@
 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
@@ -42,11 +39,7 @@
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.fused_moe import MoEWeightLoadingMode, create_moe
-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.linear import Linear, TensorParallelMode
 from ..modules.mamba.mamba2_mixer import Mamba2Mixer
 from ..modules.mlp import MLP
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
@@ -440,14 +433,10 @@ 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
-        self.is_nvfp4 = (quant_mode is not None and quant_mode.has_nvfp4()
-                         and _has_fp4_hw)
+        self.is_nvfp4 = quant_mode is not None and quant_mode.has_nvfp4()
         # 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 _has_fp4_hw
-                and model_config.quant_config_dict is not None):
+        if not self.is_nvfp4 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():
@@ -521,11 +510,6 @@ 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,
@@ -776,87 +760,6 @@ 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
-
-
-@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,
@@ -900,12 +803,10 @@ def __init__(
             }
             model_config._frozen = True
 
-        _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,
-            )
+        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()):
@@ -937,16 +838,6 @@ 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

@@ -45,9 +45,8 @@
     DeepSeekFP8BlockScalesFusedMoEMethod, FP8QDQFusedMoEMethod,
     MoEWeightLoadingMode, MXFP8CutlassFusedMoEMethod,
     NVFP4CutlassFusedMoEMethod, INT8WoqPerChannelFusedMoEMethod,
-    W4A16NVFP4CutlassFusedMoEMethod, W4A8MXFP4FP8CutlassFusedMoEMethod,
-    W4A8MXFP4MXFP8CutlassFusedMoEMethod, WFP4A16FusedMoEMethod,
-    WInt4AFP8FusedMoEMethod)
+    W4A8MXFP4FP8CutlassFusedMoEMethod, W4A8MXFP4MXFP8CutlassFusedMoEMethod,
+    WFP4A16FusedMoEMethod, WInt4AFP8FusedMoEMethod)
 # isort: on
 from .routing import BaseMoeRoutingMethod
 
@@ -822,9 +821,6 @@ 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)
@@ -985,19 +981,6 @@ 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:
@@ -1124,84 +1107,6 @@ 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],

tensorrt_llm/_torch/modules/fused_moe/quantization.py

@@ -3056,84 +3056,6 @@ def unswizzle_scales(scale_3d, N_actual, K_actual):
         module.fc2_alpha = nn.Parameter(w2_gs, requires_grad=False)
 
 
-class W4A16NVFP4CutlassFusedMoEMethod(NVFP4CutlassFusedMoEMethod):
-    """W4A16 dequant-on-the-fly variant of NVFP4 MoE for SM<100.
-
-    Loads an unmodified NVFP4 MoE ckpt; only load-time change is un-swizzling
-    per-block scales once so the per-forward dequant skips that step.
-    ``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``.
-    """
-
-    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.
-        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)
-
-        _unswizzle_inplace(module.w3_w1_weight_scale)
-        _unswizzle_inplace(module.w2_weight_scale)
-
-    def dequant_active_experts_to_hp(
-        self,
-        module: torch.nn.Module,
-        token_selected_experts: torch.Tensor,
-        out_dtype: torch.dtype,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Active-only dequant via Triton: static [E_total, N, K] workspace,
-        active-mask kernel skips dequant for experts with no routed tokens.
-        CUDA-graph capturable.
-
-        Per-expert weight scale recovered as ``alpha * input_scale`` (NVFP4
-        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
-
-        sf_vec_size = module.scaling_vector_size
-        E_total = module.w3_w1_weight.shape[0]
-
-        active_mask = build_active_expert_mask(token_selected_experts, E_total)
-
-        # FP4 weights as uint8 (2 fp4/byte); per-block scales as uint8 to
-        # expose the unswizzled [E, N_pad, K_sf_pad] e4m3 bit layout.
-        w3_w1_packed = module.w3_w1_weight.view(torch.uint8)
-        w2_packed = module.w2_weight.view(torch.uint8)
-        w3_w1_scale = module.w3_w1_weight_scale.view(torch.uint8)
-        w2_scale = module.w2_weight_scale.view(torch.uint8)
-
-        w3_w1_hp = dequant_nvfp4_active_triton(
-            w3_w1_packed,
-            w3_w1_scale,
-            fc31_w_scale_2,
-            active_mask,
-            target_dtype=out_dtype,
-            sf_vec_size=sf_vec_size,
-        )
-        w2_hp = dequant_nvfp4_active_triton(
-            w2_packed,
-            w2_scale,
-            fc2_w_scale_2,
-            active_mask,
-            target_dtype=out_dtype,
-            sf_vec_size=sf_vec_size,
-        )
-
-        return w3_w1_hp, w2_hp
-
-
 class NVFP4CuteDslFusedMoEMethod(NVFP4CutlassFusedMoEMethod):
 
     def load_expert_w3_w1_weight(self,