From e26c73c4e92ce787714cfb8b82d4b778638a61b1 Mon Sep 17 00:00:00 2001 From: sushil Dubey Date: Sat, 11 Apr 2026 12:39:02 +0530 Subject: [PATCH] [diffusion] platform: support Intel XPU (#17920) Signed-off-by: sushil.dubey Signed-off-by: Sushil Dubey Co-authored-by: Ma Mingfei --- python/pyproject_xpu.toml | 20 ++ .../diffusion/triton/scale_shift.py | 2 +- .../runtime/distributed/parallel_state.py | 14 +- .../runtime/layers/activation.py | 6 + .../layers/attention/backends/xpu_backend.py | 122 +++++++++++ .../runtime/layers/elementwise.py | 5 + .../runtime/layers/layernorm.py | 34 ++- .../runtime/models/encoders/clip.py | 6 +- .../runtime/platforms/__init__.py | 27 +++ .../runtime/platforms/interface.py | 5 + .../multimodal_gen/runtime/platforms/xpu.py | 196 ++++++++++++++++++ .../multimodal_gen/runtime/utils/common.py | 9 +- .../multimodal_gen/runtime/utils/profiler.py | 3 + 13 files changed, 437 insertions(+), 12 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py create mode 100644 python/sglang/multimodal_gen/runtime/platforms/xpu.py diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index 5a2532036..315d9e850 100644 --- a/python/pyproject_xpu.toml +++ b/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]", ] diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index 1c9ca007d..4c8c93c58 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index 234ea66ec..b3ef8f32b 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/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) ) diff --git a/python/sglang/multimodal_gen/runtime/layers/activation.py b/python/sglang/multimodal_gen/runtime/layers/activation.py index b4d457272..a930c7f95 100644 --- a/python/sglang/multimodal_gen/runtime/layers/activation.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py new file mode 100644 index 000000000..613f2df66 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/elementwise.py b/python/sglang/multimodal_gen/runtime/layers/elementwise.py index a9cafb70e..c990f7f67 100644 --- a/python/sglang/multimodal_gen/runtime/layers/elementwise.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 71f6b2ce3..63dc1fabb 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py b/python/sglang/multimodal_gen/runtime/models/encoders/clip.py index 83fdefd8c..56b4932e5 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/clip.py +++ b/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). diff --git a/python/sglang/multimodal_gen/runtime/platforms/__init__.py b/python/sglang/multimodal_gen/runtime/platforms/__init__.py index fde775378..606839d4f 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/__init__.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 370872190..9bf6e08e0 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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" diff --git a/python/sglang/multimodal_gen/runtime/platforms/xpu.py b/python/sglang/multimodal_gen/runtime/platforms/xpu.py new file mode 100644 index 000000000..68740566f --- /dev/null +++ b/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" diff --git a/python/sglang/multimodal_gen/runtime/utils/common.py b/python/sglang/multimodal_gen/runtime/utils/common.py index b5e3dc098..e2fb8a88d 100644 --- a/python/sglang/multimodal_gen/runtime/utils/common.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py index f9eebc2c4..ed76d67b8 100644 --- a/python/sglang/multimodal_gen/runtime/utils/profiler.py +++ b/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,