gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: direct-to-trainer gather — eliminates upload_batch_gpu + 6 DtoD copies per step

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-19 00:34:10 +02:00
parent 5b197151db
commit b2135d01e0
5 changed files with 371 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

266
crates/ml-dqn/src/gpu_replay_buffer.rs
View File

@@ -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::<f32>();
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).

50
crates/ml-dqn/src/replay_buffer_kernels.cu
View File

@@ -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

50
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -1427,6 +1427,18 @@ pub struct GpuDqnTrainer {
speculative_features: Vec<f32>,
/// 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() }

36
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -985,12 +985,12 @@ impl FusedTrainingCtx {
) -> Result<FusedStepResult> {
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)

28
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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(),
);
}
}
}
}