A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[BugFix] llama4 qknorm should be not shared across head (#16311)
Signed-off-by: Lu Fang <fanglu@fb.com>
Lucia Fang committed
Apr 9, 2025 at 00:59 UTC
ec7da6fcf32fc05efe5d7ba30d01d3d940f12a3c
1 file changed
+7
-12
vllm/model_executor/models/llama4.py
+7
-12
@@ -155,14 +155,8 @@ class Llama4Attention(nn.Module):
155
self.rope_theta = rope_theta
156
self.max_position_embeddings = max_position_embeddings
157
self.n_rep = self.num_heads // self.num_kv_heads
158
- self.q_norm = RMSNorm(
159
- hidden_size=self.q_size,
160
- eps=config.rms_norm_eps,
161
- has_weight=False,
162
- dtype=torch.float32,
163
- ) if self.use_qk_norm else None
164
- self.k_norm = RMSNorm(
165
- hidden_size=self.kv_size,
158
+ self.qk_norm = RMSNorm(
159
+ hidden_size=self.head_dim,
160
eps=config.rms_norm_eps,
161
has_weight=False,
162
dtype=torch.float32,
@@ -226,10 +220,11 @@ class Llama4Attention(nn.Module):
220
221
if self.rotary_emb is not None:
222
q, k = self.rotary_emb(positions, q, k)
229
- if self.q_norm is not None:
230
- q = self.q_norm(q.float()).to(q.dtype)
231
- if self.k_norm is not None:
232
- k = self.k_norm(k.float()).to(k.dtype)
223
+ if self.qk_norm is not None:
224
+ q = q.reshape(-1, self.num_heads, self.head_dim)
225
+ q = self.qk_norm(q.float()).reshape(-1, self.q_size).to(q.dtype)
226
+ k = k.reshape(-1, self.num_kv_heads, self.head_dim)
227
+ k = self.qk_norm(k.float()).reshape(-1, self.kv_size).to(k.dtype)
228
229
# We are applying temperature tuning (https://arxiv.org/abs/2501.19399)
230
# to NoPE layers, where the inference-time temperature tuning function