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>
This commit is contained in:
jgrusewski
@@ -145,6 +145,134 @@ pub struct Mamba2ForwardCacheSeq {
|
||||
pub h_enriched_seq: GpuTensor,
|
||||
}
|
||||
|
||||
/// Pre-allocated forward intermediates for [`Mamba2Block::forward_train_seq_into`].
|
||||
/// Holds all GpuTensors that the regular `forward_train_seq` allocates per
|
||||
/// call (input_2d projection output `x`, A/B projections, h_s2 residual,
|
||||
/// h_enriched_seq scan output). Constructed once per (n_batch, seq_len,
|
||||
/// in_dim, hidden_dim, state_dim) tuple and reused — eliminates the
|
||||
/// per-step `cudaMalloc` calls that block CUDA Graph capture.
|
||||
///
|
||||
/// `h_s2` is zero-initialised once and never written (no residual carry
|
||||
/// from a prior chunk in the supervised path). The scan kernel reads it
|
||||
/// as a constant addition to h_enriched_seq.
|
||||
pub struct Mamba2BlockForwardScratch {
|
||||
pub x: GpuTensor, // [n_rows, hidden_dim] (W_in output)
|
||||
pub a_proj: GpuTensor, // [n_rows, state_dim]
|
||||
pub b_proj: GpuTensor, // [n_rows, state_dim]
|
||||
pub h_s2: GpuTensor, // [n_batch, hidden_dim] (zero residual)
|
||||
pub h_enriched_seq: GpuTensor, // [n_batch, seq_len, hidden_dim]
|
||||
pub n_batch: usize,
|
||||
pub seq_len: usize,
|
||||
pub in_dim: usize,
|
||||
pub hidden_dim: usize,
|
||||
pub state_dim: usize,
|
||||
}
|
||||
|
||||
impl Mamba2BlockForwardScratch {
|
||||
pub fn new(
|
||||
stream: &Arc<CudaStream>,
|
||||
n_batch: usize,
|
||||
seq_len: usize,
|
||||
in_dim: usize,
|
||||
hidden_dim: usize,
|
||||
state_dim: usize,
|
||||
) -> Result<Self> {
|
||||
let n_rows = n_batch * seq_len;
|
||||
Ok(Self {
|
||||
x: GpuTensor::zeros(&[n_rows, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("fwd scratch x: {e}"))?,
|
||||
a_proj: GpuTensor::zeros(&[n_rows, state_dim], stream)
|
||||
.map_err(|e| anyhow!("fwd scratch a_proj: {e}"))?,
|
||||
b_proj: GpuTensor::zeros(&[n_rows, state_dim], stream)
|
||||
.map_err(|e| anyhow!("fwd scratch b_proj: {e}"))?,
|
||||
h_s2: GpuTensor::zeros(&[n_batch, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("fwd scratch h_s2: {e}"))?,
|
||||
h_enriched_seq: GpuTensor::zeros(&[n_batch, seq_len, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("fwd scratch h_enriched_seq: {e}"))?,
|
||||
n_batch, seq_len, in_dim, hidden_dim, state_dim,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Pre-allocated outputs + intermediates for the full Mamba2 seq
|
||||
/// backward — see [`Mamba2Block::backward_from_h_enriched_seq_full_into`].
|
||||
/// Holds the cuBLAS linear-backward outputs (dw_in/db_in/dw_a/db_a/
|
||||
/// dw_b/db_b plus the reshaped d_a_proj_2d / d_b_proj_2d / d_x_from_a
|
||||
/// / d_x_from_b / d_x intermediates) so the backward pass does zero
|
||||
/// per-step allocations.
|
||||
///
|
||||
/// W_out is unused in the seq path — dw_out / db_out are zero-init
|
||||
/// shells of the right shape so the Mamba2AdamW step is a no-op for
|
||||
/// those parameters.
|
||||
pub struct Mamba2BackwardGradsBuffers {
|
||||
pub d_a_proj_2d: GpuTensor, // [n_rows, state_dim] (cuBLAS-reshape of d_a_proj_flat)
|
||||
pub d_b_proj_2d: GpuTensor, // [n_rows, state_dim]
|
||||
pub d_x_from_a: GpuTensor, // [n_rows, hidden_dim]
|
||||
pub d_x_from_b: GpuTensor, // [n_rows, hidden_dim]
|
||||
pub d_x: GpuTensor, // [n_rows, hidden_dim]
|
||||
pub dw_in: GpuTensor, // [hidden_dim, in_dim]
|
||||
pub db_in: GpuTensor, // [hidden_dim]
|
||||
pub dw_a: GpuTensor, // [state_dim, hidden_dim]
|
||||
pub db_a: GpuTensor, // [state_dim]
|
||||
pub dw_b: GpuTensor, // [state_dim, hidden_dim]
|
||||
pub db_b: GpuTensor, // [state_dim]
|
||||
pub dw_c: GpuTensor, // [hidden_dim, state_dim]
|
||||
pub d_x_from_in: GpuTensor, // [n_rows, in_dim] (unused but allocated for w_in_into's dx_out)
|
||||
pub dw_out: GpuTensor, // [1, hidden_dim] (zero)
|
||||
pub db_out: GpuTensor, // [1] (zero)
|
||||
pub n_batch: usize,
|
||||
pub seq_len: usize,
|
||||
pub in_dim: usize,
|
||||
pub hidden_dim: usize,
|
||||
pub state_dim: usize,
|
||||
}
|
||||
|
||||
impl Mamba2BackwardGradsBuffers {
|
||||
pub fn new(
|
||||
stream: &Arc<CudaStream>,
|
||||
n_batch: usize,
|
||||
seq_len: usize,
|
||||
in_dim: usize,
|
||||
hidden_dim: usize,
|
||||
state_dim: usize,
|
||||
) -> Result<Self> {
|
||||
let n_rows = n_batch * seq_len;
|
||||
Ok(Self {
|
||||
d_a_proj_2d: GpuTensor::zeros(&[n_rows, state_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads d_a_proj_2d: {e}"))?,
|
||||
d_b_proj_2d: GpuTensor::zeros(&[n_rows, state_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads d_b_proj_2d: {e}"))?,
|
||||
d_x_from_a: GpuTensor::zeros(&[n_rows, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads d_x_from_a: {e}"))?,
|
||||
d_x_from_b: GpuTensor::zeros(&[n_rows, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads d_x_from_b: {e}"))?,
|
||||
d_x: GpuTensor::zeros(&[n_rows, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads d_x: {e}"))?,
|
||||
dw_in: GpuTensor::zeros(&[hidden_dim, in_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads dw_in: {e}"))?,
|
||||
db_in: GpuTensor::zeros(&[hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads db_in: {e}"))?,
|
||||
dw_a: GpuTensor::zeros(&[state_dim, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads dw_a: {e}"))?,
|
||||
db_a: GpuTensor::zeros(&[state_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads db_a: {e}"))?,
|
||||
dw_b: GpuTensor::zeros(&[state_dim, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads dw_b: {e}"))?,
|
||||
db_b: GpuTensor::zeros(&[state_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads db_b: {e}"))?,
|
||||
dw_c: GpuTensor::zeros(&[hidden_dim, state_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads dw_c: {e}"))?,
|
||||
d_x_from_in: GpuTensor::zeros(&[n_rows, in_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads d_x_from_in: {e}"))?,
|
||||
dw_out: GpuTensor::zeros(&[1, hidden_dim], stream)
|
||||
.map_err(|e| anyhow!("bwd grads dw_out: {e}"))?,
|
||||
db_out: GpuTensor::zeros(&[1], stream)
|
||||
.map_err(|e| anyhow!("bwd grads db_out: {e}"))?,
|
||||
n_batch, seq_len, in_dim, hidden_dim, state_dim,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Pre-allocated scratch for [`Mamba2Block::backward_from_h_enriched_seq_into`].
|
||||
/// Holds the 4 BIG per-step scratch buffers (the ones at ~6MB each that
|
||||
/// dominated the per-call `alloc_zeros` churn). Smaller buffers
|
||||
@@ -874,6 +1002,95 @@ impl Mamba2Block {
|
||||
Ok((h_enriched_seq, cache))
|
||||
}
|
||||
|
||||
/// Zero-allocation variant of [`forward_train_seq`] — writes all
|
||||
/// intermediates into caller-provided pre-allocated buffers
|
||||
/// ([`Mamba2BlockForwardScratch`]). No `GpuTensor::zeros` / cuBLAS
|
||||
/// alloc per call, no `cudaMalloc`. The scratch's `h_enriched_seq`
|
||||
/// IS the output (no fresh tensor returned).
|
||||
///
|
||||
/// Designed for the PerceptionTrainer hot path and CUDA Graph
|
||||
/// capture. The scratch's contents are valid until the next
|
||||
/// `forward_train_seq_into` call.
|
||||
pub fn forward_train_seq_into(
|
||||
&self,
|
||||
input: &GpuTensor,
|
||||
scratch: &mut Mamba2BlockForwardScratch,
|
||||
) -> Result<()> {
|
||||
let c = &self.config;
|
||||
let n_batch = match input.shape() {
|
||||
[b, k, d] if *k == c.seq_len && *d == c.in_dim => *b,
|
||||
shape => {
|
||||
return Err(anyhow!(
|
||||
"forward_train_seq_into: expected [B, {}, {}], got {:?}",
|
||||
c.seq_len, c.in_dim, shape
|
||||
));
|
||||
}
|
||||
};
|
||||
anyhow::ensure!(
|
||||
scratch.n_batch == n_batch
|
||||
&& scratch.seq_len == c.seq_len
|
||||
&& scratch.in_dim == c.in_dim
|
||||
&& scratch.hidden_dim == c.hidden_dim
|
||||
&& scratch.state_dim == c.state_dim,
|
||||
"fwd scratch shape mismatch: expected ({},{},{},{},{}) got ({},{},{},{},{})",
|
||||
n_batch, c.seq_len, c.in_dim, c.hidden_dim, c.state_dim,
|
||||
scratch.n_batch, scratch.seq_len, scratch.in_dim,
|
||||
scratch.hidden_dim, scratch.state_dim
|
||||
);
|
||||
let n_rows = n_batch * c.seq_len;
|
||||
|
||||
// input_2d: reshape view of input (same device storage, fresh
|
||||
// wrapper — no allocation; Arc::clone on the CudaSlice).
|
||||
let input_2d = GpuTensor::new(input.cuda_data().clone(), vec![n_rows, c.in_dim])
|
||||
.map_err(|e| anyhow!("reshape input → 2D: {e}"))?;
|
||||
|
||||
// 1. x = input_2d @ W_in.T + b_in (writes into scratch.x).
|
||||
self.w_in.inner.forward_with_slices_into(
|
||||
&input_2d, &self.w_in.weight, &self.w_in.bias,
|
||||
&self.cublas, &self.stream, &mut scratch.x,
|
||||
).map_err(|e| anyhow!("w_in fwd_into: {e}"))?;
|
||||
|
||||
// 2. a_proj = x @ W_a.T + b_a.
|
||||
self.w_a.inner.forward_with_slices_into(
|
||||
&scratch.x, &self.w_a.weight, &self.w_a.bias,
|
||||
&self.cublas, &self.stream, &mut scratch.a_proj,
|
||||
).map_err(|e| anyhow!("w_a fwd_into: {e}"))?;
|
||||
|
||||
// 3. b_proj = x @ W_b.T + b_b.
|
||||
self.w_b.inner.forward_with_slices_into(
|
||||
&scratch.x, &self.w_b.weight, &self.w_b.bias,
|
||||
&self.cublas, &self.stream, &mut scratch.b_proj,
|
||||
).map_err(|e| anyhow!("w_b fwd_into: {e}"))?;
|
||||
|
||||
// 4. scan_fwd_seq → scratch.h_enriched_seq.
|
||||
// h_s2 stays zero from construction (never written in supervised path).
|
||||
let block_threads: u32 = 32;
|
||||
let grid_y: u32 =
|
||||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, grid_y, 1),
|
||||
block_dim: (block_threads, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let n_i32 = n_batch as i32;
|
||||
let k_i32 = c.seq_len as i32;
|
||||
let sh2_i32 = c.hidden_dim as i32;
|
||||
let st_i32 = c.state_dim as i32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_fwd_seq)
|
||||
.arg(scratch.a_proj.cuda_data())
|
||||
.arg(scratch.b_proj.cuda_data())
|
||||
.arg(&self.w_c)
|
||||
.arg(scratch.h_s2.cuda_data())
|
||||
.arg(scratch.h_enriched_seq.data_mut())
|
||||
.arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32)
|
||||
.launch(cfg)
|
||||
.map_err(|e| anyhow!("scan_fwd_seq_into launch: {e}"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Backward chain paired with [`forward_train_seq`]. `d_h_enriched_seq`
|
||||
/// has shape `[N, K, hidden_dim]` matching `cache.h_enriched_seq`.
|
||||
/// Returns all nine parameter gradients (`dw_out` / `db_out` zeroed,
|
||||
@@ -1178,6 +1395,148 @@ impl Mamba2Block {
|
||||
})
|
||||
}
|
||||
|
||||
/// Fully-pre-allocated backward — companion to [`forward_train_seq_into`].
|
||||
/// Uses [`Mamba2BackwardScratch`] for the big per-channel scan
|
||||
/// buffers, [`Mamba2BackwardGradsBuffers`] for the cuBLAS linear
|
||||
/// backward outputs + reduction-result tensors. Zero allocation in
|
||||
/// the call. Caller reads the final grads from `grads_buffers`.
|
||||
///
|
||||
/// `input` is the same `[B, K, in_dim]` tensor passed to the
|
||||
/// matching `forward_train_seq_into` call. Needed for the W_in
|
||||
/// backward (which dots dY^T against the original X).
|
||||
pub fn backward_from_h_enriched_seq_full_into(
|
||||
&self,
|
||||
input: &GpuTensor,
|
||||
fwd_scratch: &Mamba2BlockForwardScratch,
|
||||
d_h_enriched_seq: &GpuTensor,
|
||||
bwd_scratch: &mut Mamba2BackwardScratch,
|
||||
grads_buffers: &mut Mamba2BackwardGradsBuffers,
|
||||
) -> Result<()> {
|
||||
let c = &self.config;
|
||||
let n_batch = fwd_scratch.n_batch;
|
||||
anyhow::ensure!(
|
||||
n_batch == bwd_scratch.n_batch
|
||||
&& n_batch == grads_buffers.n_batch
|
||||
&& c.seq_len == bwd_scratch.seq_len
|
||||
&& c.seq_len == grads_buffers.seq_len
|
||||
&& c.hidden_dim == bwd_scratch.hidden_dim
|
||||
&& c.hidden_dim == grads_buffers.hidden_dim
|
||||
&& c.state_dim == bwd_scratch.state_dim
|
||||
&& c.state_dim == grads_buffers.state_dim,
|
||||
"backward_full_into: scratch shape mismatch"
|
||||
);
|
||||
anyhow::ensure!(
|
||||
d_h_enriched_seq.shape() == [n_batch, c.seq_len, c.hidden_dim],
|
||||
"d_h_enriched_seq shape {:?} != [{}, {}, {}]",
|
||||
d_h_enriched_seq.shape(), n_batch, c.seq_len, c.hidden_dim
|
||||
);
|
||||
|
||||
// ── 5′. Scan backward kernel → bwd_scratch slots ──────────────
|
||||
let block_threads: u32 = 32;
|
||||
let grid_y_h: u32 =
|
||||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||||
let bwd_cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, grid_y_h, 1),
|
||||
block_dim: (block_threads, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let n_i32 = n_batch as i32;
|
||||
let k_i32 = c.seq_len as i32;
|
||||
let sh2_i32 = c.hidden_dim as i32;
|
||||
let st_i32 = c.state_dim as i32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_bwd_seq)
|
||||
.arg(fwd_scratch.a_proj.cuda_data())
|
||||
.arg(fwd_scratch.b_proj.cuda_data())
|
||||
.arg(d_h_enriched_seq.cuda_data())
|
||||
.arg(&self.w_c)
|
||||
.arg(&mut bwd_scratch.d_a_per_channel)
|
||||
.arg(&mut bwd_scratch.d_b_per_channel)
|
||||
.arg(&mut bwd_scratch.d_w_c_per_sample)
|
||||
.arg(&mut bwd_scratch.d_h_s2)
|
||||
.arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32)
|
||||
.launch(bwd_cfg)
|
||||
.map_err(|e| anyhow!("scan_bwd_seq_full_into launch: {e}"))?;
|
||||
}
|
||||
|
||||
// ── Reductions: per-channel → flat. Writes directly into
|
||||
// grads_buffers.{d_a_proj_2d, d_b_proj_2d, dw_c} since
|
||||
// their flat layouts match [N, K, state_d] / [sh2, state_d].
|
||||
let red_grid_z: u32 =
|
||||
((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||||
let red_cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, c.seq_len as u32, red_grid_z),
|
||||
block_dim: (block_threads, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_reduce_d_proj)
|
||||
.arg(&bwd_scratch.d_a_per_channel)
|
||||
.arg(grads_buffers.d_a_proj_2d.data_mut())
|
||||
.arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32)
|
||||
.launch(red_cfg)
|
||||
.map_err(|e| anyhow!("reduce d_a_proj full_into: {e}"))?;
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_reduce_d_proj)
|
||||
.arg(&bwd_scratch.d_b_per_channel)
|
||||
.arg(grads_buffers.d_b_proj_2d.data_mut())
|
||||
.arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32)
|
||||
.launch(red_cfg)
|
||||
.map_err(|e| anyhow!("reduce d_b_proj full_into: {e}"))?;
|
||||
}
|
||||
let red_w_c_cfg = LaunchConfig {
|
||||
grid_dim: (c.hidden_dim as u32, red_grid_z, 1),
|
||||
block_dim: (block_threads, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_reduce_d_w_c)
|
||||
.arg(&bwd_scratch.d_w_c_per_sample)
|
||||
.arg(grads_buffers.dw_c.data_mut())
|
||||
.arg(&n_i32).arg(&sh2_i32).arg(&st_i32)
|
||||
.launch(red_w_c_cfg)
|
||||
.map_err(|e| anyhow!("reduce dw_c full_into: {e}"))?;
|
||||
}
|
||||
|
||||
// ── Linear backward: w_b, w_a, w_in (all _into variants). ────
|
||||
let x_act = LinearActivations { input: fwd_scratch.x.clone() };
|
||||
self.w_b.inner.backward_with_slices_into(
|
||||
&grads_buffers.d_b_proj_2d, &x_act, &self.w_b.weight,
|
||||
&self.cublas, &self.stream,
|
||||
&mut grads_buffers.dw_b, &mut grads_buffers.db_b, &mut grads_buffers.d_x_from_b,
|
||||
).map_err(|e| anyhow!("w_b bwd_into: {e}"))?;
|
||||
|
||||
self.w_a.inner.backward_with_slices_into(
|
||||
&grads_buffers.d_a_proj_2d, &x_act, &self.w_a.weight,
|
||||
&self.cublas, &self.stream,
|
||||
&mut grads_buffers.dw_a, &mut grads_buffers.db_a, &mut grads_buffers.d_x_from_a,
|
||||
).map_err(|e| anyhow!("w_a bwd_into: {e}"))?;
|
||||
|
||||
// d_x = d_x_from_a + d_x_from_b (in-place add — zero alloc).
|
||||
grads_buffers.d_x_from_a
|
||||
.add_into(&grads_buffers.d_x_from_b, &grads_buffers.d_x, &self.stream)
|
||||
.map_err(|e| anyhow!("d_x add_into: {e}"))?;
|
||||
|
||||
// W_in backward: input_2d is a fresh wrapper around the
|
||||
// original input pointer (Arc::clone — no allocation).
|
||||
let n_rows = n_batch * c.seq_len;
|
||||
let input_2d = GpuTensor::new(input.cuda_data().clone(), vec![n_rows, c.in_dim])
|
||||
.map_err(|e| anyhow!("reshape input for w_in bwd_into: {e}"))?;
|
||||
let input_act = LinearActivations { input: input_2d };
|
||||
self.w_in.inner.backward_with_slices_into(
|
||||
&grads_buffers.d_x, &input_act, &self.w_in.weight,
|
||||
&self.cublas, &self.stream,
|
||||
&mut grads_buffers.dw_in, &mut grads_buffers.db_in, &mut grads_buffers.d_x_from_in,
|
||||
).map_err(|e| anyhow!("w_in bwd_into: {e}"))?;
|
||||
|
||||
// dw_out / db_out stay zero — W_out is unused in the seq path.
|
||||
// (The grads_buffers init already created them zero; no action.)
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Total trainable parameter count (sum of all projections + W_c).
|
||||
pub fn param_count(&self) -> usize {
|
||||
let c = &self.config;
|
||||
@@ -1351,6 +1710,67 @@ impl Mamba2AdamW {
|
||||
pub fn set_learning_rate(&mut self, lr: f32) {
|
||||
self.config.lr = lr;
|
||||
}
|
||||
|
||||
/// AdamW step that reads gradients directly from a
|
||||
/// [`Mamba2BackwardGradsBuffers`] (the pre-allocated buffer set
|
||||
/// produced by [`Mamba2Block::backward_from_h_enriched_seq_full_into`]).
|
||||
/// Avoids constructing a temporary [`Mamba2BackwardGrads`] wrapper
|
||||
/// every step.
|
||||
pub fn step_from_buffers(
|
||||
&mut self,
|
||||
block: &mut Mamba2Block,
|
||||
grads: &Mamba2BackwardGradsBuffers,
|
||||
) -> Result<()> {
|
||||
self.step_count += 1;
|
||||
let t = self.step_count;
|
||||
|
||||
let grad_scale = if let Some(max_norm) = self.config.grad_clip_max_norm {
|
||||
let mut total_sq = 0.0_f32;
|
||||
for slice in [
|
||||
grads.dw_in.cuda_data(), grads.db_in.cuda_data(),
|
||||
grads.dw_a.cuda_data(), grads.db_a.cuda_data(),
|
||||
grads.dw_b.cuda_data(), grads.db_b.cuda_data(),
|
||||
grads.dw_c.cuda_data(),
|
||||
grads.dw_out.cuda_data(),grads.db_out.cuda_data(),
|
||||
] {
|
||||
let mut host = vec![0.0_f32; slice.len()];
|
||||
self.stream.memcpy_dtoh(slice, &mut host)
|
||||
.map_err(|e| anyhow!("grad-norm dtoh: {e}"))?;
|
||||
for g in &host {
|
||||
total_sq += g * g;
|
||||
}
|
||||
}
|
||||
let norm = total_sq.sqrt();
|
||||
if norm > max_norm { max_norm / norm } else { 1.0 }
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
|
||||
let stream = &self.stream;
|
||||
let kernel = &self.kernel;
|
||||
let cfg = &self.config;
|
||||
|
||||
adamw_apply(stream, kernel, cfg, block.w_in.weight.len(),
|
||||
&mut block.w_in.weight, grads.dw_in.cuda_data(), &mut self.s_w_in, t, grad_scale)?;
|
||||
adamw_apply(stream, kernel, cfg, block.w_in.bias.len(),
|
||||
&mut block.w_in.bias, grads.db_in.cuda_data(), &mut self.s_b_in, t, grad_scale)?;
|
||||
adamw_apply(stream, kernel, cfg, block.w_a.weight.len(),
|
||||
&mut block.w_a.weight, grads.dw_a.cuda_data(), &mut self.s_w_a, t, grad_scale)?;
|
||||
adamw_apply(stream, kernel, cfg, block.w_a.bias.len(),
|
||||
&mut block.w_a.bias, grads.db_a.cuda_data(), &mut self.s_b_a, t, grad_scale)?;
|
||||
adamw_apply(stream, kernel, cfg, block.w_b.weight.len(),
|
||||
&mut block.w_b.weight, grads.dw_b.cuda_data(), &mut self.s_w_b, t, grad_scale)?;
|
||||
adamw_apply(stream, kernel, cfg, block.w_b.bias.len(),
|
||||
&mut block.w_b.bias, grads.db_b.cuda_data(), &mut self.s_b_b, t, grad_scale)?;
|
||||
adamw_apply(stream, kernel, cfg, block.w_c.len(),
|
||||
&mut block.w_c, grads.dw_c.cuda_data(), &mut self.s_w_c, t, grad_scale)?;
|
||||
adamw_apply(stream, kernel, cfg, block.w_out.weight.len(),
|
||||
&mut block.w_out.weight, grads.dw_out.cuda_data(), &mut self.s_w_out, t, grad_scale)?;
|
||||
adamw_apply(stream, kernel, cfg, block.w_out.bias.len(),
|
||||
&mut block.w_out.bias, grads.db_out.cuda_data(), &mut self.s_b_out, t, grad_scale)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// One AdamW kernel launch for a single parameter tensor. Free function
|
||||
|
||||
@@ -47,7 +47,8 @@ use rand_chacha::ChaCha8Rng;
|
||||
use crate::cfc::snap_features::{Mbp10RawInput, ES_TICK_SIZE, FEATURE_DIM, REGIME_DIM};
|
||||
use crate::heads::{HIDDEN_DIM, N_HORIZONS};
|
||||
use crate::mamba2_block::{
|
||||
Mamba2AdamW, Mamba2AdamWConfig, Mamba2BackwardScratch, Mamba2Block, Mamba2BlockConfig,
|
||||
Mamba2AdamW, Mamba2AdamWConfig, Mamba2BackwardGradsBuffers, Mamba2BackwardScratch,
|
||||
Mamba2Block, Mamba2BlockConfig, Mamba2BlockForwardScratch,
|
||||
};
|
||||
use crate::pinned_mem::{MappedF32Buffer, MappedI32Buffer, MappedI64Buffer};
|
||||
use crate::trainer::optim::AdamW;
|
||||
@@ -130,9 +131,31 @@ pub struct PerceptionTrainer {
|
||||
// Mamba2 encoder block + its optimizer
|
||||
pub mamba2: Mamba2Block,
|
||||
pub mamba2_adamw: Mamba2AdamW,
|
||||
/// Pre-allocated scratch for the Mamba2 seq backward — eliminates
|
||||
/// ~10-20ms/step of `alloc_zeros` churn in the hot path.
|
||||
/// Pre-allocated forward intermediates for Mamba2 (input projection
|
||||
/// output x, A/B projections, h_s2 residual, h_enriched_seq scan
|
||||
/// output). Construct once at trainer init; reused every step.
|
||||
mamba2_fwd_scratch: Mamba2BlockForwardScratch,
|
||||
/// Pre-allocated scratch for the Mamba2 seq backward — the four
|
||||
/// largest per-channel buffers (~6 MB each at B=8, K=96).
|
||||
mamba2_bwd_scratch: Mamba2BackwardScratch,
|
||||
/// Pre-allocated outputs for Mamba2 seq backward (cuBLAS-projection
|
||||
/// dw/db tensors + reduction-result tensors + d_x intermediates).
|
||||
/// Read by Mamba2AdamW::step_from_buffers — no temporary
|
||||
/// Mamba2BackwardGrads wrapper allocated per step.
|
||||
mamba2_grads_buffers: Mamba2BackwardGradsBuffers,
|
||||
/// Pre-allocated input window for snap_features → Mamba2 fwd.
|
||||
/// [B, K, FEATURE_DIM] — overwritten each step by the batched
|
||||
/// snap_feature kernel.
|
||||
window_tensor_d: GpuTensor,
|
||||
/// Pre-allocated transpose of Mamba2's h_enriched_seq into [K, B, H]
|
||||
/// layout for contiguous per-K slot access in the trainer loop.
|
||||
h_enriched_seq_t_d: GpuTensor,
|
||||
/// Pre-allocated per-K gradient accumulator in [K, B, H] layout.
|
||||
/// Written by the reverse-order backward K loop; transposed back
|
||||
/// to [B, K, H] for Mamba2 backward consumption.
|
||||
grad_h_enriched_seq_t_d: GpuTensor,
|
||||
/// Pre-allocated [B, K, H] grad input to Mamba2 backward.
|
||||
grad_h_enriched_seq_d: GpuTensor,
|
||||
|
||||
// CfC + heads weights + their AdamWs (6 groups — tau is trained now).
|
||||
pub w_in_d: CudaSlice<f32>,
|
||||
@@ -268,9 +291,23 @@ impl PerceptionTrainer {
|
||||
},
|
||||
)
|
||||
.context("Mamba2AdamW::new")?;
|
||||
let mamba2_fwd_scratch = Mamba2BlockForwardScratch::new(
|
||||
&stream, cfg.n_batch, cfg.seq_len, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||||
).context("Mamba2BlockForwardScratch::new")?;
|
||||
let mamba2_bwd_scratch = Mamba2BackwardScratch::new(
|
||||
&stream, cfg.n_batch, cfg.seq_len, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||||
).context("Mamba2BackwardScratch::new")?;
|
||||
let mamba2_grads_buffers = Mamba2BackwardGradsBuffers::new(
|
||||
&stream, cfg.n_batch, cfg.seq_len, FEATURE_DIM, HIDDEN_DIM, cfg.mamba2_state_dim,
|
||||
).context("Mamba2BackwardGradsBuffers::new")?;
|
||||
let window_tensor_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, FEATURE_DIM], &stream)
|
||||
.map_err(|e| anyhow::anyhow!("window_tensor_d alloc: {e}"))?;
|
||||
let h_enriched_seq_t_d = GpuTensor::zeros(&[cfg.seq_len, cfg.n_batch, HIDDEN_DIM], &stream)
|
||||
.map_err(|e| anyhow::anyhow!("h_enriched_seq_t_d alloc: {e}"))?;
|
||||
let grad_h_enriched_seq_t_d = GpuTensor::zeros(&[cfg.seq_len, cfg.n_batch, HIDDEN_DIM], &stream)
|
||||
.map_err(|e| anyhow::anyhow!("grad_h_enriched_seq_t_d alloc: {e}"))?;
|
||||
let grad_h_enriched_seq_d = GpuTensor::zeros(&[cfg.n_batch, cfg.seq_len, HIDDEN_DIM], &stream)
|
||||
.map_err(|e| anyhow::anyhow!("grad_h_enriched_seq_d alloc: {e}"))?;
|
||||
|
||||
// CfC weights (input = h_enriched [HIDDEN_DIM], output = [HIDDEN_DIM])
|
||||
let mut r = ChaCha8Rng::seed_from_u64(cfg.seed);
|
||||
@@ -386,7 +423,13 @@ impl PerceptionTrainer {
|
||||
transpose_3d_fn,
|
||||
mamba2,
|
||||
mamba2_adamw,
|
||||
mamba2_fwd_scratch,
|
||||
mamba2_bwd_scratch,
|
||||
mamba2_grads_buffers,
|
||||
window_tensor_d,
|
||||
h_enriched_seq_t_d,
|
||||
grad_h_enriched_seq_t_d,
|
||||
grad_h_enriched_seq_d,
|
||||
opt_w_in,
|
||||
opt_w_rec,
|
||||
opt_b,
|
||||
@@ -475,10 +518,8 @@ impl PerceptionTrainer {
|
||||
// one DtoD per array → device; then one kernel launch
|
||||
// with B*K threads. Output written directly into the
|
||||
// window tensor's storage.
|
||||
let mut window_tensor = GpuTensor::zeros(
|
||||
&[b_sz, k_seq, FEATURE_DIM],
|
||||
&self.stream,
|
||||
).map_err(|e| anyhow::anyhow!("window alloc: {e}"))?;
|
||||
// Use pre-allocated window_tensor_d — no per-step alloc.
|
||||
debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, FEATURE_DIM]);
|
||||
|
||||
let total_snaps = b_sz * k_seq;
|
||||
debug_assert!(total_snaps <= self.bk_capacity);
|
||||
@@ -589,23 +630,17 @@ impl PerceptionTrainer {
|
||||
.arg(&self.trade_count_all_d)
|
||||
.arg(&self.ts_ns_all_d).arg(&self.prev_ts_ns_all_d)
|
||||
.arg(&tick_size).arg(&n_total_i32)
|
||||
.arg(window_tensor.data_mut());
|
||||
.arg(self.window_tensor_d.data_mut());
|
||||
launch.launch(snap_cfg).context("snap_batched fwd")?;
|
||||
}
|
||||
self.stream.synchronize().context("snap_batched sync")?;
|
||||
|
||||
// ── 2. Mamba2 per-step forward → h_enriched_seq [B, K, HIDDEN_DIM].
|
||||
let (h_enriched_seq, cache) = self
|
||||
.mamba2
|
||||
.forward_train_seq(&window_tensor)
|
||||
.context("mamba2 forward_train_seq")?;
|
||||
// ── 2. Mamba2 per-step forward — writes into self.mamba2_fwd_scratch.
|
||||
self.mamba2
|
||||
.forward_train_seq_into(&self.window_tensor_d, &mut self.mamba2_fwd_scratch)
|
||||
.context("mamba2 forward_train_seq_into")?;
|
||||
|
||||
// ── 2b. Transpose [B, K, H] → [K, B, H] so the K-loop can
|
||||
// slice contiguous [B, H] chunks per step.
|
||||
let mut h_enriched_seq_t = GpuTensor::zeros(
|
||||
&[k_seq, b_sz, HIDDEN_DIM],
|
||||
&self.stream,
|
||||
).map_err(|e| anyhow::anyhow!("h_enriched_seq_t alloc: {e}"))?;
|
||||
// ── 2b. Transpose [B, K, H] → [K, B, H] into pre-allocated buffer.
|
||||
{
|
||||
let block_n3: u32 = 32;
|
||||
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
|
||||
@@ -619,8 +654,8 @@ impl PerceptionTrainer {
|
||||
let n3 = HIDDEN_DIM as i32;
|
||||
let mut launch = self.stream.launch_builder(&self.transpose_3d_fn);
|
||||
launch
|
||||
.arg(h_enriched_seq.cuda_data())
|
||||
.arg(h_enriched_seq_t.data_mut())
|
||||
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
|
||||
.arg(self.h_enriched_seq_t_d.data_mut())
|
||||
.arg(&n1).arg(&n2).arg(&n3);
|
||||
unsafe { launch.launch(cfg_tx).context("transpose h_enriched fwd")?; }
|
||||
}
|
||||
@@ -719,7 +754,7 @@ impl PerceptionTrainer {
|
||||
p
|
||||
};
|
||||
let henr_t_base = {
|
||||
let (p, _g) = h_enriched_seq_t.cuda_data().device_ptr(&self.stream);
|
||||
let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
let (h_per_k_base, _g_hpk) = self.h_new_per_k_d.device_ptr_mut(&self.stream);
|
||||
@@ -778,16 +813,8 @@ impl PerceptionTrainer {
|
||||
unsafe { launch.launch(bce_cfg).context("bce launch")?; }
|
||||
}
|
||||
|
||||
// ── 6. Reverse-order backward K loop. Same recurrence carry as
|
||||
// the unbatched version, but every kernel is batched over
|
||||
// B samples. grad_h_enriched_seq_t is [K, B, H] — slot k
|
||||
// contiguous, written by the cfc_step_bwd batched kernel
|
||||
// directly into its grad_x output (which IS sized [B, n_in]).
|
||||
let mut grad_h_enriched_seq_t = GpuTensor::zeros(
|
||||
&[k_seq, b_sz, HIDDEN_DIM],
|
||||
&self.stream,
|
||||
).map_err(|e| anyhow::anyhow!("grad_h_enriched_seq_t alloc: {e}"))?;
|
||||
|
||||
// ── 6. Reverse-order backward K loop using pre-allocated
|
||||
// grad_h_enriched_seq_t_d as the per-K slot output.
|
||||
self.stream.memset_zeros(&mut self.grad_h_carry_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_h_carry: {e}"))?;
|
||||
|
||||
@@ -796,13 +823,13 @@ impl PerceptionTrainer {
|
||||
p
|
||||
};
|
||||
let henr_t_base_bwd = {
|
||||
let (p, _g) = h_enriched_seq_t.cuda_data().device_ptr(&self.stream);
|
||||
let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
let (h_per_k_base_bwd, _g_hpk_bwd) = self.h_new_per_k_d.device_ptr_mut(&self.stream);
|
||||
let (probs_base_bwd, _g_probs_bwd) = self.probs_per_k_d.device_ptr_mut(&self.stream);
|
||||
let (gprobs_base_bwd, _g_gprobs_bwd) = self.grad_probs_per_k_d.device_ptr_mut(&self.stream);
|
||||
let (grad_henr_t_base, _g_ghen_t_mut) = grad_h_enriched_seq_t.data_mut().device_ptr_mut(&self.stream);
|
||||
let (grad_henr_t_base, _g_ghen_t_mut) = self.grad_h_enriched_seq_t_d.data_mut().device_ptr_mut(&self.stream);
|
||||
|
||||
for k in (0..k_seq).rev() {
|
||||
let h_new_k_ptr = h_per_k_base_bwd + (k * kb_hid_bytes) as u64;
|
||||
@@ -853,13 +880,8 @@ impl PerceptionTrainer {
|
||||
// ── 7. Sync once before optimizer step; download loss scalar.
|
||||
self.stream.synchronize().context("bwd loop sync")?;
|
||||
|
||||
// ── 7b. Transpose grad_h_enriched_seq_t [K, B, H] → [B, K, H]
|
||||
// so Mamba2.backward_from_h_enriched_seq sees the layout
|
||||
// it expects.
|
||||
let mut grad_h_enriched_seq = GpuTensor::zeros(
|
||||
&[b_sz, k_seq, HIDDEN_DIM],
|
||||
&self.stream,
|
||||
).map_err(|e| anyhow::anyhow!("grad_h_enriched_seq alloc: {e}"))?;
|
||||
// ── 7b. Transpose grad_h_enriched_seq_t_d [K, B, H] → [B, K, H]
|
||||
// (pre-allocated grad_h_enriched_seq_d).
|
||||
{
|
||||
let block_n3: u32 = 32;
|
||||
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
|
||||
@@ -873,17 +895,23 @@ impl PerceptionTrainer {
|
||||
let n3 = HIDDEN_DIM as i32;
|
||||
let mut launch = self.stream.launch_builder(&self.transpose_3d_fn);
|
||||
launch
|
||||
.arg(grad_h_enriched_seq_t.cuda_data())
|
||||
.arg(grad_h_enriched_seq.data_mut())
|
||||
.arg(self.grad_h_enriched_seq_t_d.cuda_data())
|
||||
.arg(self.grad_h_enriched_seq_d.data_mut())
|
||||
.arg(&n1).arg(&n2).arg(&n3);
|
||||
unsafe { launch.launch(cfg_tx).context("transpose grad bwd")?; }
|
||||
}
|
||||
|
||||
// ── 8. Mamba2 backward — single call consumes grad_h_enriched_seq.
|
||||
let mamba2_grads = self
|
||||
.mamba2
|
||||
.backward_from_h_enriched_seq_into(&cache, &grad_h_enriched_seq, &mut self.mamba2_bwd_scratch)
|
||||
.context("mamba2 backward_from_h_enriched_seq_into")?;
|
||||
// ── 8. Mamba2 backward — fully pre-allocated path. Writes all
|
||||
// grads into self.mamba2_grads_buffers; no allocation.
|
||||
self.mamba2
|
||||
.backward_from_h_enriched_seq_full_into(
|
||||
&self.window_tensor_d,
|
||||
&self.mamba2_fwd_scratch,
|
||||
&self.grad_h_enriched_seq_d,
|
||||
&mut self.mamba2_bwd_scratch,
|
||||
&mut self.mamba2_grads_buffers,
|
||||
)
|
||||
.context("mamba2 backward_from_h_enriched_seq_full_into")?;
|
||||
|
||||
// ── 9. Apply AdamW updates on all 7 param groups (added tau).
|
||||
self.opt_w_in.step(&mut self.w_in_d, &self.grad_w_in_d)?;
|
||||
@@ -893,8 +921,8 @@ impl PerceptionTrainer {
|
||||
self.opt_heads_w.step(&mut self.heads_w_d, &self.grad_heads_w_d)?;
|
||||
self.opt_heads_b.step(&mut self.heads_b_d, &self.grad_heads_b_d)?;
|
||||
self.mamba2_adamw
|
||||
.step(&mut self.mamba2, &mamba2_grads)
|
||||
.context("mamba2 AdamW step")?;
|
||||
.step_from_buffers(&mut self.mamba2, &self.mamba2_grads_buffers)
|
||||
.context("mamba2 AdamW step_from_buffers")?;
|
||||
|
||||
// Final: read the single loss scalar back to host. This is the
|
||||
// ONLY post-step download in the hot path.
|
||||
@@ -937,10 +965,8 @@ impl PerceptionTrainer {
|
||||
b_sz, snapshots_batch.len(), labels_batch.len()
|
||||
);
|
||||
|
||||
// Build snap_feature window [B, K, FEATURE_DIM] — fused batched
|
||||
// pack+upload+launch path (same as step_batched).
|
||||
let mut window_tensor = GpuTensor::zeros(&[b_sz, k_seq, FEATURE_DIM], &self.stream)
|
||||
.map_err(|e| anyhow::anyhow!("eval window alloc: {e}"))?;
|
||||
// Use pre-allocated window_tensor_d — no per-step alloc.
|
||||
debug_assert_eq!(self.window_tensor_d.shape(), &[b_sz, k_seq, FEATURE_DIM]);
|
||||
let total_snaps = b_sz * k_seq;
|
||||
debug_assert!(total_snaps <= self.bk_capacity);
|
||||
{
|
||||
@@ -1039,19 +1065,17 @@ impl PerceptionTrainer {
|
||||
.arg(&self.trade_count_all_d)
|
||||
.arg(&self.ts_ns_all_d).arg(&self.prev_ts_ns_all_d)
|
||||
.arg(&tick_size).arg(&n_total_i32)
|
||||
.arg(window_tensor.data_mut());
|
||||
.arg(self.window_tensor_d.data_mut());
|
||||
launch.launch(snap_cfg).context("eval snap_batched fwd")?;
|
||||
}
|
||||
self.stream.synchronize().context("eval snap_batched sync")?;
|
||||
|
||||
let (h_enriched_seq, _cache) = self
|
||||
.mamba2
|
||||
.forward_train_seq(&window_tensor)
|
||||
.context("eval mamba2 fwd")?;
|
||||
// Mamba2 fwd into pre-allocated fwd_scratch.
|
||||
self.mamba2
|
||||
.forward_train_seq_into(&self.window_tensor_d, &mut self.mamba2_fwd_scratch)
|
||||
.context("eval mamba2 fwd_into")?;
|
||||
|
||||
// Transpose Mamba2 output [B, K, H] → [K, B, H].
|
||||
let mut h_enriched_seq_t = GpuTensor::zeros(&[k_seq, b_sz, HIDDEN_DIM], &self.stream)
|
||||
.map_err(|e| anyhow::anyhow!("eval h_enriched_seq_t alloc: {e}"))?;
|
||||
// Transpose [B, K, H] → [K, B, H] into pre-allocated buffer.
|
||||
{
|
||||
let block_n3: u32 = 32;
|
||||
let grid_z = (HIDDEN_DIM as u32).div_ceil(block_n3);
|
||||
@@ -1065,8 +1089,8 @@ impl PerceptionTrainer {
|
||||
let n3 = HIDDEN_DIM as i32;
|
||||
let mut launch = self.stream.launch_builder(&self.transpose_3d_fn);
|
||||
launch
|
||||
.arg(h_enriched_seq.cuda_data())
|
||||
.arg(h_enriched_seq_t.data_mut())
|
||||
.arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
|
||||
.arg(self.h_enriched_seq_t_d.data_mut())
|
||||
.arg(&n1).arg(&n2).arg(&n3);
|
||||
unsafe { launch.launch(cfg_tx).context("eval transpose h_enriched")?; }
|
||||
}
|
||||
@@ -1112,7 +1136,7 @@ impl PerceptionTrainer {
|
||||
p
|
||||
};
|
||||
let henr_t_base = {
|
||||
let (p, _g) = h_enriched_seq_t.cuda_data().device_ptr(&self.stream);
|
||||
let (p, _g) = self.h_enriched_seq_t_d.cuda_data().device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
let (h_per_k_base, _g_hpk) = self.h_new_per_k_d.device_ptr_mut(&self.stream);
|
||||
|
||||
@@ -69,6 +69,30 @@ impl ElementwiseKernels {
|
||||
})
|
||||
}
|
||||
|
||||
/// In-place variant of [`binary`] — writes into caller-provided
|
||||
/// pre-allocated `out`. Caller is responsible for ensuring `out`
|
||||
/// has at least `n` elements. Used by hot-path callers that
|
||||
/// require zero per-step allocation for CUDA Graph capture.
|
||||
pub fn binary_into(
|
||||
&self,
|
||||
a: &CudaSlice<f32>,
|
||||
b: &CudaSlice<f32>,
|
||||
out: &CudaSlice<f32>,
|
||||
n: usize,
|
||||
op: i32,
|
||||
) -> Result<(), MLError> {
|
||||
let n_i32 = n as i32;
|
||||
let cfg = elem_cfg(n);
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.binary_fn)
|
||||
.arg(a).arg(b).arg(out).arg(&n_i32).arg(&op)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("elementwise_binary_into(op={op}): {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Element-wise binary operation on two same-shape buffers.
|
||||
///
|
||||
/// `op`: 0=add, 1=sub, 2=mul, 3=div, 4=min, 5=max
|
||||
|
||||
@@ -422,6 +422,22 @@ impl GpuTensor {
|
||||
Ok(Self { data: out, shape: self.shape.clone() })
|
||||
}
|
||||
|
||||
/// In-place variant of [`add`] — writes into caller-provided
|
||||
/// pre-allocated `out`. All three tensors must share `shape`.
|
||||
/// Used by hot-path callers that require zero per-step allocation
|
||||
/// for CUDA Graph capture.
|
||||
pub fn add_into(&self, other: &Self, out: &Self, stream: &Arc<CudaStream>) -> Result<(), MLError> {
|
||||
if self.shape != other.shape || self.shape != out.shape {
|
||||
return Err(MLError::DimensionMismatch {
|
||||
expected: self.numel(),
|
||||
actual: other.numel(),
|
||||
});
|
||||
}
|
||||
let kernels = super::elementwise::get_or_compile(stream)?;
|
||||
kernels.binary_into(&self.data, &other.data, &out.data, self.numel(), 0)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Element-wise subtraction. Shapes must match exactly.
|
||||
///
|
||||
/// GPU-native: uses `elementwise_binary` kernel (op=1).
|
||||
|
||||
@@ -462,6 +462,145 @@ impl GpuLinear {
|
||||
|
||||
Ok(LinearGrads { dw, db, dx })
|
||||
}
|
||||
|
||||
/// In-place variant of [`forward_with_slices`] — writes into
|
||||
/// caller-provided pre-allocated `y_out` instead of allocating a
|
||||
/// fresh GpuTensor per call. Used by ml-alpha's PerceptionTrainer
|
||||
/// to eliminate per-step `cudaMalloc` calls on the training hot
|
||||
/// path (CUDA Graph capture requires fixed device pointers).
|
||||
///
|
||||
/// `y_out` MUST already have shape `[batch, out_dim]`. Contents
|
||||
/// are overwritten by the cuBLAS sgemm + bias-add.
|
||||
pub fn forward_with_slices_into(
|
||||
&self,
|
||||
x: &GpuTensor,
|
||||
weight: &CudaSlice<f32>,
|
||||
bias: &CudaSlice<f32>,
|
||||
cublas: &CudaBlas,
|
||||
stream: &Arc<CudaStream>,
|
||||
y_out: &mut GpuTensor,
|
||||
) -> Result<(), MLError> {
|
||||
let batch = if x.ndim() == 1 { 1 } else { x.shape()[0] };
|
||||
let in_dim = self.in_dim;
|
||||
let out_dim = self.out_dim;
|
||||
if y_out.shape() != [batch, out_dim] {
|
||||
return Err(MLError::ModelError(format!(
|
||||
"forward_with_slices_into: y_out shape {:?} != [{}, {}]",
|
||||
y_out.shape(), batch, out_dim
|
||||
)));
|
||||
}
|
||||
|
||||
let w_ptr = raw_ptr(weight, stream);
|
||||
let x_ptr = raw_ptr(&x.data, stream);
|
||||
let y_ptr = raw_ptr_mut(&mut y_out.data, stream);
|
||||
|
||||
unsafe {
|
||||
gemm_ex_f32(
|
||||
cublas,
|
||||
cublasOperation_t::CUBLAS_OP_T,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
out_dim as i32,
|
||||
batch as i32,
|
||||
in_dim as i32,
|
||||
w_ptr,
|
||||
in_dim as i32,
|
||||
x_ptr,
|
||||
in_dim as i32,
|
||||
y_ptr,
|
||||
out_dim as i32,
|
||||
"forward_with_slices_into",
|
||||
)?;
|
||||
}
|
||||
add_bias_2d(y_out, bias, batch, out_dim, stream)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// In-place variant of [`backward_with_slices`] — writes into
|
||||
/// caller-provided pre-allocated dw / db / dx tensors. Same math
|
||||
/// as [`backward_with_slices`]; differs only in not allocating.
|
||||
///
|
||||
/// Shape contract:
|
||||
/// `dw_out` : `[out_dim, in_dim]`
|
||||
/// `db_out` : `[out_dim]`
|
||||
/// `dx_out` : `[batch, in_dim]`
|
||||
pub fn backward_with_slices_into(
|
||||
&self,
|
||||
dy: &GpuTensor,
|
||||
activations: &LinearActivations,
|
||||
weight: &CudaSlice<f32>,
|
||||
cublas: &CudaBlas,
|
||||
stream: &Arc<CudaStream>,
|
||||
dw_out: &mut GpuTensor,
|
||||
db_out: &mut GpuTensor,
|
||||
dx_out: &mut GpuTensor,
|
||||
) -> Result<(), MLError> {
|
||||
let batch = dy.shape()[0];
|
||||
let in_dim = self.in_dim;
|
||||
let out_dim = self.out_dim;
|
||||
if dw_out.shape() != [out_dim, in_dim] {
|
||||
return Err(MLError::ModelError(format!(
|
||||
"backward_with_slices_into: dw shape {:?} != [{}, {}]",
|
||||
dw_out.shape(), out_dim, in_dim
|
||||
)));
|
||||
}
|
||||
if db_out.shape() != [out_dim] {
|
||||
return Err(MLError::ModelError(format!(
|
||||
"backward_with_slices_into: db shape {:?} != [{}]",
|
||||
db_out.shape(), out_dim
|
||||
)));
|
||||
}
|
||||
if dx_out.shape() != [batch, in_dim] {
|
||||
return Err(MLError::ModelError(format!(
|
||||
"backward_with_slices_into: dx shape {:?} != [{}, {}]",
|
||||
dx_out.shape(), batch, in_dim
|
||||
)));
|
||||
}
|
||||
|
||||
let x_ptr = raw_ptr(&activations.input.data, stream);
|
||||
let dy_ptr = raw_ptr(&dy.data, stream);
|
||||
let dw_ptr = raw_ptr_mut(&mut dw_out.data, stream);
|
||||
unsafe {
|
||||
gemm_ex_f32(
|
||||
cublas,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
cublasOperation_t::CUBLAS_OP_T,
|
||||
in_dim as i32,
|
||||
out_dim as i32,
|
||||
batch as i32,
|
||||
x_ptr,
|
||||
in_dim as i32,
|
||||
dy_ptr,
|
||||
out_dim as i32,
|
||||
dw_ptr,
|
||||
in_dim as i32,
|
||||
"dW_with_slices_into",
|
||||
)?;
|
||||
}
|
||||
|
||||
reduce_sum_axis0_into(dy, batch, out_dim, stream, db_out)?;
|
||||
|
||||
let w_ptr = raw_ptr(weight, stream);
|
||||
let dy_ptr2 = raw_ptr(&dy.data, stream);
|
||||
let dx_ptr = raw_ptr_mut(&mut dx_out.data, stream);
|
||||
unsafe {
|
||||
gemm_ex_f32(
|
||||
cublas,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
in_dim as i32,
|
||||
batch as i32,
|
||||
out_dim as i32,
|
||||
w_ptr,
|
||||
in_dim as i32,
|
||||
dy_ptr2,
|
||||
out_dim as i32,
|
||||
dx_ptr,
|
||||
in_dim as i32,
|
||||
"dX_with_slices_into",
|
||||
)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Self-contained GPU linear layer that owns its weight and bias `CudaSlice`s.
|
||||
@@ -625,7 +764,27 @@ fn reduce_sum_axis0(
|
||||
cols: usize,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> Result<GpuTensor, MLError> {
|
||||
let out = GpuTensor::zeros(&[cols], stream)?;
|
||||
let mut out = GpuTensor::zeros(&[cols], stream)?;
|
||||
reduce_sum_axis0_into(x, rows, cols, stream, &mut out)?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// In-place variant — writes into caller-provided `out` (shape `[cols]`).
|
||||
/// Used by `OwnedGpuLinear::backward_with_slices_into` to eliminate
|
||||
/// per-call allocation on the training hot path.
|
||||
pub fn reduce_sum_axis0_into(
|
||||
x: &GpuTensor,
|
||||
rows: usize,
|
||||
cols: usize,
|
||||
stream: &Arc<CudaStream>,
|
||||
out: &mut GpuTensor,
|
||||
) -> Result<(), MLError> {
|
||||
if out.shape() != [cols] {
|
||||
return Err(MLError::ModelError(format!(
|
||||
"reduce_sum_axis0_into: out shape {:?} != [{}]",
|
||||
out.shape(), cols
|
||||
)));
|
||||
}
|
||||
let (_, reduce_fn) = get_bias_kernels(stream)?;
|
||||
let threads = 256_u32;
|
||||
let blocks = ((cols as u32) + threads - 1) / threads;
|
||||
@@ -646,9 +805,9 @@ fn reduce_sum_axis0(
|
||||
.arg(&rows_i32)
|
||||
.arg(&cols_i32)
|
||||
.launch(launch_cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("reduce_sum_axis0_kernel: {e}")))?;
|
||||
.map_err(|e| MLError::ModelError(format!("reduce_sum_axis0_into kernel: {e}")))?;
|
||||
}
|
||||
Ok(out)
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clone a GpuTensor by copying its data to a new allocation.
|
||||
|
||||
Reference in New Issue
Block a user