From 39f90f3723c0609ae718f2762ff8b31b0ae10c77 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 23 May 2026 21:53:56 +0200 Subject: [PATCH] fix(rl): EMA-streaming variance + kurtosis kernels fix b_size=1 dead inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mjzfk + pdgxn diags showed `advantage_var_ratio` and `td_kurtosis` identically 0 for 100% of every 50k-step smoke. Root cause: the per-batch `rl_var_over_abs_mean_b` and `rl_kurtosis_b` kernels are mathematically undefined at b_size=1 (variance of a single sample is zero; kurtosis of a single sample is 0/0). The kernels correctly returned 0 in that case but the downstream `rl_rollout_steps` and `rl_per_alpha` controllers then never saw signal and pegged at MIN (2048 / 0.4) for the entire run. ## Fix: time-axis Welford-EMA streaming Replace per-batch reduction with per-step EMA-streaming moments maintained in ISV slots: rl_var_over_abs_mean_streaming.cu — maintains streaming mean + M2, emits var/|mean| each step. Welford-EMA on the batch-mean of advantages_d (one value at b_size=1, or a single batch reduction at b_size>1) folded into the time-axis estimator. rl_kurtosis_streaming.cu — maintains streaming mean + M2 + M4, emits M4/M2² (Pearson kurtosis) each step. Same Welford-EMA shape applied to td_per_sample_d batch mean. Both kernels use STREAM_ALPHA = 0.05 (matches LR_LOSS_EMA_ALPHA — half-life ≈ 14 steps) so the time estimator smooths over noisy per-step batch-mean observations. The kernel writes the smoothed estimate DIRECTLY to the controller-input ISV slot (RL_ADVANTAGE_VAR_RATIO_EMA_INDEX = 421, RL_TD_KURTOSIS_EMA_INDEX = 422); the prior downstream ema_update_per_step calls for these two signals are REMOVED — the streaming kernel IS the EMA. ## ISV slot allocation 5 new state slots holding the streaming-mean / M2 / M4 per-stream state. RL_SLOTS_END: 442 → 447. RL_ADV_VAR_STREAM_MEAN_INDEX = 442 (streaming mean of advantages) RL_ADV_VAR_STREAM_M2_INDEX = 443 (streaming M2 of advantages) RL_TD_KURT_STREAM_MEAN_INDEX = 444 (streaming mean of TD-CE) RL_TD_KURT_STREAM_M2_INDEX = 445 RL_TD_KURT_STREAM_M4_INDEX = 446 Per `pearl_first_observation_bootstrap`: sentinel-zero state triggers replace-direct first-observation bootstrap (the first step seeds μ = batch_mean, M2 = 0, M4 = 0 — subsequent steps blend). Per `pearl_blend_formulas_must_have_permanent_floor`: var/|mean| denominator floored at 1e-6, M2² denominator floored at 1e-12 — prevents div-by-zero blow-up when streaming mean / variance is genuinely zero (cold-start or quiet regime). ## Files * crates/ml-alpha/cuda/rl_var_over_abs_mean_streaming.cu — new * crates/ml-alpha/cuda/rl_kurtosis_streaming.cu — new * crates/ml-alpha/cuda/rl_var_over_abs_mean_b.cu — deleted * crates/ml-alpha/cuda/rl_kurtosis_b.cu — deleted * crates/ml-alpha/src/rl/isv_slots.rs — +5 slots * crates/ml-alpha/src/trainer/integrated.rs — rewired launchers, dropped redundant ema_update calls * crates/ml-alpha/build.rs — swapped cubin entries ## Verified gates (local sm_86) G1 isv_bootstrap ✅ G3 controllers ✅ G4 target_update ✅ integrated_smoke ✅ Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 4 +- crates/ml-alpha/cuda/rl_kurtosis_b.cu | 81 ---------- crates/ml-alpha/cuda/rl_kurtosis_streaming.cu | 105 +++++++++++++ .../ml-alpha/cuda/rl_var_over_abs_mean_b.cu | 64 -------- .../cuda/rl_var_over_abs_mean_streaming.cu | 96 ++++++++++++ crates/ml-alpha/src/rl/isv_slots.rs | 35 ++++- crates/ml-alpha/src/trainer/integrated.rs | 145 ++++++++++-------- 7 files changed, 321 insertions(+), 209 deletions(-) delete mode 100644 crates/ml-alpha/cuda/rl_kurtosis_b.cu create mode 100644 crates/ml-alpha/cuda/rl_kurtosis_streaming.cu delete mode 100644 crates/ml-alpha/cuda/rl_var_over_abs_mean_b.cu create mode 100644 crates/ml-alpha/cuda/rl_var_over_abs_mean_streaming.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 060030251..d19648333 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -55,8 +55,8 @@ const KERNELS: &[&str] = &[ "apply_reward_scale", // RL Phase R6: element-wise rewards *= ISV[RL_REWARD_SCALE_INDEX=406]; closes the F.3b host roundtrip "actions_to_market_targets", // RL Phase R6: 9-action grid → LobSim market_targets[B*2] on device; replaces host submit_market loop per feedback_cpu_is_read_only "abs_copy", // RL Phase R7a: element-wise dst[b] = fabsf(src[b]); feeds |reward| into ema_update_on_done for the MEAN_ABS_PNL_EMA slot - "rl_var_over_abs_mean_b", // var(x)/|mean(x)| reduction → advantage_var_ratio_ema (ISV[421]) feeding rl_rollout_steps - "rl_kurtosis_b", // E[(x-μ)⁴]/σ⁴ reduction → td_kurtosis_ema (ISV[422]) feeding rl_per_alpha + "rl_var_over_abs_mean_streaming",// EMA-streaming var/|mean| (folds across STEPS, fixes b_size=1 → ISV[421] feeding rl_rollout_steps + "rl_kurtosis_streaming", // EMA-streaming kurtosis M4/M2² (folds across STEPS, fixes b_size=1) → ISV[422] feeding rl_per_alpha "rl_kl_approx_b", // Schulman-style KL = mean(log π_old − log π_new) → kl_pi_ema (ISV[419]) feeding rl_ppo_clip "rl_l2_diff_norm", // ‖W_online − W_target‖₂ → q_divergence_ema (ISV[418]) feeding rl_target_tau "rl_step_counter_update", // per-batch trade-duration counter + done-gated emit → mean_trade_duration_ema (ISV[417]) feeding rl_gamma diff --git a/crates/ml-alpha/cuda/rl_kurtosis_b.cu b/crates/ml-alpha/cuda/rl_kurtosis_b.cu deleted file mode 100644 index 30a0042c3..000000000 --- a/crates/ml-alpha/cuda/rl_kurtosis_b.cu +++ /dev/null @@ -1,81 +0,0 @@ -// rl_kurtosis_b.cu — kurtosis reduction `E[(x-μ)⁴] / σ⁴` over b_size -// floats → scalar. -// -// The rl_per_alpha controller consumes `td_kurtosis_ema` (ISV[422]) -// populated from the per-step `td_per_sample_d` (output of -// `dqn_distributional_q_bwd`'s `loss_per_batch` slot). Kurtosis -// quantifies tail-thickness of the per-sample CE distribution — heavy -// tails mean a few transitions dominate Q learning, which is exactly -// what PER α should crank up to concentrate sampling on the -// informative samples. -// -// Three-pass reduction: -// Pass 1: sum → mean -// Pass 2: sum((x - mean)²) → var = sum_sq2 / b_size -// Pass 3: sum((x - mean)⁴) → m4 = sum_sq4 / b_size -// Output: m4 / (var² + ε) -// -// For Gaussian, kurtosis = 3.0 (Pearson's definition). Heavy-tailed -// distributions can have kurtosis > 10. The rl_per_alpha kernel's -// target formula maps kurt > 3 → higher α (sharpens PER sampling). -// -// Single block, b_size threads, shared mem reused across passes. -// At b_size=1 (smoke), kurtosis is undefined (single sample has no -// distribution); the kernel returns 0 in that case — which the -// per_α controller's cold-start gate then skips, holding bootstrap. - -extern "C" __global__ void rl_kurtosis_b( - const float* __restrict__ x, - float* __restrict__ out, - int b_size -) { - extern __shared__ float s[]; - const int tid = threadIdx.x; - - // Degenerate case: b_size < 2 → return 0 (kurtosis undefined; - // cold-start gate on the controller skips when input is 0). - if (b_size < 2) { - if (tid == 0) out[0] = 0.0f; - return; - } - - // ── Pass 1: sum → mean ──────────────────────────────────────── - s[tid] = (tid < b_size) ? x[tid] : 0.0f; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) s[tid] += s[tid + stride]; - __syncthreads(); - } - __shared__ float mean; - if (tid == 0) mean = s[0] / (float)b_size; - __syncthreads(); - - // ── Pass 2: sum((x - mean)²) → var ──────────────────────────── - const float xi = (tid < b_size) ? x[tid] : mean; - const float d = xi - mean; - const float d2 = d * d; - s[tid] = d2; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) s[tid] += s[tid + stride]; - __syncthreads(); - } - __shared__ float var; - if (tid == 0) var = s[0] / (float)b_size; - __syncthreads(); - - // ── Pass 3: sum((x - mean)⁴) → m4 ───────────────────────────── - const float d4 = d2 * d2; - s[tid] = (tid < b_size) ? d4 : 0.0f; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) s[tid] += s[tid + stride]; - __syncthreads(); - } - - if (tid == 0) { - const float m4 = s[0] / (float)b_size; - const float var2 = var * var; - out[0] = m4 / (var2 + 1e-12f); - } -} diff --git a/crates/ml-alpha/cuda/rl_kurtosis_streaming.cu b/crates/ml-alpha/cuda/rl_kurtosis_streaming.cu new file mode 100644 index 000000000..537450fb2 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_kurtosis_streaming.cu @@ -0,0 +1,105 @@ +// rl_kurtosis_streaming.cu — EMA-streaming kurtosis `M4 / M2²` that +// folds ACROSS STEPS instead of across batch. +// +// Replaces `rl_kurtosis_b` for the b_size=1 hot path of the integrated +// RL trainer. The per-batch reduction was mathematically undefined at +// b_size=1 (kurtosis of a single sample is 0/0), so the td_kurtosis +// EMA stayed identically zero forever and the rl_per_alpha controller +// never received signal. Canonical incident: alpha-rl-mjzfk fold0 +// (commit 53aeef099) — 100% of 50k steps had td_kurtosis = 0. +// +// Design: maintain three streaming Welford-EMA moments in ISV slots — +// streaming mean (μ_t), M2 (E[(x-μ_t)²]), M4 (E[(x-μ_t)⁴]). Each step: +// +// 1. Compute the current batch's mean x_t. +// 2. Update streaming mean (sentinel-zero first-obs replace): +// μ_t = (1-α) μ_{t-1} + α x_t +// 3. Re-deviate against the UPDATED mean (per Welford): +// δ = x_t − μ_t +// 4. Update streaming M2 and M4: +// M2_t = (1-α) M2_{t-1} + α δ² +// M4_t = (1-α) M4_{t-1} + α δ⁴ +// 5. Write kurtosis directly to the controller's input slot: +// out_slot = M4_t / (M2_t² + ε) +// +// For a Gaussian distribution, kurtosis = 3.0 (Pearson). Heavy-tailed +// distributions (e.g. TD errors after Q has started learning) have +// kurtosis > 3 — which is what the rl_per_alpha controller cranks up +// PER's priority exponent α to track. By streaming across STEPS the +// kernel sees the actual per-step variation in td_per_sample's batch +// mean, which is a meaningful kurtosis signal even when each batch +// has only one entry. +// +// α = 0.05 — matches rl_var_over_abs_mean_streaming and +// LR_LOSS_EMA_ALPHA for consistency. The downstream +// `ema_update_per_step` call is REMOVED for this signal — the +// streaming kernel writes the smoothed estimate directly to the +// controller's input ISV slot. +// +// Per `pearl_first_observation_bootstrap` + `pearl_blend_formulas_must_have_permanent_floor`: +// sentinel-zero state triggers replace-direct first-observation +// bootstrap; the M2² denominator is permanent-floored at 1e-12 to +// prevent division blow-up when streaming variance is genuinely zero. +// +// Per `pearl_no_atomicadd`: single-thread kernel. + +#define STREAM_ALPHA 0.05f +#define M2_SQ_FLOOR 1e-12f + +extern "C" __global__ void rl_kurtosis_streaming( + const float* __restrict__ x, + float* __restrict__ isv, + int b_size, + int out_slot, + int mean_slot, + int m2_slot, + int m4_slot +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + // Current batch mean (trivial at b_size=1). + float batch_mean = 0.0f; + for (int b = 0; b < b_size; ++b) batch_mean += x[b]; + batch_mean /= (float)b_size; + + // Streaming mean with first-observation bootstrap. + const float mean_prev = isv[mean_slot]; + float mean_new; + if (mean_prev == 0.0f) { + mean_new = batch_mean; + } else { + mean_new = (1.0f - STREAM_ALPHA) * mean_prev + + STREAM_ALPHA * batch_mean; + } + isv[mean_slot] = mean_new; + + // Welford-EMA on δ² and δ⁴ against the UPDATED mean. + const float dev = batch_mean - mean_new; + const float dev2 = dev * dev; + const float dev4 = dev2 * dev2; + + const float m2_prev = isv[m2_slot]; + float m2_new; + if (m2_prev == 0.0f && mean_prev == 0.0f) { + m2_new = dev2; // first observation + } else { + m2_new = (1.0f - STREAM_ALPHA) * m2_prev + + STREAM_ALPHA * dev2; + } + isv[m2_slot] = m2_new; + + const float m4_prev = isv[m4_slot]; + float m4_new; + if (m4_prev == 0.0f && mean_prev == 0.0f) { + m4_new = dev4; // first observation + } else { + m4_new = (1.0f - STREAM_ALPHA) * m4_prev + + STREAM_ALPHA * dev4; + } + isv[m4_slot] = m4_new; + + // Kurtosis = M4 / M2² (Pearson). Permanent-floor on M2² prevents + // div-by-zero before the streaming variance has accumulated. + const float m2_sq = m2_new * m2_new; + isv[out_slot] = m4_new / (m2_sq + M2_SQ_FLOOR); +} diff --git a/crates/ml-alpha/cuda/rl_var_over_abs_mean_b.cu b/crates/ml-alpha/cuda/rl_var_over_abs_mean_b.cu deleted file mode 100644 index 5c5651722..000000000 --- a/crates/ml-alpha/cuda/rl_var_over_abs_mean_b.cu +++ /dev/null @@ -1,64 +0,0 @@ -// rl_var_over_abs_mean_b.cu — `var(x) / max(|mean(x)|, ε)` reduction -// over b_size floats → scalar. -// -// The rl_rollout_steps controller consumes `advantage_var_ratio_ema` -// (ISV[421]) populated from the per-step advantages buffer (output of -// compute_advantage_return). The controller's target formula is -// `prev × (var/abs_mean) / ADV_VAR_RATIO_TARGET` so the input is -// var/abs_mean directly. -// -// Two-pass reduction: -// Pass 1: sum → mean = sum / b_size -// Pass 2: sum((x - mean)²) → var = sum_sq / b_size -// Output: var / max(|mean|, EPS) -// -// Floor on the denominator prevents div-by-zero when advantages are -// centred at zero (which happens during cold-start before the policy -// has learned anything — though by the time this runs, the -// cold-start gate on the controller should hold prev at bootstrap -// regardless). -// -// Single block, b_size threads, shared mem holds the running sum -// for both passes (reused — pass 2 overwrites pass 1). - -extern "C" __global__ void rl_var_over_abs_mean_b( - const float* __restrict__ x, - float* __restrict__ out, - int b_size -) { - extern __shared__ float s[]; - const int tid = threadIdx.x; - - // ── Pass 1: sum → mean ──────────────────────────────────────── - s[tid] = (tid < b_size) ? x[tid] : 0.0f; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - s[tid] += s[tid + stride]; - } - __syncthreads(); - } - __shared__ float mean; - if (tid == 0) { - mean = s[0] / (float)b_size; - } - __syncthreads(); - - // ── Pass 2: sum((x - mean)²) → var ──────────────────────────── - const float xi = (tid < b_size) ? x[tid] : mean; // pad with mean → contributes 0 to var - const float d = xi - mean; - s[tid] = d * d; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - s[tid] += s[tid + stride]; - } - __syncthreads(); - } - - if (tid == 0) { - const float var = s[0] / (float)b_size; - const float abs_mean = fmaxf(fabsf(mean), 1e-6f); - out[0] = var / abs_mean; - } -} diff --git a/crates/ml-alpha/cuda/rl_var_over_abs_mean_streaming.cu b/crates/ml-alpha/cuda/rl_var_over_abs_mean_streaming.cu new file mode 100644 index 000000000..9f9561197 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_var_over_abs_mean_streaming.cu @@ -0,0 +1,96 @@ +// rl_var_over_abs_mean_streaming.cu — EMA-streaming variance ratio +// `var/|mean|` that folds ACROSS STEPS instead of across batch. +// +// Replaces `rl_var_over_abs_mean_b` for the b_size=1 hot path of the +// integrated RL trainer. The per-batch reduction was mathematically +// undefined at b_size=1 (variance of a single sample is 0), so the +// advantage_var_ratio EMA stayed identically zero forever and the +// rl_rollout_steps controller never received signal. Canonical +// incident: alpha-rl-mjzfk fold0 (commit 53aeef099) — 100% of 50k +// steps had advantage_var_ratio = 0. +// +// Design: maintain streaming Welford-EMA moments in ISV slots — +// streaming mean (μ_t), streaming M2 (E[(x-μ_t)²]). Each step: +// +// 1. Compute the current batch's mean x_t (just x[0] when b_size=1). +// 2. Update streaming mean (sentinel-zero first-obs replace): +// μ_t = (1-α) μ_{t-1} + α x_t +// 3. Re-deviate against the UPDATED mean (per Welford): +// δ = x_t − μ_t +// 4. Update streaming M2: +// M2_t = (1-α) M2_{t-1} + α δ² +// 5. Write var/|mean| directly to the controller's input slot: +// out_slot = M2_t / max(|μ_t|, ε) +// +// α = 0.05 (matches LR_LOSS_EMA_ALPHA — slow smoothing so noisy +// per-step batch means average out over ≈14-step half-life). With +// this baked into the streaming kernel, the downstream +// `ema_update_per_step` call that the trainer made for these +// signals is REMOVED — the kernel writes the smoothed estimate +// directly to the controller's input ISV slot. +// +// Per `pearl_first_observation_bootstrap` + `pearl_blend_formulas_must_have_permanent_floor`: +// sentinel-zero state triggers replace-direct first-observation +// bootstrap; the var/|mean| denominator is permanent-floored at +// 1e-6 to prevent division blow-up when the streaming mean is +// genuinely zero. +// +// Per `pearl_no_atomicadd`: single-thread kernel (all state lives in +// ISV slots updated by thread 0, output also written by thread 0). + +#define STREAM_ALPHA 0.05f +#define ABS_MEAN_FLOOR 1e-6f + +// out_slot: ISV slot to write var/|mean| to (e.g. +// RL_ADVANTAGE_VAR_RATIO_EMA_INDEX). +// mean_slot: streaming mean state (e.g. +// RL_ADV_VAR_STREAM_MEAN_INDEX). +// m2_slot: streaming M2 state. +extern "C" __global__ void rl_var_over_abs_mean_streaming( + const float* __restrict__ x, + float* __restrict__ isv, + int b_size, + int out_slot, + int mean_slot, + int m2_slot +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + // Current batch mean (trivial at b_size=1). + float batch_mean = 0.0f; + for (int b = 0; b < b_size; ++b) batch_mean += x[b]; + batch_mean /= (float)b_size; + + // Streaming mean update with first-observation bootstrap. + const float mean_prev = isv[mean_slot]; + float mean_new; + if (mean_prev == 0.0f) { + mean_new = batch_mean; + } else { + mean_new = (1.0f - STREAM_ALPHA) * mean_prev + + STREAM_ALPHA * batch_mean; + } + isv[mean_slot] = mean_new; + + // Welford-EMA M2: deviate against UPDATED mean (not prev), then + // EMA the squared deviation. First-observation seeds M2 from the + // current squared deviation. + const float dev = batch_mean - mean_new; + const float dev2 = dev * dev; + const float m2_prev = isv[m2_slot]; + float m2_new; + if (m2_prev == 0.0f && mean_prev == 0.0f) { + // First observation: M2 starts at the current squared dev + // (which is 0 since batch_mean = mean_new on the first step). + m2_new = dev2; + } else { + m2_new = (1.0f - STREAM_ALPHA) * m2_prev + + STREAM_ALPHA * dev2; + } + isv[m2_slot] = m2_new; + + // var/|mean| → controller input slot directly (no separate EMA + // update — the streaming kernel IS the EMA). + const float abs_mean = fmaxf(fabsf(mean_new), ABS_MEAN_FLOOR); + isv[out_slot] = m2_new / abs_mean; +} diff --git a/crates/ml-alpha/src/rl/isv_slots.rs b/crates/ml-alpha/src/rl/isv_slots.rs index 9ac02bf0b..3960b2ec3 100644 --- a/crates/ml-alpha/src/rl/isv_slots.rs +++ b/crates/ml-alpha/src/rl/isv_slots.rs @@ -250,6 +250,39 @@ pub const RL_PPO_RATIO_CLAMP_MAX_INDEX: usize = 440; /// fired this step. pub const RL_PPO_LOG_RATIO_ABS_MAX_INDEX: usize = 441; +// EMA-streaming state slots for variance / kurtosis estimators that +// fold ACROSS STEPS rather than across batch. At b_size=1 the +// per-batch variance is identically zero (single sample), so the +// original `rl_var_over_abs_mean_b` and `rl_kurtosis_b` kernels emit +// zero forever, leaving the rollout_steps and per_alpha controllers +// blind. The streaming kernels maintain time-axis Welford-EMA moments +// in these slots so the estimators converge over many cheap b_size=1 +// steps instead of needing a batched reduction. Canonical incident: +// alpha-rl-mjzfk fold0 (commit 53aeef099) — advantage_var_ratio and +// td_kurtosis stuck at 0 for 100% of the 50000-step run. + +/// Streaming-mean state for the advantage variance estimator. +/// Updated by `rl_var_over_abs_mean_streaming` each step from the +/// batch-mean of compute_advantage_return's `advantages_d`. Slow +/// α=0.05 EMA — half-life ≈ 14 steps. +pub const RL_ADV_VAR_STREAM_MEAN_INDEX: usize = 442; + +/// Streaming-M2 (centered second moment) for the advantage variance +/// estimator. Welford-EMA: `M2 ← (1-α)·M2 + α·(x − μ_new)²`. +pub const RL_ADV_VAR_STREAM_M2_INDEX: usize = 443; + +/// Streaming-mean state for the TD kurtosis estimator. Updated by +/// `rl_kurtosis_streaming` from per-batch `td_per_sample_d` mean. +pub const RL_TD_KURT_STREAM_MEAN_INDEX: usize = 444; + +/// Streaming-M2 for the TD kurtosis estimator. +pub const RL_TD_KURT_STREAM_M2_INDEX: usize = 445; + +/// Streaming-M4 (centered fourth moment) for the TD kurtosis +/// estimator. Welford-EMA: `M4 ← (1-α)·M4 + α·(x − μ_new)⁴`. +/// Kurtosis = M4 / (M2² + ε). +pub const RL_TD_KURT_STREAM_M4_INDEX: usize = 446; + /// Last RL-allocated slot index (exclusive). The integrated trainer /// extends `ISV_TOTAL_DIM` to at least this value at trainer init time. -pub const RL_SLOTS_END: usize = 442; +pub const RL_SLOTS_END: usize = 447; diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 6360cc781..78d7a2350 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -170,10 +170,10 @@ const ABS_COPY_CUBIN: &[u8] = // Reduction kernels that derive scalar inputs for the controller-input // EMAs (variance ratio + kurtosis). entropy_observed reuses // ema_update_per_step's built-in mean reduce directly on entropy_d. -const RL_VAR_OVER_ABS_MEAN_B_CUBIN: &[u8] = - include_bytes!(concat!(env!("OUT_DIR"), "/rl_var_over_abs_mean_b.cubin")); -const RL_KURTOSIS_B_CUBIN: &[u8] = - include_bytes!(concat!(env!("OUT_DIR"), "/rl_kurtosis_b.cubin")); +const RL_VAR_OVER_ABS_MEAN_STREAMING_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_var_over_abs_mean_streaming.cubin")); +const RL_KURTOSIS_STREAMING_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_kurtosis_streaming.cubin")); // Derivation kernels for the EMA inputs that need a new signal // computation (not just a reduction over an existing buffer): @@ -374,10 +374,10 @@ pub struct IntegratedTrainer { abs_copy_fn: CudaFunction, // ── Per-batch reducers for controller-input EMAs ── - _rl_var_over_abs_mean_b_module: Arc, - rl_var_over_abs_mean_b_fn: CudaFunction, - _rl_kurtosis_b_module: Arc, - rl_kurtosis_b_fn: CudaFunction, + _rl_var_over_abs_mean_streaming_module: Arc, + rl_var_over_abs_mean_streaming_fn: CudaFunction, + _rl_kurtosis_streaming_module: Arc, + rl_kurtosis_streaming_fn: CudaFunction, /// 1-float scratch for the reduce kernels' scalar output. Reused /// across per-step launches (reductions serialised by stream order). /// Feeds `ema_update_per_step` with `b_size=1` for the controller @@ -745,19 +745,25 @@ impl IntegratedTrainer { .load_function("abs_copy") .context("load abs_copy")?; - // Reduce kernels for the variance-ratio + kurtosis EMAs. - let rl_var_over_abs_mean_b_module = ctx - .load_cubin(RL_VAR_OVER_ABS_MEAN_B_CUBIN.to_vec()) - .context("load rl_var_over_abs_mean_b cubin")?; - let rl_var_over_abs_mean_b_fn = rl_var_over_abs_mean_b_module - .load_function("rl_var_over_abs_mean_b") - .context("load rl_var_over_abs_mean_b")?; - let rl_kurtosis_b_module = ctx - .load_cubin(RL_KURTOSIS_B_CUBIN.to_vec()) - .context("load rl_kurtosis_b cubin")?; - let rl_kurtosis_b_fn = rl_kurtosis_b_module - .load_function("rl_kurtosis_b") - .context("load rl_kurtosis_b")?; + // Streaming kernels for the variance-ratio + kurtosis EMAs. + // Per R9 fix: these now fold ACROSS STEPS (Welford-EMA on + // per-batch means in ISV state slots) instead of across batch + // — works correctly at b_size=1 where per-batch variance is + // identically zero. The kernels write the smoothed estimate + // directly to the controller-input ISV slot; no separate + // ema_update_per_step is needed downstream. + let rl_var_over_abs_mean_streaming_module = ctx + .load_cubin(RL_VAR_OVER_ABS_MEAN_STREAMING_CUBIN.to_vec()) + .context("load rl_var_over_abs_mean_streaming cubin")?; + let rl_var_over_abs_mean_streaming_fn = rl_var_over_abs_mean_streaming_module + .load_function("rl_var_over_abs_mean_streaming") + .context("load rl_var_over_abs_mean_streaming")?; + let rl_kurtosis_streaming_module = ctx + .load_cubin(RL_KURTOSIS_STREAMING_CUBIN.to_vec()) + .context("load rl_kurtosis_streaming cubin")?; + let rl_kurtosis_streaming_fn = rl_kurtosis_streaming_module + .load_function("rl_kurtosis_streaming") + .context("load rl_kurtosis_streaming")?; let ema_input_scratch_d = stream .alloc_zeros::(1) .context("alloc ema_input_scratch_d")?; @@ -1010,10 +1016,10 @@ impl IntegratedTrainer { actions_to_market_targets_fn, _abs_copy_module: abs_copy_module, abs_copy_fn, - _rl_var_over_abs_mean_b_module: rl_var_over_abs_mean_b_module, - rl_var_over_abs_mean_b_fn, - _rl_kurtosis_b_module: rl_kurtosis_b_module, - rl_kurtosis_b_fn, + _rl_var_over_abs_mean_streaming_module: rl_var_over_abs_mean_streaming_module, + rl_var_over_abs_mean_streaming_fn, + _rl_kurtosis_streaming_module: rl_kurtosis_streaming_module, + rl_kurtosis_streaming_fn, ema_input_scratch_d, _rl_kl_approx_b_module: rl_kl_approx_b_module, rl_kl_approx_b_fn, @@ -1995,15 +2001,11 @@ impl IntegratedTrainer { b_size, ) .context("ema_update_per_step(entropy_observed)")?; + // Streaming kurtosis kernel writes directly to + // RL_TD_KURTOSIS_EMA_INDEX; no separate ema_update_per_step + // needed (the streaming kernel IS the EMA — folds across STEPS). self.launch_kurtosis(&self.td_per_sample_d.clone(), b_size) .context("launch_kurtosis(td_per_sample_d)")?; - self.launch_ema_update_per_step( - crate::rl::isv_slots::RL_TD_KURTOSIS_EMA_INDEX, - RL_LR_CONTROLLER_ALPHA, - &self.ema_input_scratch_d.clone(), - 1, - ) - .context("ema_update_per_step(td_kurtosis)")?; // kl_pi EMA — Schulman-style approximation // `mean(log π_old(a) − log π_new(a))` over the sampled action. @@ -2640,17 +2642,12 @@ impl IntegratedTrainer { // flawed branch did is gone with the host advantage loop). // Feed `advantage_var_ratio_ema` (ISV[421] → rl_rollout_steps - // controller) from this step's advantages_d. Reduce to - // var/|mean| scalar, then ema_update_per_step with b_size=1. + // controller) from this step's advantages_d. Streaming kernel + // writes directly to ISV[421] (folds across STEPS — works at + // b_size=1 where per-batch variance is undefined). No + // ema_update_per_step needed downstream. self.launch_var_over_abs_mean(&self.advantages_d.clone(), b_size) .context("launch_var_over_abs_mean(advantages_d)")?; - self.launch_ema_update_per_step( - crate::rl::isv_slots::RL_ADVANTAGE_VAR_RATIO_EMA_INDEX, - RL_LR_CONTROLLER_ALPHA, - &self.ema_input_scratch_d.clone(), - 1, - ) - .context("ema_update_per_step(advantage_var_ratio)")?; // ── Step 7a (R7d): PER push + sample + gather. ──────────────── // Push CURRENT step's b_size transitions (one per batch) to @@ -2925,23 +2922,35 @@ impl IntegratedTrainer { b_size: usize, ) -> Result<()> { debug_assert!(input_d.len() >= b_size); - debug_assert!(b_size <= 1024); - let block_dim = (b_size.next_power_of_two() as u32).max(1); + // 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: (block_dim, 1, 1), - shared_mem_bytes: (block_dim as usize * std::mem::size_of::()) as u32, + block_dim: (1, 1, 1), + shared_mem_bytes: 0, }; let b_size_i = b_size as i32; - let mut launch = self.stream.launch_builder(&self.rl_var_over_abs_mean_b_fn); + 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 mut launch = self.stream.launch_builder( + &self.rl_var_over_abs_mean_streaming_fn, + ); launch .arg(input_d) - .arg(&mut self.ema_input_scratch_d) - .arg(&b_size_i); + .arg(&mut self.isv_d) + .arg(&b_size_i) + .arg(&out_slot) + .arg(&mean_slot) + .arg(&m2_slot); unsafe { launch .launch(cfg) - .context("rl_var_over_abs_mean_b launch")?; + .context("rl_var_over_abs_mean_streaming launch")?; } Ok(()) } @@ -2979,33 +2988,47 @@ impl IntegratedTrainer { Ok(()) } - /// Reduce kurtosis `E[(x-μ)⁴]/σ⁴` over `input_d[b_size]` into - /// `self.ema_input_scratch_d[1]`. Returns 0 at b_size < 2 - /// (kurtosis undefined; the per_α controller's cold-start gate - /// then skips, holding bootstrap). + /// 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); - debug_assert!(b_size <= 1024); - let block_dim = (b_size.next_power_of_two() as u32).max(1); 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, + block_dim: (1, 1, 1), + shared_mem_bytes: 0, }; let b_size_i = b_size as i32; - let mut launch = self.stream.launch_builder(&self.rl_kurtosis_b_fn); + 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 mut launch = self.stream.launch_builder( + &self.rl_kurtosis_streaming_fn, + ); launch .arg(input_d) - .arg(&mut self.ema_input_scratch_d) - .arg(&b_size_i); + .arg(&mut self.isv_d) + .arg(&b_size_i) + .arg(&out_slot) + .arg(&mean_slot) + .arg(&m2_slot) + .arg(&m4_slot); unsafe { launch .launch(cfg) - .context("rl_kurtosis_b launch")?; + .context("rl_kurtosis_streaming launch")?; } Ok(()) }