diff --git a/crates/ml-dqn/src/branching.rs b/crates/ml-dqn/src/branching.rs index e4ac88534..04622e6f5 100644 --- a/crates/ml-dqn/src/branching.rs +++ b/crates/ml-dqn/src/branching.rs @@ -529,6 +529,10 @@ impl BranchingDuelingQNetwork { /// Each branch outputs `[batch, n_d * num_atoms]`, reshaped to `[batch, n_d, num_atoms]`. /// Log-softmax is applied along the atoms dimension (dim=2). /// Expected Q = sum(softmax(logits) * z) stored in `advantages` for greedy selection. + /// + /// PERF: Cold-path implementation with GPU→CPU→GPU roundtrips for log-softmax + /// and expectation. For performance-critical forward passes, use the fused + /// CUDA kernel path in `GpuDqnTrainer::launch_cublas_forward`. fn forward_distributional( &self, h: &GpuTensor, diff --git a/crates/ml-dqn/src/quantile_regression.rs b/crates/ml-dqn/src/quantile_regression.rs index 92985ef70..334ca15eb 100644 --- a/crates/ml-dqn/src/quantile_regression.rs +++ b/crates/ml-dqn/src/quantile_regression.rs @@ -300,6 +300,7 @@ impl QuantileNetwork { /// * `target` - Target quantile values [batch, `num_quantiles`] /// * `taus` - Quantile fractions [batch, `num_quantiles`] /// * `kappa` - Huber threshold +/// * `stream` - CUDA stream from the caller (no internal context creation) /// /// # Returns /// Mean quantile Huber loss (scalar `GpuTensor`) @@ -308,34 +309,34 @@ pub fn quantile_huber_loss( target: &GpuTensor, taus: &GpuTensor, kappa: f32, + stream: &Arc, ) -> Result { - // Create CUDA context + stream ONCE, reuse for all operations. - let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| { - MLError::DeviceError(format!("CUDA context: {e}")) - })?; - let stream = ctx.new_stream().map_err(|e| { - MLError::DeviceError(format!("CUDA stream: {e}")) - })?; - // Compute quantile Huber loss via host-side calculation (cold path). // predicted, target, taus: [batch, num_quantiles] // Returns: scalar mean loss - let per_sample = quantile_huber_loss_per_sample(predicted, target, taus, kappa)?; + let per_sample = quantile_huber_loss_per_sample(predicted, target, taus, kappa, stream)?; // per_sample: [batch] // Mean over batch let batch = per_sample.numel(); if batch == 0 { - return GpuTensor::scalar(0.0, &stream); + return GpuTensor::scalar(0.0, stream); } - let host = per_sample.to_host(&stream)?; + let host = per_sample.to_host(stream)?; let mean_loss: f32 = host.iter().sum::() / batch as f32; - GpuTensor::scalar(mean_loss, &stream) + GpuTensor::scalar(mean_loss, stream) } /// Per-sample quantile Huber loss for PER importance-sampling weight correction. /// /// Identical to [`quantile_huber_loss`] except the batch dimension is preserved. /// +/// # Arguments +/// * `predicted` - Predicted quantile values [batch, `num_quantiles`] +/// * `target` - Target quantile values [batch, `num_quantiles`] +/// * `taus` - Quantile fractions [batch, `num_quantiles`] +/// * `kappa` - Huber threshold +/// * `stream` - CUDA stream from the caller (no internal context creation) +/// /// # Returns /// Per-sample loss tensor `[batch]` (mean over quantiles, **not** over batch) pub fn quantile_huber_loss_per_sample( @@ -343,6 +344,7 @@ pub fn quantile_huber_loss_per_sample( target: &GpuTensor, taus: &GpuTensor, kappa: f32, + stream: &Arc, ) -> Result { // Per-sample quantile Huber loss (cold path). // predicted, target, taus: [batch, num_quantiles] @@ -351,16 +353,9 @@ pub fn quantile_huber_loss_per_sample( MLError::ModelError(format!("quantile_huber_loss_per_sample predicted dims2: {e}")) })?; - let ctx = cudarc::driver::CudaContext::new(0).map_err(|e| { - MLError::DeviceError(format!("CUDA context: {e}")) - })?; - let stream = ctx.new_stream().map_err(|e| { - MLError::DeviceError(format!("CUDA stream: {e}")) - })?; - - let pred_host = predicted.to_host(&stream)?; - let tgt_host = target.to_host(&stream)?; - let tau_host = taus.to_host(&stream)?; + let pred_host = predicted.to_host(stream)?; + let tgt_host = target.to_host(stream)?; + let tau_host = taus.to_host(stream)?; let mut result = Vec::with_capacity(batch); for b in 0..batch { diff --git a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu index 36919dcb1..78bf651c4 100644 --- a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu @@ -17,6 +17,30 @@ #define PORTFOLIO_STATE_SIZE 8 #include "trade_physics.cuh" +// Capital floor breach: full episode restart for this window. +// 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( + 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; + step_returns[w * max_len + current_step] = step_ret; + 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] = 0.0f; + portfolio_state[ps + 2] = new_capital; + portfolio_state[ps + 3] = 0.0f; + portfolio_state[ps + 4] = new_capital; + 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 float* __restrict__ prices, // [n_windows * max_len * 4] (OHLC) @@ -99,18 +123,9 @@ extern "C" __global__ void backtest_env_step( } float liq_value = cash; float liq_ret = (value > 0.01f) ? (liq_value - value) / value : 0.0f; - step_rewards[w] = liq_ret; - step_returns[w * max_len + current_step] = liq_ret; - actions_history[w * max_len + current_step] = b0_size / 2; // Flat - portfolio_state[ps + 0] = liq_value; // value = new initial capital - portfolio_state[ps + 1] = 0.0f; // position = flat - portfolio_state[ps + 2] = liq_value; // cash = new initial capital - portfolio_state[ps + 3] = 0.0f; // entry_price = none - portfolio_state[ps + 4] = liq_value; // max_equity = RESET (was stale) - portfolio_state[ps + 5] = 0.0f; // hold_time = 0 - portfolio_state[ps + 6] = 0.0f; // cum_return = RESET (was accumulating) - portfolio_state[ps + 7] = 0.0f; // step_count = RESET (was incrementing) - done_flags[w] = 1; + 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); return; } @@ -197,18 +212,9 @@ extern "C" __global__ void backtest_env_step( entry_price = 0.0f; } float step_ret = (value > 0.01f) ? (new_value - value) / value : 0.0f; - step_rewards[w] = step_ret; - step_returns[w * max_len + current_step] = step_ret; - actions_history[w * max_len + current_step] = b0_size / 2; - portfolio_state[ps + 0] = new_value; // value = new initial capital - portfolio_state[ps + 1] = 0.0f; // position = flat - portfolio_state[ps + 2] = new_value; // cash = new initial capital - portfolio_state[ps + 3] = 0.0f; // entry_price = none - portfolio_state[ps + 4] = new_value; // max_equity = RESET - portfolio_state[ps + 5] = 0.0f; // hold_time = 0 - portfolio_state[ps + 6] = 0.0f; // cum_return = RESET - portfolio_state[ps + 7] = 0.0f; // step_count = RESET - done_flags[w] = 1; + 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); return; }