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

refactor(phase-e-4-a): window push to shift+insert (chronological layout)

Browse Source 
Switch alpha_window_push from circular-buffer-with-head_idx layout
to shift+insert layout matching production mamba2_update_history.
Slot 0 = oldest, slot K-1 = newest after each push, matching
Mamba2Block's [B, K, in_dim] input contract directly (no reorder).

Cost: O(K-1) shifts per state_dim feature per push. For K=16,
state_dim=10: 10 threads × ~15 ops each = trivial.

Kernel signature: drops head_idx, adds K. Test updated to verify
chronological shift across 3 pushes into 4-slot buffer.
docs/isv-slots.md updated per kernel-audit-doc hook.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-15 21:06:23 +02:00
parent 70df697328
commit 35dcb87709
4 changed files with 71 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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::<f32>(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<f32> = 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);

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

@@ -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::<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 states: [Vec<f32>; 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(())
}

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

@@ -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 <cuda_runtime.h>
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];
}

11
docs/isv-slots.md
View File

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