diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 5c47f8bae..1e03c74c0 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -100,6 +100,7 @@ const KERNELS: &[&str] = &[ "rl_sample_tau", // CUDA graph prereq: device-side xorshift32 tau ~ U(0,1) for IQN; replaces host ChaCha8 + mapped-pinned upload "rl_sample_noise", // CUDA graph prereq: device-side factored noise f(rand) for NoisyLinear; replaces host ChaCha8 + mapped-pinned upload "rl_write_u64", // CUDA graph prereq: single-thread u64 scalar write for device-resident ts_ns (graph-captured kernels read via pointer) + "rl_fused_reward_pipeline", // P3: 7→1 fused per-batch reward extraction + shaping + EMA + outcome update "rl_increment_step", // device-resident step counter bump (ISV[548] += 1.0); graph-safe prereq — removes scalar current_step from all downstream kernel args "rl_fused_controllers", // fused kernel: 10 RL ISV controllers in one launch (gamma, tau, ppo_clip, entropy_coef, rollout_steps, per_alpha, reward_scale, ppo_ratio_clamp, gate_threshold, q_distill_lambda) — saves 9 kernel launch overheads (~40-80μs/step) ]; diff --git a/crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu b/crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu new file mode 100644 index 000000000..5cff5f074 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu @@ -0,0 +1,259 @@ +// rl_fused_reward_pipeline.cu — fused per-batch reward pipeline kernel. +// +// Combines 7 sequential per-batch operations into a single kernel +// launch, eliminating 6 kernel launches + 6 L2 cache round-trips on +// the reward data path: +// +// 1. extract_realized_pnl_delta — read Pos, compute reward delta + done +// 2. rl_unit_state_update — detect position transitions, populate units +// 3. rl_step_counter_update — bump steps_since_done, emit trade_duration +// 4. abs_copy — write |reward| to reward_abs +// 5. rl_reward_shaping — entry cost + hold bonus + short-hold penalty +// 6. DtoD raw_rewards snapshot — raw_rewards[b] = rewards[b] (before scale) +// 7. rl_recent_outcome_update — per-batch signed outcome EMA +// +// NOT fused (they require tree-reduce across ALL batches): +// - ema_update_on_done (3 calls) — single-block tree-reduce +// - launch_rl_fused_controllers — reads ISV EMAs +// - apply_reward_scale — single-block tree-reduce +// - rl_reward_clamp_controller / rl_atom_support_update +// +// Grid: ceil(b_size/32), Block: 32. +// One thread per batch element. No shared memory needed. +// +// Per `feedback_no_atomicadd`: per-batch element-wise (b==0 writes +// pyramid_add_count to ISV diag slot — single writer, no atomic). +// Per `feedback_cpu_is_read_only`: pure device-side. +// Per `feedback_no_nvrtc`: pre-compiled cubin. + +#include + +// ISV slot indices (must match crates/ml-alpha/src/rl/isv_slots.rs) +#define RL_TRAIL_K_INIT_INDEX 496 +#define RL_MEAN_ABS_PNL_EMA_INDEX 423 +#define RL_PYRAMID_ADD_COUNT_INDEX 507 +#define RL_STEP_COUNTER_ISV_INDEX 548 +#define RL_ENTRY_COST_INDEX 532 +#define RL_SHORT_HOLD_MIN_STEPS_INDEX 533 +#define RL_SHORT_HOLD_PENALTY_INDEX 534 +#define RL_HOLD_BONUS_INDEX 535 +#define RL_OUTCOME_ALPHA_INDEX 520 + +#define MAX_UNITS 4 + +// PosFlat layout (crates/ml-backtesting/src/lob/mod.rs::PosFlat): +// offset 0: position_lots: i32 (4 bytes) +// offset 4: vwap_entry: f32 (4 bytes) +// offset 8: realized_pnl: f32 (4 bytes) +// offset 12: peak_equity: f32 (4 bytes) + +extern "C" __global__ void rl_fused_reward_pipeline( + // --- From extract_realized_pnl_delta --- + const unsigned char* __restrict__ pos_d, + float* __restrict__ prev_realized_pnl, // [B] IN/OUT + int* __restrict__ prev_position_lots, // [B] IN/OUT (extract's tracker) + float* __restrict__ rewards, // [B] OUT + float* __restrict__ dones, // [B] OUT + // --- From rl_unit_state_update --- + int* __restrict__ unit_prev_pos_lots, // [B] IN/OUT (unit's own tracker) + float* __restrict__ unit_entry_price, // [B*MAX_UNITS] OUT + int* __restrict__ unit_entry_step, // [B*MAX_UNITS] OUT + int* __restrict__ unit_lots, // [B*MAX_UNITS] OUT + float* __restrict__ unit_initial_r, // [B*MAX_UNITS] OUT + float* __restrict__ unit_trail_distance, // [B*MAX_UNITS] OUT + unsigned char* __restrict__ unit_active, // [B*MAX_UNITS] OUT + int* __restrict__ pyramid_units_count, // [B] OUT + // --- From rl_step_counter_update --- + int* __restrict__ steps_since_done, // [B] IN/OUT + float* __restrict__ trade_duration_emit, // [B] OUT + // --- From abs_copy + reward_shaping + raw_rewards + outcome --- + float* __restrict__ reward_abs, // [B] OUT + float* __restrict__ raw_rewards, // [B] OUT + float* __restrict__ outcome_ema, // [B] IN/OUT + float* __restrict__ isv, + int b_size, + int pos_bytes +) { + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= b_size) return; + + // ================================================================ + // PHASE 1: extract_realized_pnl_delta + // ================================================================ + const unsigned char* p = pos_d + b * pos_bytes; + const int current_lots = *(const int*)(p + 0); + const float current_realized = *(const float*)(p + 8); + const float vwap = *(const float*)(p + 4); + + const float prev_realized = prev_realized_pnl[b]; + const int prev_lots = prev_position_lots[b]; + + float reward = current_realized - prev_realized; + float done = (prev_lots != 0 && current_lots == 0) ? 1.0f : 0.0f; + + prev_realized_pnl[b] = current_realized; + prev_position_lots[b] = current_lots; + + // Write reward/done early — subsequent phases read from locals. + rewards[b] = reward; + dones[b] = done; + + // ================================================================ + // PHASE 2: rl_unit_state_update + // ================================================================ + const int current_step = (int)isv[RL_STEP_COUNTER_ISV_INDEX]; + const int unit_prev = unit_prev_pos_lots[b]; + + const float k_init = isv[RL_TRAIL_K_INIT_INDEX]; + const float mean_abs_pnl = isv[RL_MEAN_ABS_PNL_EMA_INDEX]; + const float bootstrap_d = k_init * (mean_abs_pnl + 1e-8f); + + const int base = b * MAX_UNITS; + + const bool was_flat_u = (unit_prev == 0); + const bool is_flat_u = (current_lots == 0); + const bool reversed_u = (!was_flat_u && !is_flat_u && + ((unit_prev > 0) != (current_lots > 0))); + + int add_count = 0; + + if (is_flat_u && !was_flat_u) { + // CLOSE: deactivate all units. + #pragma unroll + for (int u = 0; u < MAX_UNITS; ++u) { + unit_active[base + u] = 0; + unit_lots[base + u] = 0; + } + pyramid_units_count[b] = 0; + } else if (!is_flat_u && (was_flat_u || reversed_u)) { + // OPEN or REVERSE: clear all + populate slot 0. + #pragma unroll + for (int u = 0; u < MAX_UNITS; ++u) { + unit_active[base + u] = 0; + unit_lots[base + u] = 0; + } + unit_entry_price[base + 0] = vwap; + unit_entry_step[base + 0] = current_step; + unit_lots[base + 0] = current_lots; + unit_initial_r[base + 0] = bootstrap_d; + unit_trail_distance[base + 0] = bootstrap_d; + unit_active[base + 0] = 1; + pyramid_units_count[b] = 1; + add_count = 1; + } else if (!is_flat_u && !was_flat_u && !reversed_u) { + // CONTINUE: same direction. + const int abs_curr = (current_lots > 0) ? current_lots : -current_lots; + const int abs_prev_u = (unit_prev > 0) ? unit_prev : -unit_prev; + + if (abs_curr > abs_prev_u) { + // Pyramid ADD. + int count = pyramid_units_count[b]; + if (count < MAX_UNITS) { + const int slot = count; + unit_entry_price[base + slot] = vwap; + unit_entry_step[base + slot] = current_step; + const int delta_lots = abs_curr - abs_prev_u; + unit_lots[base + slot] = (current_lots > 0) ? delta_lots : -delta_lots; + unit_initial_r[base + slot] = bootstrap_d; + unit_trail_distance[base + slot] = bootstrap_d; + unit_active[base + slot] = 1; + pyramid_units_count[b] = count + 1; + add_count = 1; + } + } else if (abs_curr < abs_prev_u) { + // Partial FLAT: deactivate oldest active unit. + for (int u = 0; u < MAX_UNITS; ++u) { + if (unit_active[base + u]) { + unit_active[base + u] = 0; + unit_lots[base + u] = 0; + int count = pyramid_units_count[b]; + if (count > 0) { + pyramid_units_count[b] = count - 1; + } + break; + } + } + } + } + + unit_prev_pos_lots[b] = current_lots; + + // b==0 writes pyramid diag to ISV (single writer, no atomicAdd). + if (b == 0) { + isv[RL_PYRAMID_ADD_COUNT_INDEX] = (float)add_count; + } + + // ================================================================ + // PHASE 3: rl_step_counter_update + // ================================================================ + const int cur_counter = steps_since_done[b]; + if (done > 0.5f) { + trade_duration_emit[b] = (float)cur_counter; + steps_since_done[b] = 0; + } else { + trade_duration_emit[b] = 0.0f; + steps_since_done[b] = cur_counter + 1; + } + + // ================================================================ + // PHASE 4: abs_copy + // ================================================================ + reward_abs[b] = fabsf(reward); + + // ================================================================ + // PHASE 5: rl_reward_shaping (surfer philosophy) + // ================================================================ + const float entry_cost = isv[RL_ENTRY_COST_INDEX]; + const float min_hold = isv[RL_SHORT_HOLD_MIN_STEPS_INDEX]; + const float penalty = isv[RL_SHORT_HOLD_PENALTY_INDEX]; + const float hold_bonus = isv[RL_HOLD_BONUS_INDEX]; + // Use extract's prev_lots (captured before update) and current_lots. + const int hold_time = cur_counter; // counter BEFORE this step's increment/reset + + float r = reward; + + // 1. Entry cost: flat → positioned transition. + if (prev_lots == 0 && current_lots != 0) { + r -= entry_cost; + } + + // 2. Short-hold penalty: trade close with hold time below minimum. + if (done > 0.5f && (float)hold_time < min_hold) { + r *= penalty; + } + + // 3. Long-ride bonus: profitable close amplified by hold time. + if (done > 0.5f && r > 0.0f && hold_time > 0) { + const float ride_mult = 1.0f + hold_bonus * sqrtf((float)hold_time); + r *= ride_mult; + } + + // 4. Per-step hold bonus for staying in a profitable position. + if (prev_lots != 0 && current_lots != 0 && r > 0.0f && hold_time > 0) { + r += hold_bonus * sqrtf((float)hold_time); + } + + // Write shaped reward back. + rewards[b] = r; + + // ================================================================ + // PHASE 6: raw_rewards snapshot (before scaling) + // ================================================================ + raw_rewards[b] = r; + + // ================================================================ + // PHASE 7: rl_recent_outcome_update + // ================================================================ + if (done > 0.5f) { + const float outcome = (r > 0.0f) ? 1.0f : ((r < 0.0f) ? -1.0f : 0.0f); + const float alpha = isv[RL_OUTCOME_ALPHA_INDEX]; + const float prev_ema = outcome_ema[b]; + + if (prev_ema == 0.0f && outcome != 0.0f) { + // Pearl-A first-observation bootstrap. + outcome_ema[b] = outcome; + } else { + outcome_ema[b] = (1.0f - alpha) * prev_ema + alpha * outcome; + } + } +} diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index bbad7b2e0..e728392fe 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -241,6 +241,8 @@ const RL_SAMPLE_TAU_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_sample_tau.cubin")); const RL_WRITE_U64_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_write_u64.cubin")); +const RL_FUSED_REWARD_PIPELINE_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_fused_reward_pipeline.cubin")); // Element-wise in-place add: dst[i] += src[i]. Used by noisy // exploration to fold NoisyLinear output into ensemble_q_d before @@ -6088,133 +6090,6 @@ impl IntegratedTrainer { Ok(indices) } - /// Reduce `var(x) / max(|mean(x)|, 1e-6)` over `input_d[b_size]` - /// into `self.ema_input_scratch_d[1]`. Single block, - /// next-pow2(b_size) threads, shared mem reused across the two - /// passes inside the kernel. - fn launch_var_over_abs_mean( - &mut self, - input_d: &CudaSlice, - b_size: usize, - ) -> Result<()> { - debug_assert!(input_d.len() >= b_size); - // Single-thread streaming kernel; folds across STEPS (not batch). - // Writes var/|mean| directly into the controller-input ISV slot - // — caller does NOT need a downstream ema_update_per_step. - let cfg = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (1, 1, 1), - shared_mem_bytes: 0, - }; - let b_size_i = b_size as i32; - let out_slot: i32 = - crate::rl::isv_slots::RL_ADVANTAGE_VAR_RATIO_EMA_INDEX as i32; - let mean_slot: i32 = - crate::rl::isv_slots::RL_ADV_VAR_STREAM_MEAN_INDEX as i32; - let m2_slot: i32 = - crate::rl::isv_slots::RL_ADV_VAR_STREAM_M2_INDEX as i32; - let clamp_slot: i32 = - crate::rl::isv_slots::RL_ADV_VAR_RATIO_CLAMP_INDEX as i32; - let mut launch = self.stream.launch_builder( - &self.rl_var_over_abs_mean_streaming_fn, - ); - launch - .arg(input_d) - .arg(&mut self.isv_d) - .arg(&b_size_i) - .arg(&out_slot) - .arg(&mean_slot) - .arg(&m2_slot) - .arg(&clamp_slot); - unsafe { - launch - .launch(cfg) - .context("rl_var_over_abs_mean_streaming launch")?; - } - Ok(()) - } - - /// Reduce `‖x‖₂ = sqrt(Σ x²)` over `input_d[n]` into - /// `self.ema_input_scratch_d[1]`. Single block, 256 threads, - /// grid-stride loop covers arbitrary n via per-thread partial - /// accumulation + shared-mem tree-reduce. Used by the LR - /// controller's grad-norm input path (per-head grad_w_d L2 norm - /// fed to ema_update_per_step → ISV[424..426]). - fn launch_l2_norm( - &mut self, - input_d: &CudaSlice, - n: usize, - ) -> Result<()> { - debug_assert!(input_d.len() >= n); - if n == 0 { - return Ok(()); - } - let block_dim: u32 = 256; - let cfg = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (block_dim, 1, 1), - shared_mem_bytes: (block_dim as usize * std::mem::size_of::()) as u32, - }; - let n_i = n as i32; - let mut launch = self.stream.launch_builder(&self.rl_l2_norm_fn); - launch - .arg(input_d) - .arg(&mut self.ema_input_scratch_d) - .arg(&n_i); - unsafe { - launch.launch(cfg).context("rl_l2_norm launch")?; - } - Ok(()) - } - - /// Streaming kurtosis `M4/M2²` that folds ACROSS STEPS (Welford-EMA - /// on per-batch means) rather than across batch. Writes the - /// smoothed estimate directly into `RL_TD_KURTOSIS_EMA_INDEX` — - /// caller does NOT need a downstream ema_update_per_step. Replaces - /// the prior per-batch reduction which returned 0 at b_size=1 - /// (kurtosis undefined) and left the rl_per_alpha controller blind. - fn launch_kurtosis( - &mut self, - input_d: &CudaSlice, - b_size: usize, - ) -> Result<()> { - debug_assert!(input_d.len() >= b_size); - let cfg = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (1, 1, 1), - shared_mem_bytes: 0, - }; - let b_size_i = b_size as i32; - let out_slot: i32 = - crate::rl::isv_slots::RL_TD_KURTOSIS_EMA_INDEX as i32; - let mean_slot: i32 = - crate::rl::isv_slots::RL_TD_KURT_STREAM_MEAN_INDEX as i32; - let m2_slot: i32 = - crate::rl::isv_slots::RL_TD_KURT_STREAM_M2_INDEX as i32; - let m4_slot: i32 = - crate::rl::isv_slots::RL_TD_KURT_STREAM_M4_INDEX as i32; - let clamp_slot: i32 = - crate::rl::isv_slots::RL_TD_KURTOSIS_CLAMP_INDEX as i32; - let mut launch = self.stream.launch_builder( - &self.rl_kurtosis_streaming_fn, - ); - launch - .arg(input_d) - .arg(&mut self.isv_d) - .arg(&b_size_i) - .arg(&out_slot) - .arg(&mean_slot) - .arg(&m2_slot) - .arg(&m4_slot) - .arg(&clamp_slot); - unsafe { - launch - .launch(cfg) - .context("rl_kurtosis_streaming launch")?; - } - Ok(()) - } - /// Launch `rl_lr_controller` to emit per-head learning rates into /// `ISV[412..417]` via ReduceLROnPlateau-style monotone decay. /// Per-head observed loss (l_q / l_pi / l_v from step_synthetic's