Attention Pattern Matcher (closes NVIDIA#4404) (#88)

* WIP for attention matching: repeat_kv, eager_attention_matching

Signed-off-by: Frida Hou <[email protected]>

* works e2e with llama2 and llama3.1, eager and sdpa

Signed-off-by: Frida Hou <[email protected]>

* update for unit test test_attention_matcher

Signed-off-by: Frida Hou <[email protected]>

* minor

Signed-off-by: Frida Hou <[email protected]>

* minor

Signed-off-by: Frida Hou <[email protected]>

* unify into one transformation, update unit tests

Signed-off-by: Frida Hou <[email protected]>

* update hf_test to verify transformed output, update move_to_devide to recompile graph

Signed-off-by: Frida Hou <[email protected]>

* update after rebase

Signed-off-by: Frida Hou <[email protected]>

* minor

Signed-off-by: Frida Hou <[email protected]>

* update docstring

Signed-off-by: Frida Hou <[email protected]>

* minor

Signed-off-by: Frida Hou <[email protected]>

---------

Signed-off-by: Frida Hou <[email protected]>

Author: Fridah-nv

tensorrt_llm/_torch/auto_deploy/transformations/_graph.py

@@ -96,23 +96,24 @@ def named_graphmodules(gm: fx.GraphModule) -> Iterator[Tuple[str, fx.GraphModule
             yield name, m
 
 
-def _move_single_gm_to_device(
-    gm: GraphModule, device: torch.device, recompile_graph: bool = False
-) -> None:
+def _move_single_gm_to_device(gm: GraphModule, device: torch.device) -> None:
     """Move one GraphModule and its nodes to the specified device in-place.
     Partially inspired by https://github.com/pytorch/pytorch/blob/05cb98f91d49df9eadfcb3fc29bbd1b621d88860/torch/export/passes/__init__.py#L11
     """
     # move state dict
     gm.to(device)
+    recompile_graph = False
 
     for node in gm.graph.nodes:
         # move all the nodes kwargs with burnt-in device
         if "device" in node.kwargs:
+            recompile_graph = True
             kwargs = node.kwargs.copy()
             kwargs["device"] = device
             node.kwargs = kwargs
 
         if is_op(node, torch.ops.aten.to.device):
+            recompile_graph = True
             args = list(node.args)
             args[1] = device
             node.args = tuple(args)
@@ -135,7 +136,7 @@ def move_to_device(gm: fx.GraphModule, device: DeviceLikeType) -> fx.GraphModule
 
     for _, subgm in reversed(list(named_graphmodules(gm))):
         # recompile graph to update self generated codes in subgraph
-        _move_single_gm_to_device(subgm, device, subgm is not gm)
+        _move_single_gm_to_device(subgm, device)
 
 
 def _is_impure_node(node: Node) -> bool: