[diffusion] platform: support Intel XPU (#17920)

Signed-off-by: sushil.dubey <sushil.dubey@intel.com>
Signed-off-by: Sushil Dubey <sushil.dubey@intel.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/pyproject_xpu.toml

@@ -69,6 +69,25 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+diffusion = [
+  "PyYAML==6.0.1",
+  "cloudpickle==3.1.2",
+  "diffusers==0.36.0",
+  "imageio==2.36.0",
+  "imageio-ffmpeg==0.5.1",
+  "moviepy>=2.0.0",
+  "opencv-python==4.10.0.84",
+  "remote-pdb==2.1.0",
+  "st_attn==0.0.7 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
+  "runai_model_streamer>=0.15.5",
+  "cache-dit==1.3.0",
+  "addict==2.4.0",
+  "av==16.1.0",
+  "scikit-image==0.25.2",
+  "trimesh>=4.0.0",
+  "xatlas",
+]
+
 tracing = [
   "opentelemetry-api",
   "opentelemetry-exporter-otlp",
@@ -93,6 +112,7 @@ test = [
 dev = ["sglang[test]"]
 
 all = [
+  "sglang[diffusion]",
   "sglang[tracing]",
 ]
 

python/sglang/jit_kernel/diffusion/triton/scale_shift.py

@@ -338,7 +338,7 @@ def fuse_scale_shift_kernel(
     block_l: int = 128,
     block_c: int = 128,
 ):
-    assert x.is_cuda and scale.is_cuda
+    assert (x.is_cuda and scale.is_cuda) or (x.is_xpu and scale.is_xpu)
     assert x.is_contiguous()
 
     B, L, C = x.shape

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -183,17 +183,18 @@ def init_distributed_environment(
     rank: int = 0,
     distributed_init_method: str = "env://",
     local_rank: int = 0,
-    backend: str = "nccl",
+    backend: str | None = None,
     device_id: torch.device | None = None,
     timeout: int | None = None,
 ):
     # Determine the appropriate backend based on the platform
     from sglang.multimodal_gen.runtime.platforms import current_platform
 
-    if backend == "nccl" and not current_platform.is_cuda_alike():
-        # Use gloo backend for non-CUDA platforms (MPS, CPU)
-        backend = "gloo"
-        logger.info("Using gloo backend for %s platform", current_platform.device_name)
+    if backend is None:
+        backend = current_platform.get_torch_distributed_backend_str()
+        logger.info(
+            "Using %s backend for %s platform", backend, current_platform.device_name
+        )
 
     logger.debug(
         "world_size=%d rank=%d local_rank=%d "
@@ -211,13 +212,14 @@ def init_distributed_environment(
             "distributed environment"
         )
 
-        # For MPS and MUSA, don't pass device_id as it doesn't support device indices
+        # For MPS, MUSA, and XPU, don't pass device_id as it doesn't support device indices
         extra_args = (
             {}
             if (
                 current_platform.is_mps()
                 or current_platform.is_musa()
                 or current_platform.is_npu()
+                or current_platform.is_xpu()
             )
             else dict(device_id=device_id)
         )

python/sglang/multimodal_gen/runtime/layers/activation.py

@@ -97,6 +97,9 @@ class NewGELU(CustomOp):
     def forward_cuda(self, *args, **kwargs) -> Any:
         return self.forward_native(*args, **kwargs)
 
+    def forward_xpu(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         c = math.sqrt(2.0 / math.pi)
@@ -112,6 +115,9 @@ class QuickGELU(CustomOp):
     def forward_cuda(self, *args, **kwargs) -> Any:
         return self.forward_native(*args, **kwargs)
 
+    def forward_xpu(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         return x * torch.sigmoid(1.702 * x)

python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py

@@ -0,0 +1,122 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from functools import lru_cache
+
+import torch
+
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+try:
+    from sgl_kernel.flash_attn import flash_attn_varlen_func
+
+    flash_attn_func = flash_attn_varlen_func
+except ImportError as e:
+    raise e
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
+    FlashAttentionMetadataBuilder,
+)
+
+
+class XPUAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.FA
+
+    @staticmethod
+    def get_impl_cls() -> type["XPUAttentionImpl"]:
+        return XPUAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        """XPU backend does not require special metadata."""
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return FlashAttentionMetadataBuilder
+
+
+@lru_cache(maxsize=128)
+def _get_cu_seqlens(device_index: int, bsz: int, seqlen: int) -> torch.Tensor:
+    return torch.arange(
+        0,
+        (bsz + 1) * seqlen,
+        step=seqlen,
+        device=torch.device("xpu", device_index),
+        dtype=torch.int32,
+    )
+
+
+class XPUAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+        *,
+        return_softmax_lse: bool = False,
+    ):
+        bsz, seqlen_q, nheads_q, d = tuple(query.shape)
+        _, seqlen_k, nheads_k, _ = tuple(key.shape)
+
+        max_seqlen_q = seqlen_q
+        max_seqlen_k = seqlen_k
+
+        q_ = query.contiguous().reshape(bsz * seqlen_q, nheads_q, d)
+        k_ = key.contiguous().reshape(bsz * seqlen_k, nheads_k, d)
+        v_ = value.contiguous().reshape(bsz * seqlen_k, nheads_k, d)
+        cu_q = _get_cu_seqlens(q_.device.index, bsz, seqlen_q)
+        cu_k = _get_cu_seqlens(q_.device.index, bsz, seqlen_k)
+
+        out = flash_attn_func(
+            q=q_,
+            k=k_,
+            v=v_,
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+            return_softmax_lse=return_softmax_lse,
+        )
+
+        if return_softmax_lse:
+            out_tensor, softmax_lse = out[:2]
+            result = out_tensor.reshape(bsz, seqlen_q, nheads_q, d)
+            return result, softmax_lse
+
+        result = out.reshape(bsz, seqlen_q, nheads_q, d)
+        return result

python/sglang/multimodal_gen/runtime/layers/elementwise.py

@@ -33,3 +33,8 @@ class MulAdd(CustomOp):
         self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
     ):
         return fuse_scale_shift_kernel(a, b, c, scale_constant=k)
+
+    def forward_xpu(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ):
+        return self.forward_native(a, b, c, k=k)

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -31,7 +31,8 @@ from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
 _is_cuda = current_platform.is_cuda()
 _is_npu = current_platform.is_npu()
 _is_musa = current_platform.is_musa()
-if _is_cuda:
+_is_xpu = current_platform.is_xpu()
+if _is_cuda or _is_xpu:
     from sgl_kernel import fused_add_rmsnorm, rmsnorm
 
 if _is_npu:
@@ -212,6 +213,27 @@ class RMSNorm(CustomOp):
         out = out.view(shape)
         return out
 
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        shape = x.shape
+        x = x.reshape(-1, shape[-1])
+        if residual is not None:
+            residual_shape = residual.shape
+            residual = residual.view(-1, shape[-1])
+
+        if self.variance_size_override is not None:
+            return self.forward_native(x, residual)
+        elif residual is not None:
+            fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon)
+            return x.view(shape), residual.view(residual_shape)
+        else:
+            out = rmsnorm(x, self.weight.data, self.variance_epsilon)
+        out = out.view(shape)
+        return out
+
     def extra_repr(self) -> str:
         return f"hidden_size={self.hidden_size}, eps={self.variance_epsilon}"
 
@@ -410,6 +432,11 @@ class _ScaleResidualNormScaleShift(CustomOp):
         # so we fall back to the native PyTorch implementation.
         return self.forward_native(*args, **kwargs)
 
+    def forward_xpu(self, *args, **kwargs):
+        # XPU does not support CUDA/CUTLASS-based fused kernels yet,
+        # so we fall back to the native PyTorch implementation.
+        return self.forward_native(*args, **kwargs)
+
     def forward_native(
         self,
         residual: torch.Tensor,
@@ -514,6 +541,11 @@ class _NormScaleShift(CustomOp):
         # so we fall back to the native PyTorch implementation.
         return self.forward_native(*args, **kwargs)
 
+    def forward_xpu(self, *args, **kwargs):
+        # XPU does not support CUDA/CUTLASS-based fused kernels yet,
+        # so we fall back to the native PyTorch implementation.
+        return self.forward_native(*args, **kwargs)
+
     def forward_native(
         self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
     ) -> torch.Tensor:

python/sglang/multimodal_gen/runtime/models/encoders/clip.py

@@ -231,7 +231,11 @@ class CLIPAttention(nn.Module):
             key_states = key_states.transpose(1, 2)
             value_states = value_states.transpose(1, 2)
 
-            if current_platform.is_rocm() or current_platform.is_musa():
+            if (
+                current_platform.is_rocm()
+                or current_platform.is_musa()
+                or current_platform.is_xpu()
+            ):
                 # ROCm: Using both is_causal=True and attn_mask causes NaN.
                 # Use is_causal=True alone (padding mask not needed for CLIP
                 # since pooler_output comes from EOS token before padding).

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

@@ -138,9 +138,31 @@ def musa_platform_plugin() -> str | 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,
@@ -157,6 +179,11 @@ def resolve_current_platform_cls_qualname() -> str:
     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:

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

@@ -63,6 +63,7 @@ class PlatformEnum(enum.Enum):
     MPS = enum.auto()
     NPU = enum.auto()
     MUSA = enum.auto()
+    XPU = enum.auto()
     OOT = enum.auto()
     UNSPECIFIED = enum.auto()
 
@@ -283,6 +284,8 @@ class Platform:
             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():
@@ -300,6 +303,8 @@ class Platform:
             return "mccl"
         elif self.is_mps():
             return "gloo"
+        elif self.is_xpu():
+            return "xccl"
         else:
             raise NotImplementedError(
                 "No Accelerators(AMD/NV/MTT GPU, AMD MI instinct accelerators) available"

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

@@ -0,0 +1,196 @@
+# 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()
+
+        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"

python/sglang/multimodal_gen/runtime/utils/common.py

@@ -256,9 +256,12 @@ def is_host_cpu_x86() -> bool:
 
 
 def set_cuda_arch():
-    capability = torch.cuda.get_device_capability()
-    arch = f"{capability[0]}.{capability[1]}"
-    os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}"
+    """Set CUDA architecture for compilation. Only applies to CUDA devices."""
+    if torch.cuda.is_available():
+        capability = torch.cuda.get_device_capability()
+        arch = f"{capability[0]}.{capability[1]}"
+        os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}"
+    # For XPU or other platforms, no arch setting needed
 
 
 # musa

python/sglang/multimodal_gen/runtime/utils/profiler.py

@@ -69,6 +69,9 @@ class SGLDiffusionProfiler:
         if current_platform.is_npu():
             activities.append(torch_npu.profiler.ProfilerActivity.NPU)
 
+        if hasattr(torch, "xpu") and torch.xpu.is_available():
+            activities.append(torch.profiler.ProfilerActivity.XPU)
+
         common_torch_profiler_args = dict(
             activities=activities,
             record_shapes=True,