[XPU] Add qknorm_rope support for Flux (#30883)

Co-authored-by: Chandrakant Khandelwal <Chandrakant.Khandelwal@intel.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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)
+    # reshape (not view) so a non-contiguous q/k (e.g. a chunked qkv view) is handled.
-    k_out = k_norm(k.view(-1, head_dim)).view(k_shape)
+    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,

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)
+        cos_sin_cache = None
-        key = self.norm_k(key)
-
         if freqs_cis is not None:
             cos, sin = freqs_cis
             cos_sin_cache = torch.cat(
@@ -560,8 +557,18 @@ 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,