From 17c8a2fa5339208ccf0b94788528295d80d7e358 Mon Sep 17 00:00:00 2001 From: Kai-Hsun Chen Date: Fri, 15 May 2026 10:33:08 -0700 Subject: [PATCH] [Llama4] Use strided in-place fused QK RMSNorm to drop a redundant copy (#25089) Co-authored-by: Claude Opus 4.7 (1M context) --- python/sglang/srt/models/llama4.py | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index dd9fc349d..4b6ff13e4 100644 --- a/python/sglang/srt/models/llama4.py +++ b/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