src/cpu/swiglu.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU SwiGLU ops (CHUNK 2) ────────────────────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementations. Ports src/cuda/swiglu.cu — kernel math | ||
| 4 | // reproduced verbatim, FP32 path only. | ||
| 5 | // | ||
| 6 | // X is (B, 2D); split along the last dim into A = X[:, :D] (the value half | ||
| 7 | // that is gated through silu) and B_half = X[:, D:] (the linear half). | ||
| 8 | // forward: Y(B, D) = silu(A) * B_half | ||
| 9 | // backward: dX(B,2D): dX[:, :D] = dY * B_half * silu'(A) | ||
| 10 | // dX[:, D:] = dY * silu(A) | ||
| 11 | // | ||
| 12 | // Backward overwrites dX (the GPU kernel writes both halves directly). | ||
| 13 | |||
| 14 | #include <brotensor/tensor.h> | ||
| 15 | |||
| 16 | #include <cmath> | ||
| 17 | #include <stdexcept> | ||
| 18 | |||
| 19 | namespace brotensor::detail::cpu { | ||
| 20 | |||
| 21 | namespace { | ||
| 22 | |||
| 23 | 583 | inline float silu_scalar(float v) { | |
| 24 | 583 | return v / (1.0f + std::exp(-v)); | |
| 25 | } | ||
| 26 | |||
| 27 | // Backward needs both silu(v) and silu'(v) for the same v; both derive from | ||
| 28 | // a single sigmoid evaluation (which itself needs one std::exp), so compute | ||
| 29 | // it once here instead of calling silu_scalar plus a separate gradient | ||
| 30 | // function (which would redo the std::exp). | ||
| 31 | 583 | inline void silu_value_grad(float v, float& s_out, float& sprime_out) { | |
| 32 | 583 | const float s = 1.0f / (1.0f + std::exp(-v)); | |
| 33 | 583 | s_out = v * s; | |
| 34 | 583 | sprime_out = s * (1.0f + v * (1.0f - s)); | |
| 35 | 583 | } | |
| 36 | |||
| 37 | } // namespace | ||
| 38 | |||
| 39 | 5 | void swiglu_forward(const ::brotensor::Tensor& X, ::brotensor::Tensor& Y) { | |
| 40 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (X.cols % 2 != 0) { |
| 41 | ✗ | throw std::runtime_error("swiglu_forward: X.cols must be even (2*D)"); | |
| 42 | } | ||
| 43 | 5 | const int B = X.rows; | |
| 44 | 5 | const int D = X.cols / 2; | |
| 45 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
5 | if (Y.rows != B || Y.cols != D) Y.resize(B, D); |
| 46 | 5 | const int total = B * D; | |
| 47 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (total == 0) return; |
| 48 | 5 | const float* Xp = X.host_f32(); | |
| 49 | 5 | float* Yp = Y.host_f32_mut(); | |
| 50 | 5 | const int two_d = 2 * D; | |
| 51 |
2/2✓ Branch 0 taken 27 times.
✓ Branch 1 taken 5 times.
|
32 | for (int b = 0; b < B; ++b) { |
| 52 |
2/2✓ Branch 0 taken 583 times.
✓ Branch 1 taken 27 times.
|
610 | for (int d = 0; d < D; ++d) { |
| 53 | 583 | const float a = Xp[b * two_d + d]; | |
| 54 | 583 | const float bh = Xp[b * two_d + D + d]; | |
| 55 | 583 | Yp[b * D + d] = silu_scalar(a) * bh; | |
| 56 | 583 | } | |
| 57 | 27 | } | |
| 58 | 5 | } | |
| 59 | |||
| 60 | 5 | void swiglu_backward(const ::brotensor::Tensor& X, const ::brotensor::Tensor& dY, | |
| 61 | ::brotensor::Tensor& dX) { | ||
| 62 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (X.cols % 2 != 0) { |
| 63 | ✗ | throw std::runtime_error("swiglu_backward: X.cols must be even (2*D)"); | |
| 64 | } | ||
| 65 | 5 | const int B = X.rows; | |
| 66 | 5 | const int D = X.cols / 2; | |
| 67 |
1/4✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
5 | if (dX.rows != B || dX.cols != 2 * D) dX.resize(B, 2 * D); |
| 68 | 5 | const int total = B * D; | |
| 69 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (total == 0) return; |
| 70 | 5 | const float* Xp = X.host_f32(); | |
| 71 | 5 | const float* dYp = dY.host_f32(); | |
| 72 | 5 | float* dXp = dX.host_f32_mut(); | |
| 73 | 5 | const int two_d = 2 * D; | |
| 74 |
2/2✓ Branch 0 taken 27 times.
✓ Branch 1 taken 5 times.
|
32 | for (int b = 0; b < B; ++b) { |
| 75 |
2/2✓ Branch 0 taken 583 times.
✓ Branch 1 taken 27 times.
|
610 | for (int d = 0; d < D; ++d) { |
| 76 | 583 | const float a = Xp[b * two_d + d]; | |
| 77 | 583 | const float bh = Xp[b * two_d + D + d]; | |
| 78 | 583 | const float dy = dYp[b * D + d]; | |
| 79 | float s, sp; | ||
| 80 | 583 | silu_value_grad(a, s, sp); | |
| 81 | 583 | dXp[b * two_d + d] = dy * bh * sp; | |
| 82 | 583 | dXp[b * two_d + D + d] = dy * s; | |
| 83 | 583 | } | |
| 84 | 27 | } | |
| 85 | 5 | } | |
| 86 | |||
| 87 | } // namespace brotensor::detail::cpu | ||
| 88 |