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 4, 12:30 UTC
[Bugfix] Fix rope (#41113)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Apr 29, 2026 at 13:42 UTC
a269744e9f733ec9bac4bb6a33f70cc5af38afc3
2 files changed
+45
-36
csrc/pos_encoding_kernels.cu
+39
-34
@@ -7,23 +7,23 @@
7
8
namespace vllm {
9
10
-template <typename scalar_t, bool IS_NEOX>
10
+template <typename scalar_t, typename cache_t, bool IS_NEOX>
11
inline __device__ void apply_token_rotary_embedding(
12
- scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
13
- const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
12
+ scalar_t* __restrict__ arr, const cache_t* __restrict__ cos_ptr,
13
+ const cache_t* __restrict__ sin_ptr, int rot_offset, int embed_dim,
14
const bool inverse) {
15
int x_index, y_index;
16
float cos_f, sin_f;
17
if (IS_NEOX) {
18
x_index = rot_offset;
19
y_index = embed_dim + rot_offset;
20
- cos_f = VLLM_LDG(cos_ptr + x_index);
21
- sin_f = VLLM_LDG(sin_ptr + x_index);
20
+ cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index));
21
+ sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index));
22
} else {
23
x_index = 2 * rot_offset;
24
y_index = 2 * rot_offset + 1;
25
- cos_f = VLLM_LDG(cos_ptr + x_index / 2);
26
- sin_f = VLLM_LDG(sin_ptr + x_index / 2);
25
+ cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index / 2));
26
+ sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index / 2));
27
}
28
if (inverse) {
29
sin_f = -sin_f;
@@ -34,7 +34,7 @@ inline __device__ void apply_token_rotary_embedding(
34
arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
35
}
36
37
-template <typename scalar_t, bool IS_NEOX>
37
+template <typename scalar_t, typename cache_t, bool IS_NEOX>
38
inline __device__ void apply_rotary_embedding(
39
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
40
// head_size] or [num_tokens, num_heads,
@@ -43,14 +43,14 @@ inline __device__ void apply_rotary_embedding(
43
// [batch_size, seq_len, num_kv_heads,
44
// head_size] or [num_tokens, num_kv_heads,
45
// head_size]
46
- const float* cache_ptr, const int head_size, const int num_heads,
46
+ const cache_t* cache_ptr, const int head_size, const int num_heads,
47
const int num_kv_heads, const int rot_dim, const int token_idx,
48
const int64_t query_stride, const int64_t key_stride,
49
const int64_t head_stride, const int64_t rope_dim_offset,
50
const bool inverse) {
51
const int embed_dim = rot_dim / 2;
52
- const float* cos_ptr = cache_ptr;
53
- const float* sin_ptr = cache_ptr + embed_dim;
52
+ const cache_t* cos_ptr = cache_ptr;
53
+ const cache_t* sin_ptr = cache_ptr + embed_dim;
54
55
const int nq = num_heads * embed_dim;
56
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
@@ -58,7 +58,7 @@ inline __device__ void apply_rotary_embedding(
58
const int64_t token_head =
59
token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
60
const int rot_offset = i % embed_dim;
61
- apply_token_rotary_embedding<scalar_t, IS_NEOX>(
61
+ apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
62
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
63
}
64
@@ -69,13 +69,13 @@ inline __device__ void apply_rotary_embedding(
69
const int64_t token_head =
70
token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
71
const int rot_offset = i % embed_dim;
72
- apply_token_rotary_embedding<scalar_t, IS_NEOX>(
72
+ apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
73
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
74
}
75
}
76
}
77
78
-template <typename scalar_t, bool IS_NEOX>
78
+template <typename scalar_t, typename cache_t, bool IS_NEOX>
79
__global__ void rotary_embedding_kernel(
80
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
81
// [num_tokens]
@@ -86,15 +86,15 @@ __global__ void rotary_embedding_kernel(
86
// [batch_size, seq_len, num_kv_heads,
87
// head_size] or [num_tokens, num_kv_heads,
88
// head_size]
89
- const float* __restrict__ cos_sin_cache, // [max_position, rot_dim] fp32
89
+ const cache_t* __restrict__ cos_sin_cache, // [max_position, rot_dim]
90
const int rot_dim, const int64_t query_stride, const int64_t key_stride,
91
const int64_t head_stride, const int num_heads, const int num_kv_heads,
92
const int head_size, const int64_t rope_dim_offset, const bool inverse) {
93
const int token_idx = blockIdx.x;
94
int64_t pos = positions[token_idx];
95
- const float* cache_ptr = cos_sin_cache + pos * rot_dim;
95
+ const cache_t* cache_ptr = cos_sin_cache + pos * rot_dim;
96
97
- apply_rotary_embedding<scalar_t, IS_NEOX>(
97
+ apply_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
98
query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
99
token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
100
inverse);
@@ -168,23 +168,28 @@ void rotary_embedding(
168
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
169
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
170
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
171
- auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
171
VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
173
- if (is_neox) {
174
- vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
175
- positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
176
- key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
177
- cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
178
- head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
179
- inverse);
180
- } else {
181
- vllm::rotary_embedding_kernel<scalar_t, false>
182
- <<<grid, block, 0, stream>>>(
183
- positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
184
- key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
185
- cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
186
- head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
187
- inverse);
188
- }
172
+ using query_t = scalar_t;
173
+ VLLM_DISPATCH_FLOATING_TYPES(
174
+ cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] {
175
+ using cache_t = scalar_t;
176
+ if (is_neox) {
177
+ vllm::rotary_embedding_kernel<query_t, cache_t, true>
178
+ <<<grid, block, 0, stream>>>(
179
+ positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
180
+ key.has_value() ? key->data_ptr<query_t>() : nullptr,
181
+ cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
182
+ key_stride, head_stride, num_heads, num_kv_heads, head_size,
183
+ rope_dim_offset, inverse);
184
+ } else {
185
+ vllm::rotary_embedding_kernel<query_t, cache_t, false>
186
+ <<<grid, block, 0, stream>>>(
187
+ positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
188
+ key.has_value() ? key->data_ptr<query_t>() : nullptr,
189
+ cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
190
+ key_stride, head_stride, num_heads, num_kv_heads, head_size,
191
+ rope_dim_offset, inverse);
192
+ }
193
+ });
194
});
195
}
tests/kernels/core/test_rotary_embedding.py
+6
-2
@@ -35,6 +35,9 @@ def rotary_embedding_opcheck(
35
@pytest.mark.parametrize("seq_len", [11, 1024])
36
@pytest.mark.parametrize("use_key", [True, False])
37
@pytest.mark.parametrize("head_stride_is_contiguous", [True, False])
38
+@pytest.mark.parametrize(
39
+ "dtype", [torch.float32, torch.bfloat16]
40
+)
41
def test_rotary_embedding_opcheck(
42
default_vllm_config,
43
dist_init,
@@ -46,19 +49,20 @@ def test_rotary_embedding_opcheck(
49
seq_len,
50
use_key,
51
head_stride_is_contiguous,
52
+ dtype,
53
):
54
batch_size = 1
55
base = 10000
56
num_heads = 7
57
rot = RotaryEmbedding(
54
- head_size, rotary_dim, max_position, base, is_neox_style, torch.float32
58
+ head_size, rotary_dim, max_position, base, is_neox_style, dtype
59
)
60
61
positions = torch.randint(0, max_position, (batch_size, seq_len), device=device)
62
head_stride = head_size + (64 if head_stride_is_contiguous else 0)
63
64
query = torch.randn(
61
- batch_size, seq_len, num_heads, head_stride, dtype=torch.float32, device=device
65
+ batch_size, seq_len, num_heads, head_stride, dtype=dtype, device=device
66
)
67
key = torch.randn_like(query) if use_key else None
68
query = query[..., :head_size]