diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 57a8cc50b..477344aa7 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -67,6 +67,7 @@ const KERNELS: &[&str] = &[ "rl_streaming_clamp_init", // RL R9: device-side seeder for streaming-kernel output clamp ceilings (ISV[447], ISV[448]) — no HtoD "rl_isv_write", // RL R9: generic single-slot ISV seeder for tunable design constants — no HtoD "rl_q_pi_agree_b", // RL R9 audit: per-batch fraction (argmax Q == argmax π) → ISV[407] EMA; wires previously-dead slot + "rl_pi_action_kernel", // audit Option B: π drives action selection via multinomial sampling from softmax(pi_logits); Q becomes pure critic ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_pi_action_kernel.cu b/crates/ml-alpha/cuda/rl_pi_action_kernel.cu new file mode 100644 index 000000000..0f4e2af2b --- /dev/null +++ b/crates/ml-alpha/cuda/rl_pi_action_kernel.cu @@ -0,0 +1,103 @@ +// rl_pi_action_kernel.cu — GPU-resident multinomial sampler over the +// PPO policy head's softmax(pi_logits). Replaces rl_action_kernel +// (which sampled from Q-Thompson) as the canonical actor-side +// decision per `pearl_q_thompson_actor_makes_pi_dead_weight`. +// +// The prior architecture had Q being BOTH the actor (via Thompson) +// AND the critic (via Bellman target). π trained as a third head +// via PPO surrogate against Q's actions but never drove any +// decision — `q_pi_agree_ema` decayed to 0 across every smoke +// because π converged to Q's Thompson SAMPLING distribution mode, +// not Q's argmax mode. The actor-critic design was inverted. +// +// New architecture (this kernel): π drives action selection via +// multinomial sampling from softmax(pi_logits). Q remains the +// critic (Bellman target via `argmax_expected_q` on h_{t+1}). The +// PPO importance-ratio surrogate now has its canonical semantics — +// π_new / π_old at the action that π_old sampled. +// +// PRNG: same xorshift32-per-batch pattern as rl_action_kernel. +// Single thread per block (N_ACTIONS=9 is small enough that the +// per-thread parallelism of the Thompson kernel is unjustified; +// multinomial CDF walk is sequential anyway). +// +// Per `feedback_no_atomicadd`: single-thread kernel, no atomics. +// Per `feedback_cpu_is_read_only`: pure device-side sampling. + +#include + +#define N_ACTIONS 9 + +__device__ static uint32_t xorshift32(uint32_t* state) { + uint32_t x = *state; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *state = x; + return x; +} + +// One block per batch (grid_dim.x = b_size). Single thread per +// block (tid==0). Computes numerically-stable softmax over the +// batch's N_ACTIONS pi_logits, then multinomial-samples via CDF +// walk using the per-batch xorshift32 PRNG state. +// +// Inputs: +// pi_logits [b_size × N_ACTIONS] row-major raw logits +// prng_state [b_size] per-batch xorshift32 state (mutated) +// Outputs: +// actions [b_size] multinomial-sampled action index per batch +extern "C" __global__ void rl_pi_action_kernel( + const float* __restrict__ pi_logits, + uint32_t* __restrict__ prng_state, + int* __restrict__ actions, + int b_size +) { + const int b = blockIdx.x; + if (b >= b_size) return; + if (threadIdx.x != 0) return; + + // ── Numerically-stable softmax over pi_logits[b, :]. ────────── + const int base = b * N_ACTIONS; + float logits[N_ACTIONS]; + float max_logit = pi_logits[base]; + #pragma unroll + for (int a = 0; a < N_ACTIONS; ++a) { + logits[a] = pi_logits[base + a]; + if (logits[a] > max_logit) max_logit = logits[a]; + } + float sum_exp = 0.0f; + float probs[N_ACTIONS]; + #pragma unroll + for (int a = 0; a < N_ACTIONS; ++a) { + probs[a] = expf(logits[a] - max_logit); + sum_exp += probs[a]; + } + // Guard against pathological zero-sum (shouldn't happen with + // expf of bounded values, but defensive). + const float inv_sum = (sum_exp > 1e-9f) ? (1.0f / sum_exp) : (1.0f / (float)N_ACTIONS); + #pragma unroll + for (int a = 0; a < N_ACTIONS; ++a) probs[a] *= inv_sum; + + // ── Draw uniform u ∈ [0, 1) from per-batch PRNG. ────────────── + uint32_t state = prng_state[b]; + // A few warmup advances dispel low-quality seeding correlations. + #pragma unroll + for (int w = 0; w < 4; ++w) xorshift32(&state); + const uint32_t r = xorshift32(&state); + // Convert to [0, 1): divide by 2^32. Use 24-bit mantissa precision + // (good enough for multinomial — Q_N_ATOMS=21 actions max). + const float u = (float)(r >> 8) * (1.0f / 16777216.0f); + prng_state[b] = state; + + // ── CDF walk: pick the action whose cumulative prob crosses u. ── + int chosen = N_ACTIONS - 1; // defensive default if numerical drift loses the last bit + float cdf = 0.0f; + #pragma unroll + for (int a = 0; a < N_ACTIONS; ++a) { + cdf += probs[a]; + if (u < cdf) { chosen = a; break; } + } + + actions[b] = chosen; +} diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index cc65f19b7..89439c48b 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -111,8 +111,14 @@ struct Cli { /// Batch dimension — number of parallel LobSim backtests run in /// lockstep per step. Larger batches amortise the encoder forward /// cost across more trajectories but push GPU memory pressure - /// linearly. Smoke at 1; production sweep at 32-64 (L40S 48GB). - #[arg(long, default_value_t = 1)] + /// linearly. Default 16 — fixes the b_size=1 signal starvation + /// per `pearl_b_size_1_signal_starvation_blocks_q_learning` + /// (Q stayed at uniform baseline ln(21)=3.04 across all b_size=1 + /// smokes regardless of controller tuning; 16 parallel + /// transitions per Adam step give 16× the gradient signal). Was + /// 1 (validation-smoke default); production sweep can push to + /// 32-64 on L40S 48GB. + #[arg(long, default_value_t = 16)] n_backtests: usize, /// PER buffer capacity (R7d). Per `replay.rs` doc, naive O(N) diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 2fd6c6825..512b69195 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -199,6 +199,8 @@ const RL_ISV_WRITE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_isv_write.cubin")); const RL_Q_PI_AGREE_B_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_q_pi_agree_b.cubin")); +const RL_PI_ACTION_KERNEL_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_pi_action_kernel.cubin")); /// Per-head LR controller Wiener-α target. The controller kernel floors /// this at 0.4 per `pearl_wiener_alpha_floor_for_nonstationary` so the @@ -411,6 +413,8 @@ pub struct IntegratedTrainer { rl_isv_write_fn: CudaFunction, _rl_q_pi_agree_b_module: Arc, rl_q_pi_agree_b_fn: CudaFunction, + _rl_pi_action_kernel_module: Arc, + rl_pi_action_kernel_fn: CudaFunction, /// Per-batch step counter for `trade_duration_ema` — increments /// every step, resets on done. Zero-init at construction. pub steps_since_done_d: CudaSlice, @@ -851,6 +855,12 @@ impl IntegratedTrainer { let rl_q_pi_agree_b_fn = rl_q_pi_agree_b_module .load_function("rl_q_pi_agree_b") .context("load rl_q_pi_agree_b")?; + let rl_pi_action_kernel_module = ctx + .load_cubin(RL_PI_ACTION_KERNEL_CUBIN.to_vec()) + .context("load rl_pi_action_kernel cubin")?; + let rl_pi_action_kernel_fn = rl_pi_action_kernel_module + .load_function("rl_pi_action_kernel") + .context("load rl_pi_action_kernel")?; // Per-batch PRNG state for the Thompson sampler. Seeded // deterministically from cfg.dqn_seed via ChaCha8 host RNG so @@ -1082,6 +1092,8 @@ impl IntegratedTrainer { rl_isv_write_fn, _rl_q_pi_agree_b_module: rl_q_pi_agree_b_module, rl_q_pi_agree_b_fn, + _rl_pi_action_kernel_module: rl_pi_action_kernel_module, + rl_pi_action_kernel_fn, steps_since_done_d, trade_duration_emit_d, prev_realized_pnl_d, @@ -2497,24 +2509,34 @@ impl IntegratedTrainer { // increment that drove the old ChaCha8 host RNG is also gone. self.step_counter = self.step_counter.wrapping_add(1); + // ── π-driven action selection (Option B fix per + // `pearl_q_thompson_actor_makes_pi_dead_weight`). Replaces the + // prior Q-Thompson selector. π is now the actor: actions come + // from multinomial sample of softmax(pi_logits). Q remains the + // critic (Bellman target via argmax_expected_q on h_{t+1}). + // PPO importance-ratio surrogate now has canonical semantics — + // π_new / π_old at the action that π_old sampled. + // + // The rl_action_kernel (Q-Thompson) field is kept loaded for + // backward-compat tests + diagnostic comparison; it's no + // longer in the hot path. { let cfg = LaunchConfig { grid_dim: (b_size as u32, 1, 1), - block_dim: (N_ACTIONS as u32, 1, 1), + block_dim: (1, 1, 1), // single thread per block (multinomial CDF walk is sequential) shared_mem_bytes: 0, }; let b_size_i = b_size as i32; - let mut launch = self.stream.launch_builder(&self.rl_action_kernel_fn); + let mut launch = self.stream.launch_builder(&self.rl_pi_action_kernel_fn); launch - .arg(&q_logits_d) - .arg(&self.atom_supports_d) + .arg(&pi_logits_d) .arg(&mut self.prng_state_d) .arg(&mut self.actions_d) .arg(&b_size_i); unsafe { launch .launch(cfg) - .context("rl_action_kernel launch (Thompson)")?; + .context("rl_pi_action_kernel launch (multinomial from softmax(π))")?; } }