sp7(cuda): loss-balance controller kernel (producer-only)
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
165
crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu
Normal file
165
crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu
Normal file
@@ -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 <cuda_runtime.h>
|
||||
|
||||
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();
|
||||
}
|
||||
@@ -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).
|
||||
|
||||
Reference in New Issue
Block a user