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

fix(cuda): move q_forward_dueling_warp_shmem to common header

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-15 15:56:23 +01:00
parent 8ea400e599
commit fee3c858ae
3 changed files with 114 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/ml/src/cuda_pipeline/backtest_forward_kernel.cu
View File

@@ -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] */

107
crates/ml/src/cuda_pipeline/common_device_functions.cuh
View File

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

102
crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu
View File

@@ -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.