fix: promote entire backtest pipeline from bf16 to f32
CRITICAL CORRECTNESS FIX: bf16 has $32 resolution at $5K equity. Transaction costs ($12.50), tick PnL ($12.50), and cumulative returns were ALL below bf16 resolution — rounded to zero. Sharpe was quantized to ~0.004 resolution. All previous hyperopt evaluations used garbage metrics. Changed to f32: - backtest_env_kernel.cu: prices, portfolio_state, step_rewards, step_returns — both backtest_env_step and backtest_env_step_batch - backtest_metrics_kernel.cu: all accumulators, shared memory, metrics_out, annualization_factor - backtest_gather_kernel.cu: portfolio parameter (f32 → bf16 for model input at output stage only) - gpu_backtest_evaluator.rs: all buffer types, upload paths, metrics download, shared memory byte calculations Kept bf16: features_buf, states_buf (neural network input for tensor cores), cuBLAS forward pass buffers. Verified: Sharpe now has 6-digit precision (19.4745) vs bf16's ~0.004 resolution. All 345 tests + 3 smoke tests pass. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -1,7 +1,7 @@
|
||||
// Vectorized backtest environment step kernel.
|
||||
// One thread per walk-forward window. Each thread steps sequentially.
|
||||
//
|
||||
// Portfolio state layout per window [8 bf16]:
|
||||
// Portfolio state layout per window [8 f32]:
|
||||
// [0] value - current portfolio value
|
||||
// [1] position - current position size (-1.0 to +1.0)
|
||||
// [2] cash - cash balance
|
||||
@@ -21,9 +21,9 @@
|
||||
// cum_return and step_count are intentionally zeroed — the breached episode's
|
||||
// metrics are captured in done_flags and step_returns before this reset.
|
||||
__device__ void handle_capital_floor_breach(
|
||||
__nv_bfloat16* portfolio_state, int ps,
|
||||
__nv_bfloat16 new_capital, __nv_bfloat16 step_ret,
|
||||
__nv_bfloat16* step_rewards, __nv_bfloat16* step_returns, int* actions_history, int* done_flags,
|
||||
float* portfolio_state, int ps,
|
||||
float new_capital, float step_ret,
|
||||
float* step_rewards, float* step_returns, int* actions_history, int* done_flags,
|
||||
int w, int max_len, int current_step, int b0_size
|
||||
) {
|
||||
step_rewards[w] = step_ret;
|
||||
@@ -31,34 +31,34 @@ __device__ void handle_capital_floor_breach(
|
||||
actions_history[w * max_len + current_step] = b0_size / 2; // Flat
|
||||
// Full episode restart — metrics captured before this reset
|
||||
portfolio_state[ps + 0] = new_capital;
|
||||
portfolio_state[ps + 1] = bf16_zero();
|
||||
portfolio_state[ps + 1] = 0.0f;
|
||||
portfolio_state[ps + 2] = new_capital;
|
||||
portfolio_state[ps + 3] = bf16_zero();
|
||||
portfolio_state[ps + 3] = 0.0f;
|
||||
portfolio_state[ps + 4] = new_capital;
|
||||
portfolio_state[ps + 5] = bf16_zero();
|
||||
portfolio_state[ps + 6] = bf16_zero();
|
||||
portfolio_state[ps + 7] = bf16_zero();
|
||||
portfolio_state[ps + 5] = 0.0f;
|
||||
portfolio_state[ps + 6] = 0.0f;
|
||||
portfolio_state[ps + 7] = 0.0f;
|
||||
done_flags[w] = 1;
|
||||
}
|
||||
|
||||
extern "C" __global__ void backtest_env_step(
|
||||
// Market data (read-only, uploaded once)
|
||||
const __nv_bfloat16* __restrict__ prices, // [n_windows * max_len * 4] (OHLC)
|
||||
const float* __restrict__ prices, // [n_windows * max_len * 4] (OHLC)
|
||||
const int* __restrict__ window_lens, // [n_windows]
|
||||
|
||||
// Actions from model for current step
|
||||
const int* __restrict__ actions, // [n_windows] factored action index
|
||||
|
||||
// Portfolio state (read-write, persistent across steps)
|
||||
__nv_bfloat16* portfolio_state, // [n_windows * PORTFOLIO_STATE_SIZE]
|
||||
float* portfolio_state, // [n_windows * PORTFOLIO_STATE_SIZE]
|
||||
|
||||
// Step outputs
|
||||
__nv_bfloat16* step_rewards, // [n_windows]
|
||||
__nv_bfloat16* step_returns, // [n_windows * max_len] (accumulated)
|
||||
float* step_rewards, // [n_windows]
|
||||
float* step_returns, // [n_windows * max_len] (accumulated)
|
||||
int* done_flags, // [n_windows]
|
||||
int* actions_history, // [n_windows * max_len] (accumulated, for metrics)
|
||||
|
||||
// Config — host scalars stay float, convert on first use
|
||||
// Config — host scalars stay float
|
||||
int n_windows,
|
||||
int max_len,
|
||||
float max_position,
|
||||
@@ -72,7 +72,7 @@ extern "C" __global__ void backtest_env_step(
|
||||
float contract_multiplier, // e.g. 50.0 for ES, 20.0 for NQ
|
||||
float margin_pct // e.g. 0.06 (6% initial margin)
|
||||
) {
|
||||
__shared__ __nv_bfloat16 shmem_pf[256 * PORTFOLIO_STATE_SIZE];
|
||||
__shared__ float shmem_pf[256 * PORTFOLIO_STATE_SIZE];
|
||||
|
||||
int w = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int local_tid = threadIdx.x;
|
||||
@@ -95,38 +95,37 @@ extern "C" __global__ void backtest_env_step(
|
||||
return;
|
||||
}
|
||||
|
||||
// Read current prices — native bf16
|
||||
// Read current prices — native f32
|
||||
int price_base = (w * max_len + current_step) * 4;
|
||||
__nv_bfloat16 open = prices[price_base + 0];
|
||||
__nv_bfloat16 high = prices[price_base + 1];
|
||||
__nv_bfloat16 low = prices[price_base + 2];
|
||||
__nv_bfloat16 close = prices[price_base + 3];
|
||||
float open = prices[price_base + 0];
|
||||
float high = prices[price_base + 1];
|
||||
float low = prices[price_base + 2];
|
||||
float close = prices[price_base + 3];
|
||||
|
||||
// Read portfolio state from shared memory tile — native bf16
|
||||
__nv_bfloat16 value = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 0];
|
||||
__nv_bfloat16 position = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 1];
|
||||
__nv_bfloat16 cash = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 2];
|
||||
__nv_bfloat16 entry_price = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 3];
|
||||
__nv_bfloat16 max_equity = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 4];
|
||||
__nv_bfloat16 hold_time = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 5];
|
||||
__nv_bfloat16 cum_return = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 6];
|
||||
// Read portfolio state from shared memory tile — native f32
|
||||
float value = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 0];
|
||||
float position = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 1];
|
||||
float cash = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 2];
|
||||
float entry_price = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 3];
|
||||
float max_equity = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 4];
|
||||
float hold_time = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 5];
|
||||
float cum_return = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 6];
|
||||
|
||||
// ── Capital floor circuit breaker (shared: trade_physics.cuh) ────────
|
||||
// trade_physics expects float — promote for the call, then back to bf16
|
||||
if (check_capital_floor(__bfloat162float(value), __bfloat162float(max_equity))) {
|
||||
if (check_capital_floor(value, max_equity)) {
|
||||
// Force flat: close any open position at current price (notional model)
|
||||
if (__bfloat162float(bf16_fabs(position)) > 0.001f && close > bf16_zero()) {
|
||||
float exit_cost = compute_tx_cost(__bfloat162float(position), __bfloat162float(close),
|
||||
if (fabsf(position) > 0.001f && close > 0.0f) {
|
||||
float exit_cost = compute_tx_cost(position, close,
|
||||
tx_cost_bps, spread_cost,
|
||||
max_position, 0, -1.0f);
|
||||
cash = cash + position * close; // sell position at market (notional)
|
||||
cash = cash - bf16(exit_cost);
|
||||
position = bf16_zero();
|
||||
cash = cash - exit_cost;
|
||||
position = 0.0f;
|
||||
}
|
||||
__nv_bfloat16 liq_value = cash;
|
||||
__nv_bfloat16 liq_ret = (value > bf16(0.01f))
|
||||
float liq_value = cash;
|
||||
float liq_ret = (value > 0.01f)
|
||||
? (liq_value - value) / value
|
||||
: bf16_zero();
|
||||
: 0.0f;
|
||||
handle_capital_floor_breach(portfolio_state, ps, liq_value, liq_ret,
|
||||
step_rewards, step_returns, actions_history, done_flags,
|
||||
w, max_len, current_step, b0_size);
|
||||
@@ -136,27 +135,22 @@ extern "C" __global__ void backtest_env_step(
|
||||
// ── 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);
|
||||
// trade_physics compute_target_position returns float
|
||||
__nv_bfloat16 target_exposure = bf16(compute_target_position(exposure_idx, b0_size, max_position));
|
||||
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) ────────────
|
||||
// margin = price * multiplier * margin_pct (e.g. 5000 * 50 * 0.06 = $15K for ES)
|
||||
__nv_bfloat16 margin_per_contract = close * bf16(contract_multiplier) * bf16(margin_pct);
|
||||
target_exposure = bf16(apply_margin_cap(__bfloat162float(target_exposure),
|
||||
__bfloat162float(value),
|
||||
__bfloat162float(margin_per_contract),
|
||||
__bfloat162float(max_equity)));
|
||||
float margin_per_contract = close * contract_multiplier * margin_pct;
|
||||
target_exposure = apply_margin_cap(target_exposure, value,
|
||||
margin_per_contract, max_equity);
|
||||
|
||||
// Suppress unused variable warnings
|
||||
(void)open; (void)high; (void)low;
|
||||
|
||||
// ── Hold enforcement (shared: trade_physics.cuh) ─────────────────────
|
||||
int is_last_bar = (current_step >= wlen - 1) ? 1 : 0;
|
||||
target_exposure = bf16(enforce_hold(__bfloat162float(position),
|
||||
__bfloat162float(target_exposure),
|
||||
__bfloat162float(hold_time),
|
||||
min_hold_bars, is_last_bar));
|
||||
target_exposure = enforce_hold(position, target_exposure, hold_time,
|
||||
min_hold_bars, is_last_bar);
|
||||
|
||||
// ── Trailing stop (shared: trade_physics.cuh) ────────────────────────
|
||||
// Exit when profit retreats from peak. Uses 0.5% base distance.
|
||||
@@ -165,58 +159,50 @@ extern "C" __global__ void backtest_env_step(
|
||||
{
|
||||
// Portfolio-weighted trade return: unrealized P&L / equity
|
||||
// Matches training kernel's unrealized_trade_pnl computation
|
||||
__nv_bfloat16 trade_ret = bf16_zero();
|
||||
if (__bfloat162float(bf16_fabs(position)) > 0.001f
|
||||
&& entry_price > bf16_zero()
|
||||
&& value > bf16_one()) {
|
||||
__nv_bfloat16 unrealized = position * (close - entry_price);
|
||||
float trade_ret = 0.0f;
|
||||
if (fabsf(position) > 0.001f
|
||||
&& entry_price > 0.0f
|
||||
&& value > 1.0f) {
|
||||
float unrealized = position * (close - entry_price);
|
||||
trade_ret = unrealized / value;
|
||||
}
|
||||
if (check_trailing_stop(__bfloat162float(hold_time), min_hold_bars,
|
||||
__bfloat162float(max_equity),
|
||||
__bfloat162float(value),
|
||||
__bfloat162float(trade_ret), 0.005f, 1.0f, 1.0f)) {
|
||||
target_exposure = bf16_zero(); // Force flat — trailing stop triggered
|
||||
if (check_trailing_stop(hold_time, min_hold_bars,
|
||||
max_equity, value,
|
||||
trade_ret, 0.005f, 1.0f, 1.0f)) {
|
||||
target_exposure = 0.0f; // Force flat — trailing stop triggered
|
||||
}
|
||||
}
|
||||
|
||||
// ── Execute trade (shared: trade_physics.cuh) ────────────────────────
|
||||
__nv_bfloat16 prev_position = position;
|
||||
float prev_position = position;
|
||||
{
|
||||
float f_position = __bfloat162float(position);
|
||||
float f_cash = __bfloat162float(cash);
|
||||
float tx_cost = execute_trade(&f_position, &f_cash,
|
||||
__bfloat162float(target_exposure),
|
||||
__bfloat162float(close),
|
||||
float tx_cost = execute_trade(&position, &cash,
|
||||
target_exposure, close,
|
||||
tx_cost_bps, spread_cost, max_position,
|
||||
order_type_idx, -1.0f);
|
||||
position = bf16(f_position);
|
||||
cash = bf16(f_cash);
|
||||
(void)tx_cost;
|
||||
}
|
||||
|
||||
// Update entry_price only on NEW trade entry or reversal (for trailing stop)
|
||||
if (__bfloat162float(bf16_fabs(prev_position)) < 0.001f
|
||||
&& __bfloat162float(bf16_fabs(position)) > 0.001f) {
|
||||
if (fabsf(prev_position) < 0.001f
|
||||
&& fabsf(position) > 0.001f) {
|
||||
entry_price = close; // new entry from flat
|
||||
} else if (__bfloat162float(prev_position) * __bfloat162float(position) < 0.0f) {
|
||||
} else if (prev_position * position < 0.0f) {
|
||||
entry_price = close; // reversal
|
||||
} else if (__bfloat162float(bf16_fabs(position)) < 0.001f) {
|
||||
entry_price = bf16_zero(); // went flat
|
||||
} else if (fabsf(position) < 0.001f) {
|
||||
entry_price = 0.0f; // went flat
|
||||
}
|
||||
// On same-direction scaling (L50->L100), keep original entry_price
|
||||
|
||||
// ── Hold time tracking (shared: trade_physics.cuh) ───────────────────
|
||||
hold_time = bf16(update_hold_time(__bfloat162float(prev_position),
|
||||
__bfloat162float(position),
|
||||
__bfloat162float(hold_time)));
|
||||
hold_time = update_hold_time(prev_position, position, hold_time);
|
||||
|
||||
// Mark-to-market current position (notional model: equity = cash + position * price)
|
||||
__nv_bfloat16 new_value = cash + position * close;
|
||||
float new_value = cash + position * close;
|
||||
|
||||
// Update max_equity BEFORE floor check — prevents stale peak from
|
||||
// missing breaches or triggering false ones after profitable trades.
|
||||
max_equity = bf16_fmax(max_equity, new_value);
|
||||
max_equity = fmaxf(max_equity, new_value);
|
||||
|
||||
// ── Post-trade floor check: catch intra-step breaches ────────────
|
||||
// The pre-trade check (top of kernel) catches breaches from the
|
||||
@@ -224,21 +210,21 @@ extern "C" __global__ void backtest_env_step(
|
||||
// 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(__bfloat162float(new_value), __bfloat162float(max_equity))) {
|
||||
if (check_capital_floor(new_value, max_equity)) {
|
||||
// Emergency liquidation: close position at current price (notional model)
|
||||
if (__bfloat162float(bf16_fabs(position)) > 0.001f) {
|
||||
float f_exit_cost = compute_tx_cost(__bfloat162float(position), __bfloat162float(close),
|
||||
if (fabsf(position) > 0.001f) {
|
||||
float f_exit_cost = compute_tx_cost(position, close,
|
||||
tx_cost_bps, spread_cost,
|
||||
max_position, 0, -1.0f);
|
||||
cash = cash + position * close; // sell position at market (notional)
|
||||
cash = cash - bf16(f_exit_cost);
|
||||
position = bf16_zero();
|
||||
cash = cash - f_exit_cost;
|
||||
position = 0.0f;
|
||||
new_value = cash; // flat: equity = cash
|
||||
entry_price = bf16_zero();
|
||||
entry_price = 0.0f;
|
||||
}
|
||||
__nv_bfloat16 step_ret = (value > bf16(0.01f))
|
||||
float step_ret = (value > 0.01f)
|
||||
? (new_value - value) / value
|
||||
: bf16_zero();
|
||||
: 0.0f;
|
||||
handle_capital_floor_breach(portfolio_state, ps, new_value, step_ret,
|
||||
step_rewards, step_returns, actions_history, done_flags,
|
||||
w, max_len, current_step, b0_size);
|
||||
@@ -246,11 +232,11 @@ extern "C" __global__ void backtest_env_step(
|
||||
}
|
||||
|
||||
// Step return
|
||||
__nv_bfloat16 step_ret = (value > bf16_zero())
|
||||
float step_ret = (value > 0.0f)
|
||||
? (new_value - value) / value
|
||||
: bf16_zero();
|
||||
__nv_bfloat16 new_cum_return = cum_return + step_ret;
|
||||
__nv_bfloat16 new_max = bf16_fmax(max_equity, new_value);
|
||||
: 0.0f;
|
||||
float new_cum_return = cum_return + step_ret;
|
||||
float new_max = fmaxf(max_equity, new_value);
|
||||
|
||||
// Write portfolio state
|
||||
portfolio_state[ps + 0] = new_value;
|
||||
@@ -260,7 +246,7 @@ extern "C" __global__ void backtest_env_step(
|
||||
portfolio_state[ps + 4] = new_max;
|
||||
portfolio_state[ps + 5] = hold_time;
|
||||
portfolio_state[ps + 6] = new_cum_return;
|
||||
portfolio_state[ps + 7] = portfolio_state[ps + 7] + bf16_one();
|
||||
portfolio_state[ps + 7] = portfolio_state[ps + 7] + 1.0f;
|
||||
|
||||
// Outputs
|
||||
step_rewards[w] = step_ret;
|
||||
@@ -269,7 +255,7 @@ extern "C" __global__ void backtest_env_step(
|
||||
// Re-encode position -> exposure_idx -> factored action to ensure metrics
|
||||
// kernel counts real position changes, not model-requested actions.
|
||||
{
|
||||
float f_actual_frac = __bfloat162float(position) / fmaxf(max_position, 0.01f);
|
||||
float f_actual_frac = position / fmaxf(max_position, 0.01f);
|
||||
int actual_exp_idx = (int)roundf((f_actual_frac + 1.0f) * 0.5f * (float)(b0_size - 1));
|
||||
if (actual_exp_idx < 0) actual_exp_idx = 0;
|
||||
if (actual_exp_idx >= b0_size) actual_exp_idx = b0_size - 1;
|
||||
@@ -292,12 +278,12 @@ extern "C" __global__ void backtest_env_step(
|
||||
* Launch: grid=(ceil(n_windows/256)), block=(256).
|
||||
* ══════════════════════════════════════════════════════════════════════ */
|
||||
extern "C" __global__ void backtest_env_step_batch(
|
||||
const __nv_bfloat16* __restrict__ prices,
|
||||
const float* __restrict__ prices,
|
||||
const int* __restrict__ window_lens,
|
||||
const int* __restrict__ chunked_actions, /* [chunk_len * n_windows] */
|
||||
__nv_bfloat16* portfolio_state,
|
||||
__nv_bfloat16* step_rewards,
|
||||
__nv_bfloat16* step_returns,
|
||||
float* portfolio_state,
|
||||
float* step_rewards,
|
||||
float* step_returns,
|
||||
int* done_flags,
|
||||
int* actions_history,
|
||||
int n_windows,
|
||||
@@ -321,30 +307,30 @@ extern "C" __global__ void backtest_env_step_batch(
|
||||
int ps = w * PORTFOLIO_STATE_SIZE;
|
||||
|
||||
/* Load portfolio into registers — persistent across the step loop */
|
||||
__nv_bfloat16 value = portfolio_state[ps + 0];
|
||||
__nv_bfloat16 position = portfolio_state[ps + 1];
|
||||
__nv_bfloat16 cash = portfolio_state[ps + 2];
|
||||
__nv_bfloat16 entry_price = portfolio_state[ps + 3];
|
||||
__nv_bfloat16 max_equity = portfolio_state[ps + 4];
|
||||
__nv_bfloat16 hold_time = portfolio_state[ps + 5];
|
||||
__nv_bfloat16 cum_return = portfolio_state[ps + 6];
|
||||
float value = portfolio_state[ps + 0];
|
||||
float position = portfolio_state[ps + 1];
|
||||
float cash = portfolio_state[ps + 2];
|
||||
float entry_price = portfolio_state[ps + 3];
|
||||
float max_equity = portfolio_state[ps + 4];
|
||||
float hold_time = portfolio_state[ps + 5];
|
||||
float cum_return = portfolio_state[ps + 6];
|
||||
|
||||
for (int s = 0; s < chunk_len; s++) {
|
||||
int current_step = start_step + s;
|
||||
if (done_flags[w] || current_step >= wlen) break;
|
||||
|
||||
/* Capital floor check */
|
||||
if (check_capital_floor(__bfloat162float(value), __bfloat162float(max_equity))) {
|
||||
if (__bfloat162float(bf16_fabs(position)) > 0.001f && value > bf16_zero()) {
|
||||
float exit_cost = compute_tx_cost(__bfloat162float(position),
|
||||
__bfloat162float(prices[(w * max_len + current_step) * 4 + 3]),
|
||||
if (check_capital_floor(value, max_equity)) {
|
||||
if (fabsf(position) > 0.001f && value > 0.0f) {
|
||||
float exit_cost = compute_tx_cost(position,
|
||||
prices[(w * max_len + current_step) * 4 + 3],
|
||||
tx_cost_bps, spread_cost, max_position, 0, -1.0f);
|
||||
__nv_bfloat16 close = prices[(w * max_len + current_step) * 4 + 3];
|
||||
cash = cash + position * close - bf16(exit_cost);
|
||||
position = bf16_zero();
|
||||
float close = prices[(w * max_len + current_step) * 4 + 3];
|
||||
cash = cash + position * close - exit_cost;
|
||||
position = 0.0f;
|
||||
}
|
||||
value = cash;
|
||||
__nv_bfloat16 liq_ret = bf16_zero();
|
||||
float liq_ret = 0.0f;
|
||||
handle_capital_floor_breach(portfolio_state, ps, value, liq_ret,
|
||||
step_rewards, step_returns, actions_history, done_flags,
|
||||
w, max_len, current_step, b0_size);
|
||||
@@ -353,84 +339,78 @@ extern "C" __global__ void backtest_env_step_batch(
|
||||
|
||||
/* Read prices */
|
||||
int price_base = (w * max_len + current_step) * 4;
|
||||
__nv_bfloat16 close = prices[price_base + 3];
|
||||
float close = prices[price_base + 3];
|
||||
|
||||
/* Read action from chunked buffer: [step_offset, window] layout */
|
||||
int action_val = chunked_actions[s * n_windows + w];
|
||||
|
||||
/* Decode + constraints */
|
||||
int exposure_idx = decode_exposure_index(action_val, b0_size, b1_size, b2_size);
|
||||
__nv_bfloat16 target_exposure = bf16(compute_target_position(exposure_idx, b0_size, max_position));
|
||||
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);
|
||||
|
||||
__nv_bfloat16 margin_per = close * bf16(contract_multiplier) * bf16(margin_pct);
|
||||
target_exposure = bf16(apply_margin_cap(__bfloat162float(target_exposure),
|
||||
__bfloat162float(value), __bfloat162float(margin_per),
|
||||
__bfloat162float(max_equity)));
|
||||
float margin_per = close * contract_multiplier * margin_pct;
|
||||
target_exposure = apply_margin_cap(target_exposure, value,
|
||||
margin_per, max_equity);
|
||||
|
||||
int is_last = (current_step >= wlen - 1) ? 1 : 0;
|
||||
target_exposure = bf16(enforce_hold(__bfloat162float(position),
|
||||
__bfloat162float(target_exposure), __bfloat162float(hold_time),
|
||||
min_hold_bars, is_last));
|
||||
target_exposure = enforce_hold(position, target_exposure, hold_time,
|
||||
min_hold_bars, is_last);
|
||||
|
||||
/* Trailing stop */
|
||||
{
|
||||
__nv_bfloat16 trade_ret = bf16_zero();
|
||||
if (__bfloat162float(bf16_fabs(position)) > 0.001f
|
||||
&& entry_price > bf16_zero() && value > bf16_one()) {
|
||||
float trade_ret = 0.0f;
|
||||
if (fabsf(position) > 0.001f
|
||||
&& entry_price > 0.0f && value > 1.0f) {
|
||||
trade_ret = position * (close - entry_price) / value;
|
||||
}
|
||||
if (check_trailing_stop(__bfloat162float(hold_time), min_hold_bars,
|
||||
__bfloat162float(max_equity), __bfloat162float(value),
|
||||
__bfloat162float(trade_ret), 0.005f, 1.0f, 1.0f)) {
|
||||
target_exposure = bf16_zero();
|
||||
if (check_trailing_stop(hold_time, min_hold_bars,
|
||||
max_equity, value,
|
||||
trade_ret, 0.005f, 1.0f, 1.0f)) {
|
||||
target_exposure = 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
/* Execute trade */
|
||||
__nv_bfloat16 prev_position = position;
|
||||
float prev_position = position;
|
||||
{
|
||||
float f_pos = __bfloat162float(position);
|
||||
float f_cash = __bfloat162float(cash);
|
||||
execute_trade(&f_pos, &f_cash, __bfloat162float(target_exposure),
|
||||
__bfloat162float(close), tx_cost_bps, spread_cost,
|
||||
execute_trade(&position, &cash, target_exposure,
|
||||
close, tx_cost_bps, spread_cost,
|
||||
max_position, order_type_idx, -1.0f);
|
||||
position = bf16(f_pos);
|
||||
cash = bf16(f_cash);
|
||||
}
|
||||
|
||||
/* Entry price tracking */
|
||||
if (__bfloat162float(bf16_fabs(prev_position)) < 0.001f
|
||||
&& __bfloat162float(bf16_fabs(position)) > 0.001f) {
|
||||
if (fabsf(prev_position) < 0.001f
|
||||
&& fabsf(position) > 0.001f) {
|
||||
entry_price = close;
|
||||
} else if (__bfloat162float(prev_position) * __bfloat162float(position) < 0.0f) {
|
||||
} else if (prev_position * position < 0.0f) {
|
||||
entry_price = close;
|
||||
}
|
||||
|
||||
/* Hold time */
|
||||
hold_time = (__bfloat162float(bf16_fabs(position)) > 0.001f)
|
||||
? hold_time + bf16_one()
|
||||
: bf16_zero();
|
||||
hold_time = (fabsf(position) > 0.001f)
|
||||
? hold_time + 1.0f
|
||||
: 0.0f;
|
||||
|
||||
/* Max equity reset on flat */
|
||||
if (__bfloat162float(bf16_fabs(position)) < 0.001f
|
||||
&& __bfloat162float(bf16_fabs(prev_position)) > 0.001f) {
|
||||
if (fabsf(position) < 0.001f
|
||||
&& fabsf(prev_position) > 0.001f) {
|
||||
max_equity = cash;
|
||||
}
|
||||
|
||||
/* New value */
|
||||
__nv_bfloat16 new_value = cash + position * close * bf16(contract_multiplier);
|
||||
__nv_bfloat16 step_ret = (value > bf16_zero())
|
||||
? (new_value - value) / value : bf16_zero();
|
||||
float new_value = cash + position * close * contract_multiplier;
|
||||
float step_ret = (value > 0.0f)
|
||||
? (new_value - value) / value : 0.0f;
|
||||
cum_return = cum_return + step_ret;
|
||||
max_equity = bf16_fmax(max_equity, new_value);
|
||||
max_equity = fmaxf(max_equity, new_value);
|
||||
value = new_value;
|
||||
|
||||
/* Outputs */
|
||||
step_rewards[w] = step_ret;
|
||||
step_returns[w * max_len + current_step] = step_ret;
|
||||
{
|
||||
float f_actual = __bfloat162float(position) / fmaxf(max_position, 0.01f);
|
||||
float f_actual = position / fmaxf(max_position, 0.01f);
|
||||
int act_exp = (int)roundf((f_actual + 1.0f) * 0.5f * (float)(b0_size - 1));
|
||||
if (act_exp < 0) act_exp = 0;
|
||||
if (act_exp >= b0_size) act_exp = b0_size - 1;
|
||||
@@ -449,5 +429,5 @@ extern "C" __global__ void backtest_env_step_batch(
|
||||
portfolio_state[ps + 4] = max_equity;
|
||||
portfolio_state[ps + 5] = hold_time;
|
||||
portfolio_state[ps + 6] = cum_return;
|
||||
portfolio_state[ps + 7] = portfolio_state[ps + 7] + bf16((float)chunk_len);
|
||||
portfolio_state[ps + 7] = portfolio_state[ps + 7] + (float)chunk_len;
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
// [market_dim+8 .. market_dim+8+16) : 16 multi-timeframe features
|
||||
// [market_dim+24 .. state_dim) : zero-pad / OFI / alignment
|
||||
//
|
||||
// Portfolio state layout per window [8 bf16]:
|
||||
// Portfolio state layout per window [8 f32]:
|
||||
// [0] value - current portfolio value
|
||||
// [1] position - current position size (-1.0 to +1.0)
|
||||
// [2] cash - cash balance
|
||||
@@ -33,7 +33,7 @@
|
||||
|
||||
extern "C" __global__ void gather_states(
|
||||
const __nv_bfloat16* __restrict__ features, // [n_windows * max_len * feat_dim]
|
||||
const __nv_bfloat16* __restrict__ portfolio, // [n_windows * 8]
|
||||
const float* __restrict__ portfolio, // [n_windows * 8] (f32 from env kernel)
|
||||
__nv_bfloat16* states_out, // [n_windows * state_dim]
|
||||
int n_windows,
|
||||
int max_len,
|
||||
@@ -44,8 +44,8 @@ extern "C" __global__ void gather_states(
|
||||
float spread_cost, // host scalar — convert on first use
|
||||
int ofi_dim // 0 = no OFI, 8 = standard OFI features
|
||||
) {
|
||||
// Shared memory tile for coalesced portfolio reads — 8 BF16 values per thread
|
||||
__shared__ __nv_bfloat16 shmem_portfolio[256 * 8];
|
||||
// Shared memory tile for coalesced portfolio reads — 8 F32 values per thread
|
||||
__shared__ float shmem_portfolio[256 * 8];
|
||||
|
||||
int w = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int local_tid = threadIdx.x;
|
||||
@@ -65,13 +65,21 @@ extern "C" __global__ void gather_states(
|
||||
int feat_base = (w * max_len + current_step) * feat_dim;
|
||||
int out_base = w * state_dim;
|
||||
|
||||
// Read portfolio state from shared memory tile — native bf16
|
||||
__nv_bfloat16 value = shmem_portfolio[local_tid * 8 + 0];
|
||||
__nv_bfloat16 position = shmem_portfolio[local_tid * 8 + 1];
|
||||
__nv_bfloat16 cash = shmem_portfolio[local_tid * 8 + 2];
|
||||
__nv_bfloat16 entry_price = shmem_portfolio[local_tid * 8 + 3];
|
||||
__nv_bfloat16 max_equity = shmem_portfolio[local_tid * 8 + 4];
|
||||
__nv_bfloat16 step_count = shmem_portfolio[local_tid * 8 + 7];
|
||||
// Read portfolio state from shared memory tile — f32, convert to bf16 for output
|
||||
float f_value = shmem_portfolio[local_tid * 8 + 0];
|
||||
float f_position = shmem_portfolio[local_tid * 8 + 1];
|
||||
float f_cash = shmem_portfolio[local_tid * 8 + 2];
|
||||
float f_entry_price = shmem_portfolio[local_tid * 8 + 3];
|
||||
float f_max_equity = shmem_portfolio[local_tid * 8 + 4];
|
||||
float f_step_count = shmem_portfolio[local_tid * 8 + 7];
|
||||
|
||||
// Convert to bf16 for feature computation (output is bf16 for model input)
|
||||
__nv_bfloat16 value = bf16(f_value);
|
||||
__nv_bfloat16 position = bf16(f_position);
|
||||
__nv_bfloat16 cash = bf16(f_cash);
|
||||
__nv_bfloat16 entry_price = bf16(f_entry_price);
|
||||
__nv_bfloat16 max_equity = bf16(f_max_equity);
|
||||
__nv_bfloat16 step_count = bf16(f_step_count);
|
||||
|
||||
// Determine market_dim: if OFI is present, market features are feat_dim - ofi_dim
|
||||
int market_dim = (ofi_dim > 0) ? (feat_dim - ofi_dim) : feat_dim;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
// Per-window metrics reduction kernel — native BF16.
|
||||
// Per-window metrics reduction kernel — native F32.
|
||||
// One block per window. Threads cooperate to reduce step_returns.
|
||||
//
|
||||
// Output per window [14 bf16 values written to float metrics_out]:
|
||||
// Output per window [14 f32 values written to metrics_out]:
|
||||
// [0] sharpe_ratio (annualized)
|
||||
// [1] total_pnl (multiplicative compounding)
|
||||
// [2] max_drawdown (exact sequential scan)
|
||||
@@ -18,14 +18,14 @@
|
||||
// [13] unique_action_count
|
||||
|
||||
extern "C" __global__ void compute_backtest_metrics(
|
||||
const __nv_bfloat16* __restrict__ step_returns, // [n_windows * max_len]
|
||||
const __nv_bfloat16* __restrict__ portfolio_state, // [n_windows * 8]
|
||||
const float* __restrict__ step_returns, // [n_windows * max_len]
|
||||
const float* __restrict__ portfolio_state, // [n_windows * 8]
|
||||
const int* __restrict__ window_lens, // [n_windows]
|
||||
const int* __restrict__ actions_history, // [n_windows * max_len]
|
||||
__nv_bfloat16* metrics_out, // [n_windows * 14]
|
||||
float* metrics_out, // [n_windows * 14]
|
||||
int n_windows,
|
||||
int max_len,
|
||||
__nv_bfloat16 annualization_factor, // bf16 scalar
|
||||
float annualization_factor, // f32 scalar
|
||||
int num_actions,
|
||||
int order_actions,
|
||||
int urgency_actions
|
||||
@@ -39,33 +39,33 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
int base = w * max_len;
|
||||
|
||||
// Shared memory layout (6 reduction arrays + sort scratch):
|
||||
extern __shared__ __nv_bfloat16 shmem[];
|
||||
__nv_bfloat16* s_sum = shmem;
|
||||
__nv_bfloat16* s_sq_sum = shmem + stride;
|
||||
__nv_bfloat16* s_down_sq = shmem + 2*stride;
|
||||
__nv_bfloat16* s_buys = shmem + 3*stride;
|
||||
__nv_bfloat16* s_sells = shmem + 4*stride;
|
||||
__nv_bfloat16* s_holds = shmem + 5*stride;
|
||||
__nv_bfloat16* s_sorted = shmem + 6*stride; // [4096] for sort + boundary data
|
||||
extern __shared__ float shmem[];
|
||||
float* s_sum = shmem;
|
||||
float* s_sq_sum = shmem + stride;
|
||||
float* s_down_sq = shmem + 2*stride;
|
||||
float* s_buys = shmem + 3*stride;
|
||||
float* s_sells = shmem + 4*stride;
|
||||
float* s_holds = shmem + 5*stride;
|
||||
float* s_sorted = shmem + 6*stride; // [4096] for sort + boundary data
|
||||
|
||||
// Per-thread local accumulators — native BF16
|
||||
__nv_bfloat16 local_sum = bf16_zero();
|
||||
__nv_bfloat16 local_sq = bf16_zero();
|
||||
__nv_bfloat16 local_down = bf16_zero();
|
||||
__nv_bfloat16 local_cum = bf16_one();
|
||||
// Per-thread local accumulators — native F32
|
||||
float local_sum = 0.0f;
|
||||
float local_sq = 0.0f;
|
||||
float local_down = 0.0f;
|
||||
float local_cum = 1.0f;
|
||||
int local_buys = 0, local_sells = 0, local_holds = 0;
|
||||
int local_action_mask = 0;
|
||||
|
||||
// Boundary-aware trade tracking
|
||||
int bnd_first_action = -1;
|
||||
int bnd_last_action = -1;
|
||||
__nv_bfloat16 bnd_prefix_return = bf16_zero();
|
||||
__nv_bfloat16 bnd_suffix_return = bf16_zero();
|
||||
float bnd_prefix_return = 0.0f;
|
||||
float bnd_suffix_return = 0.0f;
|
||||
int bnd_complete_trades = 0;
|
||||
int bnd_complete_wins = 0;
|
||||
int bnd_num_changes = 0;
|
||||
int bnd_cur_action = -1;
|
||||
__nv_bfloat16 bnd_cur_return = bf16_zero();
|
||||
float bnd_cur_return = 0.0f;
|
||||
|
||||
// Consecutive chunks for correct drawdown tracking
|
||||
int chunk_size = (wlen + stride - 1) / stride;
|
||||
@@ -74,12 +74,12 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
if (chunk_end > wlen) chunk_end = wlen;
|
||||
|
||||
for (int i = chunk_start; i < chunk_end; i++) {
|
||||
__nv_bfloat16 r = step_returns[base + i];
|
||||
float r = step_returns[base + i];
|
||||
local_sum += r;
|
||||
local_sq += r * r;
|
||||
if (r < bf16_zero()) local_down += r * r;
|
||||
if (r < 0.0f) local_down += r * r;
|
||||
|
||||
local_cum *= (bf16_one() + r);
|
||||
local_cum *= (1.0f + r);
|
||||
|
||||
int act = actions_history[base + i];
|
||||
int exp_idx = act / (order_actions * urgency_actions);
|
||||
@@ -101,9 +101,9 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
bnd_prefix_return = bnd_cur_return;
|
||||
} else {
|
||||
bnd_complete_trades++;
|
||||
if (bnd_cur_return > bf16_zero()) bnd_complete_wins++;
|
||||
if (bnd_cur_return > 0.0f) bnd_complete_wins++;
|
||||
}
|
||||
bnd_cur_return = bf16_zero();
|
||||
bnd_cur_return = 0.0f;
|
||||
bnd_cur_action = exp_idx;
|
||||
}
|
||||
|
||||
@@ -120,12 +120,12 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
s_sum[tid] = local_sum;
|
||||
s_sq_sum[tid] = local_sq;
|
||||
s_down_sq[tid] = local_down;
|
||||
s_buys[tid] = bf16((float)local_buys);
|
||||
s_sells[tid] = bf16((float)local_sells);
|
||||
s_holds[tid] = bf16((float)local_holds);
|
||||
s_buys[tid] = (float)local_buys;
|
||||
s_sells[tid] = (float)local_sells;
|
||||
s_holds[tid] = (float)local_holds;
|
||||
__syncthreads();
|
||||
|
||||
// Block-level parallel reduction — native bf16 += bf16
|
||||
// Block-level parallel reduction — native f32 += f32
|
||||
for (int s = stride / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
s_sum[tid] += s_sum[tid + s];
|
||||
@@ -149,70 +149,68 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
}
|
||||
|
||||
// Export boundary trade data to shared memory.
|
||||
// Action indices fit exactly in bf16 (-1 to 8). Trade counts stored as bf16
|
||||
// (exact for ≤256, ±1 for 256-1024 — acceptable for monitoring metrics).
|
||||
{
|
||||
__nv_bfloat16* s_bnd = &s_sorted[stride];
|
||||
s_bnd[0 * stride + tid] = bf16((float)bnd_first_action);
|
||||
s_bnd[1 * stride + tid] = bf16((float)bnd_last_action);
|
||||
float* s_bnd = &s_sorted[stride];
|
||||
s_bnd[0 * stride + tid] = (float)bnd_first_action;
|
||||
s_bnd[1 * stride + tid] = (float)bnd_last_action;
|
||||
s_bnd[2 * stride + tid] = bnd_prefix_return;
|
||||
s_bnd[3 * stride + tid] = bnd_suffix_return;
|
||||
s_bnd[4 * stride + tid] = bf16((float)bnd_complete_trades);
|
||||
s_bnd[5 * stride + tid] = bf16((float)bnd_complete_wins);
|
||||
s_bnd[6 * stride + tid] = bf16((float)bnd_num_changes);
|
||||
s_bnd[4 * stride + tid] = (float)bnd_complete_trades;
|
||||
s_bnd[5 * stride + tid] = (float)bnd_complete_wins;
|
||||
s_bnd[6 * stride + tid] = (float)bnd_num_changes;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Thread 0: final metrics from fully reduced values
|
||||
if (tid == 0) {
|
||||
__nv_bfloat16 n = bf16((float)wlen);
|
||||
__nv_bfloat16 mean = s_sum[0] / n;
|
||||
__nv_bfloat16 var = s_sq_sum[0] / n - mean * mean;
|
||||
__nv_bfloat16 std_val = bf16_sqrt(bf16_fmax(var, bf16(1e-10f)));
|
||||
__nv_bfloat16 down_std = bf16_sqrt(bf16_fmax(s_down_sq[0] / n, bf16(1e-10f)));
|
||||
float n = (float)wlen;
|
||||
float mean = s_sum[0] / n;
|
||||
float var = s_sq_sum[0] / n - mean * mean;
|
||||
float std_val = sqrtf(fmaxf(var, 1e-10f));
|
||||
float down_std = sqrtf(fmaxf(s_down_sq[0] / n, 1e-10f));
|
||||
|
||||
// Exact max drawdown via sequential scan
|
||||
__nv_bfloat16 exact_equity = bf16_one();
|
||||
__nv_bfloat16 exact_peak = bf16_one();
|
||||
__nv_bfloat16 exact_max_dd = bf16_zero();
|
||||
float exact_equity = 1.0f;
|
||||
float exact_peak = 1.0f;
|
||||
float exact_max_dd = 0.0f;
|
||||
for (int i = 0; i < wlen; i++) {
|
||||
__nv_bfloat16 r = step_returns[base + i];
|
||||
exact_equity *= (bf16_one() + r);
|
||||
exact_peak = bf16_fmax(exact_peak, exact_equity);
|
||||
__nv_bfloat16 dd = (exact_peak - exact_equity) / bf16_fmax(exact_peak, bf16(1e-8f));
|
||||
exact_max_dd = bf16_fmax(exact_max_dd, dd);
|
||||
float r = step_returns[base + i];
|
||||
exact_equity *= (1.0f + r);
|
||||
exact_peak = fmaxf(exact_peak, exact_equity);
|
||||
float dd = (exact_peak - exact_equity) / fmaxf(exact_peak, 1e-8f);
|
||||
exact_max_dd = fmaxf(exact_max_dd, dd);
|
||||
}
|
||||
|
||||
// Boundary stitching: exact trade count + win rate
|
||||
__nv_bfloat16* s_bnd = &s_sorted[stride];
|
||||
float* s_bnd = &s_sorted[stride];
|
||||
int total_trades = 0;
|
||||
int total_wins = 0;
|
||||
|
||||
for (int t = 0; t < stride; t++) {
|
||||
total_trades += (int)(float)s_bnd[4 * stride + t];
|
||||
total_wins += (int)(float)s_bnd[5 * stride + t];
|
||||
total_trades += (int)s_bnd[4 * stride + t];
|
||||
total_wins += (int)s_bnd[5 * stride + t];
|
||||
}
|
||||
|
||||
__nv_bfloat16 open_return = bf16_zero();
|
||||
float open_return = 0.0f;
|
||||
int open_action = -1;
|
||||
|
||||
for (int t = 0; t < stride; t++) {
|
||||
int fa = (int)(float)s_bnd[0 * stride + t];
|
||||
int la = (int)(float)s_bnd[1 * stride + t];
|
||||
__nv_bfloat16 pr = s_bnd[2 * stride + t];
|
||||
__nv_bfloat16 sr = s_bnd[3 * stride + t];
|
||||
int nc = (int)(float)s_bnd[6 * stride + t];
|
||||
int fa = (int)s_bnd[0 * stride + t];
|
||||
int la = (int)s_bnd[1 * stride + t];
|
||||
float pr = s_bnd[2 * stride + t];
|
||||
float sr = s_bnd[3 * stride + t];
|
||||
int nc = (int)s_bnd[6 * stride + t];
|
||||
|
||||
if (fa < 0) continue;
|
||||
|
||||
if (open_action < 0) {
|
||||
open_action = fa;
|
||||
open_return = bf16_zero();
|
||||
open_return = 0.0f;
|
||||
} else if (fa != open_action) {
|
||||
total_trades++;
|
||||
if (open_return > bf16_zero()) total_wins++;
|
||||
if (open_return > 0.0f) total_wins++;
|
||||
open_action = fa;
|
||||
open_return = bf16_zero();
|
||||
open_return = 0.0f;
|
||||
}
|
||||
|
||||
if (nc == 0) {
|
||||
@@ -220,7 +218,7 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
} else {
|
||||
open_return += pr;
|
||||
total_trades++;
|
||||
if (open_return > bf16_zero()) total_wins++;
|
||||
if (open_return > 0.0f) total_wins++;
|
||||
open_action = la;
|
||||
open_return = sr;
|
||||
}
|
||||
@@ -228,19 +226,19 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
|
||||
if (open_action >= 0) {
|
||||
total_trades++;
|
||||
if (open_return > bf16_zero()) total_wins++;
|
||||
if (open_return > 0.0f) total_wins++;
|
||||
}
|
||||
|
||||
int out_base = w * 14;
|
||||
metrics_out[out_base + 0] = (mean / std_val) * annualization_factor;
|
||||
metrics_out[out_base + 1] = s_sorted[0] - bf16_one();
|
||||
metrics_out[out_base + 1] = s_sorted[0] - 1.0f;
|
||||
metrics_out[out_base + 2] = exact_max_dd;
|
||||
metrics_out[out_base + 3] = (mean / down_std) * annualization_factor;
|
||||
|
||||
__nv_bfloat16 total_trades_bf = bf16((float)total_trades);
|
||||
float total_trades_f = (float)total_trades;
|
||||
metrics_out[out_base + 4] = (total_trades > 0)
|
||||
? bf16((float)total_wins) / total_trades_bf : bf16_zero();
|
||||
metrics_out[out_base + 5] = total_trades_bf;
|
||||
? (float)total_wins / total_trades_f : 0.0f;
|
||||
metrics_out[out_base + 5] = total_trades_f;
|
||||
|
||||
metrics_out[out_base + 10] = s_buys[0];
|
||||
metrics_out[out_base + 11] = s_sells[0];
|
||||
@@ -258,7 +256,7 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
}
|
||||
if (tid == 0) {
|
||||
int out_base = w * 14;
|
||||
metrics_out[out_base + 13] = bf16((float)__popc(s_action_mask[0]));
|
||||
metrics_out[out_base + 13] = (float)__popc(s_action_mask[0]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -267,19 +265,19 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
int sample_step = (wlen + sort_len - 1) / sort_len;
|
||||
for (int i = tid; i < sort_len; i += stride) {
|
||||
int src = i * sample_step;
|
||||
s_sorted[i] = (src < wlen) ? step_returns[base + src] : bf16_zero();
|
||||
s_sorted[i] = (src < wlen) ? step_returns[base + src] : 0.0f;
|
||||
}
|
||||
|
||||
int padded_len = 1;
|
||||
while (padded_len < sort_len) padded_len <<= 1;
|
||||
|
||||
__nv_bfloat16 bf16_inf = bf16(1e30f);
|
||||
float f32_inf = 1e30f;
|
||||
for (int i = sort_len + tid; i < padded_len; i += stride) {
|
||||
s_sorted[i] = bf16_inf;
|
||||
s_sorted[i] = f32_inf;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Bitonic sort — ascending, native bf16 comparisons
|
||||
// Bitonic sort — ascending, native f32 comparisons
|
||||
for (int k = 2; k <= padded_len; k <<= 1) {
|
||||
for (int j = k >> 1; j > 0; j >>= 1) {
|
||||
for (int i = tid; i < padded_len; i += stride) {
|
||||
@@ -288,7 +286,7 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
bool ascending = ((i & k) == 0);
|
||||
if ((ascending && s_sorted[i] > s_sorted[ixj]) ||
|
||||
(!ascending && s_sorted[i] < s_sorted[ixj])) {
|
||||
__nv_bfloat16 tmp = s_sorted[i];
|
||||
float tmp = s_sorted[i];
|
||||
s_sorted[i] = s_sorted[ixj];
|
||||
s_sorted[ixj] = tmp;
|
||||
}
|
||||
@@ -305,30 +303,30 @@ extern "C" __global__ void compute_backtest_metrics(
|
||||
int var_idx = (int)(0.05f * (float)sort_len);
|
||||
if (var_idx < 0) var_idx = 0;
|
||||
if (var_idx >= sort_len) var_idx = sort_len - 1;
|
||||
__nv_bfloat16 var_95 = s_sorted[var_idx];
|
||||
float var_95 = s_sorted[var_idx];
|
||||
|
||||
int cvar_count = var_idx + 1;
|
||||
__nv_bfloat16 cvar_sum = bf16_zero();
|
||||
float cvar_sum = 0.0f;
|
||||
for (int i = 0; i < cvar_count; i++) {
|
||||
cvar_sum += s_sorted[i];
|
||||
}
|
||||
__nv_bfloat16 cvar_95 = cvar_sum / bf16((float)cvar_count);
|
||||
float cvar_95 = cvar_sum / (float)cvar_count;
|
||||
|
||||
__nv_bfloat16 daily_mean = s_sum[0] / bf16((float)wlen);
|
||||
__nv_bfloat16 max_dd = metrics_out[out_base + 2];
|
||||
__nv_bfloat16 calmar = (max_dd > bf16(0.001f))
|
||||
float daily_mean = s_sum[0] / (float)wlen;
|
||||
float max_dd = metrics_out[out_base + 2];
|
||||
float calmar = (max_dd > 0.001f)
|
||||
? (daily_mean * annualization_factor * annualization_factor) / max_dd
|
||||
: bf16_zero();
|
||||
calmar = bf16_fmax(bf16(-100.0f), bf16_fmin(bf16(100.0f), calmar));
|
||||
: 0.0f;
|
||||
calmar = fmaxf(-100.0f, fminf(100.0f, calmar));
|
||||
|
||||
__nv_bfloat16 gain_sum = bf16_zero(), loss_sum = bf16_zero();
|
||||
float gain_sum = 0.0f, loss_sum = 0.0f;
|
||||
for (int i = 0; i < sort_len; i++) {
|
||||
if (s_sorted[i] > bf16_zero()) gain_sum += s_sorted[i];
|
||||
else loss_sum -= s_sorted[i];
|
||||
if (s_sorted[i] > 0.0f) gain_sum += s_sorted[i];
|
||||
else loss_sum -= s_sorted[i];
|
||||
}
|
||||
__nv_bfloat16 omega = (loss_sum > bf16(0.001f))
|
||||
? gain_sum / loss_sum : bf16_zero();
|
||||
omega = bf16_fmax(bf16_zero(), bf16_fmin(bf16(100.0f), omega));
|
||||
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;
|
||||
|
||||
@@ -251,14 +251,14 @@ pub struct GpuBacktestEvaluator {
|
||||
gather_kernel: CudaFunction,
|
||||
|
||||
// Uploaded data (read-only; persists across the step loop)
|
||||
prices_buf: CudaSlice<half::bf16>, // [n_windows * max_len * 4]
|
||||
prices_buf: CudaSlice<f32>, // [n_windows * max_len * 4]
|
||||
features_buf: CudaSlice<half::bf16>, // [n_windows * max_len * feat_dim]
|
||||
window_lens_buf: CudaSlice<i32>, // [n_windows]
|
||||
|
||||
// Mutable state (written by kernels each step)
|
||||
portfolio_buf: CudaSlice<half::bf16>, // [n_windows * 8]
|
||||
step_rewards_buf: CudaSlice<half::bf16>, // [n_windows]
|
||||
step_returns_buf: CudaSlice<half::bf16>, // [n_windows * max_len]
|
||||
portfolio_buf: CudaSlice<f32>, // [n_windows * 8]
|
||||
step_rewards_buf: CudaSlice<f32>, // [n_windows]
|
||||
step_returns_buf: CudaSlice<f32>, // [n_windows * max_len]
|
||||
done_buf: CudaSlice<i32>, // [n_windows]
|
||||
actions_buf: CudaSlice<i32>, // [n_windows]
|
||||
actions_history_buf: CudaSlice<i32>, // [n_windows * max_len]
|
||||
@@ -267,7 +267,7 @@ pub struct GpuBacktestEvaluator {
|
||||
states_buf: CudaSlice<half::bf16>, // [n_windows * state_dim] (8-aligned)
|
||||
|
||||
// Output buffer (written by metrics kernel)
|
||||
metrics_buf: CudaSlice<half::bf16>, // [n_windows * 14]
|
||||
metrics_buf: CudaSlice<f32>, // [n_windows * 14]
|
||||
|
||||
// Dimensions and config
|
||||
n_windows: usize,
|
||||
@@ -506,7 +506,8 @@ impl GpuBacktestEvaluator {
|
||||
.map_err(|e| MLError::ModelError(format!("gather_states load: {e}")))?;
|
||||
|
||||
// ── Upload read-only data ─────────────────────────────────────────
|
||||
let prices_buf = super::clone_htod_f32_to_bf16(&stream, &flat_prices)?;
|
||||
let prices_buf = stream.clone_htod(&flat_prices)
|
||||
.map_err(|e| MLError::ModelError(format!("prices upload: {e}")))?;
|
||||
let features_buf = super::clone_htod_f32_to_bf16(&stream, &flat_features)?;
|
||||
let window_lens_buf = stream
|
||||
.clone_htod(&window_lens)
|
||||
@@ -514,13 +515,14 @@ impl GpuBacktestEvaluator {
|
||||
|
||||
// ── Allocate mutable state buffers ────────────────────────────────
|
||||
let portfolio_init = Self::init_portfolio_state(n_windows, config.initial_capital);
|
||||
let portfolio_buf = super::clone_htod_f32_to_bf16(&stream, &portfolio_init)?;
|
||||
let portfolio_buf = stream.clone_htod(&portfolio_init)
|
||||
.map_err(|e| MLError::ModelError(format!("portfolio upload: {e}")))?;
|
||||
|
||||
let step_rewards_buf = stream
|
||||
.alloc_zeros::<half::bf16>(n_windows)
|
||||
.alloc_zeros::<f32>(n_windows)
|
||||
.map_err(|e| MLError::ModelError(format!("step_rewards alloc: {e}")))?;
|
||||
let step_returns_buf = stream
|
||||
.alloc_zeros::<half::bf16>(n_windows * max_len)
|
||||
.alloc_zeros::<f32>(n_windows * max_len)
|
||||
.map_err(|e| MLError::ModelError(format!("step_returns alloc: {e}")))?;
|
||||
let done_buf = stream
|
||||
.alloc_zeros::<i32>(n_windows)
|
||||
@@ -532,7 +534,7 @@ impl GpuBacktestEvaluator {
|
||||
.alloc_zeros::<i32>(n_windows * max_len)
|
||||
.map_err(|e| MLError::ModelError(format!("actions_history alloc: {e}")))?;
|
||||
let metrics_buf = stream
|
||||
.alloc_zeros::<half::bf16>(n_windows * 14)
|
||||
.alloc_zeros::<f32>(n_windows * 14)
|
||||
.map_err(|e| MLError::ModelError(format!("metrics alloc: {e}")))?;
|
||||
|
||||
// State layout matches training experience_state_gather:
|
||||
@@ -551,7 +553,8 @@ impl GpuBacktestEvaluator {
|
||||
.map_err(|e| MLError::ModelError(format!("forward_actions alloc: {e}")))?;
|
||||
|
||||
let upload_mb =
|
||||
((flat_prices.len() + flat_features.len()) * std::mem::size_of::<half::bf16>()) as f64
|
||||
(flat_prices.len() * std::mem::size_of::<f32>()
|
||||
+ flat_features.len() * std::mem::size_of::<half::bf16>()) as f64
|
||||
/ 1_048_576.0;
|
||||
info!(
|
||||
"GpuBacktestEvaluator: {} windows × {} max_len × {} features ({:.1} MB uploaded)",
|
||||
@@ -680,12 +683,12 @@ impl GpuBacktestEvaluator {
|
||||
let state_dim = self.state_dim;
|
||||
|
||||
// Launch the gather kernel — one thread per window
|
||||
// Shared memory: 256 threads × 8 portfolio bf16 × 2 bytes = 4 KB
|
||||
// Shared memory: 256 threads × 8 portfolio f32 × 4 bytes = 8 KB
|
||||
let grid = ((self.n_windows + 255) / 256) as u32;
|
||||
let launch_cfg = LaunchConfig {
|
||||
grid_dim: (grid.max(1), 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 256 * 8 * std::mem::size_of::<half::bf16>() as u32,
|
||||
shared_mem_bytes: 256 * 8 * std::mem::size_of::<f32>() as u32,
|
||||
};
|
||||
let n_windows_i32 = self.n_windows as i32;
|
||||
let max_len_i32 = self.max_len as i32;
|
||||
@@ -930,8 +933,9 @@ impl GpuBacktestEvaluator {
|
||||
fn reset_evaluation_state(&mut self) -> Result<(), MLError> {
|
||||
// Re-init portfolio to initial capital
|
||||
let portfolio_init = Self::init_portfolio_state(self.n_windows, self.config.initial_capital);
|
||||
let portfolio_bf16 = super::clone_htod_f32_to_bf16(&self.stream, &portfolio_init)?;
|
||||
self.portfolio_buf = portfolio_bf16;
|
||||
let portfolio_f32 = self.stream.clone_htod(&portfolio_init)
|
||||
.map_err(|e| MLError::ModelError(format!("reset portfolio upload: {e}")))?;
|
||||
self.portfolio_buf = portfolio_f32;
|
||||
|
||||
// Zero done flags, step returns, step rewards, actions history
|
||||
self.stream.memset_zeros(&mut self.done_buf)
|
||||
@@ -1641,12 +1645,12 @@ impl GpuBacktestEvaluator {
|
||||
/// `evaluate_ppo` paths read directly from `states_buf`.
|
||||
fn launch_gather(&self, step: usize) -> Result<(), MLError> {
|
||||
let state_dim = self.state_dim;
|
||||
// Shared memory: 256 threads × 8 portfolio bf16 × 2 bytes = 4 KB
|
||||
// Shared memory: 256 threads × 8 portfolio f32 × 4 bytes = 8 KB
|
||||
let grid = ((self.n_windows + 255) / 256) as u32;
|
||||
let launch_cfg = LaunchConfig {
|
||||
grid_dim: (grid.max(1), 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 256 * 8 * std::mem::size_of::<half::bf16>() as u32,
|
||||
shared_mem_bytes: 256 * 8 * std::mem::size_of::<f32>() as u32,
|
||||
};
|
||||
let n_windows_i32 = self.n_windows as i32;
|
||||
let max_len_i32 = self.max_len as i32;
|
||||
@@ -1681,12 +1685,12 @@ impl GpuBacktestEvaluator {
|
||||
|
||||
/// Launch `backtest_env_step` kernel for a given step on `self.stream`.
|
||||
fn launch_env_step(&self, step: usize) -> Result<(), MLError> {
|
||||
// Shared memory: 256 threads × PORTFOLIO_STATE_SIZE(8) bf16 × 2 bytes = 4 KB
|
||||
// Shared memory: 256 threads × PORTFOLIO_STATE_SIZE(8) f32 × 4 bytes = 8 KB
|
||||
let grid = ((self.n_windows + 255) / 256) as u32;
|
||||
let env_cfg = LaunchConfig {
|
||||
grid_dim: (grid.max(1), 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 256 * 8 * std::mem::size_of::<half::bf16>() as u32,
|
||||
shared_mem_bytes: 256 * 8 * std::mem::size_of::<f32>() as u32,
|
||||
};
|
||||
let n_win_i32 = self.n_windows as i32;
|
||||
let max_len_i32 = self.max_len as i32;
|
||||
@@ -1734,7 +1738,7 @@ impl GpuBacktestEvaluator {
|
||||
// Shared memory: 6 reduction arrays × 256 threads × 4 bytes = 6 KB
|
||||
// + 4096 floats for sort scratch / boundary data = 16 KB
|
||||
// Total = 22 KB (well within the 48 KB L1/shmem limit)
|
||||
let shmem_bytes = (256_u32 * 6 + 4096) * std::mem::size_of::<half::bf16>() as u32;
|
||||
let shmem_bytes = (256_u32 * 6 + 4096) * std::mem::size_of::<f32>() as u32;
|
||||
let metrics_cfg = LaunchConfig {
|
||||
grid_dim: (self.n_windows as u32, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
@@ -1742,7 +1746,7 @@ impl GpuBacktestEvaluator {
|
||||
};
|
||||
let n_win_i32 = self.n_windows as i32;
|
||||
let max_len_i32 = self.max_len as i32;
|
||||
let annualization: half::bf16 = half::bf16::from_f32(self.annualization_factor);
|
||||
let annualization: f32 = self.annualization_factor;
|
||||
|
||||
// Safety: argument order matches `compute_backtest_metrics` signature exactly.
|
||||
let num_actions_i32 = self.b0_size;
|
||||
@@ -1767,8 +1771,8 @@ impl GpuBacktestEvaluator {
|
||||
}
|
||||
|
||||
// Single download: n_windows × 14 floats (the ONLY GPU→CPU transfer)
|
||||
let mut metrics_host = vec![0.0_f32; self.n_windows * 14];
|
||||
super::dtoh_bf16_to_f32(&self.stream, &self.metrics_buf, &mut metrics_host)?;
|
||||
let metrics_host: Vec<f32> = self.stream.clone_dtoh(&self.metrics_buf)
|
||||
.map_err(|e| MLError::ModelError(format!("metrics download: {e}")))?;
|
||||
|
||||
// Parse flat metrics into per-window structs.
|
||||
// Layout matches compute_backtest_metrics kernel output (14 floats per window):
|
||||
|
||||
746
docs/superpowers/plans/2026-04-06-per-prefix-sum.md
Normal file
746
docs/superpowers/plans/2026-04-06-per-prefix-sum.md
Normal file
@@ -0,0 +1,746 @@
|
||||
# PER Prefix Sum Replacement Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Replace the PER segment tree (hangs on H100) with a decoupled lookback prefix sum scan + binary search sampling — the NVIDIA-standard O(N) parallel scan algorithm.
|
||||
|
||||
**Architecture:** Delete `seg_tree_kernel.cu` and all segment tree code. Create `per_kernels.cu` with 4 kernels: `per_update_pa` (leaf write, accepts bf16 directly), `per_prefix_scan` (Merrill-Garland decoupled lookback), `per_sample` (binary search + gather), `per_insert_pa` (insert-time leaf write). Delete the `CastKernels` / `OnceLock` / `cast_kernels.cubin` infrastructure (source of H100 hang). Delete `update_td_f32` scratch buffer. Two kernel launches for the entire PER cycle instead of 25+.
|
||||
|
||||
**Tech Stack:** Rust 1.85, CUDA 12.4 (SM 9.0), cudarc 0.19, half (bf16)
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Write `per_kernels.cu` with All 4 Kernels
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml-dqn/src/per_kernels.cu`
|
||||
|
||||
- [ ] **Step 1: Write the `per_update_pa` kernel**
|
||||
|
||||
Create `crates/ml-dqn/src/per_kernels.cu`:
|
||||
|
||||
```cuda
|
||||
// PER (Prioritized Experience Replay) kernels — prefix sum architecture.
|
||||
//
|
||||
// Replaces the segment tree with a flat prefix sum array using NVIDIA's
|
||||
// decoupled lookback scan algorithm (Merrill & Garland, 2016).
|
||||
//
|
||||
// Four kernels:
|
||||
// 1. per_update_pa — priority update: compute pa from bf16 TD errors,
|
||||
// write to priorities_pa[] and priorities[].
|
||||
// 2. per_insert_pa — insert: write priority^alpha to priorities_pa[].
|
||||
// 3. per_prefix_scan — decoupled lookback inclusive scan of priorities_pa.
|
||||
// 4. per_sample — proportional sampling via binary search on prefix sum.
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// 1. per_update_pa — Priority leaf update from bf16 TD errors
|
||||
//
|
||||
// Reads bf16 td_errors directly (no separate bf16→f32 cast kernel).
|
||||
// Writes both raw priorities and priority^alpha values.
|
||||
// Launch: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1)
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
extern "C" __global__ void per_update_pa(
|
||||
float* __restrict__ priorities_pa,
|
||||
float* __restrict__ priorities,
|
||||
float* __restrict__ batch_max,
|
||||
const unsigned int* __restrict__ indices,
|
||||
const __nv_bfloat16* __restrict__ td_errors_bf16,
|
||||
float alpha, float epsilon,
|
||||
int capacity, int batch_size)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= batch_size) return;
|
||||
|
||||
unsigned int idx = indices[i];
|
||||
if (idx >= (unsigned int)capacity) return;
|
||||
|
||||
float td = fabsf(__bfloat162float(td_errors_bf16[i]));
|
||||
float new_prio = powf(td, alpha) + epsilon;
|
||||
if (new_prio < epsilon) new_prio = epsilon;
|
||||
if (new_prio > 1e6f) new_prio = 1e6f;
|
||||
|
||||
priorities[idx] = new_prio;
|
||||
|
||||
int ival = __float_as_int(new_prio);
|
||||
atomicMax((int*)batch_max, ival);
|
||||
|
||||
float pa = powf(new_prio, alpha);
|
||||
priorities_pa[idx] = pa;
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// 2. per_insert_pa — Insert priorities at max_priority
|
||||
//
|
||||
// Called during insert_batch. Takes raw priority values (max_priority fill),
|
||||
// computes priority^alpha, writes to priorities_pa[].
|
||||
// Launch: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1)
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
extern "C" __global__ void per_insert_pa(
|
||||
float* __restrict__ priorities_pa,
|
||||
const unsigned int* __restrict__ indices,
|
||||
const float* __restrict__ raw_priorities,
|
||||
float alpha,
|
||||
int capacity, int batch_size)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= batch_size) return;
|
||||
|
||||
unsigned int idx = indices[i];
|
||||
if (idx >= (unsigned int)capacity) return;
|
||||
|
||||
priorities_pa[idx] = powf(raw_priorities[i], alpha);
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// 3. per_prefix_scan — Decoupled lookback inclusive prefix sum
|
||||
//
|
||||
// Single-pass O(N) parallel scan using the Merrill-Garland algorithm.
|
||||
// Each block processes TILE_SIZE elements, publishes its aggregate,
|
||||
// then looks back through preceding blocks to determine its global prefix.
|
||||
//
|
||||
// Status flags per tile (packed in tile_descriptors):
|
||||
// bits [31:30] = state: 0=INVALID, 1=AGGREGATE, 2=PREFIX
|
||||
// bits [29:0] = unused (value stored separately for f32 precision)
|
||||
//
|
||||
// Launch: grid=(num_tiles, 1, 1), block=(TILE_SIZE, 1, 1)
|
||||
// where num_tiles = ceil(size / TILE_SIZE), TILE_SIZE = 256
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
#define PER_TILE_SIZE 256
|
||||
#define STATE_INVALID 0u
|
||||
#define STATE_AGGREGATE 1u
|
||||
#define STATE_PREFIX 2u
|
||||
|
||||
extern "C" __global__ void per_prefix_scan(
|
||||
const float* __restrict__ input,
|
||||
float* __restrict__ output,
|
||||
volatile unsigned int* __restrict__ tile_state,
|
||||
volatile float* __restrict__ tile_aggregate,
|
||||
volatile float* __restrict__ tile_prefix,
|
||||
int size)
|
||||
{
|
||||
int tile_idx = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
int global_idx = tile_idx * PER_TILE_SIZE + tid;
|
||||
|
||||
// ── Phase 1: Load tile into shared memory ──
|
||||
__shared__ float sdata[PER_TILE_SIZE];
|
||||
float val = (global_idx < size) ? input[global_idx] : 0.0f;
|
||||
sdata[tid] = val;
|
||||
__syncthreads();
|
||||
|
||||
// ── Phase 2: Intra-tile inclusive scan (Hillis-Steele) ──
|
||||
for (int offset = 1; offset < PER_TILE_SIZE; offset <<= 1) {
|
||||
float addend = (tid >= offset) ? sdata[tid - offset] : 0.0f;
|
||||
__syncthreads();
|
||||
sdata[tid] += addend;
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Tile aggregate = last element of scanned tile
|
||||
float tile_agg = sdata[PER_TILE_SIZE - 1];
|
||||
|
||||
// ── Phase 3: Publish aggregate and determine prefix ──
|
||||
__shared__ float s_prefix;
|
||||
|
||||
if (tid == 0) {
|
||||
// Publish this tile's aggregate
|
||||
tile_aggregate[tile_idx] = tile_agg;
|
||||
__threadfence();
|
||||
tile_state[tile_idx] = STATE_AGGREGATE;
|
||||
__threadfence();
|
||||
|
||||
if (tile_idx == 0) {
|
||||
// First tile: prefix is 0 (identity for addition)
|
||||
tile_prefix[tile_idx] = 0.0f;
|
||||
__threadfence();
|
||||
tile_state[tile_idx] = STATE_PREFIX;
|
||||
s_prefix = 0.0f;
|
||||
} else {
|
||||
// Decoupled lookback: scan backwards through preceding tiles
|
||||
float running_prefix = 0.0f;
|
||||
int lookback = tile_idx - 1;
|
||||
|
||||
while (true) {
|
||||
unsigned int state = tile_state[lookback];
|
||||
if (state == STATE_PREFIX) {
|
||||
running_prefix += tile_prefix[lookback] + tile_aggregate[lookback];
|
||||
break;
|
||||
} else if (state == STATE_AGGREGATE) {
|
||||
running_prefix += tile_aggregate[lookback];
|
||||
lookback--;
|
||||
}
|
||||
// STATE_INVALID: spin (tile hasn't started yet)
|
||||
}
|
||||
|
||||
tile_prefix[tile_idx] = running_prefix;
|
||||
__threadfence();
|
||||
tile_state[tile_idx] = STATE_PREFIX;
|
||||
s_prefix = running_prefix;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ── Phase 4: Add prefix to all elements and write output ──
|
||||
if (global_idx < size) {
|
||||
output[global_idx] = sdata[tid] + s_prefix;
|
||||
}
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// 4. per_sample — Proportional sampling via binary search
|
||||
//
|
||||
// Each thread generates a random threshold in [0, total_sum) and performs
|
||||
// binary search on the prefix sum array to find the corresponding index.
|
||||
// Also gathers priorities_pa[index] for IS weight computation.
|
||||
//
|
||||
// Launch: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1)
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
// Philox single-round helper (counter-based PRNG)
|
||||
__device__ __forceinline__
|
||||
unsigned int philox_single(unsigned int counter, unsigned int key) {
|
||||
unsigned int hi, lo;
|
||||
lo = counter * 0xD2511F53u;
|
||||
hi = __umulhi(counter, 0xD2511F53u);
|
||||
for (int r = 0; r < 10; r++) {
|
||||
unsigned int t = hi ^ key;
|
||||
hi = lo * 0xCD9E8D57u;
|
||||
lo = __umulhi(lo, 0xCD9E8D57u);
|
||||
lo ^= t;
|
||||
key += 0x9E3779B9u;
|
||||
}
|
||||
return lo ^ hi;
|
||||
}
|
||||
|
||||
extern "C" __global__ void per_sample(
|
||||
const float* __restrict__ prefix_sum,
|
||||
const float* __restrict__ priorities_pa,
|
||||
long long* __restrict__ out_indices,
|
||||
float* __restrict__ out_priorities,
|
||||
float* __restrict__ total_sum_out,
|
||||
unsigned int seed,
|
||||
int size, int batch_size)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= batch_size) return;
|
||||
|
||||
// Thread 0 writes total_sum for IS weight computation
|
||||
float total_sum = prefix_sum[size - 1];
|
||||
if (i == 0) {
|
||||
total_sum_out[0] = total_sum;
|
||||
}
|
||||
|
||||
// Generate uniform random threshold in [0, total_sum)
|
||||
unsigned int bits = philox_single((unsigned int)i, seed);
|
||||
float u = (float)(bits >> 8) * (1.0f / 16777216.0f);
|
||||
float threshold = u * total_sum;
|
||||
|
||||
// Binary search: find smallest idx where prefix_sum[idx] >= threshold
|
||||
int lo = 0, hi = size - 1;
|
||||
while (lo < hi) {
|
||||
int mid = (lo + hi) >> 1;
|
||||
if (prefix_sum[mid] < threshold) {
|
||||
lo = mid + 1;
|
||||
} else {
|
||||
hi = mid;
|
||||
}
|
||||
}
|
||||
|
||||
// Clamp to valid range
|
||||
if (lo >= size) lo = size - 1;
|
||||
if (lo < 0) lo = 0;
|
||||
|
||||
out_indices[i] = (long long)lo;
|
||||
out_priorities[i] = priorities_pa[lo];
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Add to build.rs**
|
||||
|
||||
In `crates/ml-dqn/build.rs`, add `"per_kernels.cu"` to the CUDA source files list (where `"seg_tree_kernel.cu"` currently is):
|
||||
|
||||
Replace:
|
||||
```
|
||||
"seg_tree_kernel.cu",
|
||||
```
|
||||
with:
|
||||
```
|
||||
"per_kernels.cu",
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Verify kernel compilation**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | tail -5`
|
||||
Expected: The cubin compiles (build.rs runs nvcc). Rust compilation may fail because kernel names changed — that's fixed in Task 2.
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-dqn/src/per_kernels.cu crates/ml-dqn/build.rs
|
||||
git commit -m "feat: per_kernels.cu — decoupled lookback scan + binary search sampling
|
||||
|
||||
Four PER kernels replacing the segment tree:
|
||||
1. per_update_pa — priority leaf update from bf16 TD errors directly
|
||||
2. per_insert_pa — insert priorities at max_priority
|
||||
3. per_prefix_scan — Merrill-Garland decoupled lookback inclusive scan
|
||||
4. per_sample — proportional sampling via binary search on prefix sum"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 2: Replace Segment Tree Data Structures in `GpuReplayBuffer`
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`
|
||||
- Delete: `crates/ml-dqn/src/seg_tree_kernel.cu`
|
||||
|
||||
- [ ] **Step 1: Replace `ReplayKernels` segment tree fields with PER kernel fields**
|
||||
|
||||
In `gpu_replay_buffer.rs`, replace lines 90-95:
|
||||
|
||||
```rust
|
||||
// Segment tree kernels (replace prefix_sum + searchsorted + pow_alpha pipeline)
|
||||
seg_tree_update_leaves: CudaFunction,
|
||||
seg_tree_rebuild_level: CudaFunction,
|
||||
seg_tree_insert: CudaFunction,
|
||||
seg_tree_sample: CudaFunction,
|
||||
seg_tree_gather_prios: CudaFunction,
|
||||
```
|
||||
|
||||
with:
|
||||
|
||||
```rust
|
||||
// PER prefix sum kernels (decoupled lookback scan + binary search)
|
||||
per_update_pa: CudaFunction,
|
||||
per_insert_pa: CudaFunction,
|
||||
per_prefix_scan: CudaFunction,
|
||||
per_sample: CudaFunction,
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Replace kernel loading in `ReplayKernels::compile`**
|
||||
|
||||
Replace lines 110-135 (the ST_CUBIN loading + 5 function loads):
|
||||
|
||||
```rust
|
||||
// Load precompiled segment tree kernels cubin
|
||||
static ST_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/seg_tree_kernel.cubin"));
|
||||
let st_mod = ctx.load_cubin(ST_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("seg_tree cubin load: {e}")))?;
|
||||
// ... 5 load_function calls ...
|
||||
```
|
||||
|
||||
with:
|
||||
|
||||
```rust
|
||||
// Load precompiled PER prefix sum kernels cubin
|
||||
static PER_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/per_kernels.cubin"));
|
||||
let per_mod = ctx.load_cubin(PER_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("per cubin load: {e}")))?;
|
||||
let per_ld = |n: &str| -> Result<CudaFunction, MLError> {
|
||||
per_mod.load_function(n).map_err(|e| MLError::ModelError(format!("per {n}: {e}")))
|
||||
};
|
||||
```
|
||||
|
||||
And update the struct init to use:
|
||||
```rust
|
||||
per_update_pa: per_ld("per_update_pa")?,
|
||||
per_insert_pa: per_ld("per_insert_pa")?,
|
||||
per_prefix_scan: per_ld("per_prefix_scan")?,
|
||||
per_sample: per_ld("per_sample")?,
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Replace struct fields in `GpuReplayBuffer`**
|
||||
|
||||
Replace:
|
||||
```rust
|
||||
seg_tree: CudaSlice<f32>,
|
||||
capacity_pow2: usize,
|
||||
```
|
||||
|
||||
with:
|
||||
```rust
|
||||
priorities_pa: CudaSlice<f32>,
|
||||
prefix_sum: CudaSlice<f32>,
|
||||
scan_tile_state: CudaSlice<u32>,
|
||||
scan_tile_aggregate: CudaSlice<f32>,
|
||||
scan_tile_prefix: CudaSlice<f32>,
|
||||
num_scan_tiles: usize,
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Delete cast kernel infrastructure**
|
||||
|
||||
Delete entirely from `gpu_replay_buffer.rs`:
|
||||
- `struct CastKernels` (line 48-50)
|
||||
- `static CAST_KERNELS: OnceLock<...>` (line 52)
|
||||
- `fn get_cast_kernels(...)` (lines 54-71)
|
||||
- `update_td_f32: CudaSlice<f32>` field from struct
|
||||
- The pre-init `get_cast_kernels(stream)?` call in constructor
|
||||
- `use std::sync::OnceLock;` import if now unused
|
||||
|
||||
- [ ] **Step 5: Update constructor to allocate new buffers**
|
||||
|
||||
Replace the segment tree allocation:
|
||||
```rust
|
||||
let cap_pow2 = cap.next_power_of_two();
|
||||
let seg = a32f(stream, 2 * cap_pow2, "seg_tree")?;
|
||||
```
|
||||
|
||||
with:
|
||||
```rust
|
||||
let priorities_pa = a32f(stream, cap, "priorities_pa")?;
|
||||
let prefix_sum = a32f(stream, cap, "prefix_sum")?;
|
||||
let tile_size = 256_usize; // PER_TILE_SIZE in CUDA kernel
|
||||
let num_scan_tiles = cap.div_ceil(tile_size);
|
||||
let scan_tile_state = stream.alloc_zeros::<u32>(num_scan_tiles)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc scan_tile_state: {e}")))?;
|
||||
let scan_tile_aggregate = a32f(stream, num_scan_tiles, "scan_tile_agg")?;
|
||||
let scan_tile_prefix = a32f(stream, num_scan_tiles, "scan_tile_pfx")?;
|
||||
```
|
||||
|
||||
Remove `update_td_f32` allocation. Update struct initializer to use new fields.
|
||||
|
||||
- [ ] **Step 6: Delete `seg_tree_kernel.cu`**
|
||||
|
||||
```bash
|
||||
rm crates/ml-dqn/src/seg_tree_kernel.cu
|
||||
```
|
||||
|
||||
- [ ] **Step 7: Build check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | tail -10`
|
||||
Expected: Errors in methods that still reference old fields — fixed in Task 3.
|
||||
|
||||
- [ ] **Step 8: Commit**
|
||||
|
||||
```bash
|
||||
git add -u crates/ml-dqn/src/
|
||||
git commit -m "refactor: replace segment tree with prefix sum buffers
|
||||
|
||||
Delete seg_tree (67M floats = 256MB), capacity_pow2, CastKernels,
|
||||
OnceLock, get_cast_kernels, update_td_f32 scratch buffer.
|
||||
Add priorities_pa, prefix_sum, scan scratch buffers.
|
||||
Delete seg_tree_kernel.cu entirely."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 3: Rewrite `update_priorities_gpu` — No Cast, No Tree
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`
|
||||
|
||||
- [ ] **Step 1: Delete `rebuild_tree` method**
|
||||
|
||||
Delete the entire `fn rebuild_tree(...)` method.
|
||||
|
||||
- [ ] **Step 2: Rewrite `update_priorities_gpu`**
|
||||
|
||||
Replace the entire method body with:
|
||||
|
||||
```rust
|
||||
pub fn update_priorities_gpu(
|
||||
&mut self,
|
||||
indices: &CudaSlice<u32>,
|
||||
td_errors: &CudaSlice<half::bf16>,
|
||||
bs: usize,
|
||||
ext_stream: Option<&Arc<CudaStream>>,
|
||||
) -> Result<(), MLError> {
|
||||
if bs == 0 { return Ok(()); }
|
||||
let stream_owned = ext_stream.cloned().unwrap_or_else(|| Arc::clone(&self.stream));
|
||||
let stream = &stream_owned;
|
||||
|
||||
let (al, ep, bsi) = (self.config.alpha, self.config.epsilon, bs as i32);
|
||||
let cap_i = self.config.capacity as i32;
|
||||
|
||||
stream.memset_zeros(&mut self.update_batch_max)
|
||||
.map_err(|e| MLError::ModelError(format!("zero batch_max: {e}")))?;
|
||||
|
||||
// per_update_pa: reads bf16 td_errors directly, writes priorities + priorities_pa
|
||||
unsafe {
|
||||
stream.launch_builder(&self.kernels.per_update_pa)
|
||||
.arg(&self.priorities_pa)
|
||||
.arg(&self.priorities)
|
||||
.arg(&self.update_batch_max)
|
||||
.arg(indices)
|
||||
.arg(td_errors)
|
||||
.arg(&al).arg(&ep).arg(&cap_i).arg(&bsi)
|
||||
.launch(lcfg(bs))
|
||||
.map_err(|e| MLError::ModelError(format!("per_update_pa: {e}")))?;
|
||||
}
|
||||
|
||||
// Merge batch max into pending (existing logic, unchanged)
|
||||
match self.pending_max_priority.take() {
|
||||
Some(prev) => {
|
||||
unsafe {
|
||||
stream.launch_builder(&self.kernels.max_of_two_f32)
|
||||
.arg(&self.update_max_merge).arg(&prev).arg(&self.update_batch_max)
|
||||
.launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
|
||||
.map_err(|e| MLError::ModelError(format!("max_of_two: {e}")))?;
|
||||
}
|
||||
self.pending_max_priority = Some(std::mem::replace(&mut self.update_max_merge, prev));
|
||||
}
|
||||
None => {
|
||||
let fresh = a32f(stream, 1, "bm_epoch")?;
|
||||
self.pending_max_priority = Some(std::mem::replace(&mut self.update_batch_max, fresh));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Update `update_priorities_gpu_raw` to pass bf16 directly**
|
||||
|
||||
The raw variant no longer needs the bf16→f32 conversion:
|
||||
|
||||
```rust
|
||||
pub fn update_priorities_gpu_raw(
|
||||
&mut self,
|
||||
_indices_ptr: u64,
|
||||
td_errors: &CudaSlice<half::bf16>,
|
||||
bs: usize,
|
||||
ext_stream: Option<&Arc<CudaStream>>,
|
||||
) -> Result<(), MLError> {
|
||||
let idx_slice = unsafe { &*(&self.sample_indices_u32 as *const CudaSlice<u32>) };
|
||||
self.update_priorities_gpu(idx_slice, td_errors, bs, ext_stream)
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Build check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | tail -10`
|
||||
Expected: Errors in `sample_proportional` and `insert_batch` — fixed in Tasks 4-5.
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-dqn/src/gpu_replay_buffer.rs
|
||||
git commit -m "refactor: per_update_pa accepts bf16 directly, no cast kernel
|
||||
|
||||
Eliminates the bf16→f32 cast kernel launch, OnceLock infrastructure,
|
||||
and update_td_f32 scratch buffer. The per_update_pa CUDA kernel reads
|
||||
bf16 td_errors via __bfloat162float() inline."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Rewrite `sample_proportional` — Prefix Scan + Binary Search
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`
|
||||
|
||||
- [ ] **Step 1: Rewrite the sampling section**
|
||||
|
||||
In `sample_proportional`, replace the segment tree sampling steps (tree sample + gather prios + DtoD root copy) with prefix scan + binary search:
|
||||
|
||||
Replace steps 1, 4, 5 (seg_tree_sample, seg_tree_gather_prios, DtoD tree[1] copy) with:
|
||||
|
||||
```rust
|
||||
// Step 1: Zero scan scratch and run prefix scan
|
||||
self.stream.memset_zeros(&mut self.scan_tile_state)
|
||||
.map_err(|e| MLError::ModelError(format!("zero scan state: {e}")))?;
|
||||
|
||||
let size_i = self.size as i32;
|
||||
let scan_blocks = self.num_scan_tiles as u32;
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.kernels.per_prefix_scan)
|
||||
.arg(&self.priorities_pa)
|
||||
.arg(&mut self.prefix_sum)
|
||||
.arg(&self.scan_tile_state)
|
||||
.arg(&self.scan_tile_aggregate)
|
||||
.arg(&self.scan_tile_prefix)
|
||||
.arg(&size_i)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (scan_blocks.max(1), 1, 1),
|
||||
block_dim: (256, 1, 1), // PER_TILE_SIZE
|
||||
shared_mem_bytes: 256 * 4, // PER_TILE_SIZE * sizeof(float)
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("per_prefix_scan: {e}")))?;
|
||||
}
|
||||
|
||||
// Step 2: Sample + gather priorities_pa in one kernel
|
||||
self.rng_step = self.rng_step.wrapping_add(1);
|
||||
let seed = self.rng_step;
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.kernels.per_sample)
|
||||
.arg(&self.prefix_sum)
|
||||
.arg(&self.priorities_pa)
|
||||
.arg(&mut self.sample_indices_i64)
|
||||
.arg(&mut self.sample_priorities)
|
||||
.arg(&mut self.total_sum_buf)
|
||||
.arg(&seed)
|
||||
.arg(&size_i)
|
||||
.arg(&bsi)
|
||||
.launch(lcfg(batch_size))
|
||||
.map_err(|e| MLError::ModelError(format!("per_sample: {e}")))?;
|
||||
}
|
||||
```
|
||||
|
||||
Delete the old steps:
|
||||
- `seg_tree_sample` launch
|
||||
- `seg_tree_gather_prios` launch
|
||||
- The DtoD copy of `tree[1]` → `total_sum_buf`
|
||||
|
||||
Steps 2-8 (i64_to_u32, gather, IS weights, normalize) remain **unchanged**.
|
||||
|
||||
- [ ] **Step 2: Build check**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml-dqn --lib 2>&1 | tail -10`
|
||||
Expected: Errors in `insert_batch` — fixed in Task 5.
|
||||
|
||||
- [ ] **Step 3: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-dqn/src/gpu_replay_buffer.rs
|
||||
git commit -m "feat: prefix scan + binary search sampling replaces segment tree
|
||||
|
||||
per_prefix_scan: Merrill-Garland decoupled lookback inclusive scan.
|
||||
per_sample: binary search + priority gathering in one kernel.
|
||||
2 launches instead of 3 (tree sample + gather + DtoD copy)."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Rewrite `insert_batch` / `insert_batch_bf16` + `clear()`
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`
|
||||
|
||||
- [ ] **Step 1: Replace `seg_tree_insert` + `rebuild_tree` in `insert_batch`**
|
||||
|
||||
Find the `seg_tree_insert` launch + `rebuild_tree` call in `insert_batch`. Replace with:
|
||||
|
||||
```rust
|
||||
// Write priority^alpha to priorities_pa (no tree, no propagation)
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.kernels.per_insert_pa)
|
||||
.arg(&self.priorities_pa).arg(&idx_buf).arg(&pt).arg(&al)
|
||||
.arg(&cap_i).arg(&bsi)
|
||||
.launch(lcfg(eff))
|
||||
.map_err(|e| MLError::ModelError(format!("per_insert_pa: {e}")))?;
|
||||
}
|
||||
// No rebuild needed — prefix_sum is rebuilt in sample_proportional
|
||||
```
|
||||
|
||||
where `cap_i = self.config.capacity as i32`.
|
||||
|
||||
- [ ] **Step 2: Same replacement in `insert_batch_bf16`**
|
||||
|
||||
Identical change in the bf16 insert path.
|
||||
|
||||
- [ ] **Step 3: Update `clear()`**
|
||||
|
||||
Replace:
|
||||
```rust
|
||||
self.stream.memset_zeros(&mut self.seg_tree)
|
||||
.map_err(|e| MLError::ModelError(format!("seg_tree clear: {e}")))?;
|
||||
```
|
||||
with:
|
||||
```rust
|
||||
self.stream.memset_zeros(&mut self.priorities_pa)
|
||||
.map_err(|e| MLError::ModelError(format!("priorities_pa clear: {e}")))?;
|
||||
self.stream.memset_zeros(&mut self.prefix_sum)
|
||||
.map_err(|e| MLError::ModelError(format!("prefix_sum clear: {e}")))?;
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Remove all remaining `seg_tree` / `capacity_pow2` references**
|
||||
|
||||
Search for any remaining references to `seg_tree`, `capacity_pow2`, `rebuild_tree`, `CastKernels`, `get_cast_kernels`, `update_td_f32` in `gpu_replay_buffer.rs` and delete them.
|
||||
|
||||
- [ ] **Step 5: Remove all diagnostic `eprintln` from PER code**
|
||||
|
||||
Search for `PER_DIAG` and delete all diagnostic eprintln calls in `gpu_replay_buffer.rs`.
|
||||
|
||||
- [ ] **Step 6: Build and verify**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml --lib 2>&1 | tail -5`
|
||||
Expected: Clean compilation.
|
||||
|
||||
- [ ] **Step 7: Run smoke tests**
|
||||
|
||||
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result'`
|
||||
Expected: `test result: ok. 19 passed; 0 failed;`
|
||||
|
||||
- [ ] **Step 8: Commit**
|
||||
|
||||
```bash
|
||||
git add -u crates/ml-dqn/src/
|
||||
git commit -m "feat: complete PER prefix sum replacement — zero segment tree code
|
||||
|
||||
insert_batch uses per_insert_pa (leaf write only, no rebuild).
|
||||
clear() zeroes priorities_pa + prefix_sum.
|
||||
All seg_tree, capacity_pow2, CastKernels, rebuild_tree references deleted.
|
||||
All PER_DIAG diagnostic eprintln removed."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 6: Clean Up Integration Tests + External Callers
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml/tests/gpu_per_integration_test.rs`
|
||||
- Modify: `crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs`
|
||||
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
|
||||
|
||||
- [ ] **Step 1: Fix integration tests**
|
||||
|
||||
In `gpu_per_integration_test.rs`, update all calls to `update_priorities_gpu` and `update_priorities_gpu_raw` to match the new signatures (bf16 td_errors directly, ext_stream parameter).
|
||||
|
||||
- [ ] **Step 2: Fix gpu_residency smoke test**
|
||||
|
||||
Same signature update in `gpu_residency.rs`.
|
||||
|
||||
- [ ] **Step 3: Remove diagnostic sync in fused_training.rs**
|
||||
|
||||
In `fused_training.rs`, remove the debug `cuStreamSynchronize` + `check_err` after PER update (the `H100_HANG4: syncing stream to check errors` block).
|
||||
|
||||
- [ ] **Step 4: Remove all remaining H100 diagnostic eprintln**
|
||||
|
||||
Search for `H100_STEP`, `H100_LOOP`, `H100_HANG4`, `H100_DIAG`, `adam_readback`, `replay_forward`, `PER_DIAG` eprintln calls across the entire `crates/ml/` and `crates/ml-dqn/` and delete them all. Training is confirmed working once smoke tests pass — no more diagnostic noise.
|
||||
|
||||
- [ ] **Step 5: Build, full test, commit**
|
||||
|
||||
Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` (full crate including tests)
|
||||
Expected: Clean.
|
||||
|
||||
Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep 'test result'`
|
||||
Expected: `test result: ok. 19 passed; 0 failed;`
|
||||
|
||||
```bash
|
||||
git add -u
|
||||
git commit -m "cleanup: fix tests, remove all H100 diagnostic eprintln
|
||||
|
||||
Integration tests updated for new PER API (bf16 td_errors, ext_stream).
|
||||
All diagnostic eprintln removed (H100_STEP, PER_DIAG, adam_readback, etc).
|
||||
Debug cuStreamSynchronize removed from fused training path."
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Deploy to H100 and Validate
|
||||
|
||||
**Files:** None (testing only)
|
||||
|
||||
- [ ] **Step 1: Push and deploy**
|
||||
|
||||
```bash
|
||||
git push origin main
|
||||
./scripts/argo-train.sh dqn --baseline --watch
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Monitor training progression**
|
||||
|
||||
Verify:
|
||||
- Training starts (no kernel loading hangs)
|
||||
- Step 0 completes (PER update + sample works)
|
||||
- Step 1 appears (training loop progresses)
|
||||
- Epoch 1 completes with timing
|
||||
- Q-stats and loss reported normally
|
||||
|
||||
- [ ] **Step 3: Compare epoch timing**
|
||||
|
||||
Check the `Training step breakdown` log for per-step timing. Expected: sample + fused + guard timings with no multi-second PER stalls.
|
||||
221
docs/superpowers/specs/2026-04-06-per-prefix-sum-design.md
Normal file
221
docs/superpowers/specs/2026-04-06-per-prefix-sum-design.md
Normal file
@@ -0,0 +1,221 @@
|
||||
# PER Prefix Sum — Replace Segment Tree with Decoupled Lookback Scan
|
||||
|
||||
## Problem
|
||||
|
||||
The PER segment tree (`seg_tree: CudaSlice<f32>`, 67M floats = 256MB) hangs on H100 during priority updates. The tree structure requires either:
|
||||
- atomicAdd propagation from leaves to root (8192 threads contending on root = minutes of L2 cache serialization)
|
||||
- Full tree rebuild (33M nodes × 25 levels = seconds of wasted compute for 8192 changed leaves)
|
||||
|
||||
Both approaches fail at H100 scale because the tree was designed for CPU-era single-element updates, not massively parallel GPU workloads.
|
||||
|
||||
## Solution
|
||||
|
||||
Replace the segment tree with a flat prefix sum array using NVIDIA's decoupled lookback scan algorithm (Merrill & Garland, 2016). Single-pass O(N) parallel scan + O(B log N) binary search sampling in 2 kernel launches. ~42µs per PER cycle on H100.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Data Structures
|
||||
|
||||
**Delete:**
|
||||
- `seg_tree: CudaSlice<f32>` — 67M floats (256MB). Gone.
|
||||
- `capacity_pow2: usize` — only needed for tree indexing. Gone.
|
||||
|
||||
**Keep:**
|
||||
- `priorities: CudaSlice<f32>` — `[capacity]` raw priorities `|td|^alpha + eps`
|
||||
|
||||
**Add:**
|
||||
- `priorities_pa: CudaSlice<f32>` — `[capacity]` priority^alpha values (replaces tree leaves)
|
||||
- `prefix_sum: CudaSlice<f32>` — `[capacity]` inclusive prefix sum of `priorities_pa[0..size]`
|
||||
- `scan_status: CudaSlice<u32>` — `[num_tiles]` decoupled lookback state flags (zeroed before each scan)
|
||||
- `scan_values: CudaSlice<f32>` — `[2 × num_tiles]` aggregate + prefix values (packed)
|
||||
|
||||
Where `num_tiles = ceil(capacity / TILE_SIZE)`, `TILE_SIZE = 1024`.
|
||||
|
||||
### Kernels
|
||||
|
||||
**Delete:** `seg_tree_kernel.cu` — all 5 kernels (`seg_tree_update_leaves`, `seg_tree_rebuild_level`, `seg_tree_insert`, `seg_tree_sample`, `seg_tree_gather_prios`).
|
||||
|
||||
**Add** in new file `per_kernels.cu`:
|
||||
|
||||
#### 1. `per_update_pa` — Priority leaf update (O(B), no propagation)
|
||||
|
||||
```cuda
|
||||
extern "C" __global__ void per_update_pa(
|
||||
float* __restrict__ priorities_pa, // [capacity] — write priority^alpha
|
||||
float* __restrict__ priorities, // [capacity] — write raw priority
|
||||
float* __restrict__ batch_max, // [1] atomicMax accumulator
|
||||
const unsigned int* __restrict__ indices, // [batch_size] buffer indices
|
||||
const float* __restrict__ td_errors, // [batch_size] f32 TD errors
|
||||
float alpha, float epsilon,
|
||||
int capacity, int batch_size)
|
||||
```
|
||||
|
||||
Per thread: compute `p = |td|^alpha + eps`, write `priorities[idx] = p`, compute `pa = p^alpha`, write `priorities_pa[idx] = pa`, atomicMax batch_max. No tree, no propagation, O(1) per thread.
|
||||
|
||||
Launch: `grid=(ceil(B/256), 1, 1), block=(256, 1, 1)`.
|
||||
|
||||
#### 2. `per_prefix_scan` — Decoupled lookback inclusive scan (O(N), single pass)
|
||||
|
||||
```cuda
|
||||
extern "C" __global__ void per_prefix_scan(
|
||||
const float* __restrict__ priorities_pa, // [capacity] input
|
||||
float* __restrict__ prefix_sum, // [capacity] output (inclusive scan)
|
||||
volatile unsigned int* __restrict__ tile_status, // [num_tiles] state flags
|
||||
volatile float* __restrict__ tile_values, // [2*num_tiles] [agg, prefix] pairs
|
||||
int size) // actual valid entries
|
||||
```
|
||||
|
||||
Algorithm:
|
||||
1. Each block loads a tile of TILE_SIZE elements from `priorities_pa`
|
||||
2. Intra-block inclusive scan via warp shuffles + shared memory (Blelloch/Hillis-Steele)
|
||||
3. Block 0: prefix = 0, publish tile aggregate + prefix, set status = PREFIX_AVAILABLE
|
||||
4. Block i>0: spin on `tile_status[i-1]`, look back through chain until finding PREFIX_AVAILABLE, accumulate aggregates, compute own prefix, publish, set status = PREFIX_AVAILABLE
|
||||
5. Add prefix to all scanned elements, write to `prefix_sum`
|
||||
|
||||
Launch: `grid=(ceil(size/TILE_SIZE), 1, 1), block=(TILE_SIZE, 1, 1)`, where `TILE_SIZE = 1024`.
|
||||
|
||||
For N=20M: 19532 blocks, each processing 1024 elements. At H100's 3TB/s bandwidth: 20M × 4B × 2 (read+write) = 160MB → ~53µs. Actual: ~20-30µs (L2 cache + SM occupancy).
|
||||
|
||||
#### 3. `per_sample` — Proportional sampling via binary search (O(B log N))
|
||||
|
||||
```cuda
|
||||
extern "C" __global__ void per_sample(
|
||||
const float* __restrict__ prefix_sum, // [size] inclusive prefix sum
|
||||
const float* __restrict__ priorities_pa, // [capacity] for pa gathering
|
||||
long long* __restrict__ out_indices, // [batch_size] output indices (i64)
|
||||
float* __restrict__ out_priorities, // [batch_size] sampled pa values
|
||||
float* __restrict__ total_sum_out, // [1] = prefix_sum[size-1]
|
||||
unsigned int seed, // Philox RNG seed
|
||||
int size, int batch_size)
|
||||
```
|
||||
|
||||
Per thread:
|
||||
1. Thread 0 writes `total_sum_out[0] = prefix_sum[size-1]`
|
||||
2. Generate `u = Philox(tid, seed) * prefix_sum[size-1]` (uniform in [0, total_sum))
|
||||
3. Binary search `prefix_sum[0..size]` for smallest `i` where `prefix_sum[i] >= u`
|
||||
4. Write `out_indices[tid] = i`
|
||||
5. Write `out_priorities[tid] = priorities_pa[i]` (gather for IS weights)
|
||||
|
||||
Launch: `grid=(ceil(B/256), 1, 1), block=(256, 1, 1)`.
|
||||
|
||||
For B=8192, N=20M: 8192 threads × 25 binary search steps × 1 global memory read each = 200K reads. Cache-friendly (binary search converges to same regions for similar thresholds).
|
||||
|
||||
#### 4. `per_insert_pa` — Insert priorities at max_priority (O(B), for insert_batch)
|
||||
|
||||
```cuda
|
||||
extern "C" __global__ void per_insert_pa(
|
||||
float* __restrict__ priorities_pa, // [capacity]
|
||||
const unsigned int* __restrict__ indices, // [batch_size]
|
||||
const float* __restrict__ priorities, // [batch_size] raw priority values (max_priority fill)
|
||||
float alpha,
|
||||
int capacity, int batch_size)
|
||||
```
|
||||
|
||||
Per thread: `priorities_pa[indices[i]] = powf(priorities[i], alpha)`. No propagation — caller runs `per_prefix_scan` after.
|
||||
|
||||
Launch: `grid=(ceil(B/256), 1, 1), block=(256, 1, 1)`.
|
||||
|
||||
### Per-Step Training Flow
|
||||
|
||||
```
|
||||
1. per_update_pa — 8192 leaf writes ~5µs
|
||||
2. per_prefix_scan — scan priorities_pa[0..size] ~25µs
|
||||
3. per_sample — binary search + gather pa ~10µs
|
||||
(replaces seg_tree_sample + seg_tree_gather_prios + DtoD root copy)
|
||||
4. i64_to_u32 — existing kernel ~2µs
|
||||
5. gather states/actions/rewards/dones/episodes existing kernels
|
||||
6. is_weights_f32 — existing kernel, total_sum_out ~3µs
|
||||
7. reduce_max + normalize_weights existing kernels
|
||||
Total PER: ~45µs (was: infinite hang or 10+ seconds)
|
||||
```
|
||||
|
||||
### Insert Flow
|
||||
|
||||
```
|
||||
1. fill priorities_pa slots with max_priority^alpha (per_insert_pa)
|
||||
2. fill priorities slots with max_priority (existing scatter_insert_f32)
|
||||
3. per_prefix_scan to rebuild prefix_sum (full scan)
|
||||
```
|
||||
|
||||
No `rebuild_tree`, no 25 kernel launches. One scan kernel.
|
||||
|
||||
### `sample_proportional` Changes
|
||||
|
||||
Replace steps 1, 4, 5 (tree sample, gather prios, DtoD root copy) with:
|
||||
|
||||
```rust
|
||||
// Step 0: Zero scan scratch
|
||||
stream.memset_zeros(&mut self.scan_status)?;
|
||||
|
||||
// Step 1: Prefix scan
|
||||
unsafe {
|
||||
stream.launch_builder(&self.kernels.per_prefix_scan)
|
||||
.arg(&self.priorities_pa).arg(&mut self.prefix_sum)
|
||||
.arg(&self.scan_status).arg(&self.scan_values)
|
||||
.arg(&(self.size as i32))
|
||||
.launch(scan_config)?;
|
||||
}
|
||||
|
||||
// Step 2: Sample + gather
|
||||
unsafe {
|
||||
stream.launch_builder(&self.kernels.per_sample)
|
||||
.arg(&self.prefix_sum).arg(&self.priorities_pa)
|
||||
.arg(&mut self.sample_indices_i64).arg(&mut self.sample_priorities)
|
||||
.arg(&mut self.total_sum_buf)
|
||||
.arg(&seed).arg(&(self.size as i32)).arg(&bsi)
|
||||
.launch(sample_config)?;
|
||||
}
|
||||
```
|
||||
|
||||
Steps 3-8 (gather, IS weights, normalize) remain unchanged.
|
||||
|
||||
### `update_priorities_gpu` Changes
|
||||
|
||||
Replace `seg_tree_update_leaves` + `rebuild_tree/atomicAdd` with:
|
||||
1. bf16→f32 cast (existing, uses pre-allocated scratch)
|
||||
2. `per_update_pa` kernel (leaf write only, O(B))
|
||||
3. Batch max merge (existing)
|
||||
|
||||
No tree propagation, no rebuild. The prefix sum is rebuilt in the next `sample_proportional` call.
|
||||
|
||||
### Dead Code Removal
|
||||
|
||||
| Item | Action |
|
||||
|------|--------|
|
||||
| `seg_tree_kernel.cu` | DELETE file |
|
||||
| `seg_tree: CudaSlice<f32>` field | DELETE |
|
||||
| `capacity_pow2: usize` field | DELETE |
|
||||
| `seg_tree_update_leaves` kernel | DELETE |
|
||||
| `seg_tree_rebuild_level` kernel | DELETE |
|
||||
| `seg_tree_insert` kernel | DELETE |
|
||||
| `seg_tree_sample` kernel | DELETE |
|
||||
| `seg_tree_gather_prios` kernel | DELETE |
|
||||
| `rebuild_tree()` method | DELETE |
|
||||
| `ReplayKernels::seg_tree_*` fields | DELETE (replace with `per_*` fields) |
|
||||
| `ST_CUBIN` static | DELETE (replace with `PER_CUBIN`) |
|
||||
| All `seg_tree` references in build.rs | Replace with `per_kernels.cu` |
|
||||
| All `capacity_pow2` references | Replace with `capacity` |
|
||||
|
||||
### VRAM Savings
|
||||
|
||||
| Before | After | Savings |
|
||||
|--------|-------|---------|
|
||||
| `seg_tree`: 2 × 33M × 4B = 256MB | `priorities_pa`: 20M × 4B = 80MB | |
|
||||
| | `prefix_sum`: 20M × 4B = 80MB | |
|
||||
| | `scan_scratch`: ~80KB | |
|
||||
| **Total: 256MB** | **Total: 160MB** | **96MB freed** |
|
||||
|
||||
### Testing
|
||||
|
||||
- All 19 smoke tests must pass
|
||||
- Verify correct PER sampling distribution: high-priority experiences sampled more frequently
|
||||
- Verify IS weights are correctly computed (sum to ~batch_size after normalization)
|
||||
- H100 deployment: training progresses past step 0, epochs complete
|
||||
|
||||
### Success Criteria
|
||||
|
||||
- Zero segment tree code in codebase
|
||||
- PER update + sample in 2-3 kernel launches (not 25+)
|
||||
- Total PER cycle < 50µs on H100
|
||||
- No hangs, no contention, no cross-stream races
|
||||
- Correct proportional sampling (validated by smoke tests)
|
||||
Reference in New Issue
Block a user