diff --git a/python/sglang/multimodal_gen/runtime/layers/activation.py b/python/sglang/multimodal_gen/runtime/layers/activation.py index 90e9289fe..b4d457272 100644 --- a/python/sglang/multimodal_gen/runtime/layers/activation.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/layers/custom_op.py b/python/sglang/multimodal_gen/runtime/layers/custom_op.py index fd7355a3f..c4abecb44 100644 --- a/python/sglang/multimodal_gen/runtime/layers/custom_op.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 888824f60..120730556 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py index d01a51f86..19a0ebfd8 100644 --- a/python/sglang/multimodal_gen/runtime/utils/profiler.py +++ b/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()