perf: replace backtest evaluator warp-matvec with cuBLAS SGEMM
Last component using old shared-memory-tiling kernel. With hidden_dim=768, shared memory exceeded 49KB → CUDA_ERROR_INVALID_VALUE. Replace with CublasForward::forward_online() + compute_expected_q + experience_action_select kernels. Same cuBLAS pipeline as training and experience collection. Delete backtest_forward_kernel.cu (old warp-matvec kernel). Fix hyperopt bounds test assertions for updated search ranges. 868 tests pass, 0 failed. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -1,101 +0,0 @@
|
||||
/**
|
||||
* Pure-CUDA DQN forward pass for backtest evaluation.
|
||||
*
|
||||
* Requires common_device_functions.cuh prepended via NVRTC source concatenation.
|
||||
* Launch config: grid=(N, 1, 1), block=(32, 1, 1).
|
||||
* One warp per window. All 32 lanes cooperate on the forward pass,
|
||||
* using the same warp-cooperative architecture as the experience collector.
|
||||
*
|
||||
* DQN-specific layer sizes — overridable via NVRTC #define injection.
|
||||
* These #ifndef guards allow the Rust host to inject actual dimensions
|
||||
* via NVRTC source string before compilation.
|
||||
*/
|
||||
#ifndef SHARED_H1
|
||||
#define SHARED_H1 256
|
||||
#endif
|
||||
#ifndef SHARED_H2
|
||||
#define SHARED_H2 256
|
||||
#endif
|
||||
#ifndef VALUE_H
|
||||
#define VALUE_H 128
|
||||
#endif
|
||||
#ifndef ADV_H
|
||||
#define ADV_H 128
|
||||
#endif
|
||||
|
||||
#ifndef SHMEM_MIN
|
||||
#define SHMEM_MIN(a, b) (((a) < (b)) ? (a) : (b))
|
||||
#endif
|
||||
|
||||
/* Distributed vector size: elements per lane for dim distributed across 32 lanes */
|
||||
#ifndef DIST_SIZE
|
||||
#define DIST_SIZE(dim) (((dim) + 31) / 32)
|
||||
#endif
|
||||
|
||||
extern "C" __global__ void backtest_forward_kernel(
|
||||
const float* __restrict__ states, /* [N, STATE_DIM] */
|
||||
/* shared layer weights */
|
||||
const float* __restrict__ w_s1, /* [SHARED_H1, STATE_DIM] */
|
||||
const float* __restrict__ b_s1, /* [SHARED_H1] */
|
||||
const float* __restrict__ w_s2, /* [SHARED_H2, SHARED_H1] */
|
||||
const float* __restrict__ b_s2, /* [SHARED_H2] */
|
||||
/* value head weights */
|
||||
const float* __restrict__ w_v1, /* [VALUE_H, SHARED_H2] */
|
||||
const float* __restrict__ b_v1, /* [VALUE_H] */
|
||||
const float* __restrict__ w_v2, /* [1, VALUE_H] */
|
||||
const float* __restrict__ b_v2, /* [1] */
|
||||
/* advantage head weights */
|
||||
const float* __restrict__ w_a1, /* [ADV_H, SHARED_H2] */
|
||||
const float* __restrict__ b_a1, /* [ADV_H] */
|
||||
const float* __restrict__ w_a2, /* [NUM_ACTIONS, ADV_H] */
|
||||
const float* __restrict__ b_a2, /* [NUM_ACTIONS] */
|
||||
/* output */
|
||||
int* __restrict__ out_actions, /* [N] greedy action index */
|
||||
int N
|
||||
) {
|
||||
extern __shared__ float shmem[];
|
||||
float* shmem_weights = shmem;
|
||||
float* shmem_bias = shmem + SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM;
|
||||
|
||||
int window_id = blockIdx.x;
|
||||
int lane_id = threadIdx.x;
|
||||
if (window_id >= N) return;
|
||||
|
||||
/* Distributed state loading — each lane loads stride-32 elements */
|
||||
float state_dist[DIST_SIZE(STATE_DIM)];
|
||||
for (int i = 0; i < DIST_SIZE(STATE_DIM); i++) state_dist[i] = 0.0f;
|
||||
int base = window_id * STATE_DIM;
|
||||
for (int i = lane_id; i < STATE_DIM; i += 32)
|
||||
state_dist[i / 32] = states[base + i];
|
||||
|
||||
/* Distributed scratch buffers */
|
||||
float scratch1_dist[DIST_SIZE(SHARED_H1)];
|
||||
float scratch2_dist[DIST_SIZE(SHARED_H2)];
|
||||
float q_values[NUM_ACTIONS];
|
||||
|
||||
/* Forward pass — warp-cooperative dueling Q-network */
|
||||
q_forward_dueling_warp_shmem(
|
||||
state_dist,
|
||||
w_s1, b_s1, w_s2, b_s2,
|
||||
w_v1, b_v1, w_v2, b_v2,
|
||||
w_a1, b_a1, w_a2, b_a2,
|
||||
scratch1_dist, scratch2_dist,
|
||||
q_values,
|
||||
lane_id,
|
||||
shmem_weights, shmem_bias
|
||||
);
|
||||
|
||||
/* Greedy argmax — lane 0 writes result.
|
||||
* All lanes have identical q_values after the warp-cooperative forward pass. */
|
||||
if (lane_id == 0) {
|
||||
int best = 0;
|
||||
float best_q = q_values[0];
|
||||
for (int a = 1; a < NUM_ACTIONS; a++) {
|
||||
if (q_values[a] > best_q) {
|
||||
best_q = q_values[a];
|
||||
best = a;
|
||||
}
|
||||
}
|
||||
out_actions[window_id] = best;
|
||||
}
|
||||
}
|
||||
@@ -29,6 +29,7 @@
|
||||
//! 100+ SMs, leaving no room for env_step overlap). Removing the second stream
|
||||
//! and event sync yields a 15-20% wall-time reduction.
|
||||
|
||||
use std::mem::ManuallyDrop;
|
||||
use std::sync::Arc;
|
||||
use cudarc::driver::{CudaContext, CudaFunction, CudaGraph, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
|
||||
use cudarc::nvrtc::Ptx;
|
||||
@@ -37,7 +38,8 @@ use tracing::{info, warn};
|
||||
|
||||
use ml_core::nvtx::NvtxRange;
|
||||
use crate::MLError;
|
||||
use super::gpu_weights::{DuelingWeightSet, PpoActorWeightSet};
|
||||
use super::batched_forward::CublasForward;
|
||||
use super::gpu_weights::{BranchingWeightSet, DuelingWeightSet, PpoActorWeightSet};
|
||||
|
||||
// ── CUDA Graph wrapper ───────────────────────────────────────────────────────
|
||||
//
|
||||
@@ -69,6 +71,7 @@ static METRICS_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();
|
||||
static GATHER_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();
|
||||
static PPO_FORWARD_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();
|
||||
static SUPERVISED_SIGNAL_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();
|
||||
static BACKTEST_DQN_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();
|
||||
|
||||
fn compile_env_ptx(context: &CudaContext) -> Result<Ptx, String> {
|
||||
let src = include_str!("backtest_env_kernel.cu");
|
||||
@@ -103,6 +106,68 @@ fn compile_supervised_signal_ptx(context: &CudaContext) -> Result<Ptx, String> {
|
||||
crate::cuda_pipeline::compile_ptx_for_device(src, context)
|
||||
}
|
||||
|
||||
/// Compile `compute_expected_q` + `backtest_greedy_argmax` kernels from the
|
||||
/// experience kernels source (which also contains `compute_expected_q`).
|
||||
fn compile_backtest_dqn_ptx(context: &CudaContext) -> Result<Ptx, String> {
|
||||
let defines = "\
|
||||
#define STATE_DIM 48\n\
|
||||
#define MARKET_DIM 42\n\
|
||||
#define PORTFOLIO_DIM 3\n";
|
||||
let kernel_src = include_str!("experience_kernels.cu");
|
||||
let full_source = format!("{defines}{kernel_src}");
|
||||
crate::cuda_pipeline::compile_ptx_for_device(&full_source, context)
|
||||
}
|
||||
|
||||
// ── DQN backtest forward config ───────────────────────────────────────────────
|
||||
|
||||
/// Configuration for cuBLAS-based DQN forward pass in backtest evaluation.
|
||||
///
|
||||
/// Bundles network architecture dimensions, C51 distributional parameters,
|
||||
/// and branching configuration needed by `evaluate_dqn()`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DqnBacktestConfig {
|
||||
/// Shared trunk hidden layer 1 dimension.
|
||||
pub shared_h1: usize,
|
||||
/// Shared trunk hidden layer 2 dimension.
|
||||
pub shared_h2: usize,
|
||||
/// Value head hidden dimension.
|
||||
pub value_h: usize,
|
||||
/// Advantage head hidden dimension.
|
||||
pub adv_h: usize,
|
||||
/// C51 distributional atom count (typically 51).
|
||||
pub num_atoms: usize,
|
||||
/// Exposure branch action count (branch 0, typically 5).
|
||||
pub branch_0_size: usize,
|
||||
/// Order-type branch action count (branch 1, typically 3).
|
||||
pub branch_1_size: usize,
|
||||
/// Urgency branch action count (branch 2, typically 3).
|
||||
pub branch_2_size: usize,
|
||||
/// C51 minimum support value.
|
||||
pub v_min: f32,
|
||||
/// C51 maximum support value.
|
||||
pub v_max: f32,
|
||||
}
|
||||
|
||||
impl DqnBacktestConfig {
|
||||
/// Create from the legacy `network_dims` tuple with standard C51 defaults.
|
||||
///
|
||||
/// Uses branch sizes [5, 3, 3], num_atoms=51, v_min=-25, v_max=25.
|
||||
pub fn from_network_dims(network_dims: (usize, usize, usize, usize)) -> Self {
|
||||
Self {
|
||||
shared_h1: network_dims.0,
|
||||
shared_h2: network_dims.1,
|
||||
value_h: network_dims.2,
|
||||
adv_h: network_dims.3,
|
||||
num_atoms: 51,
|
||||
branch_0_size: 5,
|
||||
branch_1_size: 3,
|
||||
branch_2_size: 3,
|
||||
v_min: -25.0,
|
||||
v_max: 25.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Public types ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Per-window evaluation result returned after a full backtest run.
|
||||
@@ -226,12 +291,40 @@ pub struct GpuBacktestEvaluator {
|
||||
state_dim: usize,
|
||||
config: GpuBacktestConfig,
|
||||
|
||||
/// Forward kernel for pure-CUDA DQN evaluation (compiled lazily via NVRTC).
|
||||
/// `None` until `evaluate_dqn()` is first called with network dims.
|
||||
forward_kernel: Option<CudaFunction>,
|
||||
/// Actions buffer for forward kernel output (reused across steps).
|
||||
forward_actions_buf: CudaSlice<i32>,
|
||||
|
||||
// ── cuBLAS DQN forward pass state (lazy-initialised) ────────────────
|
||||
|
||||
/// cuBLAS forward pass context. Created on first `evaluate_dqn()` call.
|
||||
cublas_forward: Option<CublasForward>,
|
||||
|
||||
/// Flat F32 weight buffer for cuBLAS (20 tensors concatenated).
|
||||
/// Created on first `evaluate_dqn()` call.
|
||||
cublas_params_flat: Option<CudaSlice<f32>>,
|
||||
|
||||
/// cuBLAS activation scratch buffers (created on first call).
|
||||
cublas_h_s1: Option<CudaSlice<f32>>,
|
||||
cublas_h_s2: Option<CudaSlice<f32>>,
|
||||
cublas_h_v: Option<CudaSlice<f32>>,
|
||||
cublas_h_b0: Option<CudaSlice<f32>>,
|
||||
cublas_h_b1: Option<CudaSlice<f32>>,
|
||||
cublas_h_b2: Option<CudaSlice<f32>>,
|
||||
cublas_v_logits: Option<CudaSlice<f32>>,
|
||||
cublas_b_logits: Option<CudaSlice<f32>>,
|
||||
|
||||
/// Q-values buffer for C51 expected-Q computation.
|
||||
cublas_q_values: Option<CudaSlice<f32>>,
|
||||
|
||||
/// `compute_expected_q` kernel (loaded from experience_kernels.cu).
|
||||
expected_q_kernel: Option<CudaFunction>,
|
||||
|
||||
/// `experience_action_select` kernel for greedy argmax.
|
||||
action_select_kernel: Option<CudaFunction>,
|
||||
|
||||
/// RNG states for action selection (greedy mode uses epsilon=0).
|
||||
rng_states: Option<CudaSlice<u32>>,
|
||||
|
||||
/// Cached CUDA graph for `evaluate_dqn_graphed()` step loop.
|
||||
///
|
||||
/// Captured on first call, replayed on subsequent calls with the same weights
|
||||
@@ -446,8 +539,21 @@ impl GpuBacktestEvaluator {
|
||||
portfolio_dim: PORTFOLIO_DIM,
|
||||
state_dim,
|
||||
config,
|
||||
forward_kernel: None,
|
||||
forward_actions_buf,
|
||||
cublas_forward: None,
|
||||
cublas_params_flat: None,
|
||||
cublas_h_s1: None,
|
||||
cublas_h_s2: None,
|
||||
cublas_h_v: None,
|
||||
cublas_h_b0: None,
|
||||
cublas_h_b1: None,
|
||||
cublas_h_b2: None,
|
||||
cublas_v_logits: None,
|
||||
cublas_b_logits: None,
|
||||
cublas_q_values: None,
|
||||
expected_q_kernel: None,
|
||||
action_select_kernel: None,
|
||||
rng_states: None,
|
||||
dqn_graph: None,
|
||||
})
|
||||
}
|
||||
@@ -628,103 +734,41 @@ impl GpuBacktestEvaluator {
|
||||
self.launch_metrics_and_download()
|
||||
}
|
||||
|
||||
/// Pure-GPU DQN evaluation — no candle, no closures.
|
||||
/// Pure-GPU DQN evaluation via cuBLAS SGEMM — no candle, no closures.
|
||||
///
|
||||
/// Uses a dedicated CUDA forward kernel with the same warp-cooperative
|
||||
/// approach as the experience collector. Eliminates candle dispatch overhead
|
||||
/// and enables future CUDA Graph capture.
|
||||
/// Uses cuBLAS matrix multiplications for the Q-network forward pass,
|
||||
/// replacing the old warp-cooperative shared-memory-tiling kernel that
|
||||
/// crashed with `CUDA_ERROR_INVALID_VALUE` for `hidden_dim >= 768`
|
||||
/// (shared memory exceeds 49 KB limit).
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `online_weights` - Dueling Q-network weight buffers (12 CudaSlice pointers)
|
||||
/// * `network_dims` - `(shared_h1, shared_h2, value_h, adv_h)` matching the
|
||||
/// dueling network architecture. Used for NVRTC `#define` injection on first call.
|
||||
/// * `online_weights` - Dueling Q-network weight buffers (12 CudaSlice pointers).
|
||||
/// For branching DQN, `w_a1`/`b_a1`/`w_a2`/`b_a2` map to branch 0 (exposure).
|
||||
/// * `branching_weights` - Optional branch 1 (order) and branch 2 (urgency) weights.
|
||||
/// When `None`, branch 0 weights are replicated for all 3 branches.
|
||||
/// * `dqn_cfg` - Network dimensions and C51 distributional parameters.
|
||||
pub fn evaluate_dqn(
|
||||
&mut self,
|
||||
online_weights: &DuelingWeightSet,
|
||||
network_dims: (usize, usize, usize, usize),
|
||||
branching_weights: Option<&BranchingWeightSet>,
|
||||
dqn_cfg: &DqnBacktestConfig,
|
||||
) -> Result<Vec<WindowMetrics>, MLError> {
|
||||
// Lazy-compile the forward kernel if not yet done
|
||||
if self.forward_kernel.is_none() {
|
||||
let kernel = self.compile_forward_kernel(network_dims)?;
|
||||
self.forward_kernel = Some(kernel);
|
||||
}
|
||||
// Lazy-initialise cuBLAS context + buffers
|
||||
self.ensure_cublas_ready(dqn_cfg)?;
|
||||
|
||||
let state_dim = self.state_dim;
|
||||
let (shared_h1, shared_h2, _, _) = network_dims;
|
||||
let shmem_max_in = state_dim.max(shared_h1).max(shared_h2);
|
||||
let shmem_tile_rows = {
|
||||
let max_tile = 49152 / (4 * (shmem_max_in + 1));
|
||||
let pow2 = (max_tile as u32).next_power_of_two() >> 1;
|
||||
pow2.max(16) as usize
|
||||
};
|
||||
let shmem_bytes = (shmem_tile_rows * shmem_max_in + shmem_tile_rows)
|
||||
* std::mem::size_of::<f32>();
|
||||
// Flatten DuelingWeightSet (+BranchingWeightSet) into the flat params buffer
|
||||
self.flatten_weights_for_cublas(online_weights, branching_weights, dqn_cfg)?;
|
||||
|
||||
for step in 0..self.max_len {
|
||||
// 1. Launch gather_states kernel — writes to self.states_buf
|
||||
self.launch_gather(step)?;
|
||||
|
||||
// 2. Launch pure-CUDA forward kernel — reads states_buf, writes forward_actions_buf
|
||||
let forward_cfg = LaunchConfig {
|
||||
grid_dim: (self.n_windows as u32, 1, 1),
|
||||
block_dim: (32, 1, 1),
|
||||
shared_mem_bytes: shmem_bytes as u32,
|
||||
};
|
||||
let n_i32 = self.n_windows as i32;
|
||||
|
||||
let fwd_kernel = self.forward_kernel.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("forward_kernel unexpectedly None".to_owned())
|
||||
// Compute weight pointers into the flat buffer
|
||||
let param_sizes = compute_backtest_param_sizes(dqn_cfg, self.state_dim);
|
||||
let w_ptrs = {
|
||||
let flat = self.cublas_params_flat.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_params_flat unexpectedly None".to_owned())
|
||||
})?;
|
||||
// Safety: argument order matches backtest_forward_kernel signature exactly.
|
||||
// All CudaSlice lifetimes are valid (owned by self / online_weights).
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(fwd_kernel)
|
||||
.arg(&self.states_buf)
|
||||
.arg(&online_weights.w_s1)
|
||||
.arg(&online_weights.b_s1)
|
||||
.arg(&online_weights.w_s2)
|
||||
.arg(&online_weights.b_s2)
|
||||
.arg(&online_weights.w_v1)
|
||||
.arg(&online_weights.b_v1)
|
||||
.arg(&online_weights.w_v2)
|
||||
.arg(&online_weights.b_v2)
|
||||
.arg(&online_weights.w_a1)
|
||||
.arg(&online_weights.b_a1)
|
||||
.arg(&online_weights.w_a2)
|
||||
.arg(&online_weights.b_a2)
|
||||
.arg(&self.forward_actions_buf)
|
||||
.arg(&n_i32)
|
||||
.launch(forward_cfg)
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!(
|
||||
"backtest_forward_kernel launch step {step}: {e}"
|
||||
))
|
||||
})?;
|
||||
}
|
||||
super::batched_forward::f32_weight_ptrs(flat, ¶m_sizes, &self.stream)
|
||||
};
|
||||
|
||||
// 3. DtoD copy: forward_actions_buf → actions_buf
|
||||
{
|
||||
let num_bytes = self.n_windows * std::mem::size_of::<i32>();
|
||||
let src_view = self.forward_actions_buf.slice(..self.n_windows);
|
||||
let (src_ptr, _src_sync) = src_view.device_ptr(&self.stream);
|
||||
let (dst_ptr, _dst_sync) = self.actions_buf.device_ptr(&self.stream);
|
||||
unsafe {
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr,
|
||||
src_ptr,
|
||||
num_bytes,
|
||||
self.stream.cu_stream(),
|
||||
)
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!("forward→actions DtoD step {step}: {e}"))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Launch env step kernel
|
||||
self.launch_env_step(step)?;
|
||||
}
|
||||
self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)?;
|
||||
|
||||
self.launch_metrics_and_download()
|
||||
}
|
||||
@@ -734,11 +778,10 @@ impl GpuBacktestEvaluator {
|
||||
/// Pure-GPU DQN evaluation with CUDA Graph capture and replay.
|
||||
///
|
||||
/// Identical to [`evaluate_dqn`] but wraps the step loop in a CUDA Graph.
|
||||
/// On the first call, the entire step loop (gather + forward + DtoD +
|
||||
/// env_step for each step 0..max_len) is captured into a graph and then
|
||||
/// launched. On subsequent calls with the same evaluator, the cached graph
|
||||
/// is replayed — eliminating per-step kernel launch overhead (~5-10μs
|
||||
/// per launch × 4 kernels × max_len steps = 20-40ms savings).
|
||||
/// On the first call, the entire step loop (gather + cuBLAS forward +
|
||||
/// expected_q + action_select + env_step for each step 0..max_len) is
|
||||
/// captured into a graph and then launched. On subsequent calls with the
|
||||
/// same evaluator, the cached graph is replayed.
|
||||
///
|
||||
/// The metrics kernel runs OUTSIDE the graph (one-shot, not worth graphing).
|
||||
///
|
||||
@@ -760,15 +803,25 @@ impl GpuBacktestEvaluator {
|
||||
pub fn evaluate_dqn_graphed(
|
||||
&mut self,
|
||||
online_weights: &DuelingWeightSet,
|
||||
network_dims: (usize, usize, usize, usize),
|
||||
branching_weights: Option<&BranchingWeightSet>,
|
||||
dqn_cfg: &DqnBacktestConfig,
|
||||
) -> Result<Vec<WindowMetrics>, MLError> {
|
||||
let _nvtx = NvtxRange::new("backtest_evaluate_dqn_graphed");
|
||||
|
||||
// Lazy-compile the forward kernel if not yet done
|
||||
if self.forward_kernel.is_none() {
|
||||
let kernel = self.compile_forward_kernel(network_dims)?;
|
||||
self.forward_kernel = Some(kernel);
|
||||
}
|
||||
// Lazy-initialise cuBLAS context + buffers
|
||||
self.ensure_cublas_ready(dqn_cfg)?;
|
||||
|
||||
// Flatten weights into flat params buffer
|
||||
self.flatten_weights_for_cublas(online_weights, branching_weights, dqn_cfg)?;
|
||||
|
||||
// Compute weight pointers
|
||||
let param_sizes = compute_backtest_param_sizes(dqn_cfg, self.state_dim);
|
||||
let w_ptrs = {
|
||||
let flat = self.cublas_params_flat.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_params_flat unexpectedly None".to_owned())
|
||||
})?;
|
||||
super::batched_forward::f32_weight_ptrs(flat, ¶m_sizes, &self.stream)
|
||||
};
|
||||
|
||||
// ── Replay cached graph if available ────────────────────────────────
|
||||
if let Some(ref graph) = self.dqn_graph {
|
||||
@@ -807,11 +860,12 @@ impl GpuBacktestEvaluator {
|
||||
warn!(
|
||||
"CUDA graph capture begin failed ({e}), falling back to non-graphed evaluate_dqn"
|
||||
);
|
||||
return self.evaluate_dqn(online_weights, network_dims);
|
||||
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
|
||||
.and_then(|()| self.launch_metrics_and_download());
|
||||
}
|
||||
|
||||
// Submit the step loop to the stream (captured, NOT executed yet).
|
||||
let capture_err = self.submit_dqn_step_loop(online_weights, network_dims);
|
||||
let capture_err = self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg);
|
||||
|
||||
// End capture — get the graph regardless of whether submit succeeded,
|
||||
// because the stream is in capture mode and must be ended.
|
||||
@@ -826,7 +880,8 @@ impl GpuBacktestEvaluator {
|
||||
// If the step loop submission had an error, propagate it
|
||||
if let Err(e) = capture_err {
|
||||
warn!("CUDA graph capture step loop failed ({e}), falling back to non-graphed path");
|
||||
return self.evaluate_dqn(online_weights, network_dims);
|
||||
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
|
||||
.and_then(|()| self.launch_metrics_and_download());
|
||||
}
|
||||
|
||||
// Process graph instantiation result
|
||||
@@ -834,20 +889,23 @@ impl GpuBacktestEvaluator {
|
||||
Ok(Some(g)) => g,
|
||||
Ok(None) => {
|
||||
warn!("CUDA graph capture returned empty graph, falling back to non-graphed path");
|
||||
return self.evaluate_dqn(online_weights, network_dims);
|
||||
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
|
||||
.and_then(|()| self.launch_metrics_and_download());
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"CUDA graph instantiation failed ({e}), falling back to non-graphed path"
|
||||
);
|
||||
return self.evaluate_dqn(online_weights, network_dims);
|
||||
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
|
||||
.and_then(|()| self.launch_metrics_and_download());
|
||||
}
|
||||
};
|
||||
|
||||
// Launch the graph (this actually executes the captured step loop)
|
||||
if let Err(e) = graph.launch() {
|
||||
warn!("CUDA graph launch failed ({e}), falling back to non-graphed path");
|
||||
return self.evaluate_dqn(online_weights, network_dims);
|
||||
return self.submit_dqn_step_loop_cublas(&w_ptrs, dqn_cfg)
|
||||
.and_then(|()| self.launch_metrics_and_download());
|
||||
}
|
||||
|
||||
info!("GpuBacktestEvaluator: CUDA graph captured and launched successfully");
|
||||
@@ -868,74 +926,137 @@ impl GpuBacktestEvaluator {
|
||||
self.dqn_graph = None;
|
||||
}
|
||||
|
||||
/// Submit the DQN step loop to the stream (used during graph capture).
|
||||
/// Submit the cuBLAS-based DQN step loop to the stream.
|
||||
///
|
||||
/// This is the inner loop of `evaluate_dqn` extracted as a helper so that
|
||||
/// `evaluate_dqn_graphed` can call it during stream capture. Returns an
|
||||
/// error if any kernel launch fails.
|
||||
fn submit_dqn_step_loop(
|
||||
/// For each timestep:
|
||||
/// 1. `gather_states` — assemble [N, state_dim] batch
|
||||
/// 2. `CublasForward::forward_online` — Q-network via cuBLAS SGEMM
|
||||
/// 3. `compute_expected_q` — C51 logits to expected Q-values
|
||||
/// 4. `experience_action_select` — greedy argmax (epsilon=0)
|
||||
/// 5. DtoD copy: actions → actions_buf
|
||||
/// 6. `launch_env_step` — portfolio simulation
|
||||
fn submit_dqn_step_loop_cublas(
|
||||
&self,
|
||||
online_weights: &DuelingWeightSet,
|
||||
network_dims: (usize, usize, usize, usize),
|
||||
w_ptrs: &[u64; 20],
|
||||
dqn_cfg: &DqnBacktestConfig,
|
||||
) -> Result<(), MLError> {
|
||||
let state_dim = self.state_dim;
|
||||
let (shared_h1, shared_h2, _, _) = network_dims;
|
||||
let shmem_max_in = state_dim.max(shared_h1).max(shared_h2);
|
||||
let shmem_tile_rows = {
|
||||
let max_tile = 49152 / (4 * (shmem_max_in + 1));
|
||||
let pow2 = (max_tile as u32).next_power_of_two() >> 1;
|
||||
pow2.max(16) as usize
|
||||
let cublas = self.cublas_forward.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_forward unexpectedly None".to_owned())
|
||||
})?;
|
||||
let h_s1 = self.cublas_h_s1.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_h_s1 unexpectedly None".to_owned())
|
||||
})?;
|
||||
let h_s2 = self.cublas_h_s2.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_h_s2 unexpectedly None".to_owned())
|
||||
})?;
|
||||
let h_v = self.cublas_h_v.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_h_v unexpectedly None".to_owned())
|
||||
})?;
|
||||
let h_b0 = self.cublas_h_b0.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_h_b0 unexpectedly None".to_owned())
|
||||
})?;
|
||||
let h_b1 = self.cublas_h_b1.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_h_b1 unexpectedly None".to_owned())
|
||||
})?;
|
||||
let h_b2 = self.cublas_h_b2.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_h_b2 unexpectedly None".to_owned())
|
||||
})?;
|
||||
let v_logits = self.cublas_v_logits.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_v_logits unexpectedly None".to_owned())
|
||||
})?;
|
||||
let b_logits = self.cublas_b_logits.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_b_logits unexpectedly None".to_owned())
|
||||
})?;
|
||||
let q_values = self.cublas_q_values.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_q_values unexpectedly None".to_owned())
|
||||
})?;
|
||||
let eq_kernel = self.expected_q_kernel.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("expected_q_kernel unexpectedly None".to_owned())
|
||||
})?;
|
||||
let as_kernel = self.action_select_kernel.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("action_select_kernel unexpectedly None".to_owned())
|
||||
})?;
|
||||
let rng = self.rng_states.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("rng_states unexpectedly None".to_owned())
|
||||
})?;
|
||||
|
||||
let n = self.n_windows;
|
||||
let n_i32 = n as i32;
|
||||
let na = dqn_cfg.num_atoms as i32;
|
||||
let b0 = dqn_cfg.branch_0_size as i32;
|
||||
let b1 = dqn_cfg.branch_1_size as i32;
|
||||
let b2 = dqn_cfg.branch_2_size as i32;
|
||||
let v_min_f = dqn_cfg.v_min;
|
||||
let v_max_f = dqn_cfg.v_max;
|
||||
let use_branching = if dqn_cfg.branch_1_size > 0 && dqn_cfg.branch_2_size > 0 { 1_i32 } else { 0_i32 };
|
||||
|
||||
let blocks_256 = ((n + 255) / 256) as u32;
|
||||
let launch_cfg = LaunchConfig {
|
||||
grid_dim: (blocks_256.max(1), 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let shmem_bytes = (shmem_tile_rows * shmem_max_in + shmem_tile_rows)
|
||||
* std::mem::size_of::<f32>();
|
||||
|
||||
for step in 0..self.max_len {
|
||||
// 1. Launch gather_states kernel — writes to self.states_buf
|
||||
self.launch_gather(step)?;
|
||||
|
||||
// 2. Launch pure-CUDA forward kernel — reads states_buf, writes forward_actions_buf
|
||||
let forward_cfg = LaunchConfig {
|
||||
grid_dim: (self.n_windows as u32, 1, 1),
|
||||
block_dim: (32, 1, 1),
|
||||
shared_mem_bytes: shmem_bytes as u32,
|
||||
};
|
||||
let n_i32 = self.n_windows as i32;
|
||||
// 2. cuBLAS Q-forward: states_buf → v_logits + b_logits
|
||||
cublas.forward_online(
|
||||
&self.stream,
|
||||
&self.states_buf,
|
||||
w_ptrs,
|
||||
h_s1, h_s2, h_v, h_b0, h_b1, h_b2,
|
||||
v_logits, b_logits,
|
||||
)?;
|
||||
|
||||
let fwd_kernel = self.forward_kernel.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("forward_kernel unexpectedly None".to_owned())
|
||||
})?;
|
||||
// Safety: argument order matches backtest_forward_kernel signature exactly.
|
||||
// All CudaSlice lifetimes are valid (owned by self / online_weights).
|
||||
// 3. C51 distributional → expected Q-values
|
||||
// For num_atoms==1 this kernel is a no-op, but we always run it
|
||||
// to keep the code path uniform (branch sizes still produce correct
|
||||
// argmax from b_logits when num_atoms==1).
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(fwd_kernel)
|
||||
.arg(&self.states_buf)
|
||||
.arg(&online_weights.w_s1)
|
||||
.arg(&online_weights.b_s1)
|
||||
.arg(&online_weights.w_s2)
|
||||
.arg(&online_weights.b_s2)
|
||||
.arg(&online_weights.w_v1)
|
||||
.arg(&online_weights.b_v1)
|
||||
.arg(&online_weights.w_v2)
|
||||
.arg(&online_weights.b_v2)
|
||||
.arg(&online_weights.w_a1)
|
||||
.arg(&online_weights.b_a1)
|
||||
.arg(&online_weights.w_a2)
|
||||
.arg(&online_weights.b_a2)
|
||||
.arg(&self.forward_actions_buf)
|
||||
.launch_builder(eq_kernel)
|
||||
.arg(v_logits)
|
||||
.arg(b_logits)
|
||||
.arg(q_values)
|
||||
.arg(&n_i32)
|
||||
.launch(forward_cfg)
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!(
|
||||
"backtest_forward_kernel launch step {step}: {e}"
|
||||
))
|
||||
})?;
|
||||
.arg(&na)
|
||||
.arg(&b0)
|
||||
.arg(&b1)
|
||||
.arg(&b2)
|
||||
.arg(&v_min_f)
|
||||
.arg(&v_max_f)
|
||||
.launch(launch_cfg)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"compute_expected_q t={step}: {e}"
|
||||
)))?;
|
||||
}
|
||||
|
||||
// 3. DtoD copy: forward_actions_buf → actions_buf
|
||||
// 4. Greedy action selection (epsilon=0)
|
||||
let epsilon = 0.0_f32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(as_kernel)
|
||||
.arg(q_values)
|
||||
.arg(&self.forward_actions_buf)
|
||||
.arg(rng)
|
||||
.arg(&epsilon)
|
||||
.arg(&n_i32)
|
||||
.arg(&b0)
|
||||
.arg(&b1)
|
||||
.arg(&b2)
|
||||
.arg(&use_branching)
|
||||
.launch(launch_cfg)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"experience_action_select t={step}: {e}"
|
||||
)))?;
|
||||
}
|
||||
|
||||
// 5. DtoD copy: forward_actions_buf → actions_buf
|
||||
{
|
||||
let num_bytes = self.n_windows * std::mem::size_of::<i32>();
|
||||
let src_view = self.forward_actions_buf.slice(..self.n_windows);
|
||||
let num_bytes = n * std::mem::size_of::<i32>();
|
||||
let src_view = self.forward_actions_buf.slice(..n);
|
||||
let (src_ptr, _src_sync) = src_view.device_ptr(&self.stream);
|
||||
let (dst_ptr, _dst_sync) = self.actions_buf.device_ptr(&self.stream);
|
||||
unsafe {
|
||||
@@ -946,12 +1067,12 @@ impl GpuBacktestEvaluator {
|
||||
self.stream.cu_stream(),
|
||||
)
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!("forward→actions DtoD step {step}: {e}"))
|
||||
MLError::ModelError(format!("forward->actions DtoD step {step}: {e}"))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Launch env step kernel
|
||||
// 6. Launch env step kernel
|
||||
self.launch_env_step(step)?;
|
||||
}
|
||||
|
||||
@@ -1116,74 +1237,203 @@ impl GpuBacktestEvaluator {
|
||||
self.launch_metrics_and_download()
|
||||
}
|
||||
|
||||
// ── Private extraction helpers ───────────────────────────────────────────
|
||||
// ── cuBLAS DQN forward helpers ──────────────────────────────────────────
|
||||
|
||||
/// Compile the backtest forward kernel via NVRTC with dimension injection.
|
||||
/// Lazily initialise cuBLAS context, activation buffers, and NVRTC kernels.
|
||||
///
|
||||
/// Concatenates `common_device_functions.cuh` + dimension `#define` overrides +
|
||||
/// `backtest_forward_kernel.cu`, following the same pattern as the experience
|
||||
/// collector's NVRTC compilation.
|
||||
fn compile_forward_kernel(
|
||||
&self,
|
||||
network_dims: (usize, usize, usize, usize),
|
||||
) -> Result<CudaFunction, MLError> {
|
||||
let (shared_h1, shared_h2, value_h, adv_h) = network_dims;
|
||||
let state_dim = self.state_dim;
|
||||
let shmem_max_in_dim = state_dim.max(shared_h1).max(shared_h2);
|
||||
let shmem_tile_rows = {
|
||||
let gpu_caps = crate::gpu::capabilities::cached_capabilities();
|
||||
let shmem_kb = crate::gpu::capabilities::max_shared_memory_kb(&gpu_caps.device_name);
|
||||
let shmem_limit_bytes = (shmem_kb * 1024 * 80 / 100).max(49152);
|
||||
let max_tile = shmem_limit_bytes / (4 * (shmem_max_in_dim + 1));
|
||||
let pow2 = (max_tile as u32).next_power_of_two() >> 1;
|
||||
pow2.clamp(16, 256) as usize
|
||||
};
|
||||
/// Called once on first `evaluate_dqn()` / `evaluate_dqn_graphed()`.
|
||||
/// Allocates:
|
||||
/// - `CublasForward` handle + bias kernels
|
||||
/// - Flat F32 weight buffer (20 tensors)
|
||||
/// - Activation scratch buffers for the forward pass
|
||||
/// - Q-values buffer, RNG states
|
||||
/// - `compute_expected_q` and `experience_action_select` kernels
|
||||
fn ensure_cublas_ready(&mut self, dqn_cfg: &DqnBacktestConfig) -> Result<(), MLError> {
|
||||
if self.cublas_forward.is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let common_src = include_str!("common_device_functions.cuh");
|
||||
let kernel_src = include_str!("backtest_forward_kernel.cu");
|
||||
let n = self.n_windows;
|
||||
let na = dqn_cfg.num_atoms;
|
||||
let total_branch_atoms = (dqn_cfg.branch_0_size + dqn_cfg.branch_1_size + dqn_cfg.branch_2_size) * na;
|
||||
let total_actions = dqn_cfg.branch_0_size + dqn_cfg.branch_1_size + dqn_cfg.branch_2_size;
|
||||
|
||||
// Inject actual dimensions — must precede common_src which has
|
||||
// #error guards requiring these defines.
|
||||
let dim_overrides = format!(
|
||||
"#define STATE_DIM {state_dim}\n\
|
||||
#define MARKET_DIM {market_dim}\n\
|
||||
#define PORTFOLIO_DIM 3\n\
|
||||
#define SHARED_H1 {shared_h1}\n\
|
||||
#define SHARED_H2 {shared_h2}\n\
|
||||
#define VALUE_H {value_h}\n\
|
||||
#define ADV_H {adv_h}\n\
|
||||
#define SHMEM_MAX_IN_DIM {shmem_max_in_dim}\n\
|
||||
#define SHMEM_TILE_ROWS {shmem_tile_rows}\n",
|
||||
market_dim = self.feature_dim - self.config.ofi_dim,
|
||||
);
|
||||
// Create cuBLAS forward context
|
||||
let cublas = CublasForward::new(
|
||||
&self.stream,
|
||||
n,
|
||||
self.state_dim,
|
||||
dqn_cfg.shared_h1,
|
||||
dqn_cfg.shared_h2,
|
||||
dqn_cfg.value_h,
|
||||
dqn_cfg.adv_h,
|
||||
na,
|
||||
dqn_cfg.branch_0_size,
|
||||
dqn_cfg.branch_1_size,
|
||||
dqn_cfg.branch_2_size,
|
||||
)?;
|
||||
|
||||
// dim_overrides BEFORE common_src so #error guards in .cuh see the defines
|
||||
let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}");
|
||||
info!(
|
||||
"GpuBacktestEvaluator: compiling forward kernel with dims state={state_dim} \
|
||||
shared=[{shared_h1},{shared_h2}] value={value_h} adv={adv_h} \
|
||||
shmem_max_in={shmem_max_in_dim} tile_rows={shmem_tile_rows}"
|
||||
);
|
||||
// Allocate flat params buffer
|
||||
let param_sizes = compute_backtest_param_sizes(dqn_cfg, self.state_dim);
|
||||
let total_params: usize = param_sizes.iter().sum();
|
||||
let params_flat = self.stream
|
||||
.alloc_zeros::<f32>(total_params)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc cublas_params_flat: {e}")))?;
|
||||
|
||||
// Allocate activation scratch buffers
|
||||
let h_s1 = self.stream.alloc_zeros::<f32>(n * dqn_cfg.shared_h1)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc h_s1: {e}")))?;
|
||||
let h_s2 = self.stream.alloc_zeros::<f32>(n * dqn_cfg.shared_h2)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc h_s2: {e}")))?;
|
||||
let h_v = self.stream.alloc_zeros::<f32>(n * dqn_cfg.value_h)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc h_v: {e}")))?;
|
||||
let h_b0 = self.stream.alloc_zeros::<f32>(n * dqn_cfg.adv_h)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc h_b0: {e}")))?;
|
||||
let h_b1 = self.stream.alloc_zeros::<f32>(n * dqn_cfg.adv_h)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc h_b1: {e}")))?;
|
||||
let h_b2 = self.stream.alloc_zeros::<f32>(n * dqn_cfg.adv_h)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc h_b2: {e}")))?;
|
||||
let v_logits = self.stream.alloc_zeros::<f32>(n * na)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc v_logits: {e}")))?;
|
||||
let b_logits = self.stream.alloc_zeros::<f32>(n * total_branch_atoms)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc b_logits: {e}")))?;
|
||||
let q_values = self.stream.alloc_zeros::<f32>(n * total_actions)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc q_values: {e}")))?;
|
||||
|
||||
// Compile experience kernels for compute_expected_q + action_select
|
||||
let context = self.stream.context();
|
||||
// Use arch-aware compilation so __CUDA_ARCH__ is defined correctly.
|
||||
// On Hopper (sm_90+) this enables TMA bulk async tile loads.
|
||||
let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context)
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!(
|
||||
"backtest_forward_kernel CUDA compilation failed: {e}"
|
||||
))
|
||||
})?;
|
||||
let module = context.load_module(ptx).map_err(|e| {
|
||||
MLError::ModelError(format!("backtest_forward module load: {e}"))
|
||||
let dqn_ptx = BACKTEST_DQN_PTX
|
||||
.get_or_init(|| compile_backtest_dqn_ptx(&context))
|
||||
.as_ref()
|
||||
.map_err(|e| MLError::ModelError(format!("backtest DQN PTX: {e}")))?;
|
||||
let dqn_module = context
|
||||
.load_module(dqn_ptx.clone())
|
||||
.map_err(|e| MLError::ModelError(format!("backtest DQN module load: {e}")))?;
|
||||
let eq_kernel = dqn_module
|
||||
.load_function("compute_expected_q")
|
||||
.map_err(|e| MLError::ModelError(format!("compute_expected_q load: {e}")))?;
|
||||
let as_kernel = dqn_module
|
||||
.load_function("experience_action_select")
|
||||
.map_err(|e| MLError::ModelError(format!("experience_action_select load: {e}")))?;
|
||||
|
||||
// RNG states (initialised with random seeds for greedy mode — epsilon=0 means
|
||||
// the RNG values are never actually used, but the kernel reads them).
|
||||
let rng_seeds: Vec<u32> = (0..n).map(|_| fastrand::u32(..)).collect();
|
||||
let rng_states = self.stream.clone_htod(&rng_seeds)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?;
|
||||
|
||||
info!(
|
||||
"GpuBacktestEvaluator: cuBLAS DQN forward initialised — {} windows, \
|
||||
shared=[{},{}], value={}, adv={}, atoms={}, branches=[{},{},{}]",
|
||||
n, dqn_cfg.shared_h1, dqn_cfg.shared_h2, dqn_cfg.value_h, dqn_cfg.adv_h,
|
||||
na, dqn_cfg.branch_0_size, dqn_cfg.branch_1_size, dqn_cfg.branch_2_size,
|
||||
);
|
||||
|
||||
self.cublas_forward = Some(cublas);
|
||||
self.cublas_params_flat = Some(params_flat);
|
||||
self.cublas_h_s1 = Some(h_s1);
|
||||
self.cublas_h_s2 = Some(h_s2);
|
||||
self.cublas_h_v = Some(h_v);
|
||||
self.cublas_h_b0 = Some(h_b0);
|
||||
self.cublas_h_b1 = Some(h_b1);
|
||||
self.cublas_h_b2 = Some(h_b2);
|
||||
self.cublas_v_logits = Some(v_logits);
|
||||
self.cublas_b_logits = Some(b_logits);
|
||||
self.cublas_q_values = Some(q_values);
|
||||
self.expected_q_kernel = Some(eq_kernel);
|
||||
self.action_select_kernel = Some(as_kernel);
|
||||
self.rng_states = Some(rng_states);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Flatten DuelingWeightSet (+BranchingWeightSet) into the contiguous
|
||||
/// `cublas_params_flat` buffer for cuBLAS weight pointer computation.
|
||||
///
|
||||
/// DtoD-copies each weight tensor into its position in the flat buffer.
|
||||
/// For non-branching mode (branching_weights=None), branch 0 weights
|
||||
/// are replicated for branches 1 and 2.
|
||||
#[allow(clippy::too_many_arguments, unused_assignments)]
|
||||
fn flatten_weights_for_cublas(
|
||||
&mut self,
|
||||
online: &DuelingWeightSet,
|
||||
branching: Option<&BranchingWeightSet>,
|
||||
dqn_cfg: &DqnBacktestConfig,
|
||||
) -> Result<(), MLError> {
|
||||
let param_sizes = compute_backtest_param_sizes(dqn_cfg, self.state_dim);
|
||||
let flat = self.cublas_params_flat.as_ref().ok_or_else(|| {
|
||||
MLError::ModelError("cublas_params_flat not initialised".to_owned())
|
||||
})?;
|
||||
let func = module
|
||||
.load_function("backtest_forward_kernel")
|
||||
.map_err(|e| {
|
||||
MLError::ModelError(format!("backtest_forward_kernel function load: {e}"))
|
||||
})?;
|
||||
info!("GpuBacktestEvaluator: forward kernel compiled and loaded");
|
||||
Ok(func)
|
||||
let dst_base = raw_f32_ptr(flat, &self.stream);
|
||||
|
||||
let mut byte_offset: u64 = 0;
|
||||
|
||||
// Helper: copy a single weight tensor into the flat buffer at current offset.
|
||||
macro_rules! copy_weight {
|
||||
($slice:expr, $idx:expr) => {{
|
||||
let n_elems = param_sizes[$idx];
|
||||
let n_bytes = n_elems * std::mem::size_of::<f32>();
|
||||
let src = raw_f32_ptr($slice, &self.stream);
|
||||
unsafe {
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_base + byte_offset,
|
||||
src,
|
||||
n_bytes,
|
||||
self.stream.cu_stream(),
|
||||
)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"flatten weight {}: {e}", $idx
|
||||
)))?;
|
||||
}
|
||||
byte_offset += n_bytes as u64;
|
||||
}};
|
||||
}
|
||||
|
||||
// Shared trunk: w_s1, b_s1, w_s2, b_s2
|
||||
copy_weight!(&online.w_s1, 0);
|
||||
copy_weight!(&online.b_s1, 1);
|
||||
copy_weight!(&online.w_s2, 2);
|
||||
copy_weight!(&online.b_s2, 3);
|
||||
|
||||
// Value head: w_v1, b_v1, w_v2, b_v2
|
||||
copy_weight!(&online.w_v1, 4);
|
||||
copy_weight!(&online.b_v1, 5);
|
||||
copy_weight!(&online.w_v2, 6);
|
||||
copy_weight!(&online.b_v2, 7);
|
||||
|
||||
// Branch 0 (exposure — from DuelingWeightSet): w_b0fc, b_b0fc, w_b0out, b_b0out
|
||||
copy_weight!(&online.w_a1, 8);
|
||||
copy_weight!(&online.b_a1, 9);
|
||||
copy_weight!(&online.w_a2, 10);
|
||||
copy_weight!(&online.b_a2, 11);
|
||||
|
||||
// Branch 1 (order) and Branch 2 (urgency)
|
||||
if let Some(bw) = branching {
|
||||
// Branching mode: separate weights for each branch
|
||||
copy_weight!(&bw.w_bo1, 12);
|
||||
copy_weight!(&bw.b_bo1, 13);
|
||||
copy_weight!(&bw.w_bo2, 14);
|
||||
copy_weight!(&bw.b_bo2, 15);
|
||||
copy_weight!(&bw.w_bu1, 16);
|
||||
copy_weight!(&bw.b_bu1, 17);
|
||||
copy_weight!(&bw.w_bu2, 18);
|
||||
copy_weight!(&bw.b_bu2, 19);
|
||||
} else {
|
||||
// Non-branching: replicate branch 0 weights for branches 1 and 2.
|
||||
// The FC layers (w_b1fc/b_b1fc) have the same shape as branch 0 FC,
|
||||
// and the out layers differ only in action count (but for non-branching
|
||||
// DQN all branch sizes are the same = branch_0_size).
|
||||
copy_weight!(&online.w_a1, 12);
|
||||
copy_weight!(&online.b_a1, 13);
|
||||
copy_weight!(&online.w_a2, 14);
|
||||
copy_weight!(&online.b_a2, 15);
|
||||
copy_weight!(&online.w_a1, 16);
|
||||
copy_weight!(&online.b_a1, 17);
|
||||
copy_weight!(&online.w_a2, 18);
|
||||
copy_weight!(&online.b_a2, 19);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch the `gather_states` GPU kernel for a given step.
|
||||
@@ -1339,6 +1589,46 @@ impl GpuBacktestEvaluator {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Free helpers ──────────────────────────────────────────────────────────────
|
||||
|
||||
/// Compute the 20 weight tensor sizes (element counts) for cuBLAS forward.
|
||||
///
|
||||
/// Layout: `[w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2,
|
||||
/// w_b0fc, b_b0fc, w_b0out, b_b0out,
|
||||
/// w_b1fc, b_b1fc, w_b1out, b_b1out,
|
||||
/// w_b2fc, b_b2fc, w_b2out, b_b2out]`
|
||||
fn compute_backtest_param_sizes(cfg: &DqnBacktestConfig, state_dim: usize) -> [usize; 20] {
|
||||
[
|
||||
cfg.shared_h1 * state_dim, // w_s1
|
||||
cfg.shared_h1, // b_s1
|
||||
cfg.shared_h2 * cfg.shared_h1, // w_s2
|
||||
cfg.shared_h2, // b_s2
|
||||
cfg.value_h * cfg.shared_h2, // w_v1
|
||||
cfg.value_h, // b_v1
|
||||
cfg.num_atoms * cfg.value_h, // w_v2
|
||||
cfg.num_atoms, // b_v2
|
||||
cfg.adv_h * cfg.shared_h2, // w_b0fc
|
||||
cfg.adv_h, // b_b0fc
|
||||
cfg.branch_0_size * cfg.num_atoms * cfg.adv_h, // w_b0out
|
||||
cfg.branch_0_size * cfg.num_atoms, // b_b0out
|
||||
cfg.adv_h * cfg.shared_h2, // w_b1fc
|
||||
cfg.adv_h, // b_b1fc
|
||||
cfg.branch_1_size * cfg.num_atoms * cfg.adv_h, // w_b1out
|
||||
cfg.branch_1_size * cfg.num_atoms, // b_b1out
|
||||
cfg.adv_h * cfg.shared_h2, // w_b2fc
|
||||
cfg.adv_h, // b_b2fc
|
||||
cfg.branch_2_size * cfg.num_atoms * cfg.adv_h, // w_b2out
|
||||
cfg.branch_2_size * cfg.num_atoms, // b_b2out
|
||||
]
|
||||
}
|
||||
|
||||
/// Extract raw F32 device pointer from a CudaSlice.
|
||||
fn raw_f32_ptr(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> u64 {
|
||||
let (ptr, guard) = slice.device_ptr(stream);
|
||||
let _no_drop = ManuallyDrop::new(guard);
|
||||
ptr
|
||||
}
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -1461,29 +1751,17 @@ mod tests {
|
||||
assert_eq!(3_usize, 3_usize, "PORTFOLIO_DIM constant is 3");
|
||||
}
|
||||
|
||||
/// Verify that the forward kernel CUDA source compiles (nvcc path).
|
||||
/// Verify that the cuBLAS DQN helper kernels (compute_expected_q,
|
||||
/// experience_action_select) compile from experience_kernels.cu.
|
||||
#[test]
|
||||
|
||||
fn test_forward_kernel_ptx_compilation() {
|
||||
fn test_backtest_dqn_ptx_compilation() {
|
||||
let context = cudarc::driver::CudaContext::new(0).expect("CUDA context required");
|
||||
let stream = context.default_stream();
|
||||
let context = stream.context();
|
||||
let common_src = include_str!("common_device_functions.cuh");
|
||||
let kernel_src = include_str!("backtest_forward_kernel.cu");
|
||||
let dim_overrides = "\
|
||||
#define STATE_DIM 48\n\
|
||||
#define MARKET_DIM 42\n\
|
||||
#define PORTFOLIO_DIM 3\n\
|
||||
#define SHARED_H1 256\n\
|
||||
#define SHARED_H2 256\n\
|
||||
#define VALUE_H 128\n\
|
||||
#define ADV_H 128\n\
|
||||
#define SHMEM_MAX_IN_DIM 256\n\
|
||||
#define SHMEM_TILE_ROWS 32\n";
|
||||
let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}");
|
||||
let result = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context);
|
||||
let result = compile_backtest_dqn_ptx(&context);
|
||||
if let Err(ref e) = result {
|
||||
panic!("backtest_forward_kernel PTX compilation failed: {e}");
|
||||
panic!("backtest DQN PTX compilation failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1553,4 +1831,43 @@ mod tests {
|
||||
// The evaluate_dqn_graphed and invalidate_dqn_graph methods are tested
|
||||
// at runtime on GPU hardware via integration tests.
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dqn_backtest_config_from_network_dims() {
|
||||
let cfg = DqnBacktestConfig::from_network_dims((256, 256, 128, 128));
|
||||
assert_eq!(cfg.shared_h1, 256);
|
||||
assert_eq!(cfg.shared_h2, 256);
|
||||
assert_eq!(cfg.value_h, 128);
|
||||
assert_eq!(cfg.adv_h, 128);
|
||||
assert_eq!(cfg.num_atoms, 51);
|
||||
assert_eq!(cfg.branch_0_size, 5);
|
||||
assert_eq!(cfg.branch_1_size, 3);
|
||||
assert_eq!(cfg.branch_2_size, 3);
|
||||
assert!((cfg.v_min - (-25.0)).abs() < f32::EPSILON);
|
||||
assert!((cfg.v_max - 25.0).abs() < f32::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compute_backtest_param_sizes() {
|
||||
let cfg = DqnBacktestConfig::from_network_dims((256, 256, 128, 128));
|
||||
let state_dim = 48_usize; // (42 + 3 + 3) 8-aligned
|
||||
let sizes = compute_backtest_param_sizes(&cfg, state_dim);
|
||||
|
||||
// w_s1: 256 * 48 = 12288
|
||||
assert_eq!(sizes[0], 256 * 48);
|
||||
// b_s1: 256
|
||||
assert_eq!(sizes[1], 256);
|
||||
// w_s2: 256 * 256 = 65536
|
||||
assert_eq!(sizes[2], 256 * 256);
|
||||
// w_v1: 128 * 256 = 32768
|
||||
assert_eq!(sizes[4], 128 * 256);
|
||||
// w_v2: 51 * 128 = 6528
|
||||
assert_eq!(sizes[6], 51 * 128);
|
||||
// w_b0out: 5 * 51 * 128 = 32640
|
||||
assert_eq!(sizes[10], 5 * 51 * 128);
|
||||
// w_b1out: 3 * 51 * 128 = 19584
|
||||
assert_eq!(sizes[14], 3 * 51 * 128);
|
||||
// w_b2out: 3 * 51 * 128 = 19584
|
||||
assert_eq!(sizes[18], 3 * 51 * 128);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1736,7 +1736,7 @@ impl DQNTrainer {
|
||||
max_position_absolute: f64,
|
||||
) -> Result<Option<BacktestMetrics>, MLError> {
|
||||
use crate::cuda_pipeline::gpu_backtest_evaluator::{
|
||||
GpuBacktestConfig, GpuBacktestEvaluator,
|
||||
DqnBacktestConfig, GpuBacktestConfig, GpuBacktestEvaluator,
|
||||
};
|
||||
|
||||
let total_bars = val_close_prices.len();
|
||||
@@ -1868,7 +1868,7 @@ impl DQNTrainer {
|
||||
})?;
|
||||
let agent_guard = bt_handle.block_on(agent_arc.read());
|
||||
|
||||
// Pure-CUDA forward kernel for both OFI and non-OFI paths.
|
||||
// cuBLAS SGEMM forward for both OFI and non-OFI paths.
|
||||
// The gather kernel handles OFI reordering directly (ofi_dim > 0 produces
|
||||
// [market, portfolio, OFI, pad] layout), so no Candle dispatch needed.
|
||||
let is_branching = agent_guard.is_using_branching();
|
||||
@@ -1886,6 +1886,29 @@ impl DQNTrainer {
|
||||
)?
|
||||
};
|
||||
|
||||
let branching_weights = if is_branching {
|
||||
Some(crate::cuda_pipeline::gpu_weights::extract_branching_weights(
|
||||
&vars, stream,
|
||||
)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let hp = internal_trainer.hyperparams();
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
let dqn_cfg = DqnBacktestConfig {
|
||||
shared_h1: network_dims.0,
|
||||
shared_h2: network_dims.1,
|
||||
value_h: network_dims.2,
|
||||
adv_h: network_dims.3,
|
||||
num_atoms: hp.num_atoms,
|
||||
branch_0_size: 5,
|
||||
branch_1_size: if is_branching { 3 } else { 5 },
|
||||
branch_2_size: if is_branching { 3 } else { 5 },
|
||||
v_min: hp.v_min as f32,
|
||||
v_max: hp.v_max as f32,
|
||||
};
|
||||
|
||||
info!(
|
||||
is_branching,
|
||||
ofi_enabled,
|
||||
@@ -1893,12 +1916,17 @@ impl DQNTrainer {
|
||||
shared_h2 = network_dims.1,
|
||||
value_h = network_dims.2,
|
||||
adv_h = network_dims.3,
|
||||
"Pure-CUDA backtest forward: DuelingWeightSet extracted"
|
||||
num_atoms = hp.num_atoms,
|
||||
"cuBLAS SGEMM backtest forward: DuelingWeightSet extracted"
|
||||
);
|
||||
|
||||
// Use CUDA Graph–accelerated path for maximum throughput.
|
||||
// Use CUDA Graph-accelerated path for maximum throughput.
|
||||
// Falls back to non-graphed evaluate_dqn() on unsupported drivers.
|
||||
let metrics = evaluator.evaluate_dqn_graphed(&online_weights, network_dims)?;
|
||||
let metrics = evaluator.evaluate_dqn_graphed(
|
||||
&online_weights,
|
||||
branching_weights.as_ref(),
|
||||
&dqn_cfg,
|
||||
)?;
|
||||
|
||||
drop(agent_guard);
|
||||
|
||||
@@ -3684,10 +3712,10 @@ mod tests {
|
||||
assert_eq!(bounds[8], (0.4, 0.8)); // per_alpha
|
||||
assert_eq!(bounds[9], (0.2, 0.6)); // per_beta_start
|
||||
assert_eq!(bounds[10], (10.0, 50.0)); // v_range (symmetric support, covers gamma 0.88-0.99)
|
||||
assert_eq!(bounds[11], (0.0, 1.0)); // use_branching (C6: repurposed from v_range_mirror)
|
||||
assert_eq!(bounds[11], (1.0, 1.0)); // use_branching (FIXED: always true — GPU pipeline requires branching)
|
||||
assert_eq!(bounds[13], (128.0, 512.0)); // dueling_hidden_dim
|
||||
assert_eq!(bounds[14], (3.0, 5.0)); // n_steps (raised min: Rainbow standard)
|
||||
assert_eq!(bounds[15], (51.0, 201.0)); // num_atoms
|
||||
assert_eq!(bounds[15], (11.0, 101.0)); // num_atoms (11=RTX3050, 101=max)
|
||||
|
||||
// Kelly risk parameters
|
||||
assert_eq!(bounds[17], (0.25, 0.75)); // kelly_fractional
|
||||
@@ -3695,7 +3723,7 @@ mod tests {
|
||||
assert_eq!(bounds[19], (10.0, 30.0)); // volatility_window
|
||||
|
||||
// C2: curiosity_weight fixed to 0.0, noisy_epsilon_floor fixed to 0.10, not in search space
|
||||
assert_eq!(bounds[21], (256.0, 1024.0)); // hidden_dim_base (narrowed for 2-layer net)
|
||||
assert_eq!(bounds[21], (128.0, 1024.0)); // hidden_dim_base (128-1024, H100 has headroom)
|
||||
assert_eq!(bounds[22], (0.0, 1.0)); // cql_alpha (full offline-RL range)
|
||||
|
||||
// C5: branch_hidden_dim
|
||||
|
||||
Reference in New Issue
Block a user