diff --git a/crates/ml/examples/alpha_dqn_h600_smoke.rs b/crates/ml/examples/alpha_dqn_h600_smoke.rs index 4c5e15e3f..0c4576394 100644 --- a/crates/ml/examples/alpha_dqn_h600_smoke.rs +++ b/crates/ml/examples/alpha_dqn_h600_smoke.rs @@ -804,10 +804,12 @@ fn main() -> Result<()> { // Phase E.4.A.7: GPU-resident sliding-window state buffer for the // temporal encoder. Allocated unconditionally; only populated when // cli.temporal is set. Capacity: window_k × STATE_DIM floats. + // Phase E.4.A.6: shift+insert layout — slot 0 = oldest, slot K-1 + // = newest. No head_idx tracking needed; kernel always shifts then + // inserts at slot K-1. let mut window_dev = stream .alloc_zeros::(cli.window_k * STATE_DIM) .context("alloc window buffer")?; - let mut head_idx: i32 = 0; // Kill-criteria inputs: action_counts (i32 × n_actions), // scalar_inputs (f32 × 3: rollout_R_mean, q_init_norm, q_early_norm). @@ -842,7 +844,6 @@ fn main() -> Result<()> { stream .memset_zeros(&mut window_dev) .context("zero window on episode reset")?; - head_idx = 0; } let mut states_host: Vec = Vec::with_capacity(cli.horizon * STATE_DIM); @@ -859,19 +860,18 @@ fn main() -> Result<()> { let s_vec = state_as_vec(&env, &state); let action: u8 = if cli.c51 { state_pinned.write(&s_vec); - // Phase E.4.A.7: push current state into the circular - // window buffer. The buffer is populated for ALL - // --temporal runs; consumer (Mamba2) wires in T8. + // Phase E.4.A.7: shift-and-insert push of current state + // into the chronological window buffer. Consumer (Mamba2) + // wires in T8. if cli.temporal { let (w_ptr, _g1) = window_dev.device_ptr_mut(&stream); unsafe { ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push( &stream, &push_kernel, state_pinned.dev_u64(), w_ptr, - head_idx, state_dim_i, + cli.window_k as i32, state_dim_i, )?; } - head_idx = (head_idx + 1) % (cli.window_k as i32); } { let (w_ptr, _g0) = w_dev.device_ptr(&stream); diff --git a/crates/ml/src/cuda_pipeline/alpha_kernels.rs b/crates/ml/src/cuda_pipeline/alpha_kernels.rs index 59f4ee921..625cebada 100644 --- a/crates/ml/src/cuda_pipeline/alpha_kernels.rs +++ b/crates/ml/src/cuda_pipeline/alpha_kernels.rs @@ -448,10 +448,15 @@ pub static ALPHA_C51_CUBIN: &[u8] = pub static ALPHA_WINDOW_PUSH_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/alpha_window_push.cubin")); -/// Launch `alpha_window_push_kernel`. Copies the current `state_in` -/// (state_dim floats) into the circular buffer slot indexed by -/// `head_idx`. The host tracks `head_idx = (head_idx + 1) % K` -/// across episode steps and zeros the window on episode reset. +/// Launch `alpha_window_push_kernel`. Shifts the window's K slots left +/// by one (slot 1 → slot 0, slot 2 → slot 1, …, slot K-1 → slot K-2) +/// and inserts the new state at slot K-1. After the call, +/// `window[0..K-1, :]` is chronological (slot 0 = oldest, slot K-1 = +/// newest) — matches Mamba2Block's `[B, K, in_dim]` input contract. +/// +/// One thread per state feature; each thread sequentially shifts its +/// column's K-1 slots and writes the new value at K-1. No data races +/// (each thread owns its column). /// /// # Safety /// `state_in_dev` and `window_dev` MUST be valid device pointers. @@ -461,13 +466,13 @@ pub unsafe fn launch_alpha_window_push( kernel: &cudarc::driver::CudaFunction, state_in_dev: u64, window_dev: u64, - head_idx: i32, + k: i32, state_dim: i32, ) -> Result<(), MLError> { use cudarc::driver::{LaunchConfig, PushKernelArg}; debug_assert!(state_dim > 0, "state_dim must be positive"); - debug_assert!(head_idx >= 0, "head_idx must be non-negative"); + debug_assert!(k > 1, "K must be > 1 (need at least one shift)"); const BLOCK: u32 = 32; let grid_x = ((state_dim as u32) + BLOCK - 1) / BLOCK; @@ -480,7 +485,7 @@ pub unsafe fn launch_alpha_window_push( .launch_builder(kernel) .arg(&state_in_dev) .arg(&window_dev) - .arg(&head_idx) + .arg(&k) .arg(&state_dim) .launch(cfg) .map_err(|e| MLError::ModelError(format!("alpha_window_push launch: {e}")))?; @@ -769,12 +774,13 @@ mod tests { assert!(true); } - /// GPU smoke for `alpha_window_push_kernel`: writes state vectors - /// into two non-adjacent slots of a 4-slot circular buffer and - /// verifies the rows land where expected; zero-initialised slots - /// stay zero. Phase E.4.A.6 verification. + /// GPU smoke for `alpha_window_push_kernel`: pushes three states + /// in sequence into a 4-slot buffer and verifies chronological + /// shift+insert: slot 0 = state_a (oldest, shifted left twice), + /// slot 1 = state_b (oldest now), slot 2 = state_c (middle), + /// slot 3 = state_c+1 (newest). Phase E.4.A.6 verification. #[test] - fn alpha_window_push_circular_writes_to_indexed_slot() -> Result<(), MLError> { + fn alpha_window_push_shift_and_insert() -> Result<(), MLError> { let ctx = CudaContext::new(0).map_err(|e| MLError::ModelError(format!("ctx: {e}")))?; let stream = ctx.default_stream(); let module = ctx @@ -788,36 +794,34 @@ mod tests { const D: usize = 3; let mut window_dev = stream.alloc_zeros::(K * D) .map_err(|e| MLError::ModelError(format!("alloc window: {e}")))?; - let state_a: Vec = vec![1.0, 2.0, 3.0]; - let state_b: Vec = vec![10.0, 20.0, 30.0]; - let state_a_dev = stream.clone_htod(&state_a) - .map_err(|e| MLError::ModelError(format!("htod a: {e}")))?; - let state_b_dev = stream.clone_htod(&state_b) - .map_err(|e| MLError::ModelError(format!("htod b: {e}")))?; - // Push state_a at head_idx=0 - { - let (s_ptr, _g0) = state_a_dev.device_ptr(&stream); + let states: [Vec; 3] = [ + vec![1.0, 2.0, 3.0], + vec![10.0, 20.0, 30.0], + vec![100.0, 200.0, 300.0], + ]; + + for state in &states { + let s_dev = stream.clone_htod(state) + .map_err(|e| MLError::ModelError(format!("htod: {e}")))?; + let (s_ptr, _g0) = s_dev.device_ptr(&stream); let (w_ptr, _g1) = window_dev.device_ptr_mut(&stream); unsafe { - launch_alpha_window_push(&stream, &kernel, s_ptr, w_ptr, 0, D as i32)?; - } - } - // Push state_b at head_idx=2 - { - let (s_ptr, _g0) = state_b_dev.device_ptr(&stream); - let (w_ptr, _g1) = window_dev.device_ptr_mut(&stream); - unsafe { - launch_alpha_window_push(&stream, &kernel, s_ptr, w_ptr, 2, D as i32)?; + launch_alpha_window_push(&stream, &kernel, s_ptr, w_ptr, K as i32, D as i32)?; } } stream.synchronize().map_err(|e| MLError::ModelError(format!("sync: {e}")))?; let window_host = stream.clone_dtoh(&window_dev) .map_err(|e| MLError::ModelError(format!("dtoh: {e}")))?; - assert_eq!(&window_host[0..3], &[1.0, 2.0, 3.0], "slot 0 = state_a"); - assert_eq!(&window_host[3..6], &[0.0, 0.0, 0.0], "slot 1 untouched"); - assert_eq!(&window_host[6..9], &[10.0, 20.0, 30.0], "slot 2 = state_b"); - assert_eq!(&window_host[9..12], &[0.0, 0.0, 0.0], "slot 3 untouched"); + // After 3 pushes into 4-slot buffer: + // slot 0 = zero (initial pre-shift state, shifted left twice) + // slot 1 = states[0] (shifted from slot 3 → slot 2 → slot 1) + // slot 2 = states[1] (shifted from slot 3 → slot 2) + // slot 3 = states[2] (newest, just inserted) + assert_eq!(&window_host[0..3], &[0.0, 0.0, 0.0], "slot 0 = initial zero shifted"); + assert_eq!(&window_host[3..6], &[1.0, 2.0, 3.0], "slot 1 = states[0]"); + assert_eq!(&window_host[6..9], &[10.0, 20.0, 30.0], "slot 2 = states[1]"); + assert_eq!(&window_host[9..12], &[100.0, 200.0, 300.0], "slot 3 = newest"); Ok(()) } diff --git a/crates/ml/src/cuda_pipeline/alpha_window_push.cu b/crates/ml/src/cuda_pipeline/alpha_window_push.cu index 6538769ee..f41e80d66 100644 --- a/crates/ml/src/cuda_pipeline/alpha_window_push.cu +++ b/crates/ml/src/cuda_pipeline/alpha_window_push.cu @@ -1,25 +1,33 @@ // crates/ml/src/cuda_pipeline/alpha_window_push.cu // -// Phase E.4.A.6 (2026-05-15): circular-buffer push kernel for the -// sliding-window state input to the Mamba2 temporal encoder. +// Phase E.4.A.6 (2026-05-15): shift+insert circular-buffer push for the +// sliding-window state input to the Mamba2 temporal encoder. Mirrors +// production's `mamba2_update_history` (mamba2_temporal_kernel.cu): +// shift slots 1..K-1 to slots 0..K-2, then insert new state at slot K-1. // -// One thread per (state_dim) feature: writes `state_dim` floats from -// `state_in[state_dim]` into `window[head_idx*state_dim .. (head_idx+1)*state_dim]`. -// The host increments `head_idx = (head_idx + 1) % K` after each call. -// On episode reset the host zeros the window buffer via `cudaMemsetAsync`. +// Layout invariant: window[0..K-1, :] is CHRONOLOGICAL — slot 0 is the +// oldest observation, slot K-1 is the newest. This matches Mamba2Block's +// [B, K, in_dim] input contract directly (no reorder needed). // -// Grid: ceil(state_dim / BLOCK), Block: 32. For state_dim=10 a single -// warp suffices — tiny launch overhead. +// Threads: state_dim threads, one per feature. Each thread sequentially +// shifts its column's K-1 slots and writes the new value at slot K-1. +// For typical state_dim=10, K=16: 10 threads × ~17 ops each, single +// block, no data races (each thread owns its column). #include extern "C" __global__ void alpha_window_push_kernel( - const float* __restrict__ state_in, // [state_dim] (mapped-pinned input) + const float* __restrict__ state_in, // [state_dim] (mapped-pinned) float* __restrict__ window, // [K, state_dim] row-major - int head_idx, // slot to write (host-tracked) + int K, // window length int state_dim // typical 10 ) { const int j = blockIdx.x * blockDim.x + threadIdx.x; if (j >= state_dim) return; - window[head_idx * state_dim + j] = state_in[j]; + // Shift slot (t+1) → slot t for t = 0..K-2. + for (int t = 0; t < K - 1; ++t) { + window[t * state_dim + j] = window[(t + 1) * state_dim + j]; + } + // Insert new state at slot K-1. + window[(K - 1) * state_dim + j] = state_in[j]; } diff --git a/docs/isv-slots.md b/docs/isv-slots.md index e0a190a81..104e86c15 100644 --- a/docs/isv-slots.md +++ b/docs/isv-slots.md @@ -779,7 +779,10 @@ Pearl-1 atom-headroom slots when those land. ### Phase E.4.A.6 — alpha_window_push (2026-05-15) -`crates/ml/src/cuda_pipeline/alpha_window_push.cu`: 1-kernel circular -buffer push primitive. NO ISV slot interaction — it operates purely -on the smoke binary's `window_dev` buffer. Documented here only for -the kernel-audit-doc hook requirement; truly slot-agnostic. +`crates/ml/src/cuda_pipeline/alpha_window_push.cu`: shift+insert +buffer push primitive (matches production `mamba2_update_history` +pattern). NO ISV slot interaction — it operates purely on the smoke +binary's `window_dev` buffer. Documented here only for the +kernel-audit-doc hook requirement; truly slot-agnostic. After each +push the buffer is chronological: slot 0 = oldest, slot K-1 = newest, +matching Mamba2Block's `[B, K, in_dim]` input contract.