[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)
@@ -1049,7 +1050,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

@@ -635,8 +635,8 @@ def forward(
         # Handle the attention logic
         query, key, value = qkv.chunk(3, dim=-1)
 
-        query = query.unflatten(-1, (self.local_heads, -1))
-        key = key.unflatten(-1, (self.local_heads, -1))
+        query = query.unflatten(-1, (self.local_heads, -1)).contiguous()
+        key = key.unflatten(-1, (self.local_heads, -1)).contiguous()
         value = value.unflatten(-1, (self.local_heads, -1))
 
         cos_sin_cache = None