From 64946679b5db7cca46a1138f0fc8f473086caaea Mon Sep 17 00:00:00 2001 From: Minglei Zhu Date: Sun, 18 Jan 2026 22:33:34 -0800 Subject: [PATCH] [Perf] fuse q, k norm for Flux2Attention (#17241) Co-authored-by: Minglei Zhu --- .../multimodal_gen/runtime/layers/layernorm.py | 14 +++++++++++--- .../multimodal_gen/runtime/models/dits/flux_2.py | 8 +++----- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index a9f5811e2..4bd9368ee 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -443,6 +443,15 @@ def apply_qk_norm( """ batch_size = q.size(0) + q_shape = q.shape + k_shape = k.shape + + # Ensure contiguous for view operation + if not q.is_contiguous(): + q = q.contiguous() + if not k.is_contiguous(): + k = k.contiguous() + 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. @@ -460,7 +469,8 @@ def apply_qk_norm( head_dim=head_dim, eps=q_eps, ) - return q, k + # Inplace kernel modifies q, k - return with original shape + return q.view(q_shape), k.view(k_shape) # Fallback for AMD/ROCm: apply RMSNorm separately to q and k import warnings @@ -469,8 +479,6 @@ def apply_qk_norm( "Fused QK-norm not available, using RMSNorm fallback", stacklevel=2, ) - q_shape = q.shape - k_shape = k.shape q_out = q_norm(q.view(-1, head_dim)).view(q_shape) k_out = k_norm(k.view(-1, head_dim)).view(k_shape) return q_out, k_out 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 a157b88d3..585c30cb5 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -22,7 +22,7 @@ from diffusers.models.normalization import AdaLayerNormContinuous 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 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.rotary_embedding import ( NDRotaryEmbedding, @@ -196,8 +196,7 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): key = key.unflatten(-1, (self.heads, -1)) value = value.unflatten(-1, (self.heads, -1)) - query = self.norm_q(query) - key = self.norm_k(key) + query, key = apply_qk_norm(query, key, self.norm_q, self.norm_k, self.head_dim) if self.added_kv_proj_dim is not None: encoder_query = encoder_query.unflatten(-1, (self.heads, -1)) @@ -340,8 +339,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): key = key.unflatten(-1, (self.heads, -1)) value = value.unflatten(-1, (self.heads, -1)) - query = self.norm_q(query) - key = self.norm_k(key) + query, key = apply_qk_norm(query, key, self.norm_q, self.norm_k, self.head_dim) if freqs_cis is not None: cos, sin = freqs_cis