fix(rl): remove CudaSlice::clone() from graph capture + delete dead methods + fused kernel
- Inline launch_l2_norm, launch_ema_update_per_step, launch_kurtosis, launch_var_over_abs_mean at all call sites to eliminate .clone() that allocated device memory inside CUDA graph capture regions (caused CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED at b>16). - Delete the now-unused helper methods (no hiding, no suppressing). - Add pre-commit hook guard for *_d.clone() patterns. - Add rl_fused_reward_pipeline.cu (7→1 per-batch kernel, not yet wired). - Add rl_write_u64 kernel + ts_ns device-resident for Graph B. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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)
|
||||
];
|
||||
|
||||
259
crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu
Normal file
259
crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu
Normal file
@@ -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 <stdint.h>
|
||||
|
||||
// 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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<f32>,
|
||||
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<f32>,
|
||||
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::<f32>()) 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<f32>,
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user