[None][perf] Add AutoDeploy NVFP4 RMSNorm quant fusion

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:

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()

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:

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