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

refactor: convert #define constants to kernel parameters in .cu files

Browse Source 
Replace #error guards in common_device_functions.cuh with safe defaults
(STATE_DIM=72, MARKET_DIM=42, PORTFOLIO_DIM=8). Add MAX_STATE_DIM=128
for safe stack array sizing.

Add runtime parameters to kernels that used compile-time #define:
- attention_kernel.cu: state_dim, num_heads
- attention_backward_kernel.cu: state_dim, num_heads
- backtest_forward_ppo_kernel.cu: state_dim, num_actions
- backtest_metrics_kernel.cu: num_actions, order_actions, urgency_actions
- dqn_utility_kernels.cu: state_dim
- her_relabel_kernel.cu: state_dim
- iql_value_kernel.cu: state_dim (all 3 kernel functions)
- iqn_dual_head_kernel.cu: state_dim
- monitoring_kernel.cu: num_actions, order_actions, urgency_actions

This enables single-cubin-per-kernel precompilation — no #define
injection needed at runtime.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-26 00:18:18 +01:00
parent a2bf5210ec
commit 82b595f633
10 changed files with 155 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crates/ml/src/cuda_pipeline/attention_backward_kernel.cu
View File

@@ -54,11 +54,13 @@ extern "C" __global__ void attention_backward_kernel(
const float* __restrict__ params,
float* d_input,
float* d_params,
int B
int B,
int state_dim, /* runtime state dimension (replaces ATTN_STATE_DIM) */
int num_heads /* runtime num_heads (replaces ATTN_NUM_HEADS) */
) {
const int D = ATTN_STATE_DIM;
const int H = ATTN_NUM_HEADS;
const int Dh = ATTN_HEAD_DIM;
const int D = state_dim;
const int H = num_heads;
const int Dh = D / H;
int sample = blockIdx.x;
if (sample >= B) return;

10
crates/ml/src/cuda_pipeline/attention_kernel.cu
View File

@@ -63,11 +63,13 @@ extern "C" __global__ void multihead_feature_attention(
const float* __restrict__ states,
const float* __restrict__ params,
float* output,
int B
int B,
int state_dim, /* runtime state dimension (replaces ATTN_STATE_DIM) */
int num_heads /* runtime num_heads (replaces ATTN_NUM_HEADS) */
) {
const int D = ATTN_STATE_DIM;
const int H = ATTN_NUM_HEADS;
const int Dh = ATTN_HEAD_DIM;
const int D = state_dim;
const int H = num_heads;
const int Dh = D / H;
int sample = blockIdx.x;
if (sample >= B) return;

44
crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu
View File

@@ -1,7 +1,7 @@
/**
* Pure-CUDA PPO forward for backtest evaluation.
*
* Requires common_device_functions.cuh prepended via NVRTC source concatenation.
* Requires common_device_functions.cuh prepended at build time.
* Launch config: grid=(ceil(N/32), 1, 1), block=(32, 1, 1).
* 32 independent windows per block — each thread handles its own window.
*
@@ -9,74 +9,76 @@
*
* H100 occupancy: 32 threads/block × multiple blocks/SM vs old 1 thread/block
* that left 97% of each SM idle.
*
* state_dim and num_actions are runtime parameters (not compile-time #define).
*/
/* Override NUM_ACTIONS for PPO: use full 45-action factored space */
#undef NUM_ACTIONS
#define NUM_ACTIONS 45 /* PPO_NUM_ACTIONS */
/* Actor layer sizes — overridable via NVRTC #define injection */
/* Actor layer sizes — compile-time constants (fixed architecture) */
#ifndef ACTOR_H1
#define ACTOR_H1 128
#endif
#ifndef ACTOR_H2
#define ACTOR_H2 64
#endif
/* Max num_actions for stack-allocated arrays */
#define MAX_NUM_ACTIONS 45
#define N_EXPOSURE_BINS 5
#define ACTIONS_PER_BIN 9 /* 3 order × 3 urgency */
extern "C" __global__ void backtest_forward_ppo_kernel(
const float* __restrict__ states, /* [N, STATE_DIM] */
const float* __restrict__ states, /* [N, state_dim] */
/* Actor weights (6 pointers) */
const float* __restrict__ pw1, /* [ACTOR_H1, STATE_DIM] = [128, S_DIM] */
const float* __restrict__ pw1, /* [ACTOR_H1, state_dim] = [128, S_DIM] */
const float* __restrict__ pb1, /* [ACTOR_H1] = [128] */
const float* __restrict__ pw2, /* [ACTOR_H2, ACTOR_H1] = [64, 128] */
const float* __restrict__ pb2, /* [ACTOR_H2] = [64] */
const float* __restrict__ pw3, /* [NUM_ACTIONS, ACTOR_H2] = [45, 64] */
const float* __restrict__ pb3, /* [NUM_ACTIONS] = [45] */
const float* __restrict__ pw3, /* [num_actions, ACTOR_H2] = [45, 64] */
const float* __restrict__ pb3, /* [num_actions] = [45] */
/* Output */
int* __restrict__ out_actions, /* [N] greedy exposure action (0-4) */
int N
int N,
int state_dim, /* runtime state dimension */
int num_actions /* runtime num_actions (45 for PPO) */
) {
int window_id = blockIdx.x * blockDim.x + threadIdx.x;
if (window_id >= N) return;
/* ---- Load state vector ---- */
float state[STATE_DIM];
int base = window_id * STATE_DIM;
for (int i = 0; i < STATE_DIM; i++) {
float state[MAX_STATE_DIM];
int base = window_id * state_dim;
for (int i = 0; i < state_dim; i++) {
state[i] = states[base + i];
}
/* ---- Actor forward: state → logits ---- */
float h1[ACTOR_H1];
float h2[ACTOR_H2];
float logits[NUM_ACTIONS];
float logits[MAX_NUM_ACTIONS];
/* Hidden layer 1: state → h1 with LeakyReLU */
matvec_leaky_relu(pw1, pb1, state, h1, STATE_DIM, ACTOR_H1, 1);
matvec_leaky_relu(pw1, pb1, state, h1, state_dim, ACTOR_H1, 1);
/* Hidden layer 2: h1 → h2 with LeakyReLU */
matvec_leaky_relu(pw2, pb2, h1, h2, ACTOR_H1, ACTOR_H2, 1);
/* Output layer: h2 → logits (no activation) */
matvec_leaky_relu(pw3, pb3, h2, logits, ACTOR_H2, NUM_ACTIONS, 0);
matvec_leaky_relu(pw3, pb3, h2, logits, ACTOR_H2, num_actions, 0);
/* ---- Stable softmax: logits → probabilities ---- */
float max_logit = logits[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
for (int i = 1; i < num_actions; i++) {
if (logits[i] > max_logit) max_logit = logits[i];
}
float probs[NUM_ACTIONS];
float probs[MAX_NUM_ACTIONS];
float sum_exp = 0.0f;
for (int i = 0; i < NUM_ACTIONS; i++) {
for (int i = 0; i < num_actions; i++) {
probs[i] = expf(logits[i] - max_logit);
sum_exp += probs[i];
}
float inv_sum = 1.0f / fmaxf(sum_exp, 1e-8f);
for (int i = 0; i < NUM_ACTIONS; i++) {
for (int i = 0; i < num_actions; i++) {
probs[i] *= inv_sum;
}

13
crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu
View File

@@ -25,7 +25,10 @@ extern "C" __global__ void compute_backtest_metrics(
float* metrics_out, // [n_windows * 14]
int n_windows,
int max_len,
float annualization_factor // sqrt(252) for daily
float annualization_factor, // sqrt(252) for daily
int num_actions, // runtime: DQN_NUM_ACTIONS (9)
int order_actions, // runtime: DQN_ORDER_ACTIONS (3)
int urgency_actions // runtime: DQN_URGENCY_ACTIONS (3)
) {
int w = blockIdx.x;
if (w >= n_windows) return;
@@ -96,13 +99,13 @@ extern "C" __global__ void compute_backtest_metrics(
// Factored action = exposure * (b1*b2) + order * b2 + urgency.
// exposure_idx: 0-3=sell, 4=flat, 5-8=buy (for 9-action space)
int act = actions_history[base + i];
int exp_idx = act / (DQN_ORDER_ACTIONS * DQN_URGENCY_ACTIONS);
if (exp_idx < DQN_NUM_ACTIONS / 2) local_sells++;
else if (exp_idx == DQN_NUM_ACTIONS / 2) local_holds++;
int exp_idx = act / (order_actions * urgency_actions);
if (exp_idx < num_actions / 2) local_sells++;
else if (exp_idx == num_actions / 2) local_holds++;
else local_buys++;
// Unique action tracking: set bit per exposure index (0..8)
if (exp_idx >= 0 && exp_idx < DQN_NUM_ACTIONS)
if (exp_idx >= 0 && exp_idx < num_actions)
local_action_mask |= (1 << exp_idx);
// Boundary-aware trade tracking

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

@@ -27,19 +27,29 @@ __device__ __forceinline__ __nv_bfloat16 f32_to_bf16(float x) {
/* ------------------------------------------------------------------ */
/* Constants */
/* ------------------------------------------------------------------ */
/* State layout — MUST be injected via NVRTC dim_overrides (no fallback defaults).
/* State layout — compile-time defaults for build.rs precompilation.
* STATE_DIM = tensor-core-aligned state width (e.g. 48 without OFI, 56 with OFI).
* MARKET_DIM = raw market feature count (42: 40 base + 2 regime).
* PORTFOLIO_DIM = portfolio feature count (cash, position_size, unrealized_pnl).
* OFI_DIM = order flow imbalance features from MBP-10 (0 when disabled, 8 when enabled). */
* OFI_DIM = order flow imbalance features from MBP-10 (0 when disabled, 8 when enabled).
*
* These are compile-time defaults used for stack array sizing in kernels.
* Kernels that need variable dimensions accept them as runtime parameters
* (e.g. state_dim, market_dim) and use these only for max array sizes.
* MAX_STATE_DIM provides a safe upper bound for stack-allocated arrays. */
#ifndef STATE_DIM
#error "STATE_DIM must be defined via dim_overrides before including this header"
#define STATE_DIM 72
#endif
#ifndef MARKET_DIM
#error "MARKET_DIM must be defined via dim_overrides before including this header"
#define MARKET_DIM 42
#endif
#ifndef PORTFOLIO_DIM
#error "PORTFOLIO_DIM must be defined via dim_overrides before including this header"
#define PORTFOLIO_DIM 8
#endif
/* Maximum state dimension for stack-allocated arrays in kernels.
* Must be >= any actual state_dim used at runtime. */
#ifndef MAX_STATE_DIM
#define MAX_STATE_DIM 128
#endif
#ifndef OFI_DIM
#define OFI_DIM 0 /* Default: no OFI features (safe fallback for kernels that don't use OFI) */

13
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -165,19 +165,20 @@ extern "C" __global__ void dqn_zero_kernel(
extern "C" __global__ void dqn_regime_scale_kernel(
float* __restrict__ td_errors, /* [B] scaled in-place */
const float* __restrict__ states, /* [B, STATE_DIM] */
int batch_size
const float* __restrict__ states, /* [B, state_dim] */
int batch_size,
int state_dim /* runtime state dimension */
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= batch_size) return;
/* Read target regime from first sample (GPU-only, no CPU readback).
* ADX at feature index 40, CUSUM at index 41. */
float target_adx = states[0 * STATE_DIM + 40];
float target_cusum = states[0 * STATE_DIM + 41];
float target_adx = states[0 * state_dim + 40];
float target_cusum = states[0 * state_dim + 41];
float sample_adx = states[i * STATE_DIM + 40];
float sample_cusum = states[i * STATE_DIM + 41];
float sample_adx = states[i * state_dim + 40];
float sample_cusum = states[i * state_dim + 41];
float adx_diff = sample_adx - target_adx;
float cusum_diff = sample_cusum - target_cusum;

29
crates/ml/src/cuda_pipeline/her_relabel_kernel.cu
View File

@@ -27,19 +27,20 @@
extern "C" __global__
void her_relabel_kernel(
const float* __restrict__ states, /* [capacity, STATE_DIM] */
const float* __restrict__ next_states, /* [capacity, STATE_DIM] */
const float* __restrict__ states, /* [capacity, state_dim] */
const float* __restrict__ next_states, /* [capacity, state_dim] */
const int* __restrict__ actions, /* [capacity] */
const float* __restrict__ rewards, /* [capacity] */
const float* __restrict__ dones, /* [capacity] */
const int* __restrict__ source_indices,/* [her_batch_size] */
const int* __restrict__ donor_indices, /* [her_batch_size] */
int buffer_size,
float* __restrict__ out_states, /* [her_batch_size, STATE_DIM] */
float* __restrict__ out_next_states, /* [her_batch_size, STATE_DIM] */
float* __restrict__ out_states, /* [her_batch_size, state_dim] */
float* __restrict__ out_next_states, /* [her_batch_size, state_dim] */
int* __restrict__ out_actions, /* [her_batch_size] */
float* __restrict__ out_rewards, /* [her_batch_size] */
float* __restrict__ out_dones /* [her_batch_size] */
float* __restrict__ out_dones, /* [her_batch_size] */
int state_dim /* runtime state dimension */
)
{
int sample_idx = blockIdx.x;
@@ -49,19 +50,19 @@ void her_relabel_kernel(
int donor_idx = donor_indices[sample_idx];
/* 1. Warp-cooperative coalesced copy of full state and next_state vectors.
* Each lane handles a stride-32 subset of STATE_DIM elements. */
for (int d = lane_id; d < STATE_DIM; d += 32) {
out_states[sample_idx * STATE_DIM + d] = states[src_idx * STATE_DIM + d];
out_next_states[sample_idx * STATE_DIM + d] = next_states[src_idx * STATE_DIM + d];
* Each lane handles a stride-32 subset of state_dim elements. */
for (int d = lane_id; d < state_dim; d += 32) {
out_states[sample_idx * state_dim + d] = states[src_idx * state_dim + d];
out_next_states[sample_idx * state_dim + d] = next_states[src_idx * state_dim + d];
}
__syncwarp();
/* 2. Replace goal portion (first GOAL_DIM elements) with donor's achieved goal.
* The donor's achieved goal is extracted from the donor's next_state. */
for (int d = lane_id; d < GOAL_DIM; d += 32) {
float achieved = next_states[donor_idx * STATE_DIM + d];
out_states[sample_idx * STATE_DIM + d] = achieved;
out_next_states[sample_idx * STATE_DIM + d] = achieved;
float achieved = next_states[donor_idx * state_dim + d];
out_states[sample_idx * state_dim + d] = achieved;
out_next_states[sample_idx * state_dim + d] = achieved;
}
/* 3. Scalar work: copy action/done, recompute sparse reward (lane 0 only). */
@@ -74,8 +75,8 @@ void her_relabel_kernel(
* Donor achieved = next_states[donor_idx, 0..GOAL_DIM] */
float dist_sq = 0.0f;
for (int d = 0; d < GOAL_DIM; d++) {
float src_achieved = next_states[src_idx * STATE_DIM + d];
float donor_achieved = next_states[donor_idx * STATE_DIM + d];
float src_achieved = next_states[src_idx * state_dim + d];
float donor_achieved = next_states[donor_idx * state_dim + d];
float diff = src_achieved - donor_achieved;
dist_sq += diff * diff;
}

108
crates/ml/src/cuda_pipeline/iql_value_kernel.cu
View File

@@ -31,23 +31,27 @@
#define IQL_EXPECTILE_TAU 0.7f
#endif
/* Parameter counts */
#define V_W1_SIZE (VALUE_HIDDEN_DIM * STATE_DIM)
/* Parameter sizes that don't depend on state_dim */
#define V_B1_SIZE (VALUE_HIDDEN_DIM)
#define V_W2_SIZE (VALUE_HIDDEN_DIM * VALUE_HIDDEN_DIM)
#define V_B2_SIZE (VALUE_HIDDEN_DIM)
#define V_W3_SIZE (VALUE_HIDDEN_DIM) /* output layer: 1 × H */
#define V_W3_SIZE (VALUE_HIDDEN_DIM) /* output layer: 1 x H */
#define V_B3_SIZE (1)
#define V_TOTAL_PARAMS (V_W1_SIZE + V_B1_SIZE + V_W2_SIZE + V_B2_SIZE + V_W3_SIZE + V_B3_SIZE)
/* Offsets into the flat parameter buffer */
#define V_OFF_W1 0
#define V_OFF_B1 (V_OFF_W1 + V_W1_SIZE)
#define V_OFF_W2 (V_OFF_B1 + V_B1_SIZE)
#define V_OFF_B2 (V_OFF_W2 + V_W2_SIZE)
#define V_OFF_W3 (V_OFF_B2 + V_B2_SIZE)
#define V_OFF_B3 (V_OFF_W3 + V_W3_SIZE)
/* Runtime offset computation helper — state_dim is a kernel parameter.
* W1 size = VALUE_HIDDEN_DIM * state_dim (varies with state_dim). */
__device__ __forceinline__ void iql_compute_offsets(
int state_dim,
int* off_w1, int* off_b1, int* off_w2, int* off_b2, int* off_w3, int* off_b3
) {
int w1_size = VALUE_HIDDEN_DIM * state_dim;
*off_w1 = 0;
*off_b1 = w1_size;
*off_w2 = *off_b1 + V_B1_SIZE;
*off_b2 = *off_w2 + V_W2_SIZE;
*off_w3 = *off_b2 + V_B2_SIZE;
*off_b3 = *off_w3 + V_W3_SIZE;
}
/* ------------------------------------------------------------------ */
/* SiLU activation: x * sigmoid(x) */
@@ -96,7 +100,8 @@ void iql_forward_loss_kernel(
float* __restrict__ save_pre2,
float* __restrict__ save_h1,
float* __restrict__ save_h2,
int batch_size
int batch_size,
int state_dim /* runtime state dimension */
)
{
int sample = blockIdx.x;
@@ -104,24 +109,27 @@ void iql_forward_loss_kernel(
int lane = threadIdx.x; /* 0..31 */
const float* w1 = params + V_OFF_W1;
const float* b1 = params + V_OFF_B1;
const float* w2 = params + V_OFF_W2;
const float* b2 = params + V_OFF_B2;
const float* w3 = params + V_OFF_W3;
const float* b3 = params + V_OFF_B3;
int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3;
iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3);
const float* x = states + sample * STATE_DIM;
const float* w1 = params + off_w1;
const float* b1 = params + off_b1;
const float* w2 = params + off_w2;
const float* b2 = params + off_b2;
const float* w3 = params + off_w3;
const float* b3 = params + off_b3;
/* ---- Layer 1: Linear(STATE_DIM -> H) + SiLU ---- */
const float* x = states + sample * state_dim;
/* ---- Layer 1: Linear(state_dim -> H) + SiLU ---- */
/* Each lane computes a subset of H outputs via warp reduction */
float* pre1_ptr = save_pre1 + sample * VALUE_HIDDEN_DIM;
float* h1_ptr = save_h1 + sample * VALUE_HIDDEN_DIM;
for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) {
float acc = b1[j];
for (int k = 0; k < STATE_DIM; k++) {
acc += w1[j * STATE_DIM + k] * x[k];
for (int k = 0; k < state_dim; k++) {
acc += w1[j * state_dim + k] * x[k];
}
pre1_ptr[j] = acc;
h1_ptr[j] = silu(acc);
@@ -195,7 +203,8 @@ void iql_backward_kernel(
const float* __restrict__ save_h1,
const float* __restrict__ save_h2,
float* __restrict__ grads,
int batch_size
int batch_size,
int state_dim /* runtime state dimension */
)
{
int sample = blockIdx.x;
@@ -203,18 +212,21 @@ void iql_backward_kernel(
int lane = threadIdx.x;
const float* w1 = params + V_OFF_W1;
const float* w2 = params + V_OFF_W2;
const float* w3 = params + V_OFF_W3;
int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3;
iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3);
float* gw1 = grads + V_OFF_W1;
float* gb1 = grads + V_OFF_B1;
float* gw2 = grads + V_OFF_W2;
float* gb2 = grads + V_OFF_B2;
float* gw3 = grads + V_OFF_W3;
float* gb3 = grads + V_OFF_B3;
const float* w1 = params + off_w1;
const float* w2 = params + off_w2;
const float* w3 = params + off_w3;
const float* x = states + sample * STATE_DIM;
float* gw1 = grads + off_w1;
float* gb1 = grads + off_b1;
float* gw2 = grads + off_w2;
float* gb2 = grads + off_b2;
float* gw3 = grads + off_w3;
float* gb3 = grads + off_b3;
const float* x = states + sample * state_dim;
const float* h1 = save_h1 + sample * VALUE_HIDDEN_DIM;
const float* h2 = save_h2 + sample * VALUE_HIDDEN_DIM;
const float* pre1 = save_pre1 + sample * VALUE_HIDDEN_DIM;
@@ -269,8 +281,8 @@ void iql_backward_kernel(
atomicAdd(&gb1[k], dpre1_k);
/* dL/dw1[k, d] = dpre1_k * x[d] */
for (int d = 0; d < STATE_DIM; d++) {
atomicAdd(&gw1[k * STATE_DIM + d], dpre1_k * x[d]);
for (int d = 0; d < state_dim; d++) {
atomicAdd(&gw1[k * state_dim + d], dpre1_k * x[d]);
}
}
}
@@ -388,7 +400,8 @@ void iql_forward_kernel(
const float* __restrict__ states,
const float* __restrict__ params,
float* __restrict__ v_out,
int batch_size
int batch_size,
int state_dim /* runtime state dimension */
)
{
int sample = blockIdx.x;
@@ -396,14 +409,17 @@ void iql_forward_kernel(
int lane = threadIdx.x;
const float* w1 = params + V_OFF_W1;
const float* b1 = params + V_OFF_B1;
const float* w2 = params + V_OFF_W2;
const float* b2 = params + V_OFF_B2;
const float* w3 = params + V_OFF_W3;
const float* b3 = params + V_OFF_B3;
int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3;
iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3);
const float* x = states + sample * STATE_DIM;
const float* w1 = params + off_w1;
const float* b1 = params + off_b1;
const float* w2 = params + off_w2;
const float* b2 = params + off_b2;
const float* w3 = params + off_w3;
const float* b3 = params + off_b3;
const float* x = states + sample * state_dim;
/* Shared memory for intermediate activations (reused per layer) */
__shared__ float sh1[VALUE_HIDDEN_DIM];
@@ -412,8 +428,8 @@ void iql_forward_kernel(
/* Layer 1: SiLU */
for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) {
float acc = b1[j];
for (int k = 0; k < STATE_DIM; k++) {
acc += w1[j * STATE_DIM + k] * x[k];
for (int k = 0; k < state_dim; k++) {
acc += w1[j * state_dim + k] * x[k];
}
sh1[j] = silu(acc);
}

13
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -749,13 +749,14 @@ void iqn_forward_kernel(
extern "C" __global__
void iqn_trunk_forward_kernel(
const float* __restrict__ states, /* [B, IQN_STATE_DIM] */
const float* __restrict__ w_s1, /* [IQN_SHARED_H1, IQN_STATE_DIM] */
const float* __restrict__ states, /* [B, state_dim] */
const float* __restrict__ w_s1, /* [IQN_SHARED_H1, state_dim] */
const float* __restrict__ b_s1, /* [IQN_SHARED_H1] */
const float* __restrict__ w_s2, /* [IQN_HIDDEN, IQN_SHARED_H1] */
const float* __restrict__ b_s2, /* [IQN_HIDDEN] */
float* __restrict__ h_s2_out, /* [B, IQN_HIDDEN] */
int batch_size
int batch_size,
int state_dim /* runtime state dimension (replaces IQN_STATE_DIM) */
)
{
int sample = blockIdx.x;
@@ -765,13 +766,13 @@ void iqn_trunk_forward_kernel(
extern __shared__ float shmem[];
float* h1 = shmem; /* [IQN_SHARED_H1] */
const float* my_state = states + sample * IQN_STATE_DIM;
const float* my_state = states + sample * state_dim;
/* Layer 1: h1 = leaky_relu(W_s1 @ state + b_s1) */
for (int h = lane; h < IQN_SHARED_H1; h += 32) {
float acc = b_s1[h];
const float* w_row = w_s1 + h * IQN_STATE_DIM;
for (int d = 0; d < IQN_STATE_DIM; d++)
const float* w_row = w_s1 + h * state_dim;
for (int d = 0; d < state_dim; d++)
acc += w_row[d] * my_state[d];
h1[h] = (acc > 0.0f) ? acc : 0.01f * acc;
}

9
crates/ml/src/cuda_pipeline/monitoring_kernel.cu
View File

@@ -27,7 +27,10 @@ extern "C" __global__ void monitoring_reduce(
const float* __restrict__ rewards, // [N]
const int* __restrict__ actions, // [N]
float* summary, // [16]: mean, std, min, max, sharpe, counts[9], total, _pad
int N
int N,
int num_actions, // runtime: DQN_NUM_ACTIONS (9)
int order_actions, // runtime: DQN_ORDER_ACTIONS (3)
int urgency_actions // runtime: DQN_URGENCY_ACTIONS (3)
) {
__shared__ float s_sum;
__shared__ float s_sq_sum;
@@ -60,8 +63,8 @@ extern "C" __global__ void monitoring_reduce(
/* Decode exposure index from factored action for monitoring.
* Factored action = exposure * (b1*b2) + order * b2 + urgency. */
int a = actions[i];
int exp_idx = a / (DQN_ORDER_ACTIONS * DQN_URGENCY_ACTIONS);
if (exp_idx >= 0 && exp_idx < DQN_NUM_ACTIONS) local_counts[exp_idx]++;
int exp_idx = a / (order_actions * urgency_actions);
if (exp_idx >= 0 && exp_idx < num_actions) local_counts[exp_idx]++;
}
// Warp-level reduction before atomic to shared (reduces contention 32x)