[Llama4] Use strided in-place fused QK RMSNorm to drop a redundant copy (#25089)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

python/sglang/srt/models/llama4.py

@@ -54,6 +54,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     PPProxyTensors,
 )
 from sglang.srt.models.llama import LlamaForCausalLM, LlamaMLP
+from sglang.srt.models.utils import apply_qk_norm
 from sglang.srt.utils import (
     add_prefix,
     fast_topk,
@@ -341,13 +342,27 @@ class Llama4Attention(nn.Module):
                 qk = torch.cat([q_out_unused, k_out_unused], dim=-1)
             del q_view, k_view, q_out_unused, k_out_unused
 
-        if self.qk_norm is not None:
-            # TODO there are still 2 redundant direct_copy_kernel_cuda for this `reshape` and (in attn backend) q.contiguous(), maybe we can fuse them later
-            qk = qk.reshape(-1, self.head_dim).contiguous().bfloat16()
-            qk = self.qk_norm(qk).to(torch.bfloat16)
-            qk = qk.reshape(-1, self.q_size + self.kv_size)
-
-        q, k = qk.split([self.q_size, self.kv_size], dim=-1)
+        if self.qk_norm is not None and _is_cuda:
+            # Strided in-place fused QK RMSNorm reads/writes the qkv buffer
+            # directly via the split q/k views, so the reshape-to-(N, head_dim)
+            # copy is no longer needed. The remaining redundant copy
+            # (`q.contiguous()` inside the attention backend) is unrelated.
+            q, k = qk.split([self.q_size, self.kv_size], dim=-1)
+            q, k = apply_qk_norm(
+                q=q,
+                k=k,
+                q_norm=self.qk_norm,
+                k_norm=self.qk_norm,
+                head_dim=self.head_dim,
+            )
+        else:
+            if self.qk_norm is not None:
+                # NPU/other: qk has been rebuilt via torch.cat after RoPE, so
+                # this reshape is a free view; keep the previous path.
+                qk = qk.reshape(-1, self.head_dim).contiguous().bfloat16()
+                qk = self.qk_norm(qk).to(torch.bfloat16)
+                qk = qk.reshape(-1, self.q_size + self.kv_size)
+            q, k = qk.split([self.q_size, self.kv_size], dim=-1)
 
         # We are applying temperature tuning (https://arxiv.org/abs/2501.19399) to NoPE layers, where
         # the inference-time temperature tuning function is customized to not affect short context