[Diffusion] Apply jit qk_norm to flux1 (#17296)

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

@@ -27,11 +27,12 @@ from diffusers.models.normalization import (
 )
 from torch.nn import LayerNorm as LayerNorm
 
+from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm
 from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 
 # from sglang.multimodal_gen.runtime.layers.layernorm import LayerNorm as LayerNorm
-from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, apply_qk_norm
 from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
 from sglang.multimodal_gen.runtime.layers.mlp import MLP
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
@@ -165,6 +166,20 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
         query = query.unflatten(-1, (self.heads, -1))
         key = key.unflatten(-1, (self.heads, -1))
         value = value.unflatten(-1, (self.heads, -1))
+        if (
+            query.is_cuda
+            and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon)
+            and can_use_fused_inplace_qknorm(self.head_dim, query.dtype)
+        ):
+            query, key = apply_qk_norm(
+                q=query,
+                k=key,
+                q_norm=self.norm_q,
+                k_norm=self.norm_k,
+                head_dim=self.head_dim,
+                allow_inplace=True,
+            )
+        else:
             query = self.norm_q(query)
             key = self.norm_k(key)
 
@@ -173,6 +188,23 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
             encoder_key = encoder_key.unflatten(-1, (self.heads, -1))
             encoder_value = encoder_value.unflatten(-1, (self.heads, -1))
 
+            if (
+                encoder_query.is_cuda
+                and (
+                    self.norm_added_q.variance_epsilon
+                    == self.norm_added_k.variance_epsilon
+                )
+                and can_use_fused_inplace_qknorm(self.head_dim, encoder_query.dtype)
+            ):
+                encoder_query, encoder_key = apply_qk_norm(
+                    q=encoder_query,
+                    k=encoder_key,
+                    q_norm=self.norm_added_q,
+                    k_norm=self.norm_added_k,
+                    head_dim=self.head_dim,
+                    allow_inplace=True,
+                )
+            else:
                 encoder_query = self.norm_added_q(encoder_query)
                 encoder_key = self.norm_added_k(encoder_key)