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:
jgrusewski
2026-05-15 20:52:32 +02:00
parent 5d7d4fa3c6
commit 6e86e5b435
4 changed files with 132 additions and 0 deletions

View File

@@ -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)

View File

@@ -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

View 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];
}

View File

@@ -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.