feat(dqn): wire C51 buffers to action_select in evaluator path

Threads b_logits_dir / per_sample_support / atom_positions / n_atoms
through experience_action_select kernel launch in
gpu_backtest_evaluator.rs's chunked val pipeline (line ~1722, inside
submit_dqn_step_loop_cublas).

The evaluator sources Q-values via the QValueProvider trait
(delegates forward to the trainer's CUDA-Graphed cuBLAS), so this
change extends the trait surface rather than duplicating the
forward:

- New trait method compute_q_and_b_logits_to(states_ptr, batch,
  q_out_ptr, b_logits_out_ptr) — DtoD-copies trainer's
  on_b_logits_buf into caller's chunked buffer per sub-batch
  iteration alongside the existing q_out copy.
- New trait accessors per_sample_support_ptr / atom_positions_ptr /
  num_atoms / total_branch_atoms (stable trainer-owned buffers).
- FusedTrainingCtx implements all of the above; trainer gains
  on_b_logits_buf_ptr / atom_positions_buf_ptr /
  per_sample_support_ptr_get pub accessors.
- Evaluator gains chunked_b_logits_buf field (sized
  [n_windows * CHUNK_SIZE, total_actions * num_atoms] + 32*3 tail
  safety), allocated in ensure_action_select_ready.

Phase 6's last-step plan_params forward keeps using the original
compute_q_values_to (b_logits not consumed there).

Plan C Task 4 — evaluator companion to T3 collector wire-up. After
this commit, all production callers of experience_action_select use
the amended kernel ABI (T2 amendment 5de5e546a) end-to-end
(feedback_no_partial_refactor).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

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<CudaSlice<f32>>,
 
+    /// 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<CudaSlice<f32>>,
+
     /// Batched states buffer for chunked gather: [n_windows * CHUNK_SIZE, state_dim].
     chunked_states_buf: Option<CudaSlice<f32>>,
 
@@ -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::<f32>(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::<f32>(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::<f32>(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);

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<f32> { &self.states_buf }
 

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;
 

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::<f32>();
+
+        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
     }

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<CudaSlice<f32>>`, 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`