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

fix: address all expert review findings

Browse Source 
- quantile_regression.rs: accept stream from caller (was creating context per call)
- backtest_env_kernel.cu: DRY floor breach into handle_capital_floor_breach() device fn
- backtest_env_kernel.cu: clarifying comments on intentional cum_return/step_count reset
- branching.rs: PERF comment on forward_distributional cold-path roundtrips

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-27 19:59:05 +01:00
parent be8f3310d5
commit e0fe791a50
3 changed files with 51 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/ml-dqn/src/branching.rs
View File

@@ -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,

39
crates/ml-dqn/src/quantile_regression.rs
View File

@@ -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<CudaStream>,
) -> Result<GpuTensor, MLError> {
// 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::<f32>() / 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<CudaStream>,
) -> Result<GpuTensor, MLError> {
// 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 {

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

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