From 6fc90c0bdb03b845eb21013841f2125beea8ceab Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 21 Mar 2026 10:37:52 +0100 Subject: [PATCH] refactor(ml-dqn): is_weights_f32 kernel reads total_sum from GPU pointer Change the is_weights_f32 CUDA kernel signature from taking total_sum as a scalar f32 argument (which requires a CPU readback) to reading it from a const float* GPU-resident pointer. This eliminates the last memcpy_dtoh in the PER sampling hot path. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-dqn/src/replay_buffer_kernels.cu | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index bbcc88cbe..c6ef257d8 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -130,19 +130,21 @@ void pow_alpha_f32( // ── 8. IS weights: w_i = (N * sampled_prio[i] / total_sum)^(-beta) ───────── // Then normalize by max weight. +// total_sum_buf is a GPU-resident scalar pointer (zero CPU readback). extern "C" __global__ void is_weights_f32( float* __restrict__ weights, // [batch_size] output (in-place) const float* __restrict__ sampled_prios, // [batch_size] - float total_sum, + const float* __restrict__ total_sum_buf, // [1] GPU-resident scalar float neg_beta, int n_buffer, // number of experiences in buffer int batch_size ) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= batch_size) return; - float prob = (sampled_prios[i] * (float)n_buffer) / total_sum; + float ts = total_sum_buf[0]; + float prob = (sampled_prios[i] * (float)n_buffer) / ts; weights[i] = powf(prob, neg_beta); }