diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index 0ba032fcd..9f8edc19f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -470,6 +470,15 @@ pub struct GpuBacktestEvaluator { /// Q-values buffer for chunked C51 expected-Q (n_windows * CHUNK_SIZE rows). chunked_q_values: Option>, + /// Branch-major raw C51 logits buffer for the chunked forward pass: + /// `[n_windows * CHUNK_SIZE, total_actions * num_atoms]`. Populated each + /// chunk by `QValueProvider::compute_q_and_b_logits_to`, which DtoD-copies + /// the trainer's `on_b_logits_buf`. The direction-branch slice + /// (`[N, B0, num_atoms]` at the start of each row) is consumed by the + /// Thompson direction sampler in `experience_action_select` + /// (Plan C T2 amendment, T4 evaluator wire-up 2026-04-29). + chunked_b_logits_buf: Option>, + /// Batched states buffer for chunked gather: [n_windows * CHUNK_SIZE, state_dim]. chunked_states_buf: Option>, @@ -801,6 +810,7 @@ impl GpuBacktestEvaluator { chunked_magnitude_conviction_buf: None, dqn_graph: None, chunked_q_values: None, + chunked_b_logits_buf: None, chunked_states_buf: None, chunked_actions_buf: None, chunked_intent_mag_buf: None, @@ -1572,6 +1582,9 @@ impl GpuBacktestEvaluator { let ch_magnitude_conviction = self.chunked_magnitude_conviction_buf.as_ref().ok_or_else(|| { MLError::ModelError("chunked_magnitude_conviction_buf unexpectedly None".to_owned()) })?; + let ch_b_logits = self.chunked_b_logits_buf.as_ref().ok_or_else(|| { + MLError::ModelError("chunked_b_logits_buf unexpectedly None".to_owned()) + })?; let n = self.n_windows; let state_dim = self.state_dim; @@ -1647,15 +1660,25 @@ impl GpuBacktestEvaluator { )))?; } - // ── Phase 2+3: Q-values via QValueProvider or internal cuBLAS ── + // ── Phase 2+3: Q-values + raw branch C51 logits via QValueProvider ── // ch_states_base is the device pointer to chunked_states_buf base — // computed in Phase 1 and stable across the chunk (the Option does // not reallocate). Reuse it instead of taking another borrow. + // + // Plan C T4 (2026-04-29): the val pipeline now also requires the + // trainer's `on_b_logits_buf` (branch-major raw C51 logits) for + // Thompson direction sampling in Phase 4 below. The + // compute_q_and_b_logits_to method DtoD-copies q_out_buf AND + // on_b_logits_buf per sub-batch iteration, keeping the caller's + // chunked_b_logits_buf consistent with chunked_q_values across + // the entire chunk batch (handles the case where batch > + // trainer.batch_size and the inner loop runs multiple sub-iters). let states_ptr = ch_states_base; let q_out_ptr = raw_f32_ptr(ch_q_values, &self.stream); + let b_logits_out_ptr = raw_f32_ptr(ch_b_logits, &self.stream); match q_provider { Some(ref mut qp) => { - qp.compute_q_values_to(states_ptr, batch, q_out_ptr) + qp.compute_q_and_b_logits_to(states_ptr, batch, q_out_ptr, b_logits_out_ptr) .map_err(|e| MLError::ModelError(format!( "QValueProvider chunk_start={chunk_start}: {e}" )))?; @@ -1679,6 +1702,23 @@ impl GpuBacktestEvaluator { let null_portfolio: u64 = 0; let eval_min_hold: i32 = 0; // hold enforcement removed — cost-driven let eval_max_pos: f32 = 0.0; + // Plan C T4 (2026-04-29): direction-branch Thompson buffers. + // Mirror of the gpu_experience_collector wire-up (T3, 380d0ac0d). + // Eval mode (eps_start==eps_end==0) takes the argmax-E[Q] branch + // inside the kernel (no Philox draws), but the buffers must still + // be supplied with real production data — `feedback_no_stubs`. + // ch_b_logits is BRANCH-MAJOR `[batch, total_actions * num_atoms]`; + // the kernel reads only the first `batch * b0 * num_atoms` region + // (Branch 0 = direction). per_sample_support is the trainer's + // mapped-pinned `[max_batch, 4, 3]` stride-12 buffer (sized by + // trainer.batch_size ≥ chunked batch, safe to share). Atom + // positions live on the trainer (model-global `[B0, num_atoms]`). + let qp_for_buffers = q_provider.as_ref().ok_or_else(|| MLError::ModelError( + "QValueProvider required for action_select Thompson buffers".to_owned() + ))?; + let per_sample_support_ptr = qp_for_buffers.per_sample_support_ptr(); + let atom_positions_ptr = qp_for_buffers.atom_positions_ptr(); + let n_atoms_i32 = qp_for_buffers.num_atoms() as i32; unsafe { self.stream .launch_builder(as_kernel) @@ -1708,6 +1748,11 @@ impl GpuBacktestEvaluator { .arg(ch_intent) // out_intent_mag: pure-policy magnitude intent diagnostic .arg(ch_conviction) // out_conviction: direction-branch Kelly warmup signal .arg(ch_magnitude_conviction) // out_magnitude_conviction: magnitude-branch Kelly warmup signal + // Plan C T2 amendment buffers (direction-branch Thompson): + .arg(ch_b_logits) // b_logits_dir: [batch, b0*n_atoms] used (rest is mag/ord/urg, untouched) + .arg(&per_sample_support_ptr) // per_sample_support: [N, 4, 3] stride-12 mapped pinned (trainer-owned) + .arg(&atom_positions_ptr) // atom_positions: [b0, n_atoms] adaptive C51 positions (trainer-owned) + .arg(&n_atoms_i32) // n_atoms: C51 atom count (i32 to match `int` ABI) .launch(launch_cfg) .map_err(|e| MLError::ModelError(format!( "chunked action_select chunk_start={chunk_start}: {e}" @@ -2128,6 +2173,16 @@ impl GpuBacktestEvaluator { let ch_q_values = self.stream.alloc_zeros::(cn * total_actions) .map_err(|e| MLError::ModelError(format!("alloc chunked q_values: {e}")))?; + // Plan C T4 (2026-04-29): chunked branch-major raw C51 logits buffer. + // Sized [cn, total_actions * num_atoms] — DtoD-populated each chunk + // by `compute_q_and_b_logits_to`. Direction-branch slice consumed by + // `experience_action_select` Thompson direction sampling. Padded by + // 32*3 floats to mirror the trainer's allocation pattern (cuBLAS + // tail safety; see gpu_dqn_trainer.rs:9759). + let total_branch_atoms_eval = total_actions * dqn_cfg.num_atoms; + let ch_b_logits = self.stream.alloc_zeros::(cn * total_branch_atoms_eval + 32 * 3) + .map_err(|e| MLError::ModelError(format!("alloc chunked b_logits: {e}")))?; + let state_dim_padded = (self.state_dim + 127) & !127; let ch_states_buf = self.stream.alloc_zeros::(cn * state_dim_padded) .map_err(|e| MLError::ModelError(format!("alloc chunked states_buf: {e}")))?; @@ -2193,6 +2248,7 @@ impl GpuBacktestEvaluator { self.action_select_kernel = Some(as_kernel); self.chunked_q_values = Some(ch_q_values); + self.chunked_b_logits_buf = Some(ch_b_logits); self.chunked_states_buf = Some(ch_states_buf); self.chunked_actions_buf = Some(ch_actions_buf); self.chunked_intent_mag_buf = Some(ch_intent_mag_buf); diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 5670ed004..d63752d6c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -19333,6 +19333,28 @@ impl GpuDqnTrainer { /// Raw pointer to plan_params_buf [B, 6] for trade plan integration. pub fn plan_params_buf_ptr(&self) -> u64 { self.plan_params_buf.raw_ptr() } + /// Raw pointer to on_b_logits_buf + /// [batch_size, (B0+B1+B2+B3) * num_atoms] f32, branch-major. + /// Populated by the most recent `replay_forward_for_q_values()` / + /// graphed forward; consumed by Thompson direction sampling in + /// `experience_action_select` (Plan C T2 amendment, 2026-04-29). + /// The direction-branch slice is the first + /// `batch_size * B0 * num_atoms` floats (Branch 0). + pub fn on_b_logits_buf_ptr(&self) -> u64 { self.on_b_logits_buf.raw_ptr() } + + /// Raw pointer to atom_positions_buf [B0, num_atoms] f32, the C51 + /// adaptive (non-linear) atom positions used by all branches. Stable + /// across forward calls — written by the atom-positions ISV producer + /// kernel during training. + pub fn atom_positions_buf_ptr(&self) -> u64 { self.atom_positions_buf.raw_ptr() } + + /// Raw pointer to the per-sample C51 support buffer + /// [max_batch_size, 4, 3] stride-12 (v_min, v_max, delta_z) per branch + /// per sample. Owned by the GpuIqlTrainer and registered via + /// `set_per_sample_support_ptr`; this accessor returns whatever the + /// caller registered. Stable across forward calls. + pub fn per_sample_support_ptr_get(&self) -> u64 { self.per_sample_support_ptr } + /// F32 states buffer ref for cold-path Q-value computation. pub(crate) fn states_buf_f32(&self) -> &CudaSlice { &self.states_buf } diff --git a/crates/ml/src/cuda_pipeline/q_value_provider.rs b/crates/ml/src/cuda_pipeline/q_value_provider.rs index 96e90ad5d..dcc40832f 100644 --- a/crates/ml/src/cuda_pipeline/q_value_provider.rs +++ b/crates/ml/src/cuda_pipeline/q_value_provider.rs @@ -39,6 +39,54 @@ pub trait QValueProvider { q_out_ptr: u64, ) -> Result<(), MLError>; + /// Compute Q-values AND copy out raw branch-advantage C51 logits. + /// + /// Identical to `compute_q_values_to` but additionally DtoD-copies the + /// trainer's `on_b_logits_buf` (branch-major + /// `[batch_size, (B0+B1+B2+B3) * num_atoms]`) to `b_logits_out_ptr`. The + /// direction-branch slice (first `batch_size * B0 * num_atoms` floats) is + /// consumed by the Thompson direction sampler in + /// `experience_action_select` (Plan C T2 amendment, 2026-04-29). + /// + /// Per-sample support and atom_positions can be retrieved via the + /// `per_sample_support_ptr()` / `atom_positions_ptr()` accessors — those + /// buffers are owned by the trainer and stable across calls. + /// + /// `b_logits_out_ptr`: device pointer to write + /// `[batch_size, total_branch_atoms]` f32. Must be sized for the full + /// branch-major layout (kernel reads only the direction-branch slice but + /// the underlying DtoD copy mirrors the trainer's full buffer to keep the + /// caller's buffer self-consistent for diagnostics). + fn compute_q_and_b_logits_to( + &mut self, + states_ptr: u64, + batch_size: usize, + q_out_ptr: u64, + b_logits_out_ptr: u64, + ) -> Result<(), MLError>; + + /// Total branch-atom count = `(B0 + B1 + B2 + B3) * num_atoms`. Defines + /// the row stride of the buffer written by `compute_q_and_b_logits_to`. + fn total_branch_atoms(&self) -> usize; + + /// C51 atom count (`num_atoms`). Threaded into `experience_action_select` + /// as the `n_atoms` argument so the kernel reads the right stride. + fn num_atoms(&self) -> usize; + + /// Device pointer to the per-sample C51 support buffer + /// (`[max_batch_size, 4, 3]` stride-12 floats `(v_min, v_max, delta_z)`). + /// Owned by the trainer; stable across calls. Threaded into + /// `experience_action_select` as the `per_sample_support` argument. + fn per_sample_support_ptr(&self) -> u64; + + /// Device pointer to the C51 adaptive atom positions + /// (`[B0, num_atoms]` floats). Owned by the trainer; stable across + /// calls. Threaded into `experience_action_select` as the + /// `atom_positions` argument. Pass through unchanged — the kernel + /// internally falls back to linear support when this is NULL, but + /// production always supplies the adaptive buffer. + fn atom_positions_ptr(&self) -> u64; + /// CUDA stream handle for DtoD copies. fn stream_handle(&self) -> u64; diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 71e224511..bebf2e8a9 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -3527,6 +3527,78 @@ impl crate::cuda_pipeline::q_value_provider::QValueProvider for FusedTrainingCtx Ok(()) } + fn compute_q_and_b_logits_to( + &mut self, + states_ptr: u64, + batch_size: usize, + q_out_ptr: u64, + b_logits_out_ptr: u64, + ) -> Result<(), crate::MLError> { + // Mirror of `compute_q_values_to`, but additionally DtoD-copies the + // trainer's `on_b_logits_buf` (branch-major + // `[sub_b, total_actions * num_atoms]`) into the caller's output at + // the same per-iteration offset. Plan C T2 amendment buffer feed for + // `experience_action_select` Thompson direction sampling + // (gpu_backtest_evaluator val pipeline, Plan C T4, 2026-04-29). + let max_b = self.batch_size; + let tba = self.trainer.total_actions(); + let na = self.trainer.config().num_atoms; + let row_b_logits = tba * na; // total_branch_atoms + let f32_size = std::mem::size_of::(); + + for offset in (0..batch_size).step_by(max_b) { + let sub_b = (batch_size - offset).min(max_b); + let src_ptr = states_ptr + (offset * self.padded_state_dim() * f32_size) as u64; + + let dst = self.trainer.states_buf_ptr(); + self.trainer.launch_pad_states(dst, src_ptr, sub_b) + .map_err(|e| crate::MLError::ModelError(format!("QValueProvider pad: {e}")))?; + self.trainer.replay_forward_for_q_values(sub_b) + .map_err(|e| crate::MLError::ModelError(format!("QValueProvider: {e}")))?; + + // DtoD: q_out_buf → caller's q_out at offset. + let q_src = self.trainer.q_out_buf().raw_ptr(); + let q_dst = q_out_ptr + (offset * tba * f32_size) as u64; + let q_bytes = sub_b * tba * f32_size; + unsafe { + cudarc::driver::result::memcpy_dtod_async( + q_dst, q_src, q_bytes, self.stream.cu_stream(), + ).map_err(|e| crate::MLError::ModelError(format!("QValueProvider DtoD q: {e}")))?; + } + + // DtoD: on_b_logits_buf → caller's b_logits at offset. + // Branch-major layout `[B, total_actions * num_atoms]`; the + // direction-branch slice (Branch 0) lives at the start of each + // sample's row, matching the kernel's + // `b_logits_dir + i * b0 * n_atoms` indexing. + let b_src = self.trainer.on_b_logits_buf_ptr(); + let b_dst = b_logits_out_ptr + (offset * row_b_logits * f32_size) as u64; + let b_bytes = sub_b * row_b_logits * f32_size; + unsafe { + cudarc::driver::result::memcpy_dtod_async( + b_dst, b_src, b_bytes, self.stream.cu_stream(), + ).map_err(|e| crate::MLError::ModelError(format!("QValueProvider DtoD b_logits: {e}")))?; + } + } + Ok(()) + } + + fn total_branch_atoms(&self) -> usize { + self.trainer.total_actions() * self.trainer.config().num_atoms + } + + fn num_atoms(&self) -> usize { + self.trainer.config().num_atoms + } + + fn per_sample_support_ptr(&self) -> u64 { + self.trainer.per_sample_support_ptr_get() + } + + fn atom_positions_ptr(&self) -> u64 { + self.trainer.atom_positions_buf_ptr() + } + fn stream_handle(&self) -> u64 { self.stream.cu_stream() as u64 } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 114556d4b..0ce388f06 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -1072,6 +1072,36 @@ is now end-to-end runnable from the training rollout. Evaluator path the kernel signature in `experience_kernels.cu` at commit `5de5e546a`. +Plan C T4 evaluator wire-up (2026-04-29): same four amended buffers +threaded through the validation/backtest action_select call site at +`gpu_backtest_evaluator.rs:1722` (the chunked val pipeline inside +`submit_dqn_step_loop_cublas`). Unlike the collector, the evaluator +sources Q-values via the `QValueProvider` trait (delegating forward +to the trainer's CUDA-Graphed cuBLAS), so this commit (a) extends the +trait with `compute_q_and_b_logits_to(states_ptr, batch, q_out_ptr, +b_logits_out_ptr)` plus four read-only accessors +(`per_sample_support_ptr`, `atom_positions_ptr`, `num_atoms`, +`total_branch_atoms`), (b) implements them on `FusedTrainingCtx` +(`fused_training.rs`) — the b_logits variant DtoD-copies the +trainer's `on_b_logits_buf` into the caller's chunked buffer per +sub-batch iteration, mirroring the existing q_out DtoD pattern, and +(c) adds three pub accessors on the trainer +(`on_b_logits_buf_ptr`, `atom_positions_buf_ptr`, +`per_sample_support_ptr_get`). The evaluator gains one new field — +`chunked_b_logits_buf: Option>`, sized +`[n_windows * DQN_BACKTEST_CHUNK_SIZE, total_actions * num_atoms]` +plus 32*3 cuBLAS tail-safety floats — allocated in +`ensure_action_select_ready` alongside `chunked_q_values`. The +`ch_b_logits` device pointer is appended to the action_select +launch_builder; `per_sample_support` and `atom_positions` are read +via the new trait accessors right before the launch (trainer-owned, +sized for `max_batch_size ≥ chunk batch`); `n_atoms` is +`qp.num_atoms() as i32`. Phase 6's last-step plan_params forward +continues to use `compute_q_values_to` (b_logits not consumed there +— save_h_s2 + plan MLP only). Combined with T3, all production +callers of `experience_action_select` now use the amended kernel ABI +end-to-end (`feedback_no_partial_refactor`). Compile-clean. + Persistent picked_action_history scatter (2026-04-26): `gpu_backtest_evaluator.rs` gains a window-major `picked_action_history_buf` (parallel layout to `actions_history_buf`) populated by an additional `scatter_intent_chunk`