mambacode.js 1.0.0 → 1.0.1

package/dist/kernels/conv1d.js

@@ -0,0 +1,152 @@
+// 1D Causal Convolution WGSL Kernel
+// Implements a depthwise 1D causal convolution over the sequence dimension.
+// "Causal" means the output at position t only depends on positions <= t,
+// which is enforced by left-padding with (kernel_size - 1) zeros.
+//
+// Forward:  y[b, t, d] = sum_{k=0}^{K-1} weight[d, k] * x[b, t-k, d]
+//           where x[b, t', d] = 0 for t' < 0  (causal padding)
+export const CONV1D_FORWARD_WGSL = /* wgsl */ `
+
+struct ConvParams {
+    seq_len     : u32,   // L
+    d_channels  : u32,   // D (number of depthwise channels)
+    kernel_size : u32,   // K (typically 4)
+    batch       : u32,   // B
+};
+
+@group(0) @binding(0) var<uniform>             params   : ConvParams;
+// x      (B, L, D) – input
+@group(0) @binding(1) var<storage, read>       x        : array<f32>;
+// weight (D, K)    – depthwise conv weights
+@group(0) @binding(2) var<storage, read>       weight   : array<f32>;
+// bias   (D,)      – optional bias (zeros if unused)
+@group(0) @binding(3) var<storage, read>       bias     : array<f32>;
+// y      (B, L, D) – output
+@group(0) @binding(4) var<storage, read_write> y        : array<f32>;
+
+// Dispatch: (ceil(L/16), ceil(D/16), B)
+@compute @workgroup_size(16, 16, 1)
+fn conv1d_forward(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let t  = gid.x;   // time position
+    let d  = gid.y;   // channel
+    let b  = gid.z;   // batch
+
+    if (t >= L || d >= D || b >= B) { return; }
+
+    var acc: f32 = 0.0;
+
+    // Causal: convolve over k = 0..K-1, reading position (t - k)
+    for (var k: u32 = 0u; k < K; k = k + 1u) {
+        let w_idx = d * K + k;
+        let w_val = weight[w_idx];
+
+        // t - k: use causal zero-padding for t < k
+        if (t >= k) {
+            let src = b * L * D + (t - k) * D + d;
+            acc = acc + w_val * x[src];
+        }
+        // else: zero-padding contributes 0
+    }
+
+    acc = acc + bias[d];
+
+    let out = b * L * D + t * D + d;
+    y[out] = acc;
+}
+`;
+// ---- Backward kernel for 1D convolution ----
+export const CONV1D_BACKWARD_WGSL = /* wgsl */ `
+
+struct ConvParams {
+    seq_len     : u32,
+    d_channels  : u32,
+    kernel_size : u32,
+    batch       : u32,
+};
+
+@group(0) @binding(0) var<uniform>              params   : ConvParams;
+@group(0) @binding(1) var<storage, read>        x        : array<f32>;
+@group(0) @binding(2) var<storage, read>        weight   : array<f32>;
+@group(0) @binding(3) var<storage, read>        dy       : array<f32>;
+@group(0) @binding(4) var<storage, read_write>  dx       : array<f32>;
+@group(0) @binding(5) var<storage, read_write>  dweight  : array<f32>;
+@group(0) @binding(6) var<storage, read_write>  dbias    : array<f32>;
+
+// Dispatch: (ceil(L/16), ceil(D/16), B) – computes dx
+@compute @workgroup_size(16, 16, 1)
+fn conv1d_backward_dx(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let t  = gid.x;
+    let d  = gid.y;
+    let b  = gid.z;
+
+    if (t >= L || d >= D || b >= B) { return; }
+
+    var grad: f32 = 0.0;
+
+    // dx[b, t, d] = sum_{k=0}^{K-1} dy[b, t+k, d] * weight[d, k]
+    for (var k: u32 = 0u; k < K; k = k + 1u) {
+        let tp = t + k;
+        if (tp < L) {
+            let dy_idx = b * L * D + tp * D + d;
+            let w_idx  = d * K + k;
+            grad = grad + dy[dy_idx] * weight[w_idx];
+        }
+    }
+
+    let dx_idx = b * L * D + t * D + d;
+    dx[dx_idx] = grad;
+}
+
+// Dispatch: (K, D, 1) – accumulates dweight over (B, L)
+@compute @workgroup_size(1, 1, 1)
+fn conv1d_backward_dw(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L  = params.seq_len;
+    let D  = params.d_channels;
+    let K  = params.kernel_size;
+    let B  = params.batch;
+
+    let k  = gid.x;
+    let d  = gid.y;
+
+    if (k >= K || d >= D) { return; }
+
+    var grad_w: f32 = 0.0;
+    var grad_b: f32 = 0.0;
+
+    for (var b: u32 = 0u; b < B; b = b + 1u) {
+        for (var t: u32 = 0u; t < L; t = t + 1u) {
+            let dy_idx = b * L * D + t * D + d;
+            let dy_val = dy[dy_idx];
+            if (t >= k) {
+                let x_idx = b * L * D + (t - k) * D + d;
+                grad_w = grad_w + dy_val * x[x_idx];
+            }
+            if (k == 0u) {
+                grad_b = grad_b + dy_val;
+            }
+        }
+    }
+
+    dweight[d * K + k] = grad_w;
+    if (k == 0u) {
+        dbias[d] = grad_b;
+    }
+}
+`;
+//# sourceMappingURL=conv1d.js.map

package/dist/kernels/conv1d.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"conv1d.js","sourceRoot":"","sources":["../../src/kernels/conv1d.ts"],"names":[],"mappings":"AAAA,oCAAoC;AACpC,4EAA4E;AAC5E,0EAA0E;AAC1E,kEAAkE;AAClE,EAAE;AACF,qEAAqE;AACrE,+DAA+D;AAE/D,MAAM,CAAC,MAAM,mBAAmB,GAAG,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAuD5C,CAAC;AAEF,+CAA+C;AAC/C,MAAM,CAAC,MAAM,oBAAoB,GAAG,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAsF7C,CAAC"}

package/dist/kernels/linear_projection.d.ts

@@ -0,0 +1,3 @@
+export declare const LINEAR_FORWARD_WGSL: string;
+export declare const LINEAR_BACKWARD_WGSL: string;
+//# sourceMappingURL=linear_projection.d.ts.map

package/dist/kernels/linear_projection.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"linear_projection.d.ts","sourceRoot":"","sources":["../../src/kernels/linear_projection.ts"],"names":[],"mappings":"AAUA,eAAO,MAAM,mBAAmB,EAAE,MAmEjC,CAAC;AAGF,eAAO,MAAM,oBAAoB,EAAE,MA2IlC,CAAC"}

package/dist/kernels/linear_projection.js

@@ -0,0 +1,219 @@
+// Linear Projection WGSL Kernel
+// General-purpose matrix multiplication: Y = X @ W^T + b
+// Supports the up-projection and down-projection linear layers in the Mamba block.
+//
+// Shapes:
+//   X : (batch * seq_len, in_features)   – input (rows)
+//   W : (out_features, in_features)      – weight matrix (row-major)
+//   b : (out_features,)                  – bias
+//   Y : (batch * seq_len, out_features)  – output
+export const LINEAR_FORWARD_WGSL = /* wgsl */ `
+
+struct LinearParams {
+    M : u32,   // number of rows    (batch * seq_len)
+    K : u32,   // in_features
+    N : u32,   // out_features
+};
+
+@group(0) @binding(0) var<uniform>             params : LinearParams;
+@group(0) @binding(1) var<storage, read>       X      : array<f32>;   // (M, K)
+@group(0) @binding(2) var<storage, read>       W      : array<f32>;   // (N, K)
+@group(0) @binding(3) var<storage, read>       bias   : array<f32>;   // (N,)
+@group(0) @binding(4) var<storage, read_write> Y      : array<f32>;   // (M, N)
+
+// Tiled matmul using workgroup shared memory (16x16 tiles)
+var<workgroup> tile_X : array<f32, 256>;  // 16 * 16
+var<workgroup> tile_W : array<f32, 256>;
+
+@compute @workgroup_size(16, 16, 1)
+fn linear_forward(
+    @builtin(global_invocation_id)   gid : vec3<u32>,
+    @builtin(local_invocation_id)    lid : vec3<u32>,
+    @builtin(workgroup_id)           wid : vec3<u32>,
+) {
+    let M = params.M;
+    let K = params.K;
+    let N = params.N;
+
+    let row = gid.x;   // output row (M dimension)
+    let col = gid.y;   // output col (N dimension)
+
+    var acc: f32 = 0.0;
+    let TILE: u32 = 16u;
+    let num_tiles = (K + TILE - 1u) / TILE;
+
+    for (var tile_idx: u32 = 0u; tile_idx < num_tiles; tile_idx = tile_idx + 1u) {
+        // Load X tile: shape (TILE_M, TILE_K)
+        let x_col = tile_idx * TILE + lid.y;
+        let x_row = wid.x * TILE + lid.x;
+        if (x_row < M && x_col < K) {
+            tile_X[lid.x * TILE + lid.y] = X[x_row * K + x_col];
+        } else {
+            tile_X[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        // Load W tile: shape (TILE_N, TILE_K)  — W is (N, K)
+        let w_col = tile_idx * TILE + lid.x;  // K dimension
+        let w_row = wid.y * TILE + lid.y;     // N dimension
+        if (w_row < N && w_col < K) {
+            tile_W[lid.y * TILE + lid.x] = W[w_row * K + w_col];
+        } else {
+            tile_W[lid.y * TILE + lid.x] = 0.0;
+        }
+
+        workgroupBarrier();
+
+        // Dot product within tile
+        for (var k: u32 = 0u; k < TILE; k = k + 1u) {
+            acc = acc + tile_X[lid.x * TILE + k] * tile_W[lid.y * TILE + k];
+        }
+        workgroupBarrier();
+    }
+
+    if (row < M && col < N) {
+        Y[row * N + col] = acc + bias[col];
+    }
+}
+`;
+// ---- Backward pass for linear projection ----
+export const LINEAR_BACKWARD_WGSL = /* wgsl */ `
+
+struct LinearParams {
+    M : u32,
+    K : u32,
+    N : u32,
+};
+
+@group(0) @binding(0) var<uniform>             params : LinearParams;
+@group(0) @binding(1) var<storage, read>       X      : array<f32>;   // (M, K)
+@group(0) @binding(2) var<storage, read>       W      : array<f32>;   // (N, K)
+@group(0) @binding(3) var<storage, read>       dY     : array<f32>;   // (M, N)
+@group(0) @binding(4) var<storage, read_write> dX     : array<f32>;   // (M, K)
+@group(0) @binding(5) var<storage, read_write> dW     : array<f32>;   // (N, K)
+@group(0) @binding(6) var<storage, read_write> db     : array<f32>;   // (N,)
+
+// Dispatch: (ceil(M/16), ceil(K/16), 1)  – computes dX = dY @ W
+var<workgroup> tile_dY : array<f32, 256>;
+var<workgroup> tile_W  : array<f32, 256>;
+
+@compute @workgroup_size(16, 16, 1)
+fn linear_backward_dX(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+    @builtin(local_invocation_id)  lid : vec3<u32>,
+    @builtin(workgroup_id)         wid : vec3<u32>,
+) {
+    let M = params.M;
+    let K = params.K;
+    let N = params.N;
+
+    let row = gid.x;   // M
+    let col = gid.y;   // K
+
+    var acc: f32 = 0.0;
+    let TILE: u32 = 16u;
+    let num_tiles = (N + TILE - 1u) / TILE;
+
+    for (var tile_idx: u32 = 0u; tile_idx < num_tiles; tile_idx = tile_idx + 1u) {
+        // tile_dY: (M, TILE_N) slice
+        let dy_col = tile_idx * TILE + lid.y;
+        let dy_row = wid.x * TILE + lid.x;
+        if (dy_row < M && dy_col < N) {
+            tile_dY[lid.x * TILE + lid.y] = dY[dy_row * N + dy_col];
+        } else {
+            tile_dY[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        // tile_W: (TILE_N, K) slice  — W[n, k]
+        let w_row = tile_idx * TILE + lid.x;   // N
+        let w_col = wid.y * TILE + lid.y;      // K
+        if (w_row < N && w_col < K) {
+            tile_W[lid.x * TILE + lid.y] = W[w_row * K + w_col];
+        } else {
+            tile_W[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        workgroupBarrier();
+
+        for (var n: u32 = 0u; n < TILE; n = n + 1u) {
+            acc = acc + tile_dY[lid.x * TILE + n] * tile_W[n * TILE + lid.y];
+        }
+        workgroupBarrier();
+    }
+
+    if (row < M && col < K) {
+        dX[row * K + col] = acc;
+    }
+}
+
+// Dispatch: (ceil(N/16), ceil(K/16), 1)  – computes dW = dY^T @ X
+var<workgroup> tile_dY2 : array<f32, 256>;
+var<workgroup> tile_X2  : array<f32, 256>;
+
+@compute @workgroup_size(16, 16, 1)
+fn linear_backward_dW(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+    @builtin(local_invocation_id)  lid : vec3<u32>,
+    @builtin(workgroup_id)         wid : vec3<u32>,
+) {
+    let M = params.M;
+    let K = params.K;
+    let N = params.N;
+
+    let row = gid.x;   // N
+    let col = gid.y;   // K
+
+    var acc: f32 = 0.0;
+    let TILE: u32 = 16u;
+    let num_tiles = (M + TILE - 1u) / TILE;
+
+    for (var tile_idx: u32 = 0u; tile_idx < num_tiles; tile_idx = tile_idx + 1u) {
+        // dY^T tile: [N, M] accessed as dY[m, n]
+        let m_idx = tile_idx * TILE + lid.y;
+        let n_idx = wid.x * TILE + lid.x;
+        if (n_idx < N && m_idx < M) {
+            tile_dY2[lid.x * TILE + lid.y] = dY[m_idx * N + n_idx];
+        } else {
+            tile_dY2[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        // X tile: [M, K]
+        let xm = tile_idx * TILE + lid.x;
+        let xk = wid.y * TILE + lid.y;
+        if (xm < M && xk < K) {
+            tile_X2[lid.x * TILE + lid.y] = X[xm * K + xk];
+        } else {
+            tile_X2[lid.x * TILE + lid.y] = 0.0;
+        }
+
+        workgroupBarrier();
+
+        for (var m: u32 = 0u; m < TILE; m = m + 1u) {
+            acc = acc + tile_dY2[lid.x * TILE + m] * tile_X2[m * TILE + lid.y];
+        }
+        workgroupBarrier();
+    }
+
+    if (row < N && col < K) {
+        dW[row * K + col] = acc;
+    }
+}
+
+// Dispatch: (N, 1, 1) – accumulates db = sum_M dY
+@compute @workgroup_size(64, 1, 1)
+fn linear_backward_db(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let M = params.M;
+    let N = params.N;
+
+    let n = gid.x;
+    if (n >= N) { return; }
+
+    var acc: f32 = 0.0;
+    for (var m: u32 = 0u; m < M; m = m + 1u) {
+        acc = acc + dY[m * N + n];
+    }
+    db[n] = acc;
+}
+`;
+//# sourceMappingURL=linear_projection.js.map

package/dist/kernels/linear_projection.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"linear_projection.js","sourceRoot":"","sources":["../../src/kernels/linear_projection.ts"],"names":[],"mappings":"AAAA,gCAAgC;AAChC,yDAAyD;AACzD,mFAAmF;AACnF,EAAE;AACF,UAAU;AACV,wDAAwD;AACxD,qEAAqE;AACrE,gDAAgD;AAChD,kDAAkD;AAElD,MAAM,CAAC,MAAM,mBAAmB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAmEpD,CAAC;AAEF,gDAAgD;AAChD,MAAM,CAAC,MAAM,oBAAoB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CA2IrD,CAAC"}

package/dist/kernels/selective_scan.d.ts

@@ -0,0 +1,3 @@
+export declare const SELECTIVE_SCAN_FORWARD_WGSL: string;
+export declare const SELECTIVE_SCAN_BACKWARD_WGSL: string;
+//# sourceMappingURL=selective_scan.d.ts.map

package/dist/kernels/selective_scan.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"selective_scan.d.ts","sourceRoot":"","sources":["../../src/kernels/selective_scan.ts"],"names":[],"mappings":"AAUA,eAAO,MAAM,2BAA2B,EAAE,MAsNzC,CAAC;AAKF,eAAO,MAAM,4BAA4B,EAAE,MAwH1C,CAAC"}