diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index fba09477b..5afc671d4 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -258,6 +258,245 @@ __device__ void noisy_matvec_leaky_relu_shmem( } } +/* ------------------------------------------------------------------ */ +/* Warp-Cooperative Matrix-Vector Multiply */ +/* ------------------------------------------------------------------ */ + +/** + * Warp-cooperative reduce: sum across 32 lanes using shuffle-down. + * Result is valid in ALL lanes (full butterfly reduction). + */ +__device__ float warp_reduce_sum_all(float val) { + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + +/** + * Warp-cooperative matrix-vector multiply with LeakyReLU, shared-memory tiled. + * + * Unlike the per-thread version, a full warp (32 lanes) cooperates on each + * dot product. Input and output vectors are distributed across lanes in + * strided layout: lane k holds elements k, k+32, k+64, ... + * + * This reduces per-thread local memory from in_dim floats (~2KB for dim=512) + * to ceil(in_dim/32) floats (~64B), eliminating register spills on H100. + * + * @param shmem_W Shared memory weight tile [tile_rows × in_dim] + * @param shmem_b Shared memory bias tile [tile_rows] + * @param input_dist Per-lane distributed input [ceil(in_dim/32)] + * @param output_dist Per-lane distributed output [ceil(out_dim/32)] (written) + * @param in_dim Full input dimension + * @param out_dim Full output dimension (only used for bounds context) + * @param tile_offset Global output offset for this tile + * @param tile_rows Number of output rows in this tile (≤ SHMEM_TILE_ROWS) + * @param activate 1 = apply LeakyReLU, 0 = linear + * @param lane_id threadIdx.x % 32 + */ +__device__ void warp_matvec_leaky_relu_shmem( + const float* __restrict__ shmem_W, + const float* __restrict__ shmem_b, + const float* input_dist, + float* output_dist, + int in_dim, + int out_dim, + int tile_offset, + int tile_rows, + int activate, + int lane_id +) { + (void)out_dim; + for (int j = 0; j < tile_rows; j++) { + const float* row = shmem_W + j * in_dim; + /* Each lane computes partial dot product over its elements */ + float partial = 0.0f; + for (int i = lane_id; i < in_dim; i += 32) + partial += row[i] * input_dist[i / 32]; + /* Add bias (only once, via lane 0) */ + if (lane_id == 0) + partial += shmem_b[j]; + /* Full warp reduction — result in all lanes */ + float sum = warp_reduce_sum_all(partial); + /* Activation */ + if (activate) sum = leaky_relu(sum); + /* Scatter to owning lane in strided layout */ + int gj = tile_offset + j; + if (lane_id == (gj & 31)) + output_dist[gj >> 5] = sum; + } +} + +/** + * Warp-cooperative noisy matrix-vector multiply with factorized Gaussian noise. + * + * Same as warp_matvec_leaky_relu_shmem but applies NoisyNet perturbation. + * Noise vectors are generated on-the-fly per lane (no large eps_in array needed). + * + * @param sigma_init NoisyNet sigma_0 parameter + * @param rng Per-thread RNG state pointer + */ +__device__ void warp_noisy_matvec_leaky_relu_shmem( + const float* __restrict__ shmem_W, + const float* __restrict__ shmem_b, + const float* input_dist, + float* output_dist, + int in_dim, + int out_dim, + int tile_offset, + int tile_rows, + int activate, + float sigma_init, + unsigned int* rng, + int lane_id +) { + (void)out_dim; + float sigma_scale = sigma_init / sqrtf((float)in_dim); + + /* Each lane generates its portion of input noise (strided) */ + /* With in_dim=512 and 32 lanes, each lane generates 16 noise values */ + int local_count = (in_dim + 31) / 32; + + for (int j = 0; j < tile_rows; j++) { + /* Generate output noise for this row — same value needed by all lanes */ + /* Use lane 0's RNG and broadcast */ + float eps_out_j; + if (lane_id == 0) + eps_out_j = factorized_noise_fn(gpu_random_gaussian(rng)); + eps_out_j = __shfl_sync(0xFFFFFFFF, eps_out_j, 0); + + const float* row = shmem_W + j * in_dim; + float partial = 0.0f; + for (int k = 0; k < local_count; k++) { + int i = lane_id + k * 32; + if (i < in_dim) { + float eps_in_i = factorized_noise_fn(gpu_random_gaussian(rng)); + float w_noisy = row[i] + sigma_scale * eps_out_j * eps_in_i; + partial += w_noisy * input_dist[k]; + } + } + /* Add noisy bias (lane 0 only) */ + if (lane_id == 0) + partial += shmem_b[j] + sigma_scale * eps_out_j; + /* Full warp reduction */ + float sum = warp_reduce_sum_all(partial); + if (activate) sum = leaky_relu(sum); + int gj = tile_offset + j; + if (lane_id == (gj & 31)) + output_dist[gj >> 5] = sum; + } +} + +/** + * Warp-cooperative matvec with broadcast output (non-distributed). + * + * Same as warp_matvec_leaky_relu_shmem but ALL lanes store every output + * element (no strided scatter). Used for small output layers where all + * lanes need the full result (e.g., C51 atom logits). + * + * @param output Non-distributed output array — all lanes write all elements + */ +__device__ void warp_matvec_broadcast_shmem( + const float* __restrict__ shmem_W, + const float* __restrict__ shmem_b, + const float* input_dist, + float* output, /* NOT distributed — all lanes get full copy */ + int in_dim, + int out_dim, + int tile_offset, + int tile_rows, + int activate, + int lane_id +) { + (void)out_dim; + for (int j = 0; j < tile_rows; j++) { + const float* row = shmem_W + j * in_dim; + float partial = 0.0f; + for (int i = lane_id; i < in_dim; i += 32) + partial += row[i] * input_dist[i / 32]; + if (lane_id == 0) + partial += shmem_b[j]; + float sum = warp_reduce_sum_all(partial); + if (activate) sum = leaky_relu(sum); + /* Broadcast: ALL lanes store (redundant but correct for small arrays) */ + output[tile_offset + j] = sum; + } +} + +/** + * Warp-cooperative noisy matvec with broadcast output. + * Same as warp_noisy_matvec_leaky_relu_shmem but all lanes get full output. + */ +__device__ void warp_noisy_matvec_broadcast_shmem( + const float* __restrict__ shmem_W, + const float* __restrict__ shmem_b, + const float* input_dist, + float* output, /* NOT distributed — all lanes get full copy */ + int in_dim, + int out_dim, + int tile_offset, + int tile_rows, + int activate, + float sigma_init, + unsigned int* rng, + int lane_id +) { + (void)out_dim; + float sigma_scale = sigma_init / sqrtf((float)in_dim); + int local_count = (in_dim + 31) / 32; + + for (int j = 0; j < tile_rows; j++) { + float eps_out_j; + if (lane_id == 0) + eps_out_j = factorized_noise_fn(gpu_random_gaussian(rng)); + eps_out_j = __shfl_sync(0xFFFFFFFF, eps_out_j, 0); + + const float* row = shmem_W + j * in_dim; + float partial = 0.0f; + for (int k = 0; k < local_count; k++) { + int i = lane_id + k * 32; + if (i < in_dim) { + float eps_in_i = factorized_noise_fn(gpu_random_gaussian(rng)); + float w_noisy = row[i] + sigma_scale * eps_out_j * eps_in_i; + partial += w_noisy * input_dist[k]; + } + } + if (lane_id == 0) + partial += shmem_b[j] + sigma_scale * eps_out_j; + float sum = warp_reduce_sum_all(partial); + if (activate) sum = leaky_relu(sum); + output[tile_offset + j] = sum; + } +} + +/** + * Warp-cooperative in-place RMSNorm on distributed data. + * + * Each lane holds elements in strided layout. Sum-of-squares is computed + * via warp reduction, then each lane normalizes its own elements. + * + * @param data_dist Distributed data [DIST_SIZE(dim)] per lane (modified in place) + * @param gamma Non-distributed gamma weights [dim] in global memory + * @param dim Full dimension + * @param lane_id Lane index (0..31) + */ +__device__ void warp_rmsnorm_inplace(float* data_dist, const float* __restrict__ gamma, int dim, int lane_id) { + /* Each lane computes partial sum of squares over its elements */ + float partial_sq = 0.0f; + for (int i = lane_id; i < dim; i += 32) { + float val = data_dist[i / 32]; + partial_sq += val * val; + } + /* Warp reduction for total sum of squares */ + float total_sq = warp_reduce_sum_all(partial_sq); + float rms = sqrtf(total_sq / (float)dim + 1e-6f); + float inv_rms = 1.0f / rms; + + /* Each lane normalizes its own elements */ + for (int i = lane_id; i < dim; i += 32) { + data_dist[i / 32] = data_dist[i / 32] * inv_rms * gamma[i]; + } +} + /** * In-place RMSNorm: data[i] = (data[i] / RMS(data)) * gamma[i]. * diff --git a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu index d3684f7f6..ce0c410a1 100644 --- a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu @@ -32,6 +32,9 @@ #define SHMEM_MIN(a, b) (((a) < (b)) ? (a) : (b)) #endif +/* Distributed vector size: elements per lane for dim distributed across 32 lanes */ +#define DIST_SIZE(dim) (((dim) + 31) / 32) + /* ------------------------------------------------------------------ */ /* DQN-Specific Device Functions */ /* ------------------------------------------------------------------ */ @@ -142,6 +145,65 @@ __device__ void q_forward_dueling( } \ } while (0) +/* Warp-cooperative tiling: same cooperative load, but uses warp_matvec instead of per-thread matvec */ +#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) + +#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 { \ + 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_noisy_matvec_leaky_relu_shmem(shmem_w, shmem_b, input_dist, output_dist, \ + in_d, out_d, _ts, _tr, act, sigma, rng_ptr, lane); \ + __syncwarp(0xFFFFFFFF); \ + } \ + } while (0) + +/* Warp-cooperative tiling with broadcast output (all lanes get full array). + * Used for small C51 atom output layers where all lanes need the complete result. */ +#define TILE_LAYER_WARP_BROADCAST_CLEAN(W_global, b_global, input_dist, output, 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_broadcast_shmem(shmem_w, shmem_b, input_dist, output, \ + in_d, out_d, _ts, _tr, act, lane); \ + __syncwarp(0xFFFFFFFF); \ + } \ + } while (0) + +#define TILE_LAYER_WARP_BROADCAST_NOISY(W_global, b_global, input_dist, output, in_d, out_d, act, shmem_w, shmem_b, sigma, rng_ptr, 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_noisy_matvec_broadcast_shmem(shmem_w, shmem_b, input_dist, output, \ + in_d, out_d, _ts, _tr, act, sigma, rng_ptr, lane); \ + __syncwarp(0xFFFFFFFF); \ + } \ + } while (0) + /** * Shared-memory dueling Q-network forward pass. * @@ -375,6 +437,337 @@ __device__ void q_forward_distributional_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; + } +} + +/** + * Warp-cooperative NoisyNet dueling Q-network forward pass. + * + * Same as q_forward_dueling_warp_shmem but with factorized Gaussian noise. + * Only used for the ONLINE network during exploration. + */ +__device__ void q_forward_dueling_noisy_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 */ + float sigma_init, + unsigned int* rng, + int lane_id, + float* shmem_weights, + float* shmem_bias +) { + /* Shared layers with noise: state -> scratch1 -> scratch2 */ + TILE_LAYER_WARP_NOISY(w_s1, b_s1, state_dist, scratch1_dist, + STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, + sigma_init, rng, lane_id); + TILE_LAYER_WARP_NOISY(w_s2, b_s2, scratch1_dist, scratch2_dist, + SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, + sigma_init, rng, lane_id); + + /* Value head hidden layer with noise */ + float* scratch_v_dist = scratch1_dist; + TILE_LAYER_WARP_NOISY(w_v1, b_v1, scratch2_dist, scratch_v_dist, + SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, + sigma_init, rng, lane_id); + + /* Value output: [1, VALUE_H] — noisy single neuron */ + float value; + { + cooperative_load_tile(shmem_weights, w_v2, VALUE_H); + cooperative_load_tile(shmem_bias, b_v2, 1); + __syncthreads(); + float sigma_scale = sigma_init / sqrtf((float)VALUE_H); + /* Output noise — lane 0 generates, broadcast to all */ + float eps_out_v; + if (lane_id == 0) + eps_out_v = factorized_noise_fn(gpu_random_gaussian(rng)); + eps_out_v = __shfl_sync(0xFFFFFFFF, eps_out_v, 0); + + float partial = 0.0f; + for (int i = lane_id; i < VALUE_H; i += 32) { + float eps_in_i = factorized_noise_fn(gpu_random_gaussian(rng)); + float w_noisy = shmem_weights[i] + sigma_scale * eps_out_v * eps_in_i; + partial += w_noisy * scratch_v_dist[i / 32]; + } + if (lane_id == 0) + partial += shmem_bias[0] + sigma_scale * eps_out_v; + value = warp_reduce_sum_all(partial); + __syncthreads(); + } + + /* Advantage head hidden layer with noise */ + float* scratch_a_dist = scratch1_dist + DIST_SIZE(VALUE_H); + TILE_LAYER_WARP_NOISY(w_a1, b_a1, scratch2_dist, scratch_a_dist, + SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, + sigma_init, rng, lane_id); + + /* Advantage output: [NUM_ACTIONS=5, ADV_H] — 5 noisy output neurons */ + 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(); + float sigma_scale_a = sigma_init / sqrtf((float)ADV_H); + for (int a = 0; a < NUM_ACTIONS; a++) { + /* Output noise per action — lane 0 generates, broadcast */ + float eps_out_a; + if (lane_id == 0) + eps_out_a = factorized_noise_fn(gpu_random_gaussian(rng)); + eps_out_a = __shfl_sync(0xFFFFFFFF, eps_out_a, 0); + + const float* row = shmem_weights + a * ADV_H; + float partial = 0.0f; + for (int i = lane_id; i < ADV_H; i += 32) { + float eps_in_i = factorized_noise_fn(gpu_random_gaussian(rng)); + float w_noisy = row[i] + sigma_scale_a * eps_out_a * eps_in_i; + partial += w_noisy * scratch_a_dist[i / 32]; + } + if (lane_id == 0) + partial += shmem_bias[a] + sigma_scale_a * eps_out_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; + } +} + +/** + * Warp-cooperative C51 distributional dueling Q-network forward pass. + * + * Uses distributed vectors for the heavy shared layers (s1, s2, v1, a1) + * and broadcast output for the small atom layers (v2, a2). + * All lanes get identical q_values[NUM_ACTIONS] at the end. + * + * Optionally applies NoisyNet noise and RMSNorm. + */ +__device__ void q_forward_distributional_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, + const float* __restrict__ rms_s0, const float* __restrict__ rms_s1, + const float* __restrict__ rms_v, const float* __restrict__ rms_a, + 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 num_atoms, + float v_min, float v_max, + int use_noisy, float sigma_init, int use_rmsnorm, + unsigned int* rng, + int lane_id, + float* shmem_weights, + float* shmem_bias +) { + int na = (num_atoms > NUM_ATOMS_MAX) ? NUM_ATOMS_MAX : num_atoms; + + /* ---- Shared layers (distributed I/O) ---- */ + if (use_noisy) { + TILE_LAYER_WARP_NOISY(w_s1, b_s1, state_dist, scratch1_dist, STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, sigma_init, rng, lane_id); + } else { + TILE_LAYER_WARP_CLEAN(w_s1, b_s1, state_dist, scratch1_dist, STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias, lane_id); + } + if (use_rmsnorm) warp_rmsnorm_inplace(scratch1_dist, rms_s0, SHARED_H1, lane_id); + + if (use_noisy) { + TILE_LAYER_WARP_NOISY(w_s2, b_s2, scratch1_dist, scratch2_dist, SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, sigma_init, rng, lane_id); + } else { + TILE_LAYER_WARP_CLEAN(w_s2, b_s2, scratch1_dist, scratch2_dist, SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias, lane_id); + } + if (use_rmsnorm) warp_rmsnorm_inplace(scratch2_dist, rms_s1, SHARED_H2, lane_id); + + /* ---- Value hidden layer (distributed I/O) ---- */ + /* Scratch aliasing: scratch1_dist is dead after s2. Reuse for value/adv heads. */ + float* scratch_v_dist = scratch1_dist; /* [DIST_SIZE(VALUE_H)] */ + /* scratch_a_dist starts after value portion in scratch1_dist */ + float* scratch_a_dist = scratch1_dist + DIST_SIZE(VALUE_H); + + if (use_noisy) { + TILE_LAYER_WARP_NOISY(w_v1, b_v1, scratch2_dist, scratch_v_dist, SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, sigma_init, rng, lane_id); + } else { + TILE_LAYER_WARP_CLEAN(w_v1, b_v1, scratch2_dist, scratch_v_dist, SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias, lane_id); + } + if (use_rmsnorm) warp_rmsnorm_inplace(scratch_v_dist, rms_v, VALUE_H, lane_id); + + /* ---- Value atom output (broadcast — all lanes get val_atoms[na]) ---- */ + float val_atoms[NUM_ATOMS_MAX]; + if (use_noisy) { + TILE_LAYER_WARP_BROADCAST_NOISY(w_v2, b_v2, scratch_v_dist, val_atoms, VALUE_H, na, 0, shmem_weights, shmem_bias, sigma_init, rng, lane_id); + } else { + TILE_LAYER_WARP_BROADCAST_CLEAN(w_v2, b_v2, scratch_v_dist, val_atoms, VALUE_H, na, 0, shmem_weights, shmem_bias, lane_id); + } + + /* ---- Advantage hidden layer (distributed I/O) ---- */ + if (use_noisy) { + TILE_LAYER_WARP_NOISY(w_a1, b_a1, scratch2_dist, scratch_a_dist, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, sigma_init, rng, lane_id); + } else { + TILE_LAYER_WARP_CLEAN(w_a1, b_a1, scratch2_dist, scratch_a_dist, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias, lane_id); + } + if (use_rmsnorm) warp_rmsnorm_inplace(scratch_a_dist, rms_a, ADV_H, lane_id); + + /* ---- Advantage atom output (broadcast — all lanes get adv_atoms[na*5]) ---- */ + int adv_total = NUM_ACTIONS * na; + float adv_atoms[NUM_ACTIONS * NUM_ATOMS_MAX]; + if (use_noisy) { + TILE_LAYER_WARP_BROADCAST_NOISY(w_a2, b_a2, scratch_a_dist, adv_atoms, ADV_H, adv_total, 0, shmem_weights, shmem_bias, sigma_init, rng, lane_id); + } else { + TILE_LAYER_WARP_BROADCAST_CLEAN(w_a2, b_a2, scratch_a_dist, adv_atoms, ADV_H, adv_total, 0, shmem_weights, shmem_bias, lane_id); + } + + /* ---- Distributional dueling combination + softmax → expected Q ---- + * All lanes have identical val_atoms and adv_atoms, so this is redundant + * across lanes but correct and negligible cost compared to matvec savings. */ + float delta_z = (na > 1) ? (v_max - v_min) / (float)(na - 1) : 0.0f; + + for (int a = 0; a < NUM_ACTIONS; a++) { + float logits[NUM_ATOMS_MAX]; + for (int i = 0; i < na; i++) { + float adv_mean_i = 0.0f; + for (int ap = 0; ap < NUM_ACTIONS; ap++) { + adv_mean_i += adv_atoms[ap * na + i]; + } + adv_mean_i /= (float)NUM_ACTIONS; + logits[i] = val_atoms[i] + adv_atoms[a * na + i] - adv_mean_i; + } + + float max_logit = logits[0]; + for (int i = 1; i < na; i++) { + if (logits[i] > max_logit) max_logit = logits[i]; + } + float sum_exp = 0.0f; + for (int i = 0; i < na; i++) { + logits[i] = expf(logits[i] - max_logit); + sum_exp += logits[i]; + } + float inv_sum = 1.0f / (sum_exp + 1e-8f); + + float q_expected = 0.0f; + for (int i = 0; i < na; i++) { + float z_i = v_min + (float)i * delta_z; + float p_i = logits[i] * inv_sum; + q_expected += z_i * p_i; + } + q_values[a] = q_expected; + } +} + /* ------------------------------------------------------------------ */ /* D5: NoisyNet Forward Pass */ /* ------------------------------------------------------------------ */ @@ -1199,3 +1592,721 @@ extern "C" __global__ void dqn_full_experience_kernel( rng_states[tid] = rng; } } + +/* ==================================================================== */ +/* Warp-Cooperative Experience Collection Kernel (sm_90+) */ +/* ==================================================================== */ + +/** + * Warp-cooperative DQN experience collection kernel. + * + * One warp (32 threads) per episode. All 32 lanes cooperate on each + * matrix-vector multiply, distributing vectors across lanes in strided + * layout (lane k holds elements k, k+32, k+64, ...). + * + * Per-lane stack: ~200 bytes (down from ~5.5 KB in the per-thread kernel). + * state_dist: DIST_SIZE(48) = 2 floats = 8 bytes + * scratch1_dist: DIST_SIZE(256) = 8 floats = 32 bytes + * scratch2_dist: DIST_SIZE(256) = 8 floats = 32 bytes + * q_values: 5 floats = 20 bytes + * tgt_q_values: 5 floats = 20 bytes + * lane 0 extras: cur_scratch[64] + state_full[48] + next_state_full[48] + * = 640 bytes (only lane 0) + * + * Launch config: grid=(N, 1, 1), block=(32, 1, 1). + * One block per episode. All 32 lanes must participate in __syncthreads() + * and cooperative_load_tile together. + * + * C51 distributional mode is supported via q_forward_distributional_warp_shmem + * which uses broadcast output for the small atom layers (v2, a2). + * Parameter signature is BYTE-FOR-BYTE identical for host code reuse. + */ +extern "C" __global__ void dqn_full_experience_kernel_warp( + /* Market data [total_bars, MARKET_DIM] */ + const float* __restrict__ market_features, + /* Target prices [total_bars, 4]: preproc_close, preproc_next, raw_close, raw_next */ + const float* __restrict__ targets, + /* Episode start indices [N] — global bar index where each episode begins */ + const int* __restrict__ episode_starts, + + /* ---- Online Q-network weights (12 pointers) ---- */ + const float* __restrict__ on_w_s1, + const float* __restrict__ on_b_s1, + const float* __restrict__ on_w_s2, + const float* __restrict__ on_b_s2, + const float* __restrict__ on_w_v1, + const float* __restrict__ on_b_v1, + const float* __restrict__ on_w_v2, + const float* __restrict__ on_b_v2, + const float* __restrict__ on_w_a1, + const float* __restrict__ on_b_a1, + const float* __restrict__ on_w_a2, + const float* __restrict__ on_b_a2, + + /* ---- Target Q-network weights (12 pointers) ---- */ + const float* __restrict__ tg_w_s1, + const float* __restrict__ tg_b_s1, + const float* __restrict__ tg_w_s2, + const float* __restrict__ tg_b_s2, + const float* __restrict__ tg_w_v1, + const float* __restrict__ tg_b_v1, + const float* __restrict__ tg_w_v2, + const float* __restrict__ tg_b_v2, + const float* __restrict__ tg_w_a1, + const float* __restrict__ tg_b_a1, + const float* __restrict__ tg_w_a2, + const float* __restrict__ tg_b_a2, + + /* ---- Curiosity model weights (4 pointers) ---- */ + const float* __restrict__ cur_w1, + const float* __restrict__ cur_b1, + const float* __restrict__ cur_w2, + const float* __restrict__ cur_b2, + + /* ---- Per-episode mutable state arrays ---- */ + float* portfolio_states, /* [N, PORTFOLIO_STATE_SIZE] */ + float* barrier_states, /* [N, BARRIER_STATE_SIZE] */ + int* diversity_windows, /* [N, DIVERSITY_WINDOW] */ + int* diversity_metas, /* [N, 2] */ + + /* ---- Barrier config (shared) ---- */ + const float* __restrict__ barrier_config, /* [3]: profit_mult, loss_mult, max_bars */ + + /* ---- Scalar configs ---- */ + float epsilon, + float max_position, + int episode_length, + int total_bars, + int L, /* timesteps per episode */ + float gamma, + float curiosity_max_reward, + int N, /* total number of episodes */ + float barrier_scale, + float diversity_scale, + float curiosity_scale, + float risk_weight, + float hold_reward, + float tx_cost_multiplier, + float count_bonus_coefficient, + float q_clip_min, + float q_clip_max, + float huber_kappa, + + /* ---- D5: NoisyNet config ---- */ + int use_noisy_nets, + float noisy_sigma_init, + + /* ---- D6: C51 Distributional config ---- */ + int use_distributional, + int num_atoms, + float v_min, + float v_max, + float reward_norm_alpha, + int use_rmsnorm, + + /* ---- D6: RMSNorm gamma weights (4 pointers) ---- */ + const float* __restrict__ rms_s0_gamma, /* [SHARED_H1] */ + const float* __restrict__ rms_s1_gamma, /* [SHARED_H2] */ + const float* __restrict__ rms_v_gamma, /* [VALUE_H] */ + const float* __restrict__ rms_a_gamma, /* [ADV_H] */ + + /* ---- RNG states [N] ---- */ + unsigned int* rng_states, + + /* ---- Output arrays ---- */ + float* out_states, /* [N, L, STATE_DIM] */ + int* out_actions, /* [N, L] */ + float* out_rewards, /* [N, L] */ + int* out_done, /* [N, L] */ + float* out_target_q, /* [N, L] */ + float* out_td_error /* [N, L] */ +) { + /* ---- Shared memory workspace for weight tiling ---- */ + extern __shared__ float shmem[]; + float* shmem_weights = shmem; + float* shmem_bias = shmem + SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM; + + /* ---- Thread / warp identity ---- */ + int episode_id = blockIdx.x; /* one block per episode */ + int lane_id = threadIdx.x; /* 0..31 within the warp */ + int tid_valid = (episode_id < N) ? 1 : 0; + + /* ---- Load per-episode state (lane 0 owns simulation state) ---- */ + int ps_off = tid_valid ? (episode_id * PORTFOLIO_STATE_SIZE) : 0; + /* All lanes read portfolio state — only lane 0 uses it for simulation, + * but we avoid divergent memory ops. Cost: 8 floats, negligible. */ + float cash = tid_valid ? portfolio_states[ps_off + 0] : 0.0f; + float position = tid_valid ? portfolio_states[ps_off + 1] : 0.0f; + float entry_price = tid_valid ? portfolio_states[ps_off + 2] : 0.0f; + float initial_cap = tid_valid ? portfolio_states[ps_off + 3] : 1.0f; + float spread = tid_valid ? portfolio_states[ps_off + 4] : 0.0f; + float last_price = tid_valid ? portfolio_states[ps_off + 5] : 0.0f; + float reserve_pct = tid_valid ? portfolio_states[ps_off + 6] : 0.0f; + float cum_costs = tid_valid ? portfolio_states[ps_off + 7] : 0.0f; + + int bs_off = tid_valid ? (episode_id * BARRIER_STATE_SIZE) : 0; + float barrier_st[BARRIER_STATE_SIZE]; + for (int i = 0; i < BARRIER_STATE_SIZE; i++) + barrier_st[i] = tid_valid ? barrier_states[bs_off + i] : 0.0f; + + int dw_off = tid_valid ? (episode_id * DIVERSITY_WINDOW) : 0; + int div_window[DIVERSITY_WINDOW]; + for (int i = 0; i < DIVERSITY_WINDOW; i++) + div_window[i] = tid_valid ? diversity_windows[dw_off + i] : 0; + + int dm_off = tid_valid ? (episode_id * 2) : 0; + int div_meta[2]; + div_meta[0] = tid_valid ? diversity_metas[dm_off + 0] : 0; + div_meta[1] = tid_valid ? diversity_metas[dm_off + 1] : 0; + + unsigned int rng = tid_valid ? rng_states[episode_id] : 0u; + int ep_start = tid_valid ? episode_starts[episode_id] : 0; + + /* ---- Distributed scratch buffers (per-lane, tiny) ---- */ + float state_dist[DIST_SIZE(STATE_DIM)]; + float scratch1_dist[DIST_SIZE(SHARED_H1)]; + float scratch2_dist[DIST_SIZE(SHARED_H2)]; + float q_values[NUM_ACTIONS]; + float tgt_q_values[NUM_ACTIONS]; + + /* Zero-init distributed state */ + for (int i = 0; i < DIST_SIZE(STATE_DIM); i++) state_dist[i] = 0.0f; + + int step_in_episode = 0; + + /* ---- D1: Count-bonus state (lane 0 only uses, but all lanes store to avoid branches) ---- */ + int action_counts[NUM_ACTIONS]; + int total_action_count = 0; + for (int i = 0; i < NUM_ACTIONS; i++) action_counts[i] = 0; + + /* ---- EMA reward normalizer (lane 0 only) ---- */ + float ema_mean = 0.0f; + float ema_var = 1.0f; + int ema_init = 0; + + /* ---- Main episode loop ---- */ + for (int t = 0; t < L; t++) { + int global_bar = ep_start + t; + int out_off = episode_id * L + t; + + /* Handle out-of-data: all lanes write zeros cooperatively */ + if (tid_valid && global_bar >= total_bars - 1) { + /* Cooperative state zero-write */ + int out_base = out_off * STATE_DIM; + for (int i = lane_id; i < STATE_DIM; i += 32) + out_states[out_base + i] = 0.0f; + if (lane_id == 0) { + out_actions[out_off] = 0; + out_rewards[out_off] = 0.0f; + out_done[out_off] = 1; + out_target_q[out_off] = 0.0f; + out_td_error[out_off] = 0.0f; + } + } + int skip_data = (!tid_valid || global_bar >= total_bars - 1) ? 1 : 0; + + /* ---- Step 1: Scatter market features into distributed layout ---- */ + float current_close = 0.0f, next_close = 0.0f; + float current_close_raw = 0.0f, next_close_raw = 0.0f; + float price = 1.0f, current_value = 0.0f, current_norm = 0.0f; + float max_pos_norm = 1.0f, pos_norm = 0.0f; + + if (!skip_data) { + /* Read market features in strided layout: + * lane k reads elements k, k+32, k+64, ... */ + int mf_off = global_bar * MARKET_DIM; + for (int i = lane_id; i < MARKET_DIM; i += 32) { + state_dist[i / 32] = market_features[mf_off + i]; + } + + /* ---- Step 2: Compute portfolio features (all lanes read, needed for state) ---- */ + int t_off = global_bar * 4; + current_close = targets[t_off + 0]; + next_close = targets[t_off + 1]; + current_close_raw = targets[t_off + 2]; + next_close_raw = targets[t_off + 3]; + + price = (current_close_raw != 0.0f) ? current_close_raw : current_close; + if (price <= 0.0f) price = 1.0f; + + current_value = cash + position * price; + current_norm = current_value / initial_cap; + max_pos_norm = (price > 0.0f) ? initial_cap / price : 1.0f; + pos_norm = position / max_pos_norm; + + /* Portfolio features: element idx is MARKET_DIM+k. + * In strided layout: element i lives in lane (i % 32) at slot (i / 32). + * Only the owning lane writes each element. */ + { + int idx0 = MARKET_DIM + 0; + if (lane_id == (idx0 & 31)) + state_dist[idx0 >> 5] = current_norm; + + int idx1 = MARKET_DIM + 1; + if (lane_id == (idx1 & 31)) + state_dist[idx1 >> 5] = pos_norm; + + int idx2 = MARKET_DIM + 2; + if (lane_id == (idx2 & 31)) + state_dist[idx2 >> 5] = spread; + } + + /* Zero-pad remaining dims (OFI features + tensor-core alignment padding). + * Only owning lanes write. */ + for (int i = MARKET_DIM + PORTFOLIO_DIM; i < STATE_DIM; i++) { + if (lane_id == (i & 31)) + state_dist[i >> 5] = 0.0f; + } + + /* ---- Step 3: Write state to output (all lanes cooperate) ---- */ + int out_base = out_off * STATE_DIM; + for (int i = lane_id; i < STATE_DIM; i += 32) { + out_states[out_base + i] = state_dist[i / 32]; + } + } + + /* ---- Step 4: Online Q-network forward (warp-cooperative) ---- + * All lanes participate in cooperative loads + __syncthreads(). */ + if (use_distributional && num_atoms > 1) { + q_forward_distributional_warp_shmem( + state_dist, + on_w_s1, on_b_s1, on_w_s2, on_b_s2, + on_w_v1, on_b_v1, on_w_v2, on_b_v2, + on_w_a1, on_b_a1, on_w_a2, on_b_a2, + rms_s0_gamma, rms_s1_gamma, rms_v_gamma, rms_a_gamma, + scratch1_dist, scratch2_dist, + q_values, + num_atoms, v_min, v_max, + use_noisy_nets, noisy_sigma_init, use_rmsnorm, + &rng, lane_id, + shmem_weights, shmem_bias + ); + } else if (use_noisy_nets) { + q_forward_dueling_noisy_warp_shmem( + state_dist, + on_w_s1, on_b_s1, on_w_s2, on_b_s2, + on_w_v1, on_b_v1, on_w_v2, on_b_v2, + on_w_a1, on_b_a1, on_w_a2, on_b_a2, + scratch1_dist, scratch2_dist, + q_values, + noisy_sigma_init, &rng, + lane_id, + shmem_weights, shmem_bias + ); + } else { + q_forward_dueling_warp_shmem( + state_dist, + on_w_s1, on_b_s1, on_w_s2, on_b_s2, + on_w_v1, on_b_v1, on_w_v2, on_b_v2, + on_w_a1, on_b_a1, on_w_a2, on_b_a2, + scratch1_dist, scratch2_dist, + q_values, + lane_id, + shmem_weights, shmem_bias + ); + } + + /* ---- Steps 4b through 10: simulation logic (lane 0 only) ---- + * All lanes have identical q_values from the warp-cooperative forward. + * Portfolio simulation, barriers, diversity, curiosity run on lane 0. */ + int action_idx = 0; + float online_q_selected = 0.0f; + float curiosity_reward = 0.0f; + float next_price = 1.0f; + float next_value = 0.0f; + float next_norm = 0.0f; + int barrier_label = 0; + float div_penalty = 0.0f; + + if (!skip_data) { + /* ---- Step 4b: Q-value clipping (D2) — all lanes, identical data ---- */ + for (int i = 0; i < NUM_ACTIONS; i++) { + if (q_values[i] < q_clip_min) q_values[i] = q_clip_min; + if (q_values[i] > q_clip_max) q_values[i] = q_clip_max; + } + + /* ---- Step 5: Epsilon-greedy + count-bonus action selection (lane 0) ---- */ + if (lane_id == 0) { + float r = gpu_random(&rng); + if (r < epsilon) { + action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS); + if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1; + } else { + if (count_bonus_coefficient > 0.0f && total_action_count > 0) { + float log_n = logf((float)total_action_count); + float q_bonus[NUM_ACTIONS]; + for (int i = 0; i < NUM_ACTIONS; i++) { + float bonus = count_bonus_coefficient + * sqrtf(log_n / (1.0f + (float)action_counts[i])); + q_bonus[i] = q_values[i] + bonus; + } + action_idx = argmax_q(q_bonus); + } else { + action_idx = argmax_q(q_values); + } + } + + /* D1: Update count state */ + action_counts[action_idx]++; + total_action_count++; + } + + /* Broadcast action from lane 0 to all lanes */ + action_idx = __shfl_sync(0xFFFFFFFF, action_idx, 0); + + if (lane_id == 0) { + out_actions[out_off] = action_idx; + online_q_selected = q_values[action_idx]; + + /* ---- Step 6: Portfolio simulation ---- */ + float target_exposure = action_to_exposure(action_idx); + float target_position = target_exposure * max_position; + float tx_rate = action_to_tx_cost(action_idx) * tx_cost_multiplier; + + int is_reversal = (position > 0.0f && target_position < 0.0f) || + (position < 0.0f && target_position > 0.0f); + + if (is_reversal) { + float close_cash = position * price; + float close_cost = fabsf(position) * price * tx_rate; + cash += close_cash - close_cost; + cum_costs += close_cost; + + float reserve = (reserve_pct > 0.0f) ? current_value * (reserve_pct / 100.0f) : 0.0f; + float affordable = fmaxf(cash - reserve, 0.0f); + float max_contracts = (price > 0.0f) ? affordable / (price * (1.0f + tx_rate)) : 0.0f; + max_contracts = floorf(max_contracts); + float actual = fminf(max_contracts, fabsf(target_position)); + + if (actual > 0.0f) { + float new_pos = (target_position > 0.0f) ? actual : -actual; + float open_cost = actual * price * tx_rate; + cash -= new_pos * price + open_cost; + cum_costs += open_cost; + position = new_pos; + entry_price = price; + } else { + position = 0.0f; + entry_price = 0.0f; + } + } else { + float delta = target_position - position; + if (fabsf(delta) > 0.0f) { + float trade_cost = fabsf(delta) * price * tx_rate; + cum_costs += trade_cost; + cash -= trade_cost; + + if (delta > 0.0f && reserve_pct > 0.0f) { + float pv = cash + position * price; + float reserve = pv * (reserve_pct / 100.0f); + float buy_cost = delta * price; + if (cash - buy_cost < reserve) { + float affordable = fmaxf(cash - reserve, 0.0f); + delta = fminf(delta, (price > 0.0f) ? floorf(affordable / price) : 0.0f); + } + } + + if (delta > 0.0f) { + entry_price = price; + } else if (target_position == 0.0f) { + entry_price = 0.0f; + } + cash -= delta * price; + position = position + delta; + } + } + last_price = price; + + /* ---- Step 7: Barrier tracking ---- */ + float old_barrier_entry = barrier_st[0]; + if (entry_price > 0.0f && old_barrier_entry <= 0.0f) { + barrier_init(barrier_st, barrier_config, entry_price, global_bar); + } + barrier_label = barrier_check(barrier_st, price, global_bar, position); + if (barrier_label != 0) { + barrier_reset(barrier_st); + } + + /* ---- Step 8: Diversity penalty ---- */ + div_penalty = diversity_entropy(div_window, div_meta, action_idx); + + /* ---- Step 9: Mark-to-market -> next_state ---- */ + next_price = (next_close_raw != 0.0f) ? next_close_raw : next_close; + if (next_price <= 0.0f) next_price = price; + next_value = cash + position * next_price; + next_norm = next_value / initial_cap; + } /* end lane 0 simulation block */ + } /* end if (!skip_data) before curiosity */ + + /* ---- Step 10: Curiosity inference ---- + * Requires gathering distributed state to lane 0 via warp shuffles. + * All lanes must participate in __shfl_sync even if only lane 0 uses result. */ + + /* Gather first 32 state elements (slot 0 of each lane) to all lanes. + * Lane 0 uses these for curiosity input. Cost: 32 shuffles per timestep. */ + float state_elem_0_to_31[32]; + for (int i = 0; i < 32; i++) { + state_elem_0_to_31[i] = __shfl_sync(0xFFFFFFFF, state_dist[0], i); + } + /* Gather elements 32..STATE_DIM-1 (slot 1 of lanes 0..STATE_DIM-33). + * STATE_DIM=48, so elements 32..47 are in slot 1 of lanes 0..15. */ + float state_elem_32_plus[DIST_SIZE(STATE_DIM) > 1 ? (STATE_DIM - 32) : 1]; + if (DIST_SIZE(STATE_DIM) > 1) { + for (int i = 0; i < STATE_DIM - 32; i++) { + state_elem_32_plus[i] = __shfl_sync(0xFFFFFFFF, state_dist[1], i); + } + } + + if (!skip_data && lane_id == 0 && cur_w1 != 0) { + /* Reconstruct full state and next_state for curiosity */ + float state_full[STATE_DIM]; + for (int i = 0; i < 32 && i < STATE_DIM; i++) + state_full[i] = state_elem_0_to_31[i]; + for (int i = 32; i < STATE_DIM; i++) + state_full[i] = state_elem_32_plus[i - 32]; + + /* Build next_state on lane 0 */ + float next_state_full[STATE_DIM]; + int next_bar = global_bar + 1; + if (next_bar < total_bars) { + int nmf_off = next_bar * MARKET_DIM; + for (int i = 0; i < MARKET_DIM; i++) + next_state_full[i] = market_features[nmf_off + i]; + } else { + for (int i = 0; i < MARKET_DIM; i++) + next_state_full[i] = state_full[i]; + } + float next_max_pos_norm = (next_price > 0.0f) ? initial_cap / next_price : 1.0f; + next_state_full[MARKET_DIM + 0] = next_norm; + next_state_full[MARKET_DIM + 1] = position / next_max_pos_norm; + next_state_full[MARKET_DIM + 2] = spread; + for (int i = MARKET_DIM + PORTFOLIO_DIM; i < STATE_DIM; i++) + next_state_full[i] = 0.0f; + + float cur_scratch[CUR_HIDDEN]; + curiosity_reward = curiosity_inference( + state_full, next_state_full, action_idx, + cur_w1, cur_b1, cur_w2, cur_b2, + cur_scratch, curiosity_max_reward + ); + } + + /* ---- Build next_state distributed for target network forward ---- + * All lanes cooperate to scatter next_state features. */ + float next_state_dist[DIST_SIZE(STATE_DIM)]; + for (int i = 0; i < DIST_SIZE(STATE_DIM); i++) next_state_dist[i] = 0.0f; + + if (!skip_data) { + /* Broadcast portfolio simulation results from lane 0 */ + next_price = __shfl_sync(0xFFFFFFFF, next_price, 0); + next_norm = __shfl_sync(0xFFFFFFFF, next_norm, 0); + position = __shfl_sync(0xFFFFFFFF, position, 0); + + int next_bar = global_bar + 1; + if (next_bar < total_bars) { + int nmf_off = next_bar * MARKET_DIM; + for (int i = lane_id; i < MARKET_DIM; i += 32) { + next_state_dist[i / 32] = market_features[nmf_off + i]; + } + } else { + /* Copy current state's market features */ + for (int i = lane_id; i < MARKET_DIM; i += 32) { + next_state_dist[i / 32] = state_dist[i / 32]; + } + } + + /* Next portfolio features */ + float next_max_pos_norm = (next_price > 0.0f) ? initial_cap / next_price : 1.0f; + { + int idx0 = MARKET_DIM + 0; + if (lane_id == (idx0 & 31)) + next_state_dist[idx0 >> 5] = next_norm; + + int idx1 = MARKET_DIM + 1; + if (lane_id == (idx1 & 31)) + next_state_dist[idx1 >> 5] = position / next_max_pos_norm; + + int idx2 = MARKET_DIM + 2; + if (lane_id == (idx2 & 31)) + next_state_dist[idx2 >> 5] = spread; + } + + /* Zero-pad remaining dims */ + for (int i = MARKET_DIM + PORTFOLIO_DIM; i < STATE_DIM; i++) { + if (lane_id == (i & 31)) + next_state_dist[i >> 5] = 0.0f; + } + } + + /* ---- Step 11: Target Q-network forward (warp-cooperative, always clean) ---- */ + if (use_distributional && num_atoms > 1) { + q_forward_distributional_warp_shmem( + next_state_dist, + tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2, + tg_w_v1, tg_b_v1, tg_w_v2, tg_b_v2, + tg_w_a1, tg_b_a1, tg_w_a2, tg_b_a2, + rms_s0_gamma, rms_s1_gamma, rms_v_gamma, rms_a_gamma, + scratch1_dist, scratch2_dist, + tgt_q_values, + num_atoms, v_min, v_max, + 0, 0.0f, use_rmsnorm, + &rng, lane_id, + shmem_weights, shmem_bias + ); + } else { + q_forward_dueling_warp_shmem( + next_state_dist, + tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2, + tg_w_v1, tg_b_v1, tg_w_v2, tg_b_v2, + tg_w_a1, tg_b_a1, tg_w_a2, tg_b_a2, + scratch1_dist, scratch2_dist, + tgt_q_values, + lane_id, + shmem_weights, shmem_bias + ); + } + + /* D2: Clip target Q-values (all lanes, identical) */ + for (int i = 0; i < NUM_ACTIONS; i++) { + if (tgt_q_values[i] < q_clip_min) tgt_q_values[i] = q_clip_min; + if (tgt_q_values[i] > q_clip_max) tgt_q_values[i] = q_clip_max; + } + float max_target_q = tgt_q_values[0]; + for (int i = 1; i < NUM_ACTIONS; i++) { + if (tgt_q_values[i] > max_target_q) + max_target_q = tgt_q_values[i]; + } + + /* ---- Steps 12-16: reward, done, TD error, output, reset (lane 0 only) ---- */ + if (!skip_data && lane_id == 0) { + /* ---- Step 12: Combined reward ---- */ + float pnl_reward = 0.0f; + if (current_norm > 0.0f) { + pnl_reward = (next_norm - current_norm) / current_norm; + } + + if (fabsf(position) < 0.001f) { + pnl_reward += hold_reward; + } + + float abs_pos = fabsf(pos_norm); + if (abs_pos > 0.8f) { + pnl_reward -= (abs_pos - 0.8f) * 5.0f * risk_weight; + } + + float barrier_mult = 1.0f; + if (barrier_label != 0) { + barrier_mult = 1.0f + barrier_scale * (float)barrier_label; + } + + float combined_reward = pnl_reward * barrier_mult + + diversity_scale * div_penalty + + curiosity_scale * curiosity_reward; + + /* ---- EMA reward normalization ---- */ + float raw_reward = combined_reward; + if (ema_init) { + float std = sqrtf(ema_var); + if (std > 1e-8f) + combined_reward = (combined_reward - ema_mean) / std; + combined_reward = fmaxf(-3.0f, fminf(3.0f, combined_reward)); + } + if (!ema_init) { + ema_mean = raw_reward; + ema_var = 1.0f; + ema_init = 1; + } else { + ema_mean = reward_norm_alpha * raw_reward + + (1.0f - reward_norm_alpha) * ema_mean; + float diff = raw_reward - ema_mean; + ema_var = reward_norm_alpha * diff * diff + + (1.0f - reward_norm_alpha) * ema_var; + } + + /* ---- Step 13: Episode done check ---- */ + step_in_episode++; + int next_bar = global_bar + 1; + int time_done = (step_in_episode >= episode_length) ? 1 : 0; + int barrier_done = (barrier_label != 0) ? 1 : 0; + int data_done = (next_bar >= total_bars) ? 1 : 0; + int done = (time_done || barrier_done || data_done) ? 1 : 0; + + /* ---- Step 14: TD error with Huber transformation (D3) ---- */ + float td_target = combined_reward + gamma * max_target_q * (1.0f - (float)done); + float td_delta = td_target - online_q_selected; + float td_abs = fabsf(td_delta); + float td_error; + if (huber_kappa > 0.0f) { + if (td_abs <= huber_kappa) { + td_error = 0.5f * td_delta * td_delta; + } else { + td_error = huber_kappa * (td_abs - 0.5f * huber_kappa); + } + } else { + td_error = td_abs; + } + + /* ---- Step 15: Write experience ---- */ + out_rewards[out_off] = combined_reward; + out_done[out_off] = done; + out_target_q[out_off] = max_target_q; + out_td_error[out_off] = td_error; + + /* ---- Step 16: Episode reset if done ---- */ + if (done) { + cash = initial_cap; + position = 0.0f; + entry_price = 0.0f; + cum_costs = 0.0f; + last_price = 0.0f; + step_in_episode = 0; + barrier_reset(barrier_st); + for (int i = 0; i < DIVERSITY_WINDOW; i++) + div_window[i] = 0; + div_meta[0] = 0; + div_meta[1] = 0; + for (int i = 0; i < NUM_ACTIONS; i++) action_counts[i] = 0; + total_action_count = 0; + ema_mean = 0.0f; + ema_var = 1.0f; + ema_init = 0; + } + } /* end lane 0 steps 12-16 */ + + /* ---- Broadcast simulation state from lane 0 to all lanes ---- + * Other lanes need updated cash/position/entry_price/step_in_episode + * for next iteration's portfolio features and state construction. */ + cash = __shfl_sync(0xFFFFFFFF, cash, 0); + position = __shfl_sync(0xFFFFFFFF, position, 0); + entry_price = __shfl_sync(0xFFFFFFFF, entry_price, 0); + initial_cap = __shfl_sync(0xFFFFFFFF, initial_cap, 0); + spread = __shfl_sync(0xFFFFFFFF, spread, 0); + last_price = __shfl_sync(0xFFFFFFFF, last_price, 0); + reserve_pct = __shfl_sync(0xFFFFFFFF, reserve_pct, 0); + cum_costs = __shfl_sync(0xFFFFFFFF, cum_costs, 0); + step_in_episode = __shfl_sync(0xFFFFFFFF, step_in_episode, 0); + } /* end timestep loop */ + + /* ---- Write back per-episode state (lane 0 only) ---- */ + if (tid_valid && lane_id == 0) { + portfolio_states[ps_off + 0] = cash; + portfolio_states[ps_off + 1] = position; + portfolio_states[ps_off + 2] = entry_price; + portfolio_states[ps_off + 3] = initial_cap; + portfolio_states[ps_off + 4] = spread; + portfolio_states[ps_off + 5] = last_price; + portfolio_states[ps_off + 6] = reserve_pct; + portfolio_states[ps_off + 7] = cum_costs; + + for (int i = 0; i < BARRIER_STATE_SIZE; i++) + barrier_states[bs_off + i] = barrier_st[i]; + + for (int i = 0; i < DIVERSITY_WINDOW; i++) + diversity_windows[dw_off + i] = div_window[i]; + + diversity_metas[dm_off + 0] = div_meta[0]; + diversity_metas[dm_off + 1] = div_meta[1]; + + rng_states[episode_id] = rng; + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 3ea51c4ea..c178c2d5c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -197,6 +197,10 @@ pub struct GpuExperienceCollector { state_dim: usize, /// Network hidden dims — needed for shared memory tile sizing at launch time. network_dims: (usize, usize, usize, usize), + /// Warp-cooperative kernel for sm_90+ (H100). None if not available. + warp_kernel_func: Option, + /// Whether to use warp-cooperative kernel (sm_90+ detected). + use_warp_kernel: bool, /// Number of episodes buffers were allocated for (from config, not MAX constant). alloc_episodes: usize, /// Number of timesteps buffers were allocated for (from config, not MAX constant). @@ -295,7 +299,8 @@ impl GpuExperienceCollector { let full_source = format!("{common_src}\n{dim_overrides}\n{kernel_src}"); info!( "GPU experience collector: compiling kernel with dims state={state_dim} market={market_dim} \ - shared=[{shared_h1},{shared_h2}] value={value_h} adv={adv_h} atoms_max={num_atoms_max}" + shared=[{shared_h1},{shared_h2}] value={value_h} adv={adv_h} atoms_max={num_atoms_max} \ + shmem_max_in={shmem_max_in_dim}" ); let ptx: Ptx = cudarc::nvrtc::compile_ptx(&full_source).map_err(|e| { MLError::ModelError(format!( @@ -315,6 +320,50 @@ impl GpuExperienceCollector { })?; info!("GPU experience collector: CUDA kernel compiled and loaded"); + // Detect compute capability for kernel variant selection. + // sm_90+ (H100/H200) uses warp-cooperative kernel for lower register pressure. + let sm_major = { + use cudarc::driver::sys::CUdevice_attribute; + context + .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR) + .map_err(|e| { + MLError::ModelError(format!("Failed to query compute capability: {e}")) + })? + }; + let use_warp_kernel = sm_major >= 9; + + let warp_kernel_func = if use_warp_kernel { + match module.load_function("dqn_full_experience_kernel_warp") { + Ok(f) => { + info!("GPU experience collector: warp-cooperative kernel loaded (sm_{sm_major}0+)"); + Some(f) + } + Err(e) => { + tracing::warn!( + "Warp kernel load failed on sm_{sm_major}0, falling back to standard: {e}" + ); + None + } + } + } else { + info!("GPU experience collector: using standard kernel (sm_{sm_major}0 < sm_90)"); + None + }; + + // Increase per-thread stack size for large hidden dimensions. + // With SHARED_H1=512 (H100), scratch arrays consume ~5.5KB per thread, + // exceeding the default 1KB CUDA stack limit. + { + use cudarc::driver::sys::{cuCtxSetLimit, CUlimit, CUresult}; + let stack_bytes = 16384_usize; // 16KB — enough for 5.5KB scratch + noise + call frames + let result = unsafe { cuCtxSetLimit(CUlimit::CU_LIMIT_STACK_SIZE, stack_bytes) }; + if result != CUresult::CUDA_SUCCESS { + tracing::warn!("cuCtxSetLimit(STACK_SIZE, {stack_bytes}) failed: {result:?}"); + } else { + tracing::info!("CUDA stack size set to {stack_bytes} bytes"); + } + } + // ---- Step 2: Extract weights ---- let online_weights = extract_dueling_weights(online_vars, &stream)?; let target_weights = extract_dueling_weights(target_vars, &stream)?; @@ -469,6 +518,8 @@ impl GpuExperienceCollector { kernel_func, state_dim, network_dims, + warp_kernel_func, + use_warp_kernel, alloc_episodes, alloc_timesteps, online_weights, @@ -561,26 +612,42 @@ impl GpuExperienceCollector { * std::mem::size_of::() as u32; let n = n_episodes as u32; - // Force block size to 256 for shared memory cooperative loading - let block_x = 256_u32; - let grid_x = (n + block_x - 1) / block_x; - let grid_dim = (grid_x, 1, 1); - let block_dim = (block_x, 1, 1); - let launch_config = LaunchConfig { - grid_dim, - block_dim, - shared_mem_bytes: shmem_bytes, - }; + + // Select kernel variant: warp-cooperative for sm_90+ (all modes including C51). + let use_warp = self.use_warp_kernel + && self.warp_kernel_func.is_some(); + + let (launch_config, active_kernel) = + if let (true, Some(warp_fn)) = (use_warp, &self.warp_kernel_func) { + // Warp kernel: 1 warp (32 threads) per block, 1 block per episode. + let cfg = LaunchConfig { + grid_dim: (n, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: shmem_bytes, + }; + (cfg, warp_fn) + } else { + // Standard kernel: 256 threads per block, 1 thread per episode. + let block_x = 256_u32; + let grid_x = (n + block_x - 1) / block_x; + let cfg = LaunchConfig { + grid_dim: (grid_x, 1, 1), + block_dim: (block_x, 1, 1), + shared_mem_bytes: shmem_bytes, + }; + (cfg, &self.kernel_func) + }; debug!( n_episodes, timesteps, - grid_x = grid_dim.0, - block_x = block_dim.0, + grid_x = launch_config.grid_dim.0, + block_x = launch_config.block_dim.0, shared_mem_kb = shmem_bytes / 1024, epsilon = config.epsilon, gamma = config.gamma, - "Launching dqn_full_experience_kernel (shmem weight tiling)" + warp_kernel = use_warp, + "Launching dqn experience kernel" ); // ---- Step 5: Launch kernel ---- @@ -590,9 +657,10 @@ impl GpuExperienceCollector { // Safety: kernel parameter order matches dqn_experience_kernel.cu lines 394-468 exactly. // All buffer sizes are pre-validated (MAX_EPISODES, MAX_TIMESTEPS). // All CudaSlice lifetimes are valid (owned by self). + // Both standard and warp kernels share the same parameter signature. unsafe { self.stream - .launch_builder(&self.kernel_func) + .launch_builder(active_kernel) // Market data .arg(market_features_buf) .arg(targets_buf) @@ -680,7 +748,7 @@ impl GpuExperienceCollector { .arg(&mut self.td_error_out) .launch(launch_config) .map_err(|e| { - MLError::ModelError(format!("dqn_full_experience_kernel launch failed: {e}")) + MLError::ModelError(format!("dqn experience kernel launch failed: {e}")) })?; }