From 83398cd301361d68801cc782c6affcbccb5cf5fb Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 11 Mar 2026 10:14:53 +0100 Subject: [PATCH] fix(cuda): wire epoch state into EMA/DSR accumulators and vol EMA in experience kernel The epoch_vol_ema, epoch_dsr_mean, and epoch_dsr_var variables were loaded from global memory at kernel start but never wired into the actual computation. The EMA normalizer (ema_mean/ema_var) and DSR accumulators (dsr_A/dsr_B) were initialized with hardcoded constants, so epoch state round-tripped unchanged. Changes: - Seed ema_mean/ema_var from epoch_dsr_mean/epoch_dsr_var at kernel start - Seed dsr_A/dsr_B from epoch_dsr_mean/epoch_dsr_var at kernel start - Seed ema_init/dsr_initialized from epoch_step_count > 0 (skip cold start on subsequent epochs) - Add local_vol_ema/local_median_vol seeded from epoch_vol_ema/epoch_median_vol - Update vol EMA each timestep from market feature index 3 (log close return), mirroring CPU DQNTrainer::vol_ema / median_vol logic - At writeback, write actual computed dsr_A/dsr_B or ema_mean/ema_var (conditioned on use_dsr), and computed local_vol_ema/local_median_vol, instead of unmodified loaded epoch values - Applied identically to both dqn_full_experience_kernel (standard) and dqn_full_experience_kernel_warp (warp-cooperative) variants Co-Authored-By: Claude Opus 4.6 --- .../cuda_pipeline/dqn_experience_kernel.cu | 84 +++++++++++++------ 1 file changed, 60 insertions(+), 24 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu index fd5b0a0b8..7b51694c5 100644 --- a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu @@ -1303,15 +1303,19 @@ extern "C" __global__ void dqn_full_experience_kernel( int total_action_count = 0; for (int i = 0; i < NUM_ACTIONS; i++) action_counts[i] = 0; - /* ---- EMA reward normalizer (matches CPU RewardNormalizer) ---- */ - float ema_mean = 0.0f; - float ema_var = 1.0f; - int ema_init = 0; + /* ---- EMA reward normalizer — seeded from epoch state for cross-epoch continuity ---- */ + float ema_mean = epoch_dsr_mean; + float ema_var = epoch_dsr_var; + int ema_init = (epoch_step_count > 0.0f) ? 1 : 0; - /* ---- DSR accumulators (per-episode) ---- */ - float dsr_A = 0.0f; - float dsr_B = 1e-8f; - int dsr_initialized = 0; + /* ---- DSR accumulators (per-episode) — seeded from epoch state ---- */ + float dsr_A = epoch_dsr_mean; + float dsr_B = epoch_dsr_var; + int dsr_initialized = (epoch_step_count > 0.0f) ? 1 : 0; + + /* ---- Vol EMA (epoch-persistent) — seeded from epoch state ---- */ + float local_vol_ema = epoch_vol_ema; + float local_median_vol = epoch_median_vol; /* ---- N-step ring buffer (per-episode) ---- */ float nstep_ring[N_STEPS_MAX]; @@ -1344,6 +1348,14 @@ extern "C" __global__ void dqn_full_experience_kernel( int mf_off = global_bar * MARKET_DIM; for (int i = 0; i < MARKET_DIM; i++) state[i] = market_features[mf_off + i]; + + /* Update vol EMA from market feature index 3 (log close return). + * Mirrors CPU DQNTrainer::vol_ema / median_vol update logic. */ + float vol_sample = fabsf(state[3]); + if (vol_sample > 0.0f && vol_sample < 1.0f) { + local_vol_ema = 0.99f * local_vol_ema + 0.01f * vol_sample; + local_median_vol = 0.999f * local_median_vol + 0.001f * vol_sample; + } } /* ---- Step 2: Compute portfolio features ---- */ @@ -1939,10 +1951,15 @@ extern "C" __global__ void dqn_full_experience_kernel( // Only vol_ema, dsr_mean, dsr_var, step_count are updated per-kernel. // Portfolio state is per-episode (in portfolio_states), not per-epoch. if (threadIdx.x == 0 && blockIdx.x == gridDim.x - 1) { - epoch_state[0] = epoch_vol_ema; - epoch_state[1] = epoch_median_vol; - epoch_state[5] = epoch_dsr_mean; - epoch_state[6] = epoch_dsr_var; + epoch_state[0] = local_vol_ema; + epoch_state[1] = local_median_vol; + if (use_dsr) { + epoch_state[5] = dsr_A; + epoch_state[6] = dsr_B; + } else { + epoch_state[5] = ema_mean; + epoch_state[6] = ema_var; + } epoch_state[7] = epoch_step_count + (float)L; } } @@ -2159,15 +2176,19 @@ extern "C" __global__ void dqn_full_experience_kernel_warp( int total_action_count = 0; for (int i = 0; i < NUM_ACTIONS; i++) action_counts[i] = 0; - /* ---- EMA reward normalizer (lane 0 only) ---- */ - float ema_mean = 0.0f; - float ema_var = 1.0f; - int ema_init = 0; + /* ---- EMA reward normalizer (lane 0 only) — seeded from epoch state ---- */ + float ema_mean = epoch_dsr_mean; + float ema_var = epoch_dsr_var; + int ema_init = (epoch_step_count > 0.0f) ? 1 : 0; - /* ---- DSR accumulators (per-episode, lane 0 only) ---- */ - float dsr_A = 0.0f; - float dsr_B = 1e-8f; - int dsr_initialized = 0; + /* ---- DSR accumulators (per-episode, lane 0 only) — seeded from epoch state ---- */ + float dsr_A = epoch_dsr_mean; + float dsr_B = epoch_dsr_var; + int dsr_initialized = (epoch_step_count > 0.0f) ? 1 : 0; + + /* ---- Vol EMA (epoch-persistent, lane 0 only) — seeded from epoch state ---- */ + float local_vol_ema = epoch_vol_ema; + float local_median_vol = epoch_median_vol; /* ---- N-step ring buffer (per-episode, lane 0 only) ---- */ float nstep_ring[N_STEPS_MAX]; @@ -2211,6 +2232,16 @@ extern "C" __global__ void dqn_full_experience_kernel_warp( state_dist[i / 32] = market_features[mf_off + i]; } + /* Update vol EMA from market feature index 3 (log close return), lane 0 only. + * Mirrors CPU DQNTrainer::vol_ema / median_vol update logic. */ + if (lane_id == 0) { + float vol_sample = fabsf(market_features[mf_off + 3]); + if (vol_sample > 0.0f && vol_sample < 1.0f) { + local_vol_ema = 0.99f * local_vol_ema + 0.01f * vol_sample; + local_median_vol = 0.999f * local_median_vol + 0.001f * vol_sample; + } + } + /* ---- Step 2: Compute portfolio features (all lanes read, needed for state) ---- */ int t_off = global_bar * 4; current_close = targets[t_off + 0]; @@ -2878,10 +2909,15 @@ extern "C" __global__ void dqn_full_experience_kernel_warp( // Only vol_ema, dsr_mean, dsr_var, step_count are updated per-kernel. // Portfolio state is per-episode (in portfolio_states), not per-epoch. if (threadIdx.x == 0 && blockIdx.x == gridDim.x - 1) { - epoch_state[0] = epoch_vol_ema; - epoch_state[1] = epoch_median_vol; - epoch_state[5] = epoch_dsr_mean; - epoch_state[6] = epoch_dsr_var; + epoch_state[0] = local_vol_ema; + epoch_state[1] = local_median_vol; + if (use_dsr) { + epoch_state[5] = dsr_A; + epoch_state[6] = dsr_B; + } else { + epoch_state[5] = ema_mean; + epoch_state[6] = ema_var; + } epoch_state[7] = epoch_step_count + (float)L; } }