[#4403][refactor] Move fusion, kvcache, and compile to modular inference optimizer

tensorrt_llm/_torch/auto_deploy/config/default.yaml

@@ -19,6 +19,11 @@ transforms:
     stage: post_export
   cleanup_input_constraints:
     stage: post_export
+  ############################################################################################
+  # RUN PATTERN MATCHER TRANSFORMATIONS TO STANDARDIZE GRAPH REPRESENTATION
+  ############################################################################################
+  match_moe_pattern:
+    stage: pattern_matcher
   match_repeat_kv:
     stage: pattern_matcher
   match_eager_attention:
@@ -27,12 +32,13 @@ transforms:
     stage: pattern_matcher
   match_attention_layout:
     stage: pattern_matcher
-  match_moe_pattern:
-    stage: pattern_matcher
   match_rope_pattern:
     stage: pattern_matcher
   match_rope_layout:
     stage: pattern_matcher
+  ############################################################################################
+  # RUN TRANSFORMATIONS ON STANDARDIZED GRAPH REPRESENTATION
+  ############################################################################################
   eliminate_redundant_transposes:
     stage: pattern_matcher
   # TODO (lucaslie): let's move this to perf optimization once TP sharding is improved
@@ -57,5 +63,44 @@ transforms:
   sharding_transform_executor:
     stage: sharding
     run_shape_prop: true
+  ############################################################################################
+  # MOVE MODEL AND LOAD WEIGHTS
+  ############################################################################################
   load_weights:
     stage: weight_load
+  ############################################################################################
+  # RUN POST-LOAD FUSION AND OPTIMIZATIONS
+  ############################################################################################
+  # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/4674 this is causing OOMs
+  # fuse_moe:
+  #   stage: post_load_fusion
+  # fuse_gemms:
+  #   stage: post_load_fusion
+  fuse_allreduce_residual_rmsnorm:
+    stage: post_load_fusion
+  fuse_collectives:
+    stage: post_load_fusion
+  # TODO (lucaslie): add backend selection as part of configurable inference optimizers
+  # check if we can fuse rmsnorm
+  fuse_rmsnorm:
+    stage: post_load_fusion
+    backend: flashinfer
+  ############################################################################################
+  # SWITCH TO CACHED+FLATTENED ATTENTION + INITIALIZE CACHES
+  ############################################################################################
+  update_in_out_nodes:
+    stage: cache_init
+  insert_cached_attention:
+    stage: cache_init
+  insert_cached_mla_attention:
+    stage: cache_init
+    attn_backend: MultiHeadLatentAttention
+  initialize_cache:
+    stage: cache_init
+  resize_kv_cache:
+    stage: cache_init
+  ############################################################################################
+  # COMPILE MODEL
+  ############################################################################################
+  compile_model:
+    stage: compile

tensorrt_llm/_torch/auto_deploy/transform/interface.py

@@ -54,6 +54,7 @@ class SharedConfig(BaseModel):
     sharding_config: ShardingConfig = Field(default_factory=ShardingConfig)
     local_rank: int = Field(default=0)
     world_size: int = Field(default=1)
+    attn_backend: str = Field(default="flashinfer", description="The attention backend to use.")
 
 
 class TransformConfig(BaseModel):
@@ -285,7 +286,10 @@ def __call__(
         # update + store new meta data
         history[t_name] = info
         autodeploy_meta[self._history_key] = history
-        self._set_autodeploy_meta(gm, autodeploy_meta)
+
+        if isinstance(gm, GraphModule):
+            # After compilation, gm becomes type CapturedGraph with no meta data.
+            self._set_autodeploy_meta(gm, autodeploy_meta)
 
         # return the graph module
         return gm

tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py

@@ -0,0 +1,204 @@
+import operator
+from typing import Tuple
+
+import torch
+from torch.fx import GraphModule
+
+from ...distributed.trtllm import is_trtllm_op_available
+from ...models.factory import ModelFactory
+from ...shim.interface import CachedSequenceInterface
+from ...utils.node_utils import get_op_overload_packet, get_user_if_pattern_match, is_op
+from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry
+
+# TODO: This is an overly simplified model that works well for vanilla Llama models.
+# However, we eventually want to consider more sophisticated patterns such as
+# * all_reduce(lin1(x) + lin2(x))
+# * version above with fused GEMMs (i.e. with a split node)
+# * all_reduce(pointwise_op(linear(x)))
+# * ...
+
+
+@TransformRegistry.register("fuse_collectives")
+class FuseCollectives(BaseTransform):
+    """
+    Fuses all_reduce ops with preceding (quantized) linear ops into a single fused node for improved performance.
+    """
+
+    def _apply(
+        self,
+        gm: GraphModule,
+        cm: CachedSequenceInterface,
+        factory: ModelFactory,
+        shared_config: SharedConfig,
+    ) -> Tuple[GraphModule, TransformInfo]:
+        num_gemm_collective_fusions = 0
+
+        # lookup for fused ops
+        # TODO: avoid this hardcoded lookup, e.g., by generating fused ops on the fly.
+        lookup = {
+            torch.ops.auto_deploy.torch_linear_simple: torch.ops.auto_deploy.trtllm_dist_fused_linear_all_reduce,
+            torch.ops.aten.linear: torch.ops.auto_deploy.trtllm_dist_fused_linear_all_reduce,
+            torch.ops.auto_deploy.torch_quant_fp8_linear: torch.ops.auto_deploy.torch_quant_fused_fp8_linear_all_reduce,
+        }
+
+        # go through all nodes and find all_reduce nodes
+        for node in gm.graph.nodes:
+            if not is_op(node, torch.ops.auto_deploy.torch_dist_all_reduce):
+                continue
+
+            # check if args are as expected
+            assert len(node.args) == 1 and not len(node.kwargs), (
+                "Unexpected args/kwargs for all_reduce"
+            )
+
+            # retrieve parent and check a few conditions on the parent node
+            parent_node = node.args[0]
+            if not is_op(parent_node, lookup.keys()):
+                continue
+            if len(parent_node.users) > 1:
+                continue
+
+            with gm.graph.inserting_before(node):
+                # insert fused node
+                fused_linear_collective_node = gm.graph.call_function(
+                    lookup[get_op_overload_packet(parent_node.target)],
+                    args=parent_node.args,
+                    kwargs=parent_node.kwargs,
+                )
+            node.replace_all_uses_with(fused_linear_collective_node)
+            gm.graph.erase_node(node)
+            gm.graph.erase_node(parent_node)
+            num_gemm_collective_fusions += 1
+
+        info = TransformInfo(
+            skipped=False,
+            num_matches=num_gemm_collective_fusions,
+            is_clean=False,
+            has_valid_shapes=False,
+        )
+
+        return gm, info
+
+
+@TransformRegistry.register("fuse_allreduce_residual_rmsnorm")
+class FuseAllreduceResidualRMSNorm(BaseTransform):
+    """Essentially, this transformation fuses the following operators into one allreduce trtllm implementation.
+
+    * target pattern:
+        x = all_reduce(x)
+        y = x + residual
+        return rmsnorm(y), y
+    * replacement:
+        fused_allreduce_residual_rmsnorm(x, residual, rmsnorm_weight, rmsnorm_eps)
+
+    """
+
+    def _apply(
+        self,
+        gm: GraphModule,
+        cm: CachedSequenceInterface,
+        factory: ModelFactory,
+        shared_config: SharedConfig,
+    ) -> Tuple[GraphModule, TransformInfo]:
+        if not is_trtllm_op_available():
+            return gm, TransformInfo(
+                skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True
+            )
+
+        num_ar_r_rms_fusions = 0
+
+        def trace_and_fuse(allreduce_node, graph):
+            # Check if all_reduce is followed by addition
+            users = list(allreduce_node.users.keys())
+            if len(users) != 1:
+                return  # Skip if all_reduce has more than one consumer
+            add_node = users[0]
+
+            # Traverse nodes for RMSNorm pattern which is composed of to_copy, pow, mean, add, refer
+            # the Huggingface LlamaRMSNorm implementation as example for more details
+            to_copy_1 = get_user_if_pattern_match(add_node, [torch.ops.aten.add, operator.add], 2)
+            # operand of pow and mul
+            pow_node = get_user_if_pattern_match(
+                to_copy_1, [torch.ops.aten._to_copy, torch.ops.aten.to], 2
+            )
+            mean_node = get_user_if_pattern_match(pow_node, torch.ops.aten.pow, 1)
+            add_eps_node = get_user_if_pattern_match(mean_node, torch.ops.aten.mean, 1)
+            rsqrt_node = get_user_if_pattern_match(
+                add_eps_node, [torch.ops.aten.add, operator.add], 1
+            )
+            mul_node_1 = get_user_if_pattern_match(rsqrt_node, torch.ops.aten.rsqrt, 1)
+            to_copy_2 = get_user_if_pattern_match(mul_node_1, torch.ops.aten.mul, 1)
+            mul_node_2 = get_user_if_pattern_match(
+                to_copy_2, [torch.ops.aten._to_copy, torch.ops.aten.to], 1
+            )
+            # check args of ops: pow(2) and mean(-1)
+            ARGS_MATCH = pow_node is not None and pow_node.args[1] == 2  # exponent
+            ARGS_MATCH &= mean_node is not None and mean_node.args[1] == [-1]  # dimensions
+
+            # Match found: Replace with fused operation
+            if (
+                to_copy_1
+                and pow_node
+                and mean_node
+                and add_eps_node
+                and rsqrt_node
+                and mul_node_1
+                and to_copy_2
+                and mul_node_2
+                and ARGS_MATCH
+            ):
+                # Gather the inputs for the custom operation
+                tensor = allreduce_node.args[0]
+                # Identify the residual argument in the add operation
+                # One of the args in add_node.args is the output of all_reduce
+                # The same idea also applies to norm_weight
+                residual = (
+                    add_node.args[0] if add_node.args[1] is allreduce_node else add_node.args[1]
+                )
+                norm_weight = (
+                    mul_node_2.args[0] if mul_node_2.args[1] is to_copy_2 else mul_node_2.args[1]
+                )
+                eps = add_eps_node.args[1]
+
+                # Insert nodes
+                with graph.inserting_before(allreduce_node):
+                    fused_node = graph.call_function(
+                        torch.ops.dist.fused_allreduce_residual_rmsnorm,
+                        args=(
+                            tensor,
+                            residual,
+                            norm_weight,
+                            eps,
+                        ),
+                    )
+                    # Extract outputs from the tuple returned by `fused_node`
+                    final_output_node = gm.graph.create_node(
+                        "call_function",
+                        target=operator.getitem,
+                        args=(fused_node, 0),
+                    )
+                    add_output_node = gm.graph.create_node(
+                        "call_function",
+                        target=operator.getitem,
+                        args=(fused_node, 1),
+                    )
+
+                    # Replace all uses of rmsnorm_node with final_output_node
+                    mul_node_2.replace_all_uses_with(final_output_node)
+
+                    # Replace all uses of add_node with add_output_node
+                    add_node.replace_all_uses_with(add_output_node)
+
+                nonlocal num_ar_r_rms_fusions
+                num_ar_r_rms_fusions += 1
+
+        # Traverse all nodes
+        for node in gm.graph.nodes:
+            if is_op(node, torch.ops.auto_deploy.torch_dist_all_reduce):
+                trace_and_fuse(allreduce_node=node, graph=gm.graph)
+
+        info = TransformInfo(
+            skipped=False, num_matches=num_ar_r_rms_fusions, is_clean=False, has_valid_shapes=False
+        )
+
+        return gm, info

tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py

@@ -0,0 +1,65 @@
+from typing import List, Literal, Optional, Tuple, Type
+
+from pydantic import Field
+from torch.fx import GraphModule
+
+from ...compile import compile_and_capture
+from ...models.factory import ModelFactory
+from ...shim.interface import CachedSequenceInterface
+from ..interface import (
+    BaseTransform,
+    SharedConfig,
+    TransformConfig,
+    TransformInfo,
+    TransformRegistry,
+)
+
+
+class CompileModelConfig(TransformConfig):
+    """Configuration for the compile model transform."""
+
+    cuda_graph_batch_sizes: Optional[List[int]] = Field(
+        default=None, description="The batch sizes to use for CUDA graphs."
+    )
+    num_batched_inputs: int = Field(
+        default=2, description="The number of batched inputs to use for CUDA graphs."
+    )
+    compile_backend: Literal["torch-simple", "torch-compile", "torch-cudagraph", "torch-opt"] = (
+        Field(description="The backend to use for compiling the model.")
+    )
+
+
+@TransformRegistry.register("compile_model")
+class CompileModel(BaseTransform):
+    """A transform to compile the model."""
+
+    config: CompileModelConfig
+
+    @classmethod
+    def get_config_class(cls) -> Type[TransformConfig]:
+        return CompileModelConfig
+
+    def _apply(
+        self,
+        gm: GraphModule,
+        cm: CachedSequenceInterface,
+        factory: ModelFactory,
+        shared_config: SharedConfig,
+    ) -> Tuple[GraphModule, TransformInfo]:
+        cm.info.set_generate_only_batch()
+        egm_compiled = compile_and_capture(
+            gm,
+            self.config.compile_backend,
+            args=cm.args,
+            dynamic_shapes=cm.dynamic_shapes,
+            compiler_kwargs={
+                "cuda_graph_batch_sizes": self.config.cuda_graph_batch_sizes,
+                "num_batched_inputs": self.config.num_batched_inputs,
+            },
+        )
+        cm.info.reset()
+
+        # store info object about the transform
+        info = TransformInfo(skipped=False, num_matches=1, is_clean=True, has_valid_shapes=True)
+
+        return egm_compiled, info