From 5f559ca265d32f43250e2dd81c7f02f60b1cfc70 Mon Sep 17 00:00:00 2001 From: Gurpreet Singh Dhami Date: Wed, 19 Aug 2026 15:31:30 +0000 Subject: [PATCH 1/3] [diffusion] Enable fused qk norm and rope on ROCm Enable fused qk norm and rope paths for double and single transformer blocks. Add .contiguous() to query/key unflattens in single transformer block so the in-place fused kernel receives real views instead of implicit copies. --- python/sglang/multimodal_gen/runtime/layers/layernorm.py | 7 ++++--- python/sglang/multimodal_gen/runtime/models/dits/flux_2.py | 4 ++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index a2202bec44d2..774ca865e442 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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() @@ -890,10 +891,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) @@ -1044,7 +1045,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 d3d6a74320cd..20a0549b22e3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -521,8 +521,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 From 12f48161e23f35ac1cf65fac498c54456c4ea9bb Mon Sep 17 00:00:00 2001 From: Gurpreet Singh Dhami Date: Thu, 20 Aug 2026 18:09:06 +0000 Subject: [PATCH 2/3] [diffusion] Support strided QKV views in fused qk norm rope Plumb allow_strided_qk through apply_qk_norm_with_optional_rope and enable it in the single transformer block so the fused kernel handles strided packed-QKV views natively without forcing contiguous copies. --- python/sglang/multimodal_gen/runtime/layers/layernorm.py | 2 ++ python/sglang/multimodal_gen/runtime/models/dits/flux_2.py | 5 +++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 330ff332d38c..08f63ec7f076 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -930,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.""" @@ -954,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, ) 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 61e4dc3cbeb7..2b7aaf8b199a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/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)).contiguous() - key = key.unflatten(-1, (self.local_heads, -1)).contiguous() + query = query.unflatten(-1, (self.local_heads, -1)) + key = key.unflatten(-1, (self.local_heads, -1)) value = value.unflatten(-1, (self.local_heads, -1)) cos_sin_cache = None @@ -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, From 75b1cd44c71267135a0211335df90bf0df826e0a Mon Sep 17 00:00:00 2001 From: Gurpreet Singh Dhami Date: Thu, 17 Sep 2026 20:01:28 +0000 Subject: [PATCH 3/3] fix(flux2): only enable allow_strided_qk on ROCm Co-Authored-By: Claude --- python/sglang/multimodal_gen/runtime/models/dits/flux_2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e4aaed3d691b..6203f58ba415 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -960,7 +960,7 @@ def forward( freqs_complex=complex_freqs, is_neox=False, allow_inplace=True, - allow_strided_qk=True, + allow_strided_qk=True if current_platform.is_rocm() else False, ) hidden_states = self.attn( query,