[diffusion] refactor: centralize hardware platform detection and streamline environment variable management (#15842)

python/sglang/multimodal_gen/runtime/platforms/__init__.py

@@ -139,8 +139,7 @@ def resolve_current_platform_cls_qualname() -> str:
 _current_platform: Platform | None = None
 _init_trace: str = ""
 
-if TYPE_CHECKING:
-    current_platform: Platform
+current_platform: Platform
 
 
 def __getattr__(name: str):
@@ -150,12 +149,6 @@ def __getattr__(name: str):
         #    Platform` so that they can inherit `Platform` class. Therefore,
         #    we cannot resolve `current_platform` during the import of
         #    `sglang.multimodal_gen.runtime.platforms`.
-        # 2. when users use out-of-tree platform plugins, they might run
-        #    `import sgl_diffusion`, some sgl_diffusion internal code might access
-        #    `current_platform` during the import, and we need to make sure
-        #    `current_platform` is only resolved after the plugins are loaded
-        #    (we have tests for this, if any developer violate this, they will
-        #    see the test failures).
         global _current_platform
         if _current_platform is None:
             platform_cls_qualname = resolve_current_platform_cls_qualname()

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -20,7 +20,6 @@ from sglang.multimodal_gen.runtime.platforms.interface import (
     Platform,
     PlatformEnum,
 )
-from sglang.multimodal_gen.runtime.utils.common import is_blackwell, is_sm120
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import import_pynvml
 
@@ -213,7 +212,7 @@ class CudaPlatformBase(Platform):
         elif selected_backend in [
             AttentionBackendEnum.FA,
         ]:
-            if is_blackwell():
+            if cls.is_blackwell():
                 from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
                     set_fa_ver,
                 )
@@ -224,14 +223,14 @@ class CudaPlatformBase(Platform):
             raise ValueError(f"Invalid attention backend for {cls.device_name}")
         else:
 
-            if is_blackwell():
+            if cls.is_blackwell():
                 from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
                     set_fa_ver,
                 )
 
                 set_fa_ver(4)
             target_backend = AttentionBackendEnum.FA
-            if is_sm120():
+            if cls.is_sm120():
                 try:
                     from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import (  # noqa: F401
                         SageAttention3Backend,

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 import enum
 import random
+from functools import lru_cache
 from typing import TYPE_CHECKING, NamedTuple
 
 import numpy as np
@@ -90,28 +91,79 @@ class Platform:
 
     supported_quantization: list[str] = []
 
+    @lru_cache(maxsize=1)
     def is_cuda(self) -> bool:
-        return self._enum == PlatformEnum.CUDA
+        return self.is_cuda_static()
 
+    @lru_cache(maxsize=1)
     def is_rocm(self) -> bool:
-        return self._enum == PlatformEnum.ROCM
+        return self.is_rocm_static()
 
+    @lru_cache(maxsize=1)
     def is_tpu(self) -> bool:
         return self._enum == PlatformEnum.TPU
 
+    @lru_cache(maxsize=1)
     def is_cpu(self) -> bool:
         return self._enum == PlatformEnum.CPU
 
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_blackwell(cls):
+        if not cls.is_cuda_static():
+            return False
+        return torch.cuda.get_device_capability()[0] == 10
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_sm120(cls):
+        if not cls.is_cuda_static():
+            return False
+        return torch.cuda.get_device_capability()[0] == 12
+
+    @classmethod
+    def is_cuda_static(cls) -> bool:
+        return getattr(cls, "_enum", None) == PlatformEnum.CUDA
+
+    @classmethod
+    def is_rocm_static(cls) -> bool:
+        return getattr(cls, "_enum", None) == PlatformEnum.ROCM
+
+    @lru_cache(maxsize=1)
+    def is_hpu(self) -> bool:
+        return hasattr(torch, "hpu") and torch.hpu.is_available()
+
+    @lru_cache(maxsize=1)
+    def is_xpu(self) -> bool:
+        return hasattr(torch, "xpu") and torch.xpu.is_available()
+
+    @lru_cache(maxsize=1)
+    def is_npu(self) -> bool:
+        return hasattr(torch, "npu") and torch.npu.is_available()
+
     def is_out_of_tree(self) -> bool:
         return self._enum == PlatformEnum.OOT
 
+    @lru_cache(maxsize=1)
     def is_cuda_alike(self) -> bool:
         """Stateless version of :func:`torch.cuda.is_available`."""
         return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
 
+    @lru_cache(maxsize=1)
     def is_mps(self) -> bool:
         return self._enum == PlatformEnum.MPS
 
+    @lru_cache(maxsize=1)
+    def is_musa(self):
+        try:
+            return hasattr(torch, "musa") and torch.musa.is_available()
+        except ModuleNotFoundError:
+            return False
+
+    @lru_cache(maxsize=1)
+    def is_hip(self) -> bool:
+        return self.is_rocm()
+
     @classmethod
     def get_attn_backend_cls_str(
         cls,
@@ -168,6 +220,30 @@ class Platform:
         """Get the total memory of a device in bytes."""
         raise NotImplementedError
 
+    @lru_cache(maxsize=1)
+    def get_device(self, local_rank: int) -> torch.device:
+        if self.is_cuda() or self.is_rocm():
+            return torch.device("cuda", local_rank)
+        elif self.is_musa():
+            return torch.device("musa", local_rank)
+        elif self.is_mps():
+            return torch.device("mps")
+        else:
+            return torch.device("cpu")
+
+    @lru_cache(maxsize=1)
+    def get_torch_distributed_backend_str(self) -> str:
+        if self.is_cuda_alike():
+            return "nccl"
+        elif self.is_musa():
+            return "mccl"
+        elif self.is_mps():
+            return "gloo"
+        else:
+            raise NotImplementedError(
+                "No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available"
+            )
+
     @classmethod
     def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
         """