[AMD] Keep the PTX-inline-asm diffusion norm fusions off on ROCm (fix FLUX warmup crash) (#34481)

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

python/sglang/kernels/ops/diffusion/norm/layernorm_modulate_triton.py

@@ -52,6 +52,7 @@ from sglang.kernels.ops.diffusion.common.numerics import (
     div_rn_f32,
     round_bf16_to_fp32,
 )
+from sglang.kernels.ops.diffusion.common.platform import is_cuda
 from sglang.srt.utils.custom_op import register_custom_op
 
 
@@ -338,7 +339,11 @@ def is_plain_layer_norm(norm: torch.nn.Module, hidden: int) -> bool:
 
 
 def _is_bf16_cuda(t: torch.Tensor) -> bool:
-    return t.is_cuda and t.dtype is torch.bfloat16
+    # `is_cuda` also covers ROCm, where the inline PTX below cannot compile:
+    # LLVM makes the unusable `=f` constraint a fatal error that kills the
+    # process, so this has to reject before the first launch rather than let
+    # the caller's try/except fall back.
+    return is_cuda() and t.is_cuda and t.dtype is torch.bfloat16
 
 
 def _qk_head_launch_config() -> tuple[int, int]:

python/sglang/kernels/ops/diffusion/norm/rmsnorm_scale_shift_bitexact.py

@@ -60,6 +60,7 @@ from sglang.kernels.ops.diffusion.common.numerics import (
     round_bf16_to_fp32,
     rsqrt_approx_f32,
 )
+from sglang.kernels.ops.diffusion.common.platform import is_cuda
 from sglang.srt.utils.custom_op import register_custom_op
 
 
@@ -171,7 +172,11 @@ def can_use_fused_rmsnorm_scale_shift(
     shift: torch.Tensor,
 ) -> bool:
     return (
-        x.dtype is torch.bfloat16
+        # ROCm cannot compile the inline PTX above: LLVM makes the unusable
+        # `=f` constraint a fatal error that kills the process, so reject
+        # before the first launch rather than rely on the caller's fallback.
+        is_cuda()
+        and x.dtype is torch.bfloat16
         and x.is_cuda
         and x.dim() == 3
         and x.is_contiguous()

test/registered/kernels/ops/diffusion/test_model_fast_paths.py

@@ -37,6 +37,9 @@ import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image
 import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module
 import sglang.multimodal_gen.runtime.models.dits.sana as sana
 from sglang.kernels.ops.diffusion import (
+    can_use_fused_layernorm_modulate,
+    can_use_fused_qk_head_layernorm,
+    can_use_fused_rmsnorm_scale_shift,
     can_use_wan_rmsnorm_silu,
     fused_ltx2_rms_norm_modulate,
     mark_fused_ln_modulate_site,
@@ -49,6 +52,7 @@ from sglang.kernels.ops.diffusion import (
     unmount_ltx2_rms_norm_modulate,
     wan_rmsnorm_silu,
 )
+from sglang.kernels.ops.diffusion.common.platform import is_cuda
 from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
     StableDiffusion3VAEConfig,
 )
@@ -104,6 +108,13 @@ register_amd_ci(est_time=8, suite="nightly-amd-kernel-1-gpu", nightly=True)
 
 pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 
+# The bit-exact LayerNorm/RMSNorm fusions are NVIDIA inline PTX, so their
+# guards reject on ROCm and these sites serve eager there; only the subtests
+# asserting a fused outcome are CUDA-only.
+requires_inline_ptx = pytest.mark.skipif(
+    not is_cuda(), reason="bit-exact norm fusions are NVIDIA PTX"
+)
+
 
 @pytest.fixture(autouse=True)
 def _seed_cuda():
@@ -112,6 +123,21 @@ def _seed_cuda():
     torch.cuda.manual_seed(0)
 
 
+def test_bitexact_norm_guards_follow_platform():
+    # Runs on both lanes, with shapes inside every guard's contract so only the
+    # platform decides: engaged on CUDA, rejected on ROCm.  A fatal LLVM error
+    # there kills the process, so the sites' own try/except cannot be what
+    # catches it -- the guards have to.
+    x = torch.randn(1, 256, 4096, device="cuda", dtype=torch.bfloat16)
+    row = torch.randn(1, 4096, device="cuda", dtype=torch.bfloat16)
+    vec = torch.randn(1, 1, 4096, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(4096, device="cuda", dtype=torch.bfloat16)
+    q = torch.randn(1, 256, 32, 128, device="cuda", dtype=torch.bfloat16)
+    assert can_use_fused_layernorm_modulate(x, row, row) is is_cuda()
+    assert can_use_fused_qk_head_layernorm(q, q) is is_cuda()
+    assert can_use_fused_rmsnorm_scale_shift(x, weight, vec, vec) is is_cuda()
+
+
 # -------------------------------------------------------------------------
 # FLUX.1 -- LayerNorm + adaLN modulate, and the shared-FF GELU site
 # -------------------------------------------------------------------------
@@ -131,6 +157,7 @@ def _flux_site_inputs(shape, chunks, seed):
     return norm, x, parts[0], parts[1]
 
 
+@requires_inline_ptx
 @pytest.mark.parametrize(
     "shape,chunks",
     [
@@ -151,6 +178,7 @@ def test_flux_fused_ln_modulate_is_bit_exact(shape, chunks):
     assert flux._FLUX_LN_MOD.verified
 
 
+@requires_inline_ptx
 def test_flux_norm_modulate_bitexact_supersedes_high_fold():
     # With the quality="high" affine fold mounted, the bit-exact kernel
     # still takes priority, so the site output stays lossless.
@@ -183,6 +211,7 @@ class TestFlux2EagerFusions(CustomTestCase):
         flux2._FLUX2_SWIGLU.verified = False
         flux2._FLUX2_SWIGLU_SIGS.clear()
 
+    @requires_inline_ptx
     def test_norm_modulate_is_bit_exact_across_sequence_lengths(self):
         torch.manual_seed(0)
         hidden = 256
@@ -258,6 +287,7 @@ class TestFlux2EagerFusions(CustomTestCase):
 # -------------------------------------------------------------------------
 
 
+@requires_inline_ptx
 @pytest.mark.parametrize("shape", [(1, 4096, 4096), (2, 301, 4096), (1, 1, 2560)])
 def test_glm_ln_modulate_is_bit_exact(shape):
     # (1, 4096, 4096) is the real GLM-Image image-stream shape (1024^2,
@@ -275,6 +305,7 @@ def test_glm_ln_modulate_is_bit_exact(shape):
     assert not glm_image._GLM_LN_MOD.disabled
 
 
+@requires_inline_ptx
 @pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)])
 def test_glm_qk_head_layernorm_is_bit_exact(shape):
     # (1, 4360, 32, 128) is the real GLM-Image q/k shape (text + image
@@ -297,6 +328,7 @@ def test_glm_qk_head_layernorm_is_bit_exact(shape):
 # -------------------------------------------------------------------------
 
 
+@requires_inline_ptx
 @pytest.mark.parametrize(
     "shape,nmod,transposed",
     [
@@ -346,6 +378,7 @@ def test_sana_fused_ln_modulate_is_bit_exact(shape, nmod, transposed):
 # -------------------------------------------------------------------------
 
 
+@requires_inline_ptx
 @pytest.mark.parametrize("shape", [(1, 4216, 4096), (2, 1140, 4096), (1, 128, 2048)])
 def test_ernie_norm_scale_shift_is_bit_exact(shape):
     # (1, 4216, 4096) is the real ERNIE-Image shape (1024^2 image + text
@@ -501,6 +534,7 @@ def test_ltx2_lossless_compile_keeps_expression_visible_to_inductor(monkeypatch)
     assert torch.equal(out, _ltx2_eager(rms, x, scale, shift, 1e-6))
 
 
+@requires_inline_ptx
 @pytest.mark.parametrize("hidden", [4096, 2048])
 def test_ltx2_mounted_high_uses_fused_kernel(hidden):
     block = nn.Module()