diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 1e6ab15a0..b3dc766ac 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -190,6 +190,10 @@ pub struct GpuReplayBuffer { size: usize, max_priority: f32, + // Deferred max-priority readback: accumulate batch maxima on GPU, + // flush to CPU once per epoch instead of once per train_step (~8x fewer syncs). + pending_max_priority: Option, + // Beta annealing current_step: usize, } @@ -241,6 +245,7 @@ impl GpuReplayBuffer { write_cursor: 0, size: 0, max_priority: 1.0, + pending_max_priority: None, current_step: 0, }) } @@ -280,10 +285,11 @@ impl GpuReplayBuffer { } /// Clear all experiences. Resets cursor and size but keeps allocations. - pub const fn clear(&mut self) { + pub fn clear(&mut self) { self.write_cursor = 0; self.size = 0; self.max_priority = 1.0; + self.pending_max_priority = None; self.current_step = 0; } @@ -573,12 +579,17 @@ impl GpuReplayBuffer { // Clamp to [epsilon, max_reasonable] to prevent NaN/Inf let clamped = new_prios.clamp(self.config.epsilon, 1e6)?; - // Update max_priority (single scalar readback — unavoidable for CPU-side state - // used by insert_batch to assign max_priority to new experiences). - let batch_max = clamped.max(0)?.to_vec0::()?; - if batch_max > self.max_priority { - self.max_priority = batch_max; - } + // Accumulate batch max on GPU — deferred readback via flush_max_priority(). + // insert_batch uses the PREVIOUS max_priority (slightly stale is fine for + // proportional PER: new experiences still get sampled first regardless). + let batch_max = clamped.max(0)?; + self.pending_max_priority = Some(match self.pending_max_priority.take() { + Some(prev) => { + // GPU-side max of previous accumulated max and this batch's max + Tensor::stack(&[&prev, &batch_max], 0)?.max(0)? + } + None => batch_max, + }); // Delta trick: gather old priorities at indices, compute delta = new - old, // then index_add the deltas back. Single batched GPU kernel, no CPU loop. @@ -594,6 +605,22 @@ impl GpuReplayBuffer { Ok(()) } + + /// Flush the GPU-accumulated max priority to CPU with a single scalar readback. + /// + /// Call once per epoch (after all `update_priorities_gpu` calls) instead of + /// reading back per-batch. Reduces GPU-CPU sync barriers from ~8/epoch to 1. + pub fn flush_max_priority(&mut self) -> Result<(), MLError> { + if let Some(pending) = self.pending_max_priority.take() { + let max_val = pending.to_scalar::().map_err(|e| { + MLError::ModelError(format!("Priority flush readback failed: {e}")) + })?; + if max_val > self.max_priority { + self.max_priority = max_val; + } + } + Ok(()) + } } impl std::fmt::Debug for GpuReplayBuffer { @@ -832,6 +859,69 @@ mod tests { assert!(buf.max_priority >= buf.config.epsilon); } + #[test] + fn test_flush_max_priority() { + let mut buf = make_filled_buffer(50); + let old_max = buf.max_priority; + + // Update priorities — max is now pending on GPU, not yet flushed to CPU + let batch = buf.sample_proportional(16).expect("sample"); + let td_errors = Tensor::full(5.0_f32, &[16], &Device::Cpu).expect("td"); + buf.update_priorities_gpu(&batch.indices, &td_errors).expect("update"); + + // pending_max_priority should be Some (accumulated on GPU) + assert!(buf.pending_max_priority.is_some(), "should have pending max"); + + // CPU-side max_priority unchanged until flush + assert!((buf.max_priority - old_max).abs() < 1e-6, "max_priority should not change before flush"); + + // Flush: single GPU→CPU readback + buf.flush_max_priority().expect("flush"); + assert!(buf.pending_max_priority.is_none(), "pending should be consumed"); + assert!(buf.max_priority >= old_max, "max_priority should be >= old value after flush"); + } + + #[test] + fn test_flush_max_priority_multiple_batches() { + let mut buf = make_filled_buffer(50); + + // Simulate multiple train_step calls within one epoch + for _ in 0..4 { + let batch = buf.sample_proportional(8).expect("sample"); + let td_errors = Tensor::randn(0.0_f32, 1.0, &[8], &Device::Cpu).expect("td"); + buf.update_priorities_gpu(&batch.indices, &td_errors).expect("update"); + } + + // All 4 batch maxima accumulated on GPU, no CPU readbacks yet + assert!(buf.pending_max_priority.is_some()); + + // Single flush at epoch boundary + buf.flush_max_priority().expect("flush"); + assert!(buf.pending_max_priority.is_none()); + assert!(buf.max_priority >= buf.config.epsilon); + } + + #[test] + fn test_flush_noop_when_no_pending() { + let mut buf = make_filled_buffer(50); + // No update_priorities_gpu called — flush should be a no-op + buf.flush_max_priority().expect("flush should succeed with nothing pending"); + assert!(buf.pending_max_priority.is_none()); + } + + #[test] + fn test_clear_resets_pending_max_priority() { + let mut buf = make_filled_buffer(50); + let batch = buf.sample_proportional(8).expect("sample"); + let td_errors = Tensor::randn(0.0_f32, 1.0, &[8], &Device::Cpu).expect("td"); + buf.update_priorities_gpu(&batch.indices, &td_errors).expect("update"); + assert!(buf.pending_max_priority.is_some()); + + buf.clear(); + assert!(buf.pending_max_priority.is_none(), "clear should reset pending"); + assert!((buf.max_priority - 1.0).abs() < 1e-6, "clear should reset max_priority to 1.0"); + } + #[test] fn test_proportional_sample_insufficient_data() { let buf = make_filled_buffer(5); diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index a3f8dc951..5a24e0f5e 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -65,7 +65,10 @@ impl StagedGpuBuffer { return Ok(()); } let n = self.staging.len(); - let sd = self.staging.first().map_or(0, |e| e.state.len()); + // Use the GPU buffer's configured state_dim (tensor-core-aligned) rather + // than the raw experience dimension, so the flushed tensors match the + // pre-allocated ring buffer shape and slice_scatter succeeds. + let sd = self.gpu.state_dim(); let device = self.gpu.device().clone(); let mut states_flat = Vec::with_capacity(n * sd); @@ -76,6 +79,7 @@ impl StagedGpuBuffer { for exp in self.staging.drain(..) { states_flat.extend_from_slice(&exp.state); + // Zero-pad to buffer's state_dim (handles raw→aligned padding) if exp.state.len() < sd { states_flat.resize(states_flat.len() + sd - exp.state.len(), 0.0); } @@ -376,6 +380,20 @@ impl ReplayBufferType { } } + /// Flush GPU-accumulated max priority to CPU (single scalar readback). + /// + /// Call once per epoch after all `update_priorities_gpu` calls. + /// No-op for non-GPU buffers. + #[cfg(feature = "cuda")] + pub fn flush_max_priority(&self) -> Result<(), MLError> { + match self { + Self::GpuPrioritized(buffer) => { + buffer.lock().gpu.flush_max_priority() + } + Self::Uniform(_) | Self::Prioritized(_) => Ok(()), + } + } + /// Step training counter for beta annealing (PER only) pub fn step(&self) { match self { diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 7ecf70469..d44933fe9 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -143,6 +143,115 @@ __device__ void noisy_matvec_leaky_relu( } } +/* ------------------------------------------------------------------ */ +/* Shared Memory Weight Caching */ +/* ------------------------------------------------------------------ */ + +/* Shared memory tile size: rows per tile × max input dim. + * 64 rows × 256 cols = 16384 floats = 64 KB. Fits H100's 228 KB shmem. + * Overridable via NVRTC #define injection. */ +#ifndef SHMEM_TILE_ROWS +#define SHMEM_TILE_ROWS 64 +#endif +#define SHMEM_MAX_IN_DIM 256 /* max(STATE_DIM, SHARED_H1, SHARED_H2) */ + +/** + * Cooperatively load a weight tile from global to shared memory. + * All threads in the block participate. Must call __syncthreads() after. + * Uses float4 vectorized loads where possible for 4x bandwidth. + */ +__device__ void cooperative_load_tile( + float* __restrict__ shmem_dst, + const float* __restrict__ global_src, + int count /* total floats to load */ +) { + int vec4_count = count / 4; + int remainder = count % 4; + float4* dst4 = (float4*)shmem_dst; + const float4* src4 = (const float4*)global_src; + for (int i = threadIdx.x; i < vec4_count; i += blockDim.x) { + dst4[i] = src4[i]; + } + int base = vec4_count * 4; + for (int i = threadIdx.x; i < remainder; i += blockDim.x) { + shmem_dst[base + i] = global_src[base + i]; + } +} + +/** + * Matrix-vector multiply from shared memory: output = shmem_W * input + shmem_b. + * + * Reads weights/bias from shared memory instead of global (HBM). + * Only computes `tile_rows` output rows starting at `tile_offset`. + * Weight layout: shmem_W[tile_rows, in_dim] row-major (tile-local indexing). + */ +__device__ void matvec_leaky_relu_shmem( + const float* __restrict__ shmem_W, + const float* __restrict__ shmem_b, + const float* input, + float* output, + int in_dim, + int out_dim, + int tile_offset, + int tile_rows, + int activate +) { + (void)out_dim; /* full output dim — unused, tile_rows is the active count */ + for (int j = 0; j < tile_rows; j++) { + float acc = shmem_b[j]; + const float* row = shmem_W + j * in_dim; + for (int i = 0; i < in_dim; i++) { + acc += row[i] * input[i]; + } + output[tile_offset + j] = activate ? leaky_relu(acc) : acc; + } +} + +/** + * Noisy matrix-vector multiply from shared memory with factorized Gaussian noise. + * + * Same as noisy_matvec_leaky_relu but reads W/b from shared memory. + * Only computes `tile_rows` output rows starting at `tile_offset`. + */ +__device__ void noisy_matvec_leaky_relu_shmem( + const float* __restrict__ shmem_W, + const float* __restrict__ shmem_b, + const float* input, + float* output, + int in_dim, + int out_dim, + int tile_offset, + int tile_rows, + int activate, + float sigma_init, + unsigned int* rng +) { + (void)out_dim; + float sigma_scale = sigma_init / sqrtf((float)in_dim); + + /* Generate factorized noise vectors for this tile */ + float eps_in[NOISY_MAX_DIM]; + float eps_out[NOISY_MAX_DIM]; + + for (int i = 0; i < in_dim && i < NOISY_MAX_DIM; i++) { + eps_in[i] = factorized_noise_fn(gpu_random_gaussian(rng)); + } + for (int j = 0; j < tile_rows && j < NOISY_MAX_DIM; j++) { + eps_out[j] = factorized_noise_fn(gpu_random_gaussian(rng)); + } + + for (int j = 0; j < tile_rows; j++) { + float ej = (j < NOISY_MAX_DIM) ? eps_out[j] : 0.0f; + float acc = shmem_b[j] + sigma_scale * ej; + const float* row = shmem_W + j * in_dim; + for (int i = 0; i < in_dim; i++) { + float ei = (i < NOISY_MAX_DIM) ? eps_in[i] : 0.0f; + acc += (row[i] + sigma_scale * ej * ei) * input[i]; + } + output[tile_offset + j] = activate ? leaky_relu(acc) : acc; + } +} + /** * In-place RMSNorm: data[i] = (data[i] / RMS(data)) * gamma[i]. * diff --git a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu index 6b8387d72..edcbd862f 100644 --- a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu @@ -27,6 +27,11 @@ #endif #define LEAKY_RELU_ALPHA 0.01f +/* Min helper — avoid collisions with platform min macro */ +#ifndef SHMEM_MIN +#define SHMEM_MIN(a, b) (((a) < (b)) ? (a) : (b)) +#endif + /* ------------------------------------------------------------------ */ /* DQN-Specific Device Functions */ /* ------------------------------------------------------------------ */ @@ -96,6 +101,280 @@ __device__ void q_forward_dueling( } } +/* ------------------------------------------------------------------ */ +/* Shared Memory Dueling Forward Passes */ +/* ------------------------------------------------------------------ */ + +/** + * Helper macro: tile a single layer through shared memory. + * + * Loads weight tiles from global → shared, then each thread computes + * its output rows from the tile. __syncthreads() brackets each tile. + * + * For layers smaller than SHMEM_TILE_ROWS (e.g. value output [1, 128]), + * a single tile covers the entire layer — no loop overhead. + */ +#define TILE_LAYER_CLEAN(W_global, b_global, input, output, in_d, out_d, act, shmem_w, shmem_b) \ + do { \ + for (int _tile = 0; _tile < ((out_d) + SHMEM_TILE_ROWS - 1) / SHMEM_TILE_ROWS; _tile++) { \ + int _ts = _tile * SHMEM_TILE_ROWS; \ + int _tr = SHMEM_MIN(SHMEM_TILE_ROWS, (out_d) - _ts); \ + cooperative_load_tile(shmem_w, (W_global) + _ts * (in_d), _tr * (in_d)); \ + cooperative_load_tile(shmem_b, (b_global) + _ts, _tr); \ + __syncthreads(); \ + matvec_leaky_relu_shmem(shmem_w, shmem_b, input, output, \ + in_d, out_d, _ts, _tr, act); \ + __syncthreads(); \ + } \ + } while (0) + +#define TILE_LAYER_NOISY(W_global, b_global, input, output, in_d, out_d, act, shmem_w, shmem_b, sigma, rng_ptr) \ + do { \ + for (int _tile = 0; _tile < ((out_d) + SHMEM_TILE_ROWS - 1) / SHMEM_TILE_ROWS; _tile++) { \ + int _ts = _tile * SHMEM_TILE_ROWS; \ + int _tr = SHMEM_MIN(SHMEM_TILE_ROWS, (out_d) - _ts); \ + cooperative_load_tile(shmem_w, (W_global) + _ts * (in_d), _tr * (in_d)); \ + cooperative_load_tile(shmem_b, (b_global) + _ts, _tr); \ + __syncthreads(); \ + noisy_matvec_leaky_relu_shmem(shmem_w, shmem_b, input, output, \ + in_d, out_d, _ts, _tr, act, sigma, rng_ptr); \ + __syncthreads(); \ + } \ + } while (0) + +/** + * Shared-memory dueling Q-network forward pass. + * + * Same architecture as q_forward_dueling but reads weights via cooperative + * shared memory tiling instead of direct global memory access. All threads + * in the block participate in loading each tile, amortizing HBM bandwidth. + * + * Layer tiling at default sizes (SHMEM_TILE_ROWS=64): + * s1 [256, 48]: 4 tiles + * s2 [256, 256]: 4 tiles + * v1 [128, 256]: 2 tiles + * v2 [1, 128]: 1 tile + * a1 [128, 256]: 2 tiles + * a2 [5, 128]: 1 tile + * Total: 14 tiles per forward pass (was 6 full global reads) + */ +__device__ void q_forward_dueling_shmem( + const float* state, + const float* __restrict__ w_s1, const float* __restrict__ b_s1, + const float* __restrict__ w_s2, const float* __restrict__ b_s2, + const float* __restrict__ w_v1, const float* __restrict__ b_v1, + const float* __restrict__ w_v2, const float* __restrict__ b_v2, + const float* __restrict__ w_a1, const float* __restrict__ b_a1, + const float* __restrict__ w_a2, const float* __restrict__ b_a2, + float* scratch1, float* scratch2, + float* scratch_v, float* scratch_a, + float* q_values, + float* shmem_weights, + float* shmem_bias +) { + /* Shared layers */ + TILE_LAYER_CLEAN(w_s1, b_s1, state, scratch1, STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias); + TILE_LAYER_CLEAN(w_s2, b_s2, scratch1, scratch2, SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias); + + /* Value head */ + TILE_LAYER_CLEAN(w_v1, b_v1, scratch2, scratch_v, SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias); + /* Value output: [1, VALUE_H] — single tile, trivially small */ + float value; + { + cooperative_load_tile(shmem_weights, w_v2, VALUE_H); + cooperative_load_tile(shmem_bias, b_v2, 1); + __syncthreads(); + float acc = shmem_bias[0]; + for (int i = 0; i < VALUE_H; i++) acc += shmem_weights[i] * scratch_v[i]; + value = acc; + __syncthreads(); + } + + /* Advantage head */ + TILE_LAYER_CLEAN(w_a1, b_a1, scratch2, scratch_a, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias); + float adv[NUM_ACTIONS]; + TILE_LAYER_CLEAN(w_a2, b_a2, scratch_a, adv, ADV_H, NUM_ACTIONS, 0, shmem_weights, shmem_bias); + + /* Mean advantage */ + float adv_mean = 0.0f; + for (int i = 0; i < NUM_ACTIONS; i++) adv_mean += adv[i]; + adv_mean /= (float)NUM_ACTIONS; + + /* Q(s,a) = V(s) + A(s,a) - mean(A) */ + for (int i = 0; i < NUM_ACTIONS; i++) { + q_values[i] = value + adv[i] - adv_mean; + } +} + +/** + * Shared-memory NoisyNet dueling Q-network forward pass. + * + * Same tiling as q_forward_dueling_shmem but with factorized Gaussian noise + * on each layer. Only used for the ONLINE network during exploration. + */ +__device__ void q_forward_dueling_noisy_shmem( + const float* state, + const float* __restrict__ w_s1, const float* __restrict__ b_s1, + const float* __restrict__ w_s2, const float* __restrict__ b_s2, + const float* __restrict__ w_v1, const float* __restrict__ b_v1, + const float* __restrict__ w_v2, const float* __restrict__ b_v2, + const float* __restrict__ w_a1, const float* __restrict__ b_a1, + const float* __restrict__ w_a2, const float* __restrict__ b_a2, + float* scratch1, float* scratch2, + float* scratch_v, float* scratch_a, + float* q_values, + float sigma_init, + unsigned int* rng, + float* shmem_weights, + float* shmem_bias +) { + /* Shared layers with noise */ + TILE_LAYER_NOISY(w_s1, b_s1, state, scratch1, STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, sigma_init, rng); + TILE_LAYER_NOISY(w_s2, b_s2, scratch1, scratch2, SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, sigma_init, rng); + + /* Value head with noise */ + TILE_LAYER_NOISY(w_v1, b_v1, scratch2, scratch_v, SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, sigma_init, rng); + /* Value output (1 neuron) — noisy, loaded via shmem */ + float value; + { + cooperative_load_tile(shmem_weights, w_v2, VALUE_H); + cooperative_load_tile(shmem_bias, b_v2, 1); + __syncthreads(); + float sigma_scale = sigma_init / sqrtf((float)VALUE_H); + float eps_out_v = factorized_noise_fn(gpu_random_gaussian(rng)); + float acc = shmem_bias[0] + sigma_scale * eps_out_v; + for (int i = 0; i < VALUE_H; i++) { + float eps_in_v = factorized_noise_fn(gpu_random_gaussian(rng)); + acc += (shmem_weights[i] + sigma_scale * eps_out_v * eps_in_v) * scratch_v[i]; + } + value = acc; + __syncthreads(); + } + + /* Advantage head with noise */ + TILE_LAYER_NOISY(w_a1, b_a1, scratch2, scratch_a, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, sigma_init, rng); + float adv[NUM_ACTIONS]; + TILE_LAYER_NOISY(w_a2, b_a2, scratch_a, adv, ADV_H, NUM_ACTIONS, 0, shmem_weights, shmem_bias, sigma_init, rng); + + /* Mean advantage */ + float adv_mean = 0.0f; + for (int i = 0; i < NUM_ACTIONS; i++) adv_mean += adv[i]; + adv_mean /= (float)NUM_ACTIONS; + + /* Q(s,a) = V(s) + A(s,a) - mean(A) */ + for (int i = 0; i < NUM_ACTIONS; i++) { + q_values[i] = value + adv[i] - adv_mean; + } +} + +/** + * Shared-memory C51 distributional dueling Q-network forward pass. + * + * Same tiling as q_forward_dueling_shmem with distributional atom outputs. + * Optionally applies NoisyNet noise and RMSNorm. + */ +__device__ void q_forward_distributional_shmem( + const float* state, + const float* __restrict__ w_s1, const float* __restrict__ b_s1, + const float* __restrict__ w_s2, const float* __restrict__ b_s2, + const float* __restrict__ w_v1, const float* __restrict__ b_v1, + const float* __restrict__ w_v2, const float* __restrict__ b_v2, + const float* __restrict__ w_a1, const float* __restrict__ b_a1, + const float* __restrict__ w_a2, const float* __restrict__ b_a2, + const float* __restrict__ rms_s0, const float* __restrict__ rms_s1, + const float* __restrict__ rms_v, const float* __restrict__ rms_a, + float* scratch1, float* scratch2, + float* scratch_v, float* scratch_a, + float* q_values, + int num_atoms, float v_min, float v_max, + int use_noisy, float sigma_init, int use_rmsnorm, + unsigned int* rng, + float* shmem_weights, + float* shmem_bias +) { + int na = (num_atoms > NUM_ATOMS_MAX) ? NUM_ATOMS_MAX : num_atoms; + + /* ---- Shared layers ---- */ + if (use_noisy) { + TILE_LAYER_NOISY(w_s1, b_s1, state, scratch1, STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, sigma_init, rng); + } else { + TILE_LAYER_CLEAN(w_s1, b_s1, state, scratch1, STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias); + } + if (use_rmsnorm) rmsnorm_inplace(scratch1, rms_s0, SHARED_H1); + + if (use_noisy) { + TILE_LAYER_NOISY(w_s2, b_s2, scratch1, scratch2, SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, sigma_init, rng); + } else { + TILE_LAYER_CLEAN(w_s2, b_s2, scratch1, scratch2, SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias); + } + if (use_rmsnorm) rmsnorm_inplace(scratch2, rms_s1, SHARED_H2); + + /* ---- Value head → [num_atoms] logits ---- */ + if (use_noisy) { + TILE_LAYER_NOISY(w_v1, b_v1, scratch2, scratch_v, SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, sigma_init, rng); + } else { + TILE_LAYER_CLEAN(w_v1, b_v1, scratch2, scratch_v, SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias); + } + if (use_rmsnorm) rmsnorm_inplace(scratch_v, rms_v, VALUE_H); + + float val_atoms[NUM_ATOMS_MAX]; + if (use_noisy) { + TILE_LAYER_NOISY(w_v2, b_v2, scratch_v, val_atoms, VALUE_H, na, 0, shmem_weights, shmem_bias, sigma_init, rng); + } else { + TILE_LAYER_CLEAN(w_v2, b_v2, scratch_v, val_atoms, VALUE_H, na, 0, shmem_weights, shmem_bias); + } + + /* ---- Advantage head → [NUM_ACTIONS * num_atoms] logits ---- */ + if (use_noisy) { + TILE_LAYER_NOISY(w_a1, b_a1, scratch2, scratch_a, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, sigma_init, rng); + } else { + TILE_LAYER_CLEAN(w_a1, b_a1, scratch2, scratch_a, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias); + } + if (use_rmsnorm) rmsnorm_inplace(scratch_a, rms_a, ADV_H); + + int adv_total = NUM_ACTIONS * na; + float adv_atoms[NUM_ACTIONS * NUM_ATOMS_MAX]; + if (use_noisy) { + TILE_LAYER_NOISY(w_a2, b_a2, scratch_a, adv_atoms, ADV_H, adv_total, 0, shmem_weights, shmem_bias, sigma_init, rng); + } else { + TILE_LAYER_CLEAN(w_a2, b_a2, scratch_a, adv_atoms, ADV_H, adv_total, 0, shmem_weights, shmem_bias); + } + + /* ---- Distributional dueling combination + softmax → expected Q ---- */ + float delta_z = (na > 1) ? (v_max - v_min) / (float)(na - 1) : 0.0f; + + for (int a = 0; a < NUM_ACTIONS; a++) { + float logits[NUM_ATOMS_MAX]; + for (int i = 0; i < na; i++) { + float adv_mean_i = 0.0f; + for (int ap = 0; ap < NUM_ACTIONS; ap++) { + adv_mean_i += adv_atoms[ap * na + i]; + } + adv_mean_i /= (float)NUM_ACTIONS; + logits[i] = val_atoms[i] + adv_atoms[a * na + i] - adv_mean_i; + } + + float max_logit = logits[0]; + for (int i = 1; i < na; i++) { + if (logits[i] > max_logit) max_logit = logits[i]; + } + float sum_exp = 0.0f; + for (int i = 0; i < na; i++) { + logits[i] = expf(logits[i] - max_logit); + sum_exp += logits[i]; + } + float inv_sum = 1.0f / (sum_exp + 1e-8f); + + float q_expected = 0.0f; + for (int i = 0; i < na; i++) { + float z_i = v_min + (float)i * delta_z; + float p_i = logits[i] * inv_sum; + q_expected += z_i * p_i; + } + q_values[a] = q_expected; + } +} + /* ------------------------------------------------------------------ */ /* D5: NoisyNet Forward Pass */ /* ------------------------------------------------------------------ */ @@ -313,7 +592,16 @@ __device__ __forceinline__ int argmax_q(const float* q_values) { * Full DQN experience collection kernel. * * Each thread runs one independent episode of L timesteps. - * Grid: (ceil(N/32), 1, 1), Block: (32, 1, 1). + * Grid: (ceil(N/block), 1, 1), Block: (256, 1, 1). + * + * Uses shared memory weight tiling: all threads in the block cooperatively + * load weight tiles from global memory (HBM) to shared memory (SRAM), + * then each thread reads from fast SRAM for its forward pass. This + * eliminates redundant HBM reads across threads and timesteps. + * + * Required: extern __shared__ float shmem[] of at least + * (SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM + SHMEM_TILE_ROWS) * sizeof(float) + * = 65,792 bytes at default tile sizes. */ extern "C" __global__ void dqn_full_experience_kernel( /* Market data [total_bars, MARKET_DIM] */ @@ -415,37 +703,44 @@ extern "C" __global__ void dqn_full_experience_kernel( float* out_target_q, /* [N, L] */ float* out_td_error /* [N, L] */ ) { + /* ---- Shared memory workspace for weight tiling ---- */ + extern __shared__ float shmem[]; + float* shmem_weights = shmem; + float* shmem_bias = shmem + SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM; + int tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= N) return; + /* All threads must participate in cooperative_load_tile / __syncthreads + * even if they don't own an episode. Use tid_valid to gate data work. */ + int tid_valid = (tid < N) ? 1 : 0; - /* ---- Load per-thread state ---- */ - int ps_off = tid * PORTFOLIO_STATE_SIZE; - float cash = portfolio_states[ps_off + 0]; - float position = portfolio_states[ps_off + 1]; - float entry_price = portfolio_states[ps_off + 2]; - float initial_cap = portfolio_states[ps_off + 3]; - float spread = portfolio_states[ps_off + 4]; - float last_price = portfolio_states[ps_off + 5]; - float reserve_pct = portfolio_states[ps_off + 6]; - float cum_costs = portfolio_states[ps_off + 7]; + /* ---- Load per-thread state (only for valid threads) ---- */ + int ps_off = tid_valid ? (tid * PORTFOLIO_STATE_SIZE) : 0; + float cash = tid_valid ? portfolio_states[ps_off + 0] : 0.0f; + float position = tid_valid ? portfolio_states[ps_off + 1] : 0.0f; + float entry_price = tid_valid ? portfolio_states[ps_off + 2] : 0.0f; + float initial_cap = tid_valid ? portfolio_states[ps_off + 3] : 1.0f; + float spread = tid_valid ? portfolio_states[ps_off + 4] : 0.0f; + float last_price = tid_valid ? portfolio_states[ps_off + 5] : 0.0f; + float reserve_pct = tid_valid ? portfolio_states[ps_off + 6] : 0.0f; + float cum_costs = tid_valid ? portfolio_states[ps_off + 7] : 0.0f; - int bs_off = tid * BARRIER_STATE_SIZE; + int bs_off = tid_valid ? (tid * BARRIER_STATE_SIZE) : 0; float barrier_st[BARRIER_STATE_SIZE]; for (int i = 0; i < BARRIER_STATE_SIZE; i++) - barrier_st[i] = barrier_states[bs_off + i]; + barrier_st[i] = tid_valid ? barrier_states[bs_off + i] : 0.0f; - int dw_off = tid * DIVERSITY_WINDOW; + int dw_off = tid_valid ? (tid * DIVERSITY_WINDOW) : 0; int div_window[DIVERSITY_WINDOW]; for (int i = 0; i < DIVERSITY_WINDOW; i++) - div_window[i] = diversity_windows[dw_off + i]; + div_window[i] = tid_valid ? diversity_windows[dw_off + i] : 0; - int dm_off = tid * 2; + int dm_off = tid_valid ? (tid * 2) : 0; int div_meta[2]; - div_meta[0] = diversity_metas[dm_off + 0]; - div_meta[1] = diversity_metas[dm_off + 1]; + div_meta[0] = tid_valid ? diversity_metas[dm_off + 0] : 0; + div_meta[1] = tid_valid ? diversity_metas[dm_off + 1] : 0; - unsigned int rng = rng_states[tid]; - int ep_start = episode_starts[tid]; + unsigned int rng = tid_valid ? rng_states[tid] : 0u; + int ep_start = tid_valid ? episode_starts[tid] : 0; /* Scratch buffers for Q-network forward pass (thread-local) */ float scratch1[SHARED_H1]; @@ -474,8 +769,8 @@ extern "C" __global__ void dqn_full_experience_kernel( int global_bar = ep_start + t; int out_off = tid * L + t; - /* Handle out-of-data */ - if (global_bar >= total_bars - 1) { + /* Handle out-of-data (only valid threads write) */ + if (tid_valid && global_bar >= total_bars - 1) { for (int i = 0; i < STATE_DIM; i++) out_states[out_off * STATE_DIM + i] = 0.0f; out_actions[out_off] = 0; @@ -483,47 +778,56 @@ extern "C" __global__ void dqn_full_experience_kernel( out_done[out_off] = 1; out_target_q[out_off] = 0.0f; out_td_error[out_off] = 0.0f; - continue; + /* Don't continue — must still participate in shmem syncs below. + * Use skip flag to bypass data processing. */ } + int skip_data = (!tid_valid || global_bar >= total_bars - 1) ? 1 : 0; /* ---- Step 1: Read market features ---- */ - int mf_off = global_bar * MARKET_DIM; - for (int i = 0; i < MARKET_DIM; i++) - state[i] = market_features[mf_off + i]; + if (!skip_data) { + int mf_off = global_bar * MARKET_DIM; + for (int i = 0; i < MARKET_DIM; i++) + state[i] = market_features[mf_off + i]; + } /* ---- Step 2: Compute portfolio features ---- */ - int t_off = global_bar * 4; - float current_close = targets[t_off + 0]; - float next_close = targets[t_off + 1]; - float current_close_raw = targets[t_off + 2]; - float next_close_raw = targets[t_off + 3]; + float current_close = 0.0f, next_close = 0.0f; + float current_close_raw = 0.0f, next_close_raw = 0.0f; + float price = 1.0f, current_value = 0.0f, current_norm = 0.0f; + float max_pos_norm = 1.0f, pos_norm = 0.0f; + if (!skip_data) { + int t_off = global_bar * 4; + current_close = targets[t_off + 0]; + next_close = targets[t_off + 1]; + current_close_raw = targets[t_off + 2]; + next_close_raw = targets[t_off + 3]; - float price = (current_close_raw != 0.0f) ? current_close_raw : current_close; - if (price <= 0.0f) price = 1.0f; + price = (current_close_raw != 0.0f) ? current_close_raw : current_close; + if (price <= 0.0f) price = 1.0f; - float current_value = cash + position * price; - float current_norm = current_value / initial_cap; - float max_pos_norm = (price > 0.0f) ? initial_cap / price : 1.0f; - float pos_norm = position / max_pos_norm; + current_value = cash + position * price; + current_norm = current_value / initial_cap; + max_pos_norm = (price > 0.0f) ? initial_cap / price : 1.0f; + pos_norm = position / max_pos_norm; - state[MARKET_DIM + 0] = current_norm; - state[MARKET_DIM + 1] = pos_norm; - state[MARKET_DIM + 2] = spread; + state[MARKET_DIM + 0] = current_norm; + state[MARKET_DIM + 1] = pos_norm; + state[MARKET_DIM + 2] = spread; - /* Zero-pad remaining dims (OFI features + tensor-core alignment padding). - * These dims exist in the network's input but aren't populated by the - * GPU kernel — they're zero, matching training when OFI is unavailable. */ - for (int i = MARKET_DIM + PORTFOLIO_DIM; i < STATE_DIM; i++) - state[i] = 0.0f; + /* Zero-pad remaining dims (OFI features + tensor-core alignment padding). */ + for (int i = MARKET_DIM + PORTFOLIO_DIM; i < STATE_DIM; i++) + state[i] = 0.0f; - /* ---- Step 3: Write state to output ---- */ - for (int i = 0; i < STATE_DIM; i++) - out_states[out_off * STATE_DIM + i] = state[i]; + /* ---- Step 3: Write state to output ---- */ + for (int i = 0; i < STATE_DIM; i++) + out_states[out_off * STATE_DIM + i] = state[i]; + } - /* ---- Step 4: Online Q-network forward ---- */ + /* ---- Step 4: Online Q-network forward (shmem tiling) ---- + * All threads participate in cooperative loads + __syncthreads(). + * Invalid/skipped threads still compute into scratch (harmlessly). */ if (use_distributional && num_atoms > 1) { - /* D6: C51 distributional dueling forward */ - q_forward_distributional( + q_forward_distributional_shmem( state, on_w_s1, on_b_s1, on_w_s2, on_b_s2, on_w_v1, on_b_v1, on_w_v2, on_b_v2, @@ -533,185 +837,191 @@ extern "C" __global__ void dqn_full_experience_kernel( q_values, num_atoms, v_min, v_max, use_noisy_nets, noisy_sigma_init, use_rmsnorm, - &rng + &rng, + shmem_weights, shmem_bias ); } else if (use_noisy_nets) { - /* D5: NoisyNet dueling forward (online only) */ - q_forward_dueling_noisy( + q_forward_dueling_noisy_shmem( state, on_w_s1, on_b_s1, on_w_s2, on_b_s2, on_w_v1, on_b_v1, on_w_v2, on_b_v2, on_w_a1, on_b_a1, on_w_a2, on_b_a2, scratch1, scratch2, scratch_v, scratch_a, q_values, - noisy_sigma_init, &rng + noisy_sigma_init, &rng, + shmem_weights, shmem_bias ); } else { - /* Standard dueling forward */ - q_forward_dueling( + q_forward_dueling_shmem( state, on_w_s1, on_b_s1, on_w_s2, on_b_s2, on_w_v1, on_b_v1, on_w_v2, on_b_v2, on_w_a1, on_b_a1, on_w_a2, on_b_a2, scratch1, scratch2, scratch_v, scratch_a, - q_values + q_values, + shmem_weights, shmem_bias ); } - /* ---- Step 4b: Q-value clipping (D2) ---- */ - for (int i = 0; i < NUM_ACTIONS; i++) { - if (q_values[i] < q_clip_min) q_values[i] = q_clip_min; - if (q_values[i] > q_clip_max) q_values[i] = q_clip_max; - } + /* Steps 4b through 10 only executed by valid threads with data */ + int action_idx = 0; + float online_q_selected = 0.0f; + float curiosity_reward = 0.0f; + float next_price = 1.0f; + float next_value = 0.0f; + float next_norm = 0.0f; + int barrier_label = 0; + float div_penalty = 0.0f; + float next_state[STATE_DIM]; - /* ---- Step 5: Epsilon-greedy + count-bonus action selection ---- */ - int action_idx; - float r = gpu_random(&rng); - if (r < epsilon) { - /* Random action [0, 4] (5 DQN exposure actions) */ - action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS); - if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1; - } else { - /* D1: Add UCB count bonus before greedy selection. - * bonus_a = β × sqrt(ln(N) / (1 + n_a)) - * This gives a direct exploration incentive to under-visited actions. */ - if (count_bonus_coefficient > 0.0f && total_action_count > 0) { - float log_n = logf((float)total_action_count); - float q_bonus[NUM_ACTIONS]; - for (int i = 0; i < NUM_ACTIONS; i++) { - float bonus = count_bonus_coefficient - * sqrtf(log_n / (1.0f + (float)action_counts[i])); - q_bonus[i] = q_values[i] + bonus; - } - action_idx = argmax_q(q_bonus); - } else { - action_idx = argmax_q(q_values); + if (!skip_data) { + /* ---- Step 4b: Q-value clipping (D2) ---- */ + for (int i = 0; i < NUM_ACTIONS; i++) { + if (q_values[i] < q_clip_min) q_values[i] = q_clip_min; + if (q_values[i] > q_clip_max) q_values[i] = q_clip_max; } - } - /* D1: Update count state */ - action_counts[action_idx]++; - total_action_count++; - - out_actions[out_off] = action_idx; - - float online_q_selected = q_values[action_idx]; - - /* ---- Step 6: Portfolio simulation ---- */ - float target_exposure = action_to_exposure(action_idx); - float target_position = target_exposure * max_position; - float tx_rate = action_to_tx_cost(action_idx) * tx_cost_multiplier; - - /* Detect reversal (sign change) */ - int is_reversal = (position > 0.0f && target_position < 0.0f) || - (position < 0.0f && target_position > 0.0f); - - if (is_reversal) { - /* Phase 1: Close current position */ - float close_cash = position * price; - float close_cost = fabsf(position) * price * tx_rate; - cash += close_cash - close_cost; - cum_costs += close_cost; - - /* Phase 2: Open opposite position */ - float reserve = (reserve_pct > 0.0f) ? current_value * (reserve_pct / 100.0f) : 0.0f; - float affordable = fmaxf(cash - reserve, 0.0f); - float max_contracts = (price > 0.0f) ? affordable / (price * (1.0f + tx_rate)) : 0.0f; - max_contracts = floorf(max_contracts); - float actual = fminf(max_contracts, fabsf(target_position)); - - if (actual > 0.0f) { - float new_pos = (target_position > 0.0f) ? actual : -actual; - float open_cost = actual * price * tx_rate; - cash -= new_pos * price + open_cost; - cum_costs += open_cost; - position = new_pos; - entry_price = price; + /* ---- Step 5: Epsilon-greedy + count-bonus action selection ---- */ + float r = gpu_random(&rng); + if (r < epsilon) { + /* Random action [0, 4] (5 DQN exposure actions) */ + action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS); + if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1; } else { - position = 0.0f; - entry_price = 0.0f; - } - } else { - /* Non-reversal: adjust position directly */ - float delta = target_position - position; - if (fabsf(delta) > 0.0f) { - float trade_cost = fabsf(delta) * price * tx_rate; - cum_costs += trade_cost; - cash -= trade_cost; - - /* Cash reserve check for buys */ - if (delta > 0.0f && reserve_pct > 0.0f) { - float pv = cash + position * price; - float reserve = pv * (reserve_pct / 100.0f); - float buy_cost = delta * price; - if (cash - buy_cost < reserve) { - float affordable = fmaxf(cash - reserve, 0.0f); - delta = fminf(delta, (price > 0.0f) ? floorf(affordable / price) : 0.0f); + /* D1: Add UCB count bonus before greedy selection. */ + if (count_bonus_coefficient > 0.0f && total_action_count > 0) { + float log_n = logf((float)total_action_count); + float q_bonus[NUM_ACTIONS]; + for (int i = 0; i < NUM_ACTIONS; i++) { + float bonus = count_bonus_coefficient + * sqrtf(log_n / (1.0f + (float)action_counts[i])); + q_bonus[i] = q_values[i] + bonus; } + action_idx = argmax_q(q_bonus); + } else { + action_idx = argmax_q(q_values); } + } - if (delta > 0.0f) { + /* D1: Update count state */ + action_counts[action_idx]++; + total_action_count++; + + out_actions[out_off] = action_idx; + online_q_selected = q_values[action_idx]; + + /* ---- Step 6: Portfolio simulation ---- */ + float target_exposure = action_to_exposure(action_idx); + float target_position = target_exposure * max_position; + float tx_rate = action_to_tx_cost(action_idx) * tx_cost_multiplier; + + /* Detect reversal (sign change) */ + int is_reversal = (position > 0.0f && target_position < 0.0f) || + (position < 0.0f && target_position > 0.0f); + + if (is_reversal) { + /* Phase 1: Close current position */ + float close_cash = position * price; + float close_cost = fabsf(position) * price * tx_rate; + cash += close_cash - close_cost; + cum_costs += close_cost; + + /* Phase 2: Open opposite position */ + float reserve = (reserve_pct > 0.0f) ? current_value * (reserve_pct / 100.0f) : 0.0f; + float affordable = fmaxf(cash - reserve, 0.0f); + float max_contracts = (price > 0.0f) ? affordable / (price * (1.0f + tx_rate)) : 0.0f; + max_contracts = floorf(max_contracts); + float actual = fminf(max_contracts, fabsf(target_position)); + + if (actual > 0.0f) { + float new_pos = (target_position > 0.0f) ? actual : -actual; + float open_cost = actual * price * tx_rate; + cash -= new_pos * price + open_cost; + cum_costs += open_cost; + position = new_pos; entry_price = price; - } else if (target_position == 0.0f) { + } else { + position = 0.0f; entry_price = 0.0f; } - cash -= delta * price; - position = position + delta; + } else { + /* Non-reversal: adjust position directly */ + float delta = target_position - position; + if (fabsf(delta) > 0.0f) { + float trade_cost = fabsf(delta) * price * tx_rate; + cum_costs += trade_cost; + cash -= trade_cost; + + /* Cash reserve check for buys */ + if (delta > 0.0f && reserve_pct > 0.0f) { + float pv = cash + position * price; + float reserve = pv * (reserve_pct / 100.0f); + float buy_cost = delta * price; + if (cash - buy_cost < reserve) { + float affordable = fmaxf(cash - reserve, 0.0f); + delta = fminf(delta, (price > 0.0f) ? floorf(affordable / price) : 0.0f); + } + } + + if (delta > 0.0f) { + entry_price = price; + } else if (target_position == 0.0f) { + entry_price = 0.0f; + } + cash -= delta * price; + position = position + delta; + } } - } - last_price = price; + last_price = price; - /* ---- Step 7: Barrier tracking ---- */ - /* Init barrier on position open (entry_price just set) */ - float old_barrier_entry = barrier_st[0]; - if (entry_price > 0.0f && old_barrier_entry <= 0.0f) { - barrier_init(barrier_st, barrier_config, entry_price, global_bar); - } + /* ---- Step 7: Barrier tracking ---- */ + float old_barrier_entry = barrier_st[0]; + if (entry_price > 0.0f && old_barrier_entry <= 0.0f) { + barrier_init(barrier_st, barrier_config, entry_price, global_bar); + } + barrier_label = barrier_check(barrier_st, price, global_bar, position); + if (barrier_label != 0) { + barrier_reset(barrier_st); + } - int barrier_label = barrier_check(barrier_st, price, global_bar, position); - if (barrier_label != 0) { - barrier_reset(barrier_st); - } + /* ---- Step 8: Diversity penalty ---- */ + div_penalty = diversity_entropy(div_window, div_meta, action_idx); - /* ---- Step 8: Diversity penalty ---- */ - float div_penalty = diversity_entropy(div_window, div_meta, action_idx); + /* ---- Step 9: Mark-to-market -> next_state ---- */ + next_price = (next_close_raw != 0.0f) ? next_close_raw : next_close; + if (next_price <= 0.0f) next_price = price; + next_value = cash + position * next_price; + next_norm = next_value / initial_cap; - /* ---- Step 9: Mark-to-market -> next_state ---- */ - float next_price = (next_close_raw != 0.0f) ? next_close_raw : next_close; - if (next_price <= 0.0f) next_price = price; - float next_value = cash + position * next_price; - float next_norm = next_value / initial_cap; + /* Build next_state for curiosity */ + int next_bar = global_bar + 1; + if (next_bar < total_bars) { + int nmf_off = next_bar * MARKET_DIM; + for (int i = 0; i < MARKET_DIM; i++) + next_state[i] = market_features[nmf_off + i]; + } else { + for (int i = 0; i < MARKET_DIM; i++) + next_state[i] = state[i]; + } + float next_max_pos_norm = (next_price > 0.0f) ? initial_cap / next_price : 1.0f; + next_state[MARKET_DIM + 0] = next_norm; + next_state[MARKET_DIM + 1] = position / next_max_pos_norm; + next_state[MARKET_DIM + 2] = spread; - /* Build next_state for curiosity (reuse state array) */ - float next_state[STATE_DIM]; - int next_bar = global_bar + 1; - if (next_bar < total_bars) { - int nmf_off = next_bar * MARKET_DIM; - for (int i = 0; i < MARKET_DIM; i++) - next_state[i] = market_features[nmf_off + i]; - } else { - for (int i = 0; i < MARKET_DIM; i++) - next_state[i] = state[i]; - } - float next_max_pos_norm = (next_price > 0.0f) ? initial_cap / next_price : 1.0f; - next_state[MARKET_DIM + 0] = next_norm; - next_state[MARKET_DIM + 1] = position / next_max_pos_norm; - next_state[MARKET_DIM + 2] = spread; + /* ---- Step 10: Curiosity inference ---- */ + if (cur_w1 != 0) { + curiosity_reward = curiosity_inference( + state, next_state, action_idx, + cur_w1, cur_b1, cur_w2, cur_b2, + cur_scratch, curiosity_max_reward + ); + } + } /* end if (!skip_data) */ - /* ---- Step 10: Curiosity inference ---- */ - float curiosity_reward = 0.0f; - if (cur_w1 != 0) { - curiosity_reward = curiosity_inference( - state, next_state, action_idx, - cur_w1, cur_b1, cur_w2, cur_b2, - cur_scratch, curiosity_max_reward - ); - } - - /* ---- Step 11: Target Q-network forward (always deterministic — no noise) ---- */ + /* ---- Step 11: Target Q-network forward (shmem tiling) ---- + * All threads participate in cooperative loads, same as online forward. */ if (use_distributional && num_atoms > 1) { - /* D6: C51 distributional target — clean (no noise) */ - q_forward_distributional( + q_forward_distributional_shmem( next_state, tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2, tg_w_v1, tg_b_v1, tg_w_v2, tg_b_v2, @@ -721,17 +1031,18 @@ extern "C" __global__ void dqn_full_experience_kernel( tgt_q_values, num_atoms, v_min, v_max, 0, 0.0f, use_rmsnorm, /* no noise on target */ - &rng + &rng, + shmem_weights, shmem_bias ); } else { - /* Standard dueling target — always clean */ - q_forward_dueling( + q_forward_dueling_shmem( next_state, tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2, tg_w_v1, tg_b_v1, tg_w_v2, tg_b_v2, tg_w_a1, tg_b_a1, tg_w_a2, tg_b_a2, scratch1, scratch2, scratch_v, scratch_a, - tgt_q_values + tgt_q_values, + shmem_weights, shmem_bias ); } /* D2: Clip target Q-values too */ @@ -745,132 +1056,137 @@ extern "C" __global__ void dqn_full_experience_kernel( max_target_q = tgt_q_values[i]; } - /* ---- Step 12: Combined reward ---- */ - float pnl_reward = 0.0f; - if (current_norm > 0.0f) { - pnl_reward = (next_norm - current_norm) / current_norm; - } - - /* Flat reward: when position is ~0, give a small positive signal - * so Flat isn't structurally disadvantaged (always 0 PnL otherwise). - * This breaks the zero-reward trap that causes Flat action collapse. */ - if (fabsf(position) < 0.001f) { - pnl_reward += hold_reward; - } - - /* Risk penalty: penalize excessive position (>80% exposure) */ - float abs_pos = fabsf(pos_norm); - if (abs_pos > 0.8f) { - pnl_reward -= (abs_pos - 0.8f) * 5.0f * risk_weight; - } - - /* Barrier scaling: amplify reward on barrier hit */ - float barrier_mult = 1.0f; - if (barrier_label != 0) { - barrier_mult = 1.0f + barrier_scale * (float)barrier_label; - } - - float combined_reward = pnl_reward * barrier_mult - + diversity_scale * div_penalty - + curiosity_scale * curiosity_reward; - - /* ---- EMA reward normalization (matches CPU RewardNormalizer) ---- - * Normalize-then-update: use OLD stats to normalize, then update - * stats with the raw value. Clamp to [-3, 3] like the CPU path. */ - float raw_reward = combined_reward; - if (ema_init) { - float std = sqrtf(ema_var); - if (std > 1e-8f) - combined_reward = (combined_reward - ema_mean) / std; - combined_reward = fmaxf(-3.0f, fminf(3.0f, combined_reward)); - } - /* Update EMA running stats with the raw (un-normalized) value */ - if (!ema_init) { - ema_mean = raw_reward; - ema_var = 1.0f; - ema_init = 1; - } else { - ema_mean = reward_norm_alpha * raw_reward - + (1.0f - reward_norm_alpha) * ema_mean; - float diff = raw_reward - ema_mean; - ema_var = reward_norm_alpha * diff * diff - + (1.0f - reward_norm_alpha) * ema_var; - } - - /* ---- Step 13: Episode done check ---- */ - step_in_episode++; - int time_done = (step_in_episode >= episode_length) ? 1 : 0; - int barrier_done = (barrier_label != 0) ? 1 : 0; - int data_done = (next_bar >= total_bars) ? 1 : 0; - int done = (time_done || barrier_done || data_done) ? 1 : 0; - - /* ---- Step 14: TD error with Huber transformation (D3) ---- */ - float td_target = combined_reward + gamma * max_target_q * (1.0f - (float)done); - float td_delta = td_target - online_q_selected; - float td_abs = fabsf(td_delta); - /* Huber loss: robust to outlier TD errors. - * if |δ| ≤ κ: loss = 0.5 * δ² - * if |δ| > κ: loss = κ * (|δ| - 0.5 * κ) - * When κ ≤ 0 (disabled): fall back to raw |δ| (L1). */ - float td_error; - if (huber_kappa > 0.0f) { - if (td_abs <= huber_kappa) { - td_error = 0.5f * td_delta * td_delta; - } else { - td_error = huber_kappa * (td_abs - 0.5f * huber_kappa); + /* Steps 12-16 only for valid threads with data */ + if (!skip_data) { + /* ---- Step 12: Combined reward ---- */ + float pnl_reward = 0.0f; + if (current_norm > 0.0f) { + pnl_reward = (next_norm - current_norm) / current_norm; } - } else { - td_error = td_abs; - } - /* ---- Step 15: Write experience ---- */ - out_rewards[out_off] = combined_reward; - out_done[out_off] = done; - out_target_q[out_off] = max_target_q; - out_td_error[out_off] = td_error; + /* Flat reward: when position is ~0, give a small positive signal + * so Flat isn't structurally disadvantaged (always 0 PnL otherwise). */ + if (fabsf(position) < 0.001f) { + pnl_reward += hold_reward; + } - /* ---- Step 16: Episode reset if done ---- */ - if (done) { - cash = initial_cap; - position = 0.0f; - entry_price = 0.0f; - cum_costs = 0.0f; - last_price = 0.0f; - step_in_episode = 0; - barrier_reset(barrier_st); - /* Reset diversity window */ - for (int i = 0; i < DIVERSITY_WINDOW; i++) - div_window[i] = 0; - div_meta[0] = 0; - div_meta[1] = 0; - /* D1: Reset count-bonus state */ - for (int i = 0; i < NUM_ACTIONS; i++) action_counts[i] = 0; - total_action_count = 0; - /* Reset EMA normalizer for new episode */ - ema_mean = 0.0f; - ema_var = 1.0f; - ema_init = 0; - } + /* Risk penalty: penalize excessive position (>80% exposure) */ + float abs_pos = fabsf(pos_norm); + if (abs_pos > 0.8f) { + pnl_reward -= (abs_pos - 0.8f) * 5.0f * risk_weight; + } + + /* Barrier scaling: amplify reward on barrier hit */ + float barrier_mult = 1.0f; + if (barrier_label != 0) { + barrier_mult = 1.0f + barrier_scale * (float)barrier_label; + } + + float combined_reward = pnl_reward * barrier_mult + + diversity_scale * div_penalty + + curiosity_scale * curiosity_reward; + + /* ---- EMA reward normalization (matches CPU RewardNormalizer) ---- + * Normalize-then-update: use OLD stats to normalize, then update + * stats with the raw value. Clamp to [-3, 3] like the CPU path. */ + float raw_reward = combined_reward; + if (ema_init) { + float std = sqrtf(ema_var); + if (std > 1e-8f) + combined_reward = (combined_reward - ema_mean) / std; + combined_reward = fmaxf(-3.0f, fminf(3.0f, combined_reward)); + } + /* Update EMA running stats with the raw (un-normalized) value */ + if (!ema_init) { + ema_mean = raw_reward; + ema_var = 1.0f; + ema_init = 1; + } else { + ema_mean = reward_norm_alpha * raw_reward + + (1.0f - reward_norm_alpha) * ema_mean; + float diff = raw_reward - ema_mean; + ema_var = reward_norm_alpha * diff * diff + + (1.0f - reward_norm_alpha) * ema_var; + } + + /* ---- Step 13: Episode done check ---- */ + step_in_episode++; + int next_bar = global_bar + 1; + int time_done = (step_in_episode >= episode_length) ? 1 : 0; + int barrier_done = (barrier_label != 0) ? 1 : 0; + int data_done = (next_bar >= total_bars) ? 1 : 0; + int done = (time_done || barrier_done || data_done) ? 1 : 0; + + /* ---- Step 14: TD error with Huber transformation (D3) ---- */ + float td_target = combined_reward + gamma * max_target_q * (1.0f - (float)done); + float td_delta = td_target - online_q_selected; + float td_abs = fabsf(td_delta); + /* Huber loss: robust to outlier TD errors. + * if |delta| <= kappa: loss = 0.5 * delta^2 + * if |delta| > kappa: loss = kappa * (|delta| - 0.5 * kappa) + * When kappa <= 0 (disabled): fall back to raw |delta| (L1). */ + float td_error; + if (huber_kappa > 0.0f) { + if (td_abs <= huber_kappa) { + td_error = 0.5f * td_delta * td_delta; + } else { + td_error = huber_kappa * (td_abs - 0.5f * huber_kappa); + } + } else { + td_error = td_abs; + } + + /* ---- Step 15: Write experience ---- */ + out_rewards[out_off] = combined_reward; + out_done[out_off] = done; + out_target_q[out_off] = max_target_q; + out_td_error[out_off] = td_error; + + /* ---- Step 16: Episode reset if done ---- */ + if (done) { + cash = initial_cap; + position = 0.0f; + entry_price = 0.0f; + cum_costs = 0.0f; + last_price = 0.0f; + step_in_episode = 0; + barrier_reset(barrier_st); + /* Reset diversity window */ + for (int i = 0; i < DIVERSITY_WINDOW; i++) + div_window[i] = 0; + div_meta[0] = 0; + div_meta[1] = 0; + /* D1: Reset count-bonus state */ + for (int i = 0; i < NUM_ACTIONS; i++) action_counts[i] = 0; + total_action_count = 0; + /* Reset EMA normalizer for new episode */ + ema_mean = 0.0f; + ema_var = 1.0f; + ema_init = 0; + } + } /* end if (!skip_data) */ } /* end timestep loop */ - /* ---- Write back per-thread state ---- */ - portfolio_states[ps_off + 0] = cash; - portfolio_states[ps_off + 1] = position; - portfolio_states[ps_off + 2] = entry_price; - portfolio_states[ps_off + 3] = initial_cap; - portfolio_states[ps_off + 4] = spread; - portfolio_states[ps_off + 5] = last_price; - portfolio_states[ps_off + 6] = reserve_pct; - portfolio_states[ps_off + 7] = cum_costs; + /* ---- Write back per-thread state (only valid threads) ---- */ + if (tid_valid) { + portfolio_states[ps_off + 0] = cash; + portfolio_states[ps_off + 1] = position; + portfolio_states[ps_off + 2] = entry_price; + portfolio_states[ps_off + 3] = initial_cap; + portfolio_states[ps_off + 4] = spread; + portfolio_states[ps_off + 5] = last_price; + portfolio_states[ps_off + 6] = reserve_pct; + portfolio_states[ps_off + 7] = cum_costs; - for (int i = 0; i < BARRIER_STATE_SIZE; i++) - barrier_states[bs_off + i] = barrier_st[i]; + for (int i = 0; i < BARRIER_STATE_SIZE; i++) + barrier_states[bs_off + i] = barrier_st[i]; - for (int i = 0; i < DIVERSITY_WINDOW; i++) - diversity_windows[dw_off + i] = div_window[i]; + for (int i = 0; i < DIVERSITY_WINDOW; i++) + diversity_windows[dw_off + i] = div_window[i]; - diversity_metas[dm_off + 0] = div_meta[0]; - diversity_metas[dm_off + 1] = div_meta[1]; + diversity_metas[dm_off + 0] = div_meta[0]; + diversity_metas[dm_off + 1] = div_meta[1]; - rng_states[tid] = rng; + rng_states[tid] = rng; + } } diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 04e5c3d01..91b30f7c4 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -542,14 +542,23 @@ impl GpuExperienceCollector { })?; // ---- Step 4: Launch config ---- + // Shared memory for weight tiling: tile_rows * max_in_dim floats (weights) + // + tile_rows floats (bias). Default: (64 * 256 + 64) * 4 = 65,792 bytes. + const SHMEM_TILE_ROWS: u32 = 64; + const SHMEM_MAX_IN_DIM: u32 = 256; + let shmem_bytes = (SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM + SHMEM_TILE_ROWS) + * std::mem::size_of::() as u32; + let n = n_episodes as u32; - let (grid_x, block_x) = super::optimal_launch_dims(n, 0); + // Force block size to 256 for shared memory cooperative loading + let block_x = 256_u32; + let grid_x = (n + block_x - 1) / block_x; let grid_dim = (grid_x, 1, 1); let block_dim = (block_x, 1, 1); let launch_config = LaunchConfig { grid_dim, block_dim, - shared_mem_bytes: 0, + shared_mem_bytes: shmem_bytes, }; debug!( @@ -557,9 +566,10 @@ impl GpuExperienceCollector { timesteps, grid_x = grid_dim.0, block_x = block_dim.0, + shared_mem_kb = shmem_bytes / 1024, epsilon = config.epsilon, gamma = config.gamma, - "Launching dqn_full_experience_kernel" + "Launching dqn_full_experience_kernel (shmem weight tiling)" ); // ---- Step 5: Launch kernel ---- diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 2c956bef8..c21106b64 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -853,9 +853,9 @@ mod tests { use super::gpu_experience_collector::ExperienceCollectorConfig; let cfg = ExperienceCollectorConfig::default(); - assert_eq!(cfg.n_episodes, 128); + assert_eq!(cfg.n_episodes, 2048); assert_eq!(cfg.timesteps_per_episode, 500); - assert_eq!(cfg.total_experiences(), 128 * 500); + assert_eq!(cfg.total_experiences(), 2048 * 500); assert!((cfg.epsilon - 0.1).abs() < f32::EPSILON); assert!((cfg.gamma - 0.99).abs() < f32::EPSILON); assert!((cfg.barrier_profit_mult - 1.02).abs() < f32::EPSILON); diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 9315517d9..f5d78a802 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -540,6 +540,15 @@ impl DQNAgentType { self.memory().step(); } + /// Flush GPU-accumulated max priority to CPU (single scalar readback per epoch). + /// + /// Call once at epoch boundary after all `update_priorities_gpu` calls in the + /// training loop. Reduces GPU-CPU sync barriers from ~8/epoch to 1. + #[cfg(feature = "cuda")] + pub fn flush_max_priority(&self) -> Result<(), crate::MLError> { + self.memory().flush_max_priority() + } + /// Update learning rate for the optimizer(s) by applying a decay factor. /// /// For Standard agents, updates the single optimizer. diff --git a/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs b/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs index 09885d1c4..0976d22e0 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs @@ -21,7 +21,7 @@ pub(super) fn smoke_params() -> DQNHyperparameters { // (cached_capabilities() touches CUDA even when device is CPU) p.hidden_dim_base = Some(32); // Shrink GPU experience collection to fit the small test buffer. - // Default 128×500 = 63,500 experiences per collection exceeds buffer_size=500. + // buffer_size=500, so keep total experiences per collection ≤ buffer_size. p.gpu_n_episodes = 2; p.gpu_timesteps_per_episode = 50; p diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 482d80073..5ccc5c972 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -1129,30 +1129,37 @@ impl DQNTrainer { // Forward pass to get Q-values [batch_size, num_actions] let q_values = agent.forward(&batch_tensor)?; - // Flatten to get all Q-values - let q_vec: Vec = q_values - .to_vec2::() - .map_err(|e| crate::MLError::ModelError(format!("Failed to extract Q-values: {}", e)))? - .into_iter() - .flatten() - .collect(); - - // Calculate statistics - let min = q_vec.iter().cloned().fold(f64::INFINITY, |a, b| a.min(b as f64)); - let max = q_vec.iter().cloned().fold(f64::NEG_INFINITY, |a, b| a.max(b as f64)); - let sum: f64 = q_vec.iter().map(|&v| v as f64).sum(); - let count = q_vec.len(); - let mean = sum / count as f64; - - // Calculate standard deviation - let variance: f64 = q_vec.iter() - .map(|&v| { - let diff = v as f64 - mean; - diff * diff - }) - .sum::() / count as f64; + // GPU-side statistics: flatten Q-values and compute min/max/mean/std on device. + // Only 4 scalar readbacks (16 bytes) instead of downloading the entire tensor. + let q_f32 = q_values + .to_dtype(candle_core::DType::F32) + .map_err(|e| crate::MLError::ModelError(format!("Q-value F32 cast: {}", e)))?; + let q_flat = q_f32 + .flatten_all() + .map_err(|e| crate::MLError::ModelError(format!("Q-value flatten: {}", e)))?; + let count = q_flat.elem_count(); + + let min = q_flat.min(0) + .and_then(|t| t.to_vec0::()) + .map_err(|e| crate::MLError::ModelError(format!("Q-value min: {}", e)))? as f64; + let max = q_flat.max(0) + .and_then(|t| t.to_vec0::()) + .map_err(|e| crate::MLError::ModelError(format!("Q-value max: {}", e)))? as f64; + let mean_scalar = q_flat.mean_all() + .and_then(|t| t.to_vec0::()) + .map_err(|e| crate::MLError::ModelError(format!("Q-value mean: {}", e)))?; + let mean = mean_scalar as f64; + + // std = sqrt(mean((x - mean)^2)) — all on device + let mean_tensor = q_flat.mean_all() + .map_err(|e| crate::MLError::ModelError(format!("Q-value mean tensor: {}", e)))?; + let variance = q_flat.broadcast_sub(&mean_tensor) + .and_then(|d| d.sqr()) + .and_then(|sq| sq.mean_all()) + .and_then(|v| v.to_vec0::()) + .map_err(|e| crate::MLError::ModelError(format!("Q-value variance: {}", e)))? as f64; let std = variance.sqrt(); - + Ok(QValueStats { min, max, @@ -1201,13 +1208,24 @@ impl DQNTrainer { let sample_size = self.val_data.len().min(1000); // Sample up to 1000 for speed - // WAVE 10.6: Batched processing - collect all state vectors first - let mut state_vecs = Vec::with_capacity(sample_size); + // P4: Get aligned state_dim from the agent (already tensor-core aligned at construction) + // so we can pre-allocate the flat GPU buffer without an intermediate Vec>. + let aligned_state_dim = { + let agent = self.agent.read().await; + agent.get_state_dim() + }; + + // P4: Pre-allocate a single zero-initialized flat buffer for the entire batch. + // Zero-init handles tensor core padding columns automatically -- no separate + // padding branch needed. Each state is written at offset `i * aligned_state_dim`. + // This eliminates N intermediate Vec allocations (one per sample) and + // the subsequent flat_map copy pass from the WAVE 10.6 implementation. + let mut batched_states: Vec = vec![0.0f32; sample_size * aligned_state_dim]; let mut states = Vec::with_capacity(sample_size); let mut next_states = Vec::with_capacity(sample_size); let mut actions_for_rewards = Vec::with_capacity(sample_size); - for (feature_vec, target) in self.val_data.iter().take(sample_size) { + for (i, (feature_vec, target)) in self.val_data.iter().take(sample_size).enumerate() { let current_close = if target.len() >= 2 { target[0] } else { @@ -1226,23 +1244,31 @@ impl DQNTrainer { rust_decimal::Decimal::try_from(next_close).unwrap_or(rust_decimal::Decimal::ZERO); let next_state = self.feature_vector_to_state(feature_vec, Some(next_close_price))?; - state_vecs.push(state.to_vector()); + // P4: Write state vector directly into the flat buffer at the correct row offset. + // Raw dims (e.g. 45 or 53) are shorter than aligned (48 or 56); trailing + // positions stay zero from the vec![0.0; ..] init -- no explicit pad needed. + let sv = state.to_vector(); + let row_start = i * aligned_state_dim; + let copy_len = sv.len().min(aligned_state_dim); + let buf_len = batched_states.len(); + let dst = batched_states.get_mut(row_start..row_start + copy_len) + .ok_or_else(|| anyhow::anyhow!( + "Validation buffer overrun: row_start={}, copy_len={}, buf_len={}", + row_start, copy_len, buf_len + ))?; + let src = sv.get(..copy_len).ok_or_else(|| anyhow::anyhow!( + "State vector shorter than expected: len={}, copy_len={}", + sv.len(), copy_len + ))?; + dst.copy_from_slice(src); + states.push(state); next_states.push(next_state); } - // WAVE 10.6: Single batched forward pass for all validation samples + // P4: Single batched forward pass -- tensor created directly from the flat buffer, + // no intermediate Vec> needed. let agent = self.agent.read().await; - let raw_state_dim = state_vecs[0].len(); - let aligned_state_dim = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_state_dim, &self.device); - // Zero-pad each state vector to aligned dimension for tensor core alignment - let batched_states: Vec = if aligned_state_dim > raw_state_dim { - state_vecs.iter().flat_map(|v| { - v.iter().copied().chain(std::iter::repeat(0.0f32).take(aligned_state_dim - raw_state_dim)) - }).collect() - } else { - state_vecs.iter().flat_map(|v| v.iter().copied()).collect() - }; let batch_tensor = Tensor::from_vec(batched_states, (sample_size, aligned_state_dim), &self.device) .map_err(|e| anyhow::anyhow!("Failed to create batched validation tensor: {}", e))? .to_dtype(training_dtype(&self.device)) @@ -1253,7 +1279,8 @@ impl DQNTrainer { drop(agent); // Release lock early - // WAVE 10.6: GPU-optimized argmax for action selection + // P4: GPU-side argmax -- only the u32 action indices are downloaded to CPU. + // The Q-value tensor stays on GPU until dropped (no full Q-table transfer). let greedy_action_indices = batch_q_values .argmax(1) .map_err(|e| anyhow::anyhow!("Failed to compute validation argmax: {}", e))? @@ -1273,11 +1300,17 @@ impl DQNTrainer { let recent_actions_vec: Vec = self.recent_actions.iter().copied().collect(); - for i in 0..sample_size { + for (i, action) in actions_for_rewards.iter().enumerate() { + let state_ref = states.get(i).ok_or_else(|| { + anyhow::anyhow!("Validation state index {} out of bounds (len={})", i, states.len()) + })?; + let next_ref = next_states.get(i).ok_or_else(|| { + anyhow::anyhow!("Validation next_state index {} out of bounds (len={})", i, next_states.len()) + })?; let reward_decimal = self.reward_fn.calculate_reward( - actions_for_rewards[i], - &states[i], - &next_states[i], + *action, + state_ref, + next_ref, &recent_actions_vec, )?; val_rewards.push(reward_decimal.to_f32().unwrap_or(0.0) as f64); @@ -1809,16 +1842,22 @@ impl DQNTrainer { cfg.value_hidden_dim, cfg.advantage_hidden_dim, ); - // Dynamic episode count for buffer allocation + // Dynamic episode count for buffer allocation. + // Only auto-scale when configured >= 128 (production). + // Smaller values indicate an explicit test override — respect them. let alloc_episodes = { use ml_core::memory_optimization::detect_gpu_hardware; let configured = self.hyperparams.gpu_n_episodes; - match detect_gpu_hardware() { - Ok(hw) => configured.max(hw.optimal_n_episodes( - 48, - self.hyperparams.gpu_timesteps_per_episode, - )), - Err(_) => configured, + if configured >= 128 { + match detect_gpu_hardware() { + Ok(hw) => configured.max(hw.optimal_n_episodes( + state_dim, + self.hyperparams.gpu_timesteps_per_episode, + )), + Err(_) => configured, + } + } else { + configured } }; Some(GpuExperienceCollector::new( @@ -1845,16 +1884,22 @@ impl DQNTrainer { cfg.value_hidden_dim, cfg.advantage_hidden_dim, ); - // Dynamic episode count for buffer allocation + // Dynamic episode count for buffer allocation. + // Only auto-scale when configured >= 128 (production). + // Smaller values indicate an explicit test override — respect them. let alloc_episodes = { use ml_core::memory_optimization::detect_gpu_hardware; let configured = self.hyperparams.gpu_n_episodes; - match detect_gpu_hardware() { - Ok(hw) => configured.max(hw.optimal_n_episodes( - 48, - self.hyperparams.gpu_timesteps_per_episode, - )), - Err(_) => configured, + if configured >= 128 { + match detect_gpu_hardware() { + Ok(hw) => configured.max(hw.optimal_n_episodes( + state_dim, + self.hyperparams.gpu_timesteps_per_episode, + )), + Err(_) => configured, + } + } else { + configured } }; Some(GpuExperienceCollector::new( @@ -1942,25 +1987,33 @@ impl DQNTrainer { // Dynamic episode count: use GPU hardware to maximize SM utilization. // Falls back to configured value when GPU detection fails. + // Only auto-scale when configured >= 128 (production). + // Smaller values indicate an explicit test override — respect them. + let raw_sd = if self.hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 }; + let aligned_sd = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_sd, &self.device); let n_episodes = { use ml_core::memory_optimization::detect_gpu_hardware; let configured = self.hyperparams.gpu_n_episodes; - match detect_gpu_hardware() { - Ok(hw) => { - let optimal = hw.optimal_n_episodes( - 48, - self.hyperparams.gpu_timesteps_per_episode, - ); - let chosen = configured.max(optimal); - if chosen != configured { - info!( - "GPU auto-scaled n_episodes: {} → {} (SMs={}, VRAM={:.0}MB)", - configured, chosen, hw.sm_count, hw.free_memory_mb + if configured >= 128 { + match detect_gpu_hardware() { + Ok(hw) => { + let optimal = hw.optimal_n_episodes( + aligned_sd, + self.hyperparams.gpu_timesteps_per_episode, ); + let chosen = configured.max(optimal); + if chosen != configured { + info!( + "GPU auto-scaled n_episodes: {} → {} (SMs={}, VRAM={:.0}MB)", + configured, chosen, hw.sm_count, hw.free_memory_mb + ); + } + chosen as i32 } - chosen as i32 + Err(_) => configured as i32, } - Err(_) => configured as i32, + } else { + configured as i32 } }; let timesteps = self.hyperparams.gpu_timesteps_per_episode.min(1000) as i32; @@ -2903,6 +2956,17 @@ impl DQNTrainer { } } + // Flush GPU-accumulated max priority to CPU (single readback per epoch + // instead of ~8 per-batch readbacks). Safe because insert_batch uses the + // PREVIOUS max_priority — slightly stale is fine for proportional PER. + #[cfg(feature = "cuda")] + if train_step_count > 0 { + let agent = self.agent.read().await; + if let Err(e) = agent.flush_max_priority() { + debug!("GPU max_priority flush failed (non-fatal): {}", e); + } + } + let epoch_duration = epoch_start.elapsed(); // Calculate epoch metrics (average over training steps, not samples) @@ -3132,19 +3196,17 @@ impl DQNTrainer { training_metrics::record_gradient_explosion("dqn", "current"); } - // M3: Q-value gap diagnostic (Q_best - Q_second_best) - // Small gap (<0.01) = agent uncertain, Large gap (>1.0) = possibly overfit - if let Some((gap_mean, gap_min, gap_max)) = self.compute_q_gap_for_epoch().await { + // M3: Combined Q-value diagnostics (gap + per-action averages) + // Single forward pass + readback instead of two separate ones. + if let Some(((gap_mean, gap_min, gap_max), per_action_avgs)) = + self.compute_epoch_q_diagnostics().await + { info!( "Epoch {}/{}: Q-value gap (best-2nd): mean={:.4}, min={:.4}, max={:.4}", epoch + 1, self.hyperparams.epochs, gap_mean, gap_min, gap_max ); - } - // Per-action Q-value averages (sampled from replay buffer) - // Reveals whether the agent differentiates between exposure actions - if let Some(per_action_avgs) = self.compute_per_action_q_values().await { let names = ["S100", "S50", "Flat", "L50", "L100"]; let parts: Vec = names .iter() @@ -4414,108 +4476,19 @@ impl DQNTrainer { Ok(avg_q) } - /// M3: Compute Q-value gap (Q_best - Q_second_best) across a sample of states. + /// Epoch-end Q-value diagnostics: gap analysis + per-action averages. /// - /// The Q-gap diagnostic reveals action certainty: - /// - Small gap (<0.01): agent uncertain about which action to take - /// - Large gap (>1.0): agent highly confident (possibly overfit) + /// Merges the former `compute_q_gap_for_epoch` and `compute_per_action_q_values` + /// into a single forward pass + readback, eliminating one redundant buffer sample, + /// forward pass, and `to_vec2` GPU-CPU transfer per epoch. /// - /// Returns (mean_gap, min_gap, max_gap), or None if buffer has insufficient data. - async fn compute_q_gap_for_epoch(&self) -> Option<(f64, f64, f64)> { - let agent = self.agent.read().await; - let buffer = agent.memory(); - - if buffer.len() < 10 { - return None; - } - - let sample_size = buffer.len().min(100); - let batch_sample = match buffer.sample(sample_size) { - Ok(s) => s, - Err(_) => return None, - }; - - let state_dim = agent.get_state_dim(); - - // GPU PER path: use gpu_batch.states directly (experiences vec is empty) - #[allow(unused_mut)] - let mut batch_tensor_opt: Option = None; - #[cfg(feature = "cuda")] - { - if let Some(ref gpu) = batch_sample.gpu_batch { - batch_tensor_opt = gpu.states - .to_dtype(training_dtype(agent.device())) - .ok(); - } - } - let batch_tensor = if let Some(t) = batch_tensor_opt { - t - } else { - let batched_states: Vec = batch_sample.experiences - .iter() - .flat_map(|exp| { - let mut s = exp.state.clone(); - s.resize(state_dim, 0.0); - s - }) - .collect(); - let t = match Tensor::from_vec(batched_states, (sample_size, state_dim), agent.device()) { - Ok(t) => t, - Err(_) => return None, - }; - match t.to_dtype(training_dtype(agent.device())) { - Ok(t) => t, - Err(_) => return None, - } - }; - - let batch_q_values = match agent.forward(&batch_tensor) { - Ok(q) => q, - Err(_) => return None, - }; - - // Extract Q-values as 2D: [batch_size, num_actions] - let q_2d: Vec> = match batch_q_values.to_vec2::() { - Ok(v) => v, - Err(_) => return None, - }; - - let mut gaps = Vec::with_capacity(q_2d.len()); - for row in &q_2d { - if row.len() < 2 { - continue; - } - // Find best and second-best Q-values without full sort - let mut best = f32::NEG_INFINITY; - let mut second_best = f32::NEG_INFINITY; - for &v in row { - if v > best { - second_best = best; - best = v; - } else if v > second_best { - second_best = v; - } else { - // v <= second_best, skip - } - } - if best.is_finite() && second_best.is_finite() { - gaps.push((best - second_best) as f64); - } - } - - if gaps.is_empty() { - return None; - } - - let mean = gaps.iter().sum::() / gaps.len() as f64; - let min = gaps.iter().copied().fold(f64::INFINITY, f64::min); - let max = gaps.iter().copied().fold(f64::NEG_INFINITY, f64::max); - Some((mean, min, max)) - } - - /// Compute per-action average Q-values from a sample of replay buffer states. - /// Returns `[f64; 5]` averages (one per exposure action) or `None` if insufficient data. - async fn compute_per_action_q_values(&self) -> Option<[f64; 5]> { + /// Returns (gap_stats, per_action_avgs) where: + /// - gap_stats: (mean_gap, min_gap, max_gap) of Q_best - Q_second_best + /// - per_action_avgs: `[f64; 5]` averages (one per exposure action) + async fn compute_epoch_q_diagnostics(&self) -> Option<( + (f64, f64, f64), + [f64; 5], + )> { let agent = self.agent.read().await; let buffer = agent.memory(); @@ -4523,6 +4496,7 @@ impl DQNTrainer { return None; } + // Use the larger sample size (200) for both diagnostics let sample_size = buffer.len().min(200); let batch_sample = match buffer.sample(sample_size) { Ok(s) => s, @@ -4564,16 +4538,50 @@ impl DQNTrainer { } }; + // Single forward pass for both gap and per-action diagnostics let batch_q_values = match agent.forward(&batch_tensor) { Ok(q) => q, Err(_) => return None, }; + // Single readback let q_2d: Vec> = match batch_q_values.to_vec2::() { Ok(v) => v, Err(_) => return None, }; + // --- Q-value gap (best - second_best) --- + let mut gaps = Vec::with_capacity(q_2d.len()); + for row in &q_2d { + if row.len() < 2 { + continue; + } + let mut best = f32::NEG_INFINITY; + let mut second_best = f32::NEG_INFINITY; + for &v in row { + if v > best { + second_best = best; + best = v; + } else if v > second_best { + second_best = v; + } else { + // v <= second_best, skip + } + } + if best.is_finite() && second_best.is_finite() { + gaps.push((best - second_best) as f64); + } + } + + if gaps.is_empty() { + return None; + } + + let gap_mean = gaps.iter().sum::() / gaps.len() as f64; + let gap_min = gaps.iter().copied().fold(f64::INFINITY, f64::min); + let gap_max = gaps.iter().copied().fold(f64::NEG_INFINITY, f64::max); + + // --- Per-action average Q-values --- let mut sums = [0.0_f64; 5]; let mut counts = [0_usize; 5]; for row in &q_2d { @@ -4591,7 +4599,8 @@ impl DQNTrainer { avgs[i] = sums[i] / counts[i] as f64; } } - Some(avgs) + + Some(((gap_mean, gap_min, gap_max), avgs)) } /// Get current epsilon value