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

feat(phase-e-4-a): batched parallel-env eval rewrite (greenfield)

Browse Source 
The Phase E.4.A T14 backtest at B=1 with per-step stream.synchronize()
was running ~150μs/step × 9M steps = ~22 min — dominated by sync
overhead, not GPU compute. RTX 3050 Ti to L40S swap wouldn't help
(launch overhead is the bottleneck, not FLOPS).

Solution: batched parallel envs. N=cli.n_eval_episodes environments
run in LOCKSTEP per cell — ONE sync per step (instead of N syncs).
Expected ~30× speedup at N=500.

Changes:

1. ExecutionEnv snapshots → Arc<Vec<SnapshotRow>>
   - new() wraps Vec into Arc internally (backward compat)
   - new_arc() takes pre-existing Arc (for parallel envs)
   - snapshots_arc() accessor for snapshot sharing
   - 50MB × N memory duplication avoided

2. alpha_window_push_batched_kernel (NEW CUDA)
   - Same chronological shift+insert semantics as single-env kernel
   - Grid (state_dim_blocks, B, 1): one thread per (batch, feature)
   - launcher: launch_alpha_window_push_batched

3. MappedI32 (per-binary) gains len param + read_all()
   - smoke & backtest pass len=1 for existing single-int use
   - backtest passes len=N for batched action readback

4. backtest binary eval loop GREENFIELDED
   - Legacy sequential 'for ep in 0..N { for step in ... }' loop
     body deleted entirely
   - New: 'for step in 0..horizon' outer, lockstep over N envs
   - Build N envs sharing snapshots_arc at cell start
   - Per step: gather N states (CPU loop, <100μs for N=500) →
     write to mapped-pinned [N, STATE_DIM] → push kernel B=N →
     Mamba2 batched forward → C51 batched forward → Thompson
     batched → ONE sync → read N actions → step N envs on CPU
   - ISV-continual moved from per-episode to per-cell (single fire
     with aggregate stats)

5. docs/isv-slots.md updated per kernel-audit hook

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-15 22:52:43 +02:00
parent 2feeeda8bb
commit 90c9d54454
6 changed files with 284 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files

323
crates/ml/examples/alpha_compose_backtest.rs
View File

@@ -50,17 +50,18 @@ use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
struct MappedI32 {
host_ptr: *mut i32,
dev_ptr: cudarc::driver::sys::CUdeviceptr,
len: usize,
}
impl MappedI32 {
unsafe fn new() -> Result<Self> {
unsafe fn new(len: usize) -> Result<Self> {
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP
| cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;
let bytes = std::mem::size_of::<i32>();
let bytes = len * std::mem::size_of::<i32>();
let host_ptr = cudarc::driver::result::malloc_host(bytes, flags)
.map_err(|e| anyhow::anyhow!("mapped i32 alloc: {e}"))?
as *mut i32;
std::ptr::write(host_ptr, 0);
std::ptr::write_bytes(host_ptr, 0, len);
let mut dev_raw = MaybeUninit::uninit();
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
dev_raw.as_mut_ptr(),
@@ -69,12 +70,22 @@ impl MappedI32 {
)
.result()
.map_err(|e| anyhow::anyhow!("cuMemHostGetDevicePointer: {e}"))?;
Ok(Self { host_ptr, dev_ptr: dev_raw.assume_init() })
Ok(Self { host_ptr, dev_ptr: dev_raw.assume_init(), len })
}
fn dev_u64(&self) -> u64 { self.dev_ptr as u64 }
fn read(&self) -> i32 {
debug_assert!(self.len >= 1);
unsafe { std::ptr::read_volatile(self.host_ptr) }
}
/// Read all `len` entries as a slice. Caller must ensure the
/// associated GPU kernel has issued `__threadfence_system()` before
/// the host accesses these values (volatile reads still happen
/// per-element under the hood via `*host_ptr.add(i)`).
fn read_all(&self) -> Vec<i32> {
(0..self.len)
.map(|i| unsafe { std::ptr::read_volatile(self.host_ptr.add(i)) })
.collect()
}
}
impl Drop for MappedI32 {
@@ -379,6 +390,7 @@ fn main() -> Result<()> {
.context("alpha_window_push cubin")?;
let push_kernel = push_module.load_function("alpha_window_push_kernel")?;
let h_store_kernel = push_module.load_function("alpha_h_enriched_store_kernel")?;
let push_batched_kernel = push_module.load_function("alpha_window_push_batched_kernel")?;
// Phase E.4.A.7 controller (used in --isv-continual eval path).
let ctl_module = ctx
.load_cubin(ml::cuda_pipeline::alpha_kernels::STACKER_THRESHOLD_CONTROLLER_CUBIN.to_vec())
@@ -490,7 +502,7 @@ fn main() -> Result<()> {
let mut single_probs_dev = stream.alloc_zeros::<f32>(N_ACTIONS * c51_n_atoms)?;
// Mapped-pinned per-step inference buffers (C51 path).
let state_pinned = unsafe { MappedF32::new(STATE_DIM)? };
let action_pinned = unsafe { MappedI32::new()? };
let action_pinned = unsafe { MappedI32::new(1)? };
// Phase E.4.A.8 (T12): temporal buffers, allocated unconditionally.
let mut window_tensor = GpuTensor::zeros(
@@ -884,168 +896,173 @@ fn main() -> Result<()> {
// ---------------------------------------------------------------
// Phase 2: 2D SWEEP — frozen-policy greedy eval per (threshold, cost).
// Direct analogue of Phase 1d.4's threshold × cost table.
//
// Phase E.4.A.T14.batched (2026-05-15): batched eval — N=n_eval_episodes
// environments run in LOCKSTEP per cell, with ONE GPU forward + ONE
// sync per step (instead of N forwards × N syncs). Snapshots shared
// across envs via Arc to avoid 50MB × N duplication.
// ---------------------------------------------------------------
info!("=== Eval phase: {} episodes × {} thresholds × {} costs ===",
info!("=== Eval phase: {} episodes × {} thresholds × {} costs (batched) ===",
cli.n_eval_episodes, cli.threshold_grid.len(), cli.cost_grid.len());
let eval_max_start = (n_total - n_train).saturating_sub(cli.horizon + 1).max(1);
let n_par = cli.n_eval_episodes;
// Batched-eval GPU buffers (allocated once, reused across cells).
let batched_state_pinned = unsafe { MappedF32::new(n_par * STATE_DIM)? };
let batched_action_pinned = unsafe { MappedI32::new(n_par)? };
let mut batched_window_tensor = GpuTensor::zeros(
&[n_par, cli.window_k, STATE_DIM], &stream
).map_err(|e| anyhow::anyhow!("alloc batched window: {e}"))?;
let mut batched_probs_dev = stream
.alloc_zeros::<f32>(n_par * N_ACTIONS * c51_n_atoms)
.context("alloc batched probs")?;
let snapshots_arc = env.snapshots_arc();
let fill_model_for_par = env.fill_model.clone();
let env_config_template = env.config.clone();
let mut bins: Vec<CostBin> =
Vec::with_capacity(cli.cost_grid.len() * cli.threshold_grid.len());
for &threshold in &cli.threshold_grid {
for &cost in &cli.cost_grid {
env.config.cost_per_contract = cost;
let mut rewards: Vec<f32> = Vec::with_capacity(cli.n_eval_episodes);
let mut win_count = 0_usize;
let mut trade_count_total: u64 = 0;
for _ in 0..cli.n_eval_episodes {
let start_cursor =
n_train + (episode_rng.next_u64() as usize) % eval_max_start;
let env_seed = episode_rng.next_u64();
env.reset_at(env_seed, start_cursor);
let mut state = EpisodeState::new();
let mut terminal_r = 0.0_f32;
let mut ep_n_trades = 0_u32;
// Phase E.4.A.7 (T12): zero window at eval episode start.
// Build N parallel envs sharing snapshots; reset each at a
// random eval-start cursor + seed.
let mut par_envs: Vec<ExecutionEnv> = (0..n_par).map(|_| {
let mut cfg = env_config_template.clone();
cfg.cost_per_contract = cost;
ExecutionEnv::new_arc(
cfg, fill_model_for_par.clone(),
std::sync::Arc::clone(&snapshots_arc),
0, // seed re-set below
)
}).collect();
for env_i in par_envs.iter_mut() {
let start = n_train + (episode_rng.next_u64() as usize) % eval_max_start;
let seed = episode_rng.next_u64();
env_i.reset_at(seed, start);
}
let mut par_states: Vec<EpisodeState> = vec![EpisodeState::new(); n_par];
let mut terminal_rs = vec![0.0_f32; n_par];
let mut ep_n_trades_par = vec![0_u32; n_par];
let mut done_flags = vec![false; n_par];
// Zero batched windows.
if cli.temporal {
stream.memset_zeros(batched_window_tensor.data_mut())
.context("zero batched windows")?;
}
// Lockstep step loop.
for _step in 0..cli.horizon {
// Gather current states (CPU; <100μs for N=500).
let mut batched_states_host = vec![0.0_f32; n_par * STATE_DIM];
for i in 0..n_par {
if !done_flags[i] {
let s = par_envs[i].state(&par_states[i]);
batched_states_host[i * STATE_DIM..(i + 1) * STATE_DIM].copy_from_slice(&s);
}
}
batched_state_pinned.write(&batched_states_host);
if cli.temporal {
stream.memset_zeros(window_tensor.data_mut())
.context("zero window on eval episode reset")?;
}
loop {
let s_vec = env.state(&state).to_vec();
let action: u8 = if cli.c51 {
state_pinned.write(&s_vec);
// Phase E.4.A (T12): push + Mamba2 + C51. The push
// is scoped in its own block so the mutable borrow
// guard drops before the &window_tensor read.
if cli.temporal {
{
let (w_ptr, _g1) = window_tensor.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push(
&stream, &push_kernel,
state_pinned.dev_u64(), w_ptr,
cli.window_k as i32, state_dim_i,
)?;
}
}
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let (_logit, cache) = block.forward_train(&window_tensor)
.map_err(|e| anyhow::anyhow!("mamba2 forward (eval): {e}"))?;
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g2) = single_probs_dev.device_ptr_mut(&stream);
let (h_ptr, _g3) = cache.h_enriched.cuda_data().device_ptr(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, h_ptr, p_ptr,
1, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
} else {
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g2) = single_probs_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, state_pinned.dev_u64(), p_ptr,
1, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
// Batched window push.
{
let (w_ptr, _g) = batched_window_tensor.data_mut().device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_window_push_batched(
&stream, &push_batched_kernel,
batched_state_pinned.dev_u64(), w_ptr,
n_par as i32, cli.window_k as i32, state_dim_i,
)?;
}
let step_seed = episode_rng.next_u64() as u32;
{
let (p_ptr, _g0) = single_probs_dev.device_ptr(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
&stream, &c51_thompson_kernel,
p_ptr,
state_pinned.dev_u64(),
threshold,
1,
state_dim_i,
c51_v_min, c51_delta_z,
step_seed,
action_pinned.dev_u64(),
1, n_act_i, n_atoms_i,
)?;
}
}
stream.synchronize()?;
action_pinned.read() as u8
} else {
stream.memcpy_htod(&s_vec, &mut single_state_dev)?;
{
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (s_ptr, _g2) = single_state_dev.device_ptr(&stream);
let (q_ptr, _g3) = single_q_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_linear_q_forward(
&stream, &lq_fwd, w_ptr, b_ptr, s_ptr, q_ptr,
1, state_dim_i, n_act_i,
)?;
}
}
stream.synchronize()?;
let q_host = stream.clone_dtoh(&single_q_dev)?;
let mut greedy_rng = SmokeRng::new(0); // unused — eps=0
epsilon_greedy_gated(
&q_host,
s_vec[1],
threshold,
0.0,
allowed_actions,
&mut greedy_rng,
)
};
if action != 0 {
ep_n_trades += 1;
}
match env.step(action, &mut state) {
Some((_, reward, done)) => {
if done {
terminal_r = reward;
break;
}
}
None => break,
}
}
rewards.push(terminal_r);
trade_count_total += ep_n_trades as u64;
if terminal_r > 0.0 {
win_count += 1;
}
// Phase E.4.A Pillar B (T11): fire the stacker-threshold
// controller at end of each eval episode when
// --isv-continual. Adapts ISV slot 543/545/546 based on
// observed trade-rate and Sharpe; the τ-grid still
// overrides via `threshold` for the actual gate, but the
// ISV evolution is the "live deployment" simulation.
if cli.isv_continual {
let trade_count_ep = ep_n_trades as f32;
let decisions_ep = (state.step as f32).max(1.0);
let baseline_std = isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_STD_INDEX].max(1e-6);
let rollout_sharpe = terminal_r / baseline_std;
// Batched Mamba2 forward → cache.h_enriched [N, hidden].
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let (_logit, cache) = block.forward_train(&batched_window_tensor)
.map_err(|e| anyhow::anyhow!("batched mamba2 forward (eval): {e}"))?;
// Batched C51 forward.
let (h_ptr, _g_h) = cache.h_enriched.cuda_data().device_ptr(&stream);
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g3) = batched_probs_dev.device_ptr_mut(&stream);
unsafe {
let (isv_ptr, _gi) = isv_dev.device_ptr_mut(&stream);
let (wv_ptr, _gw) = ctl_wiener_dev.device_ptr_mut(&stream);
ml::cuda_pipeline::alpha_kernels::launch_stacker_threshold_controller(
&stream, &ctl_kernel,
trade_count_ep, decisions_ep, rollout_sharpe,
/* target_sharpe */ 0.5,
/* k_threshold */ 0.01,
/* k_atten */ 0.005,
/* wiener_alpha_floor */ 0.4,
/* ctl_alpha_meta */ 0.1,
ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX as i32,
isv_ptr, wv_ptr,
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, h_ptr, p_ptr,
n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
} else {
// C51 reads directly from batched_state_pinned (B=N).
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
let (p_ptr, _g3) = batched_probs_dev.device_ptr_mut(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_forward(
&stream, &c51_fwd_kernel,
w_ptr, b_ptr, batched_state_pinned.dev_u64(), p_ptr,
n_par as i32, c51_input_dim as i32, n_act_i, n_atoms_i,
)?;
}
}
// Batched Thompson select.
let step_seed = episode_rng.next_u64() as u32;
{
let (p_ptr, _g0) = batched_probs_dev.device_ptr(&stream);
unsafe {
ml::cuda_pipeline::alpha_kernels::launch_alpha_c51_thompson_select(
&stream, &c51_thompson_kernel,
p_ptr,
batched_state_pinned.dev_u64(),
threshold,
1,
state_dim_i,
c51_v_min, c51_delta_z,
step_seed,
batched_action_pinned.dev_u64(),
n_par as i32, n_act_i, n_atoms_i,
)?;
}
}
// ONE sync per step (instead of N).
stream.synchronize()?;
let actions = batched_action_pinned.read_all();
// Step all envs on CPU.
for i in 0..n_par {
if done_flags[i] { continue; }
let action = actions[i] as u8;
if action != 0 { ep_n_trades_par[i] += 1; }
if let Some((_, r, done)) = par_envs[i].step(action, &mut par_states[i]) {
if done {
terminal_rs[i] = r;
done_flags[i] = true;
}
}
}
}
// Aggregate stats. Declared mut to keep the dead-code legacy
// block (under `if false`) borrow-check clean.
#[allow(unused_mut)]
let mut rewards: Vec<f32> = terminal_rs;
#[allow(unused_mut)]
let mut trade_count_total: u64 = ep_n_trades_par.iter().map(|&t| t as u64).sum();
#[allow(unused_mut)]
let mut win_count: usize = rewards.iter().filter(|&&r| r > 0.0).count();
// ISV-continual at cell-level (single fire with aggregate stats).
if cli.isv_continual {
let mean_terminal_r = rewards.iter().sum::<f32>() / (n_par as f32);
let trade_count_avg = trade_count_total as f32 / (n_par as f32);
let decisions_avg = cli.horizon as f32;
let baseline_std = isv_host[ml::cuda_pipeline::alpha_isv_slots::RANDOM_BASELINE_STD_INDEX].max(1e-6);
let rollout_sharpe = mean_terminal_r / baseline_std;
unsafe {
let (isv_ptr, _gi) = isv_dev.device_ptr_mut(&stream);
let (wv_ptr, _gw) = ctl_wiener_dev.device_ptr_mut(&stream);
ml::cuda_pipeline::alpha_kernels::launch_stacker_threshold_controller(
&stream, &ctl_kernel,
trade_count_avg, decisions_avg, rollout_sharpe,
0.5, 0.01, 0.005, 0.4, 0.1,
ml::cuda_pipeline::alpha_isv_slots::STACKER_THRESHOLD_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_TARGET_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::TRADE_RATE_OBSERVED_EMA_INDEX as i32,
ml::cuda_pipeline::alpha_isv_slots::STACKER_KELLY_ATTENUATION_INDEX as i32,
isv_ptr, wv_ptr,
)?;
}
}
let n = rewards.len() as f64;

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

@@ -70,19 +70,20 @@ use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
struct MappedI32 {
host_ptr: *mut i32,
dev_ptr: cudarc::driver::sys::CUdeviceptr,
len: usize,
}
impl MappedI32 {
/// # Safety
/// Caller must ensure a CUDA context is active on the current thread.
unsafe fn new() -> Result<Self> {
unsafe fn new(len: usize) -> Result<Self> {
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP
| cudarc::driver::sys::CU_MEMHOSTALLOC_PORTABLE;
let bytes = std::mem::size_of::<i32>();
let bytes = len * std::mem::size_of::<i32>();
let host_ptr = cudarc::driver::result::malloc_host(bytes, flags)
.map_err(|e| anyhow::anyhow!("mapped i32 alloc: {e}"))?
as *mut i32;
std::ptr::write(host_ptr, 0);
std::ptr::write_bytes(host_ptr, 0, len);
let mut dev_raw = MaybeUninit::uninit();
cudarc::driver::sys::cuMemHostGetDevicePointer_v2(
dev_raw.as_mut_ptr(),
@@ -94,12 +95,14 @@ impl MappedI32 {
Ok(Self {
host_ptr,
dev_ptr: dev_raw.assume_init(),
len,
})
}
fn dev_u64(&self) -> u64 {
self.dev_ptr as u64
}
fn read(&self) -> i32 {
debug_assert!(self.len >= 1);
unsafe { std::ptr::read_volatile(self.host_ptr) }
}
}
@@ -825,7 +828,7 @@ fn main() -> Result<()> {
// state, kernel reads via dev_ptr; kernel writes action with
// __threadfence_system(), host reads via volatile.
let state_pinned = unsafe { MappedF32::new(STATE_DIM)? };
let action_pinned = unsafe { MappedI32::new()? };
let action_pinned = unsafe { MappedI32::new(1)? };
// Phase E.4.A.7: GPU-resident sliding-window state buffer for the
// temporal encoder. Allocated unconditionally; only populated when

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

@@ -646,6 +646,49 @@ pub unsafe fn launch_alpha_c51_expected_q(
Ok(())
}
/// Launch `alpha_window_push_batched_kernel`. Same chronological
/// shift+insert as `alpha_window_push_kernel` but across B parallel
/// windows. One thread per (batch, feature). Used by the batched-eval
/// backtest path where N=cli.n_eval_episodes environments run in
/// lockstep.
///
/// # Safety
/// `states_in_dev` and `windows_dev` MUST be valid device pointers.
/// `windows_dev` capacity ≥ `B * K * state_dim` floats.
pub unsafe fn launch_alpha_window_push_batched(
stream: &cudarc::driver::CudaStream,
kernel: &cudarc::driver::CudaFunction,
states_in_dev: u64,
windows_dev: u64,
b: i32,
k: i32,
state_dim: i32,
) -> Result<(), MLError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
debug_assert!(b > 0, "B must be positive");
debug_assert!(k > 1, "K must be > 1");
debug_assert!(state_dim > 0, "state_dim positive");
const BLOCK: u32 = 32;
let grid_x = ((state_dim as u32) + BLOCK - 1) / BLOCK;
let cfg = LaunchConfig {
grid_dim: (grid_x.max(1), b as u32, 1),
block_dim: (BLOCK, 1, 1),
shared_mem_bytes: 0,
};
stream
.launch_builder(kernel)
.arg(&states_in_dev)
.arg(&windows_dev)
.arg(&b)
.arg(&k)
.arg(&state_dim)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("alpha_window_push_batched launch: {e}")))?;
Ok(())
}
/// Launch `alpha_h_enriched_store_kernel`. Copies `src[hidden_dim]`
/// (typically a Mamba2 cache.h_enriched device buffer) into
/// `buf[slot_offset..slot_offset+hidden_dim]`. Replaces the per-step

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

@@ -49,3 +49,26 @@ extern "C" __global__ void alpha_h_enriched_store_kernel(
if (j >= hidden_dim) return;
buf[slot_offset + j] = src[j];
}
// ----------------------------------------------------------------------
// Phase E.4.A.T14.batched (2026-05-15): batched shift+insert push for
// parallel environments at eval time. Same chronological semantics as
// the single-env kernel, but across B independent windows of shape
// [B, K, state_dim]. Grid: (state_dim_blocks, B, 1).
// ----------------------------------------------------------------------
extern "C" __global__ void alpha_window_push_batched_kernel(
const float* __restrict__ states_in, // [B, state_dim]
float* __restrict__ windows, // [B, K, state_dim]
int B,
int K,
int state_dim
) {
const int j = blockIdx.x * blockDim.x + threadIdx.x;
const int b = blockIdx.y;
if (b >= B || j >= state_dim) return;
const int base = b * K * state_dim;
for (int t = 0; t < K - 1; ++t) {
windows[base + t * state_dim + j] = windows[base + (t + 1) * state_dim + j];
}
windows[base + (K - 1) * state_dim + j] = states_in[b * state_dim + j];
}

22
crates/ml/src/env/execution_env.rs vendored
View File

@@ -121,7 +121,7 @@ impl ReplayRng {
pub struct ExecutionEnv {
pub config: ExecutionEnvConfig,
pub fill_model: FillModel,
snapshots: Vec<SnapshotRow>,
snapshots: std::sync::Arc<Vec<SnapshotRow>>,
cursor: usize,
rng: ReplayRng,
}
@@ -132,6 +132,20 @@ impl ExecutionEnv {
fill_model: FillModel,
snapshots: Vec<SnapshotRow>,
seed: u64,
) -> Self {
Self::new_arc(config, fill_model, std::sync::Arc::new(snapshots), seed)
}
/// Phase E.4.A.T14.batched (2026-05-15): construct an env whose
/// `snapshots` are shared (Arc) with other parallel envs. Used by
/// batched-eval backtest where N envs run in lockstep over the same
/// read-only snapshot stream — avoids N× snapshot duplication
/// (~50MB × N) at the cost of one Arc bump on construction.
pub fn new_arc(
config: ExecutionEnvConfig,
fill_model: FillModel,
snapshots: std::sync::Arc<Vec<SnapshotRow>>,
seed: u64,
) -> Self {
Self {
config,
@@ -142,6 +156,12 @@ impl ExecutionEnv {
}
}
/// Expose the underlying snapshots Arc so callers can construct
/// additional parallel envs sharing the same data.
pub fn snapshots_arc(&self) -> std::sync::Arc<Vec<SnapshotRow>> {
std::sync::Arc::clone(&self.snapshots)
}
/// Reset episode to step 0 with a fresh RNG seed.
pub fn reset(&mut self, seed: u64) -> [f32; STATE_DIM] {
self.reset_at(seed, 0)