diff --git a/crates/ml-alpha/cuda/rl_per_push.cu b/crates/ml-alpha/cuda/rl_per_push.cu index 71a62a226..49b349c03 100644 --- a/crates/ml-alpha/cuda/rl_per_push.cu +++ b/crates/ml-alpha/cuda/rl_per_push.cu @@ -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] */