chore: import upstream snapshot with attribution

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

@@ -0,0 +1,258 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/__init__.py
+
+import os
+import traceback
+
+# imported by other files, do not remove
+from sglang.multimodal_gen.runtime.platforms.interface import (  # noqa: F401
+    AttentionBackendEnum,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import resolve_obj_by_qualname
+
+logger = init_logger(__name__)
+
+
+def cuda_platform_plugin() -> str | None:
+    is_cuda = False
+
+    try:
+        from sglang.multimodal_gen.utils import import_pynvml
+
+        pynvml = import_pynvml()  # type: ignore[no-untyped-call]
+        pynvml.nvmlInit()
+        try:
+            # NOTE: Edge case: sgl_diffusion cpu build on a GPU machine.
+            # Third-party pynvml can be imported in cpu build,
+            # we need to check if sgl_diffusion is built with cpu too.
+            # Otherwise, sgl_diffusion will always activate cuda plugin
+            # on a GPU machine, even if in a cpu build.
+            is_cuda = pynvml.nvmlDeviceGetCount() > 0
+        finally:
+            pynvml.nvmlShutdown()
+    except Exception as e:
+        if "nvml" not in e.__class__.__name__.lower():
+            # If the error is not related to NVML, re-raise it.
+            raise e
+
+        # CUDA is supported on Jetson, but NVML may not be.
+        import os
+
+        def cuda_is_jetson() -> bool:
+            return os.path.isfile("/etc/nv_tegra_release") or os.path.exists(
+                "/sys/class/tegra-firmware"
+            )
+
+        if cuda_is_jetson():
+            is_cuda = True
+    if is_cuda:
+        logger.debug("CUDA is available")
+
+    return (
+        "sglang.multimodal_gen.runtime.platforms.cuda.CudaPlatform" if is_cuda else None
+    )
+
+
+def mps_platform_plugin() -> str | None:
+    """Detect if MPS (Metal Performance Shaders) is available on macOS."""
+    is_mps = False
+
+    try:
+        import torch
+
+        if torch.backends.mps.is_available():
+            is_mps = True
+            logger.debug("MPS (Metal Performance Shaders) is available")
+    except Exception as e:
+        logger.debug("MPS detection failed: %s", e)
+
+    return "sglang.multimodal_gen.runtime.platforms.mps.MpsPlatform" if is_mps else None
+
+
+def cpu_platform_plugin() -> str | None:
+    """Detect if CPU platform should be used."""
+    # CPU is always available as a fallback
+    return "sglang.multimodal_gen.runtime.platforms.cpu.CpuPlatform"
+
+
+def rocm_platform_plugin() -> str | None:
+    is_rocm = False
+
+    try:
+        import amdsmi
+
+        amdsmi.amdsmi_init()
+        try:
+            if len(amdsmi.amdsmi_get_processor_handles()) > 0:
+                is_rocm = True
+                logger.debug("ROCm platform is available")
+        finally:
+            amdsmi.amdsmi_shut_down()
+    except Exception as e:
+        logger.debug("ROCm platform is unavailable: %s", e)
+
+    return (
+        "sglang.multimodal_gen.runtime.platforms.rocm.RocmPlatform" if is_rocm else None
+    )
+
+
+def npu_platform_plugin() -> str | None:
+    is_npu = False
+
+    try:
+        import torch
+
+        if torch.npu.is_available():
+            is_npu = True
+            logger.debug("NPU is available")
+    except Exception as e:
+        logger.debug("NPU detection failed: %s", e)
+    return (
+        "sglang.multimodal_gen.runtime.platforms.npu.NPUPlatformBase"
+        if is_npu
+        else None
+    )
+
+
+def musa_platform_plugin() -> str | None:
+    is_musa = False
+
+    try:
+        import pymtml
+
+        pymtml.mtmlLibraryInit()
+        try:
+            is_musa = pymtml.mtmlLibraryCountDevice() > 0
+        finally:
+            pymtml.mtmlLibraryShutDown()
+    except Exception as e:
+        logger.debug("MUSA platform is unavailable: %s", e)
+
+    return (
+        "sglang.multimodal_gen.runtime.platforms.musa.MusaPlatform" if is_musa else None
+    )
+
+
+def xpu_platform_plugin() -> str | None:
+    """Detect if Intel XPU platform is available."""
+    is_xpu = False
+
+    try:
+        import torch
+
+        # Check if Intel Extension for PyTorch is available and XPU devices exist
+        if hasattr(torch, "xpu") and torch.xpu.is_available():
+            device_count = torch.xpu.device_count()
+            if device_count > 0:
+                is_xpu = True
+                logger.info(
+                    "Intel XPU platform is available with %d device(s)", device_count
+                )
+    except Exception as e:
+        logger.info("Intel XPU platform is unavailable: %s", e)
+
+    return "sglang.multimodal_gen.runtime.platforms.xpu.XpuPlatform" if is_xpu else None
+
+
+builtin_platform_plugins = {
+    "cuda": cuda_platform_plugin,
+    "rocm": rocm_platform_plugin,
+    "xpu": xpu_platform_plugin,
+    "mps": mps_platform_plugin,
+    "cpu": cpu_platform_plugin,
+    "npu": npu_platform_plugin,
+    "musa": musa_platform_plugin,
+}
+
+
+def resolve_current_platform_cls_qualname() -> str:
+    forced_platform = os.environ.get("SGLANG_DIFFUSION_PLATFORM_OVERRIDE", "").strip()
+    if forced_platform:
+        forced_map = {
+            "cpu": "sglang.multimodal_gen.runtime.platforms.cpu.CpuPlatform",
+            "cuda": "sglang.multimodal_gen.runtime.platforms.cuda.CudaPlatform",
+            "rocm": "sglang.multimodal_gen.runtime.platforms.rocm.RocmPlatform",
+            "mps": "sglang.multimodal_gen.runtime.platforms.mps.MpsPlatform",
+            "npu": "sglang.multimodal_gen.runtime.platforms.npu.NPUPlatformBase",
+            "musa": "sglang.multimodal_gen.runtime.platforms.musa.MusaPlatform",
+        }
+        qualname = forced_map.get(forced_platform.lower())
+        if qualname is None:
+            raise ValueError(
+                f"Unsupported SGLANG_DIFFUSION_PLATFORM_OVERRIDE={forced_platform!r}"
+            )
+        return qualname
+
+    # TODO(will): if we need to support other platforms, we should consider if
+    # vLLM's plugin architecture is suitable for our needs.
+
+    # Try MPS first on macOS
+    platform_cls_qualname = mps_platform_plugin()
+    if platform_cls_qualname is not None:
+        return platform_cls_qualname
+
+    # Try Intel XPU
+    platform_cls_qualname = xpu_platform_plugin()
+    if platform_cls_qualname is not None:
+        return platform_cls_qualname
+
+    # Fall back to ROCm
+    platform_cls_qualname = rocm_platform_plugin()
+    if platform_cls_qualname is not None:
+        return platform_cls_qualname
+
+    # Fall back to CUDA
+    platform_cls_qualname = cuda_platform_plugin()
+    if platform_cls_qualname is not None:
+        return platform_cls_qualname
+
+    # Fall back to NPU
+    platform_cls_qualname = npu_platform_plugin()
+    if platform_cls_qualname is not None:
+        return platform_cls_qualname
+
+    # Fall back to MUSA
+    platform_cls_qualname = musa_platform_plugin()
+    if platform_cls_qualname is not None:
+        return platform_cls_qualname
+
+    # Fall back to CPU as last resort
+    platform_cls_qualname = cpu_platform_plugin()
+    if platform_cls_qualname is not None:
+        return platform_cls_qualname
+
+    raise RuntimeError("No platform plugin found. Please check your " "installation.")
+
+
+_current_platform: Platform | None = None
+_init_trace: str = ""
+
+current_platform: Platform
+
+
+def __getattr__(name: str):
+    if name == "current_platform":
+        # lazy init current_platform.
+        # 1. out-of-tree platform plugins need `from sglang.multimodal_gen.runtime.platforms import
+        #    Platform` so that they can inherit `Platform` class. Therefore,
+        #    we cannot resolve `current_platform` during the import of
+        #    `sglang.multimodal_gen.runtime.platforms`.
+        global _current_platform
+        if _current_platform is None:
+            platform_cls_qualname = resolve_current_platform_cls_qualname()
+            _current_platform = resolve_obj_by_qualname(platform_cls_qualname)()
+            global _init_trace
+            _init_trace = "".join(traceback.format_stack())
+        return _current_platform
+    elif name in globals():
+        return globals()[name]
+    else:
+        raise AttributeError(f"No attribute named '{name}' exists in {__name__}.")
+
+
+__all__ = ["Platform", "PlatformEnum", "current_platform", "_init_trace"]

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

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.srt.utils import (
+    get_bool_env_var,
+    is_gfx95_supported,
+    is_hip,
+)
+
+USE_AITER = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
+USE_AITER_GFX95 = USE_AITER and is_gfx95_supported()

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

@@ -0,0 +1,123 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cpu.py
+
+import platform
+from functools import lru_cache
+from typing import Any
+
+import psutil
+import torch
+
+from sglang.multimodal_gen.runtime.platforms.interface import (
+    AttentionBackendEnum,
+    CpuArchEnum,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class CpuPlatform(Platform):
+    _enum = PlatformEnum.CPU
+    device_name = "CPU"
+    device_type = "cpu"
+    dispatch_key = "CPU"
+
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        return torch.device("cpu")
+
+    @classmethod
+    def get_torch_distributed_backend_str(cls) -> str:
+        return "gloo"
+
+    @classmethod
+    def get_cpu_architecture(cls) -> CpuArchEnum:
+        """Get the CPU architecture."""
+        machine = platform.machine().lower()
+        if machine in ("x86_64", "amd64", "i386", "i686"):
+            return CpuArchEnum.X86
+        elif machine in ("arm64", "aarch64"):
+            return CpuArchEnum.ARM
+        else:
+            return CpuArchEnum.UNSPECIFIED
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return platform.processor()
+
+    @classmethod
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        return platform.machine()
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+
+        return psutil.virtual_memory().total
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        return True
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        # For CPU, we can't easily get memory usage without additional libraries
+        return 0.0
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int | None = None,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+
+        total_free_memory = psutil.virtual_memory().available
+        # For simplicity, we assume 1 NUMA node for now in this platform abstraction
+        # as get_cpu_ids_by_node is not available in multimodal_gen.runtime.utils
+        n_numa_node = 1
+        free_memory = total_free_memory / n_numa_node
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_memory, dtype=torch.float32)
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_memory = float(tensor.item())
+
+        return free_memory / (1 << 30)
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        if selected_backend not in (None, AttentionBackendEnum.TORCH_SDPA):
+            logger.warning(
+                "%s is not supported on CPU; falling back to Torch SDPA.",
+                selected_backend,
+            )
+
+        logger.info("Using Torch SDPA backend for CPU.")
+        return (
+            "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+        )
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.cpu_communicator.CpuCommunicator"
+
+    @classmethod
+    def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
+        """Whether to enable DIT layerwise offload by default on the current platform."""
+        return False

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

@@ -0,0 +1,691 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cuda.py
+"""Code inside this file can safely assume cuda platform, e.g. importing
+pynvml. However, it should not initialize cuda context.
+"""
+
+import os
+from collections.abc import Callable
+from functools import lru_cache, wraps
+from typing import Any, TypeVar
+
+import psutil
+import torch
+from typing_extensions import ParamSpec
+
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.runtime.platforms.interface import (
+    AttentionBackendEnum,
+    DeviceCapability,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import import_pynvml
+
+logger = init_logger(__name__)
+
+_SDPA_BACKEND_CLS_STR = (
+    "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+)
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+pynvml = import_pynvml()  # type: ignore[no-untyped-call]
+
+# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
+# see https://github.com/huggingface/diffusers/issues/9704 for details
+torch.backends.cuda.enable_cudnn_sdp(False)
+
+
+def device_id_to_physical_device_id(device_id: int) -> int:
+    if "CUDA_VISIBLE_DEVICES" in os.environ:
+        device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
+        if device_ids == [""]:
+            msg = (
+                "CUDA_VISIBLE_DEVICES is set to empty string, which means"
+                " GPU support is disabled. If you are using ray, please unset"
+                " the environment variable `CUDA_VISIBLE_DEVICES` inside the"
+                " worker/actor. "
+                "Check https://github.com/vllm-project/vllm/issues/8402 for"
+                " more information."
+            )
+            raise RuntimeError(msg)
+        physical_device_id = device_ids[device_id]
+        return int(physical_device_id)
+    else:
+        return device_id
+
+
+def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+    @wraps(fn)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        pynvml.nvmlInit()
+        try:
+            return fn(*args, **kwargs)
+        finally:
+            pynvml.nvmlShutdown()
+
+    return wrapper
+
+
+class _CudaAttentionBackendResolver:
+    backend: AttentionBackendEnum
+
+    @classmethod
+    def resolve(cls, platform) -> str | AttentionBackendEnum:
+        raise NotImplementedError
+
+
+class _DirectCudaAttentionBackendResolver(_CudaAttentionBackendResolver):
+    backend_cls_str: str
+
+    @classmethod
+    def resolve(cls, platform) -> str:
+        return cls.backend_cls_str
+
+
+class _AITerAttentionBackendResolver(_DirectCudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.AITER
+    backend_cls_str = (
+        "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend"
+    )
+
+
+class _TorchSDPAAttentionBackendResolver(_DirectCudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.TORCH_SDPA
+    backend_cls_str = _SDPA_BACKEND_CLS_STR
+
+
+class _SparseLinearAttentionBackendResolver(_DirectCudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.SLA_ATTN
+    backend_cls_str = "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SparseLinearAttentionBackend"
+
+
+class _SageSparseLinearAttentionBackendResolver(_DirectCudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.SAGE_SLA_ATTN
+    backend_cls_str = "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SageSparseLinearAttentionBackend"
+
+
+class _SlidingTileAttentionBackendResolver(_CudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.SLIDING_TILE_ATTN
+
+    @classmethod
+    def resolve(cls, platform) -> str:
+        try:
+            from st_attn import sliding_tile_attention  # noqa: F401
+
+            from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import (  # noqa: F401
+                SlidingTileAttentionBackend,
+            )
+
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn.SlidingTileAttentionBackend"
+        except ImportError as e:
+            logger.error("Failed to import Sliding Tile Attention backend: %s", str(e))
+            raise ImportError(
+                "Sliding Tile Attention backend is not installed. "
+            ) from e
+
+
+class _SageAttentionBackendResolver(_CudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.SAGE_ATTN
+
+    @classmethod
+    def resolve(cls, platform) -> str | AttentionBackendEnum:
+        try:
+            from sageattention import sageattn  # noqa: F401
+
+            from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn import (  # noqa: F401
+                SageAttentionBackend,
+            )
+
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn.SageAttentionBackend"
+        except ImportError as e:
+            logger.info(e)
+            logger.info(
+                "Sage Attention backend is not installed (To install it, run `pip install sageattention==2.2.0 --no-build-isolation`). Falling back to Flash Attention."
+            )
+            return AttentionBackendEnum.FA
+
+
+class _SageAttention3BackendResolver(_CudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.SAGE_ATTN_3
+
+    @classmethod
+    def resolve(cls, platform) -> str | AttentionBackendEnum:
+        try:
+            from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3 import (  # noqa: F401
+                SageAttention3Backend,
+            )
+
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn3.SageAttention3Backend"
+        except ImportError as e:
+            logger.info(e)
+            logger.info(
+                "Sage Attention 3 backend is not installed (To install it, see https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell#installation). Falling back to Torch SDPA."
+            )
+            return AttentionBackendEnum.TORCH_SDPA
+
+
+class _VideoSparseAttentionBackendResolver(_CudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.VIDEO_SPARSE_ATTN
+
+    @classmethod
+    def resolve(cls, platform) -> str:
+        try:
+            from vsa import block_sparse_attn  # noqa: F401
+
+            from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import (  # noqa: F401
+                VideoSparseAttentionBackend,
+            )
+
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn.VideoSparseAttentionBackend"
+        except ImportError as e:
+            logger.error("Failed to import Video Sparse Attention backend: %s", str(e))
+            raise ImportError("Video Sparse Attention backend is not installed.") from e
+
+
+class _SparseVideoGen2AttentionBackendResolver(_CudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN
+
+    @classmethod
+    def resolve(cls, platform) -> str:
+        try:
+            from svg.kernels.triton.permute import (  # noqa: F401
+                apply_inverse_permutation_triton,
+                permute_tensor_by_labels_triton,
+            )
+            from svg.kmeans_utils import (  # noqa: F401
+                batch_kmeans_Euclid,
+                density_calculation,
+                dynamic_block_sparse_fwd_flashinfer,
+                identify_dynamic_map,
+            )
+
+            from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn import (  # noqa: F401
+                SparseVideoGen2AttentionBackend,
+            )
+
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn.SparseVideoGen2AttentionBackend"
+        except ImportError as e:
+            logger.error(
+                "Failed to import Sparse Video Gen 2 (SAP) Attention backend: %s",
+                str(e),
+            )
+            raise ImportError(
+                "Sparse Video Gen 2 (SAP) Attention backend is not installed. "
+                "Please install it by following the instructions at "
+                "https://github.com/svg-project/Sparse-VideoGen"
+            ) from e
+
+
+class _VMOBAAttentionBackendResolver(_CudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.VMOBA_ATTN
+
+    @classmethod
+    def resolve(cls, platform) -> str:
+        try:
+            from kernel.attn.vmoba_attn.vmoba import moba_attn_varlen  # noqa: F401
+
+            from sglang.multimodal_gen.runtime.layers.attention.backends.vmoba import (  # noqa: F401
+                VMOBAAttentionBackend,
+            )
+
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.vmoba.VMOBAAttentionBackend"
+        except ImportError as e:
+            logger.error("Failed to import Video MoBA Attention backend: %s", str(e))
+            raise ImportError("Video MoBA Attention backend is not installed. ") from e
+
+
+class _FlashAttention2BackendResolver(_CudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.FA2
+
+    @classmethod
+    def resolve(cls, platform) -> str:
+        from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2 import (  # noqa: F401
+            FlashAttention2Backend,
+        )
+
+        return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2.FlashAttention2Backend"
+
+
+class _FlashAttentionBackendResolver(_CudaAttentionBackendResolver):
+    backend = AttentionBackendEnum.FA
+
+    @classmethod
+    def resolve(cls, platform) -> AttentionBackendEnum:
+        if platform.is_sm120():
+            logger.info(
+                "FlashAttention is not supported on SM12.x in this build; falling back to Torch SDPA."
+            )
+            return AttentionBackendEnum.TORCH_SDPA
+        return AttentionBackendEnum.FA
+
+
+_CUDA_ATTENTION_BACKEND_RESOLVERS = {
+    resolver.backend: resolver
+    for resolver in (
+        _AITerAttentionBackendResolver,
+        _TorchSDPAAttentionBackendResolver,
+        _SparseLinearAttentionBackendResolver,
+        _SageSparseLinearAttentionBackendResolver,
+        _SlidingTileAttentionBackendResolver,
+        _SageAttentionBackendResolver,
+        _SageAttention3BackendResolver,
+        _VideoSparseAttentionBackendResolver,
+        _SparseVideoGen2AttentionBackendResolver,
+        _VMOBAAttentionBackendResolver,
+        _FlashAttention2BackendResolver,
+        _FlashAttentionBackendResolver,
+    )
+}
+
+
+class CudaPlatformBase(Platform):
+    _enum = PlatformEnum.CUDA
+    device_name: str = "cuda"
+    device_type: str = "cuda"
+    dispatch_key: str = "CUDA"
+    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
+
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        return torch.device(f"cuda:{envs.LOCAL_RANK}")
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        raise NotImplementedError
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        if enforce_eager:
+            logger.warning(
+                "To see benefits of async output processing, enable CUDA "
+                "graph. Since, enforce-eager is enabled, async output "
+                "processor cannot be used"
+            )
+            return False
+        return True
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_modelopt_fp4_quantize_op(cls) -> Callable | None:
+        try:
+            from flashinfer import fp4_quantize
+
+            return fp4_quantize
+        except ImportError:
+            pass
+
+        try:
+            from sgl_kernel import scaled_fp4_quant as fp4_quantize
+
+            return fp4_quantize
+        except ImportError:
+            return None
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_modelopt_flashinfer_fp4_backend(cls) -> str:
+        backend = envs.SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND
+        default_backend = "trtllm"
+        if backend is None:
+            return default_backend
+
+        backend = backend.lower()
+        backend = {
+            "flashinfer_cudnn": "cudnn",
+            "flashinfer_cutlass": "cutlass",
+            "flashinfer_trtllm": "trtllm",
+            "trtllm": "trtllm",
+            "cudnn": "cudnn",
+            "auto": "auto",
+        }.get(backend, backend)
+        if backend not in {"auto", "cudnn", "cutlass", "trtllm"}:
+            logger.warning(
+                "Unsupported SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=%r. "
+                "Falling back to %r.",
+                backend,
+                default_backend,
+            )
+            return default_backend
+        return backend
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_modelopt_fp4_gemm_op(cls) -> tuple[Callable | None, str | None]:
+        requested_backend = envs.SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND
+
+        try:
+            from flashinfer import mm_fp4 as flashinfer_mm_fp4
+
+            return flashinfer_mm_fp4, cls.get_modelopt_flashinfer_fp4_backend()
+        except ImportError:
+            logger.warning(
+                "Requested SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=%r "
+                "but flashinfer.mm_fp4 is unavailable. Falling back to "
+                "cutlass.",
+                requested_backend or "flashinfer_trtllm (default)",
+            )
+
+        try:
+            from sgl_kernel import cutlass_scaled_fp4_mm as cutlass_fp4_gemm
+
+            return cutlass_fp4_gemm, None
+        except ImportError:
+            return None, None
+
+    @classmethod
+    def is_full_nvlink(cls, device_ids: list[int]) -> bool:
+        raise NotImplementedError
+
+    @classmethod
+    def log_warnings(cls) -> None:
+        pass
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        torch.cuda.reset_peak_memory_stats(device)
+        return float(torch.cuda.max_memory_allocated(device))
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int | None = None,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        if empty_cache:
+            torch.cuda.empty_cache()
+
+        if device_id is None:
+            device_id = torch.cuda.current_device()
+
+        device_props = torch.cuda.get_device_properties(device_id)
+        if device_props.is_integrated:
+            free_gpu_memory = psutil.virtual_memory().available
+        else:
+            free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="cuda")
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_gpu_memory = float(tensor.item())
+
+        return free_gpu_memory / (1 << 30)
+
+    @classmethod
+    def _resolve_default_attn_backend(cls) -> AttentionBackendEnum:
+        if cls.is_sm120():
+            # On SM12.x, the sgl-kernel FlashAttention wheels may not include
+            # support yet. Default to Torch SDPA for correctness.
+            logger.info("Defaulting to Torch SDPA backend on SM12.x")
+            return AttentionBackendEnum.TORCH_SDPA
+        return AttentionBackendEnum.FA
+
+    @classmethod
+    def _prepare_flash_attention_for_blackwell(cls) -> bool:
+        if not cls.is_blackwell():
+            return True
+
+        try:
+            from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
+                set_fa_ver,
+            )
+        except ImportError:
+            logger.info(
+                "Cannot use FlashAttention backend because the "
+                "flash_attn package is not found. "
+                "Make sure that flash_attn was built and installed "
+                "(on by default)."
+            )
+            return False
+
+        set_fa_ver(4)
+        return True
+
+    @classmethod
+    def _resolve_flash_attention_backend_cls_str(
+        cls, target_backend: AttentionBackendEnum, head_size: int, dtype: torch.dtype
+    ) -> str:
+        if not cls.has_device_capability(80):
+            logger.info("Cannot use FlashAttention backend for Volta and Turing GPUs.")
+            target_backend = AttentionBackendEnum.TORCH_SDPA
+        elif dtype not in (torch.float16, torch.bfloat16):
+            logger.info(
+                "Cannot use FlashAttention backend for dtype other than "
+                "torch.float16 or torch.bfloat16."
+            )
+            target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        if (
+            target_backend == AttentionBackendEnum.FA
+            and not cls._prepare_flash_attention_for_blackwell()
+        ):
+            target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        if target_backend == AttentionBackendEnum.FA:
+            try:
+                from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (  # noqa: F401
+                    FlashAttentionBackend,
+                )
+
+                supported_sizes = FlashAttentionBackend.get_supported_head_sizes()
+                if head_size not in supported_sizes:
+                    logger.info(
+                        "Cannot use FlashAttention backend for head size %d.",
+                        head_size,
+                    )
+                    target_backend = AttentionBackendEnum.TORCH_SDPA
+            except ImportError:
+                logger.info(
+                    "Cannot use FlashAttention backend because the "
+                    "flash_attn package is not found. "
+                    "Make sure that flash_attn was built and installed "
+                    "(on by default)."
+                )
+                target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        if target_backend == AttentionBackendEnum.TORCH_SDPA:
+            return _SDPA_BACKEND_CLS_STR
+
+        return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend"
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        if selected_backend is None:
+            target_backend = cls._resolve_default_attn_backend()
+        else:
+            resolver = _CUDA_ATTENTION_BACKEND_RESOLVERS.get(selected_backend)
+            if resolver is None:
+                raise ValueError(f"Invalid attention backend for {cls.device_name}")
+
+            resolved_backend = resolver.resolve(cls)
+            if isinstance(resolved_backend, str):
+                return resolved_backend
+            target_backend = resolved_backend
+
+        return cls._resolve_flash_attention_backend_cls_str(
+            target_backend, head_size, dtype
+        )
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
+
+
+# NVML utils
+# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
+# all the related functions work on real physical device ids.
+# the major benefit of using NVML is that it will not initialize CUDA
+class NvmlCudaPlatform(CudaPlatformBase):
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        try:
+            physical_device_id = device_id_to_physical_device_id(device_id)
+            handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
+            major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
+            return DeviceCapability(major=major, minor=minor)
+        except RuntimeError:
+            return None
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def has_device_capability(
+        cls,
+        capability: tuple[int, int] | int,
+        device_id: int = 0,
+    ) -> bool:
+        try:
+            return bool(super().has_device_capability(capability, device_id))
+        except RuntimeError:
+            return False
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def get_device_name(cls, device_id: int = 0) -> str:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        return cls._get_physical_device_name(physical_device_id)
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
+        return str(pynvml.nvmlDeviceGetUUID(handle))
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
+        try:
+            return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
+        except pynvml.NVMLError_NotSupported:
+            return int(torch.cuda.get_device_properties(device_id).total_memory)
+
+    @classmethod
+    @with_nvml_context
+    def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool:
+        """
+        query if the set of gpus are fully connected by nvlink (1 hop)
+        """
+        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
+        for i, handle in enumerate(handles):
+            for j, peer_handle in enumerate(handles):
+                if i < j:
+                    try:
+                        p2p_status = pynvml.nvmlDeviceGetP2PStatus(
+                            handle,
+                            peer_handle,
+                            pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
+                        )
+                        if p2p_status != pynvml.NVML_P2P_STATUS_OK:
+                            return False
+                    except pynvml.NVMLError:
+                        logger.exception(
+                            "NVLink detection failed. This is normal if"
+                            " your machine has no NVLink equipped."
+                        )
+                        return False
+        return True
+
+    @classmethod
+    def _get_physical_device_name(cls, device_id: int = 0) -> str:
+        handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
+        return str(pynvml.nvmlDeviceGetName(handle))
+
+    @classmethod
+    @with_nvml_context
+    def log_warnings(cls) -> None:
+        device_ids: int = pynvml.nvmlDeviceGetCount()
+        if device_ids > 1:
+            device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
+            if (
+                len(set(device_names)) > 1
+                and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID"
+            ):
+                logger.warning(
+                    "Detected different devices in the system: %s. Please"
+                    " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to "
+                    "avoid unexpected behavior.",
+                    ", ".join(device_names),
+                )
+
+
+class NonNvmlCudaPlatform(CudaPlatformBase):
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
+        major, minor = torch.cuda.get_device_capability(device_id)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return str(torch.cuda.get_device_name(device_id))
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        device_props = torch.cuda.get_device_properties(device_id)
+        return int(device_props.total_memory)
+
+    @classmethod
+    def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool:
+        logger.exception(
+            "NVLink detection not possible, as context support was"
+            " not found. Assuming no NVLink available."
+        )
+        return False
+
+
+# Autodetect either NVML-enabled or non-NVML platform
+# based on whether NVML is available.
+nvml_available = False
+try:
+    try:
+        pynvml.nvmlInit()
+        nvml_available = True
+    except Exception:
+        # On Jetson, NVML is not supported.
+        nvml_available = False
+finally:
+    if nvml_available:
+        pynvml.nvmlShutdown()
+
+CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform
+
+try:
+    from sphinx.ext.autodoc.mock import _MockModule
+
+    if not isinstance(pynvml, _MockModule):
+        CudaPlatform.log_warnings()
+except ModuleNotFoundError:
+    CudaPlatform.log_warnings()

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

@@ -0,0 +1,426 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/interface.py
+from __future__ import annotations
+
+import enum
+import random
+from collections.abc import Callable
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import resolve_obj_by_qualname
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+        AttentionImpl,
+    )
+
+logger = init_logger(__name__)
+
+
+class AttentionBackendEnum(enum.Enum):
+    FA2 = enum.auto()
+    FA = enum.auto()
+    SLIDING_TILE_ATTN = enum.auto()
+    TORCH_SDPA = enum.auto()
+    SAGE_ATTN = enum.auto()
+    SAGE_ATTN_3 = enum.auto()
+    VIDEO_SPARSE_ATTN = enum.auto()
+    SPARSE_VIDEO_GEN_2_ATTN = enum.auto()
+    VMOBA_ATTN = enum.auto()
+    AITER = enum.auto()
+    AITER_SAGE = enum.auto()
+    SLA_ATTN = enum.auto()
+    SAGE_SLA_ATTN = enum.auto()
+    LASER_ATTN = enum.auto()
+    BLOCK_SPARSE_ATTN = enum.auto()
+    RAIN_FUSION_ATTN = enum.auto()
+    NO_ATTENTION = enum.auto()
+
+    def __str__(self):
+        return self.name.lower()
+
+    @property
+    def is_sparse(self) -> bool:
+        return self in {
+            AttentionBackendEnum.SLIDING_TILE_ATTN,
+            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+            AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,
+            AttentionBackendEnum.VMOBA_ATTN,
+            AttentionBackendEnum.SLA_ATTN,
+            AttentionBackendEnum.SAGE_SLA_ATTN,
+            AttentionBackendEnum.LASER_ATTN,
+            AttentionBackendEnum.BLOCK_SPARSE_ATTN,
+            AttentionBackendEnum.RAIN_FUSION_ATTN,
+        }
+
+
+class PlatformEnum(enum.Enum):
+    CUDA = enum.auto()
+    ROCM = enum.auto()
+    TPU = enum.auto()
+    CPU = enum.auto()
+    MPS = enum.auto()
+    NPU = enum.auto()
+    MUSA = enum.auto()
+    XPU = enum.auto()
+    OOT = enum.auto()
+    UNSPECIFIED = enum.auto()
+
+
+class CpuArchEnum(enum.Enum):
+    X86 = enum.auto()
+    ARM = enum.auto()
+    UNSPECIFIED = enum.auto()
+
+
+class DeviceCapability(NamedTuple):
+    major: int
+    minor: int
+
+    def as_version_str(self) -> str:
+        return f"{self.major}.{self.minor}"
+
+    def to_int(self) -> int:
+        """
+        Express device capability as an integer ``<major><minor>``.
+
+        It is assumed that the minor version is always a single digit.
+        """
+        assert 0 <= self.minor < 10
+        return self.major * 10 + self.minor
+
+
+class Platform:
+    _enum: PlatformEnum
+    device_name: str
+    device_type: str
+    device: torch.device | None = None  # Dummy attribute for compatibility
+
+    # available dispatch keys:
+    # check https://github.com/pytorch/pytorch/blob/313dac6c1ca0fa0cde32477509cce32089f8532a/torchgen/model.py#L134 # noqa
+    # use "CPU" as a fallback for platforms not registered in PyTorch
+    dispatch_key: str = "CPU"
+
+    # The torch.compile backend for compiling simple and
+    # standalone functions. The default value is "inductor" to keep
+    # the same behavior as PyTorch.
+    # NOTE: for the forward part of the model, vLLM has another separate
+    # compilation strategy.
+    simple_compile_backend: str = "inductor"
+
+    supported_quantization: list[str] = []
+
+    @lru_cache(maxsize=1)
+    def is_cuda(self) -> bool:
+        return self.is_cuda_static()
+
+    @lru_cache(maxsize=1)
+    def is_npu(self) -> bool:
+        return self._enum == PlatformEnum.NPU
+
+    @lru_cache(maxsize=1)
+    def is_rocm(self) -> bool:
+        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_hopper(cls):
+        if not cls.is_cuda_static():
+            return False
+        return torch.cuda.get_device_capability() == (9, 0)
+
+    @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, PlatformEnum.MUSA)
+
+    @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
+    @lru_cache(maxsize=1)
+    def is_amp_supported(cls) -> bool:
+        return True
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_float64_supported(cls) -> bool:
+        return True
+
+    @classmethod
+    def get_modelopt_fp4_quantize_op(cls) -> Callable | None:
+        return None
+
+    @classmethod
+    def get_modelopt_fp4_gemm_op(cls) -> tuple[Callable | None, str | None]:
+        return None, None
+
+    @classmethod
+    def get_modelopt_flashinfer_fp4_backend(cls) -> str:
+        return "auto"
+
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        raise NotImplementedError
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        """Get the attention backend class of a device."""
+        return ""
+
+    @classmethod
+    def get_device_capability(
+        cls,
+        device_id: int = 0,
+    ) -> DeviceCapability | None:
+        """Stateless version of :func:`torch.cuda.get_device_capability`."""
+        return None
+
+    @classmethod
+    def has_device_capability(
+        cls,
+        capability: tuple[int, int] | int,
+        device_id: int = 0,
+    ) -> bool:
+        """
+        Test whether this platform is compatible with a device capability.
+
+        The ``capability`` argument can either be:
+
+        - A tuple ``(major, minor)``.
+        - An integer ``<major><minor>``. (See :meth:`DeviceCapability.to_int`)
+        """
+        current_capability = cls.get_device_capability(device_id=device_id)
+        if current_capability is None:
+            return False
+
+        if isinstance(capability, tuple):
+            return current_capability >= capability
+
+        return current_capability.to_int() >= capability
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        """Get the name of a device."""
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        """Get the uuid of a device, e.g. the PCI bus ID."""
+        raise NotImplementedError
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        """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_npu():
+            return torch.device("npu", local_rank)
+        elif self.is_xpu():
+            return torch.device("xpu", 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_npu():
+            return "hccl"
+        elif self.is_musa():
+            return "mccl"
+        elif self.is_mps():
+            return "gloo"
+        elif self.is_cpu():
+            return "gloo"
+        elif self.is_xpu():
+            return "xccl"
+        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:
+        """
+        Check if the current platform supports async output.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def inference_mode(cls):
+        """A device-specific wrapper of `torch.inference_mode`.
+
+        This wrapper is recommended because some hardware backends such as TPU
+        do not support `torch.inference_mode`. In such a case, they will fall
+        back to `torch.no_grad` by overriding this method.
+        """
+        return torch.inference_mode(mode=True)
+
+    @classmethod
+    def seed_everything(cls, seed: int | None = None) -> None:
+        """
+        Set the seed of each random module.
+        `torch.manual_seed` will set seed on all devices.
+
+        Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20
+        """
+        if seed is not None:
+            random.seed(seed)
+            np.random.seed(seed)
+            torch.manual_seed(seed)
+            torch.get_device_module().manual_seed_all(seed)
+
+    @classmethod
+    def verify_model_arch(cls, model_arch: str) -> None:
+        """
+        Verify whether the current platform supports the specified model
+        architecture.
+
+        - This will raise an Error or Warning based on the model support on
+        the current platform.
+        - By default all models are considered supported.
+        """
+        pass
+
+    @classmethod
+    def verify_quantization(cls, quant: str) -> None:
+        """
+        Verify whether the quantization is supported by the current platform.
+        """
+        if cls.supported_quantization and quant not in cls.supported_quantization:
+            raise ValueError(
+                f"{quant} quantization is currently not supported in "
+                f"{cls.device_name}."
+            )
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        """
+        Return the memory usage in bytes.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int | None = None,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        """
+        Return the available memory in GiB.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        """
+        Get device specific communicator class for distributed communication.
+        """
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase"  # noqa
+
+    @classmethod
+    def get_cpu_architecture(cls) -> CpuArchEnum:
+        """Get the CPU architecture of the current platform."""
+        return CpuArchEnum.UNSPECIFIED
+
+    @classmethod
+    def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
+        """Whether to enable DIT layerwise offload by default on the current platform."""
+        return True
+
+    @classmethod
+    def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
+        """Apply platform-specific optimizations to VAE after loading."""
+        return vae
+
+    def get_attn_backend(self, *args, **kwargs) -> AttentionImpl:
+        attention_cls_str = self.get_attn_backend_cls_str(*args, **kwargs)
+        return resolve_obj_by_qualname(attention_cls_str)
+
+
+class UnspecifiedPlatform(Platform):
+    _enum = PlatformEnum.UNSPECIFIED
+    device_type = ""

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

@@ -0,0 +1,132 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+from functools import lru_cache
+from typing import Any
+
+import psutil
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+logger = init_logger(__name__)
+
+
+class MpsPlatform(Platform):
+    _enum = PlatformEnum.MPS
+    device_name: str = "mps"
+    device_type: str = "mps"
+    dispatch_key: str = "MPS"
+    device_control_env_var: str = "MPS_VISIBLE_DEVICES"
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_amp_supported(cls) -> bool:
+        return False
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_float64_supported(cls) -> bool:
+        return False
+
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        return torch.device("mps")
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+
+        return psutil.virtual_memory().total
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        if enforce_eager:
+            logger.warning(
+                "To see benefits of async output processing, enable MPS "
+                "graph. Since, enforce-eager is enabled, async output "
+                "processor cannot be used"
+            )
+            return False
+        return True
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        return 0.0
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int | None = None,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+
+        if empty_cache:
+            torch.mps.empty_cache()
+
+        # For MPS, available memory is essentially the system available memory
+        free_memory = psutil.virtual_memory().available
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_memory, dtype=torch.float32)
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_memory = float(tensor.item())
+
+        return free_memory / (1 << 30)
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible
+        logger.info("Using Torch SDPA backend for MPS.")
+        return (
+            "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+        )
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        # Use base communicator for MPS
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase"
+
+    @classmethod
+    def seed_everything(cls, seed: int | None = None) -> None:
+        """Set the seed for MPS device."""
+        if seed is not None:
+            import random
+
+            import numpy as np
+
+            random.seed(seed)
+            np.random.seed(seed)
+            torch.manual_seed(seed)
+            # MPS doesn't have manual_seed_all like CUDA
+            # The manual_seed above should be sufficient

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

@@ -0,0 +1,393 @@
+"""
+This file is a platform abstraction for MThreads (MUSA) GPUs,
+adjusted to match the structure and interface of `cuda.py`.
+"""
+
+import os
+from collections.abc import Callable
+from functools import lru_cache, wraps
+from typing import Any, TypeVar
+
+import psutil
+import pymtml
+
+# isort: off
+import torch
+import torchada  # noqa: F401
+
+# isort: on
+from typing_extensions import ParamSpec
+
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.runtime.platforms.interface import (
+    AttentionBackendEnum,
+    DeviceCapability,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def device_id_to_physical_device_id(device_id: int) -> int:
+    if "MUSA_VISIBLE_DEVICES" in os.environ:
+        device_ids = os.environ["MUSA_VISIBLE_DEVICES"].split(",")
+        if device_ids == [""]:
+            msg = (
+                "MUSA_VISIBLE_DEVICES is set to empty string, which means"
+                " GPU support is disabled. If you are using ray, please unset"
+                " the environment variable `MUSA_VISIBLE_DEVICES` inside the"
+                " worker/actor. "
+                "Check https://github.com/vllm-project/vllm/issues/8402 for"
+                " more information."
+            )
+            raise RuntimeError(msg)
+        physical_device_id = device_ids[device_id]
+        return int(physical_device_id)
+    else:
+        return device_id
+
+
+def with_mtml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+    @wraps(fn)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        pymtml.nvmlInit()
+        try:
+            return fn(*args, **kwargs)
+        finally:
+            pymtml.nvmlShutdown()
+
+    return wrapper
+
+
+class MusaPlatformBase(Platform):
+    _enum = PlatformEnum.MUSA
+    device_name: str = "musa"
+    device_type: str = "musa"
+    dispatch_key: str = "MUSA"
+    device_control_env_var: str = "MUSA_VISIBLE_DEVICES"
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_float64_supported(cls) -> bool:
+        return False
+
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        return torch.device(f"musa:{envs.LOCAL_RANK}")
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        raise NotImplementedError
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        if enforce_eager:
+            logger.warning(
+                "To see benefits of async output processing, enable MUSA "
+                "graph. Since, enforce-eager is enabled, async output "
+                "processor cannot be used"
+            )
+            return False
+        return True
+
+    @classmethod
+    def is_full_mtlink(cls, device_ids: list[int]) -> bool:
+        raise NotImplementedError
+
+    @classmethod
+    def log_warnings(cls) -> None:
+        pass
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        torch.cuda.reset_peak_memory_stats(device)
+        return float(torch.cuda.max_memory_allocated(device))
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int | None = None,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        if empty_cache:
+            torch.cuda.empty_cache()
+
+        if device_id is None:
+            device_id = torch.cuda.current_device()
+
+        device_props = torch.cuda.get_device_properties(device_id)
+        if device_props.is_integrated:
+            free_gpu_memory = psutil.virtual_memory().available
+        else:
+            free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="musa")
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_gpu_memory = float(tensor.item())
+
+        return free_gpu_memory / (1 << 30)
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        target_backend: AttentionBackendEnum | None = None
+
+        if selected_backend == AttentionBackendEnum.TORCH_SDPA:
+            logger.info("Using Torch SDPA backend")
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+        elif selected_backend == AttentionBackendEnum.SAGE_ATTN:
+            try:
+                from sageattention import sageattn  # noqa: F401
+
+                from sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn import (  # noqa: F401
+                    SageAttentionBackend,
+                )
+
+                logger.info("Using Sage Attention backend")
+
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.sage_attn.SageAttentionBackend"
+            except ImportError as e:
+                logger.info(e)
+                logger.info(
+                    "Sage Attention backend is not installed (To install it, run `pip install sageattention>=0.1.0`). Falling back to Flash Attention."
+                )
+                target_backend = AttentionBackendEnum.FA
+        elif selected_backend in [
+            AttentionBackendEnum.FA,
+        ]:
+            target_backend = AttentionBackendEnum.FA
+        elif selected_backend:
+            raise ValueError(f"Invalid attention backend for {cls.device_name}")
+        else:
+            target_backend = AttentionBackendEnum.FA
+
+        # Ensure we have a target backend selected before validation/fallback.
+        if target_backend is None:
+            target_backend = AttentionBackendEnum.FA
+
+        if dtype not in (torch.float16, torch.bfloat16):
+            logger.info(
+                "Cannot use FlashAttention backend for dtype other than "
+                "torch.float16 or torch.bfloat16."
+            )
+            target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        # FlashAttn is valid for the model, checking if the package is
+        # installed.
+        if target_backend == AttentionBackendEnum.FA:
+            try:
+                from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (  # noqa: F401
+                    FlashAttentionBackend,
+                )
+
+                supported_sizes = FlashAttentionBackend.get_supported_head_sizes()
+                if head_size not in supported_sizes:
+                    logger.info(
+                        "Cannot use FlashAttention backend for head size %d.",
+                        head_size,
+                    )
+                    target_backend = AttentionBackendEnum.TORCH_SDPA
+            except ImportError:
+                logger.info(
+                    "Cannot use FlashAttention backend because the "
+                    "flash_attn package is not found. "
+                    "Make sure that flash_attn was built and installed "
+                    "(on by default)."
+                )
+                target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        if target_backend == AttentionBackendEnum.TORCH_SDPA:
+            logger.info("Using Torch SDPA backend")
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+
+        logger.info("Using FlashAttention (FA3) backend")
+        return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend"
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
+
+
+# MTML utils
+# Note that MTML is not affected by `MUSA_VISIBLE_DEVICES`,
+# all the related functions work on real physical device ids.
+# the major benefit of using MTML is that it will not initialize MUSA
+class MtmlMusaPlatform(MusaPlatformBase):
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        try:
+            physical_device_id = device_id_to_physical_device_id(device_id)
+            handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id)
+            major, minor = pymtml.nvmlDeviceGetCudaComputeCapability(handle)
+            return DeviceCapability(major=major, minor=minor)
+        except RuntimeError:
+            return None
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def has_device_capability(
+        cls,
+        capability: tuple[int, int] | int,
+        device_id: int = 0,
+    ) -> bool:
+        try:
+            return bool(super().has_device_capability(capability, device_id))
+        except RuntimeError:
+            return False
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def get_device_name(cls, device_id: int = 0) -> str:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        return cls._get_physical_device_name(physical_device_id)
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id)
+        return str(pymtml.nvmlDeviceGetUUID(handle))
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_mtml_context
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        handle = pymtml.nvmlDeviceGetHandleByIndex(physical_device_id)
+        return int(pymtml.nvmlDeviceGetMemoryInfo(handle).total)
+
+    @classmethod
+    @with_mtml_context
+    def is_full_mtlink(cls, physical_device_ids: list[int]) -> bool:
+        """
+        query if the set of gpus are fully connected by mtlink (1 hop)
+        """
+        handles = [pymtml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
+        for i, handle in enumerate(handles):
+            for j, peer_handle in enumerate(handles):
+                if i < j:
+                    try:
+                        p2p_status = pymtml.nvmlDeviceGetP2PStatus(
+                            handle,
+                            peer_handle,
+                            pymtml.NVML_P2P_CAPS_INDEX_NVLINK,
+                        )
+                        if p2p_status != pymtml.NVML_P2P_STATUS_OK:
+                            return False
+                    except pymtml.NVMLError:
+                        logger.exception(
+                            "MTLink detection failed. This is normal if"
+                            " your machine has no MTLink equipped."
+                        )
+                        return False
+        return True
+
+    @classmethod
+    def _get_physical_device_name(cls, device_id: int = 0) -> str:
+        handle = pymtml.nvmlDeviceGetHandleByIndex(device_id)
+        return str(pymtml.nvmlDeviceGetName(handle))
+
+    @classmethod
+    @with_mtml_context
+    def log_warnings(cls) -> None:
+        device_ids: int = pymtml.nvmlDeviceGetCount()
+        if device_ids > 1:
+            device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
+            if (
+                len(set(device_names)) > 1
+                and os.environ.get("MUSA_DEVICE_ORDER") != "PCI_BUS_ID"
+            ):
+                logger.warning(
+                    "Detected different devices in the system: %s. Please"
+                    " make sure to set `MUSA_DEVICE_ORDER=PCI_BUS_ID` to "
+                    "avoid unexpected behavior.",
+                    ", ".join(device_names),
+                )
+
+
+class NonMtmlMusaPlatform(MusaPlatformBase):
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
+        major, minor = torch.cuda.get_device_capability(device_id)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return str(torch.cuda.get_device_name(device_id))
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        device_props = torch.cuda.get_device_properties(device_id)
+        return int(device_props.total_memory)
+
+    @classmethod
+    def is_full_mtlink(cls, physical_device_ids: list[int]) -> bool:
+        logger.error(
+            "MTLink detection not possible, as context support was"
+            " not found. Assuming no MTLink available."
+        )
+        return False
+
+
+# Autodetect either MTML-enabled or non-MTML platform
+# based on whether MTML is available.
+mtml_available = False
+
+if "MUSA_DISABLE_MTML" not in os.environ:
+    try:
+        try:
+            pymtml.nvmlInit()
+            mtml_available = True
+        except Exception:
+            mtml_available = False
+    finally:
+        if mtml_available:
+            pymtml.nvmlShutdown()
+
+MusaPlatform = MtmlMusaPlatform if mtml_available else NonMtmlMusaPlatform
+
+try:
+    from sphinx.ext.autodoc.mock import _MockModule
+
+    if not isinstance(pymtml, _MockModule):
+        MusaPlatform.log_warnings()
+except ModuleNotFoundError:
+    MusaPlatform.log_warnings()
+
+if __name__ == "__main__":
+    print(MusaPlatform.__name__)
+    print(MusaPlatform.get_device_name())
+    print(MusaPlatform.get_device_capability())
+    print(MusaPlatform.get_device_total_memory())
+    print(MusaPlatform.is_full_mtlink([0, 1, 2, 3, 4, 5, 6, 7]))

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

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm-ascend: https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/platform.py
+
+import os
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.runtime.platforms.interface import (
+    AttentionBackendEnum,
+    DeviceCapability,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def device_id_to_physical_device_id(device_id: int) -> int:
+    if "ASCEND_RT_VISIBLE_DEVICES" in os.environ:
+        device_ids = os.environ["ASCEND_RT_VISIBLE_DEVICES"].split(",")
+        if device_ids == [""]:
+            msg = (
+                "ASCEND_RT_VISIBLE_DEVICES is set to empty string, which means"
+                " NPU support is disabled"
+            )
+            raise RuntimeError(msg)
+        physical_device_id = device_ids[device_id]
+        return int(physical_device_id)
+    else:
+        return device_id
+
+
+class NPUPlatformBase(Platform):
+    _enum = PlatformEnum.NPU
+    device_name: str = "npu"
+    device_type: str = "npu"
+    dispatch_key: str = "NPU"
+    device_control_env_var: str = "ASCEND_RT_VISIBLE_DEVICES"
+
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        return torch.device(f"npu:{envs.LOCAL_RANK}")
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
+        return None
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return str(torch.npu.get_device_name(device_id))
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        device_props = torch.npu.get_device_properties(device_id)
+        return int(device_props.total_memory)
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        if enforce_eager:
+            logger.warning(
+                "To see benefits of async output processing, enable NPU "
+                "graph. Since, enforce-eager is enabled, async output "
+                "processor cannot be used"
+            )
+            return False
+        return True
+
+    @classmethod
+    def inference_mode(cls):
+        # npu kernels in diffusion paths may need tensor version counters
+        return torch.no_grad()
+
+    @classmethod
+    def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool:
+        logger.exception(
+            "NVLink detection not possible, as context support was"
+            " not found. Assuming no NVLink available."
+        )
+        return False
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int | None = None,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        if empty_cache:
+            torch.npu.empty_cache()
+
+        if device_id is None:
+            device_id = torch.npu.current_device()
+
+        free_gpu_memory, _ = torch.npu.mem_get_info(device_id)
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="npu")
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_gpu_memory = float(tensor.item())
+
+        return free_gpu_memory / (1 << 30)
+
+    @classmethod
+    def log_warnings(cls) -> None:
+        pass
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        torch.npu.reset_peak_memory_stats(device)
+        return float(torch.npu.max_memory_allocated(device))
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        if selected_backend == AttentionBackendEnum.FA:
+            logger.info("Using Ascend Flash Attention backend.")
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.ascend_fa.AscendFABackend"
+
+        elif selected_backend == AttentionBackendEnum.LASER_ATTN:
+            try:
+                from sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn import (  # noqa: F401
+                    LaserAttentionBackend,
+                )
+
+                logger.info("Using Laser Attention backend")
+
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.laser_attn.LaserAttentionBackend"
+            except ImportError as e:
+                logger.error(f"Failed to import Laser Attention backend: {e}")
+                raise ImportError(
+                    "Laser Attention backend is not installed. "
+                    "It requires the `attentions` module which can be installed along with sgl_kernel_npu. "
+                    "Manual installation from source is required. See https://github.com/sgl-project/sgl-kernel-npu."
+                ) from e
+
+        elif selected_backend == AttentionBackendEnum.BLOCK_SPARSE_ATTN:
+            try:
+                from sglang.multimodal_gen.runtime.layers.attention.backends.block_sparse_attn import (  # noqa: F401
+                    BlockSparseAttentionBackend,
+                )
+
+                logger.info("Using Block Sparse Attention backend")
+
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.block_sparse_attn.BlockSparseAttentionBackend"
+            except ImportError as e:
+                logger.error(f"Failed to import Block Sparse Attention backend: {e}")
+                raise ImportError(
+                    "Block Sparse Attention backend is not installed. "
+                    "It requires the `attentions` module which can be installed along with sgl_kernel_npu. "
+                    "Manual installation from source is required. See https://github.com/sgl-project/sgl-kernel-npu."
+                ) from e
+
+        elif selected_backend == AttentionBackendEnum.RAIN_FUSION_ATTN:
+            try:
+                from sglang.multimodal_gen.runtime.layers.attention.backends.rain_fusion_attn import (  # noqa: F401
+                    RainFusionAttentionBackend,
+                )
+
+                logger.info("Using Rain Fusion Attention backend")
+
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.rain_fusion_attn.RainFusionAttentionBackend"
+            except ImportError as e:
+                logger.error(f"Failed to import Rain Fusion Attention backend: {e}")
+                raise ImportError(
+                    "Rain Fusion Attention backend is not installed. "
+                    "It requires the `attentions` module which can be installed along with sgl_kernel_npu. "
+                    "Manual installation from source is required. See https://github.com/sgl-project/sgl-kernel-npu."
+                ) from e
+
+        logger.info("Using Torch SDPA backend.")
+        return (
+            "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+        )
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
+
+    @classmethod
+    def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
+        """The performance of the layerwise_offload feature depends on the device's memory size and the memory size occupied by the model. Use --dit-layerwise-offload True if it suitable for your case."""
+        return False

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

@@ -0,0 +1,402 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from rocm/vllm: https://github.com/ROCm/vllm/blob/v0.7.3%2Brocm/vllm/platforms/rocm.py
+"""
+This file is a platform abstraction for ROCm GPUs,
+adjusted to match the structure and interface of `cuda.py`.
+"""
+
+import types
+from functools import lru_cache
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import sglang.multimodal_gen.envs as envs
+from sglang.multimodal_gen.runtime.platforms.interface import (
+    AttentionBackendEnum,
+    DeviceCapability,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# ROCm uses the same torch.cuda interface
+class RocmPlatform(Platform):
+    _enum = PlatformEnum.ROCM
+    device_name: str = "rocm"
+    device_type: str = "cuda"  # torch uses 'cuda' backend string
+    dispatch_key: str = "CUDA"
+    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
+
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        return torch.device(f"cuda:{envs.LOCAL_RANK}")
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
+        major, minor = torch.cuda.get_device_capability(device_id)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return str(torch.cuda.get_device_name(device_id))
+
+    @classmethod
+    @lru_cache(maxsize=1)
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        return torch.cuda.get_device_properties(device_id).total_memory
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        if enforce_eager:
+            logger.warning(
+                "To see benefits of async output processing, enable CUDA graph. "
+                "Since enforce-eager is enabled, async output processor cannot be used"
+            )
+            return False
+        return True
+
+    @classmethod
+    def log_warnings(cls) -> None:
+        pass  # ROCm-specific warnings can be added here
+
+    @classmethod
+    def get_current_memory_usage(cls, device: torch.device | None = None) -> float:
+        torch.cuda.reset_peak_memory_stats(device)
+        return float(torch.cuda.max_memory_allocated(device))
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int | None = None,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group: Any = None,
+    ) -> float:
+        if empty_cache:
+            torch.cuda.empty_cache()
+
+        if device_id is None:
+            device_id = torch.cuda.current_device()
+
+        free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(free_gpu_memory, dtype=torch.float32, device="cuda")
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_gpu_memory = float(tensor.item())
+
+        return free_gpu_memory / (1 << 30)
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        if selected_backend == AttentionBackendEnum.TORCH_SDPA:
+            logger.info("Using Torch SDPA backend.")
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+
+        elif selected_backend in (AttentionBackendEnum.FA, None):
+            pass
+
+        elif selected_backend == AttentionBackendEnum.AITER:
+            if dtype not in (torch.float16, torch.bfloat16):
+                logger.warning(
+                    "AITer backend works best with fp16/bf16 inputs but got dtype=%s. "
+                    "Proceeding with AITer anyway.",
+                    dtype,
+                )
+            logger.info("Using AITer backend on ROCm.")
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend"
+
+        elif selected_backend == AttentionBackendEnum.AITER_SAGE:
+            if dtype in (torch.float16, torch.bfloat16):
+                logger.info("Using AITER Sage backend on ROCm.")
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter_sage.AITERSageBackend"
+            else:
+                logger.warning(
+                    "AITER Sage backend only supports bf16/fp16 inputs but got dtype=%s.",
+                    dtype,
+                )
+
+        elif selected_backend in (
+            AttentionBackendEnum.SLIDING_TILE_ATTN,
+            AttentionBackendEnum.SAGE_ATTN,
+        ):
+            raise ValueError(
+                f"{selected_backend.name} is not supported on {cls.device_name}."
+            )
+        elif selected_backend:
+            raise ValueError(
+                f"Invalid attention backend for {cls.device_name}: {selected_backend}"
+            )
+
+        target_backend = AttentionBackendEnum.FA
+        if dtype not in (torch.float16, torch.bfloat16):
+            logger.info(
+                "Cannot use FlashAttention backend for dtype other than "
+                "torch.float16 or torch.bfloat16."
+            )
+            target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        if target_backend == AttentionBackendEnum.FA:
+            try:
+                import flash_attn  # noqa: F401
+
+                from sglang.jit_kernel.flash_attention_v3 import _is_fa3_supported
+                from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (  # noqa: F401
+                    FlashAttentionBackend,
+                )
+
+                if not _is_fa3_supported():
+                    logger.info(
+                        "FlashAttention backend now dispatches through FA3 "
+                        "(CUDA-only). Using Torch SDPA backend on ROCm."
+                    )
+                    target_backend = AttentionBackendEnum.TORCH_SDPA
+
+                if target_backend == AttentionBackendEnum.FA:
+                    supported_sizes = FlashAttentionBackend.get_supported_head_sizes()
+                    if head_size not in supported_sizes:
+                        logger.info(
+                            "Cannot use FlashAttention-2 backend for head size %d.",
+                            head_size,
+                        )
+                        target_backend = AttentionBackendEnum.TORCH_SDPA
+            except ImportError:
+                logger.info(
+                    "Cannot use FlashAttention backend because the "
+                    "flash_attn package is not found. "
+                    "Make sure that flash_attn was built and installed "
+                    "(on by default)."
+                )
+                target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        if target_backend == AttentionBackendEnum.TORCH_SDPA:
+            logger.info("Using Torch SDPA backend.")
+
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+
+        logger.info("Using Flash Attention backend.")
+
+        return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend"
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # works for ROCm too
+
+    @classmethod
+    def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
+        """Apply ROCm-specific optimizations to VAE.
+
+        - Enable MIOpen benchmark mode so that the best convolution algorithm
+          is selected for each distinct input shape (benefits Conv3d-heavy VAE
+          decode).
+        - Replace nn.GroupNorm with AITer GroupNorm when available.
+        - Replace CausalConv3d (3x3x3) with temporal-unfolded batched Conv2D.
+        """
+        if envs.SGLANG_USE_ROCM_CUDNN_BENCHMARK and not torch.backends.cudnn.benchmark:
+            torch.backends.cudnn.benchmark = True
+            logger.info(
+                "Enabled cudnn.benchmark (MIOpen auto-tuning) for VAE conv layers"
+            )
+
+        if envs.SGLANG_USE_ROCM_VAE:
+            try:
+                from aiter.ops.groupnorm import GroupNorm as AiterGroupNorm
+
+                count = cls._replace_groupnorm(vae, AiterGroupNorm)
+                if count > 0:
+                    logger.info(
+                        "Replaced %d nn.GroupNorm modules with AITer GroupNorm in VAE",
+                        count,
+                    )
+            except Exception:
+                logger.warning(
+                    "Failed to apply AITer GroupNorm to VAE.",
+                    exc_info=True,
+                )
+
+        use_bf16 = envs.SGLANG_USE_ROCM_VAE_CONV2D_BF16
+        use_conv2d = envs.SGLANG_USE_ROCM_VAE_CONV2D or use_bf16
+        if use_conv2d:
+            count = cls._replace_conv3d_with_conv2d(vae, use_bf16=use_bf16)
+            if count > 0:
+                mode = "BF16" if use_bf16 else "same dtype"
+                logger.info(
+                    "Replaced %d CausalConv3d modules with batched Conv2D "
+                    "(compute=%s) in VAE",
+                    count,
+                    mode,
+                )
+
+        return vae
+
+    @staticmethod
+    def _replace_groupnorm(module: torch.nn.Module, aiter_gn_cls: type) -> int:
+        count = 0
+        for name, child in module.named_children():
+            if isinstance(child, torch.nn.GroupNorm) and child.affine:
+                replacement = aiter_gn_cls(
+                    num_groups=child.num_groups,
+                    num_channels=child.num_channels,
+                    eps=child.eps,
+                    affine=True,
+                    device=child.weight.device,
+                    dtype=child.weight.dtype,
+                )
+                replacement.weight = child.weight
+                replacement.bias = child.bias
+                setattr(module, name, replacement)
+                count += 1
+            else:
+                count += RocmPlatform._replace_groupnorm(child, aiter_gn_cls)
+        return count
+
+    @staticmethod
+    def _conv3d_as_batched_conv2d(
+        x_padded: torch.Tensor,
+        weight_2d: torch.Tensor,
+        bias: torch.Tensor | None,
+        stride: tuple[int, ...],
+        kt: int,
+        compute_bf16: bool = False,
+    ) -> torch.Tensor:
+        """Replace F.conv3d with temporal-unfolded batched Conv2D.
+
+        ``x_padded`` must already be spatially/temporally padded so that
+        ``F.conv3d(x_padded, weight, bias, stride, padding=0)`` would produce
+        the correct output.  This routine unfolds along the temporal axis,
+        reshapes into a batch of 2-D frames, runs ``F.conv2d``, and folds the
+        result back.
+
+        *weight_2d* is the pre-transformed 2-D kernel
+        ``[C_out, Kt*C_in, Kh, Kw]``, cached at patch time to avoid
+        redundant permute/reshape on every forward call.
+
+        When *compute_bf16* is True the convolution is executed in BF16 and
+        the output is cast back to the original dtype.
+        """
+        orig_dtype = x_padded.dtype
+        N, C_in, T, H, W = x_padded.shape
+        C_out = weight_2d.shape[0]
+        stride_t, stride_h, stride_w = stride
+
+        T_out = (T - kt) // stride_t + 1
+
+        # (N, C_in, T, H, W) -> (N, T_out, Kt, C_in, H, W) -> (N*T_out, Kt*C_in, H, W)
+        unfolded = x_padded.unfold(2, kt, stride_t)
+        unfolded = unfolded.permute(0, 2, 5, 1, 3, 4).reshape(
+            N * T_out, kt * C_in, H, W
+        )
+
+        w = weight_2d
+        if compute_bf16 and orig_dtype != torch.bfloat16:
+            unfolded = unfolded.to(torch.bfloat16)
+            w = w.to(torch.bfloat16)
+            b = bias.to(torch.bfloat16) if bias is not None else None
+        else:
+            b = bias
+
+        out = F.conv2d(unfolded, w, b, stride=(stride_h, stride_w))
+
+        if compute_bf16 and orig_dtype != torch.bfloat16:
+            out = out.to(orig_dtype)
+
+        _, _, H_out, W_out = out.shape
+        return out.reshape(N, T_out, C_out, H_out, W_out).permute(0, 2, 1, 3, 4)
+
+    @staticmethod
+    def _replace_conv3d_with_conv2d(
+        module: torch.nn.Module, use_bf16: bool = False
+    ) -> int:
+        """Walk *module* and patch every CausalConv3d that has a 3-D kernel.
+
+        A ``CausalConv3d`` is identified as any ``nn.Conv3d`` subclass that
+        carries a ``_padding`` attribute (set by the Wan / diffusers causal
+        conv wrapper).  Only modules whose kernel is truly 3-D (Kt>1, Kh>1,
+        Kw>1) are replaced; pointwise or 1-D-temporal convolutions are left
+        untouched.  Modules with non-default ``groups`` or ``dilation`` are
+        skipped as the 2-D decomposition assumes groups=1 and dilation=1.
+        """
+        patched = 0
+        skipped = 0
+        for _name, child in module.named_modules():
+            if not isinstance(child, nn.Conv3d):
+                continue
+            if not hasattr(child, "_padding"):
+                continue
+            kt, kh, kw = child.kernel_size
+            if kt <= 1 or kh <= 1 or kw <= 1:
+                skipped += 1
+                continue
+            if child.groups != 1 or any(d != 1 for d in child.dilation):
+                skipped += 1
+                continue
+
+            padding = child._padding
+            stride = child.stride
+
+            # Pre-compute the 2-D weight: [C_out, C_in, Kt, Kh, Kw]
+            # -> [C_out, Kt*C_in, Kh, Kw]  (cached as a buffer)
+            weight_2d = (
+                child.weight.data.permute(0, 2, 1, 3, 4)
+                .reshape(child.out_channels, kt * child.in_channels, kh, kw)
+                .contiguous()
+            )
+            child.register_buffer("_weight_2d", weight_2d)
+
+            def _patched_forward(
+                self,
+                x,
+                cache_x=None,
+                *,
+                _padding=padding,
+                _stride=stride,
+                _kt=kt,
+                _bf16=use_bf16,
+            ):
+                pad = list(_padding)
+                if cache_x is not None and _padding[4] > 0:
+                    cache_x = cache_x.to(x.device)
+                    x = torch.cat([cache_x, x], dim=2)
+                    pad[4] -= cache_x.shape[2]
+                x = F.pad(x, pad)
+                x = x.to(self.weight.dtype)
+                return RocmPlatform._conv3d_as_batched_conv2d(
+                    x,
+                    self._weight_2d,
+                    self.bias,
+                    _stride,
+                    _kt,
+                    compute_bf16=_bf16,
+                )
+
+            child.forward = types.MethodType(_patched_forward, child)
+            patched += 1
+
+        logger.info(
+            "Conv3D→Conv2D: patched %d CausalConv3d (3D kernel, compute=%s), "
+            "skipped %d (1D/pointwise/grouped)",
+            patched,
+            "BF16" if use_bf16 else "same dtype",
+            skipped,
+        )
+        return patched
+
+    @classmethod
+    def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
+        """ROCm performs better without DIT layerwise offload on Wan."""
+        return False

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

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+# Intel XPU Platform support for SGLang Diffusion
+
+import torch
+
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.runtime.platforms.interface import (
+    AttentionBackendEnum,
+    DeviceCapability,
+    Platform,
+    PlatformEnum,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class XpuPlatform(Platform):
+    """Platform implementation for Intel XPU (Data Center GPU Max, Arc, etc.)."""
+
+    _enum = PlatformEnum.XPU
+    device_name: str = "xpu"
+    device_type: str = "xpu"
+    dispatch_key: str = "XPU"
+    device_control_env_var: str = "ZE_AFFINITY_MASK"
+
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        return torch.device(f"xpu:{envs.LOCAL_RANK}")
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        device = torch.xpu.current_device()
+        major, minor = torch.ops.sgl_kernel.query_device.default(device)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        """Get the name of the Intel XPU device."""
+        return torch.xpu.get_device_name(device_id)
+
+    @classmethod
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        """Get the UUID of the Intel XPU device."""
+        props = torch.xpu.get_device_properties(device_id)
+        return str(props.uuid)
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        """Get total memory of the Intel XPU device in bytes."""
+        props = torch.xpu.get_device_properties(device_id)
+        return props.total_memory
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        """Check if async output is supported on Intel XPU."""
+        if enforce_eager:
+            logger.warning(
+                "To see benefits of async output processing, disable enforce-eager. "
+                "Since enforce-eager is enabled, async output processor cannot be used"
+            )
+            return False
+        return True
+
+    @classmethod
+    def log_warnings(cls) -> None:
+        """Log any XPU-specific warnings."""
+        pass
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        """Get current memory usage on Intel XPU."""
+        torch.xpu.reset_peak_memory_stats(device)
+        return float(torch.xpu.max_memory_allocated(device))
+
+    @classmethod
+    def get_available_gpu_memory(
+        cls,
+        device_id: int = 0,
+        distributed: bool = False,
+        empty_cache: bool = True,
+        cpu_group=None,
+    ) -> float:
+        """Return the available device memory in GiB."""
+
+        if not (hasattr(torch, "xpu") and torch.xpu.is_available()):
+            return 0.0
+
+        num_gpus = torch.xpu.device_count()
+        if device_id < 0 or device_id >= num_gpus:
+            raise ValueError(f"Invalid XPU device_id={device_id}. num_gpus={num_gpus}")
+
+        current = torch.xpu.current_device()
+        if current != device_id:
+            logger.warning(
+                "current device is not %s, but %s; this may cause useless memory allocation for torch XPU context.",
+                device_id,
+                current,
+            )
+
+        if empty_cache:
+            torch.xpu.empty_cache()
+
+        # Use mem_get_info() with a sanity cap to avoid KV-cache over-allocation
+        # on drivers that incorrectly return total memory as free memory.
+        # Consistent with the fallback: free = max(0, total - allocated).
+        try:
+            free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info(device_id)
+            used_memory = float(torch.xpu.memory_allocated(device_id))
+            free_gpu_memory = min(
+                float(free_gpu_memory),
+                max(0.0, float(total_gpu_memory) - used_memory),
+            )
+        except Exception:
+            # Fallback for devices/drivers that do not support querying free memory
+            used_memory = float(torch.xpu.memory_allocated(device_id))
+            total_gpu_memory = float(
+                torch.xpu.get_device_properties(device_id).total_memory
+            )
+            free_gpu_memory = max(0.0, total_gpu_memory - used_memory)
+
+        if distributed:
+            import torch.distributed as dist
+
+            tensor = torch.tensor(
+                free_gpu_memory,
+                dtype=torch.float32,
+                device=torch.device("xpu", device_id),
+            )
+            dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=cpu_group)
+            free_gpu_memory = float(tensor.item())
+
+        return free_gpu_memory / (1 << 30)
+
+    @classmethod
+    def get_attn_backend_cls_str(
+        cls,
+        selected_backend: AttentionBackendEnum | None,
+        head_size: int,
+        dtype: torch.dtype,
+    ) -> str:
+        """Get the attention backend class string for Intel XPU.
+
+        Defaults to XPU backend (requires fp16/bf16 and a supported head size),
+        falling back to Torch SDPA if constraints are not met.
+        """
+        if selected_backend in (AttentionBackendEnum.FA, None):
+            if dtype not in (torch.float16, torch.bfloat16):
+                logger.info(
+                    "XPU attention backend requires fp16/bf16 but got dtype=%s; falling back to Torch SDPA.",
+                    dtype,
+                )
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+
+            try:
+                from sglang.multimodal_gen.runtime.layers.attention.backends.xpu_backend import (  # noqa: F401
+                    XPUAttentionBackend,
+                )
+
+                supported_sizes = XPUAttentionBackend.get_supported_head_sizes()
+                if head_size not in supported_sizes:
+                    logger.info(
+                        "XPU attention backend does not support head_size=%d; falling back to Torch SDPA.",
+                        head_size,
+                    )
+                    return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+
+                logger.info("Using XPU attention backend on Intel XPU.")
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.xpu_backend.XPUAttentionBackend"
+            except Exception as e:
+                logger.warning(
+                    "Failed to import/use XPU attention backend (%s); falling back to Torch SDPA.",
+                    e,
+                )
+                return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+
+        if selected_backend == AttentionBackendEnum.TORCH_SDPA:
+            logger.info("Using Torch SDPA backend for Intel XPU.")
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+
+        if selected_backend in (
+            AttentionBackendEnum.SLIDING_TILE_ATTN,
+            AttentionBackendEnum.SAGE_ATTN,
+            AttentionBackendEnum.SAGE_ATTN_3,
+            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+            AttentionBackendEnum.VMOBA_ATTN,
+            AttentionBackendEnum.AITER,
+        ):
+            logger.warning(
+                f"{selected_backend.name} is not supported on Intel XPU. "
+                "Falling back to Torch SDPA backend."
+            )
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+
+        # Default fallback
+        logger.info("Using Torch SDPA backend for Intel XPU (default).")
+        return (
+            "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
+        )
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        """Get device communicator class for Intel XPU distributed communication."""
+        # Use base communicator for now; can be updated to use oneCCL-based communicator
+        return "sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase"