From 2dee23a876a02d3221b252c7978453ed08bed147 Mon Sep 17 00:00:00 2001 From: Gurpreet Singh Dhami <143527450+gurpreet-dhami@users.noreply.github.com> Date: Fri, 18 Sep 2026 04:40:32 -0400 Subject: [PATCH] [ROCm][diffusion] Enable fused qk norm and rope on ROCm (#35573) Co-authored-by: jacky.cheng Co-authored-by: Claude --- python/sglang/multimodal_gen/runtime/layers/layernorm.py | 9 ++++++--- .../sglang/multimodal_gen/runtime/models/dits/flux_2.py | 1 + 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 52baf929e..7b2d54c87 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -40,6 +40,7 @@ from sglang.multimodal_gen.runtime.platforms.aiter import USE_AITER 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() @@ -984,10 +985,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) @@ -1028,6 +1029,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.""" @@ -1053,6 +1055,7 @@ def apply_qk_norm_with_optional_rope( positions=positions, position_offset=position_offset, allow_inplace=allow_inplace, + allow_strided_qk=allow_strided_qk, ) @@ -1146,7 +1149,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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index 28314c7dd..6203f58ba 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -960,6 +960,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): freqs_complex=complex_freqs, is_neox=False, allow_inplace=True, + allow_strided_qk=True if current_platform.is_rocm() else False, ) hidden_states = self.attn( query,