[ROCm][diffusion] Enable fused qk norm and rope on ROCm

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -32,6 +32,7 @@
 from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
 
 _is_cuda = current_platform.is_cuda()
+_is_rocm = current_platform.is_rocm()
 _is_npu = current_platform.is_npu()
 _is_musa = current_platform.is_musa()
 _is_cpu = current_platform.is_cpu()
@@ -886,10 +887,10 @@ def apply_qk_norm(
     batch_size = q.size(0)
     q_eps = q_norm.variance_epsilon
     k_eps = k_norm.variance_epsilon
-    # Only try fused path on CUDA and when it won't introduce implicit copies.
+    # Only try fused path on CUDA/ROCm and when it won't introduce implicit copies.
     # The in-place kernel needs a real view (no copy), so it also requires contiguity.
     if (
-        _is_cuda
+        (_is_cuda or _is_rocm)
         and allow_inplace
         and (q_eps == k_eps)
         and q.dtype in (torch.float16, torch.bfloat16)
@@ -929,6 +930,7 @@ def apply_qk_norm_with_optional_rope(
     positions: Optional[torch.Tensor] = None,
     position_offset: int = 0,
     allow_inplace: bool = True,
+    allow_strided_qk: bool = False,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """Apply QK RMSNorm and optionally RoPE when a cos/sin cache is provided."""
 
@@ -953,6 +955,7 @@ def apply_qk_norm_with_optional_rope(
         positions=positions,
         position_offset=position_offset,
         allow_inplace=allow_inplace,
+        allow_strided_qk=allow_strided_qk,
     )
 
 
@@ -1049,7 +1052,7 @@ def apply_qk_norm_rope(
 
     if (
         fused_enabled
-        and _is_cuda
+        and (_is_cuda or _is_rocm)
         and not torch.compiler.is_compiling()
         and allow_inplace
         and (q_eps == k_eps)

python/sglang/multimodal_gen/runtime/models/dits/flux_2.py

@@ -661,6 +661,7 @@ def forward(
             cos_sin_cache=cos_sin_cache,
             is_neox=False,
             allow_inplace=True,
+            allow_strided_qk=True,
         )
         hidden_states = self.attn(
             query,