From 72626706c8a5bf22673d91fb58a47d46b8e07dea Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 1 Apr 2026 22:59:27 +0200 Subject: [PATCH] perf: GPU-resident step counter for experience collection loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../src/cuda_pipeline/experience_kernels.cu | 28 +++++++++- .../cuda_pipeline/gpu_experience_collector.rs | 55 +++++++++++++++++-- 2 files changed, 76 insertions(+), 7 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 72cb6f8b7..c19bdcdde 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -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; +} diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index a87ea3654..062010644 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -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, + /// 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, pinned_actions: PinnedHostBuf, @@ -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::(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::(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, _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.