xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-03-26 01:03:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix RoPE output shape

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-01 08:22:47 +00:00
parent 31d05f7edb
commit 46b31ed98d
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
vllm/model_executor/models/tpu/gemma.py
View File

@ -54,7 +54,8 @@ def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
x_out = torch.cat(torch.chunk(x_out, 2, dim=-1), dim=-2)
x_out = x_out.reshape(x_out.shape[0], x_out.shape[1], x_out.shape[2],
-1).transpose(1, 2)
return x_out
# Reshape the output tensor to the original shape.
return x_out.reshape(x_out.shape[0], x_out.shape[1], -1)
class RMSNorm(torch.nn.Module):