mambacode.js 1.0.0 → 1.0.2

package/dist/kernels/selective_scan.js

@@ -0,0 +1,348 @@
+// Parallel Selective Scan WGSL Kernel
+// Implements the S6 (Selective Scan) core of the Mamba architecture.
+// Uses a Kogge-Stone parallel prefix-sum approach for O(log N) time on GPU.
+//
+// Forward pass recurrence:
+//   h_t = A_t * h_{t-1} + B_t * x_t
+//   y_t = C_t * h_t + D * x_t
+//
+// where A_t, B_t, C_t are input-dependent (selective) gate matrices.
+export const SELECTIVE_SCAN_FORWARD_WGSL = /* wgsl */ `
+
+// ---- Binding layout ----
+// group 0: sequence data
+// group 1: SSM parameters
+
+struct ScanParams {
+    seq_len   : u32,   // L  – sequence length
+    d_state   : u32,   // N  – state dimension
+    d_inner   : u32,   // D  – inner (expanded) channel dimension
+    batch     : u32,   // B  – batch size
+};
+
+@group(0) @binding(0) var<uniform>             params   : ScanParams;
+// u (B, L, D)  – projected input after conv
+@group(0) @binding(1) var<storage, read>       u        : array<f32>;
+// delta (B, L, D) – time-step (Δ) after softplus
+@group(0) @binding(2) var<storage, read>       delta    : array<f32>;
+// A (D, N)  – log-space diagonal state matrix (fixed, learned)
+@group(0) @binding(3) var<storage, read>       A        : array<f32>;
+// B (B, L, N) – input projection (selective)
+@group(0) @binding(4) var<storage, read>       B        : array<f32>;
+// C (B, L, N) – output projection (selective)
+@group(0) @binding(5) var<storage, read>       C        : array<f32>;
+// D (D,) – skip-connection scale
+@group(0) @binding(6) var<storage, read>       D_vec    : array<f32>;
+// y (B, L, D) – output (written by this kernel)
+@group(0) @binding(7) var<storage, read_write> y        : array<f32>;
+// h_cache (B, L, D*N) – hidden states cache (for backward pass)
+@group(0) @binding(8) var<storage, read_write> h_cache  : array<f32>;
+
+// ---- Workgroup shared memory ----
+// Each workgroup processes one (batch, channel) slice across all time steps.
+// We store the associative pair (a_bar, bu_bar) per time step so we can run
+// a Kogge-Stone scan across the workgroup tile.
+var<workgroup> wg_a  : array<f32, 256>;   // discretised A values
+var<workgroup> wg_bu : array<f32, 256>;   // B*u values
+
+// ---- Helpers ----
+
+// Softplus: numerically stable log(1 + exp(x))
+fn softplus(x: f32) -> f32 {
+    return log(1.0 + exp(x));
+}
+
+// ZerO-Order Hold discretisation of continuous A, Δ:
+//   A_bar = exp(Δ * A)
+//   B_bar = (A_bar - 1) / A * B  ≈  Δ * B  (first-order for simplicity)
+fn discretise_A(delta_val: f32, a_log: f32) -> f32 {
+    // A is stored as -exp(a_log) to ensure A_bar < 1 (stable)
+    let a_cont = -exp(a_log);
+    return exp(delta_val * a_cont);
+}
+
+fn discretise_B(delta_val: f32, a_log: f32, b_val: f32) -> f32 {
+    let a_cont  = -exp(a_log);
+    let a_bar   = exp(delta_val * a_cont);
+    // (A_bar - 1) / A_cont * B
+    let b_bar   = (a_bar - 1.0) / a_cont * b_val;
+    return b_bar;
+}
+
+// ---- Main kernel ----
+// Dispatch: (ceil(D/8), ceil(N/8), B)
+// Each invocation is responsible for one (d, n, batch) triplet and scans
+// the entire sequence using a two-pass Kogge-Stone scan within workgroup tiles.
+
+@compute @workgroup_size(64, 1, 1)
+fn forward_scan(
+    @builtin(global_invocation_id)   gid  : vec3<u32>,
+    @builtin(local_invocation_index) lid  : u32,
+    @builtin(workgroup_id)           wgid : vec3<u32>,
+) {
+    let L = params.seq_len;
+    let N = params.d_state;
+    let D = params.d_inner;
+    let B = params.batch;
+
+    // Each workgroup handles one (batch b, channel d, state n) combination.
+    // We pack d and n into the x dimension: global d = wgid.x, global n = wgid.y
+    let d = wgid.x;
+    let n = wgid.y;
+    let b = gid.z;
+
+    if (d >= D || n >= N || b >= B) { return; }
+
+    // Tile size equals workgroup size (64).  We process TILE_SIZE steps at once.
+    let TILE: u32 = 64u;
+
+    // Running state h for this (b, d, n)
+    var h: f32 = 0.0;
+
+    var tile_start: u32 = 0u;
+    loop {
+        if (tile_start >= L) { break; }
+
+        let t = tile_start + lid;      // absolute time step handled by this lane
+        var a_bar: f32 = 1.0;
+        var bu:    f32 = 0.0;
+
+        if (t < L) {
+            // Indices
+            let delta_idx = b * L * D + t * D + d;
+            let u_idx     = b * L * D + t * D + d;
+            let A_idx     = d * N + n;
+            let B_idx     = b * L * N + t * N + n;
+
+            let dv = softplus(delta[delta_idx]);
+            a_bar  = discretise_A(dv, A[A_idx]);
+            bu     = discretise_B(dv, A[A_idx], B[B_idx]) * u[u_idx];
+        }
+
+        wg_a[lid]  = a_bar;
+        wg_bu[lid] = bu;
+        workgroupBarrier();
+
+        // ---- Kogge-Stone inclusive prefix scan within tile ----
+        // Associative operator: (a1, b1) ∘ (a2, b2) = (a1*a2, a1*b2 + b1)
+        // This computes cumulative state recurrence in log2(TILE) steps.
+        var stride: u32 = 1u;
+        loop {
+            if (stride >= TILE) { break; }
+            if (lid >= stride) {
+                let prev_a  = wg_a[lid - stride];
+                let prev_bu = wg_bu[lid - stride];
+                // Combine: new_state = prev_a * cur_a (product of A_bars)
+                //                      new_bu  = prev_a * cur_bu + prev_bu
+                let new_a  = prev_a * wg_a[lid];
+                let new_bu = prev_a * wg_bu[lid] + prev_bu;
+                workgroupBarrier();
+                wg_a[lid]  = new_a;
+                wg_bu[lid] = new_bu;
+            }
+            workgroupBarrier();
+            stride = stride << 1u;
+        }
+
+        // Incorporate the carry-in state from the previous tile.
+        // After the scan wg_bu[lid] holds the intra-tile inclusive sum.
+        // The actual h at position t = h_carry * wg_a[lid] + wg_bu[lid]
+        let h_t = h * wg_a[lid] + wg_bu[lid];
+
+        if (t < L) {
+            // Cache hidden state for backward pass
+            let h_idx = b * L * D * N + t * D * N + d * N + n;
+            h_cache[h_idx] = h_t;
+
+            // Accumulate y contribution: y_t += C_t[n] * h_t  (over all n)
+            // We use an atomic-style accumulation: each (d, n) lane adds its
+            // contribution to the same y[b, t, d].  This races without atomics,
+            // so we instead write to a full h_cache and reduce in a second pass.
+            // Here we perform direct accumulation using atomicAdd approximation:
+            // (safe because each lane writes a unique n, which is stride 1 in mem)
+            let C_idx = b * L * N + t * N + n;
+            let y_idx = b * L * D + t * D + d;
+
+            // Direct write for n == 0 (first state dim), add for the rest.
+            // Since all workgroups for the same (b,d) run concurrently we must
+            // accumulate safely: we write each partial into h_cache and reduce
+            // in a subsequent lightweight kernel (forward_reduce).
+            // (For simplicity and correctness here we directly atomically add via
+            //  f32 emulation – real deployment uses atomicAdd on f32 with spirv ext.)
+            // We store C*h contribution separately so forward_reduce can sum them.
+            // Layout: y_partial (B, L, D, N) – one slot per state dim
+            // y reused as y_partial in this kernel; forward_reduce collapses N dim.
+            let y_partial_idx = b * L * D * N + t * D * N + d * N + n;
+            // Reuse h_cache second half as y_partial (offset by B*L*D*N)
+            let offset = B * L * D * N;
+            h_cache[offset + y_partial_idx] = C[C_idx] * h_t;
+        }
+
+        // Update carry: last lane's h_t is the tile's final state
+        let last = min(TILE, L - tile_start) - 1u;
+        h = wg_a[last] * h + wg_bu[last];   // recombine carry
+
+        workgroupBarrier();
+        tile_start = tile_start + TILE;
+    }
+}
+
+// ---- Reduction kernel ----
+// Collapses the N (d_state) dimension of y_partial into y.
+// Adds the D (skip connection) term: y_t[d] += D_vec[d] * u_t[d]
+// Dispatch: (ceil(L/64), D, B)
+
+@compute @workgroup_size(64, 1, 1)
+fn forward_reduce(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L = params.seq_len;
+    let N = params.d_state;
+    let D = params.d_inner;
+    let B = params.batch;
+
+    let t = gid.x;
+    let d = gid.y;
+    let b = gid.z;
+
+    if (t >= L || d >= D || b >= B) { return; }
+
+    let offset    = B * L * D * N;
+    var sum: f32  = 0.0;
+    for (var n: u32 = 0u; n < N; n = n + 1u) {
+        let idx = offset + b * L * D * N + t * D * N + d * N + n;
+        sum = sum + h_cache[idx];
+    }
+
+    // Add skip connection
+    let u_idx = b * L * D + t * D + d;
+    sum = sum + D_vec[d] * u[u_idx];
+
+    let y_idx = b * L * D + t * D + d;
+    y[y_idx] = sum;
+}
+`;
+// ---- Backward scan kernel (for autograd) ----
+// Computes gradients w.r.t. Δ, A, B, C using the cached hidden states.
+export const SELECTIVE_SCAN_BACKWARD_WGSL = /* wgsl */ `
+
+struct ScanParams {
+    seq_len  : u32,
+    d_state  : u32,
+    d_inner  : u32,
+    batch    : u32,
+};
+
+@group(0) @binding(0) var<uniform>             params    : ScanParams;
+@group(0) @binding(1) var<storage, read>       u         : array<f32>;
+@group(0) @binding(2) var<storage, read>       delta     : array<f32>;
+@group(0) @binding(3) var<storage, read>       A         : array<f32>;
+@group(0) @binding(4) var<storage, read>       B         : array<f32>;
+@group(0) @binding(5) var<storage, read>       C         : array<f32>;
+@group(0) @binding(6) var<storage, read>       h_cache   : array<f32>;
+@group(0) @binding(7) var<storage, read>       dy        : array<f32>;  // upstream gradient
+@group(0) @binding(8) var<storage, read_write> dA        : array<f32>;
+@group(0) @binding(9) var<storage, read_write> dB        : array<f32>;
+@group(0) @binding(10) var<storage, read_write> dC       : array<f32>;
+@group(0) @binding(11) var<storage, read_write> dDelta   : array<f32>;
+@group(0) @binding(12) var<storage, read_write> du       : array<f32>;
+
+fn softplus(x: f32) -> f32 {
+    return log(1.0 + exp(x));
+}
+
+fn softplus_grad(x: f32) -> f32 {
+    // d/dx softplus(x) = sigmoid(x)
+    return 1.0 / (1.0 + exp(-x));
+}
+
+fn discretise_A(delta_val: f32, a_log: f32) -> f32 {
+    let a_cont = -exp(a_log);
+    return exp(delta_val * a_cont);
+}
+
+// Reverse scan (backward pass) – processes time from T-1 down to 0.
+// Dispatch: (D, N, B)
+@compute @workgroup_size(1, 1, 1)
+fn backward_scan(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let L = params.seq_len;
+    let N = params.d_state;
+    let D = params.d_inner;
+    let B = params.batch;
+
+    let d = gid.x;
+    let n = gid.y;
+    let b = gid.z;
+
+    if (d >= D || n >= N || b >= B) { return; }
+
+    var dh: f32 = 0.0;   // gradient of loss w.r.t. h_t, accumulated backwards
+
+    var t: u32 = L;
+    loop {
+        if (t == 0u) { break; }
+        t = t - 1u;
+
+        let delta_raw_idx = b * L * D + t * D + d;
+        let A_idx         = d * N + n;
+        let B_idx         = b * L * N + t * N + n;
+        let C_idx         = b * L * N + t * N + n;
+        let u_idx         = b * L * D + t * D + d;
+        let h_idx         = b * L * D * N + t * D * N + d * N + n;
+
+        let delta_raw = delta[delta_raw_idx];
+        let dv        = softplus(delta_raw);
+        let a_log     = A[A_idx];
+        let a_cont    = -exp(a_log);
+        let a_bar     = exp(dv * a_cont);
+        let b_val     = B[B_idx];
+        let c_val     = C[C_idx];
+        let u_val     = u[u_idx];
+        let h_t       = h_cache[h_idx];
+
+        // dy_t contribution to dh (from C * h_t in the output)
+        // y_t[d] = sum_n C[n] * h_t[n] + D * u   =>  dh_t[n] += C[n] * dy_t[d]
+        let dy_val = dy[b * L * D + t * D + d];
+        dh = dh + c_val * dy_val;
+
+        // dC[b, t, n] += dy_t[d] * h_t
+        dC[C_idx] = dC[C_idx] + dy_val * h_t;
+
+        // h_t = a_bar * h_{t-1} + b_bar * u_t
+        // b_bar = (a_bar - 1) / a_cont * b_val
+        let b_bar  = (a_bar - 1.0) / a_cont * b_val;
+        let h_prev = (t > 0u) ? h_cache[b * L * D * N + (t - 1u) * D * N + d * N + n] : 0.0;
+
+        // dh_{t-1} += a_bar * dh_t
+        // (accumulated in next iteration; here dh already contains upstream)
+        let dh_cur = dh;
+
+        // dA[d,n] += dh_t * (d a_bar/d a_cont) * (d a_cont/d a_log) * h_{t-1}
+        //          + dh_t * (d b_bar/d a_cont) * ... * b_val * u_val
+        // d(a_bar)/d(a_log) = a_bar * (-exp(a_log)) * dv = a_bar * a_cont * dv
+        let da_bar_da_log = a_bar * a_cont * dv;
+        dA[A_idx] = dA[A_idx] + dh_cur * (da_bar_da_log * h_prev);
+
+        // dB[b,t,n] += dh_t * b_bar / b_val * u_val  (since b_bar is linear in b)
+        dB[B_idx] = dB[B_idx] + dh_cur * ((a_bar - 1.0) / a_cont) * u_val;
+
+        // du[b,t,d] += dh_t * b_bar  (accumulate over n in separate kernel)
+        du[u_idx] = du[u_idx] + dh_cur * b_bar;
+
+        // dDelta[b,t,d]: chain rule through softplus and discretisation
+        // d(b_bar)/d(dv) = d/d(dv)[(a_bar-1)/a_cont * b] = a_bar * b / (a_cont ... )
+        //  actually: d(a_bar)/d(dv) = a_bar * a_cont,  d(b_bar)/d(dv) = a_bar * b_val
+        let da_bar_ddv  = a_bar * a_cont;
+        let db_bar_ddv  = a_bar * b_val;
+        let dLoss_ddv   = dh_cur * (da_bar_ddv * h_prev + db_bar_ddv * u_val);
+        let ddv_ddelta  = softplus_grad(delta_raw);
+        dDelta[delta_raw_idx] = dDelta[delta_raw_idx] + dLoss_ddv * ddv_ddelta;
+
+        // Propagate dh to previous timestep
+        dh = a_bar * dh_cur;
+    }
+}
+`;
+//# sourceMappingURL=selective_scan.js.map

package/dist/kernels/selective_scan.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"selective_scan.js","sourceRoot":"","sources":["../../src/kernels/selective_scan.ts"],"names":[],"mappings":"AAAA,sCAAsC;AACtC,qEAAqE;AACrE,4EAA4E;AAC5E,EAAE;AACF,2BAA2B;AAC3B,oCAAoC;AACpC,8BAA8B;AAC9B,EAAE;AACF,qEAAqE;AAErE,MAAM,CAAC,MAAM,2BAA2B,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAsN5D,CAAC;AAEF,gDAAgD;AAChD,uEAAuE;AAEvE,MAAM,CAAC,MAAM,4BAA4B,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAwH7D,CAAC"}

package/dist/kernels/weight_update.d.ts

@@ -0,0 +1,3 @@
+export declare const WEIGHT_UPDATE_WGSL: string;
+export declare const GRAD_CLIP_WGSL: string;
+//# sourceMappingURL=weight_update.d.ts.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"weight_update.d.ts","sourceRoot":"","sources":["../../src/kernels/weight_update.ts"],"names":[],"mappings":"AAUA,eAAO,MAAM,kBAAkB,EAAE,MAgDhC,CAAC;AAIF,eAAO,MAAM,cAAc,EAAE,MAyD5B,CAAC"}

package/dist/kernels/weight_update.js

@@ -0,0 +1,119 @@
+// Weight Update WGSL Kernel (AdamW Optimizer)
+// Implements fused AdamW parameter update on the GPU.
+//
+// AdamW update rule:
+//   m_t = beta1 * m_{t-1} + (1 - beta1) * g_t
+//   v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2
+//   m_hat = m_t / (1 - beta1^t)
+//   v_hat = v_t / (1 - beta2^t)
+//   theta_t = theta_{t-1} * (1 - lr * weight_decay) - lr * m_hat / (sqrt(v_hat) + eps)
+export const WEIGHT_UPDATE_WGSL = /* wgsl */ `
+
+struct AdamParams {
+    num_elements   : u32,
+    lr             : f32,   // learning rate
+    beta1          : f32,   // default 0.9
+    beta2          : f32,   // default 0.999
+    eps            : f32,   // default 1e-8
+    weight_decay   : f32,   // default 0.01
+    beta1_t        : f32,   // beta1^t  (precomputed bias correction term)
+    beta2_t        : f32,   // beta2^t
+};
+
+@group(0) @binding(0) var<uniform>             adam     : AdamParams;
+// param (N,)   – weight tensor (read-write: updated in-place)
+@group(0) @binding(1) var<storage, read_write> param    : array<f32>;
+// grad  (N,)   – gradient
+@group(0) @binding(2) var<storage, read>       grad     : array<f32>;
+// m     (N,)   – first moment
+@group(0) @binding(3) var<storage, read_write> m_state  : array<f32>;
+// v     (N,)   – second moment
+@group(0) @binding(4) var<storage, read_write> v_state  : array<f32>;
+
+// Dispatch: (ceil(N / 256), 1, 1)
+@compute @workgroup_size(256, 1, 1)
+fn adamw_update(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let i = gid.x;
+    if (i >= adam.num_elements) { return; }
+
+    let g = grad[i];
+    let p = param[i];
+
+    // Moment updates
+    let m_new = adam.beta1 * m_state[i] + (1.0 - adam.beta1) * g;
+    let v_new = adam.beta2 * v_state[i] + (1.0 - adam.beta2) * g * g;
+    m_state[i] = m_new;
+    v_state[i] = v_new;
+
+    // Bias-corrected estimates
+    let m_hat = m_new / (1.0 - adam.beta1_t);
+    let v_hat = v_new / (1.0 - adam.beta2_t);
+
+    // Weight decay (decoupled) + gradient step
+    param[i] = p * (1.0 - adam.lr * adam.weight_decay) -
+               adam.lr * m_hat / (sqrt(v_hat) + adam.eps);
+}
+`;
+// Gradient clipping kernel – clips global gradient norm to max_norm.
+// Run before weight updates.  Two-pass: first compute squared norm, then scale.
+export const GRAD_CLIP_WGSL = /* wgsl */ `
+
+struct ClipParams {
+    num_elements : u32,
+    max_norm_sq  : f32,   // max_norm^2
+};
+
+@group(0) @binding(0) var<uniform>             clip_p  : ClipParams;
+@group(0) @binding(1) var<storage, read_write> grad    : array<f32>;
+@group(0) @binding(2) var<storage, read_write> norm_sq : array<f32>;  // size 1, atomic accumulator
+
+var<workgroup> local_sq : array<f32, 256>;
+
+// Pass 1: reduce sum of squares into norm_sq[0]
+@compute @workgroup_size(256, 1, 1)
+fn grad_norm_reduce(
+    @builtin(global_invocation_id)   gid : vec3<u32>,
+    @builtin(local_invocation_index) lid : u32,
+) {
+    let i = gid.x;
+    local_sq[lid] = 0.0;
+    if (i < clip_p.num_elements) {
+        local_sq[lid] = grad[i] * grad[i];
+    }
+    workgroupBarrier();
+
+    // Parallel reduction within workgroup
+    var s: u32 = 128u;
+    loop {
+        if (s == 0u) { break; }
+        if (lid < s) {
+            local_sq[lid] = local_sq[lid] + local_sq[lid + s];
+        }
+        workgroupBarrier();
+        s = s >> 1u;
+    }
+
+    if (lid == 0u) {
+        // Non-atomic accumulation (single workgroup assumption for small models)
+        norm_sq[0] = norm_sq[0] + local_sq[0];
+    }
+}
+
+// Pass 2: scale gradients if norm exceeds max_norm
+@compute @workgroup_size(256, 1, 1)
+fn grad_clip_scale(
+    @builtin(global_invocation_id) gid : vec3<u32>,
+) {
+    let i = gid.x;
+    if (i >= clip_p.num_elements) { return; }
+
+    let ns = norm_sq[0];
+    if (ns > clip_p.max_norm_sq) {
+        let scale = sqrt(clip_p.max_norm_sq / ns);
+        grad[i] = grad[i] * scale;
+    }
+}
+`;
+//# sourceMappingURL=weight_update.js.map

package/dist/kernels/weight_update.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"weight_update.js","sourceRoot":"","sources":["../../src/kernels/weight_update.ts"],"names":[],"mappings":"AAAA,8CAA8C;AAC9C,sDAAsD;AACtD,EAAE;AACF,qBAAqB;AACrB,8CAA8C;AAC9C,gDAAgD;AAChD,gCAAgC;AAChC,gCAAgC;AAChC,uFAAuF;AAEvF,MAAM,CAAC,MAAM,kBAAkB,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAgDnD,CAAC;AAEF,qEAAqE;AACrE,gFAAgF;AAChF,MAAM,CAAC,MAAM,cAAc,GAAW,UAAU,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAyD/C,CAAC"}

package/dist/model/mamba_block.d.ts

@@ -0,0 +1,64 @@
+/**
+ * mamba_block.ts – Mamba Mixer Block
+ */
+export interface MambaBlockConfig {
+    dModel: number;
+    dState?: number;
+    dConv?: number;
+    expand?: number;
+    dtRank?: number;
+    biasConv?: boolean;
+}
+export interface BlockParam {
+    buf: GPUBuffer;
+    numel: number;
+    name: string;
+}
+export interface BlockCache {
+    normInv: GPUBuffer;
+    normIn: GPUBuffer;
+    normOut: GPUBuffer;
+    zBuf: GPUBuffer;
+    xConvIn: GPUBuffer;
+    convOut: GPUBuffer;
+    siluOut: GPUBuffer;
+    deltaFull: GPUBuffer;
+    B_raw: GPUBuffer;
+    C_raw: GPUBuffer;
+    hCache: GPUBuffer;
+}
+export interface BlockForwardResult {
+    output: GPUBuffer;
+    cache: BlockCache;
+}
+export declare class MambaBlock {
+    device: GPUDevice;
+    config: Required<MambaBlockConfig>;
+    dInner: number;
+    dtRank: number;
+    wInProj: Float32Array;
+    bInProj: Float32Array;
+    wConv: Float32Array;
+    bConv: Float32Array;
+    wXProj: Float32Array;
+    bXProj: Float32Array;
+    wDtProj: Float32Array;
+    bDtProj: Float32Array;
+    A_log: Float32Array;
+    D_vec: Float32Array;
+    wOutProj: Float32Array;
+    bOutProj: Float32Array;
+    normWeight: Float32Array;
+    gpuWeights: Record<string, GPUBuffer>;
+    pipelines: Record<string, GPUComputePipeline>;
+    private _wslaMode;
+    constructor(device: GPUDevice, config: MambaBlockConfig);
+    private _initWeights;
+    private _uploadWeightsToGPU;
+    private _buildPipelines;
+    forward(xBuf: GPUBuffer, batch: number, seqLen: number): BlockForwardResult;
+    parameters(): BlockParam[];
+    setWSLAMode(enabled: boolean): void;
+    getTrainableParams(): BlockParam[];
+}
+//# sourceMappingURL=mamba_block.d.ts.map

package/dist/model/mamba_block.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"mamba_block.d.ts","sourceRoot":"","sources":["../../src/model/mamba_block.ts"],"names":[],"mappings":"AAAA;;GAEG;AAiBH,MAAM,WAAW,gBAAgB;IAC/B,MAAM,EAAE,MAAM,CAAC;IACf,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,QAAQ,CAAC,EAAE,OAAO,CAAC;CACpB;AAED,MAAM,WAAW,UAAU;IACzB,GAAG,EAAE,SAAS,CAAC;IACf,KAAK,EAAE,MAAM,CAAC;IACd,IAAI,EAAE,MAAM,CAAC;CACd;AAED,MAAM,WAAW,UAAU;IACzB,OAAO,EAAE,SAAS,CAAC;IACnB,MAAM,EAAE,SAAS,CAAC;IAClB,OAAO,EAAE,SAAS,CAAC;IACnB,IAAI,EAAE,SAAS,CAAC;IAChB,OAAO,EAAE,SAAS,CAAC;IACnB,OAAO,EAAE,SAAS,CAAC;IACnB,OAAO,EAAE,SAAS,CAAC;IACnB,SAAS,EAAE,SAAS,CAAC;IACrB,KAAK,EAAE,SAAS,CAAC;IACjB,KAAK,EAAE,SAAS,CAAC;IACjB,MAAM,EAAE,SAAS,CAAC;CACnB;AAED,MAAM,WAAW,kBAAkB;IACjC,MAAM,EAAE,SAAS,CAAC;IAClB,KAAK,EAAE,UAAU,CAAC;CACnB;AAED,qBAAa,UAAU;IACnB,MAAM,EAAE,SAAS,CAAC;IAClB,MAAM,EAAE,QAAQ,CAAC,gBAAgB,CAAC,CAAC;IACnC,MAAM,EAAE,MAAM,CAAC;IACf,MAAM,EAAE,MAAM,CAAC;IACf,OAAO,EAAE,YAAY,CAAC;IACtB,OAAO,EAAE,YAAY,CAAC;IACtB,KAAK,EAAE,YAAY,CAAC;IACpB,KAAK,EAAE,YAAY,CAAC;IACpB,MAAM,EAAE,YAAY,CAAC;IACrB,MAAM,EAAE,YAAY,CAAC;IACrB,OAAO,EAAE,YAAY,CAAC;IACtB,OAAO,EAAE,YAAY,CAAC;IACtB,KAAK,EAAE,YAAY,CAAC;IACpB,KAAK,EAAE,YAAY,CAAC;IACpB,QAAQ,EAAE,YAAY,CAAC;IACvB,QAAQ,EAAE,YAAY,CAAC;IACvB,UAAU,EAAE,YAAY,CAAC;IACzB,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,SAAS,CAAC,CAAC;IACtC,SAAS,EAAE,MAAM,CAAC,MAAM,EAAE,kBAAkB,CAAC,CAAC;IAC9C,OAAO,CAAC,SAAS,CAAS;gBAEd,MAAM,EAAE,SAAS,EAAE,MAAM,EAAE,gBAAgB;IAoCvD,OAAO,CAAC,YAAY;IA2CpB,OAAO,CAAC,mBAAmB;IAqB3B,OAAO,CAAC,eAAe;IAavB,OAAO,CAAC,IAAI,EAAE,SAAS,EAAE,KAAK,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,kBAAkB;IAqL3E,UAAU,IAAI,UAAU,EAAE;IAwB1B,WAAW,CAAC,OAAO,EAAE,OAAO,GAAG,IAAI;IAInC,kBAAkB,IAAI,UAAU,EAAE;CASrC"}