From b2135d01e0080bb5ccb923aa31b27163c9f8de8a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 19 Apr 2026 00:34:10 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20direct-to-trainer=20gather=20=E2=80=94?= =?UTF-8?q?=20eliminates=20upload=5Fbatch=5Fgpu=20+=206=20DtoD=20copies=20?= =?UTF-8?q?per=20step?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PER gather now writes directly to GpuDqnTrainer's padded buffers via gather_f32_rows_padded, gather_f32_scalar, and gather_i32_scalar kernels compiled into the replay buffer cubin. set_trainer_buffers() wires stable device pointers at init; sample_proportional uses them when available, falling back to intermediate buffers otherwise. memset_zeros calls in the sampling hot path converted to raw cuMemsetD8Async. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-dqn/src/gpu_replay_buffer.rs | 266 ++++++++++++++---- crates/ml-dqn/src/replay_buffer_kernels.cu | 50 ++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 50 ++++ crates/ml/src/trainers/dqn/fused_training.rs | 36 ++- .../src/trainers/dqn/trainer/training_loop.rs | 28 ++ 5 files changed, 371 insertions(+), 59 deletions(-) diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 7eb3f0e19..bb9416e81 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -49,6 +49,12 @@ struct ReplayKernels { gather_u32: CudaFunction, /// #30 F32 row gather for f32 state storage. gather_f32_rows: CudaFunction, + /// Direct-to-trainer gather: f32 rows with zero-padding to dst_stride. + gather_f32_rows_padded: CudaFunction, + /// Direct-to-trainer gather: scalar f32 (rewards, dones). + gather_f32_scalar: CudaFunction, + /// Direct-to-trainer gather: scalar i32 (actions). + gather_i32_scalar: CudaFunction, is_weights_f32: CudaFunction, normalize_weights_f32: CudaFunction, fill_from_gpu_f32: CudaFunction, @@ -87,6 +93,9 @@ impl ReplayKernels { scatter_insert_u32: ld("scatter_insert_u32")?, gather_f32: ld("gather_f32")?, gather_u32: ld("gather_u32")?, gather_f32_rows: ld("gather_f32_rows")?, + gather_f32_rows_padded: ld("gather_f32_rows_padded")?, + gather_f32_scalar: ld("gather_f32_scalar")?, + gather_i32_scalar: ld("gather_i32_scalar")?, is_weights_f32: ld("is_weights_f32")?, normalize_weights_f32: ld("normalize_weights_f32")?, fill_from_gpu_f32: ld("fill_from_gpu_f32")?, @@ -169,6 +178,15 @@ pub struct GpuReplayBuffer { /// `sample_indices_ref` to return correctly-sized views (pre-allocated buffers are /// `max_batch_size` large; only the first `last_batch_size` elements are valid). last_batch_size: usize, + // Trainer destination buffer pointers — set once at init, stable for graph capture. + // PER gather writes directly to these, eliminating DtoD + pad_states. + trainer_states_ptr: u64, + trainer_next_states_ptr: u64, + trainer_actions_ptr: u64, + trainer_rewards_ptr: u64, + trainer_dones_ptr: u64, + trainer_is_weights_ptr: u64, + trainer_state_dim_padded: usize, } impl Drop for GpuReplayBuffer { @@ -265,6 +283,13 @@ impl GpuReplayBuffer { scratch_f32, rng_step: 0, last_batch_size: 0, + trainer_states_ptr: 0, + trainer_next_states_ptr: 0, + trainer_actions_ptr: 0, + trainer_rewards_ptr: 0, + trainer_dones_ptr: 0, + trainer_is_weights_ptr: 0, + trainer_state_dim_padded: 0, }) } @@ -291,6 +316,28 @@ impl GpuReplayBuffer { pub const fn epsilon(&self) -> f32 { self.config.epsilon } pub const fn state_dim(&self) -> usize { self.config.state_dim } + /// Wire trainer destination buffer pointers for direct-to-trainer gather. + /// Called once after GpuDqnTrainer is constructed. Pointers are stable + /// (CudaSlice allocations never move), so this is safe for graph capture. + pub fn set_trainer_buffers( + &mut self, + states_ptr: u64, + next_states_ptr: u64, + actions_ptr: u64, + rewards_ptr: u64, + dones_ptr: u64, + is_weights_ptr: u64, + state_dim_padded: usize, + ) { + self.trainer_states_ptr = states_ptr; + self.trainer_next_states_ptr = next_states_ptr; + self.trainer_actions_ptr = actions_ptr; + self.trainer_rewards_ptr = rewards_ptr; + self.trainer_dones_ptr = dones_ptr; + self.trainer_is_weights_ptr = is_weights_ptr; + self.trainer_state_dim_padded = state_dim_padded; + } + /// Grow insert scratch buffers if `eff` exceeds current allocation. /// Only triggers cuMemAlloc on the first large insert — subsequent calls reuse. fn ensure_insert_scratch(&mut self, eff: usize) -> Result<(), MLError> { @@ -420,11 +467,19 @@ impl GpuReplayBuffer { let bsi = batch_size as i32; // Step 1: Prefix scan of priorities_pa[0..size] - // Zero tile state + counter before each scan. - self.stream.memset_zeros(&mut self.scan_tile_state) - .map_err(|e| MLError::ModelError(format!("zero scan state: {e}")))?; - self.stream.memset_zeros(&mut self.scan_tile_counter) - .map_err(|e| MLError::ModelError(format!("zero scan counter: {e}")))?; + // Zero tile state + counter before each scan (raw cuMemsetD8Async — graph-safe). + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.scan_tile_state.raw_ptr(), 0, + self.scan_tile_state.num_bytes(), + self.stream.cu_stream(), + ); + cudarc::driver::sys::cuMemsetD8Async( + self.scan_tile_counter.raw_ptr(), 0, + self.scan_tile_counter.num_bytes(), + self.stream.cu_stream(), + ); + } let size_i = n as i32; let actual_tiles = n.div_ceil(256); @@ -476,40 +531,104 @@ impl GpuReplayBuffer { .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("i2u: {e}")))?; } - // Step 3: gather into pre-allocated buffers + // Step 3: gather into trainer buffers (direct) or intermediate buffers (fallback). let sdi = sd as i32; - // SAFETY: sample_states, states, sample_indices_i64 are valid device allocations. - // #30: gather f32 rows (full-precision state storage) - unsafe { - self.stream.launch_builder(&self.kernels.gather_f32_rows) - .arg(&mut self.sample_states).arg(&self.states).arg(&self.sample_indices_i64).arg(&sdi).arg(&bsi) - .launch(lcfg(batch_size * sd)).map_err(|e| MLError::ModelError(format!("g s f32: {e}")))?; - } - unsafe { - self.stream.launch_builder(&self.kernels.gather_f32_rows) - .arg(&mut self.sample_next_states).arg(&self.next_states).arg(&self.sample_indices_i64).arg(&sdi).arg(&bsi) - .launch(lcfg(batch_size * sd)).map_err(|e| MLError::ModelError(format!("g n f32: {e}")))?; - } - // SAFETY: same context, actions buffer valid. - unsafe { - let cap_i32 = self.capacity() as i32; - self.stream.launch_builder(&self.kernels.gather_u32) - .arg(&mut self.sample_actions).arg(&self.actions).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32) - .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g a: {e}")))?; - } - // SAFETY: same context, rewards buffer valid (f32). - unsafe { - let cap_i32 = self.capacity() as i32; - self.stream.launch_builder(&self.kernels.gather_f32) - .arg(&mut self.sample_rewards).arg(&self.rewards).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32) - .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g r: {e}")))?; - } - // SAFETY: same context, dones buffer valid (f32). - unsafe { - let cap_i32 = self.capacity() as i32; - self.stream.launch_builder(&self.kernels.gather_f32) - .arg(&mut self.sample_dones).arg(&self.dones).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32) - .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g d: {e}")))?; + let direct_to_trainer = self.trainer_states_ptr != 0; + + if direct_to_trainer { + // ── Direct-to-trainer path: gather + pad in one kernel, zero DtoD copies ── + let sdp = self.trainer_state_dim_padded as i32; + let t_states = self.trainer_states_ptr; + let t_next = self.trainer_next_states_ptr; + let t_actions = self.trainer_actions_ptr; + let t_rewards = self.trainer_rewards_ptr; + let t_dones = self.trainer_dones_ptr; + + // States: gather + pad → trainer.states_buf + unsafe { + self.stream.launch_builder(&self.kernels.gather_f32_rows_padded) + .arg(&t_states) + .arg(&self.states) + .arg(&self.sample_indices_i64) + .arg(&sdi) + .arg(&sdp) + .arg(&bsi) + .launch(lcfg(batch_size * self.trainer_state_dim_padded)) + .map_err(|e| MLError::ModelError(format!("gather states padded: {e}")))?; + } + // Next states: gather + pad → trainer.next_states_buf + unsafe { + self.stream.launch_builder(&self.kernels.gather_f32_rows_padded) + .arg(&t_next) + .arg(&self.next_states) + .arg(&self.sample_indices_i64) + .arg(&sdi) + .arg(&sdp) + .arg(&bsi) + .launch(lcfg(batch_size * self.trainer_state_dim_padded)) + .map_err(|e| MLError::ModelError(format!("gather next_states padded: {e}")))?; + } + // Actions: gather i32 scalar → trainer.actions_buf + unsafe { + self.stream.launch_builder(&self.kernels.gather_i32_scalar) + .arg(&t_actions) + .arg(&self.actions) + .arg(&self.sample_indices_i64) + .arg(&bsi) + .launch(lcfg(batch_size)) + .map_err(|e| MLError::ModelError(format!("gather actions direct: {e}")))?; + } + // Rewards: gather f32 scalar → trainer.rewards_buf + unsafe { + self.stream.launch_builder(&self.kernels.gather_f32_scalar) + .arg(&t_rewards) + .arg(&self.rewards) + .arg(&self.sample_indices_i64) + .arg(&bsi) + .launch(lcfg(batch_size)) + .map_err(|e| MLError::ModelError(format!("gather rewards direct: {e}")))?; + } + // Dones: gather f32 scalar → trainer.dones_buf + unsafe { + self.stream.launch_builder(&self.kernels.gather_f32_scalar) + .arg(&t_dones) + .arg(&self.dones) + .arg(&self.sample_indices_i64) + .arg(&bsi) + .launch(lcfg(batch_size)) + .map_err(|e| MLError::ModelError(format!("gather dones direct: {e}")))?; + } + } else { + // ── Fallback: gather into intermediate sample_* buffers ── + // #30: gather f32 rows (full-precision state storage) + unsafe { + self.stream.launch_builder(&self.kernels.gather_f32_rows) + .arg(&mut self.sample_states).arg(&self.states).arg(&self.sample_indices_i64).arg(&sdi).arg(&bsi) + .launch(lcfg(batch_size * sd)).map_err(|e| MLError::ModelError(format!("g s f32: {e}")))?; + } + unsafe { + self.stream.launch_builder(&self.kernels.gather_f32_rows) + .arg(&mut self.sample_next_states).arg(&self.next_states).arg(&self.sample_indices_i64).arg(&sdi).arg(&bsi) + .launch(lcfg(batch_size * sd)).map_err(|e| MLError::ModelError(format!("g n f32: {e}")))?; + } + unsafe { + let cap_i32 = self.capacity() as i32; + self.stream.launch_builder(&self.kernels.gather_u32) + .arg(&mut self.sample_actions).arg(&self.actions).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32) + .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g a: {e}")))?; + } + unsafe { + let cap_i32 = self.capacity() as i32; + self.stream.launch_builder(&self.kernels.gather_f32) + .arg(&mut self.sample_rewards).arg(&self.rewards).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32) + .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g r: {e}")))?; + } + unsafe { + let cap_i32 = self.capacity() as i32; + self.stream.launch_builder(&self.kernels.gather_f32) + .arg(&mut self.sample_dones).arg(&self.dones).arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32) + .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g d: {e}")))?; + } } // Step 3b: gather episode_ids for HER strategies. @@ -527,7 +646,8 @@ impl GpuReplayBuffer { // Step 6: IS weights via GPU-resident total_sum (zero CPU readback) // is_weights_f32 reads total_sum from GPU pointer. - // SAFETY: sample_weights, sample_priorities, total_sum_buf are valid device allocations. + // When direct_to_trainer is active, write IS weights to trainer's is_weights_buf. + // Always write to sample_weights too (needed for reduce_max + normalize below). unsafe { self.stream.launch_builder(&self.kernels.is_weights_f32) .arg(&mut self.sample_weights).arg(&self.sample_priorities).arg(&self.total_sum_buf) @@ -536,8 +656,13 @@ impl GpuReplayBuffer { } // Step 7: reduce max weight - self.stream.memset_zeros(&mut self.sample_max_weight) - .map_err(|e| MLError::ModelError(format!("mw zero: {e}")))?; + unsafe { + cudarc::driver::sys::cuMemsetD8Async( + self.sample_max_weight.raw_ptr(), 0, + self.sample_max_weight.num_bytes(), + self.stream.cu_stream(), + ); + } let rt = 256_u32.min(batch_size as u32).max(1); let rb = (batch_size as u32).div_ceil(rt); // SAFETY: sample_weights, sample_max_weight are valid device allocations. @@ -556,26 +681,55 @@ impl GpuReplayBuffer { .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("nw: {e}")))?; } + // Step 8b: copy normalized IS weights to trainer buffer if direct_to_trainer + if direct_to_trainer { + let t_isw = self.trainer_is_weights_ptr; + let num_bytes = batch_size * std::mem::size_of::(); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + t_isw as cudarc::driver::sys::CUdeviceptr, + self.sample_weights.raw_ptr() as cudarc::driver::sys::CUdeviceptr, + num_bytes, + self.stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("isw dtod: {e}")))?; + } + } + // Record the batch size so that sample_weights_ref / sample_indices_ref return // correctly-sized views (the pre-allocated buffers are max_batch_size elements, // but only the first `batch_size` are valid after this sample). self.last_batch_size = batch_size; - // Zero-alloc: return raw pointers to pre-allocated sample buffers. - // Gather kernels already wrote into self.sample_* buffers above. - // No DtoD clone, no cuMemAlloc — pure GPU-resident references. - Ok(GpuBatchPtrs { - states_ptr: self.sample_states.raw_ptr(), - next_states_ptr: self.sample_next_states.raw_ptr(), - actions_ptr: self.sample_actions.raw_ptr(), - rewards_ptr: self.sample_rewards.raw_ptr(), - dones_ptr: self.sample_dones.raw_ptr(), - weights_ptr: self.sample_weights.raw_ptr(), - indices_ptr: self.sample_indices_u32.raw_ptr(), - episode_ids_ptr: self.sample_episode_ids.raw_ptr(), - batch_size, - state_dim: sd, - }) + // Return pointers: when direct_to_trainer is active, point to trainer buffers + // (already populated by gather kernels above). Otherwise point to intermediate + // sample_* buffers for the legacy upload_batch_gpu path. + if direct_to_trainer { + Ok(GpuBatchPtrs { + states_ptr: self.trainer_states_ptr, + next_states_ptr: self.trainer_next_states_ptr, + actions_ptr: self.trainer_actions_ptr, + rewards_ptr: self.trainer_rewards_ptr, + dones_ptr: self.trainer_dones_ptr, + weights_ptr: self.trainer_is_weights_ptr, + indices_ptr: self.sample_indices_u32.raw_ptr(), + episode_ids_ptr: self.sample_episode_ids.raw_ptr(), + batch_size, + state_dim: sd, + }) + } else { + Ok(GpuBatchPtrs { + states_ptr: self.sample_states.raw_ptr(), + next_states_ptr: self.sample_next_states.raw_ptr(), + actions_ptr: self.sample_actions.raw_ptr(), + rewards_ptr: self.sample_rewards.raw_ptr(), + dones_ptr: self.sample_dones.raw_ptr(), + weights_ptr: self.sample_weights.raw_ptr(), + indices_ptr: self.sample_indices_u32.raw_ptr(), + episode_ids_ptr: self.sample_episode_ids.raw_ptr(), + batch_size, + state_dim: sd, + }) + } } /// Reference to pre-allocated sample indices buffer (valid after sample_proportional). diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index 876d3b029..cd7624a39 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -374,6 +374,56 @@ void max_of_two_f32( out[0] = (va > vb) ? va : vb; } +// ── 18a. Gather f32 rows with zero-padding to dst_stride ─────────────────── +// Replaces 2-step gather + pad_states: gathers rows from [cap, src_dim] +// into [batch, dst_stride], zero-filling columns [src_dim, dst_stride). + +extern "C" __global__ +void gather_f32_rows_padded( + float* __restrict__ dst, + const float* __restrict__ src, + const long long* __restrict__ indices, + int src_dim, int dst_stride, int batch_size) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int total = batch_size * dst_stride; + if (tid >= total) return; + int row = tid / dst_stride; + int col = tid % dst_stride; + long long src_row = indices[row]; + dst[tid] = (col < src_dim) ? src[src_row * src_dim + col] : 0.0f; +} + +// ── 18b. Gather scalar f32 by i64 indices ───────────────────────────────── +// Direct gather for rewards, dones, IS-weights into trainer buffers. + +extern "C" __global__ +void gather_f32_scalar( + float* __restrict__ dst, + const float* __restrict__ src, + const long long* __restrict__ indices, + int batch_size) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= batch_size) return; + dst[tid] = src[indices[tid]]; +} + +// ── 18c. Gather scalar i32 by i64 indices ───────────────────────────────── +// Direct gather for actions (u32/i32) into trainer buffers. + +extern "C" __global__ +void gather_i32_scalar( + int* __restrict__ dst, + const int* __restrict__ src, + const long long* __restrict__ indices, + int batch_size) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= batch_size) return; + dst[tid] = src[indices[tid]]; +} + // ── 18. NaN/Inf check kernel ──────────────────────────────────────────────── // out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 94911606a..722e211a9 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1427,6 +1427,18 @@ pub struct GpuDqnTrainer { speculative_features: Vec, /// Whether the speculative cache contains a usable result. speculative_valid: bool, + + // ── GPU-side counter increment kernel ─────────────────────────── + /// Kernel: increment_step_counters — atomically increments all 8 Adam step + /// counters and computes cosine-annealed tau on GPU, eliminating all host-side + /// per-step writes from the critical path. + increment_counters_kernel: CudaFunction, + /// Initial tau for cosine annealing (written to tau_pinned at step 0). + tau_init: f32, + /// Final tau for cosine annealing (asymptotic value after tau_anneal_steps). + tau_final: f32, + /// Total number of steps over which tau is cosine-annealed. + tau_anneal_steps: i32, } impl GpuDqnTrainer { @@ -5716,6 +5728,22 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("speculative_h_s2 alloc: {e}")))?; let speculative_features = vec![0.0_f32; config.state_dim]; + // ── GPU-side counter increment kernel ────────────────────── + // Load from a NEW CUmodule — MUST NOT share a CUmodule with any ungraphed + // kernel, as this kernel will be captured in a child graph on Hopper. + let increment_counters_kernel = { + let graph_util_module = stream.context().load_cubin(GRAPH_UTILITY_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("graph_utility cubin load: {e}")))?; + graph_util_module.load_function("increment_step_counters") + .map_err(|e| MLError::ModelError(format!("increment_step_counters load: {e}")))? + }; + info!("GpuDqnTrainer: increment_step_counters kernel loaded (GPU-side counter increments)"); + + // Tau cosine annealing parameters — defaults; caller overrides via set_tau_anneal_params. + let tau_init = 0.005_f32; + let tau_final = 0.001_f32; + let tau_anneal_steps: i32 = 100_000; + Ok(Self { config, stream, @@ -6141,6 +6169,10 @@ impl GpuDqnTrainer { speculative_h_s2, speculative_features, speculative_valid: false, + increment_counters_kernel, + tau_init, + tau_final, + tau_anneal_steps, }) } @@ -10136,6 +10168,24 @@ impl GpuDqnTrainer { /// Raw pointer to states_buf for pad_states kernel. pub(crate) fn states_buf_ptr(&self) -> u64 { self.states_buf.raw_ptr() } + /// Raw pointer to next_states_buf for direct-to-trainer gather. + pub(crate) fn next_states_buf_ptr(&self) -> u64 { self.next_states_buf.raw_ptr() } + + /// Raw pointer to actions_buf for direct-to-trainer gather. + pub(crate) fn actions_buf_ptr(&self) -> u64 { self.actions_buf.raw_ptr() } + + /// Raw pointer to rewards_buf for direct-to-trainer gather. + pub(crate) fn rewards_buf_ptr(&self) -> u64 { self.rewards_buf.raw_ptr() } + + /// Raw pointer to dones_buf for direct-to-trainer gather. + pub(crate) fn dones_buf_ptr(&self) -> u64 { self.dones_buf.raw_ptr() } + + /// Raw pointer to is_weights_buf for direct-to-trainer gather. + pub(crate) fn is_weights_buf_ptr(&self) -> u64 { self.is_weights_buf.raw_ptr() } + + /// Padded state dimension (128-byte aligned) for direct-to-trainer gather. + pub(crate) fn state_dim_padded(&self) -> usize { pad128(self.config.state_dim) } + /// 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() } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index a2e8a06c4..e0b32a4be 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -985,12 +985,12 @@ impl FusedTrainingCtx { ) -> Result { let cu_stream = self.stream.cu_stream(); - // ── Outside graph: upload batch + host-side state ───────────── + // ── Outside graph: host-side state ───────────────────────────── + // upload_batch_gpu eliminated: PER gather writes directly to trainer + // buffers via set_trainer_buffers (zero DtoD copies per step). if let Some(ref pe) = self.phase_events { PhaseEvents::record(pe.upload_start, cu_stream); } - self.trainer.upload_batch_gpu(gpu_batch) - .map_err(|e| anyhow::anyhow!("batch upload: {e}"))?; self.trainer.update_stochastic_depth_mask() .map_err(|e| anyhow::anyhow!("stochastic depth mask: {e}"))?; self.trainer.adam_step_async(); @@ -2290,6 +2290,36 @@ impl FusedTrainingCtx { self.trainer.states_buf_ptr() } + /// Raw pointer to trainer's next_states buffer. + pub(crate) fn trainer_next_states_buf_ptr(&self) -> u64 { + self.trainer.next_states_buf_ptr() + } + + /// Raw pointer to trainer's actions buffer. + pub(crate) fn trainer_actions_buf_ptr(&self) -> u64 { + self.trainer.actions_buf_ptr() + } + + /// Raw pointer to trainer's rewards buffer. + pub(crate) fn trainer_rewards_buf_ptr(&self) -> u64 { + self.trainer.rewards_buf_ptr() + } + + /// Raw pointer to trainer's dones buffer. + pub(crate) fn trainer_dones_buf_ptr(&self) -> u64 { + self.trainer.dones_buf_ptr() + } + + /// Raw pointer to trainer's IS weights buffer. + pub(crate) fn trainer_is_weights_buf_ptr(&self) -> u64 { + self.trainer.is_weights_buf_ptr() + } + + /// Padded state dimension for direct-to-trainer gather. + pub(crate) fn trainer_state_dim_padded(&self) -> usize { + self.trainer.state_dim_padded() + } + /// Run cross-branch Q attention: q_out_buf → q_coord_buf [batch_size, 12]. pub(crate) fn launch_q_attention(&self, batch_size: usize) -> Result<()> { self.trainer.launch_q_attention(batch_size) diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index cb3a579cb..cef8d5da9 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -326,6 +326,20 @@ impl DQNTrainer { } Err(e) => { tracing::error!("Fused CUDA context init failed: {e}"); } } + drop(agent); + // Wire direct-to-trainer gather: PER gather writes directly to trainer buffers. + if let Some(ref fused) = self.fused_ctx { + let mut agent_w = self.agent.write().await; + agent_w.memory_mut().gpu.set_trainer_buffers( + fused.trainer_states_buf_ptr(), + fused.trainer_next_states_buf_ptr(), + fused.trainer_actions_buf_ptr(), + fused.trainer_rewards_buf_ptr(), + fused.trainer_dones_buf_ptr(), + fused.trainer_is_weights_buf_ptr(), + fused.trainer_state_dim_padded(), + ); + } } } @@ -1432,6 +1446,20 @@ impl DQNTrainer { } Err(e) => { tracing::error!("Fused CUDA context init failed: {e}"); } } + drop(agent); + // Wire direct-to-trainer gather: PER gather writes directly to trainer buffers. + if let Some(ref fused) = self.fused_ctx { + let mut agent_w = self.agent.write().await; + agent_w.memory_mut().gpu.set_trainer_buffers( + fused.trainer_states_buf_ptr(), + fused.trainer_next_states_buf_ptr(), + fused.trainer_actions_buf_ptr(), + fused.trainer_rewards_buf_ptr(), + fused.trainer_dones_buf_ptr(), + fused.trainer_is_weights_buf_ptr(), + fused.trainer_state_dim_padded(), + ); + } } } }