src/cpu/self_attention_inference.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // CPU backend — self_attention_forward (CHUNK 6). | ||
| 2 | // | ||
| 3 | // Ground truth: src/cuda/cross_attention.cu — self_attention_forward. | ||
| 4 | // The CUDA op runs the FP16 flash path for FP16 inputs (delegating to | ||
| 5 | // flash_attention_qkvo_forward, self-attention, no biases, non-causal) and | ||
| 6 | // the FP32 mha_forward train-core for FP32 inputs. Either way the underlying | ||
| 7 | // math is identical to multi-head self-attention. | ||
| 8 | // | ||
| 9 | // The CPU backend is FP32-only (per CLAUDE.md). This op delegates straight to | ||
| 10 | // the existing CPU mha_forward (ops_impl.cpp) — the same FP32 train-core the | ||
| 11 | // CUDA op uses on its FP32 path — discarding the per-head intermediates that | ||
| 12 | // the forward-only public surface does not expose. | ||
| 13 | // | ||
| 14 | // DTYPE: GPU runs FP16 internally, CPU runs FP32. The parity test quantises | ||
| 15 | // inputs through FP16 so both backends start from identical values and | ||
| 16 | // compares with a loose FP16-scale tolerance. | ||
| 17 | // | ||
| 18 | // CONVENTIONS (inherited from the CPU mha impl, matching the CUDA mha | ||
| 19 | // kernels): Wq/Wk/Wv/Wo are all (D, D); per-head split takes contiguous | ||
| 20 | // weight rows hh*dh..hh*dh+dh; softmax scale 1/sqrt(dh); mask is a length-L | ||
| 21 | // key-validity buffer. | ||
| 22 | |||
| 23 | #include <brotensor/tensor.h> | ||
| 24 | |||
| 25 | namespace brotensor::detail::cpu { | ||
| 26 | |||
| 27 | // Implemented in ops_impl.cpp. | ||
| 28 | void mha_forward(const ::brotensor::Tensor& X, | ||
| 29 | const ::brotensor::Tensor& Wq, const ::brotensor::Tensor& Wk, | ||
| 30 | const ::brotensor::Tensor& Wv, const ::brotensor::Tensor& Wo, | ||
| 31 | const ::brotensor::Tensor* bq, const ::brotensor::Tensor* bk, | ||
| 32 | const ::brotensor::Tensor* bv, const ::brotensor::Tensor* bo, | ||
| 33 | const float* d_mask, int num_heads, | ||
| 34 | ::brotensor::Tensor& Qh, ::brotensor::Tensor& Kh, | ||
| 35 | ::brotensor::Tensor& Vh, ::brotensor::Tensor& Attnh, | ||
| 36 | ::brotensor::Tensor& Yconcat, ::brotensor::Tensor& O); | ||
| 37 | |||
| 38 | 3 | void self_attention_forward(const ::brotensor::Tensor& X, | |
| 39 | const ::brotensor::Tensor& Wq, | ||
| 40 | const ::brotensor::Tensor& Wk, | ||
| 41 | const ::brotensor::Tensor& Wv, | ||
| 42 | const ::brotensor::Tensor& Wo, | ||
| 43 | const float* d_mask, | ||
| 44 | int num_heads, | ||
| 45 | ::brotensor::Tensor& O) { | ||
| 46 | 3 | ::brotensor::Tensor Qh, Kh, Vh, Attnh, Yconcat; | |
| 47 |
2/4✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
|
6 | mha_forward(X, Wq, Wk, Wv, Wo, |
| 48 | nullptr, nullptr, nullptr, nullptr, | ||
| 49 | 3 | d_mask, num_heads, | |
| 50 | 3 | Qh, Kh, Vh, Attnh, Yconcat, O); | |
| 51 | 3 | } | |
| 52 | |||
| 53 | } // namespace brotensor::detail::cpu | ||
| 54 |