From 4949b5fccf2ce53ec50c8a5d7cda51594d6f95fe Mon Sep 17 00:00:00 2001 From: Cao E Date: Wed, 5 Aug 2026 10:12:24 +0800 Subject: [PATCH] [XPU] Add qknorm_rope support for Flux (#30883) Co-authored-by: Chandrakant Khandelwal Co-authored-by: Claude Opus 4.8 (1M context) --- .../runtime/layers/layernorm.py | 39 +++++++++++++++++-- .../runtime/models/dits/flux_2.py | 21 ++++++---- 2 files changed, 50 insertions(+), 10 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 24421b641..87134a55b 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -54,6 +54,9 @@ if USE_AITER: from aiter import rmsnorm2d_fwd as rms_norm from aiter import rmsnorm2d_fwd_with_add as fused_add_rms_norm +if _is_xpu: + from sgl_kernel import fused_inplace_qknorm_rope + if not _is_cpu: from sglang.kernels.ops.diffusion.triton.norm import norm_infer, rms_norm_fn @@ -827,6 +830,7 @@ def apply_qk_norm( 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. + # The in-place kernel needs a real view (no copy), so it also requires contiguity. if ( _is_cuda and allow_inplace @@ -834,6 +838,8 @@ def apply_qk_norm( and q.dtype in (torch.float16, torch.bfloat16) and q_norm.weight.dtype == q.dtype and k_norm.weight.dtype == k.dtype + and q.is_contiguous() + and k.is_contiguous() and can_use_fused_inplace_qknorm(head_dim, q.dtype) ): fused_inplace_qknorm( @@ -848,8 +854,9 @@ def apply_qk_norm( 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) + # reshape (not view) so a non-contiguous q/k (e.g. a chunked qkv view) is handled. + q_out = q_norm(q.reshape(-1, head_dim)).view(q_shape) + k_out = k_norm(k.reshape(-1, head_dim)).view(k_shape) return q_out, k_out @@ -905,7 +912,7 @@ def apply_qk_norm_rope( position_offset: int = 0, allow_inplace: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Apply QK RMSNorm followed by RoPE, fusing both on supported CUDA shapes.""" + """Apply QK RMSNorm followed by RoPE, fusing both on supported CUDA/XPU shapes.""" from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, @@ -985,6 +992,32 @@ def apply_qk_norm_rope( ) return q, k + # TODO: Once CUDA fused_inplace_qknorm_rope supports last-dimension-contiguous q/k, + # merge this path with the CUDA fused qknorm+rope branch. + if ( + _is_xpu + and allow_inplace + and (q_eps == k_eps) + and q.dtype in (torch.float16, torch.bfloat16) + and q_norm.weight.dtype == q.dtype + and k_norm.weight.dtype == k.dtype + and head_dim in (64, 128, 256) + and rope_dim in (32, 64, 128, 256) + ): + fused_inplace_qknorm_rope( + q=q, + k=k, + q_weight=q_norm.weight, + k_weight=k_norm.weight, + cos_sin_cache=cos_sin_cache, + positions=positions, + is_neox=is_neox, + eps=q_eps, + head_dim=head_dim, + rope_dim=rope_dim, + ) + return q, k + q, k = apply_qk_norm( q=q, k=k, 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 1409c37ae..300a65c56 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -52,7 +52,6 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ) from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( NDRotaryEmbedding, - apply_flashinfer_rope_qk_inplace, ) from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, @@ -548,9 +547,7 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): key = key.unflatten(-1, (self.local_heads, -1)) value = value.unflatten(-1, (self.local_heads, -1)) - query = self.norm_q(query) - key = self.norm_k(key) - + cos_sin_cache = None if freqs_cis is not None: cos, sin = freqs_cis cos_sin_cache = torch.cat( @@ -560,9 +557,19 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): ], dim=-1, ) - query, key = apply_flashinfer_rope_qk_inplace( - query, key, cos_sin_cache, is_neox=False - ) + + # QK-norm (+ RoPE) via the shared helper so the fused kernel path is used + # here too — the single-stream block previously ran norm and RoPE as separate ops. + query, key = apply_qk_norm_with_optional_rope( + q=query, + k=key, + q_norm=self.norm_q, + k_norm=self.norm_k, + head_dim=self.head_dim, + cos_sin_cache=cos_sin_cache, + is_neox=False, + allow_inplace=True, + ) hidden_states = self.attn( query, key,