diff --git a/crates/ml/src/cuda_pipeline/backtest_forward_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_forward_kernel.cu deleted file mode 100644 index e25332b12..000000000 --- a/crates/ml/src/cuda_pipeline/backtest_forward_kernel.cu +++ /dev/null @@ -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; - } -} diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index 0160bdd71..c8ea356c0 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -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> = OnceLock::new(); static GATHER_PTX: OnceLock> = OnceLock::new(); static PPO_FORWARD_PTX: OnceLock> = OnceLock::new(); static SUPERVISED_SIGNAL_PTX: OnceLock> = OnceLock::new(); +static BACKTEST_DQN_PTX: OnceLock> = OnceLock::new(); fn compile_env_ptx(context: &CudaContext) -> Result { let src = include_str!("backtest_env_kernel.cu"); @@ -103,6 +106,68 @@ fn compile_supervised_signal_ptx(context: &CudaContext) -> Result { 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 { + 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, /// Actions buffer for forward kernel output (reused across steps). forward_actions_buf: CudaSlice, + // ── cuBLAS DQN forward pass state (lazy-initialised) ──────────────── + + /// cuBLAS forward pass context. Created on first `evaluate_dqn()` call. + cublas_forward: Option, + + /// Flat F32 weight buffer for cuBLAS (20 tensors concatenated). + /// Created on first `evaluate_dqn()` call. + cublas_params_flat: Option>, + + /// cuBLAS activation scratch buffers (created on first call). + cublas_h_s1: Option>, + cublas_h_s2: Option>, + cublas_h_v: Option>, + cublas_h_b0: Option>, + cublas_h_b1: Option>, + cublas_h_b2: Option>, + cublas_v_logits: Option>, + cublas_b_logits: Option>, + + /// Q-values buffer for C51 expected-Q computation. + cublas_q_values: Option>, + + /// `compute_expected_q` kernel (loaded from experience_kernels.cu). + expected_q_kernel: Option, + + /// `experience_action_select` kernel for greedy argmax. + action_select_kernel: Option, + + /// RNG states for action selection (greedy mode uses epsilon=0). + rng_states: Option>, + /// 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, 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::(); + // 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::(); - 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, 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::(); 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::(); - let src_view = self.forward_actions_buf.slice(..self.n_windows); + let num_bytes = n * std::mem::size_of::(); + 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 { - 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::(total_params) + .map_err(|e| MLError::ModelError(format!("alloc cublas_params_flat: {e}")))?; + // Allocate activation scratch buffers + let h_s1 = self.stream.alloc_zeros::(n * dqn_cfg.shared_h1) + .map_err(|e| MLError::ModelError(format!("alloc h_s1: {e}")))?; + let h_s2 = self.stream.alloc_zeros::(n * dqn_cfg.shared_h2) + .map_err(|e| MLError::ModelError(format!("alloc h_s2: {e}")))?; + let h_v = self.stream.alloc_zeros::(n * dqn_cfg.value_h) + .map_err(|e| MLError::ModelError(format!("alloc h_v: {e}")))?; + let h_b0 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) + .map_err(|e| MLError::ModelError(format!("alloc h_b0: {e}")))?; + let h_b1 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) + .map_err(|e| MLError::ModelError(format!("alloc h_b1: {e}")))?; + let h_b2 = self.stream.alloc_zeros::(n * dqn_cfg.adv_h) + .map_err(|e| MLError::ModelError(format!("alloc h_b2: {e}")))?; + let v_logits = self.stream.alloc_zeros::(n * na) + .map_err(|e| MLError::ModelError(format!("alloc v_logits: {e}")))?; + let b_logits = self.stream.alloc_zeros::(n * total_branch_atoms) + .map_err(|e| MLError::ModelError(format!("alloc b_logits: {e}")))?; + let q_values = self.stream.alloc_zeros::(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 = (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::(); + 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, stream: &Arc) -> 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); + } } diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index 364c99633..e435b85b0 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -1736,7 +1736,7 @@ impl DQNTrainer { max_position_absolute: f64, ) -> Result, 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