From d18b386be31fe985effa00cd0a7e5016117880e9 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 7 Apr 2026 13:58:11 +0200 Subject: [PATCH] fix: promote entire backtest pipeline from bf16 to f32 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../src/cuda_pipeline/backtest_env_kernel.cu | 272 +++---- .../cuda_pipeline/backtest_gather_kernel.cu | 30 +- .../cuda_pipeline/backtest_metrics_kernel.cu | 178 +++-- .../cuda_pipeline/gpu_backtest_evaluator.rs | 50 +- .../plans/2026-04-06-per-prefix-sum.md | 746 ++++++++++++++++++ .../specs/2026-04-06-per-prefix-sum-design.md | 221 ++++++ 6 files changed, 1227 insertions(+), 270 deletions(-) create mode 100644 docs/superpowers/plans/2026-04-06-per-prefix-sum.md create mode 100644 docs/superpowers/specs/2026-04-06-per-prefix-sum-design.md diff --git a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu index 684a7e618..65b7946ce 100644 --- a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu @@ -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; } diff --git a/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu index 0bd315c70..0e04033b5 100644 --- a/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu index 5d68536f5..1cbe55535 100644 --- a/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index 741f32117..3a92a5d70 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -251,14 +251,14 @@ pub struct GpuBacktestEvaluator { gather_kernel: CudaFunction, // Uploaded data (read-only; persists across the step loop) - prices_buf: CudaSlice, // [n_windows * max_len * 4] + prices_buf: CudaSlice, // [n_windows * max_len * 4] features_buf: CudaSlice, // [n_windows * max_len * feat_dim] window_lens_buf: CudaSlice, // [n_windows] // Mutable state (written by kernels each step) - portfolio_buf: CudaSlice, // [n_windows * 8] - step_rewards_buf: CudaSlice, // [n_windows] - step_returns_buf: CudaSlice, // [n_windows * max_len] + portfolio_buf: CudaSlice, // [n_windows * 8] + step_rewards_buf: CudaSlice, // [n_windows] + step_returns_buf: CudaSlice, // [n_windows * max_len] done_buf: CudaSlice, // [n_windows] actions_buf: CudaSlice, // [n_windows] actions_history_buf: CudaSlice, // [n_windows * max_len] @@ -267,7 +267,7 @@ pub struct GpuBacktestEvaluator { states_buf: CudaSlice, // [n_windows * state_dim] (8-aligned) // Output buffer (written by metrics kernel) - metrics_buf: CudaSlice, // [n_windows * 14] + metrics_buf: CudaSlice, // [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::(n_windows) + .alloc_zeros::(n_windows) .map_err(|e| MLError::ModelError(format!("step_rewards alloc: {e}")))?; let step_returns_buf = stream - .alloc_zeros::(n_windows * max_len) + .alloc_zeros::(n_windows * max_len) .map_err(|e| MLError::ModelError(format!("step_returns alloc: {e}")))?; let done_buf = stream .alloc_zeros::(n_windows) @@ -532,7 +534,7 @@ impl GpuBacktestEvaluator { .alloc_zeros::(n_windows * max_len) .map_err(|e| MLError::ModelError(format!("actions_history alloc: {e}")))?; let metrics_buf = stream - .alloc_zeros::(n_windows * 14) + .alloc_zeros::(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::()) as f64 + (flat_prices.len() * std::mem::size_of::() + + flat_features.len() * std::mem::size_of::()) 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::() as u32, + shared_mem_bytes: 256 * 8 * std::mem::size_of::() 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::() as u32, + shared_mem_bytes: 256 * 8 * std::mem::size_of::() 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::() as u32, + shared_mem_bytes: 256 * 8 * std::mem::size_of::() 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::() as u32; + let shmem_bytes = (256_u32 * 6 + 4096) * std::mem::size_of::() 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 = 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): diff --git a/docs/superpowers/plans/2026-04-06-per-prefix-sum.md b/docs/superpowers/plans/2026-04-06-per-prefix-sum.md new file mode 100644 index 000000000..197505269 --- /dev/null +++ b/docs/superpowers/plans/2026-04-06-per-prefix-sum.md @@ -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 + +// ═══════════════════════════════════════════════════════════════════════ +// 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 { + 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, + capacity_pow2: usize, +``` + +with: +```rust + priorities_pa: CudaSlice, + prefix_sum: CudaSlice, + scan_tile_state: CudaSlice, + scan_tile_aggregate: CudaSlice, + scan_tile_prefix: CudaSlice, + 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` 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::(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, + td_errors: &CudaSlice, + bs: usize, + ext_stream: Option<&Arc>, + ) -> 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, + bs: usize, + ext_stream: Option<&Arc>, + ) -> Result<(), MLError> { + let idx_slice = unsafe { &*(&self.sample_indices_u32 as *const CudaSlice) }; + 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. diff --git a/docs/superpowers/specs/2026-04-06-per-prefix-sum-design.md b/docs/superpowers/specs/2026-04-06-per-prefix-sum-design.md new file mode 100644 index 000000000..fc932d937 --- /dev/null +++ b/docs/superpowers/specs/2026-04-06-per-prefix-sum-design.md @@ -0,0 +1,221 @@ +# PER Prefix Sum — Replace Segment Tree with Decoupled Lookback Scan + +## Problem + +The PER segment tree (`seg_tree: CudaSlice`, 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` — 67M floats (256MB). Gone. +- `capacity_pow2: usize` — only needed for tree indexing. Gone. + +**Keep:** +- `priorities: CudaSlice` — `[capacity]` raw priorities `|td|^alpha + eps` + +**Add:** +- `priorities_pa: CudaSlice` — `[capacity]` priority^alpha values (replaces tree leaves) +- `prefix_sum: CudaSlice` — `[capacity]` inclusive prefix sum of `priorities_pa[0..size]` +- `scan_status: CudaSlice` — `[num_tiles]` decoupled lookback state flags (zeroed before each scan) +- `scan_values: CudaSlice` — `[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` 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)