Sync v2.0 version of code to github repo

fastdeploy/model_executor/graph_optimization/graph_optimization_backend.py

@@ -16,7 +16,9 @@
 
 from typing import Callable, Optional
 
-from fastdeploy.config import LLMConfig
+from paddle.jit.dy2static.utils import Backend
+
+from fastdeploy.config import FDConfig
 from fastdeploy.model_executor.graph_optimization.cudagraph_piecewise_backend import \
     CudaGraphPiecewiseBackend
 
@@ -24,38 +26,39 @@ from fastdeploy.model_executor.graph_optimization.cudagraph_piecewise_backend im
 class GraphOptBackend:
     """ """
 
-    llm_config: LLMConfig
+    fd_config: FDConfig
     cudagraph_piecewise_backend: Optional[CudaGraphPiecewiseBackend] = None
 
-    def __init__(self, runnable: Callable, llm_config: LLMConfig):
+    def __init__(self, runnable: Callable, fd_config: FDConfig):
         self.runnable = runnable
-        self.llm_config = llm_config
+        self.fd_config = fd_config
 
-    def __call__(self, **kwargs):
-        # 1. TODO(gongshaotian): Static graph
-        if self.llm_config.graph_opt_config.graph_opt_level > 0:
+        self.max_captre_batch = fd_config.graph_opt_config.cudagraph_capture_sizes[
+            0]
+        if self.fd_config.graph_opt_config.graph_opt_level > 0:
             # 1. Prepare cuda grpah input buffers (contain output of subgraphs)
 
             # 2. Convert dynamic grpah to static graph
-            if self.llm_config.graph_opt_config.graph_opt_level > 1:
-                # with cinn
-                pass
-            else:
-                # not use cinn
-                pass
+            from paddle.jit import sot
+            backend = (Backend.CINN
+                       if self.fd_config.graph_opt_config.graph_opt_level > 1
+                       else Backend.PHI)
+            self.runnable = sot.symbolic_translate(self.runnable,
+                                                   training=False,
+                                                   backend=backend)
 
-            # 3. Split the static graph and get a list of callable obj
+    def __call__(self, **kwargs):
+        if not self.fd_config.graph_opt_config.use_cudagraph:
+            return self.runnable(**kwargs)
+        if self.cudagraph_piecewise_backend is None:
+            self.cudagraph_piecewise_backend = CudaGraphPiecewiseBackend(
+                fd_config=self.fd_config, runnable=self.runnable)
 
-            # 4. Get piecewise cuda grpah backend list
+        assert kwargs["forward_meta"].ids_remove_padding is not None
+        batch_size = kwargs["forward_meta"].ids_remove_padding.shape[0]
 
-            return self.runnable  # Fake return value
-
-        # 2. Dynamic graph
+        if ((not kwargs["forward_meta"].step_use_cudagraph)
+                or (batch_size > self.max_captre_batch)):
+            return self.runnable(**kwargs)
         else:
-            print(self.cudagraph_piecewise_backend is None)
-            if self.cudagraph_piecewise_backend is None:
-                self.cudagraph_piecewise_backend = CudaGraphPiecewiseBackend(
-                    llm_config=self.llm_config, runnable=self.runnable)
-            # TODO(gongshaotian): handling kwargs
-            assert kwargs["input_ids"] is not None
             return self.cudagraph_piecewise_backend.__call__(**kwargs)