first commit

vllm_br/config/compilation.py

@@ -0,0 +1,67 @@
+################################################################################
+# Copyright(c)2020-2025 Shanghai Biren Technology Co., Ltd. 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
+#
+#    http://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.
+#
+################################################################################
+
+import enum
+
+from vllm.config.compilation import CUDAGraphMode
+
+
+class CompilationLevel:
+    # constants for the levels of the compilation process
+    NO_COMPILATION = 0
+    DYNAMO_AS_IS = 1
+    DYNAMO_ONCE = 2
+    PIECEWISE = 3
+
+
+class SUPAGraphMode(enum.Enum):
+    """ Constants for the supagraph mode in CompilationConfig.
+    Meanwhile, the subset enum `NONE`, `PIECEWISE` and `FULL` are also
+    treated as concrete runtime mode for supagraph runtime dispatching.
+    """
+    NONE = 0
+    PIECEWISE = 1
+    FULL = 2
+    FULL_DECODE_ONLY = (FULL, NONE)
+    FULL_AND_PIECEWISE = (FULL, PIECEWISE)
+
+    def decode_mode(self) -> 'SUPAGraphMode':
+        return SUPAGraphMode(self.value[0]) if \
+            self.separate_routine() else self
+
+    def mixed_mode(self) -> 'SUPAGraphMode':
+        return SUPAGraphMode(self.value[1]) if \
+            self.separate_routine() else self
+
+    def requires_piecewise_compilation(self) -> bool:
+        return (self.decode_mode() == SUPAGraphMode.PIECEWISE
+                or self.mixed_mode() == SUPAGraphMode.PIECEWISE)
+
+    def max_supagraph_mode(self) -> 'SUPAGraphMode':
+        return SUPAGraphMode(max(
+            self.value)) if self.separate_routine() else self
+
+    def has_full_supagraphs(self) -> bool:
+        return self.max_supagraph_mode() == SUPAGraphMode.FULL
+
+    # ychun, trick for CUDAGraphMode
+    def has_full_cudagraphs(self) -> bool:
+        cuda_graph_mode = CUDAGraphMode(max(
+            self.value)) if self.separate_routine() else self
+        return cuda_graph_mode == CUDAGraphMode.FULL
+
+    def separate_routine(self) -> bool:
+        return isinstance(self.value, tuple)