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

fix: eliminate NVRTC, fix evaluator SIGSEGV, 22D search space, trade physics

Browse Source 
Major changes:
- Eliminate ALL runtime NVRTC: c51_loss + mse_loss kernels converted to
  precompiled cubins with runtime num_atoms/v_min/v_max/branch params
- Fix evaluator SIGSEGV: rng_states/q_gaps sized for chunked batch (cn)
  not n_windows; NULL pointer guard in action_select kernel
- Add trade_physics.cuh shared header for train/eval consistency
- Add equity circuit breaker (25% DD from peak) + margin-aware position cap
- Consolidate 46D→22D hyperopt search space, enable ensemble by default
- Fix trade counting: use exposure index not factored action
- Fix Calmar overflow: clamp to ±100 in kernel
- Softer CVaR penalty (cap 3.0 not 10.0) for undertrained models
- Fix win_rate display (ratio→percentage)
- Remove dead code (normalize_reward, calculate_completion_penalty)
- Add tracing subscriber to hyperopt test for visible metrics
- Per-chunk sync in evaluator for reliable error reporting

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-27 00:33:05 +01:00
parent 65cc0175ac
commit d485dde5a2
16 changed files with 1187 additions and 1489 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-hyperopt/src/optimizer.rs
View File

@@ -498,7 +498,6 @@ impl ArgminOptimizer {
}
Err(e) => {
tracing::error!("Trial {} training FAILED: {:#}", trial_num, e);
eprintln!("!!! TRIAL {} TRAINING FAILED: {:#}", trial_num, e);
eprintln!("!!! FULL ERROR CHAIN: {:?}", e);
(1e6, None)
}

8
crates/ml/build.rs
View File

@@ -12,6 +12,10 @@ fn main() {
let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
let kernel_dir = Path::new("src/cuda_pipeline");
let common_header = kernel_dir.join("common_device_functions.cuh");
let trade_physics_header = kernel_dir.join("trade_physics.cuh");
// Rebuild cubins when shared headers change
println!("cargo:rerun-if-changed={}", trade_physics_header.display());
// Detect GPU architecture from env or default to sm_80
let cuda_compute_cap = std::env::var("CUDA_COMPUTE_CAP").unwrap_or_else(|_| "80".to_string());
@@ -144,7 +148,8 @@ fn compile_kernel(
let tmp_src = out_dir.join(format!("_{kernel_name}"));
std::fs::write(&tmp_src, &full_source).unwrap();
// Compile with nvcc
// Compile with nvcc (-I kernel_dir for shared .cuh includes)
let include_flag = format!("-I{}", kernel_dir.display());
let status = Command::new(nvcc)
.args([
"-cubin",
@@ -153,6 +158,7 @@ fn compile_kernel(
"--use_fast_math",
"--ftz=true",
"--fmad=true",
&include_flag,
"-o", cubin_path.to_str().unwrap(),
tmp_src.to_str().unwrap(),
])

169
crates/ml/src/cuda_pipeline/backtest_env_kernel.cu
View File

@@ -10,8 +10,12 @@
// [5] hold_time - bars held in current position (0 = flat)
// [6] cum_return - cumulative log return
// [7] step_count - number of completed steps
//
// Trade physics (action decoding, hold enforcement, tx costs, capital floor)
// are shared with the training kernel via trade_physics.cuh.
#define PORTFOLIO_STATE_SIZE 8
#include "trade_physics.cuh"
extern "C" __global__ void backtest_env_step(
// Market data (read-only, uploaded once)
@@ -37,20 +41,17 @@ extern "C" __global__ void backtest_env_step(
float tx_cost_bps,
float spread_cost,
int current_step,
// Branching DQN action space sizes
int b0_size, // exposure branch size (default 9)
int b1_size, // order type branch size (default 3)
int b2_size, // urgency branch size (default 3)
int min_hold_bars // minimum bars to hold before exit/reversal
int b0_size,
int b1_size,
int b2_size,
int min_hold_bars
) {
// Shared memory tile for coalesced portfolio reads/writes
__shared__ float shmem_pf[256 * PORTFOLIO_STATE_SIZE];
int w = blockIdx.x * blockDim.x + threadIdx.x;
int local_tid = threadIdx.x;
// Cooperative tile load: all active threads load their portfolio slice
// in one coalesced warp transaction before divergent early-exit checks.
// Cooperative tile load
int ps = w * PORTFOLIO_STATE_SIZE;
if (w < n_windows) {
for (int i = 0; i < PORTFOLIO_STATE_SIZE; i++) {
@@ -84,64 +85,55 @@ extern "C" __global__ void backtest_env_step(
float hold_time = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 5];
float cum_return = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 6];
// Decode factored action to exposure index (matches experience_env_step).
// Factored action = exposure_idx * b1_size * b2_size + order_idx * b2_size + urgency_idx
int action_val = actions[w];
int exposure_idx;
if (b1_size > 0 && b2_size > 0) {
int denom = b1_size * b2_size;
exposure_idx = action_val / denom;
} else {
exposure_idx = action_val;
// ── Capital floor circuit breaker (shared: trade_physics.cuh) ────────
if (check_capital_floor(value, max_equity)) {
// Force flat: close any open position at current price
if (fabsf(position) > 0.001f && close > 0.0f && entry_price > 0.0f) {
float pnl = position * (close - entry_price);
cash += pnl;
}
float liq_value = cash;
float liq_ret = (value > 0.01f) ? (liq_value - value) / value : 0.0f;
step_rewards[w] = liq_ret;
step_returns[w * max_len + current_step] = liq_ret;
actions_history[w * max_len + current_step] = b0_size / 2; // Flat
portfolio_state[ps + 0] = liq_value;
portfolio_state[ps + 1] = 0.0f;
portfolio_state[ps + 2] = liq_value;
portfolio_state[ps + 3] = 0.0f;
portfolio_state[ps + 4] = max_equity;
portfolio_state[ps + 5] = 0.0f;
portfolio_state[ps + 6] = cum_return + liq_ret;
portfolio_state[ps + 7] += 1.0f;
done_flags[w] = 1;
return;
}
// Clamp to valid range
if (exposure_idx < 0) exposure_idx = 0;
if (exposure_idx >= b0_size) exposure_idx = b0_size - 1;
// Map exposure index to fraction: -1.0 + idx * (2.0 / (b0_size - 1))
// With b0_size=9: step=0.25, range [-1.0, +1.0]
float step_size = (b0_size > 1) ? 2.0f / (float)(b0_size - 1) : 0.0f;
float target_exposure = -1.0f + (float)exposure_idx * step_size;
target_exposure *= max_position;
// ── Decode action (shared: trade_physics.cuh) ────────────────────────
int action_val = actions[w];
int exposure_idx = decode_exposure_index(action_val, b0_size, b1_size, b2_size);
float target_exposure = compute_target_position(exposure_idx, b0_size, max_position);
int order_type_idx = decode_order_type(action_val, b1_size, b2_size);
// ── Margin-aware position cap (shared: trade_physics.cuh) ────────────
// ES futures: ~$15K initial margin per contract. Prevents overleveraging
// when equity is depleted — a depleted account can't hold the same
// position as a full account. close * 50 ≈ contract notional for ES.
float margin_per_contract = close * 50.0f * 0.06f; // ~6% of notional ≈ $15K for ES at $5000
target_exposure = apply_margin_cap(target_exposure, value, margin_per_contract);
// Suppress unused variable warnings
(void)open; (void)high; (void)low;
// Decode order_type and urgency from factored action (for tx cost differentiation)
int order_type_idx = (b1_size > 0 && b2_size > 0)
? (action_val / b2_size) % b1_size : 0;
// ── Hold enforcement (shared: trade_physics.cuh) ─────────────────────
int is_last_bar = (current_step >= wlen - 1) ? 1 : 0;
target_exposure = enforce_hold(position, target_exposure, hold_time, min_hold_bars, is_last_bar);
// Hold enforcement: block exits and reversals during minimum hold period.
// Matches the training kernel's hold enforcement to eliminate train/eval mismatch.
{
float prev_position = position; // position before any action is applied
int prev_sign = (prev_position > 0.001f) ? 1 : ((prev_position < -0.001f) ? -1 : 0);
int curr_sign = (target_exposure > 0.001f) ? 1 : ((target_exposure < -0.001f) ? -1 : 0);
int wants_exit = (prev_sign != 0 && curr_sign == 0);
int wants_reversal = (prev_sign != 0 && curr_sign != 0 && prev_sign != curr_sign);
if (hold_time > 0.0f && hold_time < (float)min_hold_bars
&& (wants_exit || wants_reversal)
&& current_step < wlen - 1) {
target_exposure = prev_position; // Override: keep position
}
}
// Execute trade if position changes
// ── Execute trade (shared: trade_physics.cuh) ────────────────────────
float delta = target_exposure - position;
float trade_cost = 0.0f;
float prev_position = position;
if (fabsf(delta) > 0.001f && close > 0.0f) {
/* EXACT match of training kernel (experience_env_step) tx cost formula:
* tx_cost = |delta| * close * (multiplier * 0.0001 * spread_scale * impact + premium)
* - spread_scale: sqrt(|delta|/max_pos) (Almgren-Chriss square-root impact)
* - order_premium: Market=0, IoC=+2bps, LimitMaker=-5bps rebate */
float spread_scale = sqrtf(fabsf(delta) / fmaxf(max_position, 0.01f));
if (spread_scale < 1.0f) spread_scale = 1.0f;
float impact_scale = 1.0f; /* simplified: no dynamic spread widening in backtest */
float order_premium = (order_type_idx == 0) ? 0.0f
: (order_type_idx == 1) ? 0.0002f /* IoC: +2 bps */
: -0.0005f; /* LimitMaker: -5 bps rebate */
trade_cost = fabsf(delta) * close * (tx_cost_bps * 0.0001f * spread_scale * impact_scale + order_premium)
+ fabsf(delta) * spread_cost * 0.5f;
float trade_cost = compute_tx_cost(delta, close, tx_cost_bps, spread_cost, max_position, order_type_idx);
cash -= trade_cost;
// Mark-to-market old position
@@ -154,27 +146,8 @@ extern "C" __global__ void backtest_env_step(
entry_price = close;
}
// Update hold_time based on resulting position.
// Read prev_position from local var (before trade block may have updated `position`).
// Note: `position` after the trade block reflects the new position.
{
float prev_pos_for_hold = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 1];
if (fabsf(position) > 0.001f) {
if (fabsf(prev_pos_for_hold) < 0.001f) {
hold_time = 1.0f; // Entry from flat
} else {
int prev_sign_h = (prev_pos_for_hold > 0.001f) ? 1 : -1;
int new_sign_h = (position > 0.001f) ? 1 : -1;
if (new_sign_h == prev_sign_h) {
hold_time += 1.0f; // Same direction — increment
} else {
hold_time = 1.0f; // Reversal allowed — new trade
}
}
} else {
hold_time = 0.0f; // Flat
}
}
// ── Hold time tracking (shared: trade_physics.cuh) ───────────────────
hold_time = update_hold_time(prev_position, position, hold_time);
// Mark-to-market current position
float unrealized = 0.0f;
@@ -183,23 +156,51 @@ extern "C" __global__ void backtest_env_step(
}
float new_value = cash + unrealized;
// ── Post-trade floor check: catch intra-step breaches ────────────
// The pre-trade check (top of kernel) catches breaches from the
// PREVIOUS step. This catches breaches from THIS step's price move
// or transaction costs. Without it, max_dd overshoots the 75% floor
// because the metrics kernel records the breached value before the
// next step's pre-trade check fires.
if (check_capital_floor(new_value, max_equity)) {
// Emergency liquidation: close position at current price
if (fabsf(position) > 0.001f) {
// Position already marked-to-market above; just go flat
new_value = cash + unrealized; // already computed
position = 0.0f;
cash = new_value;
entry_price = 0.0f;
}
float step_ret = (value > 0.01f) ? (new_value - value) / value : 0.0f;
step_rewards[w] = step_ret;
step_returns[w * max_len + current_step] = step_ret;
actions_history[w * max_len + current_step] = b0_size / 2;
portfolio_state[ps + 0] = new_value;
portfolio_state[ps + 1] = 0.0f;
portfolio_state[ps + 2] = new_value;
portfolio_state[ps + 3] = 0.0f;
portfolio_state[ps + 4] = max_equity;
portfolio_state[ps + 5] = 0.0f;
portfolio_state[ps + 6] = cum_return + step_ret;
portfolio_state[ps + 7] += 1.0f;
done_flags[w] = 1;
return;
}
// Step return
float step_ret = (value > 0.0f) ? (new_value - value) / value : 0.0f;
float new_cum_return = cum_return + step_ret;
// Update max equity for drawdown
float new_max = fmaxf(max_equity, new_value);
// Write portfolio state directly to global memory (no shmem write-back needed —
// threads with done_flags[w]==1 exited early and cannot participate in __syncthreads)
// Write portfolio state
portfolio_state[ps + 0] = new_value;
portfolio_state[ps + 1] = position;
portfolio_state[ps + 2] = cash;
portfolio_state[ps + 3] = entry_price;
portfolio_state[ps + 4] = new_max;
portfolio_state[ps + 5] = hold_time; // hold_time (bars in current position)
portfolio_state[ps + 5] = hold_time;
portfolio_state[ps + 6] = new_cum_return;
portfolio_state[ps + 7] += 1.0f; // step count
portfolio_state[ps + 7] += 1.0f;
// Outputs
step_rewards[w] = step_ret;

22
crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu
View File

@@ -108,13 +108,16 @@ extern "C" __global__ void compute_backtest_metrics(
if (exp_idx >= 0 && exp_idx < num_actions)
local_action_mask |= (1 << exp_idx);
// Boundary-aware trade tracking
// Boundary-aware trade tracking: count EXPOSURE changes only.
// The factored action encodes (exposure, order_type, urgency) but only
// exposure changes trigger actual position changes in env_step.
// Using exp_idx prevents order/urgency flips from inflating trade count.
if (i == chunk_start) {
bnd_first_action = act;
bnd_cur_action = act;
bnd_first_action = exp_idx;
bnd_cur_action = exp_idx;
}
if (act != bnd_cur_action) {
if (exp_idx != bnd_cur_action) {
bnd_num_changes++;
if (bnd_num_changes == 1) {
// First change: everything accumulated so far is the prefix
@@ -125,11 +128,11 @@ extern "C" __global__ void compute_backtest_metrics(
if (bnd_cur_return > 0.0f) bnd_complete_wins++;
}
bnd_cur_return = 0.0f;
bnd_cur_action = act;
bnd_cur_action = exp_idx;
}
bnd_cur_return += r;
bnd_last_action = act;
bnd_last_action = exp_idx;
}
// Post-loop: assign suffix and handle no-change edge case
@@ -393,9 +396,11 @@ extern "C" __global__ void compute_backtest_metrics(
}
float daily_mean = total_return / (float)sort_len;
float max_dd = metrics_out[out_base + 2];
float calmar = (max_dd > 1e-8f)
float calmar = (max_dd > 0.001f)
? (daily_mean * annualization_factor * annualization_factor) / max_dd
: 0.0f;
// Clamp Calmar to prevent inf/NaN from dominating the objective
calmar = fmaxf(-100.0f, fminf(100.0f, calmar));
// Omega ratio: sum of positive returns / sum of |negative returns|
float gain_sum = 0.0f, loss_sum = 0.0f;
@@ -403,7 +408,8 @@ extern "C" __global__ void compute_backtest_metrics(
if (s_sorted[i] > 0.0f) gain_sum += s_sorted[i];
else loss_sum -= s_sorted[i];
}
float omega = (loss_sum > 1e-10f) ? gain_sum / loss_sum : 0.0f;
float omega = (loss_sum > 0.001f) ? gain_sum / loss_sum : 0.0f;
omega = fmaxf(0.0f, fminf(100.0f, omega));
metrics_out[out_base + 6] = var_95;
metrics_out[out_base + 7] = cvar_95;

319
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -10,89 +10,38 @@
* - Input: value_logits + advantage_logits per branch from both online,
* target, and online_next networks (pre-computed by cuBLAS forward pass).
* - Grid: (batch_size, 1, 1) — one block per sample.
* - Block: (256, 1, 1) — 256 threads cooperate on NUM_ATOMS=51 work.
* - Block: (256, 1, 1) — 256 threads cooperate on num_atoms work.
* - Shared memory: ~2.8 KB per block (vs 46 KB/thread in fused kernel).
*
* Does NOT modify dqn_forward_loss_kernel — that kernel remains as fallback.
*
* Requires common_device_functions.cuh prepended via NVRTC source concatenation
* (same mechanism as dqn_training_kernel.cu). The following constants must be
* injected via NVRTC #define before this source is compiled:
* NUM_ATOMS, V_MIN, V_MAX, BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE
* Requires common_device_functions.cuh prepended via build.rs source
* concatenation. All variable constants (NUM_ATOMS, V_MIN, V_MAX,
* BRANCH_*_SIZE) are passed as runtime kernel parameters.
*/
/* ── Compile-time defaults (overridden by NVRTC injection) ───────────── */
/* ── Fixed compile-time constants ─────────────────────────────────────── */
#ifndef MAX_PER_SAMPLE_CE
#define MAX_PER_SAMPLE_CE 50.0f /* Layer 1: 10x converged loss — breaks PER feedback loop */
#endif
#ifndef LABEL_SMOOTHING_EPS
#define LABEL_SMOOTHING_EPS 0.01f /* Layer 2: 1% uniform mix — prevents target collapse */
#endif
#ifndef NUM_ATOMS
#define NUM_ATOMS 51
#endif
#ifndef BRANCH_0_SIZE
#define BRANCH_0_SIZE 9
#endif
#ifndef BRANCH_1_SIZE
#define BRANCH_1_SIZE 3
#endif
#ifndef BRANCH_2_SIZE
#define BRANCH_2_SIZE 3
#endif
#ifndef V_MIN
#define V_MIN (-25.0f)
#endif
#ifndef V_MAX
#define V_MAX (25.0f)
#endif
#define DELTA_Z ((V_MAX - V_MIN) / (float)(NUM_ATOMS - 1))
/* Total atoms per branch output */
#define BRANCH_0_ATOMS (BRANCH_0_SIZE * NUM_ATOMS)
#define BRANCH_1_ATOMS (BRANCH_1_SIZE * NUM_ATOMS)
#define BRANCH_2_ATOMS (BRANCH_2_SIZE * NUM_ATOMS)
/* Largest branch output size (floats) — used to size shmem regions */
#define MAX_BRANCH_ATOMS (MAX_BRANCH_SIZE * NUM_ATOMS)
/* Fixed upper bound for register-allocated arrays — covers num_atoms up to 128 */
#define MAX_ATOMS 128
/* Largest branch output size (floats) — used to size __shared__ arrays */
#define MAX_BRANCH_SIZE 9 /* max(BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE) */
#define NUM_BRANCHES 3
/* Maximum warp count for BLOCK_THREADS = 256 */
#define NUM_WARPS 8 /* 256 / 32 */
/* ── Block-level thread count (must match launch config) ─────────────── */
/* 256 threads: 8 warps, enough to saturate the issue queue for small NUM_ATOMS.
* At NUM_ATOMS=51, each thread handles at most 1 atom in most passes. */
/* 256 threads: 8 warps, enough to saturate the issue queue for small num_atoms.
* At num_atoms=51, each thread handles at most 1 atom in most passes. */
#define BLOCK_THREADS 256
/* ── Shared memory layout (all regions in floats) ───────────────────────
*
* Region Size (floats) Purpose
* ─────────────────────────────────────────────────────────────────
* shmem_support NUM_ATOMS=51 Precomputed z_j = V_MIN + j*DELTA_Z
* shmem_val NUM_ATOMS=51 Value logits for current network/state
* shmem_adv MAX_BRANCH_ATOMS=255 Adv logits for all actions in branch
* shmem_lp NUM_ATOMS=51 log_softmax result (current/target action)
* shmem_proj NUM_ATOMS=51 Bellman-projected target distribution
* shmem_current_lp NUM_ATOMS=51 Saved current action log-probs (for CE)
* shmem_blk_reduce NUM_WARPS=8 Block-level reduction scratch
* ─────────────────────────────────────────────────────────────────
* Total floats: 51+51+255+51+51+51+8 = 518
* Total bytes: 518 * 4 = 2072 bytes ≈ 2.0 KB per block
* H100 SM: 228 KB / 2.0 KB = 114 concurrent blocks per SM
*
* Note: shmem_blk_reduce is NUM_WARPS floats and reused for all block
* reductions (max-finding, sum-finding, cross-entropy). One reduction
* at a time — never overlapping.
*/
#define SHMEM_SUPPORT_OFF 0
#define SHMEM_VAL_OFF (SHMEM_SUPPORT_OFF + NUM_ATOMS)
#define SHMEM_ADV_OFF (SHMEM_VAL_OFF + NUM_ATOMS)
#define SHMEM_LP_OFF (SHMEM_ADV_OFF + MAX_BRANCH_SIZE * NUM_ATOMS)
#define SHMEM_PROJ_OFF (SHMEM_LP_OFF + NUM_ATOMS)
#define SHMEM_CURRENT_LP_OFF (SHMEM_PROJ_OFF + NUM_ATOMS)
#define SHMEM_REDUCE_OFF (SHMEM_CURRENT_LP_OFF + NUM_ATOMS)
#define SHMEM_TOTAL_FLOATS (SHMEM_REDUCE_OFF + NUM_WARPS)
/* ── Device helpers ───────────────────────────────────────────────────── */
/**
@@ -157,35 +106,36 @@ __device__ __forceinline__ float block_reduce_sum(
}
/**
* Block-cooperative log_softmax over NUM_ATOMS elements.
* Block-cooperative log_softmax over num_atoms elements.
*
* Reads from `logits[0..NUM_ATOMS]` in shared memory.
* Writes result to `out[0..NUM_ATOMS]` in shared memory.
* Reads from `logits[0..num_atoms]` in shared memory.
* Writes result to `out[0..num_atoms]` in shared memory.
* Both `logits` and `out` may alias the same region.
* `shmem_reduce[NUM_WARPS]` is used as scratch.
* Caller must __syncthreads() before reading `out`.
*/
__device__ void block_log_softmax(
const float* __restrict__ logits, /* [NUM_ATOMS] shmem */
float* __restrict__ out, /* [NUM_ATOMS] shmem output */
const float* __restrict__ logits, /* [num_atoms] shmem */
float* __restrict__ out, /* [num_atoms] shmem output */
float* __restrict__ shmem_reduce, /* [NUM_WARPS] scratch */
int tid
int tid,
int num_atoms
) {
/* Pass 1: block-max */
float local_max = -1e30f;
for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
for (int i = tid; i < num_atoms; i += BLOCK_THREADS)
local_max = fmaxf(local_max, logits[i]);
float block_max = block_reduce_max(local_max, shmem_reduce, tid);
/* Pass 2: sum(exp(x - max)) */
float local_sum = 0.0f;
for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
for (int i = tid; i < num_atoms; i += BLOCK_THREADS)
local_sum += expf(logits[i] - block_max);
float block_sum = block_reduce_sum(local_sum, shmem_reduce, tid);
float log_sum_exp = logf(block_sum + 1e-10f);
/* Pass 3: write log(softmax) = x - max - log_sum_exp */
for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
for (int i = tid; i < num_atoms; i += BLOCK_THREADS)
out[i] = logits[i] - block_max - log_sum_exp;
__syncthreads();
}
@@ -197,13 +147,14 @@ __device__ void block_log_softmax(
* Returns the scalar expected Q (identical on all threads after call).
*/
__device__ float block_expected_q(
const float* __restrict__ log_probs, /* [NUM_ATOMS] shmem */
const float* __restrict__ support, /* [NUM_ATOMS] shmem */
const float* __restrict__ log_probs, /* [num_atoms] shmem */
const float* __restrict__ support, /* [num_atoms] shmem */
float* __restrict__ shmem_reduce, /* [NUM_WARPS] scratch */
int tid
int tid,
int num_atoms
) {
float local_sum = 0.0f;
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
local_sum += expf(log_probs[j]) * support[j];
return block_reduce_sum(local_sum, shmem_reduce, tid);
}
@@ -211,44 +162,48 @@ __device__ float block_expected_q(
/**
* Block-cooperative Bellman projection.
*
* Projects `target_probs[NUM_ATOMS]` (probabilities, not log-probs) onto the
* C51 support using T_z = clamp(r + gamma * z_j * (1 - done), V_MIN, V_MAX).
* Result accumulated into `projected[NUM_ATOMS]` (must be zero-initialized).
* Projects `target_probs[num_atoms]` (probabilities, not log-probs) onto the
* C51 support using T_z = clamp(r + gamma * z_j * (1 - done), v_min, v_max).
* Result accumulated into `projected[num_atoms]` (must be zero-initialized).
*
* With BLOCK_THREADS=256 and NUM_ATOMS=51, each thread handles at most 1 source
* atom. The per-thread `local_proj[NUM_ATOMS]` register array therefore has at
* With BLOCK_THREADS=256 and num_atoms=51, each thread handles at most 1 source
* atom. The per-thread `local_proj[MAX_ATOMS]` register array therefore has at
* most 2 non-zero entries (lower and upper neighbors for that one atom).
* The warp butterfly + block accumulate then merges them.
*/
__device__ void block_bellman_project(
const float* __restrict__ target_probs, /* [NUM_ATOMS] shmem probabilities */
const float* __restrict__ target_probs, /* [num_atoms] shmem probabilities */
float reward, float done, float gamma,
float* __restrict__ projected, /* [NUM_ATOMS] shmem, zero-initialized */
const float* __restrict__ support, /* [NUM_ATOMS] shmem */
float* __restrict__ projected, /* [num_atoms] shmem, zero-initialized */
const float* __restrict__ support, /* [num_atoms] shmem */
float* __restrict__ shmem_reduce, /* [NUM_WARPS] scratch */
int tid
int tid,
int num_atoms,
float v_min,
float v_max,
float delta_z
) {
/* Phase 1: each thread accumulates its atoms' contributions in registers.
* With BLOCK_THREADS=256 and NUM_ATOMS=51, each thread processes at most
* With BLOCK_THREADS=256 and num_atoms=51, each thread processes at most
* one source atom (tid 0..50) or zero atoms (tid 51..255). */
float local_proj[NUM_ATOMS];
for (int k = 0; k < NUM_ATOMS; k++) local_proj[k] = 0.0f;
float local_proj[MAX_ATOMS];
for (int k = 0; k < num_atoms; k++) local_proj[k] = 0.0f;
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
float z_j = support[j];
float t_z = reward + gamma * z_j * (1.0f - done);
t_z = fminf(fmaxf(t_z, V_MIN), V_MAX);
t_z = fminf(fmaxf(t_z, v_min), v_max);
float b_pos = (t_z - V_MIN) / DELTA_Z;
float b_pos = (t_z - v_min) / delta_z;
/* Keep b_pos away from exact upper boundary to prevent phantom
* upper==lower when ceilf(N-1) == floorf(N-1) = N-1.
* Without this, probability mass piles on boundary atoms. */
b_pos = fminf(b_pos, (float)(NUM_ATOMS - 1) - 1e-6f);
b_pos = fminf(b_pos, (float)(num_atoms - 1) - 1e-6f);
b_pos = fmaxf(b_pos, 1e-6f);
int lower = (int)floorf(b_pos);
int upper = (int)ceilf(b_pos);
lower = max(min(lower, NUM_ATOMS - 1), 0);
upper = max(min(upper, NUM_ATOMS - 1), 0);
lower = max(min(lower, num_atoms - 1), 0);
upper = max(min(upper, num_atoms - 1), 0);
float frac = b_pos - floorf(b_pos);
float p_j = target_probs[j];
@@ -264,7 +219,7 @@ __device__ void block_bellman_project(
int lane = tid & 31;
int warp = tid >> 5;
for (int k = 0; k < NUM_ATOMS; k++) {
for (int k = 0; k < num_atoms; k++) {
/* Warp butterfly sum for atom k */
float val = local_proj[k];
for (int offset = 16; offset > 0; offset >>= 1)
@@ -310,22 +265,22 @@ __device__ void block_bellman_project(
extern "C" __global__ void c51_loss_batched(
/* ── Online network outputs on current STATES ─────────────────── */
const float* __restrict__ on_value_logits, /* [B, NUM_ATOMS] */
const float* __restrict__ on_adv_logits_b0, /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
const float* __restrict__ on_adv_logits_b1, /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
const float* __restrict__ on_adv_logits_b2, /* [B, BRANCH_2_SIZE * NUM_ATOMS] */
const float* __restrict__ on_value_logits, /* [B, num_atoms] */
const float* __restrict__ on_adv_logits_b0, /* [B, b0_size * num_atoms] */
const float* __restrict__ on_adv_logits_b1, /* [B, b1_size * num_atoms] */
const float* __restrict__ on_adv_logits_b2, /* [B, b2_size * num_atoms] */
/* ── Target network outputs on NEXT_STATES ────────────────────── */
const float* __restrict__ tg_value_logits, /* [B, NUM_ATOMS] */
const float* __restrict__ tg_adv_logits_b0, /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
const float* __restrict__ tg_adv_logits_b1, /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
const float* __restrict__ tg_adv_logits_b2, /* [B, BRANCH_2_SIZE * NUM_ATOMS] */
const float* __restrict__ tg_value_logits, /* [B, num_atoms] */
const float* __restrict__ tg_adv_logits_b0, /* [B, b0_size * num_atoms] */
const float* __restrict__ tg_adv_logits_b1, /* [B, b1_size * num_atoms] */
const float* __restrict__ tg_adv_logits_b2, /* [B, b2_size * num_atoms] */
/* ── Online network outputs on NEXT_STATES (Double DQN selector) */
const float* __restrict__ on_next_value_logits, /* [B, NUM_ATOMS] */
const float* __restrict__ on_next_adv_logits_b0, /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
const float* __restrict__ on_next_adv_logits_b1, /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
const float* __restrict__ on_next_adv_logits_b2, /* [B, BRANCH_2_SIZE * NUM_ATOMS] */
const float* __restrict__ on_next_value_logits, /* [B, num_atoms] */
const float* __restrict__ on_next_adv_logits_b0, /* [B, b0_size * num_atoms] */
const float* __restrict__ on_next_adv_logits_b1, /* [B, b1_size * num_atoms] */
const float* __restrict__ on_next_adv_logits_b2, /* [B, b2_size * num_atoms] */
/* ── Batch data ───────────────────────────────────────────────── */
const int* __restrict__ actions, /* [B] factored action indices 0-44 */
@@ -339,69 +294,92 @@ extern "C" __global__ void c51_loss_batched(
float* __restrict__ total_loss, /* [1] batch mean loss (atomicAdd) */
/* ── Saved tensors for backward pass ─────────────────────────── */
float* __restrict__ save_current_lp, /* [B, NUM_BRANCHES, NUM_ATOMS] */
float* __restrict__ save_projected, /* [B, NUM_BRANCHES, NUM_ATOMS] */
float* __restrict__ save_current_lp, /* [B, NUM_BRANCHES, num_atoms] */
float* __restrict__ save_projected, /* [B, NUM_BRANCHES, num_atoms] */
/* ── Config ───────────────────────────────────────────────────── */
float gamma,
int batch_size
int batch_size,
int num_atoms,
float v_min,
float v_max,
int b0_size,
int b1_size,
int b2_size
) {
extern __shared__ float shmem[];
/* ── Compute runtime shared memory offsets ────────────────────── */
const int max_branch = max(max(b0_size, b1_size), b2_size);
int off_support = 0;
int off_val = off_support + num_atoms;
int off_adv = off_val + num_atoms;
int off_lp = off_adv + max_branch * num_atoms;
int off_proj = off_lp + num_atoms;
int off_current_lp = off_proj + num_atoms;
int off_reduce = off_current_lp + num_atoms;
/* Named pointers into the flat shared memory region */
float* shmem_support = shmem + SHMEM_SUPPORT_OFF; /* [NUM_ATOMS] z_j */
float* shmem_val = shmem + SHMEM_VAL_OFF; /* [NUM_ATOMS] value logits */
float* shmem_adv = shmem + SHMEM_ADV_OFF; /* [MAX_BRANCH_SIZE*NUM_ATOMS] adv logits */
float* shmem_lp = shmem + SHMEM_LP_OFF; /* [NUM_ATOMS] log-prob result */
float* shmem_proj = shmem + SHMEM_PROJ_OFF; /* [NUM_ATOMS] Bellman projection */
float* shmem_current_lp = shmem + SHMEM_CURRENT_LP_OFF; /* [NUM_ATOMS] taken action lp */
float* shmem_reduce = shmem + SHMEM_REDUCE_OFF; /* [NUM_WARPS] reduction scratch */
float* shmem_support = shmem + off_support; /* [num_atoms] z_j */
float* shmem_val = shmem + off_val; /* [num_atoms] value logits */
float* shmem_adv = shmem + off_adv; /* [max_branch*num_atoms] adv logits */
float* shmem_lp = shmem + off_lp; /* [num_atoms] log-prob result */
float* shmem_proj = shmem + off_proj; /* [num_atoms] Bellman projection */
float* shmem_current_lp = shmem + off_current_lp; /* [num_atoms] taken action lp */
float* shmem_reduce = shmem + off_reduce; /* [NUM_WARPS] reduction scratch */
int sample_id = blockIdx.x;
int tid = threadIdx.x;
if (sample_id >= batch_size) return;
/* ── Compute runtime derived values ──────────────────────────── */
float delta_z = (v_max - v_min) / (float)(num_atoms - 1);
int b0_atoms = b0_size * num_atoms;
int b1_atoms = b1_size * num_atoms;
int b2_atoms = b2_size * num_atoms;
/* ── Precompute C51 support atoms once per block ─────────────── */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
shmem_support[j] = V_MIN + (float)j * DELTA_Z;
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_support[j] = v_min + (float)j * delta_z;
__syncthreads();
/* ── Decode factored action into per-branch indices ──────────── */
int factored_action = actions[sample_id];
/* Clamp to valid range — guards against uninitialized replay buffer entries */
int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE;
int max_action = b0_size * b1_size * b2_size;
if (factored_action < 0 || factored_action >= max_action)
factored_action = 0;
int a0 = factored_action / (BRANCH_1_SIZE * BRANCH_2_SIZE);
int a1 = (factored_action / BRANCH_2_SIZE) % BRANCH_1_SIZE;
int a2 = factored_action % BRANCH_2_SIZE;
int a0 = factored_action / (b1_size * b2_size);
int a1 = (factored_action / b2_size) % b1_size;
int a2 = factored_action % b2_size;
/* Per-branch metadata (compile-time sizes, runtime-selected values) */
/* Per-branch metadata (runtime sizes, runtime-selected values) */
const int branch_action[NUM_BRANCHES] = { a0, a1, a2 };
const int branch_sizes[NUM_BRANCHES] = { BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE };
const int branch_sizes[NUM_BRANCHES] = { b0_size, b1_size, b2_size };
/* Per-sample row pointers into each branch's advantage logit tensors */
const float* on_adv_row[NUM_BRANCHES] = {
on_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
on_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
on_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
on_adv_logits_b0 + (long long)sample_id * b0_atoms,
on_adv_logits_b1 + (long long)sample_id * b1_atoms,
on_adv_logits_b2 + (long long)sample_id * b2_atoms
};
const float* tg_adv_row[NUM_BRANCHES] = {
tg_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
tg_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
tg_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
tg_adv_logits_b0 + (long long)sample_id * b0_atoms,
tg_adv_logits_b1 + (long long)sample_id * b1_atoms,
tg_adv_logits_b2 + (long long)sample_id * b2_atoms
};
const float* on_next_adv_row[NUM_BRANCHES] = {
on_next_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
on_next_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
on_next_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
on_next_adv_logits_b0 + (long long)sample_id * b0_atoms,
on_next_adv_logits_b1 + (long long)sample_id * b1_atoms,
on_next_adv_logits_b2 + (long long)sample_id * b2_atoms
};
/* Value logit row pointers (same NUM_ATOMS for all branches) */
const float* on_val_row = on_value_logits + (long long)sample_id * NUM_ATOMS;
const float* tg_val_row = tg_value_logits + (long long)sample_id * NUM_ATOMS;
const float* on_next_val_row = on_next_value_logits + (long long)sample_id * NUM_ATOMS;
/* Value logit row pointers (same num_atoms for all branches) */
const float* on_val_row = on_value_logits + (long long)sample_id * num_atoms;
const float* tg_val_row = tg_value_logits + (long long)sample_id * num_atoms;
const float* on_next_val_row = on_next_value_logits + (long long)sample_id * num_atoms;
float reward = rewards[sample_id];
float done = dones[sample_id];
@@ -416,14 +394,14 @@ extern "C" __global__ void c51_loss_batched(
for (int d = 0; d < NUM_BRANCHES; d++) {
int n_d = branch_sizes[d];
int a_d = branch_action[d];
int n_atoms = n_d * NUM_ATOMS; /* total adv logit elements for this branch */
int n_atoms = n_d * num_atoms; /* total adv logit elements for this branch */
/* Global memory offset for save tensors: [sample, branch, atom] */
long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * NUM_ATOMS;
long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * num_atoms;
/* ── Helper: compute dueling logits for a single action into shmem_lp
* shmem_val must already hold value logits [NUM_ATOMS]
* shmem_adv must already hold all adv logits [n_d * NUM_ATOMS]
* shmem_val must already hold value logits [num_atoms]
* shmem_adv must already hold all adv logits [n_d * num_atoms]
* Result: dueling(a_sel)[j] = V[j] + A[a_sel,j] - mean_a(A[*,j])
* written into shmem_lp[j] */
/* (Inlined below to avoid device-function overhead for this hot path) */
@@ -434,7 +412,7 @@ extern "C" __global__ void c51_loss_batched(
* ═══════════════════════════════════════════════════════════ */
/* Load value logits */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_val[j] = on_val_row[j];
__syncthreads();
@@ -444,25 +422,25 @@ extern "C" __global__ void c51_loss_batched(
__syncthreads();
/* Dueling for action a_d: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j]) */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
float a_mean = 0.0f;
for (int a = 0; a < n_d; a++)
a_mean += shmem_adv[a * NUM_ATOMS + j];
a_mean += shmem_adv[a * num_atoms + j];
a_mean /= (float)n_d;
shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * NUM_ATOMS + j] - a_mean;
shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * num_atoms + j] - a_mean;
}
__syncthreads();
/* log_softmax → result back into shmem_lp */
block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid);
/* shmem_lp[0..NUM_ATOMS] = log_probs for taken action a_d */
block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid, num_atoms);
/* shmem_lp[0..num_atoms] = log_probs for taken action a_d */
/* Copy into shmem_current_lp for use in cross-entropy (step e) */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_current_lp[j] = shmem_lp[j];
/* Save for backward kernel */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
save_current_lp[save_off + j] = shmem_lp[j];
__syncthreads();
@@ -474,7 +452,7 @@ extern "C" __global__ void c51_loss_batched(
* ═══════════════════════════════════════════════════════════ */
/* Load online_next value and advantage logits */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_val[j] = on_next_val_row[j];
__syncthreads();
@@ -484,10 +462,10 @@ extern "C" __global__ void c51_loss_batched(
/* Compute average advantage across actions for dueling (same for all a) */
/* shmem_proj temporarily holds per-atom mean advantage — cleared next step */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
float a_mean = 0.0f;
for (int a = 0; a < n_d; a++)
a_mean += shmem_adv[a * NUM_ATOMS + j];
a_mean += shmem_adv[a * num_atoms + j];
a_mean /= (float)n_d;
shmem_proj[j] = a_mean; /* reuse proj region for mean advantage */
}
@@ -501,15 +479,15 @@ extern "C" __global__ void c51_loss_batched(
for (int a = 0; a < n_d; a++) {
/* Dueling logits for action a: V[j] + A[a,j] - mean_j(A[*,j]) */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
shmem_lp[j] = shmem_val[j] + shmem_adv[a * NUM_ATOMS + j] - shmem_proj[j];
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_lp[j] = shmem_val[j] + shmem_adv[a * num_atoms + j] - shmem_proj[j];
__syncthreads();
/* log_softmax in-place on shmem_lp */
block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid);
block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid, num_atoms);
/* Expected Q for action a */
float eq = block_expected_q(shmem_lp, shmem_support, shmem_reduce, tid);
float eq = block_expected_q(shmem_lp, shmem_support, shmem_reduce, tid, num_atoms);
if (tid == 0) eq_per_action[a] = eq;
__syncthreads();
}
@@ -530,7 +508,7 @@ extern "C" __global__ void c51_loss_batched(
* ═══════════════════════════════════════════════════════════ */
/* Load target value and advantage logits */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_val[j] = tg_val_row[j];
__syncthreads();
@@ -539,26 +517,26 @@ extern "C" __global__ void c51_loss_batched(
__syncthreads();
/* Per-atom mean advantage for target network */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
float a_mean = 0.0f;
for (int a = 0; a < n_d; a++)
a_mean += shmem_adv[a * NUM_ATOMS + j];
a_mean += shmem_adv[a * num_atoms + j];
a_mean /= (float)n_d;
shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * NUM_ATOMS + j] - a_mean;
shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * num_atoms + j] - a_mean;
}
__syncthreads();
/* log_softmax → log-probs for target's best_next action */
block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid);
block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid, num_atoms);
/* Convert to probabilities for Bellman projection */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_proj[j] = expf(shmem_lp[j]);
__syncthreads();
/* ═══════════════════════════════════════════════════════════
* STEP d: Bellman projection
* T_z_j = clamp(r + gamma * z_j * (1-done), V_MIN, V_MAX)
* T_z_j = clamp(r + gamma * z_j * (1-done), v_min, v_max)
* Project shmem_proj (target probs) onto shmem_proj (output).
*
* The projection writes back to shmem_proj in-place (the input
@@ -567,28 +545,29 @@ extern "C" __global__ void c51_loss_batched(
/* Zero-initialize projected output — block_bellman_project accumulates */
/* We use shmem_lp as the output projected array (shmem_proj still holds probs) */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_lp[j] = 0.0f;
__syncthreads();
block_bellman_project(shmem_proj, reward, done, gamma,
shmem_lp, shmem_support, shmem_reduce, tid);
shmem_lp, shmem_support, shmem_reduce, tid,
num_atoms, v_min, v_max, delta_z);
__syncthreads();
/* Layer 2: Label smoothing — mix ε uniform into projected target.
* Prevents any atom from having zero probability → no log(0) in CE. */
{
const float ls_uniform = LABEL_SMOOTHING_EPS / (float)NUM_ATOMS;
const float ls_uniform = LABEL_SMOOTHING_EPS / (float)num_atoms;
const float ls_scale = 1.0f - LABEL_SMOOTHING_EPS;
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_lp[j] = ls_scale * shmem_lp[j] + ls_uniform;
__syncthreads();
}
/* shmem_lp[0..NUM_ATOMS] = smoothed projected target distribution */
/* shmem_lp[0..num_atoms] = smoothed projected target distribution */
/* Save projected for backward kernel */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
save_projected[save_off + j] = shmem_lp[j];
__syncthreads();
@@ -598,7 +577,7 @@ extern "C" __global__ void c51_loss_batched(
* ═══════════════════════════════════════════════════════════ */
float local_ce = 0.0f;
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
local_ce -= shmem_lp[j] * shmem_current_lp[j];
float branch_ce = block_reduce_sum(local_ce, shmem_reduce, tid);

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

@@ -54,6 +54,10 @@
/* ------------------------------------------------------------------ */
#define PORTFOLIO_STRIDE 20
/* Shared trade physics (action decoding, hold enforcement, tx costs,
* capital floor). Single source of truth — also used by backtest_env_kernel. */
#include "trade_physics.cuh"
/* ------------------------------------------------------------------ */
/* Compile-time constants — overridable via NVRTC #define injection. */
/* ------------------------------------------------------------------ */
@@ -378,12 +382,16 @@ extern "C" __global__ void experience_action_select(
/* Hold enforcement: during minimum hold period, force current exposure.
* Prevents both greedy and random exploration from changing exposure.
* Portfolio state: ps[0] = position, ps[10] = hold_time. */
int ps_base = i * 20; /* PORTFOLIO_STRIDE = 20 */
float hold_time_val = portfolio_states[ps_base + 10];
float cur_position = portfolio_states[ps_base + 0];
int in_hold = (hold_time_val > 0.0f && hold_time_val < (float)min_hold_bars
&& min_hold_bars > 0);
* Portfolio state: ps[0] = position, ps[10] = hold_time.
* When portfolio_states is NULL (backtest evaluator), skip hold enforcement. */
int in_hold = 0;
float cur_position = 0.0f;
if (min_hold_bars > 0 && portfolio_states != NULL) {
int ps_base = i * 20; /* PORTFOLIO_STRIDE = 20 */
float hold_time_val = portfolio_states[ps_base + 10];
cur_position = portfolio_states[ps_base + 0];
in_hold = (hold_time_val > 0.0f && hold_time_val < (float)min_hold_bars);
}
if (in_hold) {
/* Compute current exposure index from position */
@@ -667,8 +675,17 @@ extern "C" __global__ void experience_env_step(
} /* end Kelly block */
} /* end Kelly scope */
/* Final NaN/Inf guard on target_position.
* If exposure_idx_to_fraction produced NaN, zero the position. */
/* Margin-aware position cap (shared: trade_physics.cuh).
* Prevents overleveraging when equity is depleted — a depleted account
* can't hold the same position size as a full account.
* margin ≈ 6% of notional (CME ES initial margin ~$15K per contract). */
{
float margin_per_contract = raw_close * 50.0f * 0.06f;
float portfolio_val = ps[2];
target_position = apply_margin_cap(target_position, portfolio_val, margin_per_contract);
}
/* Final NaN/Inf guard on target_position. */
if (isnan(target_position) || isinf(target_position)) {
target_position = 0.0f;
}

187
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

@@ -33,7 +33,7 @@ use std::mem::ManuallyDrop;
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaGraph, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use tracing::{info, warn};
use tracing::info;
use ml_core::nvtx::NvtxRange;
use crate::MLError;
@@ -854,6 +854,15 @@ impl GpuBacktestEvaluator {
) -> Result<Vec<WindowMetrics>, MLError> {
let _nvtx = NvtxRange::new("backtest_evaluate_dqn_graphed");
// Sync stream and drain any CUDA errors before evaluator starts.
// If training left deferred errors, log them explicitly.
if let Err(sync_err) = self.stream.synchronize() {
eprintln!("EVALUATOR PRE-SYNC: training left stale CUDA error: {sync_err}");
}
if let Err(ctx_err) = self.stream.context().check_err() {
eprintln!("EVALUATOR PRE-SYNC: context error drained: {ctx_err}");
}
// Lazy-initialise cuBLAS context + buffers
self.ensure_cublas_ready(dqn_cfg)?;
@@ -869,97 +878,19 @@ impl GpuBacktestEvaluator {
super::batched_forward::f32_weight_ptrs(flat, &param_sizes, &self.stream)
};
// ── Replay cached graph if available ────────────────────────────────
if let Some(ref graph) = self.dqn_graph {
graph.0.launch().map_err(|e| {
MLError::ModelError(format!("CUDA graph replay failed: {e}"))
})?;
return self.launch_metrics_and_download();
}
// ── First call: capture the step loop into a CUDA graph ─────────────
info!(
"GpuBacktestEvaluator: capturing CUDA graph for {} steps × {} windows",
self.max_len, self.n_windows,
);
// Synchronize the stream before capture to ensure all pending work
// is complete. CUDA Graph capture requires no in-flight work.
self.stream.synchronize().map_err(|e| {
MLError::ModelError(format!("stream sync before capture: {e}"))
})?;
// Disable cudarc's automatic event tracking during graph capture.
// After fork(), cudarc injects cuStreamWaitEvent/cuEventRecord into
// every buffer access, creating cross-stream event references that
// invalidate CUDA graph capture. Safe: single-stream, all work synced.
// SAFETY: single-stream capture, all work synchronized, re-enabled below.
unsafe { self.stream.context().disable_event_tracking(); }
// Begin stream capture in THREAD_LOCAL mode (only captures work
// submitted from this thread on this stream — safe for single-stream use).
let capture_result = self.stream.begin_capture(
cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
);
if let Err(e) = capture_result {
unsafe { self.stream.context().enable_event_tracking(); }
warn!(
"CUDA graph capture begin failed ({e}), falling back to non-graphed evaluate_dqn"
);
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
.and_then(|()| self.launch_metrics_and_download());
}
// Submit the step loop to the stream (captured, NOT executed yet).
let capture_err = self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg);
// End capture — get the graph regardless of whether submit succeeded,
// because the stream is in capture mode and must be ended.
let graph_result = self.stream.end_capture(
cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
);
// Re-enable event tracking now that capture is complete.
// SAFETY: restores normal multi-stream synchronization.
unsafe { self.stream.context().enable_event_tracking(); }
// If the step loop submission had an error, propagate it
if let Err(e) = capture_err {
warn!("CUDA graph capture step loop failed ({e}), falling back to non-graphed path");
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
.and_then(|()| self.launch_metrics_and_download());
}
// Process graph instantiation result
let graph = match graph_result {
Ok(Some(g)) => g,
Ok(None) => {
warn!("CUDA graph capture returned empty graph, falling back to non-graphed path");
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
.and_then(|()| self.launch_metrics_and_download());
}
Err(e) => {
warn!(
"CUDA graph instantiation failed ({e}), falling back to non-graphed path"
);
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
.and_then(|()| self.launch_metrics_and_download());
}
};
// Launch the graph (this actually executes the captured step loop)
if let Err(e) = graph.launch() {
warn!("CUDA graph launch failed ({e}), falling back to non-graphed path");
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
.and_then(|()| self.launch_metrics_and_download());
}
info!("GpuBacktestEvaluator: CUDA graph captured and launched successfully");
// Cache the graph for subsequent calls
self.dqn_graph = Some(SendSyncGraph(graph));
// Metrics are NOT in the graph — run them normally
// Direct kernel launches (no CUDA Graph for the backtest step loop).
//
// The step loop captures 500K+ kernel launches (10K steps × 10 windows
// × ~5 ops/step). cuGraphLaunch SIGSEGVs on graphs this large due to
// CUDA driver-internal memory limits. Direct launches are fine because:
// - Eval runs ONCE per hyperopt trial (not per training step)
// - Training is 99%+ of hyperopt wall time; eval is negligible
// - cuBLAS forward passes dominate eval throughput, not launch overhead
//
// The training pipeline (graph_forward + graph_adam in GpuDqnTrainer)
// still uses CUDA Graphs — those capture ~20 operations per replay,
// well within driver limits.
self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)?;
self.launch_metrics_and_download()
}
@@ -1057,10 +988,12 @@ impl GpuBacktestEvaluator {
let state_row_bytes = n * state_dim * std::mem::size_of::<f32>();
let action_row_bytes = n * std::mem::size_of::<i32>();
let total_chunks = (self.max_len + DQN_BACKTEST_CHUNK_SIZE - 1) / DQN_BACKTEST_CHUNK_SIZE;
for chunk_start in (0..self.max_len).step_by(DQN_BACKTEST_CHUNK_SIZE) {
let chunk_end = (chunk_start + DQN_BACKTEST_CHUNK_SIZE).min(self.max_len);
let chunk_len = chunk_end - chunk_start;
let batch = n * chunk_len; // total rows for this chunk's forward pass
let chunk_idx = chunk_start / DQN_BACKTEST_CHUNK_SIZE;
// ── Phase 1: Gather states for all steps in the chunk ────────────
//
@@ -1088,6 +1021,13 @@ impl GpuBacktestEvaluator {
}
}
// Sync after Phase 1 gather — catch any gather/DtoD errors
if chunk_idx == 0 {
self.stream.synchronize().map_err(|e| MLError::ModelError(format!(
"backtest chunk {chunk_idx}: sync after gather phase: {e}"
)))?;
}
// ── Phase 2: Single cuBLAS forward on [batch, state_dim] ─────────
//
// For the last chunk, batch may be < n * CHUNK_SIZE. The cuBLAS forward
@@ -1107,6 +1047,12 @@ impl GpuBacktestEvaluator {
ch_v_logits, ch_b_logits,
)?;
if chunk_idx == 0 {
self.stream.synchronize().map_err(|e| MLError::ModelError(format!(
"backtest chunk 0: sync after cuBLAS forward: {e}"
)))?;
}
// ── Phase 3: C51 expected Q-values on chunked output ─────────────
let batch_i32 = batch as i32;
let blocks_256 = ((batch + 255) / 256) as u32;
@@ -1135,15 +1081,15 @@ impl GpuBacktestEvaluator {
)))?;
}
if chunk_idx == 0 {
self.stream.synchronize().map_err(|e| MLError::ModelError(format!(
"backtest chunk 0: sync after expected_q: {e}"
)))?;
}
// ── Phase 4: Greedy action selection on chunked Q-values ─────────
// Kernel signature: (q_values, out_actions, rng_states, out_q_gaps, epsilon, N,
// b0, b1, b2, q_gap_threshold, portfolio_states, min_hold_bars, max_position)
// Hold enforcement handled by backtest_env_step (Layer 2), not action masking.
// The greedy argmax (epsilon=0) doesn't need masking — any hold-violating
// action is overridden by the env_step kernel. This avoids the stride-8/stride-20
// portfolio layout incompatibility between backtest and training buffers.
let null_portfolio: u64 = 0; // NULL — Layer 2 in env_step handles holds
let eval_min_hold: i32 = 0; // 0 = disabled here (enforced in env_step instead)
let null_portfolio: u64 = 0;
let eval_min_hold: i32 = 0;
let eval_max_pos: f32 = 0.0;
unsafe {
self.stream
@@ -1158,15 +1104,21 @@ impl GpuBacktestEvaluator {
.arg(&b1)
.arg(&b2)
.arg(&self.q_gap_threshold)
.arg(&null_portfolio) // portfolio_states (NULL — no hold masking in eval)
.arg(&eval_min_hold) // min_hold_bars (0 = disabled)
.arg(&eval_max_pos) // max_position (unused when disabled)
.arg(&null_portfolio)
.arg(&eval_min_hold)
.arg(&eval_max_pos)
.launch(launch_cfg)
.map_err(|e| MLError::ModelError(format!(
"chunked action_select chunk_start={chunk_start}: {e}"
)))?;
}
if chunk_idx == 0 {
self.stream.synchronize().map_err(|e| MLError::ModelError(format!(
"backtest chunk 0: sync after action_select: {e}"
)))?;
}
// ── Phase 5: Per-step env_step (sequential, stateful) ────────────
//
// For each step in the chunk, extract the n_windows actions from the
@@ -1194,6 +1146,13 @@ impl GpuBacktestEvaluator {
self.launch_env_step(step)?;
}
// Per-chunk sync: catch deferred CUDA errors from this chunk's
// kernel launches. Without this, an error in chunk N surfaces as
// a confusing failure in chunk N+1 or later.
self.stream.synchronize().map_err(|e| MLError::ModelError(format!(
"backtest chunk {chunk_idx}/{total_chunks} (steps {chunk_start}..{chunk_end}) sync: {e}"
)))?;
}
Ok(())
@@ -1428,17 +1387,6 @@ impl GpuBacktestEvaluator {
.load_function("experience_action_select")
.map_err(|e| MLError::ModelError(format!("experience_action_select load: {e}")))?;
// RNG states (initialised with random seeds for greedy mode — epsilon=0 means
// the RNG values are never actually used, but the kernel reads them).
let rng_seeds: Vec<u32> = (0..n).map(|_| fastrand::u32(..)).collect();
let rng_states = self.stream.clone_htod(&rng_seeds)
.map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?;
// Q-gap output buffer — kernel writes per-episode Q-gaps (unused by backtest
// but required by the experience_action_select kernel signature).
let q_gaps_buf = self.stream.alloc_zeros::<f32>(n)
.map_err(|e| MLError::ModelError(format!("alloc q_gaps_buf: {e}")))?;
// ── Chunked cuBLAS forward (batch_size = n_windows * CHUNK_SIZE) ─────
//
// Separate CublasForward + scratch buffers to amortise kernel-launch
@@ -1446,6 +1394,17 @@ impl GpuBacktestEvaluator {
let chunk = DQN_BACKTEST_CHUNK_SIZE;
let cn = n * chunk; // chunked batch dimension
// RNG states: must be large enough for chunked batch (n * CHUNK_SIZE),
// not just n_windows, because chunked action_select launches with batch = n * chunk.
let rng_seeds: Vec<u32> = (0..cn).map(|_| fastrand::u32(..)).collect();
let rng_states = self.stream.clone_htod(&rng_seeds)
.map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?;
// Q-gap output buffer: same sizing as rng_states — chunked action_select
// writes batch = n * chunk_len entries per call.
let q_gaps_buf = self.stream.alloc_zeros::<f32>(cn)
.map_err(|e| MLError::ModelError(format!("alloc q_gaps_buf: {e}")))?;
let chunked_cublas = CublasForward::new(
&self.stream,
cn,

87
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -3025,9 +3025,14 @@ impl GpuDqnTrainer {
// The experience collector pre-computes R_n = sum(gamma^i * r_i).
let gamma = self.config.gamma.powi(self.config.n_steps as i32);
let batch_i32 = b as i32;
let na_i32 = na as i32;
let v_min_f = self.config.v_min;
let v_max_f = self.config.v_max;
let b0_i32 = b0 as i32;
let b1_i32 = b1 as i32;
let b2_i32 = b2 as i32;
// Shared memory: SHMEM_TOTAL_FLOATS * 4 bytes
// 51+51+255+51+51+51+8 = 518 floats = 2072 bytes
// Shared memory: support + val + adv + lp + proj + current_lp + reduce
let max_branch = b0.max(b1).max(b2);
let shmem_floats = na + na + max_branch * na + na + na + na + 8;
let shmem_bytes = (shmem_floats * std::mem::size_of::<f32>()) as u32;
@@ -3062,9 +3067,15 @@ impl GpuDqnTrainer {
// ── Saved for backward (2) ──
.arg(&self.save_current_lp)
.arg(&self.save_projected)
// ── Config (2) ──
// ── Config (8) ──
.arg(&gamma)
.arg(&batch_i32)
.arg(&na_i32)
.arg(&v_min_f)
.arg(&v_max_f)
.arg(&b0_i32)
.arg(&b1_i32)
.arg(&b2_i32)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (256, 1, 1),
@@ -3159,6 +3170,12 @@ impl GpuDqnTrainer {
// N-step returns: use gamma^n for the Bellman projection.
let gamma = self.config.gamma.powi(self.config.n_steps as i32);
let batch_i32 = b as i32;
let na_i32 = na as i32;
let v_min_f = self.config.v_min;
let v_max_f = self.config.v_max;
let b0_i32 = b0 as i32;
let b1_i32 = b1 as i32;
let b2_i32 = b2 as i32;
// Shared memory: same layout as C51 kernel (support + val + adv + lp + proj + current_lp + reduce)
let max_branch = b0.max(b1).max(b2);
@@ -3196,9 +3213,15 @@ impl GpuDqnTrainer {
// save_projected = per-branch E[Q] values (3 floats per sample per branch)
.arg(&self.save_current_lp)
.arg(&self.save_projected)
// ── Config (2) ──
// ── Config (8) ──
.arg(&gamma)
.arg(&batch_i32)
.arg(&na_i32)
.arg(&v_min_f)
.arg(&v_max_f)
.arg(&b0_i32)
.arg(&b1_i32)
.arg(&b2_i32)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (256, 1, 1),
@@ -3775,35 +3798,16 @@ fn compile_per_update_kernel(stream: &Arc<CudaStream>) -> Result<CudaFunction, M
/// from cuBLAS, computing the C51 cross-entropy loss with Bellman projection.
/// Grid: (batch_size, 1, 1), Block: (256, 1, 1), Shmem: ~2 KB/block.
///
/// The kernel uses compile-time defaults (NUM_ATOMS=51, BRANCH_0_SIZE=9, etc.)
/// which match the production configuration. These are hardcoded at build time.
/// All variable constants (num_atoms, v_min, v_max, branch sizes) are passed
/// as runtime kernel parameters — no NVRTC compilation needed.
fn compile_c51_loss_kernel(
stream: &Arc<CudaStream>,
config: &GpuDqnTrainConfig,
_config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
// c51_loss uses NUM_ATOMS for BOTH loop bounds AND array sizing —
// MUST be runtime-compiled with the correct #define injected.
let defines = format!(
"#define NUM_ATOMS {na}\n\
#define V_MIN ({v_min:.1}f)\n\
#define V_MAX ({v_max:.1}f)\n\
#define BRANCH_0_SIZE {b0}\n\
#define BRANCH_1_SIZE {b1}\n\
#define BRANCH_2_SIZE {b2}\n",
na = config.num_atoms,
v_min = config.v_min,
v_max = config.v_max,
b0 = config.branch_0_size,
b1 = config.branch_1_size,
b2 = config.branch_2_size,
);
let kernel_src = include_str!("c51_loss_kernel.cu");
let full_source = format!("{defines}\n{kernel_src}");
let context = stream.context();
let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
.map_err(|e| MLError::ModelError(format!("c51_loss_kernel compilation: {e}")))?;
let ptx = Ptx::from_binary(C51_LOSS_CUBIN.to_vec());
let module = context.load_module(ptx)
.map_err(|e| MLError::ModelError(format!("c51_loss module load: {e}")))?;
.map_err(|e| MLError::ModelError(format!("c51_loss cubin load: {e}")))?;
module.load_function("c51_loss_batched")
.map_err(|e| MLError::ModelError(format!("c51_loss_batched load: {e}")))
}
@@ -3823,31 +3827,16 @@ fn compile_c51_grad_kernel(
.map_err(|e| MLError::ModelError(format!("c51_grad_kernel load: {e}")))
}
/// Compile MSE loss kernel at runtime with correct NUM_ATOMS injected.
/// MSE uses NUM_ATOMS for loop bounds — MUST match runtime config.
/// Load the MSE loss kernel from precompiled cubin.
///
/// All variable constants (num_atoms, v_min, v_max, branch sizes) are passed
/// as runtime kernel parameters — no NVRTC compilation needed.
fn compile_mse_loss_kernel(
stream: &Arc<CudaStream>,
config: &GpuDqnTrainConfig,
_config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
let defines = format!(
"#define NUM_ATOMS {na}\n\
#define V_MIN ({v_min:.1}f)\n\
#define V_MAX ({v_max:.1}f)\n\
#define BRANCH_0_SIZE {b0}\n\
#define BRANCH_1_SIZE {b1}\n\
#define BRANCH_2_SIZE {b2}\n",
na = config.num_atoms,
v_min = config.v_min,
v_max = config.v_max,
b0 = config.branch_0_size,
b1 = config.branch_1_size,
b2 = config.branch_2_size,
);
let kernel_src = include_str!("mse_loss_kernel.cu");
let full_source = format!("{defines}\n{kernel_src}");
let context = stream.context();
let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
.map_err(|e| MLError::ModelError(format!("mse_loss_kernel compilation: {e}")))?;
let ptx = Ptx::from_binary(MSE_LOSS_CUBIN.to_vec());
let module = context.load_module(ptx)
.map_err(|e| MLError::ModelError(format!("mse_loss cubin load: {e}")))?;
module.load_function("mse_loss_batched")

4
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -100,9 +100,7 @@ pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u3
/// so any change to the kernel source or network dimensions invalidates the cache.
///
/// The fused experience collector kernel (4490 lines, branching+C51+NoisyNets)
// Re-export precompilation from ml-core so existing callers
// (`crate::cuda_pipeline::compile_ptx_for_device`) keep working.
pub use ml_core::cuda_compile::compile_ptx_for_device;
// compile_ptx_for_device re-export removed: all kernels now use precompiled cubins.
/// Clone a `CudaSlice<f32>` via DtoD memcpy.
///

196
crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
View File

@@ -12,47 +12,15 @@
* Grid = (batch_size, 1, 1)
* Block = (BLOCK_THREADS=256, 1, 1)
*
* Requires #define: NUM_ATOMS, V_MIN, V_MAX, BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE
* All variable constants (NUM_ATOMS, V_MIN, V_MAX, BRANCH_*_SIZE) are
* passed as runtime kernel parameters.
*/
#ifndef NUM_ATOMS
#define NUM_ATOMS 51
#endif
#ifndef BRANCH_0_SIZE
#define BRANCH_0_SIZE 9
#endif
#ifndef BRANCH_1_SIZE
#define BRANCH_1_SIZE 3
#endif
#ifndef BRANCH_2_SIZE
#define BRANCH_2_SIZE 3
#endif
#ifndef V_MIN
#define V_MIN (-25.0f)
#endif
#ifndef V_MAX
#define V_MAX (25.0f)
#endif
#define DELTA_Z ((V_MAX - V_MIN) / (float)(NUM_ATOMS - 1))
#define BRANCH_0_ATOMS (BRANCH_0_SIZE * NUM_ATOMS)
#define BRANCH_1_ATOMS (BRANCH_1_SIZE * NUM_ATOMS)
#define BRANCH_2_ATOMS (BRANCH_2_SIZE * NUM_ATOMS)
#define MAX_BRANCH_SIZE 9
#define NUM_BRANCHES 3
#define NUM_WARPS 8
#define BLOCK_THREADS 256
/* ── Shared memory layout (same as c51_loss_kernel.cu) ─────────────── */
#define SHMEM_SUPPORT_OFF 0
#define SHMEM_VAL_OFF (SHMEM_SUPPORT_OFF + NUM_ATOMS)
#define SHMEM_ADV_OFF (SHMEM_VAL_OFF + NUM_ATOMS)
#define SHMEM_LP_OFF (SHMEM_ADV_OFF + MAX_BRANCH_SIZE * NUM_ATOMS)
#define SHMEM_PROJ_OFF (SHMEM_LP_OFF + NUM_ATOMS)
#define SHMEM_CURRENT_LP_OFF (SHMEM_PROJ_OFF + NUM_ATOMS)
#define SHMEM_REDUCE_OFF (SHMEM_CURRENT_LP_OFF + NUM_ATOMS)
#define SHMEM_TOTAL_FLOATS (SHMEM_REDUCE_OFF + NUM_WARPS)
/* ── Device helpers (same as c51_loss_kernel.cu) ───────────────────── */
__device__ __forceinline__ float block_reduce_max(
@@ -102,28 +70,29 @@ __device__ __forceinline__ float block_reduce_sum(
* Also writes softmax probs to save_probs if save_probs != NULL.
*/
__device__ float block_softmax_expected_q(
float* __restrict__ shmem_logits, /* [NUM_ATOMS] in shmem (modified in-place) */
const float* __restrict__ shmem_support, /* [NUM_ATOMS] z_j */
float* __restrict__ shmem_logits, /* [num_atoms] in shmem (modified in-place) */
const float* __restrict__ shmem_support, /* [num_atoms] z_j */
float* __restrict__ shmem_reduce, /* [NUM_WARPS] scratch */
float* __restrict__ save_probs_out, /* [NUM_ATOMS] global mem or NULL */
int tid
float* __restrict__ save_probs_out, /* [num_atoms] global mem or NULL */
int tid,
int num_atoms
) {
/* Pass 1: block-max */
float local_max = -1e30f;
for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
for (int i = tid; i < num_atoms; i += BLOCK_THREADS)
local_max = fmaxf(local_max, shmem_logits[i]);
float block_max = block_reduce_max(local_max, shmem_reduce, tid);
/* Pass 2: sum(exp(x - max)) */
float local_sum = 0.0f;
for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
for (int i = tid; i < num_atoms; i += BLOCK_THREADS)
local_sum += expf(shmem_logits[i] - block_max);
float block_sum = block_reduce_sum(local_sum, shmem_reduce, tid);
float log_sum_exp = logf(block_sum + 1e-10f);
/* Pass 3: softmax probs + expected Q */
/* Write probs to shmem_logits in-place */
for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS) {
for (int i = tid; i < num_atoms; i += BLOCK_THREADS) {
float lp = shmem_logits[i] - block_max - log_sum_exp;
float prob = expf(lp);
shmem_logits[i] = prob; /* overwrite logits with probs */
@@ -133,7 +102,7 @@ __device__ float block_softmax_expected_q(
/* E[Q] = sum_j(prob_j * z_j) */
float local_eq = 0.0f;
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
local_eq += shmem_logits[j] * shmem_support[j];
return block_reduce_sum(local_eq, shmem_reduce, tid);
}
@@ -155,22 +124,22 @@ __device__ float block_softmax_expected_q(
extern "C" __global__ void mse_loss_batched(
/* ── Online network outputs on current STATES ─────────────────── */
const float* __restrict__ on_value_logits, /* [B, NUM_ATOMS] */
const float* __restrict__ on_adv_logits_b0, /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
const float* __restrict__ on_adv_logits_b1, /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
const float* __restrict__ on_adv_logits_b2, /* [B, BRANCH_2_SIZE * NUM_ATOMS] */
const float* __restrict__ on_value_logits, /* [B, num_atoms] */
const float* __restrict__ on_adv_logits_b0, /* [B, b0_size * num_atoms] */
const float* __restrict__ on_adv_logits_b1, /* [B, b1_size * num_atoms] */
const float* __restrict__ on_adv_logits_b2, /* [B, b2_size * num_atoms] */
/* ── Target network outputs on NEXT_STATES ────────────────────── */
const float* __restrict__ tg_value_logits, /* [B, NUM_ATOMS] */
const float* __restrict__ tg_adv_logits_b0, /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
const float* __restrict__ tg_adv_logits_b1, /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
const float* __restrict__ tg_adv_logits_b2, /* [B, BRANCH_2_SIZE * NUM_ATOMS] */
const float* __restrict__ tg_value_logits, /* [B, num_atoms] */
const float* __restrict__ tg_adv_logits_b0, /* [B, b0_size * num_atoms] */
const float* __restrict__ tg_adv_logits_b1, /* [B, b1_size * num_atoms] */
const float* __restrict__ tg_adv_logits_b2, /* [B, b2_size * num_atoms] */
/* ── Online network outputs on NEXT_STATES (Double DQN selector) */
const float* __restrict__ on_next_value_logits, /* [B, NUM_ATOMS] */
const float* __restrict__ on_next_adv_logits_b0, /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
const float* __restrict__ on_next_adv_logits_b1, /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
const float* __restrict__ on_next_adv_logits_b2, /* [B, BRANCH_2_SIZE * NUM_ATOMS] */
const float* __restrict__ on_next_value_logits, /* [B, num_atoms] */
const float* __restrict__ on_next_adv_logits_b0, /* [B, b0_size * num_atoms] */
const float* __restrict__ on_next_adv_logits_b1, /* [B, b1_size * num_atoms] */
const float* __restrict__ on_next_adv_logits_b2, /* [B, b2_size * num_atoms] */
/* ── Batch data ───────────────────────────────────────────────── */
const int* __restrict__ actions, /* [B] factored action indices 0-44 */
@@ -184,64 +153,87 @@ extern "C" __global__ void mse_loss_batched(
float* __restrict__ total_loss, /* [1] batch mean loss (atomicAdd) */
/* ── Saved tensors for backward pass ─────────────────────────── */
float* __restrict__ save_current_lp, /* [B, NUM_BRANCHES, NUM_ATOMS] online probs */
float* __restrict__ save_projected, /* [B, NUM_BRANCHES, NUM_ATOMS] [td, E_Q, 0..] */
float* __restrict__ save_current_lp, /* [B, NUM_BRANCHES, num_atoms] online probs */
float* __restrict__ save_projected, /* [B, NUM_BRANCHES, num_atoms] [td, E_Q, 0..] */
/* ── Config ───────────────────────────────────────────────────── */
float gamma,
int batch_size
int batch_size,
int num_atoms,
float v_min,
float v_max,
int b0_size,
int b1_size,
int b2_size
) {
extern __shared__ float shmem[];
float* shmem_support = shmem + SHMEM_SUPPORT_OFF;
float* shmem_val = shmem + SHMEM_VAL_OFF;
float* shmem_adv = shmem + SHMEM_ADV_OFF;
float* shmem_lp = shmem + SHMEM_LP_OFF;
float* shmem_proj = shmem + SHMEM_PROJ_OFF;
float* shmem_current_lp = shmem + SHMEM_CURRENT_LP_OFF;
float* shmem_reduce = shmem + SHMEM_REDUCE_OFF;
/* ── Compute runtime shared memory offsets ────────────────────── */
const int max_branch = max(max(b0_size, b1_size), b2_size);
int off_support = 0;
int off_val = off_support + num_atoms;
int off_adv = off_val + num_atoms;
int off_lp = off_adv + max_branch * num_atoms;
int off_proj = off_lp + num_atoms;
int off_current_lp = off_proj + num_atoms;
int off_reduce = off_current_lp + num_atoms;
float* shmem_support = shmem + off_support;
float* shmem_val = shmem + off_val;
float* shmem_adv = shmem + off_adv;
float* shmem_lp = shmem + off_lp;
float* shmem_proj = shmem + off_proj;
float* shmem_current_lp = shmem + off_current_lp;
float* shmem_reduce = shmem + off_reduce;
int sample_id = blockIdx.x;
int tid = threadIdx.x;
if (sample_id >= batch_size) return;
/* ── Compute runtime derived values ──────────────────────────── */
float delta_z = (v_max - v_min) / (float)(num_atoms - 1);
int b0_atoms = b0_size * num_atoms;
int b1_atoms = b1_size * num_atoms;
int b2_atoms = b2_size * num_atoms;
/* ── Precompute C51 support atoms once per block ─────────────── */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
shmem_support[j] = V_MIN + (float)j * DELTA_Z;
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_support[j] = v_min + (float)j * delta_z;
__syncthreads();
/* ── Decode factored action into per-branch indices ──────────── */
int factored_action = actions[sample_id];
int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE;
int max_action = b0_size * b1_size * b2_size;
if (factored_action < 0 || factored_action >= max_action)
factored_action = 0;
int a0 = factored_action / (BRANCH_1_SIZE * BRANCH_2_SIZE);
int a1 = (factored_action / BRANCH_2_SIZE) % BRANCH_1_SIZE;
int a2 = factored_action % BRANCH_2_SIZE;
int a0 = factored_action / (b1_size * b2_size);
int a1 = (factored_action / b2_size) % b1_size;
int a2 = factored_action % b2_size;
const int branch_action[NUM_BRANCHES] = { a0, a1, a2 };
const int branch_sizes[NUM_BRANCHES] = { BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE };
const int branch_sizes[NUM_BRANCHES] = { b0_size, b1_size, b2_size };
const float* on_adv_row[NUM_BRANCHES] = {
on_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
on_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
on_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
on_adv_logits_b0 + (long long)sample_id * b0_atoms,
on_adv_logits_b1 + (long long)sample_id * b1_atoms,
on_adv_logits_b2 + (long long)sample_id * b2_atoms
};
const float* tg_adv_row[NUM_BRANCHES] = {
tg_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
tg_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
tg_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
tg_adv_logits_b0 + (long long)sample_id * b0_atoms,
tg_adv_logits_b1 + (long long)sample_id * b1_atoms,
tg_adv_logits_b2 + (long long)sample_id * b2_atoms
};
const float* on_next_adv_row[NUM_BRANCHES] = {
on_next_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
on_next_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
on_next_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
on_next_adv_logits_b0 + (long long)sample_id * b0_atoms,
on_next_adv_logits_b1 + (long long)sample_id * b1_atoms,
on_next_adv_logits_b2 + (long long)sample_id * b2_atoms
};
const float* on_val_row = on_value_logits + (long long)sample_id * NUM_ATOMS;
const float* tg_val_row = tg_value_logits + (long long)sample_id * NUM_ATOMS;
const float* on_next_val_row = on_next_value_logits + (long long)sample_id * NUM_ATOMS;
const float* on_val_row = on_value_logits + (long long)sample_id * num_atoms;
const float* tg_val_row = tg_value_logits + (long long)sample_id * num_atoms;
const float* on_next_val_row = on_next_value_logits + (long long)sample_id * num_atoms;
float reward = rewards[sample_id];
float done = dones[sample_id];
@@ -253,16 +245,16 @@ extern "C" __global__ void mse_loss_batched(
for (int d = 0; d < NUM_BRANCHES; d++) {
int n_d = branch_sizes[d];
int a_d = branch_action[d];
int n_atoms = n_d * NUM_ATOMS;
int n_atoms = n_d * num_atoms;
long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * NUM_ATOMS;
long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * num_atoms;
/* ═══════════════════════════════════════════════════════════
* STEP a: Online E[Q] for the taken action a_d
* ═══════════════════════════════════════════════════════════ */
/* Load online value logits */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_val[j] = on_val_row[j];
__syncthreads();
@@ -272,26 +264,26 @@ extern "C" __global__ void mse_loss_batched(
__syncthreads();
/* Dueling for action a_d: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j]) */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
float a_mean = 0.0f;
for (int a = 0; a < n_d; a++)
a_mean += shmem_adv[a * NUM_ATOMS + j];
a_mean += shmem_adv[a * num_atoms + j];
a_mean /= (float)n_d;
shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * NUM_ATOMS + j] - a_mean;
shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * num_atoms + j] - a_mean;
}
__syncthreads();
/* softmax + E[Q], saving probs to save_current_lp */
float online_eq = block_softmax_expected_q(
shmem_lp, shmem_support, shmem_reduce,
save_current_lp + save_off, tid
save_current_lp + save_off, tid, num_atoms
);
/* ═══════════════════════════════════════════════════════════
* STEP b: Double DQN — argmax action on online_next network
* ═══════════════════════════════════════════════════════════ */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_val[j] = on_next_val_row[j];
__syncthreads();
@@ -300,10 +292,10 @@ extern "C" __global__ void mse_loss_batched(
__syncthreads();
/* Per-atom mean advantage for dueling */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
float a_mean = 0.0f;
for (int a = 0; a < n_d; a++)
a_mean += shmem_adv[a * NUM_ATOMS + j];
a_mean += shmem_adv[a * num_atoms + j];
a_mean /= (float)n_d;
shmem_proj[j] = a_mean;
}
@@ -313,13 +305,13 @@ extern "C" __global__ void mse_loss_batched(
__shared__ int best_next_a_shared;
for (int a = 0; a < n_d; a++) {
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
shmem_lp[j] = shmem_val[j] + shmem_adv[a * NUM_ATOMS + j] - shmem_proj[j];
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_lp[j] = shmem_val[j] + shmem_adv[a * num_atoms + j] - shmem_proj[j];
__syncthreads();
float eq = block_softmax_expected_q(
shmem_lp, shmem_support, shmem_reduce,
(float*)0, tid /* no save needed for action selection */
(float*)0, tid, num_atoms /* no save needed for action selection */
);
if (tid == 0) eq_per_action[a] = eq;
__syncthreads();
@@ -338,7 +330,7 @@ extern "C" __global__ void mse_loss_batched(
* STEP c: Target E[Q] for best_next action
* ═══════════════════════════════════════════════════════════ */
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_val[j] = tg_val_row[j];
__syncthreads();
@@ -346,18 +338,18 @@ extern "C" __global__ void mse_loss_batched(
shmem_adv[i] = tg_adv_row[d][i];
__syncthreads();
for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
float a_mean = 0.0f;
for (int a = 0; a < n_d; a++)
a_mean += shmem_adv[a * NUM_ATOMS + j];
a_mean += shmem_adv[a * num_atoms + j];
a_mean /= (float)n_d;
shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * NUM_ATOMS + j] - a_mean;
shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * num_atoms + j] - a_mean;
}
__syncthreads();
float target_eq = block_softmax_expected_q(
shmem_lp, shmem_support, shmem_reduce,
(float*)0, tid
(float*)0, tid, num_atoms
);
/* ═══════════════════════════════════════════════════════════
@@ -379,7 +371,7 @@ extern "C" __global__ void mse_loss_batched(
save_projected[save_off + 1] = online_eq;
}
/* Zero padding */
for (int j = tid + 2; j < NUM_ATOMS; j += BLOCK_THREADS)
for (int j = tid + 2; j < num_atoms; j += BLOCK_THREADS)
save_projected[save_off + j] = 0.0f;
__syncthreads();

159
crates/ml/src/cuda_pipeline/trade_physics.cuh Normal file
View File

@@ -0,0 +1,159 @@
/**
* Shared trade physics for DQN environment step kernels.
*
* This header defines the SINGLE SOURCE OF TRUTH for:
* - Factored action decoding (exposure index extraction)
* - Exposure-to-position mapping (index → fraction → contracts)
* - Hold enforcement (block exits/reversals during min_hold_bars)
* - Transaction cost computation (Almgren-Chriss impact model)
* - Capital floor circuit breaker (75% of peak equity)
* - Hold time tracking (bars in current position)
*
* Both backtest_env_kernel.cu and experience_kernels.cu MUST include
* this header to ensure train/eval consistency. Any change here
* automatically propagates to both.
*
* All functions are __device__ __forceinline__ for zero call overhead.
*/
#ifndef TRADE_PHYSICS_CUH
#define TRADE_PHYSICS_CUH
/* ── Factored action → exposure index ─────────────────────────────────── */
__device__ __forceinline__ int decode_exposure_index(
int action_val, int b0_size, int b1_size, int b2_size
) {
int exposure_idx;
if (b1_size > 0 && b2_size > 0) {
int denom = b1_size * b2_size;
exposure_idx = action_val / denom;
} else {
exposure_idx = action_val;
}
if (exposure_idx < 0) exposure_idx = 0;
if (exposure_idx >= b0_size) exposure_idx = b0_size - 1;
return exposure_idx;
}
/* ── Exposure index → target position ─────────────────────────────────── */
__device__ __forceinline__ float compute_target_position(
int exposure_idx, int b0_size, float max_position
) {
float step_size = (b0_size > 1) ? 2.0f / (float)(b0_size - 1) : 0.0f;
float fraction = -1.0f + (float)exposure_idx * step_size;
return fraction * max_position;
}
/* ── Margin-aware position cap ───────────────────────────────────────── */
/* Real brokers enforce margin requirements: you can't hold more contracts
* than your equity supports. ES futures require ~$15K initial margin per
* contract ($12K maintenance). Without this, a depleted account holds the
* same position as a full account → catastrophic overleveraging.
*
* This caps the position at equity / margin_per_contract, ensuring the
* portfolio can't take outsized risk with depleted capital.
*
* margin_per_contract: CME initial margin per contract (~$15,000 for ES)
* Passed as 0 to disable (backward compat with training kernel). */
__device__ __forceinline__ float apply_margin_cap(
float target_position,
float equity,
float margin_per_contract
) {
if (margin_per_contract <= 0.0f || equity <= 0.0f) return target_position;
float max_contracts = equity / margin_per_contract;
if (max_contracts < 0.01f) max_contracts = 0.0f; /* can't afford any position */
float abs_target = fabsf(target_position);
if (abs_target > max_contracts) {
float sign = (target_position > 0.0f) ? 1.0f : -1.0f;
return sign * max_contracts;
}
return target_position;
}
/* ── Order type index from factored action ────────────────────────────── */
__device__ __forceinline__ int decode_order_type(
int action_val, int b1_size, int b2_size
) {
return (b1_size > 0 && b2_size > 0) ? (action_val / b2_size) % b1_size : 0;
}
/* ── Hold enforcement: block exits and reversals ─────────────────────── */
__device__ __forceinline__ float enforce_hold(
float prev_position,
float target_exposure,
float hold_time,
int min_hold_bars,
int is_last_bar /* bypass hold on last bar of episode/window */
) {
int prev_sign = (prev_position > 0.001f) ? 1 : ((prev_position < -0.001f) ? -1 : 0);
int curr_sign = (target_exposure > 0.001f) ? 1 : ((target_exposure < -0.001f) ? -1 : 0);
int wants_exit = (prev_sign != 0 && curr_sign == 0);
int wants_reversal = (prev_sign != 0 && curr_sign != 0 && prev_sign != curr_sign);
if (hold_time > 0.0f && hold_time < (float)min_hold_bars
&& (wants_exit || wants_reversal)
&& !is_last_bar) {
return prev_position; /* Override: keep position */
}
return target_exposure;
}
/* ── Transaction cost (Almgren-Chriss impact model) ──────────────────── */
__device__ __forceinline__ float compute_tx_cost(
float delta, /* position change (signed) */
float close, /* current price */
float tx_cost_bps, /* base cost multiplier */
float spread_cost, /* bid-ask spread cost per unit */
float max_position, /* for impact scaling */
int order_type_idx /* 0=Market, 1=IoC, 2=LimitMaker */
) {
float abs_delta = fabsf(delta);
float spread_scale = sqrtf(abs_delta / fmaxf(max_position, 0.01f));
if (spread_scale < 1.0f) spread_scale = 1.0f;
float impact_scale = 1.0f;
float order_premium = (order_type_idx == 0) ? 0.0f
: (order_type_idx == 1) ? 0.0002f /* IoC: +2 bps */
: -0.0005f; /* LimitMaker: -5 bps rebate */
return abs_delta * close * (tx_cost_bps * 0.0001f * spread_scale * impact_scale + order_premium)
+ abs_delta * spread_cost * 0.5f;
}
/* ── Capital floor circuit breaker (PDT $25K rule) ───────────────────── */
/* Triggers when equity drops 25% from peak (floor = 75% of peak). */
/* With $35K initial: floor = $26.25K — above $25K PDT minimum. */
/* Matches training kernel (experience_env_step) capital_floor logic. */
__device__ __forceinline__ int check_capital_floor(
float value, float max_equity
) {
float capital_floor = max_equity * 0.75f;
return (value < capital_floor && max_equity > 1.0f) ? 1 : 0;
}
/* ── Hold time tracking ──────────────────────────────────────────────── */
__device__ __forceinline__ float update_hold_time(
float prev_position, float new_position, float hold_time
) {
if (fabsf(new_position) > 0.001f) {
if (fabsf(prev_position) < 0.001f) {
return 1.0f; /* Entry from flat */
}
int prev_sign = (prev_position > 0.001f) ? 1 : -1;
int new_sign = (new_position > 0.001f) ? 1 : -1;
if (new_sign == prev_sign) {
return hold_time + 1.0f; /* Same direction — increment */
}
return 1.0f; /* Reversal — new trade */
}
return 0.0f; /* Flat */
}
#endif /* TRADE_PHYSICS_CUH */

1447
crates/ml/src/hyperopt/adapters/dqn.rs
View File

File diff suppressed because it is too large Load Diff

6
crates/ml/src/hyperopt/campaign.rs
View File

@@ -310,6 +310,12 @@ mod tests {
#[test]
#[ignore] // GPU + real data, ~30 min on RTX 3050
fn test_local_hyperopt() {
// Initialize tracing so info!/warn!/error! output is visible in test runs
let _ = tracing_subscriber::fmt()
.with_env_filter("info")
.with_test_writer()
.try_init();
let data_dir = std::env::var("FOXHUNT_TEST_DATA")
.unwrap_or_else(|_| "test_data/futures-baseline".to_owned());
let data_path = std::path::Path::new(&data_dir);

2
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -759,7 +759,7 @@ impl FusedTrainingCtx {
// All auxiliary gradients (IQN, attention, ensemble, CQL) are now in grad_buf.
// The single Adam update sees: C51 + iqn_lambda*IQN + attention + diversity + CQL.
let fused_result = self.trainer.replay_adam_and_readback()
.map_err(|e| anyhow::anyhow!("graph_adam replay: {e}"))?;
.map_err(|e| { eprintln!("!!! ADAM REPLAY FAILED: {e}"); anyhow::anyhow!("graph_adam replay: {e}") })?;
// ── Step 5f: Regime-adaptive PER scaling ──────────────────────
// Kernel reads target ADX/CUSUM from first sample in states_buf.

27
crates/ml/src/trainers/dqn/trainer/tests.rs
View File

@@ -551,24 +551,27 @@ fn test_c2_five_actions_produce_distinct_exposures() {
assert_eq!(unique.len(), 5, "All 5 exposure levels must be distinct");
}
/// Verify hyperopt search space is 45D (Task 11: +4 new dims).
/// Verify hyperopt search space is 22D (consolidated from 46D).
#[test]
fn test_c3_search_space_is_40d() {
fn test_c3_search_space_is_22d() {
let bounds = crate::hyperopt::adapters::dqn::DQNParams::continuous_bounds();
assert_eq!(bounds.len(), 46, "Search space must be 46D (45 base + min_hold_bars)");
assert_eq!(bounds.len(), 22, "Search space must be 22D (consolidated from 46D)");
let names = crate::hyperopt::adapters::dqn::DQNParams::param_names();
assert_eq!(names.len(), 46);
assert!(names.contains(&"count_bonus_coefficient"), "count_bonus_coefficient must be in search space (C3)");
assert!(names.contains(&"sharpe_weight"), "sharpe_weight must be in search space (C4)");
assert!(names.contains(&"branch_hidden_dim"), "branch_hidden_dim must be in search space (L2)");
assert_eq!(names.len(), 22);
assert!(names.contains(&"c51_warmup_epochs"), "c51_warmup_epochs must be in search space");
// Task 11: curiosity_weight is now in the search space
assert!(names.contains(&"curiosity_weight"), "curiosity_weight must be in search space (Task 11)");
assert!(names.contains(&"her_ratio"), "her_ratio must be in search space (Task 11)");
assert!(names.contains(&"use_cvar_action_selection"), "use_cvar_action_selection must be in search space (Task 11)");
assert!(names.contains(&"cvar_alpha"), "cvar_alpha must be in search space (Task 11)");
assert!(names.contains(&"iqn_lambda"), "iqn_lambda must be in search space");
assert!(names.contains(&"min_hold_bars"), "min_hold_bars must be in search space");
assert!(names.contains(&"hidden_dim_base"), "hidden_dim_base must be in search space");
assert!(!names.contains(&"noisy_epsilon_floor"), "noisy_epsilon_floor must not be in search space");
// These are now FIXED (not in search space):
assert!(!names.contains(&"count_bonus_coefficient"), "count_bonus_coefficient should be fixed, not in search space");
assert!(!names.contains(&"sharpe_weight"), "sharpe_weight should be fixed, not in search space");
assert!(!names.contains(&"branch_hidden_dim"), "branch_hidden_dim should be fixed, not in search space");
assert!(!names.contains(&"curiosity_weight"), "curiosity_weight should be fixed, not in search space");
assert!(!names.contains(&"her_ratio"), "her_ratio should be fixed, not in search space");
assert!(!names.contains(&"use_cvar_action_selection"), "use_cvar_action_selection should be fixed, not in search space");
assert!(!names.contains(&"cvar_alpha"), "cvar_alpha should be fixed, not in search space");
}
/// Verify noisy_epsilon_floor is fixed to 0.10 (prevents action collapse).

9
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1070,7 +1070,7 @@ impl DQNTrainer {
// cross-stream deadlock or data corruption.
if let Some(ref stream) = self.cuda_stream {
stream.synchronize()
.map_err(|e| anyhow::anyhow!("Stream sync before PER insert: {e}"))?;
.map_err(|e| { eprintln!("!!! SYNC FAILED: {e}"); anyhow::anyhow!("Stream sync before PER insert: {e}") })?;
}
if let Some(ref mut mon) = self.gpu_monitoring {
@@ -1243,9 +1243,12 @@ impl DQNTrainer {
let batch_data = explicit_batch.as_ref().ok_or_else(|| {
anyhow::anyhow!("No batch data for fused training step")
})?;
fused.run_full_step(batch_data, &mut *agent, &self.device)
let result = fused.run_full_step(batch_data, &mut *agent, &self.device)
.map_err(|e| { eprintln!("!!! FUSED STEP ERROR: {:#}", e); e })
.context("Fused CUDA training step failed")?
.context("Fused CUDA training step failed")?;
if train_step_count % 50 == 0 {
}
result
} else {
agent.train_step(explicit_batch)
.map_err(|e| anyhow::anyhow!("Train step FAILED: {e}"))?