From 2733afe54e4efe142cfdd01efb672eabf603c9a4 Mon Sep 17 00:00:00 2001 From: Chetan Kumar Verma <39086835+ckvermaAI@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:45:07 +0530 Subject: [PATCH] [diffusion] Fix the XPU capability gates that broke the Wan2.2 A14B DiT path (#36825) Co-authored-by: Claude Opus 5 (1M context) --- .../sglang/kernels/ops/diffusion/__init__.py | 2 + .../kernels/ops/diffusion/common/platform.py | 6 + .../layout/wan_causal_cache_triton.py | 7 +- .../diffusion/modulate/scale_shift_triton.py | 5 +- .../scale_residual_norm_scale_shift_triton.py | 160 ++++++++++++++++++ .../runtime/layers/custom_op.py | 3 + .../runtime/layers/elementwise.py | 2 +- .../runtime/layers/layernorm.py | 53 +++++- .../loader/component_loaders/vae_loader.py | 6 +- .../runtime/managers/gpu_worker.py | 40 +++-- .../runtime/models/vaes/wanvae.py | 12 +- .../test/unit/test_vae_loader.py | 30 ++++ 12 files changed, 298 insertions(+), 28 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/norm/scale_residual_norm_scale_shift_triton.py diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 53ddd3c5e..fdc0e2fbf 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/common/platform.py b/python/sglang/kernels/ops/diffusion/common/platform.py index 8479d7266..be476c713 100644 --- a/python/sglang/kernels/ops/diffusion/common/platform.py +++ b/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 diff --git a/python/sglang/kernels/ops/diffusion/layout/wan_causal_cache_triton.py b/python/sglang/kernels/ops/diffusion/layout/wan_causal_cache_triton.py index 0e00c7faa..436b22086 100644 --- a/python/sglang/kernels/ops/diffusion/layout/wan_causal_cache_triton.py +++ b/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 diff --git a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py b/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py index 830a8823c..e791eae5e 100644 --- a/python/sglang/kernels/ops/diffusion/modulate/scale_shift_triton.py +++ b/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], diff --git a/python/sglang/kernels/ops/diffusion/norm/scale_residual_norm_scale_shift_triton.py b/python/sglang/kernels/ops/diffusion/norm/scale_residual_norm_scale_shift_triton.py new file mode 100644 index 000000000..fad7e15a8 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/custom_op.py b/python/sglang/multimodal_gen/runtime/layers/custom_op.py index 50df04146..4151cc2b0 100644 --- a/python/sglang/multimodal_gen/runtime/layers/custom_op.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/elementwise.py b/python/sglang/multimodal_gen/runtime/layers/elementwise.py index 3acc516f4..a53be8c75 100644 --- a/python/sglang/multimodal_gen/runtime/layers/elementwise.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index ab2e58fe9..52baf929e 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 58f796fd4..370bb9ef3 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 563f8fee2..cedd9f3c6 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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.") diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index c3865c5bb..d50d65bcb 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/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() ): diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index 5821a2caf..889c6e83c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/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" )