From 3f5de202de89a259851ce02516c24ba4d486dca7 Mon Sep 17 00:00:00 2001 From: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com> Date: Sun, 10 May 2026 04:33:15 -0700 Subject: [PATCH 1/3] [None][perf] Add AutoDeploy NVFP4 RMSNorm quant fusion Signed-off-by: Tal Cherckez <127761168+tcherckez-nvidia@users.noreply.github.com> --- .../_torch/auto_deploy/config/default.yaml | 4 + .../custom_ops/distributed/trtllm_dist.py | 94 ++++ .../custom_ops/normalization/rms_norm.py | 169 ++++++ .../library/fuse_rmsnorm_quant_nvfp4.py | 531 ++++++++++++++++++ .../auto_deploy/utils/multi_stream_utils.py | 23 +- .../custom_ops/test_multi_stream_moe.py | 56 ++ .../library/test_quant_fusion.py | 300 ++++++++++ 7 files changed, 1173 insertions(+), 4 deletions(-) create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 166ff6b965f4..0f7f61cf7967 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -217,6 +217,10 @@ transforms: rmsnorm_backend: flashinfer gated_rmsnorm_backend: triton requires_shape_prop: true + fuse_rmsnorm_quant_nvfp4: + stage: post_load_fusion + enabled: true + requires_shape_prop: true fuse_gdn_gating: stage: post_load_fusion fuse_l2norm: diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/distributed/trtllm_dist.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/distributed/trtllm_dist.py index 136cebaf7c8f..a3573fca1281 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/distributed/trtllm_dist.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/distributed/trtllm_dist.py @@ -23,6 +23,7 @@ import torch +import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.distributed import AllReduce, allgather from tensorrt_llm._torch.distributed.symm_mem_allgather import SymmetricMemoryAllGather from tensorrt_llm._torch.modules.linear import AllReduceFusionOp, AllReduceParams, AllReduceStrategy @@ -200,6 +201,99 @@ def trtllm_fused_allreduce_residual_rmsnorm_fake( return torch.empty_like(tensor), torch.empty_like(tensor) +@torch.library.custom_op( + "dist::trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4", + mutates_args=(), + device_types="cuda", +) +def trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4( + tensor: torch.Tensor, + residual: torch.Tensor, + norm_weight: torch.Tensor, + scale: torch.Tensor, + eps: float, + strategy: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused allreduce + residual + RMSNorm + NVFP4 quantization.""" + all_reduce_params = AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, + bias=None, + residual=residual, + norm_weight=norm_weight, + scale=scale, + eps=eps, + ) + quant_fp4, scale_factor, residual_out = trtllm_allreduce( + tensor, ReduceOp.SUM, strategy=strategy, all_reduce_params=all_reduce_params + ) + return quant_fp4, scale_factor, residual_out + + +@trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4.register_fake +def trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4_fake( + tensor: torch.Tensor, + residual: torch.Tensor, + norm_weight: torch.Tensor, + scale: torch.Tensor, + eps: float, + strategy: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del norm_weight, scale, eps, strategy + fp4_shape, scale_shape = fp4_utils.get_fp4_shape(tensor.shape, 16) + return ( + tensor.new_empty(fp4_shape, dtype=torch.uint8), + tensor.new_empty((scale_shape,), dtype=torch.uint8), + torch.empty_like(residual), + ) + + +@torch.library.custom_op( + "dist::trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4", + mutates_args=(), + device_types="cuda", +) +def trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4( + tensor: torch.Tensor, + residual: torch.Tensor, + norm_weight: torch.Tensor, + scale: torch.Tensor, + eps: float, + strategy: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused allreduce + residual + RMSNorm with both BF16 and NVFP4 norm outputs.""" + all_reduce_params = AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4, + bias=None, + residual=residual, + norm_weight=norm_weight, + scale=scale, + eps=eps, + ) + norm_out, quant_fp4, scale_factor, residual_out = trtllm_allreduce( + tensor, ReduceOp.SUM, strategy=strategy, all_reduce_params=all_reduce_params + ) + return norm_out, quant_fp4, scale_factor, residual_out + + +@trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4.register_fake +def trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4_fake( + tensor: torch.Tensor, + residual: torch.Tensor, + norm_weight: torch.Tensor, + scale: torch.Tensor, + eps: float, + strategy: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del norm_weight, scale, eps, strategy + fp4_shape, scale_shape = fp4_utils.get_fp4_shape(tensor.shape, 16) + return ( + torch.empty_like(tensor), + tensor.new_empty(fp4_shape, dtype=torch.uint8), + tensor.new_empty((scale_shape,), dtype=torch.uint8), + torch.empty_like(residual), + ) + + def is_trtllm_op_available(): """Check if TRT-LLM ops are available and running with MPI.""" return is_ompi() diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py index 6c750c46eee0..8691674b16c8 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py @@ -259,6 +259,175 @@ def _triton_rmsnorm_gated_meta( return x.new_empty(x.shape, dtype=x.dtype) +@torch.library.custom_op("auto_deploy::trtllm_fused_gated_rmsnorm_quant_nvfp4", mutates_args=()) +def trtllm_fused_gated_rmsnorm_quant_nvfp4( + x: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + eps: float, + group_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse gated RMSNorm and NVFP4 quantization using the TRT-LLM Torch kernel.""" + if weight.dtype in (torch.float16, torch.bfloat16): + kernel_dtype = weight.dtype + elif gate.dtype in (torch.float16, torch.bfloat16): + kernel_dtype = gate.dtype + else: + kernel_dtype = x.dtype + + if x.dtype != kernel_dtype: + x = x.to(kernel_dtype) + if gate.dtype != kernel_dtype: + gate = gate.to(kernel_dtype) + if weight.dtype != kernel_dtype: + weight = weight.to(kernel_dtype) + + x_shape = x.shape + hidden_size = x_shape[-1] + x_2d = x.reshape(-1, hidden_size) + if x_2d.stride(-1) != 1: + x_2d = x_2d.contiguous() + + gate_2d = gate.reshape(-1, hidden_size) + if gate_2d.stride(-1) != 1: + gate_2d = gate_2d.contiguous() + + fp4_i32, scale_factors = torch.ops.trtllm.fused_gated_rmsnorm_quant( + x_2d, gate_2d, weight.contiguous(), group_size, eps, scale.contiguous() + ) + fp4_u8 = fp4_i32.view(torch.uint8) + return fp4_u8.reshape(*x_shape[:-1], hidden_size // 2), scale_factors + + +@trtllm_fused_gated_rmsnorm_quant_nvfp4.register_fake +def _trtllm_fused_gated_rmsnorm_quant_nvfp4_fake( + x: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + eps: float, + group_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + del gate, weight, scale, eps, group_size + output_shape = (*x.shape[:-1], x.shape[-1] // 2) + m = int(x.numel() // x.shape[-1]) + sf_vec_size = 16 + sf_size = ((m + 127) // 128) * 128 * ((x.shape[-1] // sf_vec_size + 3) // 4) * 4 + return x.new_empty(output_shape, dtype=torch.uint8), x.new_empty((sf_size,), dtype=torch.uint8) + + +def _run_trtllm_fused_add_rmsnorm_quant_nvfp4( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + eps: float, + output_hp_norm: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + x_shape = x.shape + hidden_size = x_shape[-1] + x_2d = x.reshape(-1, hidden_size) + if x_2d.stride(-1) != 1: + x_2d = x_2d.contiguous() + + residual_2d = residual.reshape(-1, hidden_size) + if residual_2d.dtype != x_2d.dtype: + residual_2d = residual_2d.to(x_2d.dtype) + if residual_2d.stride(-1) != 1: + residual_2d = residual_2d.contiguous() + + if weight.dtype != x_2d.dtype: + weight = weight.to(x_2d.dtype) + + fp4_i32, residual_out, scale_factors, norm_out = torch.ops.trtllm.fused_add_rms_norm_quant( + x_2d, + residual_2d, + weight.contiguous(), + scale.contiguous(), + True, + eps, + output_hp_norm, + ) + fp4_u8 = fp4_i32.view(torch.uint8).reshape(*x_shape[:-1], hidden_size // 2) + residual_out = residual_out.reshape(x_shape) + if norm_out is not None: + norm_out = norm_out.reshape(x_shape) + return fp4_u8, residual_out, scale_factors, norm_out + + +@torch.library.custom_op("auto_deploy::trtllm_fused_add_rmsnorm_quant_nvfp4", mutates_args=()) +def trtllm_fused_add_rmsnorm_quant_nvfp4( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fuse residual add, RMSNorm, and NVFP4 quantization using a TRT-LLM kernel.""" + fp4_out, residual_out, scale_factors, _ = _run_trtllm_fused_add_rmsnorm_quant_nvfp4( + x, residual, weight, scale, eps, False + ) + return fp4_out, residual_out, scale_factors + + +@trtllm_fused_add_rmsnorm_quant_nvfp4.register_fake +def _trtllm_fused_add_rmsnorm_quant_nvfp4_fake( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del residual, weight, scale, eps + output_shape = (*x.shape[:-1], x.shape[-1] // 2) + m = int(x.numel() // x.shape[-1]) + sf_vec_size = 16 + sf_size = ((m + 127) // 128) * 128 * ((x.shape[-1] // sf_vec_size + 3) // 4) * 4 + return ( + x.new_empty(output_shape, dtype=torch.uint8), + torch.empty_like(x), + x.new_empty((sf_size,), dtype=torch.uint8), + ) + + +@torch.library.custom_op("auto_deploy::trtllm_fused_add_rmsnorm_out_quant_nvfp4", mutates_args=()) +def trtllm_fused_add_rmsnorm_out_quant_nvfp4( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fuse residual add, RMSNorm, and NVFP4 quantization while keeping BF16 norm output.""" + fp4_out, residual_out, scale_factors, norm_out = _run_trtllm_fused_add_rmsnorm_quant_nvfp4( + x, residual, weight, scale, eps, True + ) + assert norm_out is not None + return norm_out, fp4_out, residual_out, scale_factors + + +@trtllm_fused_add_rmsnorm_out_quant_nvfp4.register_fake +def _trtllm_fused_add_rmsnorm_out_quant_nvfp4_fake( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del residual, weight, scale, eps + output_shape = (*x.shape[:-1], x.shape[-1] // 2) + m = int(x.numel() // x.shape[-1]) + sf_vec_size = 16 + sf_size = ((m + 127) // 128) * 128 * ((x.shape[-1] // sf_vec_size + 3) // 4) * 4 + return ( + torch.empty_like(x), + x.new_empty(output_shape, dtype=torch.uint8), + torch.empty_like(x), + x.new_empty((sf_size,), dtype=torch.uint8), + ) + + # Forked from: # https://github.com/state-spaces/mamba/blob/6b32be06d026e170b3fdaf3ae6282c5a6ff57b06/mamba_ssm/ops/triton/layernorm_gated.py # NOTES: diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py new file mode 100644 index 000000000000..856b45904428 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py @@ -0,0 +1,531 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fuse Torch-backend NVFP4 norm-quant kernels into AutoDeploy graphs.""" + +import operator +from typing import List, Tuple, Type + +import torch +from torch.fx import GraphModule, Node + +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils._graph import eliminate_dead_code +from ...utils.node_utils import extract_op_args, extract_output_tuple, is_any_view_op, is_op +from ..interface import ( + BaseTransform, + SharedConfig, + TransformConfig, + TransformInfo, + TransformRegistry, +) + + +def _is_supported_nvfp4_linear(node: Node) -> bool: + return is_op(node, torch.ops.auto_deploy.torch_quant_nvfp4_linear) + + +def _extract_nvfp4_linear_args(node: Node): + return extract_op_args( + node, "input", "weight_fp4", "bias", "input_scale", "weight_scale", "alpha" + ) + + +def _same_scale_node(lhs: Node, rhs: Node) -> bool: + if lhs is rhs: + return True + return lhs.op == "get_attr" and rhs.op == "get_attr" and lhs.target == rhs.target + + +def _is_view_like(node: Node) -> bool: + return is_any_view_op(node) + + +def _is_dtype_cast(node: Node) -> bool: + return is_op(node, torch.ops.aten.to.dtype) + + +def _is_supported_passthrough(node: Node) -> bool: + return _is_view_like(node) or _is_dtype_cast(node) + + +def _unwrap_post_norm_nodes(node: Node) -> Tuple[Node, list[Node]]: + current = node + post_nodes: list[Node] = [] + while isinstance(current, Node) and _is_supported_passthrough(current): + post_nodes.append(current) + current = current.args[0] + return current, post_nodes + + +def _has_unsupported_post_norm_nodes(post_nodes: list[Node]) -> bool: + return any(not _is_dtype_cast(node) for node in post_nodes) + + +def _collect_terminal_users_through_supported_passthrough( + source_node: Node, + *, + max_traversal_nodes: int = 256, +) -> Tuple[List[Node], bool]: + terminal_users: List[Node] = [] + data_nodes = {source_node} + stack = list(source_node.users) + seen = set() + while stack: + user = stack.pop() + if user in seen: + continue + seen.add(user) + if len(seen) > max_traversal_nodes: + return [], False + if _is_supported_passthrough(user): + if user.args and isinstance(user.args[0], Node) and user.args[0] in data_nodes: + data_nodes.add(user) + stack.extend(list(user.users)) + continue + terminal_users.append(user) + return terminal_users, True + + +def _extract_dtype_from_meta(node: Node) -> torch.dtype | None: + val = node.meta.get("val") + if hasattr(val, "dtype"): + return val.dtype + + tensor_meta = node.meta.get("tensor_meta") + if hasattr(tensor_meta, "dtype"): + return tensor_meta.dtype + + return None + + +def _get_out_dtype(linear_node: Node, source_node: Node) -> torch.dtype: + return ( + _extract_dtype_from_meta(linear_node) + or _extract_dtype_from_meta(source_node) + or torch.bfloat16 + ) + + +def _get_last_dim_from_meta(node: Node) -> int | None: + val = node.meta.get("val") + if hasattr(val, "shape") and len(val.shape) > 0: + return int(val.shape[-1]) + + tensor_meta = node.meta.get("tensor_meta") + if hasattr(tensor_meta, "shape") and len(tensor_meta.shape) > 0: + return int(tensor_meta.shape[-1]) + + return None + + +def _supports_trtllm_fused_add_rmsnorm_quant_nvfp4(node: Node) -> bool: + hidden_size = _get_last_dim_from_meta(node) + if hidden_size is None: + return False + return 2048 <= hidden_size <= 16384 and hidden_size % 16 == 0 + + +def _get_arg_defined_before( + graph, + arg: Node, + insertion_node: Node, + node_order: dict[Node, int], +) -> Node | None: + if node_order.get(arg, -1) < node_order.get(insertion_node, float("inf")): + return arg + if arg.op != "get_attr": + return None + + # Model parameters/buffers are order-independent logically, but FX lint requires + # the get_attr node to appear before every consumer. + with graph.inserting_before(insertion_node): + cloned_arg = graph.get_attr(arg.target) + cloned_arg.meta.update(arg.meta) + return cloned_arg + + +def _collect_grouped_nvfp4_linear_users( + source_node: Node, + seed_user: Node, + seed_scale: Node, + processed_users: set[int], +) -> List[Node]: + terminal_users, traversal_ok = _collect_terminal_users_through_supported_passthrough( + source_node + ) + if not traversal_ok: + return [] + + grouped_users: List[Node] = [] + for user in terminal_users: + if not _is_supported_nvfp4_linear(user) or id(user) in processed_users: + continue + + input_arg, _, _, input_scale, _, _ = _extract_nvfp4_linear_args(user) + if not isinstance(input_arg, Node) or not isinstance(input_scale, Node): + continue + + user_source, post_nodes = _unwrap_post_norm_nodes(input_arg) + if user_source is not source_node or _has_unsupported_post_norm_nodes(post_nodes): + continue + if not _same_scale_node(seed_scale, input_scale): + continue + + grouped_users.append(user) + + if seed_user not in grouped_users: + grouped_users.append(seed_user) + + return grouped_users + + +def _all_terminal_users_are_grouped(source_node: Node, grouped_users: List[Node]) -> bool: + terminal_users, traversal_ok = _collect_terminal_users_through_supported_passthrough( + source_node + ) + if not traversal_ok: + return False + grouped_user_set = set(grouped_users) + return all(user in grouped_user_set for user in terminal_users) + + +def _has_terminal_users_outside_group(source_node: Node, grouped_users: List[Node]) -> bool: + terminal_users, traversal_ok = _collect_terminal_users_through_supported_passthrough( + source_node + ) + if not traversal_ok: + return True + grouped_user_set = set(grouped_users) + return any(user not in grouped_user_set for user in terminal_users) + + +def _is_getitem(node: Node, idx: int) -> bool: + return node.op == "call_function" and node.target == operator.getitem and node.args[1] == idx + + +def _extract_nonquant_allreduce_norm(node: Node): + if not _is_getitem(node, 0): + return None + + source_node = node.args[0] + if not isinstance(source_node, Node) or not is_op( + source_node, torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm + ): + return None + + tensor, residual, norm_weight, eps, strategy = source_node.args + _, residual_out = extract_output_tuple(source_node, count=2) + return source_node, tensor, residual, norm_weight, eps, strategy, residual_out + + +def _extract_add_rmsnorm(node: Node): + if not is_op( + node, + [ + torch.ops.auto_deploy.flashinfer_rms_norm, + torch.ops.auto_deploy.torch_rmsnorm, + torch.ops.auto_deploy.triton_rms_norm, + ], + ): + return None + + norm_input, norm_weight, eps = node.args + pre_norm_cast = None + if isinstance(norm_input, Node) and _is_dtype_cast(norm_input): + pre_norm_cast = norm_input + norm_input = norm_input.args[0] + + if not isinstance(norm_input, Node) or not is_op(norm_input, torch.ops.aten.add.Tensor): + return None + + add_lhs, add_rhs = norm_input.args[:2] + if not isinstance(add_lhs, Node) or not isinstance(add_rhs, Node): + return None + + return norm_input, pre_norm_cast, add_lhs, add_rhs, norm_weight, eps + + +def _extract_gated_rmsnorm(node: Node): + if not is_op( + node, + [ + torch.ops.auto_deploy.torch_rmsnorm_gated, + torch.ops.auto_deploy.triton_rmsnorm_gated, + ], + ): + return None + + x, weight, gate, eps, group_size, norm_before_gate = extract_op_args( + node, "x", "weight", "gate", "eps", "group_size", "norm_before_gate" + ) + if gate is None or norm_before_gate: + return None + return x, weight, gate, eps, group_size + + +def _insert_prequant_linear( + graph, + linear_node: Node, + fp4_input: Node, + scale_factors: Node, + out_dtype: torch.dtype, +) -> Node: + _, weight_fp4, bias, _, weight_scale, alpha = _extract_nvfp4_linear_args(linear_node) + return graph.call_function( + torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear.default, + args=(fp4_input, weight_fp4, scale_factors, weight_scale, alpha), + kwargs={"bias": bias, "out_dtype": out_dtype}, + ) + + +@TransformRegistry.register("fuse_rmsnorm_quant_nvfp4") +class FuseRMSNormQuantNVFP4(BaseTransform): + """Fuse NVFP4 quantization into RMSNorm producers where TRT-LLM kernels exist.""" + + config: TransformConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return TransformConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + graph = gm.graph + cnt = 0 + processed_nvfp4_users: set[int] = set() + original_nodes = list(graph.nodes) + node_order = {n: i for i, n in enumerate(original_nodes)} + + for node in original_nodes: + if not _is_supported_nvfp4_linear(node) or id(node) in processed_nvfp4_users: + continue + + input_arg, _, _, input_scale, _, _ = _extract_nvfp4_linear_args(node) + if not isinstance(input_arg, Node) or not isinstance(input_scale, Node): + continue + + norm_node, post_nodes = _unwrap_post_norm_nodes(input_arg) + if _has_unsupported_post_norm_nodes(post_nodes): + continue + + nvfp4_linear_users = _collect_grouped_nvfp4_linear_users( + norm_node, node, input_scale, processed_nvfp4_users + ) + if not nvfp4_linear_users: + continue + + earliest_user = min(nvfp4_linear_users, key=lambda n: node_order.get(n, float("inf"))) + out_dtype = _get_out_dtype(earliest_user, norm_node) + + allreduce_info = _extract_nonquant_allreduce_norm(norm_node) + if allreduce_info is not None: + ( + allreduce_node, + tensor, + residual, + norm_weight, + eps, + strategy, + residual_out_node, + ) = allreduce_info + + needs_norm_output = _has_terminal_users_outside_group(norm_node, nvfp4_linear_users) + insertion_node = earliest_user + if needs_norm_output: + insertion_node = min( + list(norm_node.users), key=lambda n: node_order.get(n, float("inf")) + ) + + with graph.inserting_before(insertion_node): + if needs_norm_output: + fused_quant = graph.call_function( + torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4.default, + args=(tensor, residual, norm_weight, input_scale, eps, strategy), + ) + new_norm_out = graph.call_function(operator.getitem, args=(fused_quant, 0)) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 2)) + new_residual_out = graph.call_function( + operator.getitem, args=(fused_quant, 3) + ) + new_norm_out.meta.update(norm_node.meta) + else: + fused_quant = graph.call_function( + torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4.default, + args=(tensor, residual, norm_weight, input_scale, eps, strategy), + ) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) + new_residual_out = graph.call_function( + operator.getitem, args=(fused_quant, 2) + ) + if residual_out_node is not None: + new_residual_out.meta.update(residual_out_node.meta) + fp4_node.meta.update(norm_node.meta) + + if needs_norm_output: + norm_node.replace_all_uses_with(new_norm_out) + if residual_out_node is not None: + residual_out_node.replace_all_uses_with(new_residual_out) + + cnt += self._replace_nvfp4_linears( + nvfp4_linear_users, + fp4_node, + scale_node, + out_dtype, + processed_nvfp4_users, + ) + + if residual_out_node is not None and len(residual_out_node.users) == 0: + graph.erase_node(residual_out_node) + if len(norm_node.users) == 0: + graph.erase_node(norm_node) + if len(allreduce_node.users) == 0: + graph.erase_node(allreduce_node) + continue + + add_norm_info = _extract_add_rmsnorm(norm_node) + if add_norm_info is not None: + if not _supports_trtllm_fused_add_rmsnorm_quant_nvfp4(norm_node): + continue + add_node, pre_norm_cast, add_lhs, add_rhs, norm_weight, eps = add_norm_info + needs_norm_output = _has_terminal_users_outside_group(norm_node, nvfp4_linear_users) + insertion_candidates = list(add_node.users) + list(norm_node.users) + if pre_norm_cast is not None: + insertion_candidates.extend(list(pre_norm_cast.users)) + insertion_node = min( + insertion_candidates, + key=lambda n: node_order.get(n, float("inf")), + ) + fused_input_scale = _get_arg_defined_before( + graph, input_scale, insertion_node, node_order + ) + if fused_input_scale is None: + continue + + with graph.inserting_before(insertion_node): + if needs_norm_output: + fused_quant = graph.call_function( + torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_out_quant_nvfp4.default, + args=(add_lhs, add_rhs, norm_weight, fused_input_scale, eps), + ) + new_norm_out = graph.call_function(operator.getitem, args=(fused_quant, 0)) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) + new_residual_out = graph.call_function( + operator.getitem, args=(fused_quant, 2) + ) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 3)) + new_norm_out.meta.update(norm_node.meta) + else: + fused_quant = graph.call_function( + torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_quant_nvfp4.default, + args=(add_lhs, add_rhs, norm_weight, fused_input_scale, eps), + ) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) + new_residual_out = graph.call_function( + operator.getitem, args=(fused_quant, 1) + ) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 2)) + new_residual_out.meta.update(add_node.meta) + fp4_node.meta.update(norm_node.meta) + + if needs_norm_output: + norm_node.replace_all_uses_with(new_norm_out) + + cnt += self._replace_nvfp4_linears( + nvfp4_linear_users, + fp4_node, + scale_node, + out_dtype, + processed_nvfp4_users, + ) + + add_node.replace_all_uses_with(new_residual_out) + + if len(norm_node.users) == 0: + graph.erase_node(norm_node) + if pre_norm_cast is not None and len(pre_norm_cast.users) == 0: + graph.erase_node(pre_norm_cast) + if len(add_node.users) == 0: + graph.erase_node(add_node) + continue + + gated_info = _extract_gated_rmsnorm(norm_node) + if gated_info is None: + continue + if not _all_terminal_users_are_grouped(norm_node, nvfp4_linear_users): + continue + + x, weight, gate, eps, group_size = gated_info + with graph.inserting_before(earliest_user): + fused_quant = graph.call_function( + torch.ops.auto_deploy.trtllm_fused_gated_rmsnorm_quant_nvfp4.default, + args=(x, gate, weight, input_scale, eps, group_size), + ) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) + fp4_node.meta.update(norm_node.meta) + + cnt += self._replace_nvfp4_linears( + nvfp4_linear_users, + fp4_node, + scale_node, + out_dtype, + processed_nvfp4_users, + ) + + if len(norm_node.users) == 0: + graph.erase_node(norm_node) + + if cnt > 0: + eliminate_dead_code(gm) + gm.recompile() + + info = TransformInfo( + skipped=(cnt == 0), + num_matches=cnt, + is_clean=(cnt == 0), + has_valid_shapes=(cnt == 0), + ) + return gm, info + + def _replace_nvfp4_linears( + self, + nvfp4_linear_users: List[Node], + fp4_node: Node, + scale_node: Node, + out_dtype: torch.dtype, + processed_nvfp4_users: set[int], + ) -> int: + cnt = 0 + graph = fp4_node.graph + for nvfp4_user in nvfp4_linear_users: + with graph.inserting_before(nvfp4_user): + gemm_node = _insert_prequant_linear( + graph, nvfp4_user, fp4_node, scale_node, out_dtype + ) + gemm_node.meta.update(nvfp4_user.meta) + nvfp4_user.replace_all_uses_with(gemm_node) + graph.erase_node(nvfp4_user) + processed_nvfp4_users.add(id(nvfp4_user)) + cnt += 1 + return cnt diff --git a/tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py index be55a586ee65..02f32565a6ba 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/multi_stream_utils.py @@ -31,12 +31,14 @@ """ from threading import RLock -from typing import Any, Callable, Dict, List +from typing import Any, Callable, Dict, List, TypeVar import torch from .logger import ad_logger +T = TypeVar("T") + # --------------------------------------------------------------------------- # Runtime enable/disable flag # --------------------------------------------------------------------------- @@ -168,6 +170,19 @@ def wait_event(device: int, stream_name: str) -> None: # --------------------------------------------------------------------------- +def _record_stream_for_tensor_outputs(x: object, stream: torch.cuda.Stream) -> None: + if isinstance(x, torch.Tensor): + x.record_stream(stream) + return + if isinstance(x, (list, tuple)): + for item in x: + _record_stream_for_tensor_outputs(item, stream) + return + if isinstance(x, dict): + for item in x.values(): + _record_stream_for_tensor_outputs(item, stream) + + @torch._dynamo.disable def record_event_passthrough( x: torch.Tensor, @@ -191,10 +206,10 @@ def record_event_passthrough( @torch._dynamo.disable def begin_aux_stream_passthrough( - x: torch.Tensor, + x: T, *, device: int = -1, -) -> torch.Tensor: +) -> T: """Record a CUDA event on the main stream, switch to aux, and wait for it. After this function returns the thread-local current stream is the @@ -226,7 +241,7 @@ def begin_aux_stream_passthrough( # NOTE: skip during CUDA graph capture — passthrough partitions are # reclassified as dynamic and won't be captured anyway. if not torch.cuda.is_current_stream_capturing(): - x.record_stream(aux_stream) + _record_stream_for_tensor_outputs(x, aux_stream) torch.cuda.set_stream(aux_stream) # Make aux wait for the main-stream event before executing any work. aux_stream.wait_event(main_event) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py index 569713af2c4f..272cf19f7b66 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py @@ -111,6 +111,17 @@ def _mock_fused_add_res_norm_fake(shared_out, routed_out, residual, weight, eps) return torch.nn.functional.layer_norm(combined, (combined.shape[-1],), weight=weight, eps=eps) +@torch.library.custom_op("auto_deploy::mock_tuple_fork_moe_test", mutates_args=()) +def mock_tuple_fork(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Mock tuple-producing op used to simulate fused RMSNorm + quant outputs.""" + return x + 1, x + 2 + + +@mock_tuple_fork.register_fake +def _mock_tuple_fork_fake(x): + return torch.empty_like(x), torch.empty_like(x) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -589,6 +600,29 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) +class MockTupleForkNemotronHMoELayer(nn.Module): + """Nemotron-H pattern where shared and routed paths fork from tuple outputs.""" + + def __init__(self, hidden_dim: int, intermediate_dim: int, num_experts: int = 8): + super().__init__() + self.gate = nn.Linear(hidden_dim, num_experts, bias=False) + self.shared_experts = _SimpleMLP(hidden_dim, intermediate_dim) + self.expert_weight = nn.Parameter(torch.randn(hidden_dim, hidden_dim)) + self.layernorm = nn.LayerNorm(hidden_dim) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + shared_input, routed_input = torch.ops.auto_deploy.mock_tuple_fork_moe_test(hidden_states) + logits = self.gate(routed_input) + routing_weights, selected_experts = torch.topk(logits, k=2, dim=-1) + + shared_out = self.shared_experts(shared_input) + moe_out = torch.ops.auto_deploy.mock_fused_moe_moe_test( + routed_input, selected_experts, routing_weights, self.expert_weight + ) + + return self.layernorm(shared_out + moe_out) + + def test_fused_merge_pattern_and_correctness(): """Fused merge node (MLIR-like): pattern + graph + correctness.""" hidden_dim, intermediate_dim = 128, 256 @@ -639,3 +673,25 @@ def test_fused_merge_multi_layer(): assert num == 2, f"Expected 2 replacements, got {num}" _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) + + +def test_tuple_fork_pattern_and_correctness(): + """Tuple fork point: begin_aux must preserve and record all tuple tensors.""" + hidden_dim, intermediate_dim = 128, 256 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockTupleForkNemotronHMoELayer(hidden_dim, intermediate_dim).eval().to("cuda") + example = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example) + + gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS) + + assert num == 1, f"Expected 1 replacement, got {num}" + _assert_stream_nodes_present(gm) + begin_nodes = [ + n + for n in gm.graph.nodes + if n.op == "call_function" and n.target is begin_aux_stream_passthrough + ] + assert begin_nodes[0].args[0].target is torch.ops.auto_deploy.mock_tuple_fork_moe_test.default + _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py index 3be43e8c70cd..be40894725a8 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py @@ -31,6 +31,9 @@ FuseRMSNormQuantFP8, _get_out_dtype_str, ) +from tensorrt_llm._torch.auto_deploy.transform.library.fuse_rmsnorm_quant_nvfp4 import ( + FuseRMSNormQuantNVFP4, +) from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import fp4_global_scale, fp8_scale @@ -748,3 +751,300 @@ def test_get_out_dtype_str_returns_none_when_norm_meta_missing(): norm.meta.pop("val", None) assert _get_out_dtype_str(norm) is None + + +def _make_nvfp4_graph_root(hidden_size=64): + root = nn.Module() + root.register_buffer("norm_weight", torch.empty(hidden_size, dtype=torch.bfloat16)) + root.register_buffer("weight_fp4", torch.empty(32, hidden_size // 2, dtype=torch.uint8)) + root.register_buffer("input_scale", torch.empty(1, dtype=torch.float32)) + root.register_buffer("weight_scale", torch.empty(128, 4, dtype=torch.uint8)) + root.register_buffer("alpha", torch.empty(1, dtype=torch.float32)) + return root + + +def test_fuse_gated_rmsnorm_quant_nvfp4_rewrites_graph(): + root = _make_nvfp4_graph_root() + graph = torch.fx.Graph() + x = graph.placeholder("x") + gate = graph.placeholder("gate") + norm_weight = graph.get_attr("norm_weight") + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + + norm_out = graph.call_function( + torch.ops.auto_deploy.triton_rmsnorm_gated.default, + args=(x, norm_weight, gate, 1e-5, 256, False), + ) + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output(gemm_out) + norm_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.auto_deploy.trtllm_fused_gated_rmsnorm_quant_nvfp4) + for n in gm.graph.nodes + ) + assert any(is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.triton_rmsnorm_gated) for n in gm.graph.nodes) + + +def test_fuse_gated_rmsnorm_quant_nvfp4_accepts_dtype_cast(): + root = _make_nvfp4_graph_root() + graph = torch.fx.Graph() + x = graph.placeholder("x") + gate = graph.placeholder("gate") + norm_weight = graph.get_attr("norm_weight") + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + + norm_out = graph.call_function( + torch.ops.auto_deploy.triton_rmsnorm_gated.default, + args=(x, norm_weight, gate, 1e-5, 256, False), + ) + cast_out = graph.call_function(torch.ops.aten.to.dtype, args=(norm_out, torch.bfloat16)) + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(cast_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output(gemm_out) + norm_out.meta["val"] = torch.empty((2, 64), dtype=torch.float32) + cast_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.auto_deploy.trtllm_fused_gated_rmsnorm_quant_nvfp4) + for n in gm.graph.nodes + ) + assert any(is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.aten.to.dtype) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) + + +def test_fuse_allreduce_rmsnorm_quant_nvfp4_rewrites_graph(): + root = _make_nvfp4_graph_root() + graph = torch.fx.Graph() + x = graph.placeholder("x") + residual = graph.placeholder("residual") + norm_weight = graph.get_attr("norm_weight") + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + + fused_allreduce = graph.call_function( + torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm.default, + args=(x, residual, norm_weight, 1e-5, "AUTO"), + ) + norm_out = graph.call_function(operator.getitem, args=(fused_allreduce, 0)) + residual_out = graph.call_function(operator.getitem, args=(fused_allreduce, 1)) + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output((gemm_out, residual_out)) + norm_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + residual_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4) + for n in gm.graph.nodes + ) + assert any(is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) for n in gm.graph.nodes) + assert not any( + is_op(n, torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm) for n in gm.graph.nodes + ) + assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) + + +def test_fuse_allreduce_rmsnorm_quant_nvfp4_keeps_norm_for_mixed_consumers(): + root = _make_nvfp4_graph_root() + graph = torch.fx.Graph() + x = graph.placeholder("x") + residual = graph.placeholder("residual") + norm_weight = graph.get_attr("norm_weight") + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + + fused_allreduce = graph.call_function( + torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm.default, + args=(x, residual, norm_weight, 1e-5, "AUTO"), + ) + norm_out = graph.call_function(operator.getitem, args=(fused_allreduce, 0)) + residual_out = graph.call_function(operator.getitem, args=(fused_allreduce, 1)) + extra_consumer = graph.call_function(torch.ops.aten.add.Tensor, args=(norm_out, 1.0)) + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output((gemm_out, extra_consumer, residual_out)) + norm_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + residual_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + extra_consumer.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4) + for n in gm.graph.nodes + ) + assert any(is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) for n in gm.graph.nodes) + assert any(is_op(n, torch.ops.aten.add.Tensor) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) + + +def test_fuse_add_rmsnorm_quant_nvfp4_rewrites_graph(): + hidden_size = 2048 + root = _make_nvfp4_graph_root(hidden_size) + graph = torch.fx.Graph() + x = graph.placeholder("x") + residual = graph.placeholder("residual") + norm_weight = graph.get_attr("norm_weight") + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + + add_out = graph.call_function(torch.ops.aten.add.Tensor, args=(x, residual)) + norm_out = graph.call_function( + torch.ops.auto_deploy.flashinfer_rms_norm.default, + args=(add_out, norm_weight, 1e-5), + ) + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output((gemm_out, add_out)) + add_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + norm_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_quant_nvfp4) for n in gm.graph.nodes + ) + assert any(is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.flashinfer_rms_norm) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.aten.add.Tensor) for n in gm.graph.nodes) + + +def test_fuse_add_cast_rmsnorm_quant_nvfp4_rewrites_graph(): + hidden_size = 2048 + root = _make_nvfp4_graph_root(hidden_size) + graph = torch.fx.Graph() + x = graph.placeholder("x") + residual = graph.placeholder("residual") + norm_weight = graph.get_attr("norm_weight") + weight_fp4 = graph.get_attr("weight_fp4") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + + add_out = graph.call_function(torch.ops.aten.add.Tensor, args=(x, residual)) + cast_out = graph.call_function(torch.ops.aten.to.dtype, args=(add_out, torch.bfloat16)) + norm_out = graph.call_function( + torch.ops.auto_deploy.flashinfer_rms_norm.default, + args=(cast_out, norm_weight, 1e-5), + ) + input_scale = graph.get_attr("input_scale") + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output((gemm_out, add_out)) + add_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + cast_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + norm_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + gm.graph.lint() + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_quant_nvfp4) for n in gm.graph.nodes + ) + assert any(is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.flashinfer_rms_norm) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.aten.to.dtype) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.aten.add.Tensor) for n in gm.graph.nodes) + + +def test_fuse_add_rmsnorm_quant_nvfp4_keeps_norm_for_mixed_consumers(): + hidden_size = 2048 + root = _make_nvfp4_graph_root(hidden_size) + graph = torch.fx.Graph() + x = graph.placeholder("x") + residual = graph.placeholder("residual") + norm_weight = graph.get_attr("norm_weight") + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + + add_out = graph.call_function(torch.ops.aten.add.Tensor, args=(x, residual)) + norm_out = graph.call_function( + torch.ops.auto_deploy.flashinfer_rms_norm.default, + args=(add_out, norm_weight, 1e-5), + ) + extra_consumer = graph.call_function(torch.ops.aten.mul.Tensor, args=(norm_out, 2.0)) + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output((gemm_out, extra_consumer, add_out)) + add_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + norm_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + extra_consumer.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_out_quant_nvfp4) + for n in gm.graph.nodes + ) + assert any(is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) for n in gm.graph.nodes) + assert any(is_op(n, torch.ops.aten.mul.Tensor) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.flashinfer_rms_norm) for n in gm.graph.nodes) + assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) From 302b48d6f922a2bb98a155cdf6719d9ab4a1d1e3 Mon Sep 17 00:00:00 2001 From: Tal Cherckez Date: Sun, 24 May 2026 02:38:09 -0700 Subject: [PATCH 2/3] Fix NVFP4 RMSNorm quant fusion Signed-off-by: Tal Cherckez --- .../custom_ops/normalization/rms_norm.py | 24 ++-- .../library/fuse_rmsnorm_quant_nvfp4.py | 43 ++++-- .../custom_ops/test_multi_stream_moe.py | 1 + .../library/test_quant_fusion.py | 124 ++++++++++++++++++ 4 files changed, 169 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py index 8691674b16c8..2d042705423f 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py @@ -21,6 +21,10 @@ import torch.nn.functional as F from einops import rearrange +import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils + +from ..quantization.quant import TRTLLM_NVFP4_SCALING_VECTOR_SIZE + try: from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl except (ModuleNotFoundError, ImportError): @@ -37,6 +41,11 @@ def get_env_enable_pdl() -> bool: from .triton_rms_norm import rms_norm +def _get_nvfp4_fake_shapes(x: torch.Tensor) -> tuple[tuple[int, ...], int]: + output_shape, sf_size = fp4_utils.get_fp4_shape(x.shape, TRTLLM_NVFP4_SCALING_VECTOR_SIZE) + return tuple(output_shape), sf_size + + @torch.library.custom_op("auto_deploy::flashinfer_rms_norm", mutates_args=()) def flashinfer_rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: """Custom operator for FlashInfer RMSNorm implementation. @@ -310,10 +319,7 @@ def _trtllm_fused_gated_rmsnorm_quant_nvfp4_fake( group_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: del gate, weight, scale, eps, group_size - output_shape = (*x.shape[:-1], x.shape[-1] // 2) - m = int(x.numel() // x.shape[-1]) - sf_vec_size = 16 - sf_size = ((m + 127) // 128) * 128 * ((x.shape[-1] // sf_vec_size + 3) // 4) * 4 + output_shape, sf_size = _get_nvfp4_fake_shapes(x) return x.new_empty(output_shape, dtype=torch.uint8), x.new_empty((sf_size,), dtype=torch.uint8) @@ -380,10 +386,7 @@ def _trtllm_fused_add_rmsnorm_quant_nvfp4_fake( eps: float, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: del residual, weight, scale, eps - output_shape = (*x.shape[:-1], x.shape[-1] // 2) - m = int(x.numel() // x.shape[-1]) - sf_vec_size = 16 - sf_size = ((m + 127) // 128) * 128 * ((x.shape[-1] // sf_vec_size + 3) // 4) * 4 + output_shape, sf_size = _get_nvfp4_fake_shapes(x) return ( x.new_empty(output_shape, dtype=torch.uint8), torch.empty_like(x), @@ -416,10 +419,7 @@ def _trtllm_fused_add_rmsnorm_out_quant_nvfp4_fake( eps: float, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: del residual, weight, scale, eps - output_shape = (*x.shape[:-1], x.shape[-1] // 2) - m = int(x.numel() // x.shape[-1]) - sf_vec_size = 16 - sf_size = ((m + 127) // 128) * 128 * ((x.shape[-1] // sf_vec_size + 3) // 4) * 4 + output_shape, sf_size = _get_nvfp4_fake_shapes(x) return ( torch.empty_like(x), x.new_empty(output_shape, dtype=torch.uint8), diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py index 856b45904428..bff680ca1018 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py @@ -33,6 +33,8 @@ TransformRegistry, ) +_VIEW_METHOD_TARGETS = {"view", "reshape", "transpose", "permute", "contiguous"} + def _is_supported_nvfp4_linear(node: Node) -> bool: return is_op(node, torch.ops.auto_deploy.torch_quant_nvfp4_linear) @@ -51,7 +53,12 @@ def _same_scale_node(lhs: Node, rhs: Node) -> bool: def _is_view_like(node: Node) -> bool: - return is_any_view_op(node) + return is_any_view_op(node) or ( + node.op == "call_method" + and node.target in _VIEW_METHOD_TARGETS + and bool(node.args) + and isinstance(node.args[0], Node) + ) def _is_dtype_cast(node: Node) -> bool: @@ -66,6 +73,8 @@ def _unwrap_post_norm_nodes(node: Node) -> Tuple[Node, list[Node]]: current = node post_nodes: list[Node] = [] while isinstance(current, Node) and _is_supported_passthrough(current): + if not current.args or not isinstance(current.args[0], Node): + break post_nodes.append(current) current = current.args[0] return current, post_nodes @@ -113,8 +122,11 @@ def _extract_dtype_from_meta(node: Node) -> torch.dtype | None: def _get_out_dtype(linear_node: Node, source_node: Node) -> torch.dtype: + input_arg, _, _, _, _, _ = _extract_nvfp4_linear_args(linear_node) + input_dtype = _extract_dtype_from_meta(input_arg) if isinstance(input_arg, Node) else None return ( _extract_dtype_from_meta(linear_node) + or input_dtype or _extract_dtype_from_meta(source_node) or torch.bfloat16 ) @@ -334,7 +346,6 @@ def _apply( continue earliest_user = min(nvfp4_linear_users, key=lambda n: node_order.get(n, float("inf"))) - out_dtype = _get_out_dtype(earliest_user, norm_node) allreduce_info = _extract_nonquant_allreduce_norm(norm_node) if allreduce_info is not None: @@ -355,11 +366,17 @@ def _apply( list(norm_node.users), key=lambda n: node_order.get(n, float("inf")) ) + fused_input_scale = _get_arg_defined_before( + graph, input_scale, insertion_node, node_order + ) + if fused_input_scale is None: + continue + with graph.inserting_before(insertion_node): if needs_norm_output: fused_quant = graph.call_function( torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4.default, - args=(tensor, residual, norm_weight, input_scale, eps, strategy), + args=(tensor, residual, norm_weight, fused_input_scale, eps, strategy), ) new_norm_out = graph.call_function(operator.getitem, args=(fused_quant, 0)) fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) @@ -371,7 +388,7 @@ def _apply( else: fused_quant = graph.call_function( torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4.default, - args=(tensor, residual, norm_weight, input_scale, eps, strategy), + args=(tensor, residual, norm_weight, fused_input_scale, eps, strategy), ) fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) scale_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) @@ -391,7 +408,7 @@ def _apply( nvfp4_linear_users, fp4_node, scale_node, - out_dtype, + norm_node, processed_nvfp4_users, ) @@ -419,14 +436,17 @@ def _apply( fused_input_scale = _get_arg_defined_before( graph, input_scale, insertion_node, node_order ) - if fused_input_scale is None: + fused_norm_weight = _get_arg_defined_before( + graph, norm_weight, insertion_node, node_order + ) + if fused_input_scale is None or fused_norm_weight is None: continue with graph.inserting_before(insertion_node): if needs_norm_output: fused_quant = graph.call_function( torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_out_quant_nvfp4.default, - args=(add_lhs, add_rhs, norm_weight, fused_input_scale, eps), + args=(add_lhs, add_rhs, fused_norm_weight, fused_input_scale, eps), ) new_norm_out = graph.call_function(operator.getitem, args=(fused_quant, 0)) fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) @@ -438,7 +458,7 @@ def _apply( else: fused_quant = graph.call_function( torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_quant_nvfp4.default, - args=(add_lhs, add_rhs, norm_weight, fused_input_scale, eps), + args=(add_lhs, add_rhs, fused_norm_weight, fused_input_scale, eps), ) fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) new_residual_out = graph.call_function( @@ -455,7 +475,7 @@ def _apply( nvfp4_linear_users, fp4_node, scale_node, - out_dtype, + norm_node, processed_nvfp4_users, ) @@ -489,7 +509,7 @@ def _apply( nvfp4_linear_users, fp4_node, scale_node, - out_dtype, + norm_node, processed_nvfp4_users, ) @@ -513,12 +533,13 @@ def _replace_nvfp4_linears( nvfp4_linear_users: List[Node], fp4_node: Node, scale_node: Node, - out_dtype: torch.dtype, + source_node: Node, processed_nvfp4_users: set[int], ) -> int: cnt = 0 graph = fp4_node.graph for nvfp4_user in nvfp4_linear_users: + out_dtype = _get_out_dtype(nvfp4_user, source_node) with graph.inserting_before(nvfp4_user): gemm_node = _insert_prequant_linear( graph, nvfp4_user, fp4_node, scale_node, out_dtype diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py index 272cf19f7b66..41c0214ac977 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py @@ -693,5 +693,6 @@ def test_tuple_fork_pattern_and_correctness(): for n in gm.graph.nodes if n.op == "call_function" and n.target is begin_aux_stream_passthrough ] + assert len(begin_nodes) == 1, f"Expected exactly one begin_aux node, got {len(begin_nodes)}" assert begin_nodes[0].args[0].target is torch.ops.auto_deploy.mock_tuple_fork_moe_test.default _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda")) diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py index be40894725a8..81ab47b2d669 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py @@ -839,6 +839,51 @@ def test_fuse_gated_rmsnorm_quant_nvfp4_accepts_dtype_cast(): assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) +def test_fuse_gated_rmsnorm_quant_nvfp4_preserves_mixed_consumer_dtypes(): + root = _make_nvfp4_graph_root() + graph = torch.fx.Graph() + x = graph.placeholder("x") + gate = graph.placeholder("gate") + norm_weight = graph.get_attr("norm_weight") + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + + norm_out = graph.call_function( + torch.ops.auto_deploy.triton_rmsnorm_gated.default, + args=(x, norm_weight, gate, 1e-5, 256, False), + ) + direct_gemm = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + cast_out = graph.call_function(torch.ops.aten.to.dtype, args=(norm_out, torch.bfloat16)) + casted_gemm = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(cast_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output((direct_gemm, casted_gemm)) + norm_out.meta["val"] = torch.empty((2, 64), dtype=torch.float32) + cast_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + direct_gemm.meta["val"] = torch.empty((2, 32), dtype=torch.float32) + casted_gemm.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + out_dtypes = { + n.kwargs["out_dtype"] + for n in gm.graph.nodes + if is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) + } + + assert info.num_matches == 2 + assert out_dtypes == {torch.float32, torch.bfloat16} + assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) + + def test_fuse_allreduce_rmsnorm_quant_nvfp4_rewrites_graph(): root = _make_nvfp4_graph_root() graph = torch.fx.Graph() @@ -923,6 +968,46 @@ def test_fuse_allreduce_rmsnorm_quant_nvfp4_keeps_norm_for_mixed_consumers(): assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) +def test_fuse_allreduce_rmsnorm_quant_nvfp4_clones_late_input_scale(): + root = _make_nvfp4_graph_root() + graph = torch.fx.Graph() + x = graph.placeholder("x") + residual = graph.placeholder("residual") + norm_weight = graph.get_attr("norm_weight") + + fused_allreduce = graph.call_function( + torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm.default, + args=(x, residual, norm_weight, 1e-5, "AUTO"), + ) + norm_out = graph.call_function(operator.getitem, args=(fused_allreduce, 0)) + residual_out = graph.call_function(operator.getitem, args=(fused_allreduce, 1)) + extra_consumer = graph.call_function(torch.ops.aten.add.Tensor, args=(norm_out, 1.0)) + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output((gemm_out, extra_consumer, residual_out)) + norm_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + residual_out.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + extra_consumer.meta["val"] = torch.empty((2, 64), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + gm.graph.lint() + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4) + for n in gm.graph.nodes + ) + + def test_fuse_add_rmsnorm_quant_nvfp4_rewrites_graph(): hidden_size = 2048 root = _make_nvfp4_graph_root(hidden_size) @@ -1007,6 +1092,45 @@ def test_fuse_add_cast_rmsnorm_quant_nvfp4_rewrites_graph(): assert not any(is_op(n, torch.ops.aten.add.Tensor) for n in gm.graph.nodes) +def test_fuse_add_cast_rmsnorm_quant_nvfp4_clones_late_norm_weight(): + hidden_size = 2048 + root = _make_nvfp4_graph_root(hidden_size) + graph = torch.fx.Graph() + x = graph.placeholder("x") + residual = graph.placeholder("residual") + + add_out = graph.call_function(torch.ops.aten.add.Tensor, args=(x, residual)) + cast_out = graph.call_function(torch.ops.aten.to.dtype, args=(add_out, torch.bfloat16)) + norm_weight = graph.get_attr("norm_weight") + norm_out = graph.call_function( + torch.ops.auto_deploy.flashinfer_rms_norm.default, + args=(cast_out, norm_weight, 1e-5), + ) + weight_fp4 = graph.get_attr("weight_fp4") + input_scale = graph.get_attr("input_scale") + weight_scale = graph.get_attr("weight_scale") + alpha = graph.get_attr("alpha") + gemm_out = graph.call_function( + torch.ops.auto_deploy.torch_quant_nvfp4_linear.default, + args=(norm_out, weight_fp4, None, input_scale, weight_scale, alpha), + ) + graph.output((gemm_out, add_out)) + add_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + cast_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + norm_out.meta["val"] = torch.empty((2, hidden_size), dtype=torch.bfloat16) + gemm_out.meta["val"] = torch.empty((2, 32), dtype=torch.bfloat16) + + gm = torch.fx.GraphModule(root, graph) + transform = FuseRMSNormQuantNVFP4(TransformConfig(stage="post_load_fusion")) + gm, info = transform._apply(gm, None, None, None) + + gm.graph.lint() + assert info.num_matches == 1 + assert any( + is_op(n, torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_quant_nvfp4) for n in gm.graph.nodes + ) + + def test_fuse_add_rmsnorm_quant_nvfp4_keeps_norm_for_mixed_consumers(): hidden_size = 2048 root = _make_nvfp4_graph_root(hidden_size) From e28e20630508bff320f9068137f3304623ff8885 Mon Sep 17 00:00:00 2001 From: Tal Cherckez Date: Wed, 27 May 2026 01:06:11 -0700 Subject: [PATCH 3/3] Refine NVFP4 RMSNorm quant fusion review fixes Move view and dtype passthrough handling into shared node utilities, including method-style view ops used around RMSNorm outputs. Set packed NVFP4 output and scale metadata explicitly so the transform no longer needs its own pre-shape-prop pass. Split the fusion dispatcher into allreduce, add, and gated RMSNorm handlers and extend quant-fusion tests for passthrough traversal and metadata. Signed-off-by: Tal Cherckez --- .../_torch/auto_deploy/config/default.yaml | 1 - .../library/fuse_rmsnorm_quant_fp8.py | 30 +- .../library/fuse_rmsnorm_quant_nvfp4.py | 507 ++++++++++-------- .../_torch/auto_deploy/utils/node_utils.py | 85 ++- .../library/test_quant_fusion.py | 58 +- 5 files changed, 419 insertions(+), 262 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 0f7f61cf7967..875f4ad2668d 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -220,7 +220,6 @@ transforms: fuse_rmsnorm_quant_nvfp4: stage: post_load_fusion enabled: true - requires_shape_prop: true fuse_gdn_gating: stage: post_load_fusion fuse_l2norm: diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_fp8.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_fp8.py index 435e80579da0..7f0a32b811dc 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_fp8.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_fp8.py @@ -29,9 +29,9 @@ collect_terminal_users_through_passthrough, extract_op_args, extract_output_tuple, - is_any_view_op, is_op, is_trivial_passthrough_user, + unwrap_input_through_passthrough, ) from ..interface import ( BaseTransform, @@ -91,26 +91,24 @@ def _collect_grouped_fp8_linear_users( return grouped_users -def _is_view_like(node: Node) -> bool: - return is_any_view_op(node) - - def _unwrap_post_norm_nodes(node: Node) -> Tuple[Node, list[Node]]: - current = node - post_nodes: list[Node] = [] - while isinstance(current, Node) and _is_view_like(current): - post_nodes.append(current) - current = current.args[0] - return current, post_nodes + return unwrap_input_through_passthrough(node) def _reapply_post_norm_nodes(graph, current: Node, post_nodes: list[Node]) -> Node: for post_node in reversed(post_nodes): - current = graph.call_function( - post_node.target, - args=(current, *post_node.args[1:]), - kwargs=post_node.kwargs, - ) + if post_node.op == "call_method": + current = graph.call_method( + post_node.target, + args=(current, *post_node.args[1:]), + kwargs=post_node.kwargs, + ) + else: + current = graph.call_function( + post_node.target, + args=(current, *post_node.args[1:]), + kwargs=post_node.kwargs, + ) current.meta.update(post_node.meta) return current diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py index bff680ca1018..2858adf90191 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py @@ -21,10 +21,21 @@ import torch from torch.fx import GraphModule, Node +import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils + +from ...custom_ops.quantization.quant import TRTLLM_NVFP4_SCALING_VECTOR_SIZE from ...models.factory import ModelFactory from ...shim.interface import CachedSequenceInterface from ...utils._graph import eliminate_dead_code -from ...utils.node_utils import extract_op_args, extract_output_tuple, is_any_view_op, is_op +from ...utils.logger import ad_logger +from ...utils.node_utils import ( + collect_terminal_users_through_passthrough, + extract_op_args, + extract_output_tuple, + is_dtype_cast_op, + is_op, + unwrap_input_through_passthrough, +) from ..interface import ( BaseTransform, SharedConfig, @@ -33,8 +44,6 @@ TransformRegistry, ) -_VIEW_METHOD_TARGETS = {"view", "reshape", "transpose", "permute", "contiguous"} - def _is_supported_nvfp4_linear(node: Node) -> bool: return is_op(node, torch.ops.auto_deploy.torch_quant_nvfp4_linear) @@ -52,61 +61,12 @@ def _same_scale_node(lhs: Node, rhs: Node) -> bool: return lhs.op == "get_attr" and rhs.op == "get_attr" and lhs.target == rhs.target -def _is_view_like(node: Node) -> bool: - return is_any_view_op(node) or ( - node.op == "call_method" - and node.target in _VIEW_METHOD_TARGETS - and bool(node.args) - and isinstance(node.args[0], Node) - ) - - -def _is_dtype_cast(node: Node) -> bool: - return is_op(node, torch.ops.aten.to.dtype) - - -def _is_supported_passthrough(node: Node) -> bool: - return _is_view_like(node) or _is_dtype_cast(node) - - def _unwrap_post_norm_nodes(node: Node) -> Tuple[Node, list[Node]]: - current = node - post_nodes: list[Node] = [] - while isinstance(current, Node) and _is_supported_passthrough(current): - if not current.args or not isinstance(current.args[0], Node): - break - post_nodes.append(current) - current = current.args[0] - return current, post_nodes + return unwrap_input_through_passthrough(node, allow_dtype_cast=True) def _has_unsupported_post_norm_nodes(post_nodes: list[Node]) -> bool: - return any(not _is_dtype_cast(node) for node in post_nodes) - - -def _collect_terminal_users_through_supported_passthrough( - source_node: Node, - *, - max_traversal_nodes: int = 256, -) -> Tuple[List[Node], bool]: - terminal_users: List[Node] = [] - data_nodes = {source_node} - stack = list(source_node.users) - seen = set() - while stack: - user = stack.pop() - if user in seen: - continue - seen.add(user) - if len(seen) > max_traversal_nodes: - return [], False - if _is_supported_passthrough(user): - if user.args and isinstance(user.args[0], Node) and user.args[0] in data_nodes: - data_nodes.add(user) - stack.extend(list(user.users)) - continue - terminal_users.append(user) - return terminal_users, True + return any(not is_dtype_cast_op(node) for node in post_nodes) def _extract_dtype_from_meta(node: Node) -> torch.dtype | None: @@ -144,6 +104,45 @@ def _get_last_dim_from_meta(node: Node) -> int | None: return None +def _get_shape_from_meta(node: Node) -> tuple | None: + val = node.meta.get("val") + if hasattr(val, "shape"): + return tuple(val.shape) + + tensor_meta = node.meta.get("tensor_meta") + if hasattr(tensor_meta, "shape"): + return tuple(tensor_meta.shape) + + return None + + +def _new_empty_from_meta(source_node: Node, shape: tuple, dtype: torch.dtype) -> torch.Tensor: + val = source_node.meta.get("val") + if hasattr(val, "new_empty"): + return val.new_empty(shape, dtype=dtype) + return torch.empty(shape, dtype=dtype, device="meta") + + +def _set_tensor_val_meta( + node: Node, + source_node: Node, + shape: tuple, + dtype: torch.dtype, +) -> None: + node.meta["val"] = _new_empty_from_meta(source_node, shape, dtype) + node.meta.pop("tensor_meta", None) + + +def _set_nvfp4_quant_meta(fp4_node: Node, scale_node: Node, source_node: Node) -> None: + source_shape = _get_shape_from_meta(source_node) + if source_shape is None: + return + + fp4_shape, scale_shape = fp4_utils.get_fp4_shape(source_shape, TRTLLM_NVFP4_SCALING_VECTOR_SIZE) + _set_tensor_val_meta(fp4_node, source_node, tuple(fp4_shape), torch.uint8) + _set_tensor_val_meta(scale_node, source_node, (scale_shape,), torch.uint8) + + def _supports_trtllm_fused_add_rmsnorm_quant_nvfp4(node: Node) -> bool: hidden_size = _get_last_dim_from_meta(node) if hidden_size is None: @@ -176,8 +175,8 @@ def _collect_grouped_nvfp4_linear_users( seed_scale: Node, processed_users: set[int], ) -> List[Node]: - terminal_users, traversal_ok = _collect_terminal_users_through_supported_passthrough( - source_node + terminal_users, traversal_ok = collect_terminal_users_through_passthrough( + source_node, allow_dtype_cast=True ) if not traversal_ok: return [] @@ -206,8 +205,8 @@ def _collect_grouped_nvfp4_linear_users( def _all_terminal_users_are_grouped(source_node: Node, grouped_users: List[Node]) -> bool: - terminal_users, traversal_ok = _collect_terminal_users_through_supported_passthrough( - source_node + terminal_users, traversal_ok = collect_terminal_users_through_passthrough( + source_node, allow_dtype_cast=True ) if not traversal_ok: return False @@ -216,8 +215,8 @@ def _all_terminal_users_are_grouped(source_node: Node, grouped_users: List[Node] def _has_terminal_users_outside_group(source_node: Node, grouped_users: List[Node]) -> bool: - terminal_users, traversal_ok = _collect_terminal_users_through_supported_passthrough( - source_node + terminal_users, traversal_ok = collect_terminal_users_through_passthrough( + source_node, allow_dtype_cast=True ) if not traversal_ok: return True @@ -257,7 +256,7 @@ def _extract_add_rmsnorm(node: Node): norm_input, norm_weight, eps = node.args pre_norm_cast = None - if isinstance(norm_input, Node) and _is_dtype_cast(norm_input): + if isinstance(norm_input, Node) and is_dtype_cast_op(norm_input): pre_norm_cast = norm_input norm_input = norm_input.args[0] @@ -347,187 +346,247 @@ def _apply( earliest_user = min(nvfp4_linear_users, key=lambda n: node_order.get(n, float("inf"))) - allreduce_info = _extract_nonquant_allreduce_norm(norm_node) - if allreduce_info is not None: - ( - allreduce_node, - tensor, - residual, - norm_weight, - eps, - strategy, - residual_out_node, - ) = allreduce_info - - needs_norm_output = _has_terminal_users_outside_group(norm_node, nvfp4_linear_users) - insertion_node = earliest_user - if needs_norm_output: - insertion_node = min( - list(norm_node.users), key=lambda n: node_order.get(n, float("inf")) - ) - - fused_input_scale = _get_arg_defined_before( - graph, input_scale, insertion_node, node_order - ) - if fused_input_scale is None: - continue - - with graph.inserting_before(insertion_node): - if needs_norm_output: - fused_quant = graph.call_function( - torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4.default, - args=(tensor, residual, norm_weight, fused_input_scale, eps, strategy), - ) - new_norm_out = graph.call_function(operator.getitem, args=(fused_quant, 0)) - fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) - scale_node = graph.call_function(operator.getitem, args=(fused_quant, 2)) - new_residual_out = graph.call_function( - operator.getitem, args=(fused_quant, 3) - ) - new_norm_out.meta.update(norm_node.meta) - else: - fused_quant = graph.call_function( - torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4.default, - args=(tensor, residual, norm_weight, fused_input_scale, eps, strategy), - ) - fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) - scale_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) - new_residual_out = graph.call_function( - operator.getitem, args=(fused_quant, 2) - ) - if residual_out_node is not None: - new_residual_out.meta.update(residual_out_node.meta) - fp4_node.meta.update(norm_node.meta) - - if needs_norm_output: - norm_node.replace_all_uses_with(new_norm_out) - if residual_out_node is not None: - residual_out_node.replace_all_uses_with(new_residual_out) - - cnt += self._replace_nvfp4_linears( - nvfp4_linear_users, - fp4_node, - scale_node, - norm_node, - processed_nvfp4_users, - ) - - if residual_out_node is not None and len(residual_out_node.users) == 0: - graph.erase_node(residual_out_node) - if len(norm_node.users) == 0: - graph.erase_node(norm_node) - if len(allreduce_node.users) == 0: - graph.erase_node(allreduce_node) - continue - - add_norm_info = _extract_add_rmsnorm(norm_node) - if add_norm_info is not None: - if not _supports_trtllm_fused_add_rmsnorm_quant_nvfp4(norm_node): - continue - add_node, pre_norm_cast, add_lhs, add_rhs, norm_weight, eps = add_norm_info - needs_norm_output = _has_terminal_users_outside_group(norm_node, nvfp4_linear_users) - insertion_candidates = list(add_node.users) + list(norm_node.users) - if pre_norm_cast is not None: - insertion_candidates.extend(list(pre_norm_cast.users)) - insertion_node = min( - insertion_candidates, - key=lambda n: node_order.get(n, float("inf")), - ) - fused_input_scale = _get_arg_defined_before( - graph, input_scale, insertion_node, node_order - ) - fused_norm_weight = _get_arg_defined_before( - graph, norm_weight, insertion_node, node_order - ) - if fused_input_scale is None or fused_norm_weight is None: - continue - - with graph.inserting_before(insertion_node): - if needs_norm_output: - fused_quant = graph.call_function( - torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_out_quant_nvfp4.default, - args=(add_lhs, add_rhs, fused_norm_weight, fused_input_scale, eps), - ) - new_norm_out = graph.call_function(operator.getitem, args=(fused_quant, 0)) - fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) - new_residual_out = graph.call_function( - operator.getitem, args=(fused_quant, 2) - ) - scale_node = graph.call_function(operator.getitem, args=(fused_quant, 3)) - new_norm_out.meta.update(norm_node.meta) - else: - fused_quant = graph.call_function( - torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_quant_nvfp4.default, - args=(add_lhs, add_rhs, fused_norm_weight, fused_input_scale, eps), - ) - fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) - new_residual_out = graph.call_function( - operator.getitem, args=(fused_quant, 1) - ) - scale_node = graph.call_function(operator.getitem, args=(fused_quant, 2)) - new_residual_out.meta.update(add_node.meta) - fp4_node.meta.update(norm_node.meta) - - if needs_norm_output: - norm_node.replace_all_uses_with(new_norm_out) - - cnt += self._replace_nvfp4_linears( - nvfp4_linear_users, - fp4_node, - scale_node, - norm_node, - processed_nvfp4_users, - ) - - add_node.replace_all_uses_with(new_residual_out) - - if len(norm_node.users) == 0: - graph.erase_node(norm_node) - if pre_norm_cast is not None and len(pre_norm_cast.users) == 0: - graph.erase_node(pre_norm_cast) - if len(add_node.users) == 0: - graph.erase_node(add_node) + num_matches = self._try_fuse_allreduce_rmsnorm_quant( + norm_node, + input_scale, + nvfp4_linear_users, + earliest_user, + node_order, + processed_nvfp4_users, + ) + if num_matches is not None: + cnt += num_matches continue - gated_info = _extract_gated_rmsnorm(norm_node) - if gated_info is None: - continue - if not _all_terminal_users_are_grouped(norm_node, nvfp4_linear_users): + num_matches = self._try_fuse_add_rmsnorm_quant( + norm_node, + input_scale, + nvfp4_linear_users, + node_order, + processed_nvfp4_users, + ) + if num_matches is not None: + cnt += num_matches continue - x, weight, gate, eps, group_size = gated_info - with graph.inserting_before(earliest_user): - fused_quant = graph.call_function( - torch.ops.auto_deploy.trtllm_fused_gated_rmsnorm_quant_nvfp4.default, - args=(x, gate, weight, input_scale, eps, group_size), - ) - fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) - scale_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) - fp4_node.meta.update(norm_node.meta) - - cnt += self._replace_nvfp4_linears( - nvfp4_linear_users, - fp4_node, - scale_node, + num_matches = self._try_fuse_gated_rmsnorm_quant( norm_node, + input_scale, + nvfp4_linear_users, + earliest_user, processed_nvfp4_users, ) - - if len(norm_node.users) == 0: - graph.erase_node(norm_node) + if num_matches is None: + continue + cnt += num_matches if cnt > 0: eliminate_dead_code(gm) + gm.graph.lint() gm.recompile() info = TransformInfo( skipped=(cnt == 0), num_matches=cnt, - is_clean=(cnt == 0), - has_valid_shapes=(cnt == 0), + is_clean=True, + has_valid_shapes=True, ) return gm, info + def _try_fuse_allreduce_rmsnorm_quant( + self, + norm_node: Node, + input_scale: Node, + nvfp4_linear_users: List[Node], + earliest_user: Node, + node_order: dict[Node, int], + processed_nvfp4_users: set[int], + ) -> int | None: + allreduce_info = _extract_nonquant_allreduce_norm(norm_node) + if allreduce_info is None: + return None + + ( + allreduce_node, + tensor, + residual, + norm_weight, + eps, + strategy, + residual_out_node, + ) = allreduce_info + graph = norm_node.graph + needs_norm_output = _has_terminal_users_outside_group(norm_node, nvfp4_linear_users) + insertion_node = earliest_user + if needs_norm_output: + insertion_node = min( + list(norm_node.users), key=lambda n: node_order.get(n, float("inf")) + ) + + fused_input_scale = _get_arg_defined_before(graph, input_scale, insertion_node, node_order) + if fused_input_scale is None: + return 0 + + with graph.inserting_before(insertion_node): + if needs_norm_output: + fused_quant = graph.call_function( + torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_out_quant_nvfp4.default, + args=(tensor, residual, norm_weight, fused_input_scale, eps, strategy), + ) + new_norm_out = graph.call_function(operator.getitem, args=(fused_quant, 0)) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 2)) + new_residual_out = graph.call_function(operator.getitem, args=(fused_quant, 3)) + new_norm_out.meta.update(norm_node.meta) + else: + fused_quant = graph.call_function( + torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm_quant_nvfp4.default, + args=(tensor, residual, norm_weight, fused_input_scale, eps, strategy), + ) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) + new_residual_out = graph.call_function(operator.getitem, args=(fused_quant, 2)) + if residual_out_node is not None: + new_residual_out.meta.update(residual_out_node.meta) + _set_nvfp4_quant_meta(fp4_node, scale_node, norm_node) + + if needs_norm_output: + norm_node.replace_all_uses_with(new_norm_out) + if residual_out_node is not None: + residual_out_node.replace_all_uses_with(new_residual_out) + + num_matches = self._replace_nvfp4_linears( + nvfp4_linear_users, + fp4_node, + scale_node, + norm_node, + processed_nvfp4_users, + ) + + if residual_out_node is not None and len(residual_out_node.users) == 0: + graph.erase_node(residual_out_node) + if len(norm_node.users) == 0: + graph.erase_node(norm_node) + if len(allreduce_node.users) == 0: + graph.erase_node(allreduce_node) + return num_matches + + def _try_fuse_add_rmsnorm_quant( + self, + norm_node: Node, + input_scale: Node, + nvfp4_linear_users: List[Node], + node_order: dict[Node, int], + processed_nvfp4_users: set[int], + ) -> int | None: + add_norm_info = _extract_add_rmsnorm(norm_node) + if add_norm_info is None: + return None + + if not _supports_trtllm_fused_add_rmsnorm_quant_nvfp4(norm_node): + hidden_size = _get_last_dim_from_meta(norm_node) + ad_logger.debug( + "fuse_rmsnorm_quant_nvfp4: skipping add+norm at " + f"{norm_node.name}, hidden_size={hidden_size} outside supported " + "[2048, 16384] range or not divisible by 16" + ) + return 0 + + add_node, pre_norm_cast, add_lhs, add_rhs, norm_weight, eps = add_norm_info + graph = norm_node.graph + needs_norm_output = _has_terminal_users_outside_group(norm_node, nvfp4_linear_users) + insertion_candidates = list(add_node.users) + list(norm_node.users) + if pre_norm_cast is not None: + insertion_candidates.extend(list(pre_norm_cast.users)) + insertion_node = min( + insertion_candidates, + key=lambda n: node_order.get(n, float("inf")), + ) + fused_input_scale = _get_arg_defined_before(graph, input_scale, insertion_node, node_order) + fused_norm_weight = _get_arg_defined_before(graph, norm_weight, insertion_node, node_order) + if fused_input_scale is None or fused_norm_weight is None: + return 0 + + with graph.inserting_before(insertion_node): + if needs_norm_output: + fused_quant = graph.call_function( + torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_out_quant_nvfp4.default, + args=(add_lhs, add_rhs, fused_norm_weight, fused_input_scale, eps), + ) + new_norm_out = graph.call_function(operator.getitem, args=(fused_quant, 0)) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) + new_residual_out = graph.call_function(operator.getitem, args=(fused_quant, 2)) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 3)) + new_norm_out.meta.update(norm_node.meta) + else: + fused_quant = graph.call_function( + torch.ops.auto_deploy.trtllm_fused_add_rmsnorm_quant_nvfp4.default, + args=(add_lhs, add_rhs, fused_norm_weight, fused_input_scale, eps), + ) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) + new_residual_out = graph.call_function(operator.getitem, args=(fused_quant, 1)) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 2)) + new_residual_out.meta.update(add_node.meta) + _set_nvfp4_quant_meta(fp4_node, scale_node, norm_node) + + if needs_norm_output: + norm_node.replace_all_uses_with(new_norm_out) + + num_matches = self._replace_nvfp4_linears( + nvfp4_linear_users, + fp4_node, + scale_node, + norm_node, + processed_nvfp4_users, + ) + + add_node.replace_all_uses_with(new_residual_out) + + if len(norm_node.users) == 0: + graph.erase_node(norm_node) + if pre_norm_cast is not None and len(pre_norm_cast.users) == 0: + graph.erase_node(pre_norm_cast) + if len(add_node.users) == 0: + graph.erase_node(add_node) + return num_matches + + def _try_fuse_gated_rmsnorm_quant( + self, + norm_node: Node, + input_scale: Node, + nvfp4_linear_users: List[Node], + earliest_user: Node, + processed_nvfp4_users: set[int], + ) -> int | None: + gated_info = _extract_gated_rmsnorm(norm_node) + if gated_info is None: + return None + if not _all_terminal_users_are_grouped(norm_node, nvfp4_linear_users): + return 0 + + x, weight, gate, eps, group_size = gated_info + graph = norm_node.graph + with graph.inserting_before(earliest_user): + fused_quant = graph.call_function( + torch.ops.auto_deploy.trtllm_fused_gated_rmsnorm_quant_nvfp4.default, + args=(x, gate, weight, input_scale, eps, group_size), + ) + fp4_node = graph.call_function(operator.getitem, args=(fused_quant, 0)) + scale_node = graph.call_function(operator.getitem, args=(fused_quant, 1)) + _set_nvfp4_quant_meta(fp4_node, scale_node, norm_node) + + num_matches = self._replace_nvfp4_linears( + nvfp4_linear_users, + fp4_node, + scale_node, + norm_node, + processed_nvfp4_users, + ) + + if len(norm_node.users) == 0: + graph.erase_node(norm_node) + return num_matches + def _replace_nvfp4_linears( self, nvfp4_linear_users: List[Node], diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py index 90937d355efd..f89634df51af 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py @@ -439,27 +439,48 @@ def is_op(node: Node, ops: Union[OperatorLike, Iterable[OperatorLike]]) -> bool: return is_match -def is_trivial_passthrough_user(node: Node) -> bool: +_VIEW_METHOD_TARGETS = {"view", "reshape", "transpose", "permute", "contiguous"} + + +def is_call_method_view_op(node: Node) -> bool: + """Check whether a node is a method-style tensor view/layout op.""" + return ( + node.op == "call_method" + and node.target in _VIEW_METHOD_TARGETS + and bool(node.args) + and isinstance(node.args[0], Node) + ) + + +def is_view_like_op(node: Node) -> bool: + """Check whether a node is a tensor view/layout passthrough op.""" + if is_call_method_view_op(node): + return True + return ( + is_op(node, torch.ops.aten.view) + or is_op(node, torch.ops.aten.reshape) + or is_op(node, torch.ops.aten.transpose) + or is_op(node, torch.ops.aten.permute) + or is_op(node, torch.ops.aten.contiguous) + or is_op(node, torch.ops.auto_deploy.view) + ) + + +def is_dtype_cast_op(node: Node) -> bool: + """Check whether a node casts only tensor dtype.""" + return is_op(node, torch.ops.aten.to.dtype) + + +def is_trivial_passthrough_user(node: Node, *, allow_dtype_cast: bool = False) -> bool: """Check whether a node is a trivial layout/index passthrough op.""" - if node.op == "call_method": - return node.target in { - "view", - "reshape", - "transpose", - "permute", - "contiguous", - "__getitem__", - } + if allow_dtype_cast and is_dtype_cast_op(node): + return True + if is_view_like_op(node): + return True + if node.op == "call_method" and node.target == "__getitem__": + return True if node.op == "call_function": - if node.target is operator.getitem: - return True - return ( - is_op(node, torch.ops.aten.view) - or is_op(node, torch.ops.aten.reshape) - or is_op(node, torch.ops.aten.transpose) - or is_op(node, torch.ops.aten.permute) - or is_op(node, torch.ops.aten.contiguous) - ) + return node.target is operator.getitem return False @@ -467,6 +488,7 @@ def collect_terminal_users_through_passthrough( source_node: Node, *, max_traversal_nodes: int = 256, + allow_dtype_cast: bool = False, ) -> Tuple[List[Node], bool]: """Collect terminal users while traversing trivial passthrough users. @@ -489,7 +511,7 @@ def collect_terminal_users_through_passthrough( seen.add(user) if len(seen) > max_traversal_nodes: return [], False - if is_trivial_passthrough_user(user): + if is_trivial_passthrough_user(user, allow_dtype_cast=allow_dtype_cast): if user.args and isinstance(user.args[0], Node) and user.args[0] in data_nodes: data_nodes.add(user) stack.extend(list(user.users)) @@ -498,6 +520,29 @@ def collect_terminal_users_through_passthrough( return terminal_users, True +def unwrap_input_through_passthrough( + node: Node, + *, + allow_dtype_cast: bool = False, +) -> Tuple[Node, List[Node]]: + """Walk backward through view-like passthrough ops from a consumer input. + + The returned post_nodes are ordered from consumer input back toward the + source producer. Transforms that can absorb dtype casts may opt in with + allow_dtype_cast=True. + """ + current = node + post_nodes: List[Node] = [] + while isinstance(current, Node) and ( + is_view_like_op(current) or (allow_dtype_cast and is_dtype_cast_op(current)) + ): + if not current.args or not isinstance(current.args[0], Node): + break + post_nodes.append(current) + current = current.args[0] + return current, post_nodes + + def get_shared_input_scale_for_fp8_linears( nodes: Iterable[Node], ) -> Tuple[List[Node], Optional[Node]]: diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py index 81ab47b2d669..a80c3eaaa386 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_quant_fusion.py @@ -35,7 +35,11 @@ FuseRMSNormQuantNVFP4, ) from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer -from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op +from tensorrt_llm._torch.auto_deploy.utils.node_utils import ( + collect_terminal_users_through_passthrough, + is_op, + unwrap_input_through_passthrough, +) from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import fp4_global_scale, fp8_scale @@ -753,6 +757,33 @@ def test_get_out_dtype_str_returns_none_when_norm_meta_missing(): assert _get_out_dtype_str(norm) is None +def test_passthrough_helpers_handle_method_views_and_optional_dtype_cast(): + graph = torch.fx.Graph() + x = graph.placeholder("x") + view = graph.call_method("reshape", args=(x, (2, 4))) + cast = graph.call_function(torch.ops.aten.to.dtype, args=(view, torch.bfloat16)) + user = graph.call_function(torch.ops.aten.neg.default, args=(cast,)) + graph.output(user) + + source, post_nodes = unwrap_input_through_passthrough(view) + assert source is x + assert post_nodes == [view] + + source, post_nodes = unwrap_input_through_passthrough(cast, allow_dtype_cast=True) + assert source is x + assert post_nodes == [cast, view] + + terminal_users, traversal_ok = collect_terminal_users_through_passthrough(x) + assert traversal_ok + assert terminal_users == [cast] + + terminal_users, traversal_ok = collect_terminal_users_through_passthrough( + x, allow_dtype_cast=True + ) + assert traversal_ok + assert terminal_users == [user] + + def _make_nvfp4_graph_root(hidden_size=64): root = nn.Module() root.register_buffer("norm_weight", torch.empty(hidden_size, dtype=torch.bfloat16)) @@ -763,6 +794,21 @@ def _make_nvfp4_graph_root(hidden_size=64): return root +def _get_fused_getitem(gm, fused_op, index): + matches = [ + n + for n in gm.graph.nodes + if n.op == "call_function" + and n.target is operator.getitem + and len(n.args) == 2 + and isinstance(n.args[0], torch.fx.Node) + and is_op(n.args[0], fused_op) + and n.args[1] == index + ] + assert len(matches) == 1 + return matches[0] + + def test_fuse_gated_rmsnorm_quant_nvfp4_rewrites_graph(): root = _make_nvfp4_graph_root() graph = torch.fx.Graph() @@ -798,6 +844,16 @@ def test_fuse_gated_rmsnorm_quant_nvfp4_rewrites_graph(): assert any(is_op(n, torch.ops.auto_deploy.trtllm_nvfp4_prequant_linear) for n in gm.graph.nodes) assert not any(is_op(n, torch.ops.auto_deploy.torch_quant_nvfp4_linear) for n in gm.graph.nodes) assert not any(is_op(n, torch.ops.auto_deploy.triton_rmsnorm_gated) for n in gm.graph.nodes) + fp4_node = _get_fused_getitem( + gm, torch.ops.auto_deploy.trtllm_fused_gated_rmsnorm_quant_nvfp4, 0 + ) + scale_node = _get_fused_getitem( + gm, torch.ops.auto_deploy.trtllm_fused_gated_rmsnorm_quant_nvfp4, 1 + ) + assert tuple(fp4_node.meta["val"].shape) == (2, 32) + assert fp4_node.meta["val"].dtype == torch.uint8 + assert tuple(scale_node.meta["val"].shape) == (512,) + assert scale_node.meta["val"].dtype == torch.uint8 def test_fuse_gated_rmsnorm_quant_nvfp4_accepts_dtype_cast():