Skip to content

[AMD] Disable PTX-inline-asm diffusion LayerNorm fast path on ROCm (fix FLUX warmup crash) #34352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kangwangamd wants to merge 1 commit into from
+12 −1
Closed

[AMD] Disable PTX-inline-asm diffusion LayerNorm fast path on ROCm (fix FLUX warmup crash) #34352

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,8 +336,19 @@ def is_plain_layer_norm(norm: torch.nn.Module, hidden: int) -> bool:
)


# The fused fast paths below reproduce NVIDIA aten LayerNorm bit-for-bit
# using PTX inline asm (rcp.approx.f32 / div.rn.f32 / rsqrt.approx.f32 with
# "=f" float-register constraints via tl.inline_asm_elementwise). PTX is
# CUDA-only: on ROCm the AMDGPU backend cannot allocate the "f" constraint
# and the Triton JIT aborts at warmup ("couldn't allocate output register
# for constraint 'f'"), which hangs FLUX/diffusion bring-up. Disable the
# fused paths on ROCm so callers fall back to the eager aten chain (the
# module already guarantees bit-exact-or-eager-fallback).
_IS_ROCM = getattr(torch.version, "hip", None) is not None


def _is_bf16_cuda(t: torch.Tensor) -> bool:
return t.is_cuda and t.dtype is torch.bfloat16
return (not _IS_ROCM) and t.is_cuda and t.dtype is torch.bfloat16


def _mod_row_stride(t: torch.Tensor, batch: int, hidden: int) -> int | None:
Expand Down
Loading