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

refactor(ml-dqn): is_weights_f32 kernel reads total_sum from GPU pointer

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-21 10:37:52 +01:00
parent 1eca11ef77
commit 6fc90c0bdb
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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