cwenzi/neuroflow-cpp

1#include "test_framework.hpp"
2#include "neuroflow/rope.hpp"
3#include <cmath>
4
5using namespace neuroflow;
6
7TEST(RoPE, ConstructionAndFreqShape) {
8    RoPE rope(64, 128);
9    EXPECT_EQ(rope.head_dim_, 64u);
10    EXPECT_EQ(rope.max_seq_len_, 128u);
11    EXPECT_EQ(rope.freqs_cos_.shape_.size(), 2u);
12    EXPECT_EQ(rope.freqs_cos_.shape_[0], 128u);
13    EXPECT_EQ(rope.freqs_cos_.shape_[1], 32u);
14}
15
16TEST(RoPE, FreqValuesInRange) {
17    RoPE rope(64, 128);
18    const float* cos_data = rope.freqs_cos_.as_fp32();
19    const float* sin_data = rope.freqs_sin_.as_fp32();
20    size_t n = rope.freqs_cos_.numel();
21    for (size_t i = 0; i < n; ++i) {
22        EXPECT_TRUE(cos_data[i] >= -1.0f - 1e-6f && cos_data[i] <= 1.0f + 1e-6f);
23        EXPECT_TRUE(sin_data[i] >= -1.0f - 1e-6f && sin_data[i] <= 1.0f + 1e-6f);
24    }
25}
26
27TEST(RoPE, ApplySingleModifiesQK) {
28    size_t head_dim = 64;
29    size_t n_heads = 4;
30    size_t seq_len = 8;
31    size_t total = seq_len * n_heads * head_dim;
32
33    RoPE rope(head_dim, 128);
34
35    Tensor q({seq_len, n_heads * head_dim}, QuantType::FP32);
36    float* qp = q.as_fp32();
37    for (size_t i = 0; i < total; ++i) qp[i] = 1.0f;
38
39    Tensor q_copy({seq_len, n_heads * head_dim}, QuantType::FP32);
40    memcpy(q_copy.as_fp32(), qp, total * sizeof(float));
41
42    rope.apply_single(q, seq_len, n_heads, 0);
43
44    bool changed = false;
45    const float* qp_after = q.as_fp32();
46    const float* qcp = q_copy.as_fp32();
47    for (size_t i = 0; i < total; ++i) {
48        if (std::abs(qp_after[i] - qcp[i]) > 1e-6f) { changed = true; break; }
49    }
50    EXPECT_TRUE(changed);
51}
52
53TEST(RoPE, PositionZeroCosOneSinZero) {
54    RoPE rope(64, 128);
55    const float* cos_data = rope.freqs_cos_.as_fp32();
56    const float* sin_data = rope.freqs_sin_.as_fp32();
57    for (size_t d = 0; d < rope.freqs_cos_.shape_[1]; ++d) {
58        EXPECT_NEAR(cos_data[0 * rope.freqs_cos_.shape_[1] + d], 1.0f, 1e-5f);
59        EXPECT_NEAR(sin_data[0 * rope.freqs_sin_.shape_[1] + d], 0.0f, 1e-5f);
60    }
61}
62
63TEST(RoPE, YarnScaleUpdatesFreqs) {
64    RoPE rope(64, 128);
65    float orig_cos_1 = rope.freqs_cos_.as_fp32()[rope.freqs_cos_.shape_[1]];
66    rope.set_yarn_scale(2.0f);
67    float new_cos_1 = rope.freqs_cos_.as_fp32()[rope.freqs_cos_.shape_[1]];
68    EXPECT_NE(orig_cos_1, new_cos_1);
69}
70
71int main() { RUN_ALL_TESTS(); }
72