[MUSA][Diffusion] Improve wan model inference speed using torch.compile (#25256)

Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>

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,

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
     ):

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:

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]