From 52c6e27e7e59405a6db38a871df229cc4a512ac6 Mon Sep 17 00:00:00 2001 From: sushil Dubey Date: Mon, 6 Jul 2026 23:45:57 +0530 Subject: [PATCH] [AMD][diffusion] fix: disable layernorm torch.compile decorator in eager mode on ROCm to avoid memory-access fault (#29673) --- .../sglang/multimodal_gen/runtime/layers/layernorm.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) 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: