File size: 1,196 Bytes
26d5b81 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | #ifndef NEUROFLOW_ROPE_HPP
#define NEUROFLOW_ROPE_HPP
#include <cstddef>
#include <vector>
#include "tensor.hpp"
#ifdef USE_CUDA
#include <cuda_runtime.h>
#endif
namespace neuroflow {
class RoPE {
public:
size_t head_dim_;
size_t max_seq_len_;
float yarn_scale_factor_ = 1.0f;
Tensor freqs_cos_;
Tensor freqs_sin_;
RoPE(size_t head_dim, size_t max_seq_len);
void apply(Tensor& qkv, size_t seq_len, size_t n_heads, size_t d_model, size_t offset = 0) const;
void apply_single(Tensor& x, size_t seq_len, size_t n_heads, size_t offset = 0) const;
void set_yarn_scale(float scale_factor);
};
#ifdef USE_CUDA
void launch_rope(float* qkv, const float* cos_table, const float* sin_table,
size_t seq_len, size_t n_heads, size_t d_model, size_t head_dim,
size_t offset, cudaStream_t stream);
void launch_rope_single(float* data, const float* cos_table, const float* sin_table,
size_t seq_len, size_t n_heads, size_t stride, size_t head_dim,
size_t offset, cudaStream_t stream);
#endif
} // namespace neuroflow
#endif // NEUROFLOW_ROPE_HPP |