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

perf: GPU-resident step counter for experience collection loop

Browse Source 
Replace host-side `current_t` scalar parameter in experience_env_step
with a GPU-resident counter buffer.  The kernel now reads the timestep
index from device memory (step_counter_gpu[0]) instead of receiving it
as a kernel argument that changes every iteration.  A tiny single-thread
step_counter_advance kernel increments the counter after each env_step.

This eliminates per-timestep host→GPU parameter variation in the 100-
iteration experience collection loop, making all iterations dispatch
identical kernel argument sets — a prerequisite for future CUDA Graph
capture of the timestep sequence.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-01 22:59:27 +02:00
parent 6767241b17
commit 72626706c8
2 changed files with 76 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -826,7 +826,7 @@ extern "C" __global__ void experience_env_step(
int b0_size,
int b1_size,
int b2_size,
int current_t,
const int* __restrict__ current_t_ptr, /* [1] GPU-resident timestep counter */
const __nv_bfloat16* __restrict__ cvar_scales, /* [N] or NULL — CVaR position scaling */
const __nv_bfloat16* __restrict__ q_gaps, /* [N] or NULL — Q-gap conviction scaling */
__nv_bfloat16* raw_returns_out, /* [N, L] output: true per-bar portfolio return */
@@ -869,7 +869,12 @@ extern "C" __global__ void experience_env_step(
int t = current_timesteps[i];
int bar_idx = episode_starts[i] + t;
/* Flat index into [N, L] output arrays. */
/* Flat index into [N, L] output arrays.
* current_t_ptr is a GPU-resident scalar counter [1], advanced by
* step_counter_advance kernel after each timestep. Reading from device
* memory instead of a host-side scalar eliminates per-timestep
* host→GPU parameter variation, enabling future CUDA Graph capture. */
int current_t = current_t_ptr[0];
long long out_off = (long long)i * L + current_t;
/* ---- Write current state to replay buffer (f32) ---- */
@@ -1936,3 +1941,22 @@ extern "C" __global__ void classify_regime_gpu(
}
regimes[idx] = regime;
}
/* ═════════════════════════════════════════════════════════════════════════ */
/* GPU-resident step counter advance */
/* ═════════════════════════════════════════════════════════════════════════ */
/**
* Advances the GPU-resident timestep counter by 1.
*
* Launched as a single-thread kernel (grid=1, block=1) after each
* experience_env_step call. Keeping the step counter on-device
* eliminates per-timestep host→GPU parameter variation in the
* experience collection loop, enabling future CUDA Graph capture
* of the entire timestep sequence.
*
* @param counter [1] int — incremented in-place.
*/
extern "C" __global__ void step_counter_advance(int* counter) {
counter[0] += 1;
}

55
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -588,6 +588,14 @@ pub struct GpuExperienceCollector {
/// #31 Bottleneck tanh+concat kernel (loaded from utility cubin). None when bn_dim=0.
bn_tanh_concat_fn: CudaFunction,
// ── GPU-resident step counter (Task 8: CUDA Graph prep) ────────────
/// GPU-resident timestep counter [1] i32 — replaces host-side `t_i32`
/// parameter in the experience collection loop. Zeroed before each
/// epoch and advanced by `step_counter_kernel` after each env_step.
step_counter_gpu: CudaSlice<i32>,
/// Single-thread kernel that increments `step_counter_gpu` by 1.
step_counter_kernel: CudaFunction,
// Pre-allocated pinned host buffers for DtoH transfers.
pinned_states: PinnedHostBuf<f32>,
pinned_actions: PinnedHostBuf<i32>,
@@ -655,7 +663,7 @@ impl GpuExperienceCollector {
let total_branch_actions = branch_sizes[0] + branch_sizes[1] + branch_sizes[2]; // 11
// ── Step 1: Compile experience kernels ──────────────────────────
let (state_gather_kernel, action_select_kernel, expert_override_kernel, env_step_kernel, expected_q_kernel) =
let (state_gather_kernel, action_select_kernel, expert_override_kernel, env_step_kernel, expected_q_kernel, step_counter_kernel) =
compile_experience_kernels(&stream, state_dim, market_dim)?;
let (nstep_kernel, reward_norm_kernel) = compile_nstep_kernel(&stream)?;
let fill_episode_ids_kernel = compile_fill_episode_ids_kernel(&stream)?;
@@ -1018,6 +1026,10 @@ impl GpuExperienceCollector {
let exp_bn_concat = stream.alloc_zeros::<f32>(alloc_episodes * (bn_alloc + portfolio_dim_bn))
.map_err(|e| MLError::ModelError(format!("alloc exp_bn_concat: {e}")))?;
// Task 8: GPU-resident step counter for experience collection loop
let step_counter_gpu = stream.alloc_zeros::<i32>(1)
.map_err(|e| MLError::ModelError(format!("alloc step_counter_gpu: {e}")))?;
Ok(Self {
stream,
state_dim,
@@ -1107,6 +1119,8 @@ impl GpuExperienceCollector {
bottleneck_dim: bn_dim_from_params,
market_dim_bn: market_dim,
bn_tanh_concat_fn,
step_counter_gpu,
step_counter_kernel,
pinned_states,
pinned_actions,
pinned_rewards,
@@ -1503,6 +1517,19 @@ impl GpuExperienceCollector {
.map_err(|e| MLError::ModelError(format!("memset position_histogram: {e}")))?;
}
// Task 8: Zero GPU-resident step counter before timestep loop.
// env_step reads current_t from this device buffer instead of a
// host-side scalar, eliminating per-timestep parameter variation.
self.stream.memset_zeros(&mut self.step_counter_gpu)
.map_err(|e| MLError::ModelError(format!("memset step_counter_gpu: {e}")))?;
// Single-thread launch config for step_counter_advance kernel.
let step_counter_cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
for t in 0..timesteps {
// ── 1. Gather states: [N, state_dim] ────────────────────────
unsafe {
@@ -1748,7 +1775,8 @@ impl GpuExperienceCollector {
let tx_cost = config.tx_cost_multiplier;
let rw_loss_av = config.loss_aversion;
let l_i32 = timesteps as i32;
let t_i32 = t as i32;
// Task 8: current_t is now GPU-resident — pass device pointer
// instead of host scalar. step_counter_gpu[0] == t at this point.
unsafe {
self.stream
.launch_builder(&self.env_step_kernel)
@@ -1774,7 +1802,7 @@ impl GpuExperienceCollector {
.arg(&b0)
.arg(&b1)
.arg(&b2)
.arg(&t_i32)
.arg(&self.step_counter_gpu) // GPU-resident current_t [1]
.arg(&self.cvar_scales_ptr) // CVaR position scaling (0 = NULL = no scaling)
.arg(&self.q_gaps_buf) // Q-gap conviction scaling
.arg(&mut self.raw_returns_out) // Raw portfolio returns (unshaped) for Sharpe/MaxDD
@@ -1803,6 +1831,20 @@ impl GpuExperienceCollector {
"experience_env_step t={t}: {e}"
)))?;
}
// ── 6. Advance GPU-resident step counter ────────────────────
// Single-thread kernel: step_counter_gpu[0] += 1.
// Runs on the same stream — serialized after env_step, before
// the next iteration's state_gather.
unsafe {
self.stream
.launch_builder(&self.step_counter_kernel)
.arg(&mut self.step_counter_gpu)
.launch(step_counter_cfg)
.map_err(|e| MLError::ModelError(format!(
"step_counter_advance t={t}: {e}"
)))?;
}
}
// #33 Saboteur: select best (worst for trader) params from this epoch
@@ -2127,7 +2169,7 @@ fn compile_experience_kernels(
stream: &Arc<CudaStream>,
_state_dim: usize,
_market_dim: usize,
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
// Load precompiled cubin (state_dim/market_dim are runtime kernel params)
let context = stream.context();
let module = context
@@ -2149,8 +2191,11 @@ fn compile_experience_kernels(
let expected_q = module
.load_function("compute_expected_q")
.map_err(|e| MLError::ModelError(format!("load compute_expected_q: {e}")))?;
let step_counter = module
.load_function("step_counter_advance")
.map_err(|e| MLError::ModelError(format!("load step_counter_advance: {e}")))?;
Ok((state_gather, action_select, expert_override, env_step, expected_q))
Ok((state_gather, action_select, expert_override, env_step, expected_q, step_counter))
}
/// Compile the n-step return accumulation kernel.