diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 574c74d9b190..f643eff0ef1e 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/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.""" 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 5b9574a10065..faa584c6570e 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/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], diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index e6769b241646..10c09634f308 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -2892,6 +2892,84 @@ def process_weights_after_loading(self, module: torch.nn.Module): super().process_weights_after_loading(module) +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, diff --git a/tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py b/tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py new file mode 100644 index 000000000000..da5e9d3afce7 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py @@ -0,0 +1,334 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + +Pipeline: scatter active_mask[E] from routing -> static (E, N/BN, K/BK) grid +-> inactive blocks early-return, active blocks dequant their tile. Downstream +fused_moe only reads rows in token_selected_experts, so leaving inactive +rows uninitialized is safe. +""" + +import torch +import triton # type: ignore[import] +import triton.language as tl # type: ignore[import] + +# 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 +# the device tensor lazily on first use. +_E2M1_CODEBOOK = [ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +] # yapf: disable + +# Per-device cache so we don't reallocate the table on every call. +_E2M1_CODEBOOK_CACHE: "dict[torch.device, torch.Tensor]" = {} + + +def _get_e2m1_codebook(device: torch.device) -> torch.Tensor: + table = _E2M1_CODEBOOK_CACHE.get(device) + if table is None: + table = torch.tensor(_E2M1_CODEBOOK, dtype=torch.float32, device=device) + _E2M1_CODEBOOK_CACHE[device] = table + return table + + +def build_active_expert_mask( + token_selected_experts: torch.Tensor, num_experts: int +) -> torch.Tensor: + """Sync-free, CUDA-graph-safe active-expert mask via scalar ``scatter_`` + (the fancy-index form ``mask[ids] = 1`` is illegal under stream capture). + + Caller must pre-clamp ids into ``[0, num_experts)``. + """ + mask = torch.zeros(num_experts, dtype=torch.uint8, device=token_selected_experts.device) + mask.scatter_(0, token_selected_experts.reshape(-1).long(), 1) + return mask + + +@triton.jit +def _dequant_nvfp4_active_kernel( + # Inputs + packed_weight_ptr, + scale_ptr, + weight_scale_2_ptr, + active_mask_ptr, + e2m1_table_ptr, # [16] fp32 codebook + # Output + out_ptr, + # Strides (element counts, not bytes) + pw_stride_e, + pw_stride_n, + sc_stride_e, + sc_stride_n, + out_stride_e, + out_stride_n, + # Shapes (runtime) + N, + K, + # Compile-time + SF_VEC: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """Per-block dequant of one tile of one expert's weight.""" + pid_e = tl.program_id(0) + pid_n = tl.program_id(1) + pid_k = tl.program_id(2) + + # ---- Active-mask early exit ---- + is_active = tl.load(active_mask_ptr + pid_e) + if is_active == 0: + return + + # Output element offsets + n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) # [BLOCK_N] + k_offs = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) # [BLOCK_K] + n_mask = n_offs < N + k_mask = k_offs < K + full_mask = n_mask[:, None] & k_mask[None, :] # [BLOCK_N, BLOCK_K] + + # ---- Load packed FP4 bytes ---- + # Each k position corresponds to byte k//2 with nibble shift (k%2)*4. + # Adjacent k positions share a byte: redundant loads, but cache friendly. + packed_idx = k_offs // 2 # [BLOCK_K] + nibble_shift = (k_offs % 2) * 4 + w_offs = pid_e * pw_stride_e + n_offs[:, None] * pw_stride_n + packed_idx[None, :] + packed = tl.load(packed_weight_ptr + w_offs, mask=full_mask, other=0).to(tl.int32) + nibble = (packed >> nibble_shift[None, :]) & 0xF # [BLOCK_N, BLOCK_K] + + # ---- E2M1 codebook lookup via gather ---- + # 16-element table lives in L1 / constant cache after first touch. + # nibble has shape [BLOCK_N, BLOCK_K]; the load fans out per lane and + # the hardware broadcasts duplicate indices. + val = tl.load(e2m1_table_ptr + nibble) # [BLOCK_N, BLOCK_K] fp32 + + # ---- Per-block FP8 (e4m3) scale ---- + # Each sf_vec_size consecutive k positions share one scale byte. + sf_idx = k_offs // SF_VEC # [BLOCK_K] + s_offs = pid_e * sc_stride_e + n_offs[:, None] * sc_stride_n + sf_idx[None, :] + scale_byte = tl.load(scale_ptr + s_offs, mask=full_mask, other=0) + # Reinterpret uint8 bits as fp8 e4m3, then convert to fp32. + scale_fp = scale_byte.to(tl.float8e4nv, bitcast=True).to(tl.float32) + + # ---- Per-tensor scale (scalar per expert) ---- + s2 = tl.load(weight_scale_2_ptr + pid_e) + + # ---- Combine and store ---- + out_val = val * scale_fp * s2 + o_offs = pid_e * out_stride_e + n_offs[:, None] * out_stride_n + k_offs[None, :] + tl.store(out_ptr + o_offs, out_val.to(out_ptr.dtype.element_ty), mask=full_mask) + + +def dequant_nvfp4_active_triton( + packed_weight: torch.Tensor, + scale_linear: torch.Tensor, + weight_scale_2: torch.Tensor, + active_mask: torch.Tensor, + *, + target_dtype: torch.dtype = torch.bfloat16, + sf_vec_size: int = 16, + block_n: int = 32, + block_k: int = 64, +) -> torch.Tensor: + """Triton-based active-only NVFP4 weight dequant. + + Args: + packed_weight: ``[E, N, K_packed]`` uint8 -- FP4 nibbles, two per byte. + scale_linear: ``[E, N_pad, K_sf_pad]`` uint8 -- per-block FP8 (e4m3) + scale in **linear** (un-swizzled) layout. The kernel uses tensor + strides directly, so padding on N/K_sf is fine as long as the + stride math points to the right elements. + weight_scale_2: ``[E]`` float32 -- per-tensor scale. + active_mask: ``[E]`` uint8 -- 1 for experts that need dequanting. + target_dtype: ``torch.bfloat16`` or ``torch.float16``. + sf_vec_size: NVFP4 per-block scale vector size (fixed at 16). + block_n, block_k: Triton tile shape. ``block_k`` should be a + multiple of ``sf_vec_size`` so each tile covers an integer + number of scale blocks. + + Returns: + ``[E, N, K]`` (with ``K = K_packed * 2``) in ``target_dtype``. + 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}" + ) + + E, N, K_packed = packed_weight.shape + K = K_packed * 2 + device = packed_weight.device + + if active_mask.dtype != torch.uint8: + active_mask = active_mask.to(torch.uint8) + + out = torch.empty(E, N, K, dtype=target_dtype, device=device) + e2m1_table = _get_e2m1_codebook(device) + + grid = (E, triton.cdiv(N, block_n), triton.cdiv(K, block_k)) + _dequant_nvfp4_active_kernel[grid]( + packed_weight, + scale_linear, + weight_scale_2, + active_mask, + e2m1_table, + out, + # strides (in elements) + packed_weight.stride(0), + packed_weight.stride(1), + scale_linear.stride(0), + scale_linear.stride(1), + out.stride(0), + out.stride(1), + # shapes + N, + K, + # constexpr + SF_VEC=sf_vec_size, + BLOCK_N=block_n, + BLOCK_K=block_k, + ) + return out + + +@triton.jit +def _dequant_nvfp4_linear_kernel( + # Inputs + packed_weight_ptr, + scale_ptr, + weight_scale_2_ptr, # pointer to one fp32 scalar (per-tensor) + e2m1_table_ptr, + # Output + out_ptr, + # Strides (in elements) + pw_stride_n, + sc_stride_n, + out_stride_n, + # Shapes (runtime) + N, + K, + # Compile-time + SF_VEC: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """Per-block dequant of one tile of a single 2D weight matrix. + + Dedicated 2D path for ``NVFP4LinearMethod``-style weights (one matrix, + one per-tensor scale, no expert dim or active mask). Codebook gather and + FP8 e4m3 -> fp32 conversion match the MoE kernel. + """ + pid_n = tl.program_id(0) + pid_k = tl.program_id(1) + + n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + k_offs = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + n_mask = n_offs < N + k_mask = k_offs < K + full_mask = n_mask[:, None] & k_mask[None, :] + + # ---- Load packed FP4 bytes ---- + packed_idx = k_offs // 2 + nibble_shift = (k_offs % 2) * 4 + w_offs = n_offs[:, None] * pw_stride_n + packed_idx[None, :] + packed = tl.load(packed_weight_ptr + w_offs, mask=full_mask, other=0).to(tl.int32) + nibble = (packed >> nibble_shift[None, :]) & 0xF + + # ---- E2M1 codebook gather ---- + val = tl.load(e2m1_table_ptr + nibble) + + # ---- Per-block FP8 (e4m3) scale ---- + sf_idx = k_offs // SF_VEC + s_offs = n_offs[:, None] * sc_stride_n + sf_idx[None, :] + scale_byte = tl.load(scale_ptr + s_offs, mask=full_mask, other=0) + scale_fp = scale_byte.to(tl.float8e4nv, bitcast=True).to(tl.float32) + + # ---- Per-tensor scale (single scalar, broadcast to the tile) ---- + s2 = tl.load(weight_scale_2_ptr) + + out_val = val * scale_fp * s2 + o_offs = n_offs[:, None] * out_stride_n + k_offs[None, :] + tl.store(out_ptr + o_offs, out_val.to(out_ptr.dtype.element_ty), mask=full_mask) + + +def dequant_nvfp4_2d_triton( + packed_weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_scale_2: torch.Tensor, + *, + target_dtype: torch.dtype = torch.bfloat16, + sf_vec_size: int = 16, + block_n: int = 32, + block_k: int = 64, +) -> torch.Tensor: + """2D (Linear) NVFP4 dequant via a dedicated Triton kernel. + + Distinct from the MoE 3D path: no expert dim, no active mask, the + per-tensor scale is a single scalar. + + Args: + packed_weight: ``[N, K_packed]`` uint8 -- FP4 nibbles, two per byte. + weight_scale: per-block FP8 (e4m3) scale in **linear** (un-swizzled) + layout. Accepted as either: + + * a flat 1-D buffer of length ``pad_up(N, 128) * pad_up(K/sf, 4)`` + (what ``NVFP4LinearMethod.create_weights`` allocates), or + * a 2-D buffer of shape ``[pad_rows, pad_cols]``. + weight_scale_2: per-tensor FP32 scale (any shape with a single + element; only ``data_ptr()`` is consumed by the kernel). + target_dtype: BF16 or FP16. + sf_vec_size: NVFP4 per-block size (16). + block_n, block_k: Triton tile shape. ``block_k`` must be a multiple + of ``sf_vec_size``. + + 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}" + ) + + N, K_packed = packed_weight.shape + K = K_packed * 2 + device = packed_weight.device + + # 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)}") + + 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. + weight_scale_2 = weight_scale_2.reshape(-1) + + grid = (triton.cdiv(N, block_n), triton.cdiv(K, block_k)) + _dequant_nvfp4_linear_kernel[grid]( + packed_weight, + weight_scale, + weight_scale_2, + e2m1_table, + out, + packed_weight.stride(0), + weight_scale.stride(0), + out.stride(0), + N, + K, + SF_VEC=sf_vec_size, + BLOCK_N=block_n, + BLOCK_K=block_k, + ) + return out diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 2d3482fe1b34..84a695b1f974 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1861,6 +1861,87 @@ def post_load_weights(self, module: Linear): module.rebuild_tensor_metadata) +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. + """ + + def post_load_weights(self, module: Linear): + # Skip parent's 32x16 weight padding (apply() accepts [N, K/2] as-is) + # and un-swizzle per-block scale once at load. + LinearMethodBase.post_load_weights(self, module) + 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)) + + def apply(self, module: Linear, input: torch.Tensor, + bias: Optional[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)") + + ## FP8 input from upstream FMHA pre-quant: invert by / module.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" + input = (input.to(module.dtype) / module.inv_input_scale).to( + module.dtype) + + 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") + 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, + module.weight_scale_2, + target_dtype=module.dtype, + sf_vec_size=module.scaling_vector_size, + ) + + if module.use_custom_cublas_mm: + output_buffer_kind = ( + int(BufferKind.NCCL_WINDOW) + if self.supports_nccl_symmetric_memory_window_output + and module.all_reduce is not None + and module.all_reduce.uses_nccl_symmetric_memory_window() else + int(BufferKind.DEFAULT)) + group = (module.mapping.tp_group + if output_buffer_kind == int(BufferKind.NCCL_WINDOW) + and module.mapping is not None else None) + output = torch.ops.trtllm.cublas_mm( + input, + weight_deq.t(), + bias, + out_dtype=None, + output_buffer_kind=output_buffer_kind, + group=group, + ) + else: + output = F.linear(input, weight_deq, bias) + + if original_shape is not None: + output = output.reshape(*original_shape[:-1], output.shape[-1]) + return output + + class W4A8NVFP4FP8LinearMethod(LinearMethodBase): def create_weights(self, module: Linear, in_features: int, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ae173a79e0e1..68f4c8edc31d 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7141,6 +7141,35 @@ def test_nvfp4_4gpus_online_eplb(self, moe_backend): layer_updates_per_iter=2) self._run_nvfp4_4gpus_eplb(moe_backend, eplb_config, model_path) + @skip_pre_hopper + @skip_post_blackwell + @pytest.mark.skip_less_mpi_world_size(4) + @pytest.mark.skip_less_device_memory(80000) + def test_nvfp4_4gpus_hopper_w4a16(self): + """W4A16 NVFP4 dequant fallback on Hopper (SM 90), MTP draft_len=4.""" + model_path = f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4" + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + mamba_ssm_cache_dtype="float16", + free_gpu_memory_fraction=0.5, + ) + max_batch_size = 32 + cuda_graph_config = CudaGraphConfig(max_batch_size=max_batch_size, + enable_padding=True) + mtp_config = MTPDecodingConfig(max_draft_len=4) + pytorch_config = dict(cuda_graph_config=cuda_graph_config) + with LLM( + model_path, + kv_cache_config=kv_cache_config, + max_batch_size=max_batch_size, + tensor_parallel_size=4, + speculative_config=mtp_config, + **pytorch_config, + ) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + @skip_pre_hopper @pytest.mark.skip_less_mpi_world_size(4) @pytest.mark.skip_less_device_memory(40000) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 29e0090d0c0d..34110b214d1d 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -102,6 +102,7 @@ l0_dgx_h100: - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_dp_off-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_fp8_4gpus[attention_dp_on-cpp_mamba_cache] + - accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_4gpus_hopper_w4a16 - test_e2e.py::test_ptp_quickstart_advanced_bs1 - test_e2e.py::test_ptp_quickstart_advanced_deepseek_v3_lite_4gpus_adp_balance[DeepSeek-V3-Lite-FP8-DeepSeek-V3-Lite/fp8] - test_e2e.py::test_trtllm_bench_llmapi_launch[pytorch_backend-llama-v3-llama3-8b]