fix backend bugs (#2781)

Author: Galaxy1458

src/flag_gems/runtime/backend/__init__.py

@@ -40,7 +40,6 @@ def __init__(self):
 
 # Global singleton instance
 _state = BackendState()
-vendor_module = _state.vendor_module
 
 
 class BackendArchEvent:

src/flag_gems/runtime/backend/_cambricon/ops/cat.py

@@ -85,10 +85,10 @@ def gen_imports(self):
         self.writeline("import triton.language as tl")
         self.newline()
         self.writeline("from flag_gems.runtime import torch_device_fn")
-        self.writeline("from flag_gems.runtime.backend import vendor_module")
+        self.writeline("from flag_gems.runtime.backend import _state")
         self.writeline("from flag_gems.utils import libentry, libtuner")
         self.newline()
-        self.writeline("TOTAL_CORE_NUM = vendor_module.TOTAL_CORE_NUM")
+        self.writeline("TOTAL_CORE_NUM = _state.vendor_module.TOTAL_CORE_NUM")
         self.newline()
         self.newline()
 

src/flag_gems/runtime/backend/_cambricon/ops/flip.py

@@ -81,9 +81,9 @@ def __imports(self):
 from triton import language as tl
 
 from flag_gems.utils import libentry
-from flag_gems.runtime.backend import vendor_module
-TOTAL_CORE_NUM = vendor_module.utils.TOTAL_CORE_NUM
-MAX_NRAM_SIZE = vendor_module.utils.MAX_NRAM_SIZE
+from flag_gems.runtime.backend import _state
+TOTAL_CORE_NUM = _state.vendor_module.utils.TOTAL_CORE_NUM
+MAX_NRAM_SIZE = _state.vendor_module.utils.MAX_NRAM_SIZE
 
 
         """

src/flag_gems/runtime/backend/_cambricon/ops/repeat.py

@@ -65,8 +65,8 @@ def generate_imports(code: IndentedBuffer) -> IndentedBuffer:
     code.writeline("from flag_gems.runtime import torch_device_fn")
     code.writeline("from flag_gems.utils.shape_utils import volume")
     code.writeline("from flag_gems.utils import libentry")
-    code.writeline("from flag_gems.runtime.backend import vendor_module")
-    code.writeline("MAX_GRID_SIZE_X = vendor_module.MAX_GRID_SIZE_X")
+    code.writeline("from flag_gems.runtime.backend import _state")
+    code.writeline("MAX_GRID_SIZE_X = _state.vendor_module.MAX_GRID_SIZE_X")
     code.writeline("from flag_gems.utils.type_utils import type_promotion")
     code.writeline("from flag_gems.utils import triton_lang_extension as tle")
     code.newline()

src/flag_gems/runtime/backend/_cambricon/ops/stack.py

@@ -52,9 +52,9 @@ def __imports(self):
             from triton import language as tl
             from typing import List, Tuple, Union
             from flag_gems.utils import libentry
-            from flag_gems.runtime.backend import vendor_module
-            TOTAL_CORE_NUM = vendor_module.TOTAL_CORE_NUM
-            MAX_NRAM_SIZE = vendor_module.MAX_NRAM_SIZE
+            from flag_gems.runtime.backend import _state
+            TOTAL_CORE_NUM = _state.vendor_module.TOTAL_CORE_NUM
+            MAX_NRAM_SIZE = _state.vendor_module.MAX_NRAM_SIZE
 
             """
         self.tpl(textwrap.dedent(tpl))

src/flag_gems/runtime/backend/_cambricon/ops/tile.py

@@ -65,8 +65,8 @@ def generate_imports(code: IndentedBuffer) -> IndentedBuffer:
     code.writeline("from flag_gems.runtime import torch_device_fn")
     code.writeline("from flag_gems.utils.shape_utils import volume")
     code.writeline("from flag_gems.utils import libentry, libtuner")
-    code.writeline("from flag_gems.runtime.backend import vendor_module")
-    code.writeline("MAX_GRID_SIZE_X = vendor_module.MAX_GRID_SIZE_X")
+    code.writeline("from flag_gems.runtime.backend import _state")
+    code.writeline("MAX_GRID_SIZE_X = _state.vendor_module.MAX_GRID_SIZE_X")
     code.writeline("from flag_gems.utils.type_utils import type_promotion")
     code.writeline("from flag_gems.utils import triton_lang_extension as tle")
     code.newline()

src/flag_gems/runtime/backend/_cambricon/ops/vstack.py

@@ -77,7 +77,7 @@ def __imports(self):
 from triton import language as tl
 from flag_gems.runtime import torch_device_fn
 from flag_gems.utils import libentry, libtuner
-from flag_gems.runtime.backend import vendor_module
+from flag_gems.runtime.backend import _state
 TOTAL_CORE_NUM = vendor_module.TOTAL_CORE_NUM
 MAX_NRAM_SIZE = vendor_module.MAX_NRAM_SIZE
         """

src/flag_gems/runtime/backend/_enflame/gcu300/utils/codegen_config_utils.py

@@ -5,7 +5,7 @@
 import triton
 
 from flag_gems.runtime import device
-from flag_gems.runtime.backend import vendor_module
+from flag_gems.runtime.backend import _state
 from flag_gems.runtime.common import vendors
 
 ENFLAME_GCU300_4SIPS = int(os.getenv("ENFLAME_GCU300_4SIPS", "0"))
@@ -67,12 +67,12 @@ def __post_init__(self):
     vendors.CAMBRICON: (
         CodeGenConfig(
             8192,
-            tuple([vendor_module.TOTAL_CORE_NUM, 1, 1]),
+            tuple([_state.vendor_module.TOTAL_CORE_NUM, 1, 1]),
             32,
             False,
             prefer_1d_tile=int(triton.__version__[0]) < 3,
         )
-        if vendor_module.vendor_info.vendor_name == "cambricon"
+        if _state.vendor_module.vendor_info.vendor_name == "cambricon"
         else None
     ),
     vendors.METAX: CodeGenConfig(

src/flag_gems/runtime/backend/_enflame/gcu400/utils/codegen_config_utils.py

@@ -4,7 +4,7 @@
 import triton
 
 from flag_gems.runtime import device
-from flag_gems.runtime.backend import vendor_module
+from flag_gems.runtime.backend import _state
 from flag_gems.runtime.common import vendors
 
 
@@ -61,12 +61,12 @@ def __post_init__(self):
     vendors.CAMBRICON: (
         CodeGenConfig(
             8192,
-            tuple([vendor_module.TOTAL_CORE_NUM, 1, 1]),
+            tuple([_state.vendor_module.TOTAL_CORE_NUM, 1, 1]),
             32,
             False,
             prefer_1d_tile=int(triton.__version__[0]) < 3,
         )
-        if vendor_module.vendor_info.vendor_name == "cambricon"
+        if _state.vendor_module.vendor_info.vendor_name == "cambricon"
         else None
     ),
     vendors.METAX: CodeGenConfig(

src/flag_gems/runtime/backend/_kunlunxin/utils/codegen_config_utils.py

@@ -4,7 +4,7 @@
 import triton
 
 from flag_gems.runtime import device
-from flag_gems.runtime.backend import vendor_module
+from flag_gems.runtime.backend import _state
 from flag_gems.runtime.common import vendors
 
 
@@ -65,12 +65,12 @@ def __post_init__(self):
     ),
     vendors.CAMBRICON: CodeGenConfig(
         8192,
-        tuple([vendor_module.TOTAL_CORE_NUM, 1, 1]),
+        tuple([_state.vendor_module.TOTAL_CORE_NUM, 1, 1]),
         32,
         False,
         prefer_1d_tile=int(triton.__version__[0]) < 3,
     )
-    if vendor_module.vendor_info.vendor_name == "cambricon"
+    if _state.vendor_module.vendor_info.vendor_name == "cambricon"
     else None,
     vendors.METAX: CodeGenConfig(
         2048,