[SOT] Add unsafe cache fast path (#73372)

DrRyanHuang · web-flow · commit ffaab1d0e27d · 2025-06-18T10:04:35.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -26,6 +26,7 @@
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
     ENV_SOT_ENABLE_GUARD_TREE,
     ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
+    ENV_SOT_UNSAFE_CACHE_FASTPATH,
     BreakGraphError,
     CompileCountInfo,
     ConditionalFallbackError,
@@ -73,18 +74,23 @@ class OpcodeExecutorCache(metaclass=Singleton):
     MAX_CACHE_SIZE = 20
     MAX_COMPILE_TIME_PER_CODE = 40
     MAX_COMPILE_TIME_TOTAL = 15 * 60
+    CACHE_HIT_FASTPATH_THRESHOLD = 32
     cache: dict[
         types.CodeType, tuple[GuardedFunctions, paddle.framework.core.GuardTree]
     ]
     translate_count: int
     code_symbolic_inputs: dict[types.CodeType, dict[str, None | dict[int, int]]]
     compile_time_stats: dict[types.CodeType, float]
+    consecutive_cache_hit_count: dict[types.CodeType, int]
+    last_cache_index: dict[types.CodeType, int]
 
     def __init__(self):
         self.cache = {}
         self.translate_count = 0
         self.code_symbolic_inputs = {}
         self.compile_time_stats = {}
+        self.consecutive_cache_hit_count = {}
+        self.last_cache_index = {}
 
     def get_symbolic_inputs(
         self, code: types.CodeType
@@ -141,6 +147,14 @@ def __call__(self, frame: types.FrameType, **kwargs) -> CustomCode:
             **kwargs,
         )
 
+    def is_fastpath_threshold_reached(self, code):
+        # Returns True if the number of consecutive cache hits for the given code
+        # exceeds the UNSAFE_CACHE_FASTPATH threshold.
+        return (
+            self.consecutive_cache_hit_count.get(code, 0)
+            >= self.CACHE_HIT_FASTPATH_THRESHOLD
+        )
+
     @event_register("lookup")
     def lookup(
         self,
@@ -161,13 +175,21 @@ def lookup(
         Returns:
             CustomCode: The custom code object if a matching guard function is found, otherwise None.
         """
+        code: types.CodeType = frame.f_code
 
         if len(guarded_fns) >= self.MAX_CACHE_SIZE:
             log(2, "[Cache] Exceed max cache size, skip it\n")
             return CustomCode(None, False)
 
         enable_strict_guard = ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get()
         enable_guard_tree = ENV_SOT_ENABLE_GUARD_TREE.get()
+        enable_unsafe_cache_fastpath = ENV_SOT_UNSAFE_CACHE_FASTPATH.get()
+
+        if enable_unsafe_cache_fastpath and (
+            self.is_fastpath_threshold_reached(code)
+        ):
+            # NOTE: In inference scenarios, cache misses are generally rare, so we can enable this short path.
+            return guarded_fns[self.last_cache_index[code]][0]
 
         cache_index = None
         if enable_strict_guard or enable_guard_tree:
@@ -232,6 +254,13 @@ def lookup(
                     assert (
                         cache_index is None or index == cache_index
                     ), f"cache_index({cache_index}) is not equal to index({index})"
+
+                    if self.last_cache_index.get(code, None) == index:
+                        self.consecutive_cache_hit_count[code] += 1
+                    else:
+                        self.last_cache_index[code] = index
+                        self.consecutive_cache_hit_count[code] = 1
+
                     return custom_code
                 else:
                     log_do(
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -26,6 +26,7 @@
     ENV_SOT_LOG_LEVEL,
     ENV_SOT_SERIALIZE_INFO,
     ENV_SOT_TRACE_NUMPY,
+    ENV_SOT_UNSAFE_CACHE_FASTPATH,
     ENV_SOT_WITH_CONTROL_FLOW,
     ENV_STRICT_MODE,
     PEP508LikeEnvironmentVariable,
diff --git a/python/paddle/jit/sot/utils/envs.py b/python/paddle/jit/sot/utils/envs.py
@@ -150,6 +150,9 @@ def parse_parameterized_key(input_str: str) -> dict[str, list[str]]:
     "SOT_FORCE_FALLBACK_SIR_IDS", ""
 )
 ENV_SOT_TRACE_NUMPY = BooleanEnvironmentVariable("ENV_SOT_TRACE_NUMPY", True)
+ENV_SOT_UNSAFE_CACHE_FASTPATH = BooleanEnvironmentVariable(
+    "SOT_UNSAFE_CACHE_FASTPATH", False
+)
 
 
 def update_ce_flags():
diff --git a/test/sot/test_guard_fastpath_strategy.py b/test/sot/test_guard_fastpath_strategy.py
@@ -0,0 +1,86 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+from test_case_base import (
+    TestCaseBase,
+)
+
+import paddle
+from paddle.jit.sot.opcode_translator.executor.executor_cache import (
+    OpcodeExecutorCache,
+)
+from paddle.jit.sot.utils import ENV_SOT_UNSAFE_CACHE_FASTPATH
+from paddle.utils.environments import (
+    EnvironmentVariableGuard,
+)
+
+
+def add(x, y):
+    return x + y
+
+
+def subtract(x, y):
+    return x - y
+
+
+class TestUnsafeCacheFastPath(TestCaseBase):
+    def test_guard(self):
+        # NOTE: When UNSAFE CACHE FASTPATH is enabled, if the same cache entry is hit consecutively
+        # for 32 times (this threshold is configurable), the cache is considered stable and
+        # subsequent guard checks will be skipped to improve performance.
+        # The related logic is implemented in the OpcodeExecutorCache class.
+        with EnvironmentVariableGuard(ENV_SOT_UNSAFE_CACHE_FASTPATH, True):
+
+            self.assertTrue(ENV_SOT_UNSAFE_CACHE_FASTPATH.get())
+
+            self.assertFalse(
+                OpcodeExecutorCache().is_fastpath_threshold_reached(
+                    add.__code__
+                )
+            )
+            for _ in range(33):
+                self.assert_results(add, 1, paddle.ones([4]))
+            self.assertTrue(
+                OpcodeExecutorCache().is_fastpath_threshold_reached(
+                    add.__code__
+                )
+            )
+            self.assertFalse(
+                OpcodeExecutorCache().is_fastpath_threshold_reached(
+                    subtract.__code__
+                )
+            )
+
+        with EnvironmentVariableGuard(ENV_SOT_UNSAFE_CACHE_FASTPATH, False):
+            self.assertFalse(ENV_SOT_UNSAFE_CACHE_FASTPATH.get())
+            self.assertFalse(
+                OpcodeExecutorCache().is_fastpath_threshold_reached(
+                    subtract.__code__
+                )
+            )
+            for _ in range(35):
+                self.assert_results(add, 1, 2)
+            self.assertFalse(
+                OpcodeExecutorCache().is_fastpath_threshold_reached(
+                    subtract.__code__
+                )
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -150,6 +150,9 @@ def parse_parameterized_key(input_str: str) -> dict[str, list[str]]:`
`150`	`150`	`"SOT_FORCE_FALLBACK_SIR_IDS", ""`
`151`	`151`	`)`
`152`	`152`	`ENV_SOT_TRACE_NUMPY = BooleanEnvironmentVariable("ENV_SOT_TRACE_NUMPY", True)`
	`153`	`+ENV_SOT_UNSAFE_CACHE_FASTPATH = BooleanEnvironmentVariable(`
	`154`	`+ "SOT_UNSAFE_CACHE_FASTPATH", False`
	`155`	`+)`
`153`	`156`
`154`	`157`
`155`	`158`	`def update_ce_flags():`