# 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"