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