minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[AMD][diffusion] fix: disable layernorm torch.compile decorator in eager mode on ROCm to avoid memory-access fault (#29673)

Browse Source
This commit is contained in:
sushil Dubey
2026-07-06 11:15:57 -07:00
committed by GitHub
parent 3abdbab9bb
commit 52c6e27e7e
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/layernorm.py
+7 -4
View File
@@ -366,7 +366,10 @@ class LayerNorm(CustomOp):
x = x.view(-1, self.hidden_size) x = x.view(-1, self.hidden_size)
return self.forward_triton(x).view(shape) 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( def forward_native(
self, self,
x: torch.Tensor, x: torch.Tensor,
@@ -566,7 +569,7 @@ class _ScaleResidualNormScaleShift(CustomOp):
# so we fall back to the native PyTorch implementation. # so we fall back to the native PyTorch implementation.
return self.forward_native(*args, **kwargs) 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( def forward_native(
self, self,
residual: torch.Tensor, residual: torch.Tensor,
@@ -734,7 +737,7 @@ class _NormScaleShift(CustomOp):
# so we fall back to the native PyTorch implementation. # so we fall back to the native PyTorch implementation.
return self.forward_native(*args, **kwargs) 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( def forward_native(
self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
) -> torch.Tensor: ) -> torch.Tensor:
@@ -831,7 +834,7 @@ class _NormTanhMulAdd(CustomOp):
# Fallback to native because ROCm does not support CuTeDSL. # Fallback to native because ROCm does not support CuTeDSL.
return self.forward_native(*args, **kwargs) 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( def forward_native(
self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
) -> torch.Tensor: ) -> torch.Tensor: