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, 13:30 UTC
[Perf]Optimize rotary_emb implementation to use Triton operator for improved inference performance (#16457)
Signed-off-by: cynthieye <yexin93@qq.com> Co-authored-by: MagnetoWang <magnetowang@outlook.com>
yexin(叶鑫) committed
Apr 25, 2025 at 14:52 UTC
b22980a1dc894d98de64119af4a580881f16b2f9
2 files changed
+24
-12
cmake/external_projects/vllm_flash_attn.cmake
+1
-1
@@ -38,7 +38,7 @@ else()
38
FetchContent_Declare(
39
vllm-flash-attn
40
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
41
- GIT_TAG e93779c59ba4905e56e5c39dc2c1904ada71fa21
41
+ GIT_TAG 8798f27777fb57f447070301bf33a9f9c607f491
42
GIT_PROGRESS TRUE
43
# Don't share the vllm-flash-attn build between build types
44
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
vllm/model_executor/layers/rotary_embedding.py
+23
-11
@@ -46,20 +46,12 @@ def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
46
return x.flatten(-2)
47
48
49
-def _apply_rotary_emb(
49
+def _apply_rotary_emb_torch(
50
x: torch.Tensor,
51
cos: torch.Tensor,
52
sin: torch.Tensor,
53
is_neox_style: bool,
54
) -> torch.Tensor:
55
- """
56
- Args:
57
- x: [num_tokens, num_heads, head_size]
58
- cos: [num_tokens, head_size // 2]
59
- sin: [num_tokens, head_size // 2]
60
- is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
61
- positional embeddings.
62
- """
55
cos = cos.unsqueeze(-2).to(x.dtype)
56
sin = sin.unsqueeze(-2).to(x.dtype)
57
if is_neox_style:
@@ -75,6 +67,24 @@ def _apply_rotary_emb(
67
return torch.stack((o1, o2), dim=-1).flatten(-2)
68
69
70
+def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
71
+ is_neox_style: bool) -> torch.Tensor:
72
+ """
73
+ Args:
74
+ x: [num_tokens, num_heads, head_size]
75
+ cos: [num_tokens, head_size // 2]
76
+ sin: [num_tokens, head_size // 2]
77
+ is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
78
+ positional embeddings.
79
+ """
80
+ if current_platform.is_cuda_alike():
81
+ from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
82
+ return apply_rotary_emb(x.unsqueeze(0), cos, sin,
83
+ not is_neox_style).squeeze(0)
84
+ else:
85
+ return _apply_rotary_emb_torch(x, cos, sin, is_neox_style)
86
+
87
+
88
@CustomOp.register("rotary_embedding")
89
class RotaryEmbedding(CustomOp):
90
"""Original rotary positional embedding."""
@@ -141,14 +151,16 @@ class RotaryEmbedding(CustomOp):
151
query = query.view(num_tokens, -1, self.head_size)
152
query_rot = query[..., :self.rotary_dim]
153
query_pass = query[..., self.rotary_dim:]
144
- query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
154
+ query_rot = _apply_rotary_emb_torch(query_rot, cos, sin,
155
+ self.is_neox_style)
156
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
157
158
key_shape = key.shape
159
key = key.view(num_tokens, -1, self.head_size)
160
key_rot = key[..., :self.rotary_dim]
161
key_pass = key[..., self.rotary_dim:]
151
- key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
162
+ key_rot = _apply_rotary_emb_torch(key_rot, cos, sin,
163
+ self.is_neox_style)
164
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
165
return query, key
166