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

perf(cuda): vectorize rl_per_push — float4 loads (4× fewer memory transactions)

Browse Source 
Replace 3 sequential 128-float copy loops with 32 float4 vectorized
loads each. Reduces per-thread memory transactions 4× for the heaviest
kernel (46% of GPU time at 123μs/call → expected ~30μs).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 01:14:17 +02:00
parent 5f3e3578a0
commit e716d57d6f
1 changed files with 17 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
crates/ml-alpha/cuda/rl_per_push.cu
View File

@@ -16,11 +16,14 @@
* ===================================================================== */
#define HIDDEN_DIM 128
#define HIDDEN_DIM_F4 (HIDDEN_DIM / 4) /* 32 float4 loads per copy */
#define N_STEP_MAX 16
#define SCALARS_PER_TRANSITION 7
#define RL_GAMMA_INDEX 400
#define RL_N_STEP_INDEX 403
static_assert(HIDDEN_DIM % 4 == 0, "HIDDEN_DIM must be divisible by 4 for float4 vectorization");
extern "C" __global__ void rl_per_push(
/* Current step data */
const float* __restrict__ h_t_current, /* [B, 128] */
@@ -71,8 +74,10 @@ extern "C" __global__ void rl_per_push(
nstep_scalars[ring_base_scalars + 2] = dones[b];
const int ring_base_h = (b * N_STEP_MAX + ring_idx) * HIDDEN_DIM;
for (int i = 0; i < HIDDEN_DIM; ++i) {
nstep_h_t[ring_base_h + i] = h_t_current[b * HIDDEN_DIM + i];
{
const float4* src4 = (const float4*)(h_t_current + b * HIDDEN_DIM);
float4* dst4 = (float4*)(nstep_h_t + ring_base_h);
for (int i = 0; i < HIDDEN_DIM_F4; ++i) dst4[i] = src4[i];
}
nstep_write_idx[b] = (ring_idx + 1) % N_STEP_MAX;
@@ -150,15 +155,19 @@ extern "C" __global__ void rl_per_push(
n_step_gamma = gamma_power; /* γ^actual_n */
}
/* Copy oldest h_t from ring → replay_h_t[my_slot] */
/* Copy oldest h_t from ring → replay_h_t[my_slot] (float4 coalesced) */
const int oldest_h_base = (b * N_STEP_MAX + oldest_ring) * HIDDEN_DIM;
for (int i = 0; i < HIDDEN_DIM; ++i) {
replay_h_t[my_slot * HIDDEN_DIM + i] = nstep_h_t[oldest_h_base + i];
{
const float4* src4 = (const float4*)(nstep_h_t + oldest_h_base);
float4* dst4 = (float4*)(replay_h_t + my_slot * HIDDEN_DIM);
for (int i = 0; i < HIDDEN_DIM_F4; ++i) dst4[i] = src4[i];
}
/* Copy h_tp1_current[b] → replay_h_tp1[my_slot] */
for (int i = 0; i < HIDDEN_DIM; ++i) {
replay_h_tp1[my_slot * HIDDEN_DIM + i] = h_tp1_current[b * HIDDEN_DIM + i];
/* Copy h_tp1_current[b] → replay_h_tp1[my_slot] (float4 coalesced) */
{
const float4* src4 = (const float4*)(h_tp1_current + b * HIDDEN_DIM);
float4* dst4 = (float4*)(replay_h_tp1 + my_slot * HIDDEN_DIM);
for (int i = 0; i < HIDDEN_DIM_F4; ++i) dst4[i] = src4[i];
}
/* Write scalars: [action_f32, scaled_reward, raw_reward, done, log_pi_old, n_step_gamma, n_step_return_raw] */