[Core][Bugfix] New way for full cudagraph, add support for FA2 and FlashInfer

vllm/compilation/backends.py

@@ -563,10 +563,6 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:
 
         self._called = True
 
-        if not self.compilation_config.use_cudagraph or \
-            not self.compilation_config.cudagraph_copy_inputs:
-            return self.split_gm
-
         # if we need to copy input buffers for cudagraph
         from torch._guards import detect_fake_mode
         fake_mode = detect_fake_mode()
@@ -585,6 +581,19 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:
                 any(is_symbolic(d) for d in x.size())
         ]
 
+        if self.compilation_config.full_cuda_graph:
+            assert self.compilation_config.use_cudagraph, \
+                "full_cuda_graph mode requires use_cudagraph to be True"
+            fullgraph_wrapper = resolve_obj_by_qualname(
+                current_platform.get_fullgraph_wrapper_cls())
+            self.split_gm = fullgraph_wrapper(self.split_gm, self.vllm_config,
+                                              self.graph_pool,
+                                              self.sym_tensor_indices)
+
+        if not self.compilation_config.use_cudagraph or \
+            not self.compilation_config.cudagraph_copy_inputs:
+            return self.split_gm
+
         # compiler managed cudagraph input buffers
         # we assume the first run with symbolic shapes
         # has the maximum size among all the tensors

vllm/compilation/base_piecewise_backend.py

@@ -70,3 +70,46 @@ def __call__(self, *args) -> Any:
             or a replayed static graph.
         """
         raise NotImplementedError
+
+
+class AbstractFullgraphWrapper(Protocol):
+    """
+    FullgraphWrapper interface that allows platforms to wrap the piecewise graph
+    to be viewed or captured as a full graph.
+    """
+
+    def __init__(self, graph: fx.GraphModule, vllm_config: VllmConfig,
+                 graph_pool: Any, sym_shape_indices: list[int], **kwargs):
+        """
+        Initializes the FullgraphWrapper class with compilation and 
+        execution-related configurations.
+
+        Args:
+            graph (fx.GraphModule): The graph represented in fx.
+            vllm_config (VllmConfig): Global configuration for vLLM.
+            graph_pool (Any): 
+                Graph memory pool handle, e.g., 
+                    `torch.cuda.graph_pool_handle()`.
+            sym_shape_indices (list[int]): 
+                Indices of symbolic shape.
+
+        Keyword Args:
+            kwargs: Additional keyword arguments reserved for future 
+                extensions or custom platforms.
+
+        """
+        raise NotImplementedError
+
+    def __call__(self, *args) -> Any:
+        """
+        Executes the wrapped graph for given input args.
+
+        Args:
+            *args: Variable length input arguments to be passed into the 
+                graph. The symbolic shape is expected to be in position 
+                `sym_shape_indices[0]`.
+
+        Returns:
+            Any: Output of the executed wrapped graph.
+        """
+        raise NotImplementedError