From 005fe591dbe3d5c70fc672e9dcf7444e86e6b93e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 31 Mar 2026 00:39:59 +0200 Subject: [PATCH] feat: add PPO CUDA kernels and curiosity inference kernel Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-ppo/src/*.cu | 0 .../curiosity_inference_kernel.cu | 77 +++++++++++++++++++ 2 files changed, 77 insertions(+) create mode 100644 crates/ml-ppo/src/*.cu create mode 100644 crates/ml/src/cuda_pipeline/curiosity_inference_kernel.cu diff --git a/crates/ml-ppo/src/*.cu b/crates/ml-ppo/src/*.cu new file mode 100644 index 000000000..e69de29bb diff --git a/crates/ml/src/cuda_pipeline/curiosity_inference_kernel.cu b/crates/ml/src/cuda_pipeline/curiosity_inference_kernel.cu new file mode 100644 index 000000000..ac3bd41e1 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/curiosity_inference_kernel.cu @@ -0,0 +1,77 @@ +/** + * Curiosity inference-only kernel for Q-penalty computation. + * + * Runs the curiosity forward dynamics model (2-layer MLP) on a training batch + * and outputs per-sample MSE prediction error. This error is used by the C51 + * loss kernel to scale down Bellman targets for novel/uncertain states. + * + * Architecture: [MARKET_DIM + 3] -> [CUR_HIDDEN] LeakyReLU -> [CUR_OUTPUT=MARKET_DIM] + * One thread per sample. Output: float prediction error per sample. + * + * Requires common_device_functions.cuh to be prepended (provides MARKET_DIM, + * CUR_INPUT, CUR_HIDDEN, CUR_OUTPUT, DQN_ORDER_ACTIONS, DQN_URGENCY_ACTIONS, + * DQN_NUM_ACTIONS, bf16 helpers). + */ + +extern "C" __global__ void curiosity_inference_error( + const __nv_bfloat16* __restrict__ states, /* [N, state_dim] */ + const int* __restrict__ actions, /* [N] */ + const __nv_bfloat16* __restrict__ next_states, /* [N, state_dim] */ + const __nv_bfloat16* __restrict__ w1, /* [CUR_HIDDEN, CUR_INPUT] */ + const __nv_bfloat16* __restrict__ b1, /* [CUR_HIDDEN] */ + const __nv_bfloat16* __restrict__ w2, /* [CUR_OUTPUT, CUR_HIDDEN] */ + const __nv_bfloat16* __restrict__ b2, /* [CUR_OUTPUT] */ + float* __restrict__ prediction_errors, /* [N] output: per-sample MSE */ + int N, + int state_dim +) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= N) return; + + const __nv_bfloat16* state_bf = states + tid * state_dim; + int action_idx = actions[tid]; + const __nv_bfloat16* next_state_bf = next_states + tid * state_dim; + + /* ── Build input: first MARKET_DIM state features + 3-class action one-hot ── */ + __nv_bfloat16 input[CUR_INPUT]; + for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf[i]; } + input[MARKET_DIM + 0] = bf16_zero(); + input[MARKET_DIM + 1] = bf16_zero(); + input[MARKET_DIM + 2] = bf16_zero(); + + /* Action to category one-hot — decode exposure from factored action. + * Factored action = exposure_idx * (b1_size * b2_size) + order_idx * b2_size + urgency_idx. + * With b0=9, b1=3, b2=3: exposure_idx = action / 9. */ + int exposure_idx = action_idx / (DQN_ORDER_ACTIONS * DQN_URGENCY_ACTIONS); + int category; + if (exposure_idx < DQN_NUM_ACTIONS / 2) category = 0; /* Short */ + else if (exposure_idx == DQN_NUM_ACTIONS / 2) category = 1; /* Flat */ + else category = 2; /* Long */ + input[MARKET_DIM + category] = bf16_one(); + + /* ── Layer 1: hidden = LeakyReLU(w1 * input + b1, alpha=0.01) ── */ + float hidden[CUR_HIDDEN]; + for (int h = 0; h < CUR_HIDDEN; h++) { + float sum = (float)b1[h]; + for (int i = 0; i < CUR_INPUT; i++) { + sum += (float)w1[h * CUR_INPUT + i] * (float)input[i]; + } + hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum; + } + + /* ── Layer 2: pred = w2 * hidden + b2 (no activation) ── */ + float mse = 0.0f; + for (int o = 0; o < CUR_OUTPUT; o++) { + float sum = (float)b2[o]; + for (int h = 0; h < CUR_HIDDEN; h++) { + sum += (float)w2[o * CUR_HIDDEN + h] * hidden[h]; + } + /* Accumulate MSE against actual next state */ + float diff = sum - (float)next_state_bf[o]; + mse += diff * diff; + } + mse /= (float)CUR_OUTPUT; + + /* Clamp to prevent extreme penalties from corrupting training */ + prediction_errors[tid] = fminf(mse, 10.0f); +}