diff --git a/crates/ml/build.rs b/crates/ml/build.rs index b97df3561..f47ae6082 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -1621,6 +1621,10 @@ fn main() { // atom support, single network, no per-branch / Adam machinery — // testing the calibration hypothesis from `pearl_action_pruning_falsified`. "alpha_c51.cu", + // Phase E.4.A.6 (2026-05-15): circular-buffer push kernel for the + // sliding-window state input to the Mamba2 temporal encoder + // (Mamba2Block from ml-alpha). + "alpha_window_push.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/alpha_kernels.rs b/crates/ml/src/cuda_pipeline/alpha_kernels.rs index 4ac323e1b..59f4ee921 100644 --- a/crates/ml/src/cuda_pipeline/alpha_kernels.rs +++ b/crates/ml/src/cuda_pipeline/alpha_kernels.rs @@ -443,6 +443,50 @@ pub static STACKER_THRESHOLD_CONTROLLER_CUBIN: &[u8] = pub static ALPHA_C51_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/alpha_c51.cubin")); +/// Precompiled circular-buffer push kernel for the sliding-window state +/// input to Mamba2 temporal encoder. Phase E.4.A.6. +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. +/// +/// # Safety +/// `state_in_dev` and `window_dev` MUST be valid device pointers. +/// `window_dev` capacity ≥ `K * state_dim` floats. +pub unsafe fn launch_alpha_window_push( + stream: &cudarc::driver::CudaStream, + kernel: &cudarc::driver::CudaFunction, + state_in_dev: u64, + window_dev: u64, + head_idx: 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"); + + const BLOCK: u32 = 32; + let grid_x = ((state_dim as u32) + BLOCK - 1) / BLOCK; + let cfg = LaunchConfig { + grid_dim: (grid_x.max(1), 1, 1), + block_dim: (BLOCK, 1, 1), + shared_mem_bytes: 0, + }; + stream + .launch_builder(kernel) + .arg(&state_in_dev) + .arg(&window_dev) + .arg(&head_idx) + .arg(&state_dim) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("alpha_window_push launch: {e}")))?; + Ok(()) +} + /// Launch `alpha_c51_forward_kernel`. One thread per (sample, action); /// each thread sequentially computes n_atoms logits + softmax. Per-thread /// scratch lives in local memory (MAX_N_ATOMS = 64 floats hardcoded in @@ -725,6 +769,58 @@ 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. + #[test] + fn alpha_window_push_circular_writes_to_indexed_slot() -> Result<(), MLError> { + let ctx = CudaContext::new(0).map_err(|e| MLError::ModelError(format!("ctx: {e}")))?; + let stream = ctx.default_stream(); + let module = ctx + .load_cubin(ALPHA_WINDOW_PUSH_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("cubin: {e}")))?; + let kernel = module + .load_function("alpha_window_push_kernel") + .map_err(|e| MLError::ModelError(format!("kfn: {e}")))?; + + const K: usize = 4; + 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 (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)?; + } + } + 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"); + Ok(()) + } + /// End-to-end GPU smoke for the Munchausen target launcher. /// /// Loads the cubin, allocates synthetic inputs (batch=2, n_actions=3 with diff --git a/crates/ml/src/cuda_pipeline/alpha_window_push.cu b/crates/ml/src/cuda_pipeline/alpha_window_push.cu new file mode 100644 index 000000000..6538769ee --- /dev/null +++ b/crates/ml/src/cuda_pipeline/alpha_window_push.cu @@ -0,0 +1,25 @@ +// 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. +// +// 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`. +// +// Grid: ceil(state_dim / BLOCK), Block: 32. For state_dim=10 a single +// warp suffices — tiny launch overhead. + +#include + +extern "C" __global__ void alpha_window_push_kernel( + const float* __restrict__ state_in, // [state_dim] (mapped-pinned input) + float* __restrict__ window, // [K, state_dim] row-major + int head_idx, // slot to write (host-tracked) + 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]; +} diff --git a/docs/isv-slots.md b/docs/isv-slots.md index 5e6fef421..e0a190a81 100644 --- a/docs/isv-slots.md +++ b/docs/isv-slots.md @@ -776,3 +776,10 @@ inference; effective policy adapts via ISV slot 543 (threshold), 545 (observed-rate EMA), 546 (Kelly attenuation). No new ISV slots allocated for the E.3 follow-up; E.4.B will add MoE-gate-entropy and 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.