cwenzi/neuroflow-cpp
1
1#include "neuroflow/rope.hpp"
2
3#include <cmath>
4#include <iostream>
5
6#ifdef USE_CUDA
7#include "cuda_context.hpp"
8#endif
9
10namespace neuroflow {
11
12RoPE::RoPE(size_t head_dim, size_t max_seq_len)
13 : head_dim_(head_dim), max_seq_len_(max_seq_len),
14 freqs_cos_({max_seq_len, head_dim / 2}, QuantType::FP32),
15 freqs_sin_({max_seq_len, head_dim / 2}, QuantType::FP32) {
16 float* cos_p = freqs_cos_.as_fp32();
17 float* sin_p = freqs_sin_.as_fp32();
18 float base = 10000.0f;
19 size_t half_dim = head_dim / 2;
20
21 for (size_t pos = 0; pos < max_seq_len; ++pos) {
22 for (size_t d = 0; d < half_dim; ++d) {
23 float freq = 1.0f / std::pow(base, static_cast<float>(2 * d) / static_cast<float>(head_dim));
24 float angle = static_cast<float>(pos) * freq;
25 cos_p[pos * half_dim + d] = std::cos(angle);
26 sin_p[pos * half_dim + d] = std::sin(angle);
27 }
28 }
29
30#ifdef USE_CUDA
31 if (CudaContext::instance().is_available()) {
32 freqs_cos_.to_gpu();
33 freqs_sin_.to_gpu();
34 }
35#endif
36}
37
38void RoPE::apply(Tensor& qkv, size_t seq_len, size_t n_heads, size_t d_model, size_t offset) const {
39 size_t half_dim = head_dim_ / 2;
40
41#ifdef USE_CUDA
42 if (CudaContext::instance().is_available() && qkv.is_on_gpu()) {
43 launch_rope(qkv.as_gpu_fp32(), freqs_cos_.as_gpu_fp32(), freqs_sin_.as_gpu_fp32(),
44 seq_len, n_heads, d_model, head_dim_, offset,
45 CudaContext::instance().stream());
46 qkv.gpu_dirty_ = true;
47 return;
48 }
49#endif
50
51 float* qkvp = qkv.as_fp32();
52 const float* cos_p = freqs_cos_.as_fp32();
53 const float* sin_p = freqs_sin_.as_fp32();
54
55 for (size_t h = 0; h < n_heads; ++h) {
56 size_t q_off = h * head_dim_;
57 size_t k_off = d_model + h * head_dim_;
58
59 for (size_t i = 0; i < seq_len; ++i) {
60 size_t pos = offset + i;
61 if (pos >= max_seq_len_) {
62 if (yarn_scale_factor_ > 1.0f) {
63 pos = pos % max_seq_len_;
64 } else {
65 pos = max_seq_len_ - 1;
66 }
67 }
68
69 for (size_t d = 0; d < half_dim; ++d) {
70 float c = cos_p[pos * half_dim + d];
71 float s = sin_p[pos * half_dim + d];
72
73 size_t qi = i * 3 * d_model + q_off + d;
74 size_t qi2 = i * 3 * d_model + q_off + half_dim + d;
75 float q0 = qkvp[qi];
76 float q1 = qkvp[qi2];
77 qkvp[qi] = q0 * c - q1 * s;
78 qkvp[qi2] = q0 * s + q1 * c;
79
80 size_t ki = i * 3 * d_model + k_off + d;
81 size_t ki2 = i * 3 * d_model + k_off + half_dim + d;
82 float k0 = qkvp[ki];
83 float k1 = qkvp[ki2];
84 qkvp[ki] = k0 * c - k1 * s;
85 qkvp[ki2] = k0 * s + k1 * c;
86 }
87 }
88 }
89}
90
91void RoPE::apply_single(Tensor& x, size_t seq_len, size_t n_heads, size_t offset) const {
92 size_t half_dim = head_dim_ / 2;
93 size_t stride = n_heads * head_dim_;
94
95#ifdef USE_CUDA
96 if (CudaContext::instance().is_available() && x.is_on_gpu()) {
97 launch_rope_single(x.as_gpu_fp32(), freqs_cos_.as_gpu_fp32(), freqs_sin_.as_gpu_fp32(),
98 seq_len, n_heads, stride, head_dim_, offset,
99 CudaContext::instance().stream());
100 x.gpu_dirty_ = true;
101 return;
102 }
103#endif
104
105 float* xp = x.as_fp32();
106 const float* cos_p = freqs_cos_.as_fp32();
107 const float* sin_p = freqs_sin_.as_fp32();
108
109 for (size_t h = 0; h < n_heads; ++h) {
110 size_t h_off = h * head_dim_;
111 for (size_t i = 0; i < seq_len; ++i) {
112 size_t pos = offset + i;
113 if (pos >= max_seq_len_) {
114 pos = yarn_scale_factor_ > 1.0f ? pos % max_seq_len_ : max_seq_len_ - 1;
115 }
116 for (size_t d = 0; d < half_dim; ++d) {
117 float c = cos_p[pos * half_dim + d];
118 float s = sin_p[pos * half_dim + d];
119 size_t idx0 = i * stride + h_off + d;
120 size_t idx1 = i * stride + h_off + half_dim + d;
121 float x0 = xp[idx0];
122 float x1 = xp[idx1];
123 xp[idx0] = x0 * c - x1 * s;
124 xp[idx1] = x0 * s + x1 * c;
125 }
126 }
127 }
128}
129
130void RoPE::set_yarn_scale(float scale_factor) {
131 if (scale_factor <= 1.0f) return;
132 if (scale_factor > 100.0f) {
133 std::cerr << "[YaRN WARNING] Extreme scale_factor=" << scale_factor
134 << ", results may be unreliable" << std::endl;
135 }
136 yarn_scale_factor_ = scale_factor;
137
138 float* cos_p = freqs_cos_.as_fp32();
139 float* sin_p = freqs_sin_.as_fp32();
140 float base = 10000.0f;
141 size_t half_dim = head_dim_ / 2;
142
143 for (size_t pos = 0; pos < max_seq_len_; ++pos) {
144 for (size_t d = 0; d < half_dim; ++d) {
145 float freq = 1.0f / std::pow(base, static_cast<float>(2 * d) / static_cast<float>(head_dim_));
146 float scaled_freq = freq / scale_factor;
147 float angle = static_cast<float>(pos) * scaled_freq;
148 cos_p[pos * half_dim + d] = std::cos(angle);
149 sin_p[pos * half_dim + d] = std::sin(angle);
150 }
151 }
152
153#ifdef USE_CUDA
154 if (CudaContext::instance().is_available()) {
155 freqs_cos_.to_gpu();
156 freqs_sin_.to_gpu();
157 }
158#endif
159
160 std::cerr << "[YaRN] Applied scale_factor=" << scale_factor
161 << ", temp_scale=" << (std::sqrt(std::log(scale_factor)) + 1.0f) << std::endl;
162}
163
164#ifdef USE_CUDA
165__global__ void kernel_rope_impl(float* qkv, const float* cos_table, const float* sin_table,
166 size_t seq_len, size_t n_heads, size_t d_model,
167 size_t head_dim, size_t half_dim, size_t offset) {
168 size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
169 size_t total = seq_len * n_heads * half_dim;
170 if (idx >= total) return;
171
172 size_t d = idx % half_dim;
173 size_t h = (idx / half_dim) % n_heads;
174 size_t i = idx / (n_heads * half_dim);
175
176 size_t pos = offset + i;
177
178 float c = cos_table[pos * half_dim + d];
179 float s = sin_table[pos * half_dim + d];
180
181 size_t q_off = h * head_dim;
182 size_t k_off = d_model + h * head_dim;
183
184 size_t qi = i * 3 * d_model + q_off + d;
185 size_t qi2 = i * 3 * d_model + q_off + half_dim + d;
186 float q0 = qkv[qi];
187 float q1 = qkv[qi2];
188 qkv[qi] = q0 * c - q1 * s;
189 qkv[qi2] = q0 * s + q1 * c;
190
191 size_t ki = i * 3 * d_model + k_off + d;
192 size_t ki2 = i * 3 * d_model + k_off + half_dim + d;
193 float k0 = qkv[ki];
194 float k1 = qkv[ki2];
195 qkv[ki] = k0 * c - k1 * s;
196 qkv[ki2] = k0 * s + k1 * c;
197}
198
199void launch_rope(float* qkv, const float* cos_table, const float* sin_table,
200 size_t seq_len, size_t n_heads, size_t d_model, size_t head_dim,
201 size_t offset, cudaStream_t stream) {
202 size_t half_dim = head_dim / 2;
203 size_t total = seq_len * n_heads * half_dim;
204 int block = 256;
205 int grid = (static_cast<int>(total) + block - 1) / block;
206 kernel_rope_impl<<<grid, block, 0, stream>>>(qkv, cos_table, sin_table,
207 seq_len, n_heads, d_model, head_dim, half_dim, offset);
208}
209__global__ void kernel_rope_single_impl(float* data, const float* cos_table, const float* sin_table,
210 size_t seq_len, size_t n_heads, size_t stride,
211 size_t head_dim, size_t half_dim, size_t offset) {
212 size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
213 size_t total = seq_len * n_heads * half_dim;
214 if (idx >= total) return;
215
216 size_t d = idx % half_dim;
217 size_t h = (idx / half_dim) % n_heads;
218 size_t i = idx / (n_heads * half_dim);
219
220 size_t pos = offset + i;
221 float c = cos_table[pos * half_dim + d];
222 float s = sin_table[pos * half_dim + d];
223
224 size_t h_off = h * head_dim;
225 size_t idx0 = i * stride + h_off + d;
226 size_t idx1 = i * stride + h_off + half_dim + d;
227 float x0 = data[idx0];
228 float x1 = data[idx1];
229 data[idx0] = x0 * c - x1 * s;
230 data[idx1] = x0 * s + x1 * c;
231}
232
233void launch_rope_single(float* data, const float* cos_table, const float* sin_table,
234 size_t seq_len, size_t n_heads, size_t stride, size_t head_dim,
235 size_t offset, cudaStream_t stream) {
236 size_t half_dim = head_dim / 2;
237 size_t total = seq_len * n_heads * half_dim;
238 int block = 256;
239 int grid = (static_cast<int>(total) + block - 1) / block;
240 kernel_rope_single_impl<<<grid, block, 0, stream>>>(data, cos_table, sin_table,
241 seq_len, n_heads, stride, head_dim, half_dim, offset);
242}
243#endif
244
245} // namespace neuroflow