[diffusion] hardware: SiluAndMul/RMSNorm/LayerNorm MUSA implementations (custom ops, 12/N) (#18583)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
Co-authored-by: Qingfu Wen <qingfu.wen@mthreads.com>

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

@@ -55,6 +55,9 @@ class SiluAndMul(CustomOp):
         out = torch_npu.npu_swiglu(x)
         return out
 
+    def forward_musa(self, x: torch.Tensor) -> torch.Tensor:
+        return nn.SwishGLU()(x)
+
 
 @CustomOp.register("gelu_and_mul")
 class GeluAndMul(CustomOp):

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

@@ -54,10 +54,8 @@ class CustomOp(nn.Module):
         return self.forward_native(*args, **kwargs)
 
     def forward_musa(self, *args, **kwargs) -> Any:
-        # XXX (MUSA): MUSA kernels follow the CUDA path by default.
-        # At this stage, sgl-kernel support for MUSA is still under active
-        # development, so we fall back to the PyTorch-native implementation.
-        return self.forward_native(*args, **kwargs)
+        # MUSA kernels follow the CUDA path by default.
+        return self.forward_cuda(*args, **kwargs)
 
     def forward_oot(self, *args, **kwargs) -> Any:
         # By default, we assume that OOT ops are compatible with the

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

@@ -14,12 +14,16 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
 
 _is_cuda = current_platform.is_cuda()
 _is_npu = current_platform.is_npu()
+_is_musa = current_platform.is_musa()
 if _is_cuda:
     from sgl_kernel import fused_add_rmsnorm, rmsnorm
 
 if _is_npu:
     import torch_npu
 
+if _is_musa:
+    from sgl_kernel import fused_add_rmsnorm
+
 from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
 from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
 from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
@@ -165,6 +169,45 @@ class RMSNorm(CustomOp):
         # ROCm builds of sgl-kernel do not expose rmsnorm custom ops yet.
         return self.forward_native(x, residual)
 
+    def _get_weight(self, dtype: torch.dtype) -> torch.Tensor:
+        """Return weight matched to *dtype*.
+
+        MUSA kernels require input and weight to share the same dtype,
+        unlike CUDA kernels which may handle mixed dtypes internally.
+        """
+        weight = self.weight.data
+        if weight.dtype != dtype:
+            weight = weight.to(dtype=dtype)
+        return weight
+
+    def forward_musa(
+        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 requires contiguous inputs.
+            if not x.is_contiguous():
+                x = x.contiguous()
+            if not residual.is_contiguous():
+                residual = residual.contiguous()
+            weight = self._get_weight(x.dtype)
+            fused_add_rmsnorm(x, residual, weight, self.variance_epsilon)
+            return x.view(shape), residual.view(residual_shape)
+        else:
+            weight = self._get_weight(x.dtype)
+            out = F.rms_norm(x, (self.hidden_size,), weight, self.variance_epsilon)
+        out = out.view(shape)
+        return out
+
     def extra_repr(self) -> str:
         s = f"hidden_size={self.weight.data.size(0)}"
         s += f", eps={self.variance_epsilon}"
@@ -254,6 +297,9 @@ class LayerNorm(CustomOp):
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         return self.forward_native(x, residual)
 
+    def forward_musa(self, x: torch.Tensor):
+        return F.layer_norm(x, (self.hidden_size,), self.weight, self.bias, self.eps)
+
     def extra_repr(self) -> str:
         s = f"hidden_size={self.weight.data.size(0)}"
         s += f", eps={self.variance_epsilon}"
@@ -357,6 +403,11 @@ class _ScaleResidualNormScaleShift(CustomOp):
         # so we fall back to the native PyTorch implementation.
         return self.forward_native(*args, **kwargs)
 
+    def forward_musa(self, *args, **kwargs):
+        # MUSA 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,
@@ -456,6 +507,11 @@ class _NormScaleShift(CustomOp):
         # so we fall back to the native PyTorch implementation.
         return self.forward_native(*args, **kwargs)
 
+    def forward_musa(self, *args, **kwargs):
+        # MUSA 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:

python/sglang/multimodal_gen/runtime/utils/profiler.py

@@ -55,7 +55,9 @@ class SGLDiffusionProfiler:
             pass
 
         activities = [torch.profiler.ProfilerActivity.CPU]
-        if torch.cuda.is_available():
+        if torch.cuda.is_available() or (
+            hasattr(torch, "musa") and torch.musa.is_available()
+        ):
             activities.append(torch.profiler.ProfilerActivity.CUDA)
         if current_platform.is_npu():
             activities.append(torch_npu.profiler.ProfilerActivity.NPU)
@@ -127,7 +129,9 @@ class SGLDiffusionProfiler:
             return
         self.has_stopped = True
         logger.info("Stopping Profiler...")
-        if torch.cuda.is_available():
+        if torch.cuda.is_available() or (
+            hasattr(torch, "musa") and torch.musa.is_available()
+        ):
             torch.cuda.synchronize()
         if current_platform.is_npu():
             torch.npu.synchronize()