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 14, 05:30 UTC
Fix cuda illegal mem access with Llama4 TP8 + rms_norm custom op (#22701)
Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
Po-Han Huang (NVIDIA) committed
Aug 13, 2025 at 12:21 UTC
4f0f844b1675419fd2171bc5e981a82386ec552b
1 file changed
+6
-2
vllm/model_executor/models/llama4.py
+6
-2
@@ -224,10 +224,14 @@ class Llama4Attention(nn.Module):
224
225
if self.rotary_emb is not None:
226
q, k = self.rotary_emb(positions, q, k)
227
+
228
if self.qk_norm is not None:
228
- q = q.reshape(-1, self.num_heads, self.head_dim)
229
+ # Normalization is applied on the head_dim dimension. The rest of
230
+ # the dimensions are collapsed into a single dimension to support
231
+ # custom rms_norm cuda kernel.
232
+ q = q.reshape(-1, self.head_dim)
233
q = self.qk_norm(q.float()).reshape(-1, self.q_size).to(q.dtype)
230
- k = k.reshape(-1, self.num_kv_heads, self.head_dim)
234
+ k = k.reshape(-1, self.head_dim)
235
k = self.qk_norm(k.float()).reshape(-1, self.kv_size).to(k.dtype)
236
237
# We are applying temperature tuning (https://arxiv.org/abs/2501.19399)