biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Kernel] Use fused rmsnorm for some models like qwen3 series (#17735)

Browse Source 
Signed-off-by: evian <eviantai@u.nus.edu>
Co-authored-by: evian <eviantai@u.nus.edu>
This commit is contained in:
Wanrui Dai
2025-05-07 14:10:02 +08:00
committed by GitHub
parent 1a45a61387
commit f80ae5bdcf
7 changed files with 19 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
vllm/model_executor/models/intern_vit.py
View File

@@ -190,8 +190,8 @@ class InternParallelAttention(nn.Module):
if self.tp_size > 1:
q = tensor_model_parallel_all_gather(q.contiguous())
k = tensor_model_parallel_all_gather(k.contiguous())
q = self.q_norm.forward_native(q)
k = self.k_norm.forward_native(k)
q = self.q_norm(q)
k = self.k_norm(k)
if self.tp_size > 1:
splitter = partial(split_tensor_along_last_dim,
num_partitions=self.tp_size)
@@ -264,10 +264,8 @@ class InternSdpaAttention(nn.Module):
if self.qk_normalization:
B_, N_, H_, D_ = q.shape
q = self.q_norm.forward_native(q.flatten(-2,
-1)).view(B_, N_, H_, D_)
k = self.k_norm.forward_native(k.flatten(-2,
-1)).view(B_, N_, H_, D_)
q = self.q_norm(q.flatten(-2, -1)).view(B_, N_, H_, D_)
k = self.k_norm(k.flatten(-2, -1)).view(B_, N_, H_, D_)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)