[diffusion] Fix the XPU capability gates that broke the Wan2.2 A14B DiT path (#36825)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/kernels/ops/diffusion/__init__.py

@@ -512,7 +512,9 @@ _EXPORTS: dict[str, str] = {
     "can_use_flux2_gated_resnorm": "norm.flux2_gated_resnorm_jit",
     "flux2_gated_resnorm_raw": "norm.flux2_gated_resnorm_jit",
     "FLYDSL_NORM_MIN_ALIGNED_DIM": "norm.fused_residual_norm_flydsl",
+    "can_use_fused_scale_residual_norm_scale_shift_triton": "norm.scale_residual_norm_scale_shift_triton",
     "flydsl_fused_residual_norm_scale_shift": "norm.fused_residual_norm_flydsl",
+    "fused_scale_residual_norm_scale_shift_triton": "norm.scale_residual_norm_scale_shift_triton",
     "flydsl_norm_scale_shift": "norm.fused_residual_norm_flydsl",
     "apply_group_norm_silu": "norm.group_norm_silu",
     "triton_group_norm_silu": "norm.group_norm_silu_triton",

python/sglang/kernels/ops/diffusion/common/platform.py

@@ -66,6 +66,12 @@ def is_hip() -> bool:
     return current_platform.is_hip()
 
 
+def is_xpu() -> bool:
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    return current_platform.is_xpu()
+
+
 def has_triton() -> bool:
     """True when the live device runs the Triton implementations."""
     return platform_key() in _CUDA_LIKE

python/sglang/kernels/ops/diffusion/layout/wan_causal_cache_triton.py

@@ -117,7 +117,7 @@ def cat_pad_channels_last_3d(
     pw_l, pw_r, ph_t, ph_b, pt_front, pt_back = padding
     if pw_l != pw_r or ph_t != ph_b or pt_back != 0:
         return None
-    if x.dim() != 5 or not x.is_cuda:
+    if x.dim() != 5 or x.device.type not in ("cuda", "xpu"):
         return None
     cache_t = 0
     if cache_x is not None:
@@ -318,7 +318,10 @@ def dup_up3d_add(
         return None
     if repeats <= 0 or repeats & (repeats - 1):
         return None
-    if not main.is_cuda or not src.is_cuda:
+    if main.device.type not in ("cuda", "xpu") or src.device.type not in (
+        "cuda",
+        "xpu",
+    ):
         return None
     if main.dtype != src.dtype or main.device != src.device:
         return None

python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py

@@ -6,6 +6,7 @@ from sglang.kernels.ops.diffusion.common.numerics import mul_rn_f32
 from sglang.kernels.ops.diffusion.common.platform import (
     is_cuda,
     is_hip,
+    is_xpu,
     lazy_fallback,
     select_impl,
 )
@@ -420,8 +421,8 @@ def fuse_scale_shift_kernel(
 
         # Compact scale [B, F, 1, C] -> [B*F, C] (per-frame)
         scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous()
-        if shift.dim() == 4 and is_hip():
-            # ROCm has no fused CUTLASS scale-shift kernel, so this native path
+        if shift.dim() == 4 and (is_hip() or is_xpu()):
+            # ROCm and XPU lack a fused CUTLASS scale-shift kernel, so this path
             # handles the causal Wan / LingBot output AdaLN, which passes a
             # per-frame shift [B, F, 1, C]. Broadcast it across each frame's
             # tokens to per-token [B, L, C] before flattening to [B*L, C],

python/sglang/kernels/ops/diffusion/norm/scale_residual_norm_scale_shift_triton.py

@@ -0,0 +1,160 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+MAX_FUSED_HIDDEN = 8192
+
+
+@triton.jit
+def _scale_residual_norm_scale_shift_kernel(
+    residual_out_ptr,
+    out_ptr,  # outputs, x.dtype
+    residual_ptr,
+    x_ptr,  # inputs, x.dtype
+    gate_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale_ptr,
+    shift_ptr,
+    frame_seqlen,
+    eps,
+    D: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    HAS_AFFINE: tl.constexpr,
+    HAS_GATE: tl.constexpr,
+    GATE_PER_FRAME: tl.constexpr,
+    SCALE_VEC: tl.constexpr,
+    SHIFT_VEC: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_D)
+    mask = cols < D
+    off = row * D + cols
+
+    res = tl.load(residual_ptr + off, mask=mask, other=0.0).to(tl.float32)
+    xv = tl.load(x_ptr + off, mask=mask, other=0.0).to(tl.float32)
+    if HAS_GATE:
+        if GATE_PER_FRAME:
+            gate_off = (row // frame_seqlen) * D + cols
+        else:
+            gate_off = cols
+        g = tl.load(gate_ptr + gate_off, mask=mask, other=0.0)
+        residual_output = res + xv * g
+    else:
+        residual_output = res + xv
+
+    tl.store(
+        residual_out_ptr + off,
+        residual_output.to(residual_out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+    mean = tl.sum(residual_output, axis=0) / D
+    centered = tl.where(mask, residual_output - mean, 0.0)
+    var = tl.sum(centered * centered, axis=0) / D
+    normed = centered * (1.0 / tl.sqrt(var + eps))
+    if HAS_AFFINE:
+        normed = normed * tl.load(weight_ptr + cols, mask=mask, other=0.0).to(
+            tl.float32
+        ) + tl.load(bias_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+
+    if SCALE_VEC:
+        sc = tl.load(scale_ptr + cols, mask=mask, other=0.0)
+    else:
+        sc = tl.load(scale_ptr)
+    if SHIFT_VEC:
+        sh = tl.load(shift_ptr + cols, mask=mask, other=0.0)
+    else:
+        sh = tl.load(shift_ptr)
+    tl.store(
+        out_ptr + off,
+        (normed * (1.0 + sc) + sh).to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def can_use_fused_scale_residual_norm_scale_shift_triton(
+    *,
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: torch.Tensor | int,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+) -> bool:
+    if x.device.type != "xpu" or x.dtype not in (torch.bfloat16, torch.float16):
+        return False
+    if x.dim() != 3 or x.shape[0] != 1 or not x.is_contiguous():
+        return False
+    for operand in (residual, gate, shift, scale, weight, bias):
+        if isinstance(operand, torch.Tensor) and (
+            operand.device != x.device or not operand.is_contiguous()
+        ):
+            return False
+    if residual.shape != x.shape or residual.dtype != x.dtype:
+        return False
+    hidden = x.shape[-1]
+    if hidden > MAX_FUSED_HIDDEN:
+        return False
+    if isinstance(gate, torch.Tensor):
+        if gate.dim() not in (3, 4) or gate.shape[0] != 1 or gate.shape[-1] != hidden:
+            return False
+        if gate.dim() == 3:
+            if gate.shape[1] != 1:
+                return False
+        elif gate.shape[2] != 1 or x.shape[1] % gate.shape[1] != 0:
+            return False
+    elif gate != 1:
+        return False
+    for modulation in (scale, shift):
+        if not isinstance(modulation, torch.Tensor):
+            return False
+        if modulation.numel() not in (1, hidden):
+            return False
+    if (weight is None) != (bias is None):
+        return False
+    if weight is not None and (weight.numel() != hidden or bias.numel() != hidden):
+        return False
+    return True
+
+
+def fused_scale_residual_norm_scale_shift_triton(
+    *,
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: torch.Tensor | int,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    seq_len, hidden = x.shape[1], x.shape[2]
+    residual_output = torch.empty_like(x)
+    out = torch.empty_like(x)
+    has_gate = isinstance(gate, torch.Tensor)
+    gate_per_frame = has_gate and gate.dim() == 4
+    frame_seqlen = seq_len // gate.shape[1] if gate_per_frame else seq_len
+    _scale_residual_norm_scale_shift_kernel[(seq_len,)](
+        residual_output,
+        out,
+        residual,
+        x,
+        gate.reshape(-1) if has_gate else x,
+        weight if weight is not None else x,
+        bias if bias is not None else x,
+        scale.reshape(-1),
+        shift.reshape(-1),
+        frame_seqlen,
+        eps,
+        D=hidden,
+        BLOCK_D=triton.next_power_of_2(hidden),
+        HAS_AFFINE=weight is not None,
+        HAS_GATE=has_gate,
+        GATE_PER_FRAME=gate_per_frame,
+        SCALE_VEC=scale.numel() == hidden,
+        SHIFT_VEC=shift.numel() == hidden,
+        num_warps=8,
+    )
+    return out, residual_output

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

@@ -69,6 +69,9 @@ class CustomOp(nn.Module):
         # PyTorch-native implementation.
         return self.forward_native(*args, **kwargs)
 
+    def forward_xpu(self, *args, **kwargs) -> Any:
+        return self.forward_native(*args, **kwargs)
+
     def dispatch_forward(self) -> Callable:
         if _is_cuda:
             return self.forward_cuda

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

@@ -37,7 +37,7 @@ class MulAdd(CustomOp):
     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)
+        return self.forward_cuda(a, b, c, k=k)
 
     @torch.compile
     def forward_musa(

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

@@ -13,8 +13,10 @@ import torch.nn.functional as F
 
 from sglang.kernels.ops.diffusion import (
     can_use_fused_inplace_qknorm_rope,
+    can_use_fused_scale_residual_norm_scale_shift_triton,
     fuse_scale_shift_kernel,
     fused_inplace_qknorm_rope,
+    fused_scale_residual_norm_scale_shift_triton,
     triton_one_pass_rms_norm,
 )
 from sglang.kernels.ops.diffusion.modulate.scale_shift_triton import (
@@ -521,6 +523,22 @@ class FP32LayerNorm(CustomOp, nn.LayerNorm):
         )
         return output.to(origin_dtype)
 
+    def forward_xpu(self, inputs: torch.Tensor) -> torch.Tensor:
+        def matches_input(param: torch.Tensor | None) -> bool:
+            return param is None or (
+                param.dtype == inputs.dtype and param.device == inputs.device
+            )
+
+        if not (matches_input(self.weight) and matches_input(self.bias)):
+            return self.forward_native(inputs)
+        return F.layer_norm(
+            inputs,
+            self.normalized_shape,
+            self.weight,
+            self.bias,
+            self.eps,
+        )
+
 
 ################################################################################
 # Fused norm kernel
@@ -667,10 +685,37 @@ 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_xpu(
+        self,
+        residual: torch.Tensor,
+        x: torch.Tensor,
+        gate: torch.Tensor | int,
+        shift: torch.Tensor,
+        scale: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.norm_type == "layer":
+            weight = self.norm.weight
+            bias = self.norm.bias
+            if can_use_fused_scale_residual_norm_scale_shift_triton(
+                residual=residual,
+                x=x,
+                gate=gate,
+                shift=shift,
+                scale=scale,
+                weight=weight,
+                bias=bias,
+            ):
+                return fused_scale_residual_norm_scale_shift_triton(
+                    residual=residual,
+                    x=x,
+                    gate=gate,
+                    shift=shift,
+                    scale=scale,
+                    weight=weight,
+                    bias=bias,
+                    eps=self.eps,
+                )
+        return self.forward_native(residual, x, gate, shift, scale)
 
     @torch.compile(disable=current_platform.is_npu() or current_platform.is_rocm())
     def forward_native(

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -137,7 +137,11 @@ def _should_use_channels_last_3d(
     if component_type not in (
         "vae",
         "video_vae",
-    ) or not (current_platform.is_cuda() or current_platform.is_rocm()):
+    ) or not (
+        current_platform.is_cuda()
+        or current_platform.is_rocm()
+        or current_platform.is_xpu()
+    ):
         return False
 
     override = os.getenv(VAE_CHANNELS_LAST_3D_ENV)

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -106,6 +106,24 @@ from sglang.srt.utils.network import NetworkAddress
 logger = init_logger(__name__)
 
 
+def _device_has_allocator_cache() -> bool:
+    return (
+        current_platform.is_cuda()
+        or current_platform.is_rocm()
+        or current_platform.is_xpu()
+    )
+
+
+def _device_module():
+    return torch.get_device_module(current_platform.device_type)
+
+
+def _device_initialized() -> bool:
+    if not _device_has_allocator_cache():
+        return False
+    return _device_module().is_initialized()
+
+
 @dataclass
 class _ExpandedOutputParts:
     tensor_outputs: list[torch.Tensor] = field(default_factory=list)
@@ -256,8 +274,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
 
         released = self._realtime_sessions.release(session_id)
         if released:
-            if torch.cuda.is_initialized():
-                torch.cuda.empty_cache()
+            if _device_initialized():
+                _device_module().empty_cache()
         return OutputBatch(output={"released": released, "session_id": session_id})
 
     def _configure_persistent_torch_compile_cache(self) -> None:
@@ -916,9 +934,9 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
 
         if (
             os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
-            and torch.cuda.is_initialized()
+            and _device_initialized()
         ):
-            torch.cuda.synchronize()
+            _device_module().synchronize()
         start_time = time.perf_counter()
         output_batch.output = [
             self._materialize_frame_output(output, output_batch, req)
@@ -927,9 +945,9 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
         if output_batch.metrics is not None:
             if (
                 os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1"
-                and torch.cuda.is_initialized()
+                and _device_initialized()
             ):
-                torch.cuda.synchronize()
+                _device_module().synchronize()
             output_batch.metrics.record_stage(
                 "GPUWorker.frame_materialize_for_return",
                 time.perf_counter() - start_time,
@@ -1524,11 +1542,7 @@ OOM detected. Possible solutions:
 
 
 def _oom_exceptions():
-    # torch.OutOfMemoryError exists only in some PyTorch builds
-    types = [torch.cuda.OutOfMemoryError]
-    if hasattr(torch, "OutOfMemoryError"):
-        types.append(torch.OutOfMemoryError)
-    return tuple(types)
+    return (torch.OutOfMemoryError,)
 
 
 def run_scheduler_process(
@@ -1575,8 +1589,8 @@ def run_scheduler_process(
         if "scheduler" in locals():
             del scheduler
         gc.collect()
-        if torch.cuda.is_initialized():
-            torch.cuda.empty_cache()
+        if _device_initialized():
+            _device_module().empty_cache()
         if torch.distributed.is_available() and torch.distributed.is_initialized():
             torch.distributed.destroy_process_group()
         logger.info(f"Worker {rank}: Shutdown complete.")

python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py

@@ -54,7 +54,7 @@ from sglang.multimodal_gen.runtime.models.vaes.common import (
 )
 from sglang.multimodal_gen.runtime.platforms import current_platform
 
-if current_platform.is_cuda():
+if current_platform.is_cuda() or current_platform.is_xpu():
     try:
         from sglang.kernels.ops.diffusion import cat_pad_channels_last_3d, dup_up3d_add
     except ImportError:  # pragma: no cover
@@ -74,7 +74,9 @@ first_chunk = contextvars.ContextVar("first_chunk", default=None)
 
 def _channels_last_3d_supported_by_platform() -> bool:
     return hasattr(torch, "channels_last_3d") and (
-        current_platform.is_cuda() or current_platform.is_rocm()
+        current_platform.is_cuda()
+        or current_platform.is_rocm()
+        or current_platform.is_xpu()
     )
 
 
@@ -103,7 +105,7 @@ def _fused_conv_cache_supported(conv: nn.Module, x: torch.Tensor) -> bool:
         cat_pad_channels_last_3d is not None
         and type(conv) is WanCausalConv3d
         and x.dim() == 5
-        and x.is_cuda
+        and x.device.type in ("cuda", "xpu")
         and current_platform.is_amp_supported()
         and _conv3d_weight_is_channels_last_3d(conv.weight)
         and not torch.compiler.is_compiling()
@@ -560,8 +562,8 @@ def residual_up_block_forward(self, x):
         if (
             dup_up3d_add is not None
             and type(shortcut) is DupUp3D
-            and x.is_cuda
-            and x_copy.is_cuda
+            and x.device.type in ("cuda", "xpu")
+            and x_copy.device.type in ("cuda", "xpu")
             and x.dtype == x_copy.dtype
             and not torch.compiler.is_compiling()
         ):

python/sglang/multimodal_gen/test/unit/test_vae_loader.py

@@ -747,10 +747,21 @@ class TestVAELoader(unittest.TestCase):
             patch.dict("os.environ", {}, clear=True),
             patch.object(vae_loader.current_platform, "is_cuda", return_value=False),
             patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
+            patch.object(vae_loader.current_platform, "is_xpu", return_value=False),
         ):
             server_args = _FakeServerArgs(QwenImagePipelineConfig())
             self.assertFalse(_should_use_channels_last_3d(server_args, "vae"))
 
+    def test_channels_last_3d_selected_on_xpu(self):
+        with (
+            patch.dict("os.environ", {}, clear=True),
+            patch.object(vae_loader.current_platform, "is_cuda", return_value=False),
+            patch.object(vae_loader.current_platform, "is_rocm", return_value=False),
+            patch.object(vae_loader.current_platform, "is_xpu", return_value=True),
+        ):
+            server_args = _FakeServerArgs(QwenImagePipelineConfig())
+            self.assertTrue(_should_use_channels_last_3d(server_args, "vae"))
+
     @unittest.skipUnless(
         hasattr(torch, "channels_last_3d"), "channels_last_3d is unavailable"
     )
@@ -763,11 +774,30 @@ class TestVAELoader(unittest.TestCase):
         with (
             patch.object(wanvae.current_platform, "is_cuda", return_value=False),
             patch.object(wanvae.current_platform, "is_rocm", return_value=False),
+            patch.object(wanvae.current_platform, "is_xpu", return_value=False),
         ):
             out = wanvae.match_conv3d_input_format(x, weight)
 
         self.assertIs(out, x)
 
+    @unittest.skipUnless(
+        hasattr(torch, "channels_last_3d"), "channels_last_3d is unavailable"
+    )
+    def test_match_conv3d_input_format_uses_channels_last_3d_on_xpu(self):
+        x = torch.randn(1, 3, 2, 4, 4)
+        weight = torch.randn(3, 3, 1, 1, 1).contiguous(
+            memory_format=torch.channels_last_3d
+        )
+
+        with (
+            patch.object(wanvae.current_platform, "is_cuda", return_value=False),
+            patch.object(wanvae.current_platform, "is_rocm", return_value=False),
+            patch.object(wanvae.current_platform, "is_xpu", return_value=True),
+        ):
+            out = wanvae.match_conv3d_input_format(x, weight)
+
+        self.assertTrue(out.is_contiguous(memory_format=torch.channels_last_3d))
+
     @unittest.skipUnless(
         hasattr(torch, "channels_last_3d"), "channels_last_3d is unavailable"
     )