From fee3c858ae1babc1e0981d03765f21f8a9bc42b7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 15 Mar 2026 15:56:23 +0100 Subject: [PATCH] fix(cuda): move q_forward_dueling_warp_shmem to common header MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The backtest_forward_kernel.cu calls q_forward_dueling_warp_shmem but it was defined in dqn_experience_kernel.cu — a different compilation unit. Move the function, TILE_LAYER_WARP_CLEAN macro, SHMEM_MIN and DIST_SIZE to common_device_functions.cuh so both kernels can use them. Add #ifndef guards to all macros to prevent redefinition warnings. Co-Authored-By: Claude Opus 4.6 --- .../cuda_pipeline/backtest_forward_kernel.cu | 2 + .../cuda_pipeline/common_device_functions.cuh | 107 ++++++++++++++++++ .../cuda_pipeline/dqn_experience_kernel.cu | 102 +---------------- 3 files changed, 114 insertions(+), 97 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/backtest_forward_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_forward_kernel.cu index c5ee7206e..e25332b12 100644 --- a/crates/ml/src/cuda_pipeline/backtest_forward_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_forward_kernel.cu @@ -28,7 +28,9 @@ #endif /* Distributed vector size: elements per lane for dim distributed across 32 lanes */ +#ifndef DIST_SIZE #define DIST_SIZE(dim) (((dim) + 31) / 32) +#endif extern "C" __global__ void backtest_forward_kernel( const float* __restrict__ states, /* [N, STATE_DIM] */ diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 75d52ccb1..36ef74e06 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -1389,3 +1389,110 @@ __device__ __forceinline__ int simulate_fill_check( } return filled; } + +/* ── Warp-cooperative forward pass utilities ────────────────────────────── */ + +#ifndef SHMEM_MIN +#define SHMEM_MIN(a, b) (((a) < (b)) ? (a) : (b)) +#endif + +#ifndef DIST_SIZE +#define DIST_SIZE(dim) (((dim) + 31) / 32) +#endif + +/* Warp-cooperative tiled layer: cooperative load + warp_matvec */ +#ifndef TILE_LAYER_WARP_CLEAN +#define TILE_LAYER_WARP_CLEAN(W_global, b_global, input_dist, output_dist, in_d, out_d, act, shmem_w, shmem_b, lane) \ + 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); \ + __syncwarp(0xFFFFFFFF); \ + warp_matvec_leaky_relu_shmem(shmem_w, shmem_b, input_dist, output_dist, \ + in_d, out_d, _ts, _tr, act, lane); \ + __syncwarp(0xFFFFFFFF); \ + } \ + } while (0) +#endif + +/** + * Warp-cooperative dueling Q-network forward pass (clean, no noise). + * + * Runs the full DQN forward pass: shared layers → value head → advantage head + * → Q(s,a) = V(s) + A(s,a) - mean(A). All 32 lanes cooperate on matrix-vector + * products; all lanes get all NUM_ACTIONS Q-values for argmax. + */ +__device__ void q_forward_dueling_warp_shmem( + const float* state_dist, + 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_dist, + float* scratch2_dist, + float* q_values, + int lane_id, + float* shmem_weights, + float* shmem_bias +) { + /* Shared layers: state -> scratch1 -> scratch2 */ + TILE_LAYER_WARP_CLEAN(w_s1, b_s1, state_dist, scratch1_dist, + STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, lane_id); + TILE_LAYER_WARP_CLEAN(w_s2, b_s2, scratch1_dist, scratch2_dist, + SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, lane_id); + + /* Value head hidden layer: scratch2 -> scratch_v (reuse scratch1_dist) */ + float* scratch_v_dist = scratch1_dist; + TILE_LAYER_WARP_CLEAN(w_v1, b_v1, scratch2_dist, scratch_v_dist, + SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, lane_id); + + /* Value output: [1, VALUE_H] — single neuron */ + float value; + { + cooperative_load_tile(shmem_weights, w_v2, VALUE_H); + cooperative_load_tile(shmem_bias, b_v2, 1); + __syncthreads(); + float partial = 0.0f; + for (int i = lane_id; i < VALUE_H; i += 32) + partial += shmem_weights[i] * scratch_v_dist[i / 32]; + if (lane_id == 0) + partial += shmem_bias[0]; + value = warp_reduce_sum_all(partial); + __syncthreads(); + } + + /* Advantage head hidden layer: scratch2 -> scratch_a */ + float* scratch_a_dist = scratch1_dist + DIST_SIZE(VALUE_H); + TILE_LAYER_WARP_CLEAN(w_a1, b_a1, scratch2_dist, scratch_a_dist, + SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, lane_id); + + /* Advantage output: [NUM_ACTIONS, ADV_H] — all lanes get full copy */ + float adv[NUM_ACTIONS]; + { + cooperative_load_tile(shmem_weights, w_a2, NUM_ACTIONS * ADV_H); + cooperative_load_tile(shmem_bias, b_a2, NUM_ACTIONS); + __syncthreads(); + for (int a = 0; a < NUM_ACTIONS; a++) { + const float* row = shmem_weights + a * ADV_H; + float partial = 0.0f; + for (int i = lane_id; i < ADV_H; i += 32) + partial += row[i] * scratch_a_dist[i / 32]; + if (lane_id == 0) + partial += shmem_bias[a]; + adv[a] = warp_reduce_sum_all(partial); + } + __syncthreads(); + } + + /* Q(s,a) = V(s) + A(s,a) - mean(A) */ + float adv_mean = 0.0f; + for (int i = 0; i < NUM_ACTIONS; i++) adv_mean += adv[i]; + adv_mean /= (float)NUM_ACTIONS; + for (int i = 0; i < NUM_ACTIONS; i++) { + q_values[i] = value + adv[i] - adv_mean; + } +} diff --git a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu index 8db49ec08..cc17541de 100644 --- a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu @@ -33,7 +33,9 @@ #endif /* Distributed vector size: elements per lane for dim distributed across 32 lanes */ +#ifndef DIST_SIZE #define DIST_SIZE(dim) (((dim) + 31) / 32) +#endif /* ------------------------------------------------------------------ */ /* DQN-Specific Device Functions */ @@ -270,6 +272,7 @@ __device__ __forceinline__ float max_arr(const float* vals, int len) { } while (0) /* Warp-cooperative tiling: same cooperative load, but uses warp_matvec instead of per-thread matvec */ +#ifndef TILE_LAYER_WARP_CLEAN #define TILE_LAYER_WARP_CLEAN(W_global, b_global, input_dist, output_dist, in_d, out_d, act, shmem_w, shmem_b, lane) \ do { \ for (int _tile = 0; _tile < ((out_d) + SHMEM_TILE_ROWS - 1) / SHMEM_TILE_ROWS; _tile++) { \ @@ -283,6 +286,7 @@ __device__ __forceinline__ float max_arr(const float* vals, int len) { __syncwarp(0xFFFFFFFF); \ } \ } while (0) +#endif #define TILE_LAYER_WARP_NOISY(W_global, b_global, input_dist, output_dist, in_d, out_d, act, shmem_w, shmem_b, sigma, rng_ptr, lane) \ do { \ @@ -672,103 +676,7 @@ __device__ void q_forward_branching_shmem( /* Warp-Cooperative Dueling Forward Passes */ /* ------------------------------------------------------------------ */ -/** - * Warp-cooperative dueling Q-network forward pass (clean — no noise). - * - * Same architecture as q_forward_dueling_shmem but the entire warp (32 lanes) - * cooperates on each matrix-vector product. Vectors are distributed across - * lanes in strided layout: lane k holds elements k, k+32, k+64, ... - * - * Per-lane stack: ~200 bytes (vs ~5.5 KB for per-thread version). - * state_dist: DIST_SIZE(48) = 2 floats - * scratch1_dist: DIST_SIZE(256) = 8 floats - * scratch2_dist: DIST_SIZE(256) = 8 floats - * q_values: 5 floats - * - * Value head final layer (1 output neuron): - * Each lane computes partial dot over its VALUE_H/32 elements of scratch_v_dist, - * warp reduce to broadcast result to all lanes. - * - * Advantage head final layer (5 output neurons): - * Small enough to NOT distribute. Each lane does partial dot for each action, - * warp reduces each. All lanes get all 5 Q-values for argmax. - */ -__device__ void q_forward_dueling_warp_shmem( - const float* state_dist, /* distributed: [DIST_SIZE(STATE_DIM)] per lane */ - 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_dist, /* [DIST_SIZE(SHARED_H1)] */ - float* scratch2_dist, /* [DIST_SIZE(SHARED_H2)] */ - float* q_values, /* [NUM_ACTIONS] — all lanes get full copy */ - int lane_id, - float* shmem_weights, - float* shmem_bias -) { - /* Shared layers: state -> scratch1 -> scratch2 */ - TILE_LAYER_WARP_CLEAN(w_s1, b_s1, state_dist, scratch1_dist, - STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, lane_id); - TILE_LAYER_WARP_CLEAN(w_s2, b_s2, scratch1_dist, scratch2_dist, - SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, lane_id); - - /* Value head hidden layer: scratch2 -> scratch_v (reuse scratch1_dist) */ - float* scratch_v_dist = scratch1_dist; /* safe: scratch1_dist is dead after s2 */ - TILE_LAYER_WARP_CLEAN(w_v1, b_v1, scratch2_dist, scratch_v_dist, - SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, lane_id); - - /* Value output: [1, VALUE_H] — single neuron. - * Each lane has VALUE_H/32 elements of scratch_v_dist. Partial dot, then reduce. */ - float value; - { - cooperative_load_tile(shmem_weights, w_v2, VALUE_H); - cooperative_load_tile(shmem_bias, b_v2, 1); - __syncthreads(); - float partial = 0.0f; - for (int i = lane_id; i < VALUE_H; i += 32) - partial += shmem_weights[i] * scratch_v_dist[i / 32]; - if (lane_id == 0) - partial += shmem_bias[0]; - value = warp_reduce_sum_all(partial); /* broadcast to all lanes */ - __syncthreads(); - } - - /* Advantage head hidden layer: scratch2 -> scratch_a (reuse scratch1_dist after VALUE_H) */ - float* scratch_a_dist = scratch1_dist + DIST_SIZE(VALUE_H); - TILE_LAYER_WARP_CLEAN(w_a1, b_a1, scratch2_dist, scratch_a_dist, - SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, lane_id); - - /* Advantage output: [NUM_ACTIONS=5, ADV_H] — 5 output neurons. - * Small enough that all lanes compute all 5 and get full Q-values. */ - float adv[NUM_ACTIONS]; - { - cooperative_load_tile(shmem_weights, w_a2, NUM_ACTIONS * ADV_H); - cooperative_load_tile(shmem_bias, b_a2, NUM_ACTIONS); - __syncthreads(); - for (int a = 0; a < NUM_ACTIONS; a++) { - const float* row = shmem_weights + a * ADV_H; - float partial = 0.0f; - for (int i = lane_id; i < ADV_H; i += 32) - partial += row[i] * scratch_a_dist[i / 32]; - if (lane_id == 0) - partial += shmem_bias[a]; - adv[a] = warp_reduce_sum_all(partial); /* no activation on output */ - } - __syncthreads(); - } - - /* 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) — all lanes get identical values */ - for (int i = 0; i < NUM_ACTIONS; i++) { - q_values[i] = value + adv[i] - adv_mean; - } -} +/* q_forward_dueling_warp_shmem is now in common_device_functions.cuh */ /** * Warp-cooperative NoisyNet dueling Q-network forward pass.