[Bugfix] Add custom Triton cache manager to resolve MoE MP issue

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -272,6 +272,12 @@ def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
     )
 
 
+def maybe_set_triton_cache_manager(module: str) -> None:
+    cache_manger = os.environ.get("TRITON_CACHE_MANAGER", None)
+    if cache_manger != module:
+        os.environ["TRITON_CACHE_MANAGER"] = module
+
+
 def get_config_file_name(E: int, N: int, dtype: Optional[str]) -> str:
     device_name = torch.cuda.get_device_name().replace(" ", "_")
     dtype_selector = "" if not dtype else f",dtype={dtype}"
@@ -428,6 +434,10 @@ def fused_experts(hidden_states: torch.Tensor,
     CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
     M = min(num_tokens, CHUNK_SIZE)
 
+    # workaround for https://github.com/vllm-project/vllm/issues/6103
+    maybe_set_triton_cache_manager(
+        "vllm.triton_utils.custom_cache_manager:CustomCacheManager")
+
     if override_config:
         config = override_config
     else:

vllm/triton_utils/custom_cache_manager.py

@@ -0,0 +1,30 @@
+import os
+
+from triton.runtime.cache import (FileCacheManager, default_cache_dir,
+                                  default_dump_dir, default_override_dir)
+
+
+class CustomCacheManager(FileCacheManager):
+
+    def __init__(self, key, override=False, dump=False):
+        self.key = key
+        self.lock_path = None
+        if dump:
+            self.cache_dir = default_dump_dir()
+            self.cache_dir = os.path.join(self.cache_dir, self.key)
+            self.lock_path = os.path.join(self.cache_dir, "lock")
+            os.makedirs(self.cache_dir, exist_ok=True)
+        elif override:
+            self.cache_dir = default_override_dir()
+            self.cache_dir = os.path.join(self.cache_dir, self.key)
+        else:
+            # create cache directory if it doesn't exist
+            self.cache_dir = os.getenv("TRITON_CACHE_DIR",
+                                       "").strip() or default_cache_dir()
+            if self.cache_dir:
+                self.cache_dir = f"{self.cache_dir}_{os.getpid()}"
+                self.cache_dir = os.path.join(self.cache_dir, self.key)
+                self.lock_path = os.path.join(self.cache_dir, "lock")
+                os.makedirs(self.cache_dir, exist_ok=True)
+            else:
+                raise RuntimeError("Could not create or locate cache dir")