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

[kangwangamd]

Motivation

On ROCm (MI300X gfx942), bringing up FLUX.1-dev on the native SGLang diffusion backend crashes during warmup:

[aiter] import [module_rmsnorm_quant] ...
error: couldn't allocate output register for constraint 'f'

The server never becomes ready (ready to roll never prints, /health_generate = 503). In CI this is multimodal-gen-test-1-gpu-amd* (flux_image_t2i) timing out at the 90-min step limit (e.g. run 31327057626, 08/09; run 31415663809, 08/10). Deterministic and chronic. Root-cause detail in #34351.

Root cause

The diffusion fast-path helpers embed NVIDIA PTX inline asm with float-register constraints ("=f") via tl.inline_asm_elementwise:

• numerics.py: rcp.approx.f32, div.rn.f32, rsqrt.approx.f32 / mul.rn.f32 (all constraints="=f,...")
• layernorm_modulate.py _rcp4: rcp.approx.f32 / fma.rn.f32 / sub.ftz.f32

PTX is CUDA-only; the AMDGPU backend can't allocate the 'f' constraint, so the Triton JIT aborts at the first warmup forward. These helpers exist to reproduce NVIDIA aten LayerNorm bit-for-bit (SASS-level), so porting the PTX would break the documented contract. The module already guarantees bit-exact-or-eager-fallback (torch.equal gate + can_use_*), so the correct AMD behavior is to take the eager path.

Fix

Gate the fused fast paths off on ROCm: _is_bf16_cuda() (used by both can_use_fused_layernorm_modulate and can_use_fused_qk_head_layernorm) now returns False when torch.version.hip is not None, so ROCm falls back to the eager aten LayerNorm. CUDA behavior is unchanged.

Test plan (offline, real hardware)

• HW: MI300X (gfx942), image rocm/sgl-dev:v0.5.17-rocm720-mi30x-20260809.
• Before: sglang serve --model-path black-forest-labs/FLUX.1-dev --model-type diffusion --num-gpus 1 → couldn't allocate output register for constraint 'f' at warmup, health 503, never ready.
• After (this patch): warmup completes (1024x1024 native FLUX denoise), "The server is fired up and ready to roll!", /health_generate = 200. No allocate output register error.
• CUDA path untouched (guard is ROCm-only).

Fixes #34351

cc @BBuf

---

CI States

Latest PR Test (Base): ❌ Run #31456475803
Latest PR Test (Extra): ❌ Run #31456475694

[kangwangamd]

Closing in favour of #34481, which supersedes this. That PR covers the same ROCm PTX-inline-asm problem but more completely: it gates the fast paths in numerics.py and in both norm kernels (layernorm-modulate and the bit-exact rmsnorm scale/shift), and adds a test, where this PR only guarded layernorm_modulate.py.

For the record, the fix here was hardware-verified on MI300X (gfx942) with image rocm/sgl-dev:v0.5.17-rocm720-mi30x-20260809: before, FLUX warmup aborts with couldn't allocate output register for constraint 'f' and health stays 503; after, warmup reaches 100% and the server is ready. Same evidence should apply to #34481's wider guard.

[michaelzhang-ai]

Thanks @kangwangamd — carried your MI300X evidence into #34481 as corroboration of the same root cause.

One note since #34481 shrank after you closed this: it no longer touches numerics.py, and instead reuses the is_cuda() predicate try_fused_scaled_residual_add_exact already applies to its own mul_rn_f32 kernel. Coverage is the same and checkable — exactly three modules consume the PTX helpers, and all three are now gated. The LayerNorm guard is the shape you had here; the additions are the RMSNorm/LTX-2 half plus test skips for the AMD nightly lane.