src/cpu/self_attention.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // CPU backend — self-attention training ops (CHUNK 5). | ||
| 2 | // | ||
| 3 | // Ground truth: src/cuda/cross_attention.cu — self_attention_forward_train | ||
| 4 | // and self_attention_backward are thin wrappers over mha_forward / | ||
| 5 | // mha_backward (the self-attention degenerate case: Lq == Lk, Ctx == X, | ||
| 6 | // D_ctx == D). The CPU backend follows the same delegation: these forward | ||
| 7 | // straight to the existing CPU mha_forward / mha_backward in ops_impl.cpp. | ||
| 8 | // | ||
| 9 | // DTYPE: FP32 on both backends — straightforward FP32<->FP32 parity. | ||
| 10 | // | ||
| 11 | // CONVENTIONS (inherited from the CPU mha impl, which matches the CUDA mha | ||
| 12 | // kernels): | ||
| 13 | // * Weight layout: Wq/Wk/Wv/Wo are all (D, D); per-head projection takes | ||
| 14 | // contiguous weight rows hh*dh..hh*dh+dh. | ||
| 15 | // * Per-head split: Qh/Kh/Vh are (H*L, dh) row-major by (head, token); | ||
| 16 | // Attnh is (H*L, L); Yconcat is (L, D) with head hh in columns | ||
| 17 | // hh*dh..hh*dh+dh; O is (L, D). | ||
| 18 | // * Softmax scale 1/sqrt(dh); mask is length-L (self-attention so a single | ||
| 19 | // buffer gates both query rows and key columns). | ||
| 20 | // * Backward: dWq/dWk/dWv/dWo ACCUMULATE (+=); dX is OVERWRITTEN. | ||
| 21 | |||
| 22 | #include <brotensor/tensor.h> | ||
| 23 | |||
| 24 | namespace brotensor::detail::cpu { | ||
| 25 | |||
| 26 | // Implemented in ops_impl.cpp. | ||
| 27 | void mha_forward(const ::brotensor::Tensor& X, | ||
| 28 | const ::brotensor::Tensor& Wq, const ::brotensor::Tensor& Wk, | ||
| 29 | const ::brotensor::Tensor& Wv, const ::brotensor::Tensor& Wo, | ||
| 30 | const ::brotensor::Tensor* bq, const ::brotensor::Tensor* bk, | ||
| 31 | const ::brotensor::Tensor* bv, const ::brotensor::Tensor* bo, | ||
| 32 | const float* d_mask, int num_heads, | ||
| 33 | ::brotensor::Tensor& Qh, ::brotensor::Tensor& Kh, | ||
| 34 | ::brotensor::Tensor& Vh, ::brotensor::Tensor& Attnh, | ||
| 35 | ::brotensor::Tensor& Yconcat, ::brotensor::Tensor& O); | ||
| 36 | void mha_backward(const ::brotensor::Tensor& dO, | ||
| 37 | const ::brotensor::Tensor& X, | ||
| 38 | const ::brotensor::Tensor& Qh, const ::brotensor::Tensor& Kh, | ||
| 39 | const ::brotensor::Tensor& Vh, const ::brotensor::Tensor& Attnh, | ||
| 40 | const ::brotensor::Tensor& Yconcat, | ||
| 41 | const ::brotensor::Tensor& Wq, const ::brotensor::Tensor& Wk, | ||
| 42 | const ::brotensor::Tensor& Wv, const ::brotensor::Tensor& Wo, | ||
| 43 | const float* d_mask, int num_heads, | ||
| 44 | ::brotensor::Tensor& dX, | ||
| 45 | ::brotensor::Tensor& dWq, ::brotensor::Tensor& dWk, | ||
| 46 | ::brotensor::Tensor& dWv, ::brotensor::Tensor& dWo, | ||
| 47 | ::brotensor::Tensor* dbq, ::brotensor::Tensor* dbk, | ||
| 48 | ::brotensor::Tensor* dbv, ::brotensor::Tensor* dbo); | ||
| 49 | |||
| 50 | 9 | void self_attention_forward_train(const ::brotensor::Tensor& X, | |
| 51 | const ::brotensor::Tensor& Wq, | ||
| 52 | const ::brotensor::Tensor& Wk, | ||
| 53 | const ::brotensor::Tensor& Wv, | ||
| 54 | const ::brotensor::Tensor& Wo, | ||
| 55 | const float* d_mask, | ||
| 56 | int num_heads, | ||
| 57 | ::brotensor::Tensor& Qh, | ||
| 58 | ::brotensor::Tensor& Kh, | ||
| 59 | ::brotensor::Tensor& Vh, | ||
| 60 | ::brotensor::Tensor& Attnh, | ||
| 61 | ::brotensor::Tensor& Yconcat, | ||
| 62 | ::brotensor::Tensor& O) { | ||
| 63 | 18 | mha_forward(X, Wq, Wk, Wv, Wo, | |
| 64 | nullptr, nullptr, nullptr, nullptr, | ||
| 65 | 9 | d_mask, num_heads, | |
| 66 | 9 | Qh, Kh, Vh, Attnh, Yconcat, O); | |
| 67 | 9 | } | |
| 68 | |||
| 69 | 7 | void self_attention_backward(const ::brotensor::Tensor& dO, | |
| 70 | const ::brotensor::Tensor& X, | ||
| 71 | const ::brotensor::Tensor& Qh, | ||
| 72 | const ::brotensor::Tensor& Kh, | ||
| 73 | const ::brotensor::Tensor& Vh, | ||
| 74 | const ::brotensor::Tensor& Attnh, | ||
| 75 | const ::brotensor::Tensor& Yconcat, | ||
| 76 | const ::brotensor::Tensor& Wq, | ||
| 77 | const ::brotensor::Tensor& Wk, | ||
| 78 | const ::brotensor::Tensor& Wv, | ||
| 79 | const ::brotensor::Tensor& Wo, | ||
| 80 | const float* d_mask, | ||
| 81 | int num_heads, | ||
| 82 | ::brotensor::Tensor& dX, | ||
| 83 | ::brotensor::Tensor& dWq, | ||
| 84 | ::brotensor::Tensor& dWk, | ||
| 85 | ::brotensor::Tensor& dWv, | ||
| 86 | ::brotensor::Tensor& dWo) { | ||
| 87 | 14 | mha_backward(dO, X, Qh, Kh, Vh, Attnh, Yconcat, | |
| 88 | 7 | Wq, Wk, Wv, Wo, d_mask, num_heads, | |
| 89 | 7 | dX, dWq, dWk, dWv, dWo, | |
| 90 | nullptr, nullptr, nullptr, nullptr); | ||
| 91 | 7 | } | |
| 92 | |||
| 93 | } // namespace brotensor::detail::cpu | ||
| 94 |