MambaBlock::forward already had a working, tested CUDA dispatch
(forward_cuda: cuBLAS matmuls for projections, CPU for the scan).
backward() had none — it silently ran entirely CPU-serial on GPU
tensors via to_vec()/from_vec() D2H/H2D round-trips. This adds the
missing acceleration, mirroring forward_cuda's hybrid split: the
four large projection-parameter gradients (in_proj, x_proj, dt_proj,
out_proj) now go through batched GPU matmuls; the inherently
sequential scan reverse-pass and small per-channel grads stay CPU.
Extracted CpuWeights::pull and recompute_forward_cpu out of the old
inline per-batch forward-recompute block inside backward() (pure
refactor, gradient-checked unchanged by real_selective_scan.rs's
existing 6 tests) so CPU backward and the new CUDA backward share
identical forward math and can never numerically diverge on it.
New CUDA-vs-CPU gradient-check test (mamba_cuda_backward_matches_cpu,
#[ignore]-gated, GPU-only) caught a real bug during development:
Tensor::contiguous() is a no-op stub in this rtx-tensor version, and
cuda_matmul reads raw GPU storage by shape.dims() ignoring
strides/offset, so .transpose(..).matmul(..) on a GPU tensor silently
computed garbage (80-200x relative error on 3 of 4 accelerated
gradients). Fixed by building already-transposed [dim, b*l] buffers
on CPU before upload instead of transposing GPU-side. All 9 gradients
now match CPU backward within ~2.2e-5 max relative error (tolerance
1e-4).
Co-Authored-By: Claude Sonnet 5 <[email protected]>
Adds d350_gpu_backbone_training — 200-step Adam loop on a 4-regime
corpus that prefers Device::cuda(0) and falls back gracefully to CPU.
Measured numbers:
- CPU (DIM=16): 886 steps/s, MSE 0.2163 → 0.0024
- GPU (DIM=16): 803 steps/s, MSE 0.2163 → 0.0024
GPU is marginally slower at DIM=16 because the SSM scan and conv1d
remain on CPU in both paths; cuBLAS only helps the four linear
projections, which are tiny at dim=16. The GPU advantage emerges at
larger dims (≥256) where the projections dominate. Correctness is
identical on both devices.
Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
MambaRecurrence reconstructs MambaBlock's selective-scan as an explicit
step(state, x_t) -> (state', y_t) recurrence over host f32 weights, with the
recurrent state (h, conv-ring) carried between steps. This is the substrate the
SMT memory updater is behaviorally cloned on (E3b) — no BPTT, one step at a time.
Dimensions are derived from the persisted tensor shapes, so mamba.rs is left
untouched (it is near the 1250-LOC cap). silu/softplus are byte-identical copies
of the forward's.
Equivalence pin (the plan's highest-risk item): stepping a window one token at a
time from a zero state reproduces the full-sequence forward EXACTLY —
max_abs_diff = 0.0 (bit-identical), on an active_block with wide Delta so the
scan genuinely drives the output. Plus a fresh-state determinism test.
clippy(-D warnings) + fmt clean on the new module and test.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
End-to-end demonstration that training the real selective-scan Mamba buys
a genuine temporal-modeling win — the payoff of M1–M3.
Task: next-token prediction on a multi-regime sequence (x[t] = μ_regime +
noise). Predicting x[t+1] inside a regime requires integrating recent
history to average out the noise — a memoryless model can't.
Result (same sequence, all three):
persistence (memoryless) MSE = 0.0895
untrained Mamba MSE = 0.2167
trained Mamba (400 Adam) MSE = 0.0002
The trained backbone integrates history to de-noise the regime mean,
beating the memoryless persistence baseline by ~450× and improving
~1000× over its untrained self. (Single-sequence fit: demonstrates the
SSM's temporal-modeling capacity, not held-out generalization.)
This replaces the old non-result ("random Mamba 6.5% vs linear 48%") with
a real "trained SSM exploits temporal structure" demonstration.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Plain randn(0,1) init made A=-exp(A_log) and Δ wildly large, overflowing
the real exp(Δ·A) scan to NaN (the stub never cared — it discarded these
weights). new()/new_seeded() now share a build() with canonical S6 init:
- A_log = ln(1..=d_state) ⇒ A = -(1..=d_state), bounded
- dt_bias so softplus(dt_bias) ≈ 0.01 (small, stable Δ; near-identity
scan at init — intentional for gradient flow)
- D = 1, zero conv bias, projections scaled by 1/√fan_in (capped 0.5)
This fixes the NaN that broke omni-cortex's d231 action-conditioned
predictor training (now green). Seeded determinism preserved.
Tests: active_block helper (Δ overridden to ≈0.69) exercises the
scan-active regime so the liveness check can observe each weight; the
training test asserts a seed-varying backbone weight (conv1d_weight)
moves. All 6 selective-scan tests green incl. the finite-diff grad check.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Demonstrates end-to-end trainability of the real selective scan. A
self-contained Adam loop (the rtx-transformers AdamOptimizer has no
public gradient setter — the spec permits a bespoke loop) fits a teacher
block's output on a fixed input: forward → MSE → analytic backward →
Adam step → rebuild. Over 200 steps the loss drops >50% and the backbone
weight A_log moves, confirming gradients actually train the model (not
just the head). All 6 selective-scan tests green.
The production AdamOptimizer can be wired once it exposes a gradient
setter; the M2 backward already returns grads in its HashMap shape.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Hand-written VJP of the selective-scan forward (Phase-3 spec §5/§7):
MambaBlock::backward(x, d_out) -> per-parameter gradients keyed by
persistence name, summed over the batch. Differentiates the scan
analytically (reverse-time recurrence over the cached h trajectory)
rather than through the immature rtx-tensor autograd tape.
Covers every parameter: in_proj, conv1d_weight, conv1d_bias, A_log
(via A=-exp(A_log) ⇒ dA_log = dA·A), x_proj, dt_proj, dt_bias, D,
out_proj. Adds stable sigmoid_f32 / silu_grad_f32 helpers.
New test analytic_gradients_match_finite_differences: on a small
well-conditioned instance, ≥30 sampled grad elements across all 9
params match central finite differences within (5e-3 + 5e-2·|fd|).
All 5 selective-scan tests green.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
The MambaBlock forward was a stub — SelectiveScan::forward passed input
through, discretize returned zeros, conv1d was a no-op, and B/C were
randn per call, leaving conv1d_weight/dt_proj/A_log as dead weights with
zero temporal mixing. This implements a genuine S6 selective-scan:
- New params: x_proj [d_inner, dt_rank+2*d_state] (data-dependent dt,B,C),
dt_bias [d_inner], D [d_inner] (skip). Added to new()/new_seeded() and
the persistence contract (persistence_tensors/from_persistence_tensors).
- Real forward (CPU f32, looped — backbone is small): in_proj -> causal
depthwise conv1d -> SiLU -> x_proj->(dt,B,C) -> delta=softplus(dt.dt_proj
+dt_bias) -> A=-exp(A_log) -> sequential scan h=dA.h+dBu, y=sum C.h + D.u
-> gate by SiLU(z) -> out_proj. Residual moved OUT (canonical).
Numerically-stable silu_f32/softplus_f32 helpers.
- The scan runs inline (not via the immature rtx-tensor autograd tape);
the analytic backward lands in M2 per docs/phase3_real_ssm_spec.md.
New tests/real_selective_scan.rs (4 cases, all green): liveness (each
formerly-dead weight now moves the output), causality (no future
leakage), seeded determinism, and finite/non-constant output.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
