cudagraph

optimization.py

@@ -13,6 +13,7 @@ from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
 
 from optimization_utils import capture_component_call
 from optimization_utils import aoti_compile
+from optimization_utils import cudagraph
 
 
 P = ParamSpec('P')
@@ -57,7 +58,7 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             dynamic_shapes=dynamic_shapes,
         )
 
-        return aoti_compile(exported, INDUCTOR_CONFIGS)
+        return cudagraph(aoti_compile(exported, INDUCTOR_CONFIGS))
 
     transformer_config = pipeline.transformer.config
     pipeline.transformer = compile_transformer()

optimization_utils.py

@@ -4,16 +4,24 @@ import contextlib
 from contextvars import ContextVar
 from io import BytesIO
 from typing import Any
+from typing import Callable
+from typing import ParamSpec
+from typing import TypeVar
 from typing import cast
 from unittest.mock import patch
 
 import torch
+from torch.utils._pytree import tree_map_only
 from torch._inductor.package.package import package_aoti
 from torch.export.pt2_archive._package import AOTICompiledModel
 from torch.export.pt2_archive._package_weights import TensorProperties
 from torch.export.pt2_archive._package_weights import Weights
 
 
+P = ParamSpec('P')
+T = TypeVar('T')
+
+
 INDUCTOR_CONFIGS_OVERRIDES = {
     'aot_inductor.package_constants_in_so': False,
     'aot_inductor.package_constants_on_disk': True,
@@ -64,6 +72,48 @@ def aoti_compile(
     return ZeroGPUCompiledModel(archive_file, weights)
 
 
+def cudagraph(fn: Callable[P, list[torch.Tensor]]):
+
+    graphs = {}
+
+    def fn_(*args: P.args, **kwargs: P.kwargs):
+
+        key = hash(tuple(
+            tuple(kwarg.shape)
+            for a in sorted(kwargs.keys())
+            if isinstance((kwarg := kwargs[a]), torch.Tensor)
+        ))
+
+        if key in graphs:
+            wrapped, *_ = graphs[key]
+            return wrapped(*args, **kwargs)
+
+        graph = torch.cuda.CUDAGraph()
+        in_args, in_kwargs = tree_map_only(torch.Tensor, lambda t: t.clone(), (args, kwargs))
+        in_args, in_kwargs = _cast_as((args, kwargs), (in_args, in_kwargs))
+
+        fn(*in_args, **in_kwargs)
+        with torch.cuda.graph(graph):
+            out_tensors = fn(*in_args, **in_kwargs)
+
+        def wrapped(*args: P.args, **kwargs: P.kwargs):
+            for a, b in zip(in_args, args):
+                if isinstance(a, torch.Tensor):
+                    assert isinstance(b, torch.Tensor)
+                    a.copy_(b)
+            for key in kwargs:
+                if isinstance((kwarg := kwargs[key]), torch.Tensor):
+                    assert isinstance((in_kwarg := in_kwargs[key]), torch.Tensor)
+                    in_kwarg.copy_(kwarg)
+            graph.replay()
+            return [tensor.clone() for tensor in out_tensors]
+
+        graphs[key] = (wrapped, graph, in_args, in_kwargs, out_tensors)
+        return wrapped(*args, **kwargs)
+
+    return fn_
+
+
 @contextlib.contextmanager
 def capture_component_call(
     pipeline: Any,
@@ -94,3 +144,7 @@ def capture_component_call(
         except CapturedCallException as e:
             captured_call.args = e.args
             captured_call.kwargs = e.kwargs
+
+
+def _cast_as(type_from: T, value: Any) -> T:
+    return value