[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()