From 63226f17cbfc90ba864cad7205bfed02bed0c2ea Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 3 May 2026 01:12:46 +0200 Subject: [PATCH] sp7(cuda): loss-balance controller kernel (producer-only) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-block 8-thread kernel: 2 heads (CQL, C51) × 4 branches. Reads 3-float views into grad_decomp_result_pinned for IQN/CQL_SX/C51, plus FLATNESS_BASE for the per-branch target modulator. Writes new budgets + raw Wiener observations (diff² and sample²) to producer scratch. Per-branch flatness-modulated target ratio (CQL → ANCHOR×(1-flatness), C51 → ANCHOR×flatness). Wiener-optimal α from per-branch EMA state slots, clamped to [1e-4, 0.5]. Cold-start seeds at consumer-side bootstrap (0.02 CQL, 0.05 C51) on first observation; sentinel-aware. Producer-only commit; build.rs entry, kernel slot, launch site land in subsequent tasks. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../loss_balance_controller_kernel.cu | 165 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 7 +- 2 files changed, 171 insertions(+), 1 deletion(-) create mode 100644 crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu diff --git a/crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu b/crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu new file mode 100644 index 000000000..34402981f --- /dev/null +++ b/crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu @@ -0,0 +1,165 @@ +// crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu +// +// SP7 (2026-05-03): per-branch loss-balance controller for CQL and C51 budgets. +// +// Reads per-loss decomp pinned slots [mag_norm, dir_norm, trunk_norm] for +// IQN (reference), CQL_SX (managed, post-budget), C51 (managed, pre-budget). +// Reads per-branch flatness from ISV[FLATNESS_BASE..+4]. Reads prior +// budgets from ISV[BUDGET_{CQL,C51}_BASE..+4] and prior Wiener state from +// ISV[LB_*_VAR_{CQL,C51}_BASE..+4]. Writes new budgets + new Wiener +// observations to the producer scratch buffer (downstream +// `apply_pearls_ad_kernel` smooths into ISV). +// +// Math (per branch b ∈ {dir, mag, ord, urg}, per head h ∈ {CQL, C51}): +// target_ratio[CQL][b] = ANCHOR_CQL_RATIO * (1 - flatness[b]) +// target_ratio[C51][b] = ANCHOR_C51_RATIO * flatness[b] +// slice_idx = (b == 0) ? 1 : (b == 1) ? 0 : 2 // dir, mag, trunk +// actual_ratio = h_norm[slice] / max(iqn_norm[slice], EPS_DIV) +// correction = target_ratio[h][b] / max(actual_ratio, EPS_DIV) +// candidate = old_budget * correction +// +// diff = candidate - old_budget +// sample_sq = h_norm[slice]^2 +// alpha_eff = old_diff_var / (old_diff_var + old_sample_var + EPS_DIV) +// alpha_eff = clamp(alpha_eff, ALPHA_FLOOR, ALPHA_CEIL) +// new_budget = clamp(old_budget + alpha_eff * diff, EPS_DIV, MAX_BUDGET) +// +// Cold-start (per spec section "Cold start"): +// if h_norm[slice] < EPS_DIV || iqn_norm[slice] < EPS_DIV: +// new_budget = cold_start_basis (matches consumer bootstrap) +// else if old_budget < EPS_DIV: +// new_budget = clamp(cold_start_basis * correction, EPS_DIV, MAX_BUDGET) +// +// Single block, 8 threads (2 heads × 4 branches). No atomicAdd. +// __threadfence_system() after writes. + +#include + +extern "C" __global__ void loss_balance_controller_update( + // Pointers into the shared `grad_decomp_result_pinned` buffer at the + // 3-float component offsets (the launcher computes these by adding + // (offset_elems * sizeof(float)) to grad_decomp_result_dev_ptr): + // iqn_decomp = base + 0 (slot iqn, layout [mag, dir, trunk]) + // cql_decomp = base + 24 (slot cql_sx, layout [mag, dir, trunk]) + // c51_decomp = base + 36 (slot c51, layout [mag, dir, trunk]) + const float* __restrict__ iqn_decomp, + const float* __restrict__ cql_decomp, + const float* __restrict__ c51_decomp, + // ISV signal bus (read-only). + const float* __restrict__ isv_signals, + int flatness_isv_base, // FLATNESS_BASE = 206 + int budget_cql_isv_base, // BUDGET_CQL_BASE = 198 + int budget_c51_isv_base, // BUDGET_C51_BASE = 190 + int diff_var_cql_isv_base, // LB_DIFF_VAR_CQL_BASE = 297 + int sample_var_cql_isv_base, // LB_SAMPLE_VAR_CQL_BASE = 301 + int diff_var_c51_isv_base, // LB_DIFF_VAR_C51_BASE = 305 + int sample_var_c51_isv_base, // LB_SAMPLE_VAR_C51_BASE = 309 + // Producer scratch buffer (kernel writes here; apply_pearls_ad_kernel + // smooths these into ISV downstream). + float* __restrict__ scratch_out, + int scratch_budget_cql_base, + int scratch_budget_c51_base, + int scratch_diff_var_cql_base, + int scratch_sample_var_cql_base, + int scratch_diff_var_c51_base, + int scratch_sample_var_c51_base +) { + // ── Invariant 1 anchors ────────────────────────────────────────── + // Modulated per-branch by flatness (spec section "Math"). + const float ANCHOR_CQL_RATIO = 2.0f; + const float ANCHOR_C51_RATIO = 1.0f; + // Numerical-stability anchors (Wiener α can mathematically span [0,1]). + const float ALPHA_FLOOR = 1e-4f; + const float ALPHA_CEIL = 0.5f; + const float EPS_DIV = 1e-8f; + const float MAX_BUDGET = 1.0f; + // Cold-start basis — must match consumer-side bootstrap in + // fused_training.rs (CQL_BOOTSTRAP_BUDGET, C51_BOOTSTRAP_BUDGET). + const float COLD_START_FLOOR_CQL = 0.02f; + const float COLD_START_FLOOR_C51 = 0.05f; + + // ── Thread layout: head × branch ───────────────────────────────── + int tid = threadIdx.x; + if (blockIdx.x != 0 || tid >= 8) return; + + int head = tid / 4; // 0 = CQL, 1 = C51 + int branch = tid % 4; // 0..3 (dir, mag, ord, urg) + + // ── Slice mapping (spec "Slice→branch mapping") ────────────────── + // grad_decomp pinned layout is [mag, dir, trunk] per component. + // branch=0 → dir slice [1]; branch=1 → mag slice [0]; + // branch=2,3 → trunk slice [2]. + int slice_idx = (branch == 0) ? 1 : (branch == 1) ? 0 : 2; + + // ── Read the slice norms once ──────────────────────────────────── + float iqn_n = iqn_decomp[slice_idx]; + float h_n = (head == 0) ? cql_decomp[slice_idx] : c51_decomp[slice_idx]; + + // ── Read flatness for this branch ──────────────────────────────── + float flat_b = isv_signals[flatness_isv_base + branch]; + flat_b = fminf(1.0f, fmaxf(0.0f, flat_b)); // safety clamp + + // ── Compute per-branch flatness-modulated target ratio ─────────── + float target_ratio = (head == 0) + ? ANCHOR_CQL_RATIO * (1.0f - flat_b) + : ANCHOR_C51_RATIO * flat_b; + + // ── Read prior budget + per-branch Wiener state ────────────────── + int budget_isv_base = (head == 0) ? budget_cql_isv_base : budget_c51_isv_base; + int diff_var_isv_base = (head == 0) ? diff_var_cql_isv_base : diff_var_c51_isv_base; + int sample_var_isv_base = (head == 0) ? sample_var_cql_isv_base : sample_var_c51_isv_base; + int scratch_budget_base = (head == 0) ? scratch_budget_cql_base : scratch_budget_c51_base; + int scratch_diff_base = (head == 0) ? scratch_diff_var_cql_base : scratch_diff_var_c51_base; + int scratch_sample_base = (head == 0) ? scratch_sample_var_cql_base : scratch_sample_var_c51_base; + float cold_start_basis = (head == 0) ? COLD_START_FLOOR_CQL : COLD_START_FLOOR_C51; + + float old_budget = isv_signals[budget_isv_base + branch]; + float old_diff_var = isv_signals[diff_var_isv_base + branch]; + float old_sample_var = isv_signals[sample_var_isv_base + branch]; + + // ── Branch on cold-start regime ────────────────────────────────── + float new_budget; + float diff; + float sample_sq = h_n * h_n; + + if (h_n < EPS_DIV || iqn_n < EPS_DIV) { + // Warmup hasn't fired or the loss head is gated off. Seed budget + // at the consumer-side bootstrap value (effectively a no-op vs + // the existing behavior on the very first step). + new_budget = cold_start_basis; + diff = 0.0f; + } else { + float actual_ratio = h_n / iqn_n; + float correction = target_ratio / fmaxf(actual_ratio, EPS_DIV); + + if (old_budget < EPS_DIV) { + // Sentinel-0 read (cold start or fold boundary). The norm + // we observed was produced under cql/c51_budget = cold_start_basis + // (consumer-side bootstrap). Seed the controller at the value + // that achieves target_ratio next step. + float seed = cold_start_basis * correction; + new_budget = fminf(MAX_BUDGET, fmaxf(EPS_DIV, seed)); + diff = new_budget - cold_start_basis; + } else { + float candidate = old_budget * correction; + diff = candidate - old_budget; + + // Wiener α from prior state (no in-step state mutation — + // we update the EMAs via scratch_out + apply_pearls_ad). + float wiener_num = old_diff_var; + float wiener_den = old_diff_var + old_sample_var + EPS_DIV; + float alpha_eff = wiener_num / wiener_den; + alpha_eff = fminf(ALPHA_CEIL, fmaxf(ALPHA_FLOOR, alpha_eff)); + + new_budget = old_budget + alpha_eff * diff; + new_budget = fminf(MAX_BUDGET, fmaxf(EPS_DIV, new_budget)); + } + } + + // ── Write the three outputs to scratch (apply_pearls_ad smooths) ── + scratch_out[scratch_budget_base + branch] = new_budget; + scratch_out[scratch_diff_base + branch] = diff * diff; // raw observation + scratch_out[scratch_sample_base + branch] = sample_sq; // raw observation + + __threadfence_system(); +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 081aec433..ab40223c9 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -3885,7 +3885,12 @@ extended. No producer kernel yet — that arrives in the next commit. description (was claiming a non-existent "regime_stability allocator") + clarified `sp5_budget_c51`/`sp5_budget_ens` to reflect SP7 ownership. T2 polish landed (commit `60804788e`) — registry descriptions corrected to reflect current state vs T6/T7 future changes. -- T3: kernel `.cu` file. +- T3 (commit ⟨pending⟩): kernel `loss_balance_controller_kernel.cu` + (~165 LOC). 8-thread single block (2 heads × 4 branches). Reads three + 3-float views into `grad_decomp_result_pinned`, FLATNESS_BASE, prior + budgets, prior Wiener state. Writes new budget + raw Wiener observations + to producer scratch. Cold-start sentinel-aware. Producer-only — no + call site yet. - T4: build.rs cubin manifest. - T5: trainer struct + launcher fn. - T6: Pearl 2 contract change (drop CQL/C51/ENS args).