feat(crt-train): add ISV-driven smoothness_lambda_controller kernel
This commit is contained in:
83
crates/ml-alpha/cuda/smoothness_lambda_controller.cu
Normal file
83
crates/ml-alpha/cuda/smoothness_lambda_controller.cu
Normal file
@@ -0,0 +1,83 @@
|
||||
// smoothness_lambda_controller.cu — ISV-driven per-horizon λ for the
|
||||
// output_smoothness regularizer.
|
||||
//
|
||||
// Reads `raw_per_h[5]` (emitted by output_smoothness_loss_and_grad),
|
||||
// maintains a per-horizon Wiener-α-floor EMA of observed jitter,
|
||||
// derives per-horizon target by anchoring on observed h30 jitter
|
||||
// scaled by HORIZONS[0]/HORIZONS[h], and emits next-step λ[h] with a
|
||||
// permanent floor.
|
||||
//
|
||||
// Per `pearl_controller_anchors_isv_driven`: target is signal-derived,
|
||||
// not a constant. Per `pearl_first_observation_bootstrap`: first
|
||||
// observation replaces EMA directly. Per
|
||||
// `pearl_blend_formulas_must_have_permanent_floor`: λ has a permanent
|
||||
// floor so the controller never fully self-disables. Per
|
||||
// `pearl_wiener_alpha_floor_for_nonstationary`: α floored at 0.5 since
|
||||
// the controller's target drifts as the policy co-adapts.
|
||||
//
|
||||
// Per `feedback_no_atomicadd`: 5 threads, single block, single writer
|
||||
// per (h) slot; no atomics. Per `pearl_no_host_branches_in_captured_graph`:
|
||||
// no host branching; thread-id gating only.
|
||||
|
||||
#define SLC_N_HORIZONS 5
|
||||
#define SLC_LAMBDA_FLOOR 1.0e-4f
|
||||
#define SLC_TARGET_EPS 1.0e-9f
|
||||
#define SLC_ALPHA_FLOOR 0.5f
|
||||
|
||||
// HORIZONS = {30, 100, 300, 1000, 6000}.
|
||||
// Ratios HORIZONS[0]/HORIZONS[h] = {1, 0.3, 0.1, 0.03, 0.005}.
|
||||
// Constant array known at compile time.
|
||||
__device__ __constant__ float TARGET_K_RATIO[SLC_N_HORIZONS] = {
|
||||
1.0f,
|
||||
30.0f / 100.0f,
|
||||
30.0f / 300.0f,
|
||||
30.0f / 1000.0f,
|
||||
30.0f / 6000.0f,
|
||||
};
|
||||
|
||||
extern "C" __global__ void smoothness_lambda_controller(
|
||||
const float* __restrict__ raw_per_h, // [5] emitted by output_smoothness
|
||||
float* __restrict__ jitter_ema, // [5] in/out — EMA state
|
||||
int* __restrict__ first_obs, // [1] in/out — sentinel
|
||||
float base_lambda, // scalar — amplitude knob
|
||||
float* __restrict__ lambda_out // [5] output — λ for next step
|
||||
) {
|
||||
const int h = threadIdx.x;
|
||||
if (h >= SLC_N_HORIZONS) return;
|
||||
|
||||
__shared__ float s_jitter_after[SLC_N_HORIZONS];
|
||||
|
||||
// Pass 1: EMA update with sentinel bootstrap.
|
||||
const float raw_h = raw_per_h[h];
|
||||
const int sentinel = first_obs[0];
|
||||
float jitter_h;
|
||||
if (sentinel == 0) {
|
||||
jitter_h = raw_h;
|
||||
} else {
|
||||
jitter_h = (1.0f - SLC_ALPHA_FLOOR) * jitter_ema[h] + SLC_ALPHA_FLOOR * raw_h;
|
||||
}
|
||||
jitter_ema[h] = jitter_h;
|
||||
s_jitter_after[h] = jitter_h;
|
||||
|
||||
// Single-writer of sentinel — thread h=0 only.
|
||||
if (h == 0 && sentinel == 0) {
|
||||
first_obs[0] = 1;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Pass 2: derive target and update λ.
|
||||
// target[h] = jitter_ema[0] * TARGET_K_RATIO[h]
|
||||
// = jitter_ema[0] for h=0 (self-target)
|
||||
// < jitter_ema[0] for h>0
|
||||
const float jitter_h0 = s_jitter_after[0];
|
||||
const float target_h = jitter_h0 * TARGET_K_RATIO[h];
|
||||
|
||||
// Excess controller: ratio - 1, clamped at zero (only push UP).
|
||||
// When observed > target: excess > 0 → λ grows
|
||||
// When observed ≤ target: excess = 0 → λ relaxes toward base_lambda × 1 = base_lambda
|
||||
// Floor: λ ≥ LAMBDA_FLOOR.
|
||||
const float safe_target = fmaxf(target_h, SLC_TARGET_EPS);
|
||||
const float excess_ratio = fmaxf(0.0f, jitter_h / safe_target - 1.0f);
|
||||
const float lambda_new = base_lambda * (1.0f + excess_ratio);
|
||||
lambda_out[h] = fmaxf(SLC_LAMBDA_FLOOR, lambda_new);
|
||||
}
|
||||
Reference in New Issue
Block a user