gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(phase-e-4-a): T10 Mamba2 backward chain + KC calibration

Browse Source 
T10 wires the C51 → Mamba2 backward chain in both binaries:
- New alpha_train_window_store_batched_kernel captures per-step windows
  for end-of-epoch re-forward (Mamba2 cache required for backward).
- Mamba2Block::backward_from_h_enriched lets the C51 grad-input feed
  directly into Mamba2 backward, skipping the unused W_out projection.
- alpha_c51_grad_input → backward_from_h_enriched → Mamba2AdamW.step
  closes the loop in alpha_dqn_h600_smoke and alpha_compose_backtest.

Smoke (--temporal --c51): all 4 KCs PASS. R_mean -6.3 → +4.2 vs
Phase E.3 close R_mean -4.7 (no-temporal). EARLY_Q_MOVEMENT
calibration (mamba2_snapshot + mamba2_weight_distance) lifts the
diagnostic from 0.0023 (head-only) to 0.0590 (head + encoder),
giving an honest learning signal when the encoder absorbs gradient.

Backtest (--c51 --temporal --window-k 16 --isv-continual):
  cost=0.0000  best τ=0.250  Sharpe_ann=+34.56  (was +10.41 head-only,
                                                  -22.54 frozen-Mamba2)
  cost=0.0625  best τ=0.250  Sharpe_ann=+33.22
  cost=0.1250  best τ=0.250  Sharpe_ann=+30.85  (Phase 1d.4 baseline: -4.0)
  cost=0.2500  best τ=0.250  Sharpe_ann=+27.73
  cost=0.5000  best τ=0.250  Sharpe_ann=+15.68

Caveat: in-sample results; OOS gate next.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-15 23:49:25 +02:00
parent 0ab54dc8ef
commit 3dc022b843
6 changed files with 572 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
crates/ml/examples/alpha_compose_backtest.rs
View File

@@ -42,7 +42,10 @@ use cudarc::driver::{CudaContext, DevicePtr, DevicePtrMut};
use tracing::info;
// Phase E.4.A.8/T12: Mamba2 temporal encoder (ml-alpha Phase 1d.1).
use ml_alpha::mamba2_block::{Mamba2Block, Mamba2BlockConfig};
// Phase E.4.A.T10: Mamba2AdamW for training Mamba2 weights.
use ml_alpha::mamba2_block::{
Mamba2Block, Mamba2BlockConfig, Mamba2AdamW, Mamba2AdamWConfig,
};
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
// ── Mapped-pinned helpers (mirror gpu_training_guard.rs MappedBuffer) ──
@@ -521,7 +524,7 @@ fn main() -> Result<()> {
let mut h_enriched_buf_dev = stream
.alloc_zeros::<f32>(h_enriched_buf_capacity)
.context("alloc h_enriched buffer")?;
let mamba2_block: Option<Mamba2Block> = if cli.temporal {
let mut mamba2_block: Option<Mamba2Block> = if cli.temporal {
let cfg = Mamba2BlockConfig {
in_dim: STATE_DIM,
hidden_dim: cli.mamba2_hidden_dim,
@@ -535,6 +538,22 @@ fn main() -> Result<()> {
} else {
None
};
// Phase E.4.A.T10: AdamW for Mamba2 weights (only when --temporal).
let mut mamba2_adamw: Option<Mamba2AdamW> = if let Some(block) = mamba2_block.as_ref() {
Some(Mamba2AdamW::new(block, Mamba2AdamWConfig::default())
.map_err(|e| anyhow::anyhow!("Mamba2AdamW init: {e}"))?)
} else {
None
};
// Load the C51 grad-input kernel (for backward chain into Mamba2).
let c51_grad_input_kernel = c51_module
.load_function("alpha_c51_grad_input_kernel")
.context("c51 grad_input load")?;
// T10: train-time window store (captures windows during inference for
// the end-of-epoch batched Mamba2 forward).
let train_window_store_kernel = push_module
.load_function("alpha_train_window_store_batched_kernel")
.context("train_window_store kernel load")?;
let n_atoms_i = c51_n_atoms as i32;
// ISV buffer + Wiener state for the eval-time controller path.
let mut isv_host: Vec<f32> = vec![0.0; 552];
@@ -569,6 +588,13 @@ fn main() -> Result<()> {
let mut batched_window_tensor_train = GpuTensor::zeros(
&[train_n_par, cli.window_k, STATE_DIM], &stream
).map_err(|e| anyhow::anyhow!("alloc batched window train: {e}"))?;
// T10 buffers: all-step train windows and d_h_enriched for the backward chain.
let mut train_windows_tensor_train = GpuTensor::zeros(
&[train_batch, cli.window_k, STATE_DIM], &stream
).map_err(|e| anyhow::anyhow!("alloc train windows train: {e}"))?;
let mut d_h_enriched_tensor_train = GpuTensor::zeros(
&[train_batch, cli.mamba2_hidden_dim], &stream
).map_err(|e| anyhow::anyhow!("alloc d_h_enriched train: {e}"))?;
let mut batched_probs_inference_train = stream
.alloc_zeros::<f32>(train_n_par * N_ACTIONS * c51_n_atoms)?;
let snapshots_arc_train = env.snapshots_arc();
@@ -597,6 +623,8 @@ fn main() -> Result<()> {
let mut done_flags = vec![false; train_n_par];
if cli.temporal {
stream.memset_zeros(batched_window_tensor_train.data_mut())?;
stream.memset_zeros(train_windows_tensor_train.data_mut())?;
stream.memset_zeros(d_h_enriched_tensor_train.data_mut())?;
}
stream.memset_zeros(&mut h_enriched_train_dev)?;
let mut states_host_train: Vec<f32> = vec![0.0; train_batch * STATE_DIM];
@@ -626,6 +654,19 @@ fn main() -> Result<()> {
)?;
}
}
// T10: capture post-push windows into [H*N_par, K, in_dim] buffer
// at step row offset `step * N_par`.
{
let (src_ptr, _g_src) = batched_window_tensor_train.data().device_ptr(&stream);
let (dst_ptr, _g_dst) = train_windows_tensor_train.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_train_window_store_batched(
&stream, &train_window_store_kernel,
src_ptr, dst_ptr, (step * train_n_par) as i32,
train_n_par as i32, cli.window_k as i32, state_dim_i,
)?;
}
}
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let (_logit, cache) = block.forward_train(&batched_window_tensor_train)
.map_err(|e| anyhow::anyhow!("train mamba2: {e}"))?;
@@ -708,15 +749,28 @@ fn main() -> Result<()> {
stream.memcpy_htod(&dones_host_train, &mut dones_train_dev)?;
let bt = train_batch as i32;
// T10: re-forward Mamba2 on collected train windows to get a cache
// that backward_from_h_enriched can consume. cache.h_enriched
// replaces h_enriched_train_dev curr offset (bit-identical since
// Mamba2 weights haven't been updated yet this epoch).
let cache_train: Option<ml_alpha::mamba2_block::Mamba2ForwardCache> = if cli.temporal {
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let (_logit, cache) = block.forward_train(&train_windows_tensor_train)
.map_err(|e| anyhow::anyhow!("train mamba2 (epoch forward): {e}"))?;
Some(cache)
} else {
None
};
let (curr_input_ptr_guard, next_input_ptr_guard);
let (curr_input_ptr, next_input_ptr): (u64, u64) = if cli.temporal {
let (h_base, g) = h_enriched_train_dev.device_ptr(&stream);
let cache = cache_train.as_ref().expect("cache_train missing");
let (h_curr, g) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (h_base_next, g_next) = h_enriched_train_dev.device_ptr(&stream);
curr_input_ptr_guard = g;
next_input_ptr_guard = g_next;
let _ = (&curr_input_ptr_guard, &next_input_ptr_guard);
(
h_base,
h_curr,
h_base_next + (train_n_par as u64) * (cli.mamba2_hidden_dim as u64) * 4u64,
)
} else {
@@ -783,6 +837,35 @@ fn main() -> Result<()> {
)?;
}
}
// T10 backward chain: C51 grad_input → Mamba2 backward → AdamW step.
if cli.temporal {
{
let (p_ptr, _g0) = probs_curr_train_dev.device_ptr(&stream);
let (m_ptr, _g1) = m_train_dev.device_ptr(&stream);
let (a_ptr, _g2) = actions_train_dev.device_ptr(&stream);
let (w_ptr, _g3) = w_dev.device_ptr(&stream);
let (dh_ptr, _g4) = d_h_enriched_tensor_train.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad_input(
&stream, &c51_grad_input_kernel,
p_ptr, m_ptr, a_ptr, w_ptr, dh_ptr,
bt, c51_input_dim as i32, n_act_i, n_atoms_i,
1.0 / bt as f32,
)?;
}
}
let cache = cache_train.as_ref().expect("cache_train missing");
let block_ref = mamba2_block.as_ref().expect("Mamba2Block missing");
let grads = block_ref
.backward_from_h_enriched(cache, &d_h_enriched_tensor_train)
.map_err(|e| anyhow::anyhow!("Mamba2 backward: {e}"))?;
let _ = cache_train;
let block_mut = mamba2_block.as_mut().expect("Mamba2Block missing");
let adamw = mamba2_adamw.as_mut().expect("Mamba2AdamW missing");
adamw
.step(block_mut, &grads)
.map_err(|e| anyhow::anyhow!("Mamba2AdamW step: {e}"))?;
}
{
let (dw_ptr, _g0) = dw_dev.device_ptr_mut(&stream);
unsafe {

209
crates/ml/examples/alpha_dqn_h600_smoke.rs
View File

@@ -49,7 +49,10 @@ use cudarc::driver::{CudaContext, DevicePtr, DevicePtrMut};
use tracing::{info, warn};
// Phase E.4.A.8: Mamba2 temporal encoder (ml-alpha Phase 1d.1).
use ml_alpha::mamba2_block::{Mamba2Block, Mamba2BlockConfig};
// Phase E.4.A.T10: Mamba2AdamW for training Mamba2 weights.
use ml_alpha::mamba2_block::{
Mamba2Block, Mamba2BlockConfig, Mamba2AdamW, Mamba2AdamWConfig,
};
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
// ── Mapped-pinned helpers ──────────────────────────────────────────
@@ -481,6 +484,43 @@ fn weight_distance_from_init(w_now: &[f32], b_now: &[f32], w_init: &[f32], b_ini
s.sqrt()
}
/// Snapshot of all trainable Mamba2 parameters at init time. Concatenated
/// flat Vec<f32> in deterministic order so a later snapshot subtracts
/// element-wise. Used by the EARLY_Q_MOVEMENT_EMA diagnostic so the
/// kill criterion captures encoder movement (not just C51 head), which
/// would otherwise under-report system learning when --temporal is on.
fn mamba2_snapshot(
stream: &std::sync::Arc<cudarc::driver::CudaStream>,
block: &Mamba2Block,
) -> anyhow::Result<Vec<f32>> {
let mut out = Vec::new();
out.extend(stream.clone_dtoh(&block.w_in.weight).context("dtoh w_in")?);
out.extend(stream.clone_dtoh(&block.w_in.bias).context("dtoh w_in.b")?);
out.extend(stream.clone_dtoh(&block.w_a.weight).context("dtoh w_a")?);
out.extend(stream.clone_dtoh(&block.w_a.bias).context("dtoh w_a.b")?);
out.extend(stream.clone_dtoh(&block.w_b.weight).context("dtoh w_b")?);
out.extend(stream.clone_dtoh(&block.w_b.bias).context("dtoh w_b.b")?);
out.extend(stream.clone_dtoh(&block.w_c).context("dtoh w_c")?);
out.extend(stream.clone_dtoh(&block.w_out.weight).context("dtoh w_out")?);
out.extend(stream.clone_dtoh(&block.w_out.bias).context("dtoh w_out.b")?);
Ok(out)
}
/// L2 distance ||M_now − M_init||_F across the same concatenated layout
/// produced by `mamba2_snapshot`. Returns 0.0 if init has different
/// length (shouldn't happen unless model topology changed).
fn mamba2_weight_distance(now: &[f32], init: &[f32]) -> f32 {
if now.len() != init.len() {
return 0.0;
}
let mut s: f32 = 0.0;
for (a, b) in now.iter().zip(init.iter()) {
let d = a - b;
s += d * d;
}
s.sqrt()
}
/// SplitMix64 RNG — same as ExecutionEnv's ReplayRng so seeds compose
/// cleanly when the smoke is reproduced.
struct SmokeRng {
@@ -855,7 +895,7 @@ fn main() -> Result<()> {
// --temporal is set. Hidden_dim = c51_input_dim (so the encoder's
// output flows directly into the C51 head without an additional
// projection). State_dim ≤ 16 (kernel max).
let mamba2_block: Option<Mamba2Block> = if cli.temporal {
let mut mamba2_block: Option<Mamba2Block> = if cli.temporal {
let cfg = Mamba2BlockConfig {
in_dim: STATE_DIM,
hidden_dim: cli.mamba2_hidden_dim,
@@ -871,6 +911,40 @@ fn main() -> Result<()> {
} else {
None
};
// Phase E.4.A.T10: AdamW for Mamba2 weights (active when --temporal).
let mut mamba2_adamw: Option<Mamba2AdamW> = if let Some(block) = mamba2_block.as_ref() {
Some(Mamba2AdamW::new(block, Mamba2AdamWConfig::default())
.map_err(|e| anyhow::anyhow!("Mamba2AdamW init: {e}"))?)
} else {
None
};
// Snapshot Mamba2 weights at init so EARLY_Q_MOVEMENT can include
// encoder movement (not just C51 head). Without this, the diagnostic
// under-reports system learning when --temporal: the encoder absorbs
// the gradient signal, the head moves little, and the KC reads ~0.
let mamba2_init_snap: Option<Vec<f32>> = if let Some(block) = mamba2_block.as_ref() {
Some(mamba2_snapshot(&stream, block)?)
} else {
None
};
// d_h_enriched as GpuTensor so it can feed `backward_from_h_enriched`
// after the C51 grad-input kernel writes per-row gradients into it.
let mut d_h_enriched_tensor_smoke = GpuTensor::zeros(
&[cli.horizon, cli.mamba2_hidden_dim], &stream
).map_err(|e| anyhow::anyhow!("alloc d_h_enriched smoke: {e}"))?;
// Per-step train-time windows, written by the new
// `alpha_train_window_store_batched_kernel` during inference and
// re-forwarded as one big batched Mamba2 call at end-of-episode.
let mut train_windows_tensor_smoke = GpuTensor::zeros(
&[cli.horizon, cli.window_k, STATE_DIM], &stream
).map_err(|e| anyhow::anyhow!("alloc train windows smoke: {e}"))?;
// Load the C51 grad-input kernel and the train-window store kernel.
let c51_grad_input_kernel = c51_module
.load_function("alpha_c51_grad_input_kernel")
.context("c51 grad_input load")?;
let train_window_store_kernel = push_module
.load_function("alpha_train_window_store_batched_kernel")
.context("train_window_store kernel load")?;
// Kill-criteria inputs: action_counts (i32 × n_actions),
// scalar_inputs (f32 × 3: rollout_R_mean, q_init_norm, q_early_norm).
@@ -901,10 +975,19 @@ fn main() -> Result<()> {
// Phase E.4.A.7: clear the temporal window on episode start so
// Mamba2 sees zero-history for the first window_k-1 steps.
// Phase E.4.A.T10: also zero the train-time window buffer and the
// d_h_enriched buffer so terminal slots (past ep_len) contribute
// zero gradient through the Mamba2 backward.
if cli.temporal {
stream
.memset_zeros(window_tensor.data_mut())
.context("zero window on episode reset")?;
stream
.memset_zeros(train_windows_tensor_smoke.data_mut())
.context("zero train windows on episode reset")?;
stream
.memset_zeros(d_h_enriched_tensor_smoke.data_mut())
.context("zero d_h_enriched on episode reset")?;
}
let mut states_host: Vec<f32> = Vec::with_capacity(cli.horizon * STATE_DIM);
@@ -923,13 +1006,27 @@ fn main() -> Result<()> {
state_pinned.write(&s_vec);
// Phase E.4.A.7: shift-and-insert push of current state.
if cli.temporal {
let (w_ptr, _g1) = window_tensor.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push(
&stream, &push_kernel,
state_pinned.dev_u64(), w_ptr,
cli.window_k as i32, state_dim_i,
)?;
{
let (w_ptr, _g1) = window_tensor.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push(
&stream, &push_kernel,
state_pinned.dev_u64(), w_ptr,
cli.window_k as i32, state_dim_i,
)?;
}
}
// T10: capture post-push window for end-of-episode re-forward.
{
let (src_ptr, _g_src) = window_tensor.data().device_ptr(&stream);
let (dst_ptr, _g_dst) = train_windows_tensor_smoke.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_train_window_store_batched(
&stream, &train_window_store_kernel,
src_ptr, dst_ptr, state.step as i32,
1, cli.window_k as i32, state_dim_i,
)?;
}
}
}
// Phase E.4.A.8: Mamba2 forward over the window → h_enriched
@@ -1086,24 +1183,44 @@ fn main() -> Result<()> {
// Run ONE extra Mamba2 forward on the terminal window to
// fill slot ep_len (the "next-state" representation for
// the last training transition).
if cli.temporal {
//
// Phase E.4.A.T10: also run a second Mamba2 forward on the
// collected train_windows_tensor_smoke to get a cache that
// covers all ep_len visited windows. We use the cache's
// h_enriched as the curr_input for C51 (so the C51 grad can
// flow back into Mamba2 via backward_from_h_enriched), while
// keeping h_enriched_buf_dev offset for next_input (frozen
// target side, no backward).
let cache_train: Option<ml_alpha::mamba2_block::Mamba2ForwardCache> = if cli.temporal {
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let (_logit, cache) = block.forward_train(&window_tensor)
.map_err(|e| anyhow::anyhow!("mamba2 terminal forward: {e}"))?;
let term_offset = ((ep_len as usize) * cli.mamba2_hidden_dim) as i32;
let (src_ptr, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (buf_ptr, _g_buf) = h_enriched_buf_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store(
&stream, &h_store_kernel,
src_ptr, buf_ptr, term_offset,
cli.mamba2_hidden_dim as i32,
)?;
// Terminal forward fills h_enriched_buf_dev[ep_len].
{
let (_logit, cache) = block.forward_train(&window_tensor)
.map_err(|e| anyhow::anyhow!("mamba2 terminal forward: {e}"))?;
let term_offset = ((ep_len as usize) * cli.mamba2_hidden_dim) as i32;
let (src_ptr, _g_src) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (buf_ptr, _g_buf) = h_enriched_buf_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_h_enriched_store(
&stream, &h_store_kernel,
src_ptr, buf_ptr, term_offset,
cli.mamba2_hidden_dim as i32,
)?;
}
}
}
// Training-time forward on the full [horizon, K, in_dim]
// batch (slots ep_len..horizon are zero, contribute zero
// gradient since d_h_enriched stays zero there).
let (_logit, cache) = block.forward_train(&train_windows_tensor_smoke)
.map_err(|e| anyhow::anyhow!("mamba2 train forward: {e}"))?;
Some(cache)
} else {
None
};
let (curr_input_ptr_guard, next_input_ptr_guard);
let (curr_input_ptr, next_input_ptr): (u64, u64) = if cli.temporal {
let (h_ptr_curr, g_curr) = h_enriched_buf_dev.device_ptr(&stream);
let cache = cache_train.as_ref().expect("cache_train missing");
let (h_ptr_curr, g_curr) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (h_ptr_next_base, g_next) = h_enriched_buf_dev.device_ptr(&stream);
curr_input_ptr_guard = g_curr;
next_input_ptr_guard = g_next;
@@ -1183,6 +1300,39 @@ fn main() -> Result<()> {
)?;
}
}
// T10 backward chain: C51 grad_input → Mamba2 backward → AdamW.
// The C51 grad_input kernel writes dL/d_h_enriched into the
// first ep_len rows of d_h_enriched_tensor_smoke; remaining
// rows stay zero (zeroed on episode reset).
if cli.temporal {
{
let (p_ptr, _g0) = probs_current_dev.device_ptr(&stream);
let (m_ptr, _g1) = m_dev.device_ptr(&stream);
let (a_ptr, _g2) = actions_dev.device_ptr(&stream);
let (w_ptr, _g3) = w_dev.device_ptr(&stream);
let (dh_ptr, _g4) = d_h_enriched_tensor_smoke.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_grad_input(
&stream, &c51_grad_input_kernel,
p_ptr, m_ptr, a_ptr, w_ptr,
dh_ptr,
ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
1.0 / ep_len as f32,
)?;
}
}
let cache = cache_train.as_ref().expect("cache_train missing");
let block_ref = mamba2_block.as_ref().expect("Mamba2Block missing");
let grads = block_ref
.backward_from_h_enriched(cache, &d_h_enriched_tensor_smoke)
.map_err(|e| anyhow::anyhow!("Mamba2 backward: {e}"))?;
let _ = cache_train;
let block_mut = mamba2_block.as_mut().expect("Mamba2Block missing");
let adamw = mamba2_adamw.as_mut().expect("Mamba2AdamW missing");
adamw
.step(block_mut, &grads)
.map_err(|e| anyhow::anyhow!("Mamba2AdamW step: {e}"))?;
}
} else {
// ── Linear-Q batched compute (Munchausen target) ────────
// Forward Q_current on states
@@ -1361,7 +1511,18 @@ fn main() -> Result<()> {
// q_early = q_init + ||W_now − W_init||_F so the kernel's
// |q_early − q_init| / |q_init| ratio yields the relative
// weight-space distance from init (direction-sensitive).
let dist = weight_distance_from_init(&w_now, &b_now, &w_init, &b_init);
//
// When --temporal is on, include Mamba2 weight movement too:
// the encoder absorbs most of the gradient signal so the C51
// head moves little on its own. Without this addition, the
// diagnostic systematically under-reports system learning.
let mut dist = weight_distance_from_init(&w_now, &b_now, &w_init, &b_init);
if let (Some(block), Some(init_snap)) =
(mamba2_block.as_ref(), mamba2_init_snap.as_ref())
{
let now_snap = mamba2_snapshot(&stream, block)?;
dist += mamba2_weight_distance(&now_snap, init_snap);
}
let q_early_norm = q_init_norm + dist;
let rollout_r_mean: f32 = if recent_returns.is_empty() {
0.0

41
crates/ml/src/cuda_pipeline/alpha_kernels.rs
View File

@@ -684,6 +684,47 @@ pub unsafe fn launch_alpha_h_enriched_store_batched(
Ok(())
}
/// Launch `alpha_train_window_store_batched_kernel`. Copies a batch of
/// `B` windows of shape `[K, state_dim]` into a step-aligned slice of a
/// destination buffer of shape `[H*B, K, state_dim]`. Lets the training
/// loop re-run a single batched Mamba2 forward over all visited windows
/// (T10 backward chain).
///
/// # Safety
/// `src_dev` and `buf_dev` must be valid device pointers.
pub unsafe fn launch_alpha_train_window_store_batched(
stream: &cudarc::driver::CudaStream,
kernel: &cudarc::driver::CudaFunction,
src_dev: u64,
buf_dev: u64,
step_row_offset: i32,
b: i32,
k: i32,
state_dim: i32,
) -> Result<(), MLError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
debug_assert!(b > 0 && k > 0 && state_dim > 0);
debug_assert!(step_row_offset >= 0);
const BLOCK: u32 = 32;
let grid_x = ((state_dim as u32) + BLOCK - 1) / BLOCK;
let cfg = LaunchConfig {
grid_dim: (grid_x.max(1), k as u32, b as u32),
block_dim: (BLOCK, 1, 1),
shared_mem_bytes: 0,
};
stream
.launch_builder(kernel)
.arg(&src_dev)
.arg(&buf_dev)
.arg(&step_row_offset)
.arg(&b)
.arg(&k)
.arg(&state_dim)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("alpha_train_window_store_batched launch: {e}")))?;
Ok(())
}
/// Launch `alpha_window_push_batched_kernel`. Same chronological
/// shift+insert as `alpha_window_push_kernel` but across B parallel
/// windows. One thread per (batch, feature). Used by the batched-eval

25
crates/ml/src/cuda_pipeline/alpha_window_push.cu
View File

@@ -94,3 +94,28 @@ extern "C" __global__ void alpha_h_enriched_store_batched_kernel(
const int src_idx = b * hidden_dim + j;
buf[dst_idx] = src[src_idx];
}
// ----------------------------------------------------------------------
// Phase E.4.A.T10 (2026-05-15): batched train-window store. Copies the
// per-step windows tensor [B, K, state_dim] into an all-steps buffer
// [H*B, K, state_dim] at row offset step*B. Used so the training loop
// can re-run a single batched Mamba2 forward over all visited windows
// and recover a cache for backward.
// One thread per (env, slot, feature). Grid: (sd_blocks, K, B).
// ----------------------------------------------------------------------
extern "C" __global__ void alpha_train_window_store_batched_kernel(
const float* __restrict__ src, // [B, K, state_dim]
float* __restrict__ buf, // [H*B, K, state_dim]
int step_row_offset, // step * B (row count of width K*state_dim)
int B,
int K,
int state_dim
) {
const int j = blockIdx.x * blockDim.x + threadIdx.x;
const int k = blockIdx.y;
const int b = blockIdx.z;
if (b >= B || k >= K || j >= state_dim) return;
const int src_idx = (b * K + k) * state_dim + j;
const int dst_idx = ((step_row_offset + b) * K + k) * state_dim + j;
buf[dst_idx] = src[src_idx];
}