diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 166ff6b965f4..875f4ad2668d 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -217,6 +217,9 @@ transforms: rmsnorm_backend: flashinfer gated_rmsnorm_backend: triton requires_shape_prop: true + fuse_rmsnorm_quant_nvfp4: + stage: post_load_fusion + enabled: 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..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. @@ -259,6 +268,166 @@ 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, sf_size = _get_nvfp4_fake_shapes(x) + 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, sf_size = _get_nvfp4_fake_shapes(x) + 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, sf_size = _get_nvfp4_fake_shapes(x) + 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_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 new file mode 100644 index 000000000000..2858adf90191 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_rmsnorm_quant_nvfp4.py @@ -0,0 +1,611 @@ +# 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 + +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.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, + 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 _unwrap_post_norm_nodes(node: Node) -> Tuple[Node, list[Node]]: + 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_op(node) for node in post_nodes) + + +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: + 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 + ) + + +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 _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: + 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_passthrough( + source_node, allow_dtype_cast=True + ) + 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_passthrough( + source_node, allow_dtype_cast=True + ) + 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_passthrough( + source_node, allow_dtype_cast=True + ) + 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_op(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"))) + + 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 + + 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 + + num_matches = self._try_fuse_gated_rmsnorm_quant( + norm_node, + input_scale, + nvfp4_linear_users, + earliest_user, + processed_nvfp4_users, + ) + 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=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], + fp4_node: Node, + scale_node: Node, + 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 + ) + 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/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/custom_ops/test_multi_stream_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py index 569713af2c4f..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 @@ -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,26 @@ 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 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 3be43e8c70cd..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 @@ -31,8 +31,15 @@ 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.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 @@ -748,3 +755,476 @@ 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 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)) + 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 _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() + 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) + 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(): + 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_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() + 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_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) + 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_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) + 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)