GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 82.7% 62 / 0 / 75
Functions: 100.0% 3 / 0 / 3
Branches: 30.4% 42 / 0 / 138
src/cpu/rope_mrope.cpp
Line Branch Exec Source
1 // ─── CPU M-RoPE (Qwen2.5-VL / Qwen3-VL) ────────────────────────────────────
2 //
3 // Three per-axis (t, h, w) position-ID streams index three per-axis cos/sin
4 // tables and apply pair-wise rotations to disjoint sub-ranges of each head.
5 // head_dim is split into three contiguous scalar sub-ranges of widths 2*d_t,
6 // 2*d_h, 2*d_w (in order t,h,w), with 2*(d_t + d_h + d_w) == head_dim.
7 //
8 // X / Y : (L, num_heads*head_dim)
9 // cos_a/sin_a: (max_pos_a, d_a) FP32 — width = pairs per axis, NOT 2*d_a;
10 // this matches rope_apply's convention (one cos per pair) and
11 // lets the M-RoPE op reuse the same math without doubling table
12 // width.
13 // pos_a : length-L INT32 host pointer on CPU; device pointer on GPU.
14 //
15 // CPU is FP32-only; X must be FP32.
16
17 #include <brotensor/tensor.h>
18
19 #include <cstdint>
20 #include <stdexcept>
21 #include <string>
22
23 namespace brotensor::detail::cpu {
24
25 namespace {
26
27 18 inline void check_axis_tbl(const ::brotensor::Tensor& cos_a,
28 const ::brotensor::Tensor& sin_a,
29 const char* axis, int d_a) {
30
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (d_a < 0) {
31 ✗ throw std::runtime_error(std::string("rope_apply_mrope: d_") + axis +
32 " must be non-negative");
33 }
34
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 4 times.
 18 if (d_a == 0) {
35 // Empty axis: tables are not consulted; accept any (even empty) shape.
36 4 return;
37 }
38
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (cos_a.dtype != Dtype::FP32 || sin_a.dtype != Dtype::FP32) {
39 ✗ throw std::runtime_error(std::string("rope_apply_mrope: cos_") + axis +
40 ✗ " / sin_" + axis + " must be FP32");
41 }
42
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (cos_a.cols != d_a || sin_a.cols != d_a) {
43 ✗ throw std::runtime_error(std::string("rope_apply_mrope: cos_") + axis +
44 ✗ " / sin_" + axis +
45 ✗ " must each have cols == d_" + axis);
46 }
47
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (cos_a.rows != sin_a.rows || cos_a.rows < 1) {
48 ✗ throw std::runtime_error(std::string("rope_apply_mrope: cos_") + axis +
49 ✗ " / sin_" + axis + " row count mismatch");
50 }
51 18 }
52
53 18 inline void apply_axis(const float* Xp, float* Yp,
54 const float* cos_a, const float* sin_a,
55 const int32_t* pos_a, int max_pos_a,
56 int L, int D, int head_dim, int num_heads,
57 int pair_offset_in_head, int d_a) {
58
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 4 times.
 18 if (d_a == 0) return;
59
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 310 times.
 324 for (int row = 0; row < L; ++row) {
60 310 const int pos = pos_a[row];
61
1/2
✓ Branch 0 taken 310 times.
✗ Branch 1 not taken.
 310 if (pos < 0 || pos >= max_pos_a) {
62 ✗ throw std::runtime_error(
63 "rope_apply_mrope: position id out of range");
64 }
65
2/2
✓ Branch 0 taken 680 times.
✓ Branch 1 taken 310 times.
 990 for (int h = 0; h < num_heads; ++h) {
66 1360 const int base_off = row * D + h * head_dim
67 680 + 2 * pair_offset_in_head;
68
2/2
✓ Branch 0 taken 5088 times.
✓ Branch 1 taken 680 times.
 5768 for (int i = 0; i < d_a; ++i) {
69 5088 const float c = cos_a[pos * d_a + i];
70 5088 const float s = sin_a[pos * d_a + i];
71 5088 const float x0 = Xp[base_off + 2 * i];
72 5088 const float x1 = Xp[base_off + 2 * i + 1];
73 5088 Yp[base_off + 2 * i] = x0 * c - x1 * s;
74 5088 Yp[base_off + 2 * i + 1] = x0 * s + x1 * c;
75 5088 }
76 680 }
77 310 }
78 18 }
79
80 } // namespace
81
82 7 void rope_apply_mrope(const ::brotensor::Tensor& X,
83 const ::brotensor::Tensor& cos_t,
84 const ::brotensor::Tensor& sin_t,
85 const ::brotensor::Tensor& cos_h,
86 const ::brotensor::Tensor& sin_h,
87 const ::brotensor::Tensor& cos_w,
88 const ::brotensor::Tensor& sin_w,
89 const int32_t* pos_t, const int32_t* pos_h,
90 const int32_t* pos_w,
91 int head_dim, int num_heads,
92 int d_t, int d_h, int d_w,
93 ::brotensor::Tensor& Y) {
94
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (head_dim <= 0 || (head_dim & 1) != 0) {
95 ✗ throw std::runtime_error("rope_apply_mrope: head_dim must be a "
96 "positive even integer");
97 }
98
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (num_heads <= 0) {
99 ✗ throw std::runtime_error("rope_apply_mrope: num_heads must be positive");
100 }
101
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (X.dtype != Dtype::FP32) {
102 ✗ throw std::runtime_error("rope_apply_mrope: X must be FP32 "
103 "(CPU backend is FP32-only)");
104 }
105
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (X.cols != num_heads * head_dim) {
106 ✗ throw std::runtime_error("rope_apply_mrope: X.cols != "
107 "num_heads * head_dim");
108 }
109
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 1 time.
 7 if (2 * (d_t + d_h + d_w) != head_dim) {
110
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 throw std::runtime_error("rope_apply_mrope: 2*(d_t + d_h + d_w) "
111 "must equal head_dim");
112 }
113 6 check_axis_tbl(cos_t, sin_t, "t", d_t);
114 6 check_axis_tbl(cos_h, sin_h, "h", d_h);
115 6 check_axis_tbl(cos_w, sin_w, "w", d_w);
116 6 const int L = X.rows;
117
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (Y.rows != L || Y.cols != X.cols || Y.dtype != Dtype::FP32) {
118 6 Y.resize(L, X.cols, Dtype::FP32);
119 6 }
120
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 6 if (L == 0 || num_heads == 0 || head_dim == 0) return;
121
2/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (d_t > 0 && !pos_t) throw std::runtime_error("rope_apply_mrope: pos_t null");
122
3/6
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (d_h > 0 && !pos_h) throw std::runtime_error("rope_apply_mrope: pos_h null");
123
3/6
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (d_w > 0 && !pos_w) throw std::runtime_error("rope_apply_mrope: pos_w null");
124
125 6 const int D = num_heads * head_dim;
126 6 const float* Xp = X.host_f32();
127 6 float* Yp = Y.host_f32_mut();
128
129 // Y may not alias X here in general (output adopt allocates fresh), but
130 // copy first so any pair we skip (e.g. d_a=0 sub-range) is identity-copied.
131
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (Xp != Yp) {
132 6 const int n = L * D;
133
2/2
✓ Branch 0 taken 10176 times.
✓ Branch 1 taken 6 times.
 10182 for (int i = 0; i < n; ++i) Yp[i] = Xp[i];
134 6 }
135
136 // Sub-range pair offsets within each head: [0, d_t) for t, then
137 // [d_t, d_t + d_h) for h, then [d_t + d_h, d_t + d_h + d_w) for w.
138 12 apply_axis(Xp, Yp, cos_t.host_f32(), sin_t.host_f32(),
139
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 pos_t, d_t > 0 ? cos_t.rows : 0,
140 6 L, D, head_dim, num_heads, /*pair_off=*/0, d_t);
141 12 apply_axis(Xp, Yp, cos_h.host_f32(), sin_h.host_f32(),
142
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 pos_h, d_h > 0 ? cos_h.rows : 0,
143 6 L, D, head_dim, num_heads, /*pair_off=*/d_t, d_h);
144 12 apply_axis(Xp, Yp, cos_w.host_f32(), sin_w.host_f32(),
145
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 pos_w, d_w > 0 ? cos_w.rows : 0,
146 6 L, D, head_dim, num_heads, /*pair_off=*/d_t + d_h, d_w);
147 6 }
148
149 } // namespace brotensor::detail::cpu
150