1+ // Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
2+ //
3+ // Licensed under the Apache License, Version 2.0 (the "License");
4+ // you may not use this file except in compliance with the License.
5+ // You may obtain a copy of the License at
6+ //
7+ // http://www.apache.org/licenses/LICENSE-2.0
8+ //
9+ // Unless required by applicable law or agreed to in writing, software
10+ // distributed under the License is distributed on an "AS IS" BASIS,
11+ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ // See the License for the specific language governing permissions and
13+ // limitations under the License.
14+
15+ #include " helper.h"
16+ #include " paddle/extension.h"
17+
18+ template <typename T, bool IS_NEOX>
19+ inline __device__ void apply_token_rotary_embedding_kernel (
20+ T* __restrict__ arr,
21+ const T* __restrict__ cos_ptr,
22+ const T* __restrict__ sin_ptr,
23+ int rot_offset,
24+ int embed_dim) {
25+ int x_index, y_index;
26+ T cos , sin ;
27+ if (IS_NEOX) {
28+ x_index = rot_offset;
29+ y_index = embed_dim + rot_offset;
30+ cos = cos_ptr[x_index];
31+ sin = sin_ptr[x_index];
32+ } else {
33+ x_index = 2 * rot_offset;
34+ y_index = 2 * rot_offset + 1 ;
35+ cos = cos_ptr[x_index / 2 ];
36+ sin = sin_ptr[x_index / 2 ];
37+ }
38+
39+ const T x = arr[x_index];
40+ const T y = arr[y_index];
41+ arr[x_index] = x * cos - y * sin ;
42+ arr[y_index] = y * cos + x * sin ;
43+ }
44+
45+
46+ template <typename T, bool IS_NEOX>
47+ __global__ void apply_rotary_embedding_kernel (
48+ T* __restrict__ query, // [num_tokens, num_heads, head_size]
49+ T* __restrict__ key, // [num_tokens, num_kv_heads, head_size]
50+ const int * __restrict__ position_ids, // [num_tokens]
51+ const T* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim // 2]
52+ const int rot_dim,
53+ const int64_t query_stride,
54+ const int64_t key_stride,
55+ const int num_heads,
56+ const int num_kv_heads,
57+ const int head_size) {
58+ // Each thread block is responsible for one token.
59+ const int token_idx = blockIdx .x ;
60+ int pos = position_ids[token_idx];
61+ const T* cache_ptr = cos_sin_cache + pos * rot_dim;
62+
63+ const int embed_dim = rot_dim / 2 ;
64+ const T* cos_ptr = cache_ptr;
65+ const T* sin_ptr = cache_ptr + embed_dim;
66+
67+ const int nq = num_heads * embed_dim;
68+ for (int i = threadIdx .x ; i < nq; i += blockDim .x ) {
69+ const int head_idx = i / embed_dim;
70+ const int64_t token_head = token_idx * query_stride + head_idx * head_size;
71+ const int rot_offset = i % embed_dim;
72+ apply_token_rotary_embedding_kernel<T, IS_NEOX>(
73+ query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
74+ }
75+
76+ const int nk = num_kv_heads * embed_dim;
77+ for (int i = threadIdx .x ; i < nk; i += blockDim .x ) {
78+ const int head_idx = i / embed_dim;
79+ const int64_t token_head = token_idx * key_stride + head_idx * head_size;
80+ const int rot_offset = i % embed_dim;
81+ apply_token_rotary_embedding_kernel<T, IS_NEOX>(
82+ key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
83+ }
84+ }
85+
86+
87+ void FusedRotaryPositionEncoding (
88+ paddle::Tensor& query, // [num_tokens, num_heads, head_size] or
89+ // [num_tokens, num_heads * head_size]
90+ paddle::Tensor& key,
91+ // [num_tokens, num_kv_heads, head_size] or [num_tokens, num_kv_heads *
92+ // head_size]
93+ const paddle::Tensor& position_ids, // [num_tokens]
94+ const paddle::Tensor& cos_sin_cache, // [max_position, rot_dim]
95+ int head_size,
96+ bool is_neox) {
97+ int64_t num_tokens = query.dims ()[0 ];
98+ int num_heads = query.numel () / num_tokens / head_size;
99+ int num_kv_heads = key.numel () / num_tokens / head_size;
100+ int rot_dim = cos_sin_cache.dims ()[1 ];
101+ int64_t query_stride = num_heads * head_size;
102+ int64_t key_stride = num_kv_heads * head_size;
103+
104+ dim3 grid (num_tokens);
105+ dim3 block (std::min<int64_t >(num_heads * rot_dim / 2 , 512 ));
106+ PD_DISPATCH_FLOATING_TYPES (
107+ query.dtype (), " apply_rotary_embedding_kernel"
108+ if (is_neox) {
109+ apply_rotary_embedding_kernel<data_t , true >
110+ <<<grid, block, 0 , query.stream()>>> (query.data <data_t >(),
111+ key.data <data_t >(),
112+ position_ids.data <int >(),
113+ cos_sin_cache.data <data_t >(),
114+ rot_dim,
115+ query_stride,
116+ key_stride,
117+ num_heads,
118+ num_kv_heads,
119+ head_size);
120+ } else {
121+ apply_rotary_embedding_kernel<data_t , false >
122+ <<<grid, block, 0 , query.stream()>>> (query.data <data_t >(),
123+ key.data <data_t >(),
124+ position_ids.data <int >(),
125+ cos_sin_cache.data <data_t >(),
126+ rot_dim,
127+ query_stride,
128+ key_stride,
129+ num_heads,
130+ num_kv_heads,
131+ head_size);
132+ }
133+ });
134+ }
135+
136+ PD_BUILD_OP (fused_rotary_position_encoding)
137+ .Inputs({" query" " key" " position_ids" " cos_sin_cache"
138+ .Outputs({" query_out" " key_out"
139+ .Attrs({" head_size: int" " is_neox: bool"
140+ .SetInplaceMap({{" query" " query_out" " key" " key_out"
141+ .SetKernelFn(PD_KERNEL(FusedRotaryPositionEncoding));
0 commit comments