File size: 5,938 Bytes
e92f76f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | // ============================================================
// fsl_selective_mamba_step.cpp — Selective SSM (Mamba-2)
// ============================================================
// Implements the selective state-space model step:
// 1. Depthwise convolution: z_t = Conv_{d_c}(x_t; W)
// 2. Split and SiLU gating: u_t = z1 ⊙ silu(z2)
// 3. SSM update: s_{t+1} = A * s_t + B * u_t
// 4. Zero output: y_t = 0_m
//
// YAML parameters:
// d_state = 16 (n)
// d_model = 512 (m)
// d_conv = 4 (d_c)
//
// A is diagonal: A = diag(-exp(A_log))
// B is fixed (full n×m matrix)
#include <cstddef>
#include <cstring>
#include <cmath>
// ============================================================
// SiLU activation: silu(x) = x * sigmoid(x)
// ============================================================
static inline float silu(float x) {
return x / (1.0f + std::exp(-x));
}
// ============================================================
// Depthwise 1D convolution (causal padding)
// ============================================================
static void depthwise_conv1d(
const float* input, // [m] input signal
const float* W_conv, // [m * d_c] convolution weights
float* output, // [m] output signal
size_t m, // model dimension
size_t d_c // convolution width
) {
for (size_t i = 0; i < m; ++i) {
float sum = 0.0f;
for (size_t k = 0; k < d_c; ++k) {
// Causal padding: pad with zeros on the left
size_t idx = i + k - (d_c / 2);
float x_val = (idx < m) ? input[idx] : 0.0f;
sum += W_conv[i * d_c + k] * x_val;
}
output[i] = sum;
}
}
// ============================================================
// Selective Mamba Step
// ============================================================
extern "C" void fsl_selective_mamba_step(
const float* state, // [n] current SSM state
const float* input, // [m] raw token (pre-convolution)
const float* A_log, // [n] log-space diagonal matrix
const float* B_full, // [n*m] full input matrix (row-major)
const float* W_conv, // [m*d_c] depthwise conv kernel
float* next_state, // [n] next SSM state (output)
float* output, // [m] intermediate output (zeroed)
size_t n, // d_state = 16
size_t m, // d_model = 512
size_t d_c // d_conv = 4
) {
// Temporary buffers (stack-allocated for small sizes)
float z[m]; // Conv output
float z1[m / 2]; // First half
float z2[m / 2]; // Second half
float u[m]; // Selective input
float As[n]; // A * state
float Bu[n]; // B * u
// Step 1: Depthwise convolution
depthwise_conv1d(input, W_conv, z, m, d_c);
// Step 2: Split and apply SiLU gating (selectivity)
std::memcpy(z1, z, (m / 2) * sizeof(float));
std::memcpy(z2, z + (m / 2), (m / 2) * sizeof(float));
// u = z1 ⊙ silu(z2)
for (size_t i = 0; i < m / 2; ++i) {
u[i] = z1[i] * silu(z2[i]);
}
// Zero-pad u to full size m
std::memset(u + (m / 2), 0, (m / 2) * sizeof(float));
// Step 3: Compute A * state (A = diag(-exp(A_log)))
for (size_t i = 0; i < n; ++i) {
As[i] = -std::exp(A_log[i]) * state[i];
}
// Step 4: Compute B * u (B_full is n×m, row-major)
std::memset(Bu, 0, n * sizeof(float));
for (size_t i = 0; i < n; ++i) {
for (size_t j = 0; j < m; ++j) {
Bu[i] += B_full[i * m + j] * u[j];
}
}
// Step 5: State update: s_{t+1} = A*s + B*u
for (size_t i = 0; i < n; ++i) {
next_state[i] = As[i] + Bu[i];
}
// Step 6: Zero output (per FSTK semantics)
std::memset(output, 0, m * sizeof(float));
}
// ============================================================
// Output projection: y_t = C * s_t + D * u_t
// ============================================================
extern "C" void fsl_output_projection(
const float* state, // [n] SSM state
const float* input, // [m] convolved input
const float* matrix_c, // [m*n] output matrix C (row-major)
const float* matrix_d, // [m*m] skip matrix D (row-major)
float* output, // [m] output token
size_t n, // d_state = 16
size_t m // d_model = 512
) {
// y = C * s
for (size_t i = 0; i < m; ++i) {
output[i] = 0.0f;
for (size_t j = 0; j < n; ++j) {
output[i] += matrix_c[i * n + j] * state[j];
}
}
// y += D * u
for (size_t i = 0; i < m; ++i) {
for (size_t j = 0; j < m; ++j) {
output[i] += matrix_d[i * m + j] * input[j];
}
}
}
// ============================================================
// FSM transition: evaluate condition and update state
// ============================================================
extern "C" int fsl_fsm_transition(
int from_state, // current FSM state (integer ID)
int to_state, // target FSM state (integer ID)
int condition // boolean condition flag
) {
// If condition is true, transition to to_state
// Otherwise, stay in from_state
return condition ? to_state : from_state;
}
// ============================================================
// Scan complete check: ||s||_2 < epsilon
// ============================================================
extern "C" int fsl_scan_complete(
const float* state, // [n] SSM state
size_t n, // state dimension
float epsilon // convergence threshold
) {
float norm = 0.0f;
for (size_t i = 0; i < n; ++i) {
norm += state[i] * state[i];
}
norm = std::sqrt(norm);
return (norm < epsilon) ? 1 : 0;
}
|