diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 5c2e60c4706a..89a3b279c8db 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/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.""" 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 af612a4ff0d3..80289dd56a61 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/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], diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index a39b53dcdbdb..30da85e3efb7 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/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, diff --git a/tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py b/tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py deleted file mode 100644 index da5e9d3afce7..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/triton_dequant_nvfp4.py +++ /dev/null @@ -1,334 +0,0 @@ -# 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 0d77f4c57fe4..599af9d834b4 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1901,87 +1901,6 @@ def transform_weights(self, module: Linear) -> None: 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 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. - LinearMethodBase.transform_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 550ff8641d80..94500e8a1ad9 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6898,34 +6898,6 @@ 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_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 @skip_post_hopper @pytest.mark.skip_less_device_memory(80000) 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 b11214e16256..f969e3ec5a1c 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -106,7 +106,6 @@ l0_dgx_h100: - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp2pp2-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp2pp2-fp8kv=True-attn_backend=TRTLLM-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_guided_decoding_4gpus[xgrammar] - - 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]