diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 5017535f7..d7fde1f32 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -366,7 +366,10 @@ class LayerNorm(CustomOp): x = x.view(-1, self.hidden_size) return self.forward_triton(x).view(shape) - @torch.compile(backend="inductor", disable=current_platform.is_npu()) + @torch.compile( + backend="inductor", + disable=current_platform.is_npu() or current_platform.is_rocm(), + ) def forward_native( self, x: torch.Tensor, @@ -566,7 +569,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()) + @torch.compile(disable=current_platform.is_npu() or current_platform.is_rocm()) def forward_native( self, residual: torch.Tensor, @@ -734,7 +737,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()) + @torch.compile(disable=current_platform.is_npu() or current_platform.is_rocm()) def forward_native( self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor ) -> torch.Tensor: @@ -831,7 +834,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()) + @torch.compile(disable=current_platform.is_npu() or current_platform.is_rocm()) def forward_native( self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor ) -> torch.Tensor: