From 3bf7e346fc44119f5421b78c36025bedd878c8ef Mon Sep 17 00:00:00 2001 From: Qingfu Wen Date: Sun, 17 May 2026 22:10:24 +0800 Subject: [PATCH] [MUSA][Diffusion] Improve wan model inference speed using torch.compile (#25256) Co-authored-by: R0CKSTAR --- python/sglang/jit_kernel/diffusion/triton/scale_shift.py | 5 +++++ python/sglang/multimodal_gen/runtime/layers/elementwise.py | 6 ++++++ python/sglang/multimodal_gen/runtime/layers/layernorm.py | 3 +++ .../sglang/multimodal_gen/runtime/models/dits/wanvideo.py | 2 +- 4 files changed, 15 insertions(+), 1 deletion(-) diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index fc0746613..8f594b894 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -664,6 +664,11 @@ if current_platform.is_mps(): fuse_scale_shift_kernel = fuse_scale_shift_kernel_native +if current_platform.is_musa(): + from .torch_fallback import fuse_scale_shift_kernel_native + + fuse_scale_shift_kernel = fuse_scale_shift_kernel_native + if current_platform.is_cpu(): from .torch_fallback import ( fuse_scale_shift_kernel_native, diff --git a/python/sglang/multimodal_gen/runtime/layers/elementwise.py b/python/sglang/multimodal_gen/runtime/layers/elementwise.py index 181ed143d..87eec7600 100644 --- a/python/sglang/multimodal_gen/runtime/layers/elementwise.py +++ b/python/sglang/multimodal_gen/runtime/layers/elementwise.py @@ -39,6 +39,12 @@ class MulAdd(CustomOp): ): return self.forward_native(a, b, c, k=k) + @torch.compile + def forward_musa( + self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0 + ): + return self.forward_native(a, b, c, k=k) + def forward_npu( self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0 ): diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index d5e16c488..bbb54b178 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -490,6 +490,7 @@ class _ScaleResidualNormScaleShift(CustomOp): # so we fall back to the native PyTorch implementation. return self.forward_native(*args, **kwargs) + @torch.compile(disable=current_platform.is_npu()) def forward_native( self, residual: torch.Tensor, @@ -631,6 +632,7 @@ class _NormScaleShift(CustomOp): # so we fall back to the native PyTorch implementation. return self.forward_native(*args, **kwargs) + @torch.compile(disable=current_platform.is_npu()) def forward_native( self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor ) -> torch.Tensor: @@ -715,6 +717,7 @@ class _NormTanhMulAdd(CustomOp): # Fallback to native because ROCm does not support CuTeDSL. return self.forward_native(*args, **kwargs) + @torch.compile(disable=current_platform.is_npu()) def forward_native( self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor ) -> torch.Tensor: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 427ab3bc4..89f5c14e9 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -1050,7 +1050,7 @@ class WanTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): ) hidden_states = self.patch_embedding(hidden_states) - hidden_states = hidden_states.flatten(2).transpose(1, 2) + hidden_states = hidden_states.flatten(2).transpose(1, 2).contiguous() # shape is [B, T' * H' * W', C] seq_len_orig = hidden_states.shape[1]