File size: 3,108 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 | // ============================================================
// fsl_mamba_step.cpp — Basic SSM state transition kernel
// ============================================================
// Implements: s_{t+1} = A * s_t + B * u_t
// Fixed A, B matrices (non-selective version).
//
// YAML parameters:
// d_state = 16 (n)
// d_model = 512 (m)
//
// This is a hand-rolled C implementation targeting the lowest
// publicly inspectable layer (C/C++ callable kernel).
#include <cstddef>
#include <cstring>
#include <cmath>
// ============================================================
// Basic Mamba Step: s_{t+1} = A * s_t + B * u_t
// ============================================================
extern "C" void fsl_mamba_step(
const float* state, // [n] current SSM state
const float* input, // [m] convolved input
const float* matrix_a, // [n*n] state matrix A (row-major)
const float* matrix_b, // [n*m] input matrix B (row-major)
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
) {
// Accumulator for A*s + B*u
float acc[n];
std::memset(acc, 0, n * sizeof(float));
// Compute v1 = A * state
for (size_t i = 0; i < n; ++i) {
for (size_t j = 0; j < n; ++j) {
acc[i] += matrix_a[i * n + j] * state[j];
}
}
// Compute v2 = B * input and accumulate into acc
for (size_t i = 0; i < n; ++i) {
for (size_t j = 0; j < m; ++j) {
acc[i] += matrix_b[i * m + j] * input[j];
}
}
// Store next_state = A*s + B*u
std::memcpy(next_state, acc, n * sizeof(float));
// Zero output (per FSTK semantics)
std::memset(output, 0, m * sizeof(float));
}
// ============================================================
// Vector operations for testing
// ============================================================
extern "C" void fsl_vec_add(
const float* a,
const float* b,
float* result,
size_t n
) {
for (size_t i = 0; i < n; ++i) {
result[i] = a[i] + b[i];
}
}
extern "C" void fsl_vec_scale(
const float* a,
float scalar,
float* result,
size_t n
) {
for (size_t i = 0; i < n; ++i) {
result[i] = a[i] * scalar;
}
}
extern "C" float fsl_vec_norm(
const float* a,
size_t n
) {
float sum = 0.0f;
for (size_t i = 0; i < n; ++i) {
sum += a[i] * a[i];
}
return std::sqrt(sum);
}
// ============================================================
// Matrix-vector multiply (for testing)
// ============================================================
extern "C" void fsl_matvec(
const float* matrix, // [n*m] row-major
const float* vec, // [m]
float* result, // [n]
size_t n,
size_t m
) {
for (size_t i = 0; i < n; ++i) {
result[i] = 0.0f;
for (size_t j = 0; j < m; ++j) {
result[i] += matrix[i * m + j] * vec[j];
}
}
}
|