feat(phase-e-4-a): alpha_window_push circular-buffer kernel + GPU test
Phase E.4.A Task 6: tiny CUDA kernel that pushes a state[state_dim] vector into slot `head_idx` of a circular window buffer [K, state_dim]. Host tracks head_idx and zeros buffer on episode reset. This is the GPU-side append primitive that the smoke binary's per-step inference will call (T7) before Mamba2 over the buffer (T8). GPU smoke (alpha_window_push_circular_writes_to_indexed_slot): writes state_a at slot 0, state_b at slot 2, verifies non-targeted slots stay zero. PASS. docs/isv-slots.md: documents kernel as slot-agnostic per kernel-audit-doc hook requirement. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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::<f32>(K * D)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc window: {e}")))?;
|
||||
let state_a: Vec<f32> = vec![1.0, 2.0, 3.0];
|
||||
let state_b: Vec<f32> = 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
|
||||
|
||||
25
crates/ml/src/cuda_pipeline/alpha_window_push.cu
Normal file
25
crates/ml/src/cuda_pipeline/alpha_window_push.cu
Normal file
@@ -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 <cuda_runtime.h>
|
||||
|
||||
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];
|
||||
}
|
||||
Reference in New Issue
Block a user