From 82b595f6330199b756e3add9472bebbdb96ea4e6 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 26 Mar 2026 00:18:18 +0100 Subject: [PATCH] refactor: convert #define constants to kernel parameters in .cu files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../attention_backward_kernel.cu | 10 +- .../ml/src/cuda_pipeline/attention_kernel.cu | 10 +- .../backtest_forward_ppo_kernel.cu | 44 +++---- .../cuda_pipeline/backtest_metrics_kernel.cu | 13 ++- .../cuda_pipeline/common_device_functions.cuh | 20 +++- .../src/cuda_pipeline/dqn_utility_kernels.cu | 13 ++- .../src/cuda_pipeline/her_relabel_kernel.cu | 29 ++--- .../ml/src/cuda_pipeline/iql_value_kernel.cu | 108 ++++++++++-------- .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 13 ++- .../ml/src/cuda_pipeline/monitoring_kernel.cu | 9 +- 10 files changed, 155 insertions(+), 114 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu index 76360b2f3..51bd0f3e5 100644 --- a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/attention_kernel.cu b/crates/ml/src/cuda_pipeline/attention_kernel.cu index 1cc530c1e..51c517b79 100644 --- a/crates/ml/src/cuda_pipeline/attention_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_kernel.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu index 816f76e71..1d9386eed 100644 --- a/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu @@ -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; } diff --git a/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu index 67f9299e3..9f9a1b94e 100644 --- a/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu @@ -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 diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 29a3b7b7d..75fd65e08 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -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) */ diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 77732dfa7..9835e1966 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -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; diff --git a/crates/ml/src/cuda_pipeline/her_relabel_kernel.cu b/crates/ml/src/cuda_pipeline/her_relabel_kernel.cu index 201ba6bb2..e10ec4c65 100644 --- a/crates/ml/src/cuda_pipeline/her_relabel_kernel.cu +++ b/crates/ml/src/cuda_pipeline/her_relabel_kernel.cu @@ -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; } diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index 66edddcb0..56a6ccbba 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -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); } diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu index 74e38f018..59dc59522 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -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; } diff --git a/crates/ml/src/cuda_pipeline/monitoring_kernel.cu b/crates/ml/src/cuda_pipeline/monitoring_kernel.cu index 018312fbb..67e8cf753 100644 --- a/crates/ml/src/cuda_pipeline/monitoring_kernel.cu +++ b/crates/ml/src/cuda_pipeline/monitoring_kernel.cu @@ -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)