gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(rl): EMA-streaming variance + kurtosis kernels fix b_size=1 dead inputs

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-23 21:53:56 +02:00
parent 6a58ac9465
commit 39f90f3723
7 changed files with 321 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/ml-alpha/build.rs
View File

@@ -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

81
crates/ml-alpha/cuda/rl_kurtosis_b.cu
View File

@@ -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);
}
}

105
crates/ml-alpha/cuda/rl_kurtosis_streaming.cu Normal file
View File

@@ -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);
}

64
crates/ml-alpha/cuda/rl_var_over_abs_mean_b.cu
View File

@@ -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;
}
}

96
crates/ml-alpha/cuda/rl_var_over_abs_mean_streaming.cu Normal file
View File

@@ -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;
}

35
crates/ml-alpha/src/rl/isv_slots.rs
View File

@@ -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;

145
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -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<CudaModule>,
rl_var_over_abs_mean_b_fn: CudaFunction,
_rl_kurtosis_b_module: Arc<CudaModule>,
rl_kurtosis_b_fn: CudaFunction,
_rl_var_over_abs_mean_streaming_module: Arc<CudaModule>,
rl_var_over_abs_mean_streaming_fn: CudaFunction,
_rl_kurtosis_streaming_module: Arc<CudaModule>,
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::<f32>(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::<f32>()) 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<f32>,
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::<f32>()) 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(())
}