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

fix(cuda): wire epoch state into EMA/DSR accumulators and vol EMA in experience kernel

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-11 10:14:53 +01:00
parent 7e098777be
commit 83398cd301
1 changed files with 60 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu
View File

@@ -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;
}
}