[CINN] Remove the restriction of automatically falling back to SOT after enabling CINN

fastdeploy/config.py

@@ -1442,10 +1442,6 @@ def __init__(
         else:
             self.graph_opt_config.init_with_cudagrpah_size(max_capture_size=self.scheduler_config.max_num_seqs)
 
-        # TODO(wangmingkai02): change graph_opt_level=2 when using static mode with cinn
-        if self.graph_opt_config.graph_opt_level == 2:
-            self.graph_opt_config.graph_opt_level = 1
-
         self.tokenizer = tokenizer
         self.ips = ips
         self.tool_parser = tool_parser

tests/graph_optimization/test_graph_opt_backend.py

@@ -32,6 +32,7 @@
 from fastdeploy.model_executor.graph_optimization.decorator import (
     support_graph_optimization,
 )
+from fastdeploy.model_executor.graph_optimization.utils import sot_warmup_guard
 
 
 @support_graph_optimization
@@ -46,7 +47,7 @@ def __init__(self, fd_config: FDConfig) -> None:
 
     def forward(
         self,
-        ids_remove_padding,
+        ids_remove_padding: paddle.Tensor,
         forward_meta: ForwardMeta,
     ):
         hidden_states = self.embed_tokens(forward_meta.ids_remove_padding)
@@ -58,7 +59,7 @@ def forward(
 
     def forward_dynamic(
         self,
-        ids_remove_padding,
+        ids_remove_padding: paddle.Tensor,
         forward_meta: ForwardMeta,
     ):
         hidden_states = self.embed_tokens(forward_meta.ids_remove_padding)
@@ -164,15 +165,23 @@ def _run_model_test(self, fd_config, test_name, compare_with_baseline=True):
         """
         test_model = Attention(fd_config=fd_config, **self.model_config)
 
+        with sot_warmup_guard(True):
+            _ = test_model(ids_remove_padding=self.input_tensor, forward_meta=self.forward_meta)
+
         # Run model test
         output = test_model(ids_remove_padding=self.input_tensor, forward_meta=self.forward_meta)
 
         # Validate results if comparison is requested
         if compare_with_baseline:
             np.testing.assert_allclose(
-                self.baseline_result, output.numpy(), err_msg=f"Test {test_name} failed: output mismatch"
+                self.baseline_result,
+                output.numpy(),
+                err_msg=f"Test {test_name} failed: output mismatch",
+                atol=1e-6,  # for CINN
             )
 
+        paddle.jit.sot.opcode_translator.executor.executor_cache.OpcodeExecutorCache().clear()
+
     def test_dynamic_graph(self):
         """Test dynamic graph mode"""
         fd_config = self._setup_test_config(graph_opt_level=0, use_cudagraph=False)

tests/graph_optimization/test_static_graph_cuda_graph_split.py

@@ -14,17 +14,17 @@
 # limitations under the License.
 """
 
-import os
-
-os.environ["FLAGS_cuda_graph_blacklist"] = "pd_op.matmul,pd_op.transpose"
-
-
 import unittest
 from unittest.mock import Mock
 
 import paddle
 import paddle.nn as nn
 
+from fastdeploy.model_executor.graph_optimization.utils import sot_warmup_guard
+
+paddle.set_flags({"FLAGS_cuda_graph_blacklist": "pd_op.matmul,pd_op.transpose"})
+
+
 from fastdeploy.config import (
     CacheConfig,
     FDConfig,
@@ -77,10 +77,10 @@ def __init__(self, fd_config: FDConfig, **kwargs):
         super().__init__()
         self.model = Attention(fd_config)
 
-    def forward(self, ids_remove_padding, forward_meta: ForwardMeta):
+    def forward(self, ids_remove_padding: paddle.Tensor, forward_meta: ForwardMeta):
         return self.model(ids_remove_padding=ids_remove_padding, forward_meta=forward_meta)
 
-    def forward_correct(self, ids_remove_padding, forward_meta: ForwardMeta):
+    def forward_correct(self, ids_remove_padding: paddle.Tensor, forward_meta: ForwardMeta):
         return self.model.forward_dynamic(ids_remove_padding=ids_remove_padding, forward_meta=forward_meta)
 
 
@@ -111,7 +111,8 @@ def test(self):
         forward_meta1 = ForwardMeta(input_ids=x, ids_remove_padding=x, step_use_cudagraph=True)
 
         # Trigger Capture
-        _ = test_model1(x, forward_meta=forward_meta1)
+        with sot_warmup_guard(True):
+            _ = test_model1(x, forward_meta=forward_meta1)
 
         # Replay
         _ = test_model1(x, forward_meta=forward_meta1)