Files
foxhunt/crates/ml-core
jgrusewski eb0e4b6328 perf(ml-alpha): full zero-alloc training step (#3 foundation)
Eliminates ALL per-step allocations from the training hot path —
foundation for CUDA Graph capture (next commit). Before this commit,
each step_batched call allocated:

  Mamba2 forward:    input_2d view, x, a_proj, b_proj, h_s2, h_enriched_seq
  Mamba2 backward:   d_a_per_channel/d_b_per_channel/d_w_c/d_h_s2 (#2 covered)
                     d_a_proj_flat, d_b_proj_flat, dw_c
                     LinearGrads.{dw,db,dx} × 3 projections (cuBLAS internal)
                     d_x_from_a + d_x_from_b + d_x (elementwise add)
                     dw_out, db_out (zero-init shells)
  Trainer wrapper:   window_tensor, h_enriched_seq_t, grad_h_enriched_seq_t,
                     grad_h_enriched_seq

~20-25 cudaMalloc / GpuTensor::zeros calls per step × 2000 steps/epoch =
40-50K allocations per epoch.

This commit adds zero-alloc `_into` variants throughout the chain:

  ml-core/cuda_autograd/linear.rs:
    OwnedGpuLinear::forward_with_slices_into
    OwnedGpuLinear::backward_with_slices_into
    reduce_sum_axis0_into

  ml-core/cuda_autograd/elementwise.rs + gpu_tensor.rs:
    ElementwiseKernels::binary_into
    GpuTensor::add_into

  ml-alpha/mamba2_block.rs:
    Mamba2BlockForwardScratch (pre-allocated forward cache)
    Mamba2BackwardGradsBuffers (pre-allocated backward outputs)
    Mamba2Block::forward_train_seq_into (zero-alloc forward)
    Mamba2Block::backward_from_h_enriched_seq_full_into (zero-alloc backward)
    Mamba2AdamW::step_from_buffers (reads grads_buffers directly)

  ml-alpha/trainer/perception.rs:
    PerceptionTrainer pre-allocates: window_tensor_d, h_enriched_seq_t_d,
      grad_h_enriched_seq_t_d, grad_h_enriched_seq_d, mamba2_fwd_scratch,
      mamba2_grads_buffers
    step_batched + evaluate_batched fully wired through _into variants

Original `forward_with_slices` / `backward_with_slices` / `binary` / `add` /
`backward_from_h_enriched_seq` paths preserved unchanged — Phase E.3
ml/examples callers unaffected.

The captured-graph commit (next) only needs to wrap this zero-alloc
training step in cuGraph capture/replay; no further refactoring of
buffer management.

77 ml-alpha tests pass. Synthetic overfit converges identically
(0.29 → 0.0007 in 250 steps) — gradients are bit-identical to the
allocating path.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-17 13:34:56 +02:00
..