src/cpu/attention_moments.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // CPU backend — attention_token_moments (CHUNK 5). | ||
| 2 | // | ||
| 3 | // Ground truth: src/cuda/attention_moments.cu. | ||
| 4 | // | ||
| 5 | // Per-text-token spatial moments of a cross-attention map. Attn is (Lq, Lk) | ||
| 6 | // where Lq = h_lat * w_lat spatial query tokens and Lk text tokens. For each | ||
| 7 | // text token k it computes: | ||
| 8 | // mass[k] = sum_q Attn(q, k) | ||
| 9 | // centroid[k,0] = (sum_q y(q) * Attn(q, k)) / mass[k] (y = q / w_lat) | ||
| 10 | // centroid[k,1] = (sum_q x(q) * Attn(q, k)) / mass[k] (x = q % w_lat) | ||
| 11 | // If mass[k] <= MOM_EPS, both centroid components are set to 0. | ||
| 12 | // | ||
| 13 | // DTYPE: the CUDA kernel reads FP16 Attn and writes FP32 mass/centroid. The | ||
| 14 | // CPU backend is FP32-only — it reads FP32 Attn. The parity test quantises | ||
| 15 | // the attention map through FP16 so CPU (FP32) and GPU (FP16) start from the | ||
| 16 | // same input bit patterns, then compares with a loose tolerance (the GPU's | ||
| 17 | // long FP16 reduction over up to ~1024 spatial tokens accumulates noise). | ||
| 18 | // | ||
| 19 | // mass and centroid are OVERWRITTEN (computed fresh, not accumulated). | ||
| 20 | |||
| 21 | #include <brotensor/tensor.h> | ||
| 22 | |||
| 23 | #include <stdexcept> | ||
| 24 | |||
| 25 | namespace brotensor::detail::cpu { | ||
| 26 | |||
| 27 | namespace { | ||
| 28 | constexpr float MOM_EPS = 1e-8f; | ||
| 29 | } | ||
| 30 | |||
| 31 | 5 | void attention_token_moments(const ::brotensor::Tensor& Attn, | |
| 32 | int h_lat, int w_lat, | ||
| 33 | ::brotensor::Tensor& mass, | ||
| 34 | ::brotensor::Tensor& centroid) { | ||
| 35 | using ::brotensor::Dtype; | ||
| 36 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (h_lat <= 0 || w_lat <= 0) { |
| 37 | ✗ | throw std::runtime_error( | |
| 38 | "attention_token_moments: h_lat and w_lat must be positive"); | ||
| 39 | } | ||
| 40 | 5 | const int Lq = h_lat * w_lat; | |
| 41 | 5 | const int Lk = Attn.cols; | |
| 42 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (Attn.rows != Lq) { |
| 43 | ✗ | throw std::runtime_error( | |
| 44 | "attention_token_moments: Attn.rows must equal h_lat * w_lat"); | ||
| 45 | } | ||
| 46 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
5 | if (mass.rows != Lk || mass.cols != 1 || mass.dtype != Dtype::FP32) |
| 47 | 5 | mass.resize(Lk, 1, Dtype::FP32); | |
| 48 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
5 | if (centroid.rows != Lk || centroid.cols != 2 || centroid.dtype != Dtype::FP32) |
| 49 | 5 | centroid.resize(Lk, 2, Dtype::FP32); | |
| 50 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (Lk == 0) return; |
| 51 | |||
| 52 | 5 | const float* Ap = Attn.host_f32(); | |
| 53 | 5 | float* Mp = mass.host_f32_mut(); | |
| 54 | 5 | float* Cp = centroid.host_f32_mut(); | |
| 55 | |||
| 56 |
2/2✓ Branch 0 taken 5 times.
✓ Branch 1 taken 108 times.
|
113 | for (int k = 0; k < Lk; ++k) { |
| 57 | 108 | float am = 0.0f, ay = 0.0f, ax = 0.0f; | |
| 58 |
2/2✓ Branch 0 taken 83248 times.
✓ Branch 1 taken 108 times.
|
83356 | for (int q = 0; q < Lq; ++q) { |
| 59 | 83248 | const float a = Ap[static_cast<std::size_t>(q) * Lk + k]; | |
| 60 | 83248 | const int y = q / w_lat; | |
| 61 | 83248 | const int x = q - y * w_lat; | |
| 62 | 83248 | am += a; | |
| 63 | 83248 | ay += static_cast<float>(y) * a; | |
| 64 | 83248 | ax += static_cast<float>(x) * a; | |
| 65 | 83248 | } | |
| 66 | 108 | Mp[k] = am; | |
| 67 |
1/2✓ Branch 0 taken 108 times.
✗ Branch 1 not taken.
|
108 | if (am > MOM_EPS) { |
| 68 | 108 | const float inv = 1.0f / am; | |
| 69 | 108 | Cp[static_cast<std::size_t>(k) * 2 + 0] = ay * inv; | |
| 70 | 108 | Cp[static_cast<std::size_t>(k) * 2 + 1] = ax * inv; | |
| 71 | 108 | } else { | |
| 72 | ✗ | Cp[static_cast<std::size_t>(k) * 2 + 0] = 0.0f; | |
| 73 | ✗ | Cp[static_cast<std::size_t>(k) * 2 + 1] = 0.0f; | |
| 74 | } | ||
| 75 | 108 | } | |
| 76 | 5 | } | |
| 77 | |||
| 78 | } // namespace brotensor::detail::cpu | ||
| 79 |