The four biggest per-call scratch buffers in
Mamba2Block::backward_from_h_enriched_seq were allocated fresh on
every training step:
d_a_per_channel [N, sh2, K, state_d] ~6 MB at B=8, K=96
d_b_per_channel [N, sh2, K, state_d] ~6 MB
d_w_c_per_sample [N, sh2, state_d] ~64 KB
d_h_s2 [N, sh2] ~4 KB
For 2000 optimizer steps/epoch × 15 epochs = 30 000 alloc_zeros
calls per training run, all on the hot path.
New `Mamba2BackwardScratch` struct holds these as device-resident
buffers, constructed once per (n_batch, seq_len, hidden_dim,
state_dim) at trainer init. New `backward_from_h_enriched_seq_into`
method takes the scratch by &mut and reuses the buffers each call.
The smaller per-call buffers (d_a_proj_flat, d_b_proj_flat, dw_c)
still allocate per call — they feed into ownership-transferring
LinearGrads outputs, where pre-allocation would require refactoring
ml-core's cuBLAS wrappers without proportional gain.
The original `backward_from_h_enriched_seq` is preserved (Phase E.3
callers still use it). Trainer switches to the `_into` variant.
Expected per-step savings: ~10-20 ms on L40S (4 cudaMalloc latencies
per call × 2.5-5 μs each + cache pressure reduction). Over 2000
steps/epoch that's 20-40 sec/epoch.
77 ml-alpha tests pass. Synthetic overfit unchanged (0.27 → 0.0006).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
ebae67cb6b