xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 02:55:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix cuda illegal mem access with Llama4 TP8 + rms_norm custom op (#22701)

Browse Source 
Signed-off-by: Po-Han Huang <pohanh@nvidia.com>
This commit is contained in:
Po-Han Huang (NVIDIA) 2025-08-13 12:21:50 +08:00 committed by GitHub
parent c5830381af
commit 4f0f844b16
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/model_executor/models/llama4.py
View File

@ -224,10 +224,14 @@ class Llama4Attention(nn.Module):
if self.rotary_emb is not None:
q, k = self.rotary_emb(positions, q, k)
if self.qk_norm is not None:
q = q.reshape(-1, self.num_heads, self.head_dim)
# Normalization is applied on the head_dim dimension. The rest of
# the dimensions are collapsed into a single dimension to support
# custom rms_norm cuda kernel.
q = q.reshape(-1, self.head_dim)
q = self.qk_norm(q.float()).reshape(-1, self.q_size).to(q.dtype)
k = k.reshape(-1, self.num_kv_heads, self.head_dim)
k = k.reshape(-1, self.head_dim)
k = self.qk_norm(k.float()).reshape(-1, self.kv_size).to(k.dtype)
# We are applying temperature tuning (https://arxiv.org/abs/2501.19399)