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, 11:00 UTC
Make key optional for rotary embedding (#17566)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
May 7, 2025 at 00:11 UTC
98c89e16ff834f1c9f1c465e7342c5353e0eb627
10 files changed
+221
-151
csrc/cpu/pos_encoding.cpp
+24
-15
@@ -9,7 +9,8 @@ void rotary_embedding_impl(
9
scalar_t* __restrict__ query, /// [batch_size, seq_len, num_heads,
10
/// head_size] or [num_tokens, num_heads,
11
/// head_size]
12
- scalar_t* __restrict__ key, // [batch_size, seq_len, num_kv_heads,
12
+ scalar_t* __restrict__ key, // nullptr (optional) or
13
+ // [batch_size, seq_len, num_kv_heads,
14
// head_size] or [num_tokens, num_kv_heads,
15
// head_size]
16
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
@@ -85,10 +86,13 @@ void rotary_embedding_impl(
86
compute_loop(token_head, cache_ptr, query);
87
}
88
88
- for (int i = 0; i < num_kv_heads; ++i) {
89
- const int head_idx = i;
90
- const int64_t token_head = token_idx * key_stride + head_idx * head_size;
91
- compute_loop(token_head, cache_ptr, key);
89
+ if (key != nullptr) {
90
+ for (int i = 0; i < num_kv_heads; ++i) {
91
+ const int head_idx = i;
92
+ const int64_t token_head =
93
+ token_idx * key_stride + head_idx * head_size;
94
+ compute_loop(token_head, cache_ptr, key);
95
+ }
96
}
97
}
98
}
@@ -100,7 +104,8 @@ void rotary_embedding_gptj_impl(
104
scalar_t* __restrict__ query, /// [batch_size, seq_len, num_heads,
105
/// head_size] or [num_tokens, num_heads,
106
/// head_size]
103
- scalar_t* __restrict__ key, // [batch_size, seq_len, num_kv_heads,
107
+ scalar_t* __restrict__ key, // nullptr (optional) or
108
+ // [batch_size, seq_len, num_kv_heads,
109
// head_size] or [num_tokens, num_kv_heads,
110
// head_size]
111
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
@@ -138,6 +143,10 @@ void rotary_embedding_gptj_impl(
143
}
144
}
145
146
+ if (key == nullptr) {
147
+ return;
148
+ }
149
+
150
#pragma omp parallel for collapse(2)
151
for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
152
for (int i = 0; i < num_kv_heads; ++i) {
@@ -168,13 +177,13 @@ void rotary_embedding_gptj_impl(
177
}; // namespace
178
179
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
171
- torch::Tensor& key, int64_t head_size,
180
+ std::optional<torch::Tensor> key, int64_t head_size,
181
torch::Tensor& cos_sin_cache, bool is_neox) {
182
int num_tokens = positions.numel();
183
int rot_dim = cos_sin_cache.size(1);
184
int num_heads = query.size(-1) / head_size;
176
- int num_kv_heads = key.size(-1) / head_size;
177
- int64_t key_stride = key.stride(-2);
185
+ int num_kv_heads = key.has_value() ? key->size(-1) / head_size : num_heads;
186
+ int64_t key_stride = key.has_value() ? key->stride(-2) : 0;
187
int64_t query_stride = query.stride(-2);
188
189
VLLM_DISPATCH_FLOATING_TYPES(
@@ -183,15 +192,15 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
192
if (is_neox) {
193
rotary_embedding_impl(
194
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
186
- key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(),
187
- rot_dim, query_stride, key_stride, num_heads, num_kv_heads,
188
- head_size, num_tokens);
195
+ key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
196
+ cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
197
+ key_stride, num_heads, num_kv_heads, head_size, num_tokens);
198
} else {
199
rotary_embedding_gptj_impl(
200
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
192
- key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(),
193
- rot_dim, query_stride, key_stride, num_heads, num_kv_heads,
194
- head_size, num_tokens);
201
+ key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
202
+ cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
203
+ key_stride, num_heads, num_kv_heads, head_size, num_tokens);
204
}
205
206
CPU_KERNEL_GUARD_OUT(rotary_embedding_impl)
csrc/cpu/torch_bindings.cpp
+1
-1
@@ -117,7 +117,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
117
// Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
118
ops.def(
119
"rotary_embedding(Tensor positions, Tensor! query,"
120
- " Tensor! key, int head_size,"
120
+ " Tensor!? key, int head_size,"
121
" Tensor cos_sin_cache, bool is_neox) -> ()");
122
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
123
csrc/ops.h
+4
-4
@@ -86,13 +86,13 @@ void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
86
std::optional<torch::Tensor> residual);
87
88
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
89
- torch::Tensor& key, int64_t head_size,
89
+ std::optional<torch::Tensor> key, int64_t head_size,
90
torch::Tensor& cos_sin_cache, bool is_neox);
91
92
void batched_rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
93
- torch::Tensor& key, int64_t head_size,
94
- torch::Tensor& cos_sin_cache, bool is_neox,
95
- int64_t rot_dim,
93
+ std::optional<torch::Tensor> key,
94
+ int64_t head_size, torch::Tensor& cos_sin_cache,
95
+ bool is_neox, int64_t rot_dim,
96
torch::Tensor& cos_sin_cache_offsets);
97
98
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
csrc/pos_encoding_kernels.cu
+55
-43
@@ -38,7 +38,8 @@ inline __device__ void apply_rotary_embedding(
38
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
39
// head_size] or [num_tokens, num_heads,
40
// head_size]
41
- scalar_t* __restrict__ key, // [batch_size, seq_len, num_kv_heads,
41
+ scalar_t* __restrict__ key, // nullptr or
42
+ // [batch_size, seq_len, num_kv_heads,
43
// head_size] or [num_tokens, num_kv_heads,
44
// head_size]
45
const scalar_t* cache_ptr, const int head_size, const int num_heads,
@@ -57,13 +58,15 @@ inline __device__ void apply_rotary_embedding(
58
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
59
}
60
60
- const int nk = num_kv_heads * embed_dim;
61
- for (int i = threadIdx.x; i < nk; i += blockDim.x) {
62
- const int head_idx = i / embed_dim;
63
- const int64_t token_head = token_idx * key_stride + head_idx * head_size;
64
- const int rot_offset = i % embed_dim;
65
- apply_token_rotary_embedding<scalar_t, IS_NEOX>(
66
- key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
61
+ if (key != nullptr) {
62
+ const int nk = num_kv_heads * embed_dim;
63
+ for (int i = threadIdx.x; i < nk; i += blockDim.x) {
64
+ const int head_idx = i / embed_dim;
65
+ const int64_t token_head = token_idx * key_stride + head_idx * head_size;
66
+ const int rot_offset = i % embed_dim;
67
+ apply_token_rotary_embedding<scalar_t, IS_NEOX>(
68
+ key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
69
+ }
70
}
71
}
72
@@ -74,7 +77,8 @@ __global__ void rotary_embedding_kernel(
77
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
78
// head_size] or [num_tokens, num_heads,
79
// head_size]
77
- scalar_t* __restrict__ key, // [batch_size, seq_len, num_kv_heads,
80
+ scalar_t* __restrict__ key, // nullptr or
81
+ // [batch_size, seq_len, num_kv_heads,
82
// head_size] or [num_tokens, num_kv_heads,
83
// head_size]
84
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
@@ -98,7 +102,8 @@ __global__ void batched_rotary_embedding_kernel(
102
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
103
// head_size] or [num_tokens, num_heads,
104
// head_size]
101
- scalar_t* __restrict__ key, // [batch_size, seq_len, num_kv_heads,
105
+ scalar_t* __restrict__ key, // nullptr or
106
+ // [batch_size, seq_len, num_kv_heads,
107
// head_size] or [num_tokens, num_kv_heads,
108
// head_size]
109
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
@@ -127,10 +132,12 @@ void rotary_embedding(
132
// [num_tokens, num_heads * head_size] or
133
// [batch_size, seq_len, num_heads, head_size] or
134
// [num_tokens, num_heads, head_size]
130
- torch::Tensor& key, // [batch_size, seq_len, num_kv_heads * head_size] or
131
- // [num_tokens, num_kv_heads * head_size] or
132
- // [batch_size, seq_len, num_heads, head_size] or
133
- // [num_tokens, num_heads, head_size]
135
+ std::optional<torch::Tensor> key,
136
+ // null or
137
+ // [batch_size, seq_len, num_kv_heads * head_size] or
138
+ // [num_tokens, num_kv_heads * head_size] or
139
+ // [batch_size, seq_len, num_heads, head_size] or
140
+ // [num_tokens, num_heads, head_size]
141
int64_t head_size,
142
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
143
bool is_neox) {
@@ -138,40 +145,40 @@ void rotary_embedding(
145
int64_t num_tokens = positions.numel();
146
int positions_ndim = positions.dim();
147
141
- // Make sure num_tokens dim is consistent across positions, query, and key.
148
+ // Make sure num_tokens dim is consistent across positions, query, and key
149
TORCH_CHECK(
150
positions_ndim == 1 || positions_ndim == 2,
151
"positions must have shape [num_tokens] or [batch_size, seq_len]");
152
if (positions_ndim == 1) {
146
- TORCH_CHECK(
147
- query.size(0) == positions.size(0) && key.size(0) == positions.size(0),
148
- "query, key and positions must have the same number of tokens");
153
+ TORCH_CHECK(query.size(0) == positions.size(0) &&
154
+ (!key.has_value() || key->size(0) == positions.size(0)),
155
+ "query, key and positions must have the same number of tokens");
156
}
157
if (positions_ndim == 2) {
158
TORCH_CHECK(
159
query.size(0) == positions.size(0) &&
153
- key.size(0) == positions.size(0) &&
160
+ (!key.has_value() || key->size(0) == positions.size(0)) &&
161
query.size(1) == positions.size(1) &&
155
- key.size(1) == positions.size(1),
162
+ (!key.has_value() || key->size(1) == positions.size(1)),
163
"query, key and positions must have the same batch_size and seq_len");
164
}
165
166
// Make sure head_size is valid for query and key
167
// hidden_size = num_heads * head_size
168
int query_hidden_size = query.numel() / num_tokens;
162
- int key_hidden_size = key.numel() / num_tokens;
169
+ int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
170
TORCH_CHECK(query_hidden_size % head_size == 0);
171
TORCH_CHECK(key_hidden_size % head_size == 0);
172
173
// Make sure query and key have consistent number of heads
174
int num_heads = query_hidden_size / head_size;
168
- int num_kv_heads = key_hidden_size / head_size;
175
+ int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
176
TORCH_CHECK(num_heads % num_kv_heads == 0);
177
178
int rot_dim = cos_sin_cache.size(1);
179
int seq_dim_idx = positions_ndim - 1;
180
int64_t query_stride = query.stride(seq_dim_idx);
174
- int64_t key_stride = key.stride(seq_dim_idx);
181
+ int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
182
183
dim3 grid(num_tokens);
184
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
@@ -181,15 +188,16 @@ void rotary_embedding(
188
if (is_neox) {
189
vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
190
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
184
- key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(), rot_dim,
185
- query_stride, key_stride, num_heads, num_kv_heads, head_size);
191
+ key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
192
+ cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
193
+ num_heads, num_kv_heads, head_size);
194
} else {
195
vllm::rotary_embedding_kernel<scalar_t, false>
196
<<<grid, block, 0, stream>>>(
197
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
190
- key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(),
191
- rot_dim, query_stride, key_stride, num_heads, num_kv_heads,
192
- head_size);
198
+ key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
199
+ cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
200
+ key_stride, num_heads, num_kv_heads, head_size);
201
}
202
});
203
}
@@ -204,10 +212,12 @@ void batched_rotary_embedding(
212
// [num_tokens, num_heads * head_size] or
213
// [batch_size, seq_len, num_heads, head_size] or
214
// [num_tokens, num_heads, head_size]
207
- torch::Tensor& key, // [batch_size, seq_len, num_kv_heads * head_size] or
208
- // [num_tokens, num_kv_heads * head_size] or
209
- // [batch_size, seq_len, num_heads, head_size] or
210
- // [num_tokens, num_heads, head_size]
215
+ std::optional<torch::Tensor>
216
+ key, // null or
217
+ // [batch_size, seq_len, num_kv_heads * head_size] or
218
+ // [num_tokens, num_kv_heads * head_size] or
219
+ // [batch_size, seq_len, num_heads, head_size] or
220
+ // [num_tokens, num_heads, head_size]
221
int64_t head_size,
222
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
223
bool is_neox, int64_t rot_dim,
@@ -221,38 +231,38 @@ void batched_rotary_embedding(
231
"cos_sin_cache_offsets");
232
233
int positions_ndim = positions.dim();
224
- // Make sure num_tokens dim is consistent across positions, query, and key.
234
+ // Make sure num_tokens dim is consistent across positions, query, and key
235
TORCH_CHECK(
236
positions_ndim == 1 || positions_ndim == 2,
237
"positions must have shape [num_tokens] or [batch_size, seq_len]");
238
if (positions_ndim == 1) {
229
- TORCH_CHECK(
230
- query.size(0) == positions.size(0) && key.size(0) == positions.size(0),
231
- "query, key and positions must have the same number of tokens");
239
+ TORCH_CHECK(query.size(0) == positions.size(0) &&
240
+ (!key.has_value() || key->size(0) == positions.size(0)),
241
+ "query, key and positions must have the same number of tokens");
242
}
243
if (positions_ndim == 2) {
244
TORCH_CHECK(
245
query.size(0) == positions.size(0) &&
236
- key.size(0) == positions.size(0) &&
246
+ (!key.has_value() || key->size(0) == positions.size(0)) &&
247
query.size(1) == positions.size(1) &&
238
- key.size(1) == positions.size(1),
248
+ (!key.has_value() || key->size(1) == positions.size(1)),
249
"query, key and positions must have the same batch_size and seq_len");
250
}
251
252
// Make sure head_size is valid for query and key
253
int query_hidden_size = query.numel() / num_tokens;
244
- int key_hidden_size = key.numel() / num_tokens;
254
+ int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
255
TORCH_CHECK(query_hidden_size % head_size == 0);
256
TORCH_CHECK(key_hidden_size % head_size == 0);
257
258
// Make sure query and key have concistent number of heads
259
int num_heads = query_hidden_size / head_size;
250
- int num_kv_heads = key_hidden_size / head_size;
260
+ int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
261
TORCH_CHECK(num_heads % num_kv_heads == 0);
262
263
int seq_dim_idx = positions_ndim - 1;
264
int64_t query_stride = query.stride(seq_dim_idx);
255
- int64_t key_stride = key.stride(seq_dim_idx);
265
+ int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
266
267
dim3 grid(num_tokens);
268
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
@@ -263,14 +273,16 @@ void batched_rotary_embedding(
273
vllm::batched_rotary_embedding_kernel<scalar_t, true>
274
<<<grid, block, 0, stream>>>(
275
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
266
- key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(),
276
+ key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
277
+ cos_sin_cache.data_ptr<scalar_t>(),
278
cos_sin_cache_offsets.data_ptr<int64_t>(), rot_dim, query_stride,
279
key_stride, num_heads, num_kv_heads, head_size);
280
} else {
281
vllm::batched_rotary_embedding_kernel<scalar_t, false>
282
<<<grid, block, 0, stream>>>(
283
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
273
- key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(),
284
+ key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
285
+ cos_sin_cache.data_ptr<scalar_t>(),
286
cos_sin_cache_offsets.data_ptr<int64_t>(), rot_dim, query_stride,
287
key_stride, num_heads, num_kv_heads, head_size);
288
}
csrc/torch_bindings.cpp
+2
-2
@@ -176,7 +176,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
176
// Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
177
ops.def(
178
"rotary_embedding(Tensor positions, Tensor! query,"
179
- " Tensor! key, int head_size,"
179
+ " Tensor!? key, int head_size,"
180
" Tensor cos_sin_cache, bool is_neox) -> ()");
181
ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
182
@@ -184,7 +184,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
184
// (supports multiple loras).
185
ops.def(
186
"batched_rotary_embedding(Tensor positions, Tensor! query,"
187
- " Tensor! key, int head_size,"
187
+ " Tensor!? key, int head_size,"
188
" Tensor cos_sin_cache, bool is_neox,"
189
" int rot_dim,"
190
" Tensor cos_sin_cache_offsets) -> ()");
tests/kernels/core/test_pos_encoding.py
+34
-15
@@ -21,6 +21,7 @@ SEEDS = [0]
21
CUDA_DEVICES = [
22
f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
23
]
24
+USE_KEY = [True, False]
25
26
27
def _get_flat_tensor_shape(batch_size: int, seq_len: int, num_heads: int,
@@ -46,6 +47,7 @@ TENSORS_SHAPES_FN = [_get_batch_tensor_shape, _get_flat_tensor_shape]
47
@pytest.mark.parametrize("dtype", DTYPES)
48
@pytest.mark.parametrize("seed", SEEDS)
49
@pytest.mark.parametrize("device", CUDA_DEVICES)
50
+@pytest.mark.parametrize("use_key", USE_KEY)
51
@torch.inference_mode()
52
def test_rotary_embedding(
53
is_neox_style: bool,
@@ -58,6 +60,7 @@ def test_rotary_embedding(
60
dtype: torch.dtype,
61
seed: int,
62
device: str,
63
+ use_key: bool,
64
max_position: int = 8192,
65
base: int = 10000,
66
) -> None:
@@ -74,7 +77,7 @@ def test_rotary_embedding(
77
positions = torch.randint(0, max_position, (batch_size, seq_len))
78
query_shape = tensor_shape_fn(batch_size, seq_len, num_heads, head_size)
79
query = torch.randn(query_shape, dtype=dtype)
77
- key = torch.randn_like(query)
80
+ key = torch.randn_like(query) if use_key else None
81
82
# NOTE(woosuk): The reference implementation should be executed first
83
# because the custom kernel is in-place.
@@ -85,10 +88,14 @@ def test_rotary_embedding(
88
ref_query,
89
atol=get_default_atol(out_query),
90
rtol=get_default_rtol(out_query))
88
- torch.testing.assert_close(out_key,
89
- ref_key,
90
- atol=get_default_atol(out_key),
91
- rtol=get_default_rtol(out_key))
91
+ if use_key:
92
+ torch.testing.assert_close(out_key,
93
+ ref_key,
94
+ atol=get_default_atol(out_key),
95
+ rtol=get_default_rtol(out_key))
96
+ else:
97
+ assert ref_key is None and out_key is None, \
98
+ "expected returned key to be None"
99
100
101
@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
@@ -101,6 +108,7 @@ def test_rotary_embedding(
108
@pytest.mark.parametrize("dtype", DTYPES)
109
@pytest.mark.parametrize("seed", SEEDS)
110
@pytest.mark.parametrize("device", CUDA_DEVICES)
111
+@pytest.mark.parametrize("use_key", USE_KEY)
112
@torch.inference_mode()
113
def test_batched_rotary_embedding(
114
is_neox_style: bool,
@@ -113,6 +121,7 @@ def test_batched_rotary_embedding(
121
dtype: torch.dtype,
122
seed: int,
123
device: str,
124
+ use_key: bool,
125
max_position: int = 8192,
126
base: int = 10000,
127
) -> None:
@@ -129,7 +138,7 @@ def test_batched_rotary_embedding(
138
positions = torch.randint(0, max_position, (batch_size, seq_len))
139
query_shape = tensor_shape_fn(batch_size, seq_len, num_heads, head_size)
140
query = torch.randn(query_shape, dtype=dtype)
132
- key = torch.randn_like(query)
141
+ key = torch.randn_like(query) if use_key else None
142
143
# NOTE(woosuk): The reference implementation should be executed first
144
# because the custom kernel is in-place.
@@ -145,10 +154,14 @@ def test_batched_rotary_embedding(
154
ref_query,
155
atol=get_default_atol(out_query),
156
rtol=get_default_rtol(out_query))
148
- torch.testing.assert_close(out_key,
149
- ref_key,
150
- atol=get_default_atol(out_key),
151
- rtol=get_default_rtol(out_key))
157
+ if use_key:
158
+ torch.testing.assert_close(out_key,
159
+ ref_key,
160
+ atol=get_default_atol(out_key),
161
+ rtol=get_default_rtol(out_key))
162
+ else:
163
+ assert ref_key is None and out_key is None, \
164
+ "expected returned key to be None"
165
166
167
@pytest.mark.parametrize("is_neox_style", IS_NEOX_STYLE)
@@ -160,6 +173,7 @@ def test_batched_rotary_embedding(
173
@pytest.mark.parametrize("dtype", DTYPES)
174
@pytest.mark.parametrize("seed", SEEDS)
175
@pytest.mark.parametrize("device", CUDA_DEVICES)
176
+@pytest.mark.parametrize("use_key", USE_KEY)
177
@torch.inference_mode()
178
def test_batched_rotary_embedding_multi_lora(
179
is_neox_style: bool,
@@ -171,6 +185,7 @@ def test_batched_rotary_embedding_multi_lora(
185
dtype: torch.dtype,
186
seed: int,
187
device: str,
188
+ use_key: bool,
189
max_position: int = 8192,
190
base: int = 10000,
191
) -> None:
@@ -190,7 +205,7 @@ def test_batched_rotary_embedding_multi_lora(
205
seq_len,
206
num_heads * head_size,
207
dtype=dtype)
193
- key = torch.randn_like(query)
208
+ key = torch.randn_like(query) if use_key else None
209
210
offset_map = torch.tensor(
211
list(
@@ -214,10 +229,14 @@ def test_batched_rotary_embedding_multi_lora(
229
ref_query,
230
atol=get_default_atol(out_query),
231
rtol=get_default_rtol(out_query))
217
- torch.testing.assert_close(out_key,
218
- ref_key,
219
- atol=get_default_atol(out_key),
220
- rtol=get_default_rtol(out_key))
232
+ if use_key:
233
+ torch.testing.assert_close(out_key,
234
+ ref_key,
235
+ atol=get_default_atol(out_key),
236
+ rtol=get_default_rtol(out_key))
237
+ else:
238
+ assert ref_key is None and out_key is None, \
239
+ "expected returned key to be None"
240
241
242
@torch.inference_mode()
tests/kernels/core/test_rotary_embedding.py
+4
-3
@@ -15,7 +15,7 @@ from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
15
def rotary_embedding_opcheck(rot,
16
positions: torch.Tensor,
17
query: torch.Tensor,
18
- key: torch.Tensor,
18
+ key: Optional[torch.Tensor] = None,
19
offsets: Optional[torch.Tensor] = None):
20
cos_sin_cache = rot.cos_sin_cache.to(query.device, dtype=query.dtype)
21
@@ -37,9 +37,10 @@ def rotary_embedding_opcheck(rot,
37
@pytest.mark.parametrize("rotary_dim", [32])
38
@pytest.mark.parametrize("head_size", [32, 108])
39
@pytest.mark.parametrize("seq_len", [11, 1024])
40
+@pytest.mark.parametrize("use_key", [True, False])
41
def test_rotary_embedding_opcheck(dist_init, device, max_position,
42
is_neox_style, rotary_dim, head_size,
42
- seq_len):
43
+ seq_len, use_key):
44
batch_size = 1
45
base = 10000
46
num_heads = 7
@@ -54,7 +55,7 @@ def test_rotary_embedding_opcheck(dist_init, device, max_position,
55
num_heads * head_size,
56
dtype=torch.float32,
57
device=device)
57
- key = torch.randn_like(query)
58
+ key = torch.randn_like(query) if use_key else None
59
60
rotary_embedding_opcheck(rot, positions, query, key)
61
offsets = torch.zeros(batch_size * seq_len,
tests/neuron/1_core/test_rotary_embedding.py
+21
-12
@@ -11,14 +11,16 @@ from vllm.platforms import current_platform
11
12
13
@pytest.mark.parametrize(
14
- "max_position,is_neox_style,rotary_dim,head_size,seq_len", [
15
- (16, False, 32, 32, 1024),
16
- (16, False, 32, 128, 1024),
17
- (16, True, 32, 32, 1024),
18
- (16, True, 32, 128, 1024),
14
+ "max_position,is_neox_style,rotary_dim,head_size,seq_len,use_key", [
15
+ (16, False, 32, 32, 1024, True),
16
+ (16, False, 32, 128, 1024, True),
17
+ (16, True, 32, 32, 1024, True),
18
+ (16, True, 32, 128, 1024, True),
19
+ (16, False, 32, 128, 1024, False),
20
+ (16, True, 32, 128, 1024, False),
21
])
22
def test_rotary_embedding_opcheck(max_position, is_neox_style, rotary_dim,
21
- head_size, seq_len):
23
+ head_size, seq_len, use_key):
24
import torch_xla.core.xla_model as xm
25
26
device = xm.xla_device()
@@ -40,19 +42,26 @@ def test_rotary_embedding_opcheck(max_position, is_neox_style, rotary_dim,
42
num_heads * head_size,
43
dtype=torch.float32,
44
device="cpu")
43
- key = torch.randn_like(query)
44
-
45
+ key = torch.randn_like(query) if use_key else None
46
assert positions.is_cpu, \
47
"reference input tensor is expected to be CPU tensor."
48
ref_query, ref_key = rot.to(device="cpu").forward_native(
49
positions, query, key)
50
out_query, out_key = rot.to(device=device).forward_neuron(
51
positions.to(device=device), query.to(device=device),
51
- key.to(device=device))
52
- assert out_query.is_xla and out_key.is_xla, \
53
- "output tensor is expected to be XLA tensor"
52
+ key.to(device=device) if key is not None else None)
53
+ if use_key:
54
+ assert out_query.is_xla and out_key.is_xla, \
55
+ "output tensor is expected to be XLA tensor"
56
+ torch.testing.assert_close(out_key.cpu(),
57
+ ref_key,
58
+ atol=1e-2,
59
+ rtol=1e-2)
60
+ else:
61
+ assert out_key is None, "expected returned key to be None"
62
+ assert out_query.is_xla, \
63
+ "output tensor is expected to be XLA tensor"
64
torch.testing.assert_close(out_query.cpu(),
65
ref_query,
66
atol=1e-2,
67
rtol=1e-2)
58
- torch.testing.assert_close(out_key.cpu(), ref_key, atol=1e-2, rtol=1e-2)
vllm/_custom_ops.py
+8
-6
@@ -153,34 +153,36 @@ def merge_attn_states(output: torch.Tensor,
153
def rotary_embedding(
154
positions: torch.Tensor,
155
query: torch.Tensor,
156
- key: torch.Tensor,
156
+ key: Optional[torch.Tensor],
157
head_size: int,
158
cos_sin_cache: torch.Tensor,
159
is_neox: bool,
160
) -> None:
161
# TODO: Remove this contiguous call when the kernel is updated to support tensor slices
162
query_contiguous = query.contiguous()
163
- key_contiguous = key.contiguous()
163
+ key_contiguous = key.contiguous() if key is not None else None
164
torch.ops._C.rotary_embedding(positions, query_contiguous, key_contiguous,
165
head_size, cos_sin_cache, is_neox)
166
query.copy_(query_contiguous)
167
- key.copy_(key_contiguous)
167
+ if key is not None:
168
+ key.copy_(key_contiguous)
169
170
171
def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,
171
- key: torch.Tensor, head_size: int,
172
+ key: Optional[torch.Tensor], head_size: int,
173
cos_sin_cache: torch.Tensor, is_neox: bool,
174
rot_dim: int,
175
cos_sin_cache_offsets: torch.Tensor) -> None:
176
# TODO: Remove this contiguous call when the kernel is updated to support tensor slices
177
query_contiguous = query.contiguous()
177
- key_contiguous = key.contiguous()
178
+ key_contiguous = key.contiguous() if key is not None else None
179
torch.ops._C.batched_rotary_embedding(positions, query_contiguous,
180
key_contiguous, head_size,
181
cos_sin_cache, is_neox, rot_dim,
182
cos_sin_cache_offsets)
183
query.copy_(query_contiguous)
183
- key.copy_(key_contiguous)
184
+ if key is not None:
185
+ key.copy_(key_contiguous)
186
187
188
# layer norm ops
vllm/model_executor/layers/rotary_embedding.py
+68
-50
@@ -138,9 +138,9 @@ class RotaryEmbedding(CustomOp):
138
self,
139
positions: torch.Tensor,
140
query: torch.Tensor,
141
- key: torch.Tensor,
141
+ key: Optional[torch.Tensor] = None,
142
offsets: Optional[torch.Tensor] = None,
143
- ) -> Tuple[torch.Tensor, torch.Tensor]:
143
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
144
"""A PyTorch-native implementation of forward()."""
145
if offsets is not None:
146
positions = positions + offsets
@@ -157,22 +157,24 @@ class RotaryEmbedding(CustomOp):
157
self.is_neox_style)
158
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
159
160
- key_shape = key.shape
161
- key = key.view(num_tokens, -1, self.head_size)
162
- key_rot = key[..., :self.rotary_dim]
163
- key_pass = key[..., self.rotary_dim:]
164
- key_rot = _apply_rotary_emb_torch(key_rot, cos, sin,
165
- self.is_neox_style)
166
- key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
160
+ # key may be None in some cases, e.g. cross-layer KV sharing
161
+ if key is not None:
162
+ key_shape = key.shape
163
+ key = key.view(num_tokens, -1, self.head_size)
164
+ key_rot = key[..., :self.rotary_dim]
165
+ key_pass = key[..., self.rotary_dim:]
166
+ key_rot = _apply_rotary_emb_torch(key_rot, cos, sin,
167
+ self.is_neox_style)
168
+ key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
169
return query, key
170
171
def forward_cuda(
172
self,
173
positions: torch.Tensor,
174
query: torch.Tensor,
173
- key: torch.Tensor,
175
+ key: Optional[torch.Tensor] = None,
176
offsets: Optional[torch.Tensor] = None,
175
- ) -> Tuple[torch.Tensor, torch.Tensor]:
177
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
178
from vllm import _custom_ops as ops
179
180
# __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
@@ -198,32 +200,39 @@ class RotaryEmbedding(CustomOp):
200
self,
201
positions: torch.Tensor,
202
query: torch.Tensor,
201
- key: torch.Tensor,
203
+ key: Optional[torch.Tensor] = None,
204
offsets: Optional[torch.Tensor] = None,
203
- ) -> Tuple[torch.Tensor, torch.Tensor]:
205
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
206
from vllm._ipex_ops import ipex_ops as ops
207
208
self.cos_sin_cache = self.cos_sin_cache.to(positions.device,
209
dtype=query.dtype)
210
# ops.rotary_embedding()/batched_rotary_embedding()
211
# are in-place operations that update the query and key tensors.
210
- if offsets is not None:
211
- ops.batched_rotary_embedding(positions, query, key, self.head_size,
212
- self.cos_sin_cache,
213
- self.is_neox_style, self.rotary_dim,
214
- offsets)
212
+ if key is None:
213
+ # XPU kernel doesn't support key=None so fall back to native impl
214
+ # TODO(sarckk): add support for optional key in
215
+ # ipex.llm.functional.rotary_embedding_batched
216
+ return self.forward_native(positions, query, key, offsets)
217
else:
216
- ops.rotary_embedding(positions, query, key, self.head_size,
217
- self.cos_sin_cache, self.is_neox_style)
218
+ if offsets is not None:
219
+ ops.batched_rotary_embedding(positions, query, key,
220
+ self.head_size,
221
+ self.cos_sin_cache,
222
+ self.is_neox_style,
223
+ self.rotary_dim, offsets)
224
+ else:
225
+ ops.rotary_embedding(positions, query, key, self.head_size,
226
+ self.cos_sin_cache, self.is_neox_style)
227
return query, key
228
229
def forward_hpu(
230
self,
231
positions: torch.Tensor,
232
query: torch.Tensor,
224
- key: torch.Tensor,
233
+ key: Optional[torch.Tensor] = None,
234
offsets: Optional[torch.Tensor] = None,
226
- ) -> Tuple[torch.Tensor, torch.Tensor]:
235
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
236
from habana_frameworks.torch.hpex.kernels import (
237
RotaryPosEmbeddingMode, apply_rotary_pos_emb)
238
if offsets is not None:
@@ -265,21 +274,23 @@ class RotaryEmbedding(CustomOp):
274
rope_mode)
275
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
276
268
- key_shape = key.shape
269
- key = key.view(num_tokens, -1, self.head_size)
270
- key_rot = key[..., :self.rotary_dim]
271
- key_pass = key[..., self.rotary_dim:]
272
- key_rot = apply_rotary_pos_emb(key_rot, cos, sin, None, 0, rope_mode)
273
- key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
277
+ if key is not None:
278
+ key_shape = key.shape
279
+ key = key.view(num_tokens, -1, self.head_size)
280
+ key_rot = key[..., :self.rotary_dim]
281
+ key_pass = key[..., self.rotary_dim:]
282
+ key_rot = apply_rotary_pos_emb(key_rot, cos, sin, None, 0,
283
+ rope_mode)
284
+ key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
285
return query, key
286
287
def forward_neuron(
288
self,
289
positions: torch.Tensor,
290
query: torch.Tensor,
280
- key: torch.Tensor,
291
+ key: Optional[torch.Tensor] = None,
292
offsets: Optional[torch.Tensor] = None,
282
- ) -> Tuple[torch.Tensor, torch.Tensor]:
293
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
294
295
def _apply_rotary_emb_neuron(
296
x: torch.Tensor,
@@ -319,14 +330,16 @@ class RotaryEmbedding(CustomOp):
330
331
query_shape = query.shape
332
query = query.view(num_tokens, -1, self.head_size)
322
- key_shape = key.shape
323
- key = key.view(num_tokens, -1, self.head_size)
333
+ if key is not None:
334
+ key_shape = key.shape
335
+ key = key.view(num_tokens, -1, self.head_size)
336
337
if self.rotary_dim == self.head_size:
338
query = _apply_rotary_emb(query, cos, sin, self.is_neox_style)
339
query = query.reshape(query_shape)
328
- key = _apply_rotary_emb(key, cos, sin, self.is_neox_style)
329
- key = key.reshape(key_shape)
340
+ if key is not None:
341
+ key = _apply_rotary_emb(key, cos, sin, self.is_neox_style)
342
+ key = key.reshape(key_shape)
343
else:
344
head_size = query.shape[-1]
345
query_reshaped = query.view(-1, head_size)
@@ -339,14 +352,15 @@ class RotaryEmbedding(CustomOp):
352
query = torch.cat((query_rot, query_pass),
353
dim=-1).reshape(query_shape)
354
342
- key_reshaped = key.view(-1, head_size)
343
- key_pass = key_reshaped[:, self.rotary_dim:].view(
344
- *key.shape[:-1], head_size - self.rotary_dim)
345
- key_rot = key_reshaped[:, :self.rotary_dim].view(
346
- *key.shape[:-1], self.rotary_dim)
347
- key_rot = _apply_rotary_emb_neuron(key_rot, cos, sin,
348
- self.is_neox_style)
349
- key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
355
+ if key is not None:
356
+ key_reshaped = key.view(-1, head_size)
357
+ key_pass = key_reshaped[:, self.rotary_dim:].view(
358
+ *key.shape[:-1], head_size - self.rotary_dim)
359
+ key_rot = key_reshaped[:, :self.rotary_dim].view(
360
+ *key.shape[:-1], self.rotary_dim)
361
+ key_rot = _apply_rotary_emb_neuron(key_rot, cos, sin,
362
+ self.is_neox_style)
363
+ key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
364
return query, key
365
366
def extra_repr(self) -> str:
@@ -672,9 +686,10 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):
686
self,
687
positions: torch.Tensor,
688
query: torch.Tensor,
675
- key: torch.Tensor,
689
+ key: Optional[torch.Tensor] = None,
690
offsets: Optional[torch.Tensor] = None,
677
- ) -> Tuple[torch.Tensor, torch.Tensor]:
691
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
692
+ assert key is not None
693
query = query.view(*query.shape[:-1], -1, self.head_size)
694
key = key.view(*key.shape[:-1], -1, self.head_size)
695
@@ -782,10 +797,11 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
797
self,
798
positions: torch.Tensor,
799
query: torch.Tensor,
785
- key: torch.Tensor,
800
+ key: Optional[torch.Tensor] = None,
801
offsets: Optional[torch.Tensor] = None,
787
- ) -> Tuple[torch.Tensor, torch.Tensor]:
802
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
803
"""PyTorch-native implementation equivalent to forward()."""
804
+ assert key is not None
805
query_rot = query[..., :self.rotary_dim]
806
key_rot = key[..., :self.rotary_dim]
807
if self.rotary_dim < self.head_size:
@@ -912,8 +928,9 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):
928
def forward(
929
self,
930
query: torch.Tensor,
915
- key: torch.Tensor,
916
- ) -> Tuple[torch.Tensor, torch.Tensor]:
931
+ key: Optional[torch.Tensor] = None,
932
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
933
+ assert key is not None
934
self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device)
935
query_ = torch.view_as_complex(query.float().reshape(
936
*query.shape[:-1], -1, 2))
@@ -957,8 +974,8 @@ class MRotaryEmbedding(RotaryEmbedding):
974
self,
975
positions: torch.Tensor,
976
query: torch.Tensor,
960
- key: torch.Tensor,
961
- ) -> Tuple[torch.Tensor, torch.Tensor]:
977
+ key: Optional[torch.Tensor] = None,
978
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
979
"""PyTorch-native implementation equivalent to forward().
980
981
Args:
@@ -969,6 +986,7 @@ class MRotaryEmbedding(RotaryEmbedding):
986
key: [num_tokens, num_kv_heads * head_size]
987
"""
988
assert positions.ndim == 1 or positions.ndim == 2
989
+ assert key is not None
990
991
num_tokens = positions.shape[-1]
992
cos_sin = self.cos_sin_cache[positions]