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 | #include "cpu_jit.h" | ||
| 16 | |||
| 17 | #include <cmath> | ||
| 18 | #include <stdexcept> | ||
| 19 | |||
| 20 | namespace brotensor::detail::cpu { | ||
| 21 | |||
| 22 | namespace { | ||
| 23 | |||
| 24 | ✗ | inline float silu_scalar(float v) { | |
| 25 | ✗ | return v / (1.0f + std::exp(-v)); | |
| 26 | } | ||
| 27 | |||
| 28 | // Backward needs both silu(v) and silu'(v) for the same v; both derive from | ||
| 29 | // a single sigmoid evaluation (which itself needs one std::exp), so compute | ||
| 30 | // it once here instead of calling silu_scalar plus a separate gradient | ||
| 31 | // function (which would redo the std::exp). | ||
| 32 | 583 | inline void silu_value_grad(float v, float& s_out, float& sprime_out) { | |
| 33 | 583 | const float s = 1.0f / (1.0f + std::exp(-v)); | |
| 34 | 583 | s_out = v * s; | |
| 35 | 583 | sprime_out = s * (1.0f + v * (1.0f - s)); | |
| 36 | 583 | } | |
| 37 | |||
| 38 | } // namespace | ||
| 39 | |||
| 40 | 247 | void swiglu_forward(const ::brotensor::Tensor& X, ::brotensor::Tensor& Y) { | |
| 41 |
1/2✓ Branch 0 taken 247 times.
✗ Branch 1 not taken.
|
247 | if (X.cols % 2 != 0) { |
| 42 | ✗ | throw std::runtime_error("swiglu_forward: X.cols must be even (2*D)"); | |
| 43 | } | ||
| 44 | 247 | const int B = X.rows; | |
| 45 | 247 | const int D = X.cols / 2; | |
| 46 |
3/4✓ Branch 0 taken 240 times.
✓ Branch 1 taken 7 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 240 times.
|
247 | if (Y.rows != B || Y.cols != D) Y.resize(B, D); |
| 47 | 247 | const int total = B * D; | |
| 48 |
1/2✓ Branch 0 taken 247 times.
✗ Branch 1 not taken.
|
247 | if (total == 0) return; |
| 49 | 247 | const float* Xp = X.host_f32(); | |
| 50 | 247 | float* Yp = Y.host_f32_mut(); | |
| 51 | |||
| 52 |
1/2✓ Branch 0 taken 247 times.
✗ Branch 1 not taken.
|
247 | if (jit::is_jit_available()) { |
| 53 | 247 | jit::swiglu_forward(Xp, Yp, B, D); | |
| 54 | 247 | return; | |
| 55 | } | ||
| 56 | ✗ | const int two_d = 2 * D; | |
| 57 | ✗ | for (int b = 0; b < B; ++b) { | |
| 58 | ✗ | for (int d = 0; d < D; ++d) { | |
| 59 | ✗ | const float a = Xp[b * two_d + d]; | |
| 60 | ✗ | const float bh = Xp[b * two_d + D + d]; | |
| 61 | ✗ | Yp[b * D + d] = silu_scalar(a) * bh; | |
| 62 | ✗ | } | |
| 63 | ✗ | } | |
| 64 | 247 | } | |
| 65 | |||
| 66 | 5 | void swiglu_backward(const ::brotensor::Tensor& X, const ::brotensor::Tensor& dY, | |
| 67 | ::brotensor::Tensor& dX) { | ||
| 68 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (X.cols % 2 != 0) { |
| 69 | ✗ | throw std::runtime_error("swiglu_backward: X.cols must be even (2*D)"); | |
| 70 | } | ||
| 71 | 5 | const int B = X.rows; | |
| 72 | 5 | const int D = X.cols / 2; | |
| 73 |
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); |
| 74 | 5 | const int total = B * D; | |
| 75 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (total == 0) return; |
| 76 | 5 | const float* Xp = X.host_f32(); | |
| 77 | 5 | const float* dYp = dY.host_f32(); | |
| 78 | 5 | float* dXp = dX.host_f32_mut(); | |
| 79 | 5 | const int two_d = 2 * D; | |
| 80 |
2/2✓ Branch 0 taken 27 times.
✓ Branch 1 taken 5 times.
|
32 | for (int b = 0; b < B; ++b) { |
| 81 |
2/2✓ Branch 0 taken 583 times.
✓ Branch 1 taken 27 times.
|
610 | for (int d = 0; d < D; ++d) { |
| 82 | 583 | const float a = Xp[b * two_d + d]; | |
| 83 | 583 | const float bh = Xp[b * two_d + D + d]; | |
| 84 | 583 | const float dy = dYp[b * D + d]; | |
| 85 | float s, sp; | ||
| 86 | 583 | silu_value_grad(a, s, sp); | |
| 87 | 583 | dXp[b * two_d + d] = dy * bh * sp; | |
| 88 | 583 | dXp[b * two_d + D + d] = dy * s; | |
| 89 | 583 | } | |
| 90 | 27 | } | |
| 91 | 5 | } | |
| 92 | |||
| 93 | } // namespace brotensor::detail::cpu | ||
| 94 |