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

feat(sp4): Tasks A14+A15 — Pearl C engagement counter in 5 Adam kernels

Browse Source 
Layer A scaffolding for post-clamp feedback-loop self-correction.

A14: All 5 Adam kernels (dqn/iqn/iql/attn/curiosity) extended with
register-then-tree-reduce engagement counter. Per-thread `local_engage`
set on clamp engagement; block tree-reduce (no atomicAdd, mirrors
dqn_grad_norm_kernel pattern); single block-leader writes per-block
count to clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x].
`engage_buf_offset == -1 (SP4_ENGAGE_OFFSET_DISABLED)` skips writeback.

A15: Host-side rate-deficit check in pearl_c_post_adam_engagement_check.
Sums per-block counts via mapped-pinned zero-copy reads, computes
engagement_rate, rate_deficit = rate - 0.01. Applies Pearls A+D via
pearl_c_rate_deficit_ema (host-only [f32;8]) + pearl_c_rate_deficit_state_buf
(MappedF32Buffer[24] = 8 groups × 3 Wiener floats). Wired post-graph
in FusedTrainingCtx::run_full_step for groups 0/3/4/5/6 and
post-curiosity_train in training_loop for group 7.

3 design issues resolved per Layer A scope:
1. Curiosity sub-launches: SP4_ENGAGE_BUF_LEN extended 2048->2816
   (= 11 x 256). Curiosity sub-launches use offsets 1792/2048/2304/2560
   (W1/b1/W2/b2). Host-side check sums all 4. Allocation +
   reset-registry updated.
2. TLOB/Attn sharing attn_adam_kernel: TLOB passes
   SP4_ENGAGE_OFFSET_DISABLED=-1 (silently skips Pearl C). Attn writes
   to its slot (offset 6x256=1536). Layer B can refactor if needed.
3. DQN main Adam covers groups 0/1/2 in single launch: Layer A accounts
   only under group 0 (DqnTrunk). Groups 1/2 Pearl C tracking deferred
   to Layer B's per-group sub-launch decision.

Wiring path: GpuDqnTrainer exposes nan_flags_buf_ptr() +
clamp_engage_per_block_buf_dev_ptr(). FusedTrainingCtx adds
wire_aux_trainer_pearl_c_buffers() for IQN/IQL hi+lo/Attn.
GpuExperienceCollector adds set_curiosity_pearl_c_buffers() for the
curiosity trainer. Wiring fires once in init_gpu_experience_collector.

State reset registry (Task A12 follow-up):
- sp4_clamp_engage_counters description updated to reflect 2816 buffer
  length and 4 distinct curiosity sub-launch offsets
- New entries: sp4_pearl_c_rate_deficit_state (24 mapped-pinned floats,
  Pearls A+D Wiener state) and sp4_pearl_c_rate_deficit_ema (host-only
  [f32; 8] EMA surrogate). Both reset to zero at fold boundary so Pearl
  A's first-observation sentinel fires on the new fold's first
  engagement-rate-deficit observation. Dispatch arms wired in
  reset_named_state alongside existing sp4_* entries.

Layer A: Pearl C is observability scaffolding. Mech 9's clamp still
uses hardcoded 100xQ_ABS_REF. Engagement counters fire correctly,
rate_deficit EMA tracks. Force-bump branch logs via tracing::debug only;
Layer B's atomic flip activates the actual ISV mutation.

cargo check --lib --tests clean. state_reset_registry tests pass.
sp4_isv_slots::tests::pearl_c_engage_buf_layout pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-01 01:42:40 +02:00
parent 91535deb1a
commit ea31ebfe38
19 changed files with 1291 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
crates/ml/src/cuda_pipeline/attention_backward_kernel.cu
View File

@@ -81,52 +81,92 @@ extern "C" __global__ void attn_adam_kernel(
* — same pattern wired into all five Adam kernels. Used by both
* GpuAttention and GpuTlob (TLOB shares this kernel via the same cubin).
* 0=disabled (debug only). */
float weight_clamp_max_abs
float weight_clamp_max_abs,
/* SP4 Task A14 (2026-04-30): Pearl C engagement counter via warp+
* block tree-reduce (no atomicAdd). GpuAttention + GpuTlob both pass
* `engage_buf_offset = ParamGroup::Attn × MAX_BLOCKS_PER_ADAM` so
* Pearl C tracks the kernel-level engagement rate (TLOB and core
* attention share the same SP4 group). `engage_buf_offset < 0`
* skips the writeback. */
int* __restrict__ nan_flags_buf,
int diag_slot,
int* __restrict__ clamp_engage_per_block_buf,
int engage_buf_offset
) {
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= total_params) return;
bool active = (tid < total_params);
int local_engage = 0;
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged. */
{
if (active) {
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged. */
float g_check = (float)grads[tid];
if (!isfinite(g_check)) { grads[tid] = 0.0f; return; }
if (!isfinite(g_check)) {
grads[tid] = 0.0f;
} else {
/* Gradient clipping */
float grad_norm = norm[0];
float clip_scale = (grad_norm > max_grad_norm && grad_norm > 0.0f)
? max_grad_norm / grad_norm : 1.0f;
float g = grads[tid] * clip_scale;
/* AdamW: weight decay applied to params directly */
params[tid] = params[tid] * (1.0f - lr * weight_decay);
/* Moment updates */
float m_val = beta1 * m[tid] + (1.0f - beta1) * g;
float v_val = beta2 * v[tid] + (1.0f - beta2) * g * g;
m[tid] = m_val;
v[tid] = v_val;
/* Bias correction */
int adam_t = adam_t_buf[0];
float f_t = (float)adam_t;
float bc1 = 1.0f - powf(beta1, f_t);
float bc2 = 1.0f - powf(beta2, f_t);
float m_hat = m_val / (bc1 + 1e-12f);
float v_hat = v_val / (bc2 + 1e-12f);
/* Parameter update */
float p = params[tid] - lr * m_hat / (sqrtf(v_hat) + eps);
/* SP3 Mech 9 + SP4 Task A14: post-Adam |p| clamp +
* engagement counter — last step before writeback. */
if (weight_clamp_max_abs > 0.0f) {
if (fabsf(p) > weight_clamp_max_abs) {
if (nan_flags_buf != nullptr && diag_slot >= 0) {
nan_flags_buf[diag_slot] = 1;
}
local_engage = 1;
}
p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
}
params[tid] = p;
/* Zero gradient for next step */
grads[tid] = 0.0f;
}
}
/* Gradient clipping */
float grad_norm = norm[0];
float clip_scale = (grad_norm > max_grad_norm && grad_norm > 0.0f)
? max_grad_norm / grad_norm : 1.0f;
float g = grads[tid] * clip_scale;
/* SP4 Task A14: warp + block tree-reduce of `local_engage`. */
int e = local_engage;
for (int offset = 16; offset > 0; offset >>= 1)
e += __shfl_xor_sync(0xFFFFFFFF, e, offset);
/* AdamW: weight decay applied to params directly */
params[tid] = params[tid] * (1.0f - lr * weight_decay);
__shared__ int warp_engage[8]; /* blockDim.x = 256 ⇒ 8 warps */
int warp_id = threadIdx.x / 32;
int warp_lane = threadIdx.x % 32;
if (warp_lane == 0) warp_engage[warp_id] = e;
__syncthreads();
/* Moment updates */
float m_val = beta1 * m[tid] + (1.0f - beta1) * g;
float v_val = beta2 * v[tid] + (1.0f - beta2) * g * g;
m[tid] = m_val;
v[tid] = v_val;
/* Bias correction */
int adam_t = adam_t_buf[0];
float f_t = (float)adam_t;
float bc1 = 1.0f - powf(beta1, f_t);
float bc2 = 1.0f - powf(beta2, f_t);
float m_hat = m_val / (bc1 + 1e-12f);
float v_hat = v_val / (bc2 + 1e-12f);
/* Parameter update */
float p = params[tid] - lr * m_hat / (sqrtf(v_hat) + eps);
/* SP3 Mech 9: post-Adam |p| clamp — last step before writeback. */
if (weight_clamp_max_abs > 0.0f) {
p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
if (warp_id == 0) {
int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
}
}
params[tid] = p;
/* Zero gradient for next step */
grads[tid] = 0.0f;
}
/* =====================================================================

84
crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu
View File

@@ -323,34 +323,80 @@ extern "C" __global__ void curiosity_adam_step(
int step, /* 1-based step counter */
/* SP3 Mech 9 (close-out v2): ISV-driven |p| bound (100 × Q_ABS_REF.max(1.0))
* — same pattern wired into all five Adam kernels. 0=disabled (debug only). */
float weight_clamp_max_abs
float weight_clamp_max_abs,
/* SP4 Task A14 (2026-04-30): Pearl C engagement counter via warp+
* block tree-reduce. Curiosity has 4 sub-buffer launches (W1, b1,
* W2, b2) per step. Each sub-launch's grid is sized for its own
* param count, so distinct engage_buf_offsets are required to
* avoid overwriting peers' counts:
* W1 → SP4_ENGAGE_OFFSET_CURIOSITY_W1 (= 7 × MAX_BLOCKS = 1792)
* b1 → SP4_ENGAGE_OFFSET_CURIOSITY_B1 (= 8 × MAX_BLOCKS = 2048)
* W2 → SP4_ENGAGE_OFFSET_CURIOSITY_W2 (= 9 × MAX_BLOCKS = 2304)
* b2 → SP4_ENGAGE_OFFSET_CURIOSITY_B2 (= 10 × MAX_BLOCKS = 2560)
* Host-side `pearl_c_post_adam_engagement_check(Curiosity, N)` sums
* all four ranges. `engage_buf_offset < 0` skips the writeback. */
int* __restrict__ nan_flags_buf,
int diag_slot,
int* __restrict__ clamp_engage_per_block_buf,
int engage_buf_offset
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= num_params) return;
bool active = (i < num_params);
int local_engage = 0;
float lr_bf = lr;
float beta1_bf = beta1;
float beta2_bf = beta2;
float eps_bf = eps;
float one_bf = 1.0f;
if (active) {
float lr_bf = lr;
float beta1_bf = beta1;
float beta2_bf = beta2;
float eps_bf = eps;
float one_bf = 1.0f;
float g = grads[i] / (float)batch_size;
float g = grads[i] / (float)batch_size;
/* Skip NaN/Inf gradients */
if (!isfinite((float)g)) return;
/* Skip NaN/Inf gradients — fall through to engagement reduce so
* the per-block count stays correct. */
if (isfinite((float)g)) {
m[i] = beta1_bf * m[i] + (one_bf - beta1_bf) * g;
v[i] = beta2_bf * v[i] + (one_bf - beta2_bf) * g * g;
float m_hat = m[i] / (1.0f - powf(beta1, (float)step));
float v_hat = v[i] / (1.0f - powf(beta2, (float)step));
float p = params[i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf);
m[i] = beta1_bf * m[i] + (one_bf - beta1_bf) * g;
v[i] = beta2_bf * v[i] + (one_bf - beta2_bf) * g * g;
float m_hat = m[i] / (1.0f - powf(beta1, (float)step));
float v_hat = v[i] / (1.0f - powf(beta2, (float)step));
float p = params[i] - lr_bf * m_hat / (sqrtf(v_hat) + eps_bf);
/* SP3 Mech 9 + SP4 Task A14: post-Adam |p| clamp +
* engagement counter — last step before writeback. */
if (weight_clamp_max_abs > 0.0f) {
if (fabsf(p) > weight_clamp_max_abs) {
if (nan_flags_buf != nullptr && diag_slot >= 0) {
nan_flags_buf[diag_slot] = 1;
}
local_engage = 1;
}
p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
}
/* SP3 Mech 9: post-Adam |p| clamp — last step before writeback. */
if (weight_clamp_max_abs > 0.0f) {
p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
params[i] = p;
}
}
params[i] = p;
/* SP4 Task A14: warp + block tree-reduce of `local_engage`. */
int e = local_engage;
for (int offset = 16; offset > 0; offset >>= 1)
e += __shfl_xor_sync(0xFFFFFFFF, e, offset);
__shared__ int warp_engage[8]; /* blockDim.x = 256 ⇒ 8 warps */
int warp_id = threadIdx.x / 32;
int warp_lane = threadIdx.x % 32;
if (warp_lane == 0) warp_engage[warp_id] = e;
__syncthreads();
if (warp_id == 0) {
int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
}
}
}
/* ------------------------------------------------------------------ */

12
crates/ml/src/cuda_pipeline/decision_transformer.rs
View File

@@ -984,6 +984,13 @@ impl DecisionTransformer {
// the clamp here — the kernel branch is `if (bound > 0.0f)`,
// so Adam behaves exactly as before for DT.
let weight_clamp_max_abs_dt: f32 = 0.0_f32;
// SP4 Task A14: DT is a separate offline-RL trainer outside the
// SP4 8-group taxonomy — Pearl C engagement counter disabled.
let dt_nan_flags_null: u64 = 0;
let dt_diag_slot_disabled: i32 = -1;
let dt_engage_buf_null: u64 = 0;
let dt_engage_off_disabled: i32 =
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_DISABLED;
unsafe {
stream.launch_builder(&adam.adam_update)
.arg(&params_ptr_mut)
@@ -999,6 +1006,11 @@ impl DecisionTransformer {
.arg(&l1_end_dt)
.arg(&l1_lambda_dt)
.arg(&weight_clamp_max_abs_dt) // SP3 Mech 9: disabled (no ISV in DT scope)
// SP4 Task A14: disabled (DT outside SP4 group taxonomy).
.arg(&dt_nan_flags_null)
.arg(&dt_diag_slot_disabled)
.arg(&dt_engage_buf_null)
.arg(&dt_engage_off_disabled)
.launch(lc)
.map_err(|e| MLError::ModelError(format!("DT adam: {e}")))?;
}

136
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -109,6 +109,19 @@ extern "C" __global__ void dqn_grad_norm_finalize(
* headroom. Disabled when `weight_clamp_max_abs == 0.0` (debug
* fast-path; production never sets 0).
*
* SP4 Task A14 (2026-04-30): Pearl C engagement counter. Each thread
* sets `local_engage = 1` when its parameter overshoots the bound;
* block-wide warp+block tree-reduce (mirrors `dqn_grad_norm_kernel`
* pattern, no atomicAdd per `feedback_no_atomicadd`); single block-leader
* writes per-block count to
* `clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x]`. Host
* sums across blocks → engagement_rate; rate-deficit Wiener-EMA detects
* post-clamp feedback-loop saturation. `engage_buf_offset < 0` skips
* the counter write (used by aux trainers outside the SP4 8-group
* taxonomy — DT, OFI embed, denoise, recursive_conf, sel).
* `nan_flags_buf[diag_slot] = 1` is a sticky idempotent write capturing
* "any thread overshot the bound this step" for HEALTH_DIAG.
*
* Launch config: grid=(ceil(TOTAL_PARAMS/256), 1, 1), block=(256, 1, 1).
* ══════════════════════════════════════════════════════════════════════ */
@@ -129,62 +142,105 @@ extern "C" __global__ void dqn_adam_update_kernel(
const float* __restrict__ weight_decay_mask, /* [TOTAL_PARAMS] 1.0=decay, 0.0=no decay */
int l1_end, /* param index where L1 stops (end of w_s1) */
float l1_lambda, /* L1 penalty (1e-3). 0=disabled */
float weight_clamp_max_abs /* SP3 Mech 9: ISV-driven |p_val| bound (100 * Q_ABS_REF.max(1.0)). 0=disabled. */
float weight_clamp_max_abs, /* SP3 Mech 9: ISV-driven |p_val| bound (100 * Q_ABS_REF.max(1.0)). 0=disabled. */
/* SP4 Task A14: Pearl C engagement counter. */
int* __restrict__ nan_flags_buf, /* sticky overflow flag store; nan_flags_buf[diag_slot]=1 on first thread that overshoots */
int diag_slot, /* index into nan_flags_buf (one per Adam kernel) */
int* __restrict__ clamp_engage_per_block_buf, /* [SP4_PARAM_GROUP_COUNT × MAX_BLOCKS_PER_ADAM] per-block engagement counts */
int engage_buf_offset /* per-group offset (group*MAX_BLOCKS); -1 skips writeback (aux trainers) */
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total_params) return;
/* Out-of-range threads still participate in the block reduce (with
* local_engage=0) so the block-wide tree-reduce stays correct. */
bool active = (idx < total_params);
float lr = *lr_ptr;
int t = *t_ptr;
float g_f = grads[idx];
int local_engage = 0;
float lr = active ? *lr_ptr : 0.0f;
int t = active ? *t_ptr : 1;
float g_f = active ? grads[idx] : 0.0f;
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged.
* Non-deterministic NaN can arise from IQN/CQL auxiliary gradient
* injection during early training with extreme quantile estimates. */
if (!isfinite(g_f)) return;
bool grad_finite = active && isfinite(g_f);
/* Clip using the COMPLETED L2 norm (single buffer, written by grad_norm_finalize).
* max_grad_norm read from pinned device-mapped buffer (EMA-based adaptive clip). */
float max_grad_norm = *max_grad_norm_buf;
float norm_f = *grad_norm_f32; /* Already L2 norm (sqrt applied by finalize) */
float clip_scale_f = (norm_f > max_grad_norm) ? (max_grad_norm / norm_f) : 1.0f;
float clipped_g_f = g_f * clip_scale_f;
if (grad_finite) {
/* Clip using the COMPLETED L2 norm (single buffer, written by grad_norm_finalize).
* max_grad_norm read from pinned device-mapped buffer (EMA-based adaptive clip). */
float max_grad_norm = *max_grad_norm_buf;
float norm_f = *grad_norm_f32; /* Already L2 norm (sqrt applied by finalize) */
float clip_scale_f = (norm_f > max_grad_norm) ? (max_grad_norm / norm_f) : 1.0f;
float clipped_g_f = g_f * clip_scale_f;
/* Adam bias correction: native float — no bf16 rounding issues. */
float beta1_t = fmaxf(1.0f - powf(beta1, (float)t), 1e-8f);
float beta2_t = fmaxf(1.0f - powf(beta2, (float)t), 1e-8f);
/* Adam bias correction: native float — no bf16 rounding issues. */
float beta1_t = fmaxf(1.0f - powf(beta1, (float)t), 1e-8f);
float beta2_t = fmaxf(1.0f - powf(beta2, (float)t), 1e-8f);
/* Adam moment updates in native float — full precision, no bf16 conversion. */
float m_i = beta1 * m[idx] + (1.0f - beta1) * clipped_g_f;
float v_i = beta2 * v[idx] + (1.0f - beta2) * clipped_g_f * clipped_g_f;
m[idx] = m_i;
v[idx] = v_i;
/* Adam moment updates in native float — full precision, no bf16 conversion. */
float m_i = beta1 * m[idx] + (1.0f - beta1) * clipped_g_f;
float v_i = beta2 * v[idx] + (1.0f - beta2) * clipped_g_f * clipped_g_f;
m[idx] = m_i;
v[idx] = v_i;
float m_hat = m_i / beta1_t;
float v_hat = v_i / beta2_t;
float m_hat = m_i / beta1_t;
float v_hat = v_i / beta2_t;
/* AdamW weight decay (decoupled) — masked: trunk+value=1.0, branches=0.0. */
float p_val = params[idx];
float wd = weight_decay * weight_decay_mask[idx];
p_val = p_val - lr * (m_hat / (sqrtf(v_hat) + epsilon) + wd * p_val);
/* AdamW weight decay (decoupled) — masked: trunk+value=1.0, branches=0.0. */
float p_val = params[idx];
float wd = weight_decay * weight_decay_mask[idx];
p_val = p_val - lr * (m_hat / (sqrtf(v_hat) + epsilon) + wd * p_val);
/* G8: L1 proximal step on w_s1 — automatic feature selection */
if (idx < l1_end && l1_lambda > 0.0f) {
float sign_val = (p_val > 0.0f) ? 1.0f : -1.0f;
p_val = sign_val * fmaxf(fabsf(p_val) - lr * l1_lambda, 0.0f);
/* G8: L1 proximal step on w_s1 — automatic feature selection */
if (idx < l1_end && l1_lambda > 0.0f) {
float sign_val = (p_val > 0.0f) ? 1.0f : -1.0f;
p_val = sign_val * fmaxf(fabsf(p_val) - lr * l1_lambda, 0.0f);
}
/* SP3 Mech 9 + SP4 Task A14: post-Adam weight clamp + Pearl C
* engagement counter. ISV-driven bound (100 × Q_ABS_REF.max(1.0))
* computed host-side and passed in. Targets cuBLAS sgemm f32
* accumulator overflow at slots 26+32 — bounds weights 1 OOM below
* slot 44-45 diagnostic threshold (1e3 × Q_ABS_REF) so the
* diagnostic retains regression-sentinel headroom. Disabled when
* bound==0 (debug fast-path; production never sets 0).
*
* `local_engage = 1` when |p_val| would overshoot — captured into
* the per-block tree-reduce below. The sticky `nan_flags_buf` write
* is idempotent (always writes 1, never 0) so no race exists. */
if (weight_clamp_max_abs > 0.0f) {
if (fabsf(p_val) > weight_clamp_max_abs) {
if (nan_flags_buf != nullptr && diag_slot >= 0) {
nan_flags_buf[diag_slot] = 1;
}
local_engage = 1;
}
p_val = fminf(fmaxf(p_val, -weight_clamp_max_abs), weight_clamp_max_abs);
}
params[idx] = p_val;
}
/* SP3 Mech 9: post-Adam weight clamp. ISV-driven bound
* (100 × Q_ABS_REF.max(1.0)) computed host-side and passed in. Targets
* cuBLAS sgemm f32 accumulator overflow at slots 26+32 — bounds weights
* 1 OOM below slot 44-45 diagnostic threshold (1e3 × Q_ABS_REF) so the
* diagnostic retains regression-sentinel headroom. Disabled when
* bound==0 (debug fast-path; production never sets 0). */
if (weight_clamp_max_abs > 0.0f) {
p_val = fminf(fmaxf(p_val, -weight_clamp_max_abs), weight_clamp_max_abs);
}
/* SP4 Task A14: warp + block tree-reduce of `local_engage`. Mirrors
* `dqn_grad_norm_kernel`'s pattern — no atomicAdd. blockDim.x=256 ⇒ 8
* warps; warp_sums fits in __shared__[8]. */
int e = local_engage;
for (int offset = 16; offset > 0; offset >>= 1)
e += __shfl_xor_sync(0xFFFFFFFF, e, offset);
params[idx] = p_val;
__shared__ int warp_engage[8];
int warp_id = threadIdx.x / 32;
int warp_lane = threadIdx.x % 32;
if (warp_lane == 0) warp_engage[warp_id] = e;
__syncthreads();
if (warp_id == 0) {
int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
}
}
}
/* ══════════════════════════════════════════════════════════════════════

58
crates/ml/src/cuda_pipeline/gpu_attention.rs
View File

@@ -230,6 +230,13 @@ pub struct GpuAttention {
t_dev_ptr: u64,
/// Total attention parameters count.
total_params: usize,
// ── SP4 Task A14 Pearl C engagement-counter buffer pointers ────────
/// Mapped-pinned `nan_flags_buf[55]` device pointer (owned by
/// `GpuDqnTrainer`). 0 = unset → kernel skips per-block writeback.
pearl_c_nan_flags_ptr: u64,
/// Mapped-pinned `clamp_engage_per_block_buf` device pointer.
pearl_c_engage_buf_ptr: u64,
}
impl GpuAttention {
@@ -534,9 +541,32 @@ impl GpuAttention {
t_pinned,
t_dev_ptr,
total_params,
// SP4 Task A14: Pearl C buffers wired post-construction by
// FusedTrainingCtx via `set_pearl_c_buffers`. Until then,
// null pointer + disabled offset signals the kernel to skip
// per-block engagement-counter writeback.
pearl_c_nan_flags_ptr: 0,
pearl_c_engage_buf_ptr: 0,
})
}
/// SP4 Task A14: wire the trainer's mapped-pinned Pearl C buffers
/// (owned by `GpuDqnTrainer`) into this attention head. Called once
/// post-construction by `FusedTrainingCtx`. After wiring, every
/// `adam_step` records per-block clamp-engagement counts at offset
/// `ParamGroup::Attn × MAX_BLOCKS_PER_ADAM` (= 6 × 256 = 1536).
pub fn set_pearl_c_buffers(&mut self, nan_flags: u64, engage_buf: u64) {
self.pearl_c_nan_flags_ptr = nan_flags;
self.pearl_c_engage_buf_ptr = engage_buf;
}
/// SP4 Task A15: parameter count accessor for the host-side Pearl C
/// engagement-rate check. Returns `config.total_params()` cached at
/// construction.
pub fn total_params(&self) -> usize {
self.total_params
}
/// Apply multi-head attention to batch of states with OFI embedding.
///
/// Input is constructed by concatenating states[D,B] + ofi_embed[10,B] into
@@ -1084,6 +1114,30 @@ impl GpuAttention {
let wd = 1e-5_f32;
let t_ptr = self.t_dev_ptr;
// SP4 Task A14: Pearl C engagement-counter args. Both
// `GpuAttention` (this module) and `GpuTlob` share the same
// `attn_adam_kernel` cubin. For Layer A, GpuAttention writes to
// its slot at offset 6×256=1536 (ParamGroup::Attn). GpuTlob
// passes SP4_ENGAGE_OFFSET_DISABLED (= -1) so it silently
// skips Pearl C bookkeeping until Layer B refactors them apart.
// If `set_pearl_c_buffers` has not been called (FusedTrainingCtx
// wiring), pass null + disabled offset so the kernel skips.
let pearl_c_active = self.pearl_c_engage_buf_ptr != 0
&& self.pearl_c_nan_flags_ptr != 0;
let nan_flags_ptr: u64 = self.pearl_c_nan_flags_ptr;
let diag_slot: i32 = if pearl_c_active {
crate::cuda_pipeline::SP4_ATTN_ADAM_DIAG_SLOT
} else {
-1
};
let engage_buf_ptr: u64 = self.pearl_c_engage_buf_ptr;
let engage_buf_offset: i32 = if pearl_c_active {
(crate::cuda_pipeline::ParamGroup::Attn as i32)
* (crate::cuda_pipeline::MAX_BLOCKS_PER_ADAM as i32)
} else {
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_DISABLED
};
unsafe {
stream
.launch_builder(&self.adam_kernel)
@@ -1101,6 +1155,10 @@ impl GpuAttention {
.arg(&t_ptr)
.arg(&tp_i32)
.arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp
.arg(&nan_flags_ptr) // SP4 A14: sticky engage flag store
.arg(&diag_slot) // SP4 A14: diag slot (53 = Attn)
.arg(&engage_buf_ptr) // SP4 A14: per-block engage buf
.arg(&engage_buf_offset) // SP4 A14: group offset (6 = Attn)
.launch(adam_cfg)
.map_err(|e| MLError::ModelError(format!("attention adam kernel: {e}")))?;
}

108
crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs
View File

@@ -494,6 +494,13 @@ pub struct GpuCuriosityTrainer {
// Maximum number of samples the buffers were allocated for
buf_capacity: usize,
// ── SP4 Task A14 Pearl C engagement-counter buffer pointers ────────
/// Mapped-pinned `nan_flags_buf[55]` device pointer (owned by
/// `GpuDqnTrainer`). 0 = unset → kernel skips per-block writeback.
pearl_c_nan_flags_ptr: u64,
/// Mapped-pinned `clamp_engage_per_block_buf` device pointer.
pearl_c_engage_buf_ptr: u64,
}
/// Launch Adam optimizer step for one parameter group.
@@ -502,6 +509,7 @@ pub struct GpuCuriosityTrainer {
/// bound for `curiosity_adam_step`. Caller threads `100 × Q_ABS_REF.max(1.0)`
/// from training_loop down through `train_on_collector_buffers`. 0=disabled
/// (debug only).
#[allow(clippy::too_many_arguments)]
fn launch_adam_step(
stream: &CudaStream,
adam_func: &CudaFunction,
@@ -513,9 +521,25 @@ fn launch_adam_step(
batch_size: usize,
step: i32,
weight_clamp_max_abs: f32,
/* SP4 Task A14: per-sub-launch Pearl C engagement-counter args.
* `engage_buf_offset` is unique per sub-buffer (W1/b1/W2/b2) so
* the 4 sub-launches don't overwrite each other's per-block counts.
* Pass `nan_flags_ptr=0` + `engage_buf_ptr=0` +
* `engage_buf_offset=SP4_ENGAGE_OFFSET_DISABLED` to skip writeback. */
nan_flags_ptr: u64,
diag_slot: i32,
engage_buf_ptr: u64,
engage_buf_offset: i32,
) -> Result<(), MLError> {
let blocks = ((num_params as u32) + 255) / 256;
debug_assert!(
(blocks as usize) <= crate::cuda_pipeline::MAX_BLOCKS_PER_ADAM,
"curiosity_adam blocks={} exceeds MAX_BLOCKS_PER_ADAM={} — Pearl C \
clamp_engage_per_block_buf cannot accommodate.",
blocks, crate::cuda_pipeline::MAX_BLOCKS_PER_ADAM
);
let cfg = LaunchConfig {
grid_dim: (((num_params as u32) + 255) / 256, 1, 1),
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
@@ -536,6 +560,10 @@ fn launch_adam_step(
.arg(&ADAM_EPS)
.arg(&step)
.arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp
.arg(&nan_flags_ptr) // SP4 A14: sticky engage flag store
.arg(&diag_slot) // SP4 A14: diag slot (54 = Curiosity)
.arg(&engage_buf_ptr) // SP4 A14: per-block engage buf
.arg(&engage_buf_offset) // SP4 A14: per-sub-launch offset
.launch(cfg)
.map_err(|e| {
MLError::ModelError(format!("curiosity_adam_step launch: {e}"))
@@ -681,9 +709,34 @@ impl GpuCuriosityTrainer {
step: 0,
state_dim,
buf_capacity: max_samples,
// SP4 Task A14: Pearl C buffers wired post-construction by
// FusedTrainingCtx via `set_pearl_c_buffers`. Until then,
// null pointer + disabled offset signals each curiosity
// sub-launch to skip per-block engagement writeback.
pearl_c_nan_flags_ptr: 0,
pearl_c_engage_buf_ptr: 0,
})
}
/// SP4 Task A14: thread the main DQN trainer's mapped-pinned Pearl C
/// buffers into this curiosity trainer. Called once post-construction
/// by `FusedTrainingCtx`. After wiring, the 4 sub-launches (W1, b1,
/// W2, b2) record per-block clamp-engagement counts at distinct
/// offsets `SP4_ENGAGE_OFFSET_CURIOSITY_{W1,B1,W2,B2}` so they don't
/// overwrite each other. Pass 0 for either pointer to disable.
pub fn set_pearl_c_buffers(&mut self, nan_flags: u64, engage_buf: u64) {
self.pearl_c_nan_flags_ptr = nan_flags;
self.pearl_c_engage_buf_ptr = engage_buf;
}
/// SP4 Task A15: total parameter count summed across the 4 sub-buffers
/// (W1 + b1 + W2 + b2). Used by the host-side Pearl C engagement-rate
/// check, which sums per-block counts across all 4 sub-launch ranges
/// before dividing by total params.
pub const fn total_params() -> usize {
CUR_W1_LEN + CUR_B1_LEN + CUR_W2_LEN + CUR_B2_LEN
}
/// Train the curiosity forward model on GPU-resident experience data.
///
/// Performs one training step via cuBLAS GEMM pipeline (forward + backward + Adam).
@@ -920,10 +973,55 @@ impl GpuCuriosityTrainer {
self.step += 1;
let step = self.step;
launch_adam_step(stream, &self.adam_func, &mut weights.w1, &self.grad_w1, CUR_W1_LEN, &mut self.adam_m_w1, &mut self.adam_v_w1, n_train, step, weight_clamp_max_abs)?;
launch_adam_step(stream, &self.adam_func, &mut weights.b1, &self.grad_b1, CUR_B1_LEN, &mut self.adam_m_b1, &mut self.adam_v_b1, n_train, step, weight_clamp_max_abs)?;
launch_adam_step(stream, &self.adam_func, &mut weights.w2, &self.grad_w2, CUR_W2_LEN, &mut self.adam_m_w2, &mut self.adam_v_w2, n_train, step, weight_clamp_max_abs)?;
launch_adam_step(stream, &self.adam_func, &mut weights.b2, &self.grad_b2, CUR_B2_LEN, &mut self.adam_m_b2, &mut self.adam_v_b2, n_train, step, weight_clamp_max_abs)?;
// SP4 Task A14: each sub-launch gets its own engage_buf_offset so
// the 4 launches' per-block counts don't overlap. Pearl C is
// active only when both buffer pointers are wired by
// FusedTrainingCtx via `set_pearl_c_buffers`.
let pearl_c_active = self.pearl_c_engage_buf_ptr != 0
&& self.pearl_c_nan_flags_ptr != 0;
let nan_flags_ptr = self.pearl_c_nan_flags_ptr;
let engage_buf_ptr = self.pearl_c_engage_buf_ptr;
let diag_slot = if pearl_c_active {
crate::cuda_pipeline::SP4_CURIOSITY_ADAM_DIAG_SLOT
} else {
-1
};
let (off_w1, off_b1, off_w2, off_b2) = if pearl_c_active {
(
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_CURIOSITY_W1,
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_CURIOSITY_B1,
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_CURIOSITY_W2,
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_CURIOSITY_B2,
)
} else {
let off = crate::cuda_pipeline::SP4_ENGAGE_OFFSET_DISABLED;
(off, off, off, off)
};
launch_adam_step(
stream, &self.adam_func, &mut weights.w1, &self.grad_w1,
CUR_W1_LEN, &mut self.adam_m_w1, &mut self.adam_v_w1,
n_train, step, weight_clamp_max_abs,
nan_flags_ptr, diag_slot, engage_buf_ptr, off_w1,
)?;
launch_adam_step(
stream, &self.adam_func, &mut weights.b1, &self.grad_b1,
CUR_B1_LEN, &mut self.adam_m_b1, &mut self.adam_v_b1,
n_train, step, weight_clamp_max_abs,
nan_flags_ptr, diag_slot, engage_buf_ptr, off_b1,
)?;
launch_adam_step(
stream, &self.adam_func, &mut weights.w2, &self.grad_w2,
CUR_W2_LEN, &mut self.adam_m_w2, &mut self.adam_v_w2,
n_train, step, weight_clamp_max_abs,
nan_flags_ptr, diag_slot, engage_buf_ptr, off_w2,
)?;
launch_adam_step(
stream, &self.adam_func, &mut weights.b2, &self.grad_b2,
CUR_B2_LEN, &mut self.adam_m_b2, &mut self.adam_v_b2,
n_train, step, weight_clamp_max_abs,
nan_flags_ptr, diag_slot, engage_buf_ptr, off_b2,
)?;
debug!(step, n_train, "curiosity GPU training step complete (cuBLAS GEMM pipeline)");

322
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -66,6 +66,14 @@ use super::gpu_aux_heads::{
use super::gpu_health_diag::GpuHealthDiag;
use super::gpu_moe_head::GpuMoeHead;
use super::mapped_pinned::{MappedF32Buffer, MappedI32Buffer, MappedU32Buffer, MappedU64Buffer};
// SP4 Task A14: shared constants (engagement-counter buffer layout +
// per-Adam-kernel diag-slot allocations) used by `new()` constructor +
// the 5 Adam launch sites + the host-side Pearl C rate-deficit check.
use super::sp4_isv_slots::{
SP4_PARAM_GROUP_COUNT, MAX_BLOCKS_PER_ADAM, SP4_ENGAGE_BUF_LEN,
SP4_ENGAGE_OFFSET_DISABLED, SP4_NAN_FLAGS_END,
SP4_DQN_ADAM_DIAG_SLOT, ParamGroup,
};
// ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ──────
pub(crate) static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin"));
@@ -2907,6 +2915,27 @@ pub struct GpuDqnTrainer {
/// = 2048 ints. Mapped-pinned i32.
pub(crate) clamp_engage_per_block_buf: MappedI32Buffer,
/// SP4 Task A15: Pearl C rate-deficit Wiener-state buffer.
/// `SP4_PARAM_GROUP_COUNT × 3 = 24` floats laid out as
/// `[sample_var, diff_var, x_lag]` per param group. Tracks the
/// engagement-rate-deficit (= engagement_rate - 0.01 theoretical p99
/// baseline) via Pearls A+D over time. Mapped-pinned f32 — host-side
/// `pearl_c_post_adam_engagement_check` reads/writes via host_ptr.
/// Reset to all-zeros at fold boundary via the state-reset registry
/// (Task A12). Layer A: this state is computed and tracked but not
/// yet consumed — Layer B will read sustained-positive rate-deficit
/// to force-bump `ISV[WEIGHT_BOUND[group]]`.
pub(crate) pearl_c_rate_deficit_state_buf: MappedF32Buffer,
/// SP4 Task A15: Pearl C rate-deficit EMA (host-only, one f32 per
/// param group). Acts as the `prev_x_mean` for `pearls_ad_update`
/// applied to engagement-rate deficits each Adam step. Pearls A+D
/// state for those EMAs lives in `pearl_c_rate_deficit_state_buf`
/// (mapped pinned). Host-only because rate-deficit reduce + EMA is
/// post-Adam diagnostic, never read on device. Reset to all-zeros
/// at fold boundary via the state-reset registry (Task A12).
pub(crate) pearl_c_rate_deficit_ema: [f32; SP4_PARAM_GROUP_COUNT],
/// SP4 producer step-observation scratch — each of the 47 producers
/// (40 SP4 + 7 retrofit) writes its per-step `step_observation` to
/// its dedicated slot here. Host then applies Pearls A+D
@@ -6047,6 +6076,12 @@ impl GpuDqnTrainer {
// SP3 Mech 9: ISV-driven post-Adam weight clamp (see `launch_adam_update`).
let q_abs_ref_eff = self.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32);
let weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff;
// SP4 Task A14: aux trainer (OFI embed) is outside the SP4 8-group
// taxonomy — disable Pearl C engagement counter for this launch.
let _aux_nan_flags_null: u64 = 0;
let _aux_diag_slot_disabled: i32 = -1;
let _aux_engage_buf_null: u64 = 0;
let _aux_engage_off_disabled: i32 = SP4_ENGAGE_OFFSET_DISABLED;
let blocks = ((OFI_EMBED_TOTAL_PARAMS as u32 + 255) / 256).max(1);
unsafe {
self.stream
@@ -6068,6 +6103,13 @@ impl GpuDqnTrainer {
.arg(&l1_end_aux)
.arg(&l1_lambda_aux)
.arg(&weight_clamp_max_abs) // SP3 Mech 9: post-Adam |p_val| clamp
// SP4 Task A14: aux trainers (OFI embed) live outside the
// SP4 8-group taxonomy — pass null nan_flags + disabled
// engage_buf_offset so the kernel skips per-block writeback.
.arg(&_aux_nan_flags_null)
.arg(&_aux_diag_slot_disabled)
.arg(&_aux_engage_buf_null)
.arg(&_aux_engage_off_disabled)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
@@ -12723,9 +12765,12 @@ impl GpuDqnTrainer {
// the F1 step-3540 NaN root). SP3 Mech 10 (slot 49): one additional
// slot for post-trunk h_s2 activation max-abs sentinel — fires at
// 1e3 × H_S2_RMS_EMA.max(1.0) and stays quiet under normal Mech 10
// clamp operation (clamp at 100×, diagnostic at 1000×). Buffer size
// reviewable by SP2 framework codification.
let nan_flags_buf = stream.alloc_zeros::<i32>(50)
// clamp operation (clamp at 100×, diagnostic at 1000×). SP4 Task
// A14 (slots 50-54): per-Adam-kernel sticky engagement-overflow
// flags (one per of dqn/iqn/iql/attn/curiosity Adam kernels).
// Sized via SP4_NAN_FLAGS_END = 55. Buffer size reviewable by SP2
// framework codification.
let nan_flags_buf = stream.alloc_zeros::<i32>(SP4_NAN_FLAGS_END)
.map_err(|e| MLError::ModelError(format!("nan_flags alloc: {e}")))?;
// SP2 + SP3 Mech 5 + Mech 10: fused NaN/threshold-check (ptr, len) tables.
@@ -12763,12 +12808,12 @@ impl GpuDqnTrainer {
// Per `feedback_no_htod_htoh_only_mapped_pinned`: mapped pinned
// (`cuMemHostAlloc(DEVICEMAP|PORTABLE)`) is the only allowed CPU↔GPU
// path for these state buffers.
const SP4_PARAM_GROUP_COUNT: usize = 8;
const SP4_PRODUCER_COUNT: usize = 47; // 40 SP4 + 7 retrofit existing producers
const SP4_WIENER_FLOATS_PER_SLOT: usize = 3; // sample_var, diff_var, x_lag
const SP4_WIENER_TOTAL_FLOATS: usize = SP4_PRODUCER_COUNT * SP4_WIENER_FLOATS_PER_SLOT; // 141
const MAX_BLOCKS_PER_ADAM: usize = 256;
const SP4_ENGAGE_BUF_LEN: usize = SP4_PARAM_GROUP_COUNT * MAX_BLOCKS_PER_ADAM; // 2048
// SP4 Task A14: SP4_PARAM_GROUP_COUNT, MAX_BLOCKS_PER_ADAM, and
// SP4_ENGAGE_BUF_LEN are now public constants in `sp4_isv_slots`
// so launch sites + helpers can reference them without re-declaring.
let wiener_state_buf = unsafe { MappedF32Buffer::new(SP4_WIENER_TOTAL_FLOATS) }
.map_err(|e| MLError::ModelError(format!("SP4 wiener_state_buf alloc: {e}")))?;
@@ -12779,6 +12824,18 @@ impl GpuDqnTrainer {
let producer_step_scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
.map_err(|e| MLError::ModelError(format!("SP4 producer_step_scratch_buf alloc: {e}")))?;
// SP4 Task A15: dedicated Pearl C rate-deficit Wiener-state
// buffer (8 groups × 3 floats: sample_var, diff_var, x_lag).
// Separate from `wiener_state_buf` to keep SP4-bound producer
// state isolated from Pearl C diagnostic state — clean reset
// semantics + independent fold-boundary handling.
const PEARL_C_RATE_DEFICIT_FLOATS: usize = SP4_PARAM_GROUP_COUNT * SP4_WIENER_FLOATS_PER_SLOT;
let pearl_c_rate_deficit_state_buf =
unsafe { MappedF32Buffer::new(PEARL_C_RATE_DEFICIT_FLOATS) }
.map_err(|e| MLError::ModelError(format!(
"SP4 pearl_c_rate_deficit_state_buf alloc: {e}"
)))?;
// SP4 Task A7 fix-up #2: oracle sub-buffer descriptor tables.
// `SP4_ORACLE_TABLE_MAX_SUB = 4` covers Curiosity's 4 sub-buffers
// (groups 0-6 use only the first slot). Tables are reused across
@@ -14188,6 +14245,8 @@ impl GpuDqnTrainer {
nan_flags_buf,
wiener_state_buf,
clamp_engage_per_block_buf,
pearl_c_rate_deficit_state_buf,
pearl_c_rate_deficit_ema: [0.0_f32; SP4_PARAM_GROUP_COUNT],
producer_step_scratch_buf,
oracle_subbuf_table_buf,
oracle_subbuf_counts_buf,
@@ -18353,6 +18412,12 @@ impl GpuDqnTrainer {
// SP3 Mech 9: ISV-driven post-Adam weight clamp (see `launch_adam_update`).
let q_abs_ref_eff = self.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32);
let weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff;
// SP4 Task A14: aux trainer (denoise) is outside the SP4 8-group
// taxonomy — Pearl C engagement counter disabled for this launch.
let _aux_nan_flags_null: u64 = 0;
let _aux_diag_slot_disabled: i32 = -1;
let _aux_engage_buf_null: u64 = 0;
let _aux_engage_off_disabled: i32 = SP4_ENGAGE_OFFSET_DISABLED;
let blocks = ((DENOISE_TOTAL_PARAMS as u32 + 255) / 256).max(1);
unsafe {
self.stream
@@ -18374,6 +18439,11 @@ impl GpuDqnTrainer {
.arg(&l1_end_aux)
.arg(&l1_lambda_aux)
.arg(&weight_clamp_max_abs) // SP3 Mech 9: post-Adam |p_val| clamp
// SP4 Task A14: disabled — denoise outside SP4 group taxonomy.
.arg(&_aux_nan_flags_null)
.arg(&_aux_diag_slot_disabled)
.arg(&_aux_engage_buf_null)
.arg(&_aux_engage_off_disabled)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
@@ -20895,7 +20965,20 @@ impl GpuDqnTrainer {
let weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff;
let count = self.total_params as i32;
// SP4 Task A14: cap blocks at MAX_BLOCKS_PER_ADAM (256). Per-block
// engagement-counter buffer is sized for that ceiling; running with
// more blocks would have the kernel write past the per-group slice.
// For DQN main-trainer total_params can exceed 256×256=65 536, so we
// assert the bound here. If it ever fires, the buffer needs sizing
// up before this Adam launch can grow.
let blocks = ((self.total_params + 255) / 256) as u32;
debug_assert!(
(blocks as usize) <= MAX_BLOCKS_PER_ADAM,
"dqn_adam blocks={} exceeds MAX_BLOCKS_PER_ADAM={} — Pearl C \
clamp_engage_per_block_buf cannot accommodate; resize buffer \
or chunk the launch.",
blocks, MAX_BLOCKS_PER_ADAM
);
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
@@ -20906,6 +20989,17 @@ impl GpuDqnTrainer {
let m_ptr = self.ptrs.m_buf;
let v_ptr = self.ptrs.v_buf;
// SP4 Task A14: Pearl C engagement-counter args. The single main
// DQN Adam launch covers param-groups 0/1/2 (trunk/value/branches);
// we account engagement under group 0 (DqnTrunk) for Layer A — see
// `pearl_c_post_adam_engagement_check` doc-comment for the
// sub-launch deferral rationale.
let nan_flags_ptr = self.nan_flags_buf.raw_ptr();
let diag_slot: i32 = SP4_DQN_ADAM_DIAG_SLOT;
let engage_buf_ptr: u64 = self.clamp_engage_per_block_buf.dev_ptr;
let engage_buf_offset: i32 =
(ParamGroup::DqnTrunk.idx() * MAX_BLOCKS_PER_ADAM) as i32;
// Safety: argument order matches the extern "C" kernel signature exactly.
// grad_norm_buf contains L2 norm from launch_grad_norm_finalize().
unsafe {
@@ -20928,6 +21022,10 @@ impl GpuDqnTrainer {
.arg(&l1_end) // G8: L1 boundary (end of w_s1)
.arg(&l1_lambda) // G8: L1 penalty strength
.arg(&weight_clamp_max_abs) // SP3 Mech 9: post-Adam |p_val| clamp
.arg(&nan_flags_ptr) // SP4 A14: sticky engage flag store
.arg(&diag_slot) // SP4 A14: diag slot (50 for dqn_adam)
.arg(&engage_buf_ptr) // SP4 A14: per-block engagement buf
.arg(&engage_buf_offset) // SP4 A14: group offset (0 = DqnTrunk)
.launch(cfg)
.map_err(|e| {
MLError::ModelError(format!("dqn_adam_update_kernel launch: {e}"))
@@ -20937,6 +21035,141 @@ impl GpuDqnTrainer {
Ok(())
}
/// SP4 Task A15: host-side Pearl C engagement-rate-deficit check.
///
/// Reads per-block engagement counts from
/// `clamp_engage_per_block_buf` (mapped pinned, zero-copy host
/// access), reduces to a single rate, computes the deficit relative
/// to the theoretical p99 baseline (1%), then applies Pearls A + D
/// (`pearls_ad_update`) to maintain the rate-deficit EMA per param
/// group. For Layer A this method only updates the diagnostic EMA
/// and emits a `tracing::debug!` line when the EMA exceeds 0.005
/// (sustained over-engagement). Layer B will read the EMA and
/// force-bump `ISV[WEIGHT_BOUND[group]]`.
///
/// Curiosity reads four sub-launch ranges
/// (`SP4_ENGAGE_OFFSET_CURIOSITY_{W1,B1,W2,B2}`) so the four
/// internal Adam updates' counts are summed into a single group
/// engagement rate. All other groups read the single contiguous
/// `[group_idx × MAX_BLOCKS_PER_ADAM, +MAX_BLOCKS_PER_ADAM)` range.
///
/// **Layer A limitation #3** (groups 1/2 accounting): the main DQN
/// Adam covers param groups 0 (Trunk), 1 (Value), and 2 (Branches)
/// in a single fused launch — engagement is recorded under group 0
/// only. Layer B will split per-group sub-launches.
pub(crate) fn pearl_c_post_adam_engagement_check(
&mut self,
group: ParamGroup,
param_count: usize,
) -> Result<(), MLError> {
use super::sp4_wiener_ema::{pearls_ad_update, WienerState};
if param_count == 0 {
return Ok(());
}
let group_idx = group.idx();
// Curiosity sums 4 sub-launch ranges (see launch_adam_step in
// gpu_curiosity_trainer.rs). All other groups read a single
// contiguous range starting at `group_idx × MAX_BLOCKS_PER_ADAM`.
let block_offsets: &[usize] = match group {
ParamGroup::Curiosity => &[
7 * MAX_BLOCKS_PER_ADAM, // W1
8 * MAX_BLOCKS_PER_ADAM, // b1
9 * MAX_BLOCKS_PER_ADAM, // W2
10 * MAX_BLOCKS_PER_ADAM, // b2
],
_ => &[
/* group_idx * MAX_BLOCKS_PER_ADAM — fall back to
* runtime computation in the loop body. */
0
],
};
// Reduce per-block counts → total engagement count.
let mut total_engage: i64 = 0;
// Safety: `clamp_engage_per_block_buf.host_ptr` is a mapped-pinned
// i32 buffer of `SP4_ENGAGE_BUF_LEN` elements — valid for the
// lifetime of the trainer. Reads are zero-copy via the device-
// mapped page; the kernel writes were ordered by Adam-launch
// synchronization upstream of this call.
unsafe {
let p = self.clamp_engage_per_block_buf.host_ptr as *const i32;
if matches!(group, ParamGroup::Curiosity) {
for &base in block_offsets {
for b in 0..MAX_BLOCKS_PER_ADAM {
total_engage += *p.add(base + b) as i64;
}
}
} else {
let base = group_idx * MAX_BLOCKS_PER_ADAM;
for b in 0..MAX_BLOCKS_PER_ADAM {
total_engage += *p.add(base + b) as i64;
}
}
}
let engagement_rate = (total_engage as f32) / (param_count as f32);
let rate_deficit = engagement_rate - 0.01_f32;
// Apply Pearls A + D to the rate-deficit time series. The EMA
// surrogate (`prev_x_mean`) lives in `pearl_c_rate_deficit_ema`
// (host-only), and the (sample_var, diff_var, x_lag) Wiener
// triple lives in `pearl_c_rate_deficit_state_buf` at offset
// `group_idx * 3`.
let prev_x_mean = self.pearl_c_rate_deficit_ema[group_idx];
let mut state = unsafe {
let p = self.pearl_c_rate_deficit_state_buf.host_ptr as *const f32;
let off = group_idx * 3;
WienerState {
sample_var: *p.add(off),
diff_var: *p.add(off + 1),
x_lag: *p.add(off + 2),
}
};
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, rate_deficit);
self.pearl_c_rate_deficit_ema[group_idx] = new_x_mean;
unsafe {
let p = self.pearl_c_rate_deficit_state_buf.host_ptr as *mut f32;
let off = group_idx * 3;
*p.add(off) = state.sample_var;
*p.add(off + 1) = state.diff_var;
*p.add(off + 2) = state.x_lag;
}
// Layer A: log only on sustained over-engagement. Layer B will
// mutate `ISV[WEIGHT_BOUND[group]]` instead.
if new_x_mean > 0.005_f32 {
tracing::debug!(
target: "sp4.pearl_c",
group = ?group,
param_count,
engagement_rate,
rate_deficit,
rate_deficit_ema = new_x_mean,
"SP4 Pearl C: sustained engagement-rate deficit \
(>0.005 EMA) — Layer A diagnostic; Layer B will \
force-bump WEIGHT_BOUND."
);
}
// Zero per-block counts post-read so next step starts clean.
// Curiosity must zero all 4 sub-ranges; others zero one range.
unsafe {
let p = self.clamp_engage_per_block_buf.host_ptr as *mut i32;
if matches!(group, ParamGroup::Curiosity) {
for &base in block_offsets {
std::ptr::write_bytes(p.add(base), 0, MAX_BLOCKS_PER_ADAM);
}
} else {
let base = group_idx * MAX_BLOCKS_PER_ADAM;
std::ptr::write_bytes(p.add(base), 0, MAX_BLOCKS_PER_ADAM);
}
}
Ok(())
}
// ═══════════════════════════════════════════════════════════════════
// Weight flattening / unflattening (GPU d2d — zero host roundtrip)
// ═══════════════════════════════════════════════════════════════════
@@ -21941,14 +22174,18 @@ impl GpuDqnTrainer {
/// SP4 Layer A Task A12 (2026-04-30): bulk-zero the Pearl C
/// engagement-counter buffer at fold boundary.
///
/// `clamp_engage_per_block_buf` is a 2048-int mapped-pinned buffer
/// (`SP4_PARAM_GROUP_COUNT × MAX_BLOCKS_PER_ADAM = 8 × 256`) counting
/// post-Adam clamp triggers per (param-group, Adam-launch-block). Host
/// reduces across blocks → engagement_rate; rate-deficit Wiener EMA
/// detects post-clamp feedback-loop saturation. At fold boundary all
/// counters reset to 0 so the new fold's engagement-rate tracking
/// restarts from a clean baseline (no inherited fold-N saturation
/// state biasing the new fold's first per-step rate readings).
/// `clamp_engage_per_block_buf` is a 2816-int mapped-pinned buffer
/// (`(SP4_PARAM_GROUP_COUNT + 3 curiosity sub-launches) ×
/// MAX_BLOCKS_PER_ADAM = 11 × 256`) counting post-Adam clamp
/// triggers per (param-group, Adam-launch-block). Host reduces
/// across blocks → engagement_rate; rate-deficit Wiener EMA detects
/// post-clamp feedback-loop saturation. The buffer length grew
/// 2048→2816 in Task A14/A15 to give Curiosity's 4 sub-launches
/// (W1/b1/W2/b2) distinct offset slots avoiding per-block index
/// collisions. At fold boundary all counters reset to 0 so the new
/// fold's engagement-rate tracking restarts from a clean baseline
/// (no inherited fold-N saturation state biasing the new fold's
/// first per-step rate readings).
///
/// Constructor zero-init via `MappedI32Buffer::new` is the cold-start;
/// this method re-applies the same byte-zero at every subsequent fold
@@ -21957,7 +22194,7 @@ impl GpuDqnTrainer {
// Safety: `host_ptr` is a mapped-pinned `*mut i32` of length
// `self.clamp_engage_per_block_buf.len` returned by `cuMemHostAlloc`.
// i32 zero == bytewise-zero, so `write_bytes(.., 0, ..)` produces
// 2048 valid i32 zeros.
// SP4_ENGAGE_BUF_LEN=2816 valid i32 zeros.
unsafe {
std::ptr::write_bytes(
self.clamp_engage_per_block_buf.host_ptr as *mut u8,
@@ -21967,6 +22204,50 @@ impl GpuDqnTrainer {
}
}
/// SP4 Task A15: bulk-zero the Pearl C rate-deficit Wiener-state
/// buffer at fold boundary.
///
/// `pearl_c_rate_deficit_state_buf` is a 24-float mapped-pinned
/// buffer = `SP4_PARAM_GROUP_COUNT × {sample_var, diff_var, x_lag}`.
/// Reset to all-zeros so Pearl A's first-observation sentinel fires
/// on the new fold's first engagement-rate-deficit observation
/// (`prev_x_mean == 0 AND state.x_lag == 0` → first-observation
/// replacement). Companion to `reset_sp4_pearl_c_rate_deficit_ema`
/// — both halves of Pearl A's sentinel contract reset in lockstep.
pub fn reset_sp4_pearl_c_rate_deficit_state(&mut self) {
unsafe {
std::ptr::write_bytes(
self.pearl_c_rate_deficit_state_buf.host_ptr as *mut u8,
0u8,
self.pearl_c_rate_deficit_state_buf.len * std::mem::size_of::<f32>(),
);
}
}
/// SP4 Task A15: zero the host-only Pearl C rate-deficit EMA at
/// fold boundary. Companion to
/// `reset_sp4_pearl_c_rate_deficit_state` — both halves of Pearl
/// A's sentinel contract reset together.
pub fn reset_sp4_pearl_c_rate_deficit_ema(&mut self) {
self.pearl_c_rate_deficit_ema = [0.0_f32; SP4_PARAM_GROUP_COUNT];
}
/// SP4 Task A14: device pointer to the trainer's `nan_flags_buf`
/// (mapped via `CudaSlice::raw_ptr()`). Used by aux trainers (IQN,
/// IQL hi/lo, Attn, Curiosity) to share the Pearl C sticky-flag
/// store via `set_pearl_c_buffers`.
pub fn nan_flags_buf_ptr(&self) -> u64 {
self.nan_flags_buf.raw_ptr()
}
/// SP4 Task A14: device pointer to the trainer's
/// `clamp_engage_per_block_buf` (mapped pinned). Used by aux
/// trainers to share the Pearl C per-block engagement counter
/// buffer via `set_pearl_c_buffers`.
pub fn clamp_engage_per_block_buf_dev_ptr(&self) -> u64 {
self.clamp_engage_per_block_buf.dev_ptr
}
/// Read-only accessor for diagnostics / HEALTH_DIAG mirror.
pub fn grad_norm_fast_ema_value(&self) -> f32 { unsafe { *self.grad_norm_fast_ema_pinned } }
/// Read-only accessor for diagnostics / HEALTH_DIAG mirror.
@@ -22309,6 +22590,12 @@ impl GpuDqnTrainer {
// SP3 Mech 9: ISV-driven post-Adam weight clamp (see `launch_adam_update`).
let q_abs_ref_eff = self.read_isv_signal_at(Q_ABS_REF_INDEX).max(1.0_f32);
let weight_clamp_max_abs: f32 = 100.0_f32 * q_abs_ref_eff;
// SP4 Task A14: aux trainer (sel/recursive_conf) is outside the SP4
// 8-group taxonomy — Pearl C engagement counter disabled.
let _aux_nan_flags_null: u64 = 0;
let _aux_diag_slot_disabled: i32 = -1;
let _aux_engage_buf_null: u64 = 0;
let _aux_engage_off_disabled: i32 = SP4_ENGAGE_OFFSET_DISABLED;
let blocks = ((sel_dim as u32 + 255) / 256).max(1);
unsafe {
self.stream
@@ -22330,6 +22617,11 @@ impl GpuDqnTrainer {
.arg(&l1_end_aux)
.arg(&l1_lambda_aux)
.arg(&weight_clamp_max_abs) // SP3 Mech 9: post-Adam |p_val| clamp
// SP4 Task A14: disabled — sel/recursive_conf outside SP4 group taxonomy.
.arg(&_aux_nan_flags_null)
.arg(&_aux_diag_slot_disabled)
.arg(&_aux_engage_buf_null)
.arg(&_aux_engage_off_disabled)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),

10
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -4281,6 +4281,16 @@ impl GpuExperienceCollector {
self.curiosity_trainer.as_ref()
}
/// SP4 Task A14: thread the main DQN trainer's mapped-pinned Pearl C
/// buffers (`nan_flags_buf`, `clamp_engage_per_block_buf`) into the
/// curiosity trainer. Called once post-construction by
/// `FusedTrainingCtx`. No-op when curiosity is disabled.
pub fn set_curiosity_pearl_c_buffers(&mut self, nan_flags: u64, engage_buf: u64) {
if let Some(ref mut trainer) = self.curiosity_trainer {
trainer.set_pearl_c_buffers(nan_flags, engage_buf);
}
}
/// Disable curiosity training (e.g. after an async kernel crash).
pub fn disable_curiosity(&mut self) {
self.curiosity_trainer = None;

83
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -257,6 +257,26 @@ pub struct GpuIqlTrainer {
grad_norm_blocks: usize,
#[allow(dead_code)]
state_dim_padded: usize,
/// SP4 Task A14: Pearl C engagement-counter buffer pointers, shared
/// from the main DQN trainer's mapped-pinned buffers via
/// `set_pearl_c_buffers`. The IQL Adam writes
/// `clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x]` and
/// `nan_flags_buf[52]` (SP4_IQL_ADAM_DIAG_SLOT). High- and low-tau
/// IQL trainers share the same diag slot but use distinct
/// `engage_buf_offset` values (high=4 × MAX_BLOCKS, low=5 × MAX_BLOCKS)
/// — the offset is supplied as a constructor arg via `param_group`.
/// Both pointers default to 0 → kernel skips the write.
pearl_c_nan_flags_ptr: u64,
pearl_c_engage_buf_ptr: u64,
/// SP4 Task A14: which `ParamGroup` this IQL trainer's params live in.
/// `IqlHigh` (= group 4) for the standard high-tau trainer,
/// `IqlLow` (= group 5) for the low-tau companion. Used to compute
/// `engage_buf_offset = group.idx() * MAX_BLOCKS_PER_ADAM` at launch.
/// Defaults to `IqlHigh` (smoke tests / standalone runs); set via
/// `set_pearl_c_param_group` after construction.
pearl_c_param_group: crate::cuda_pipeline::ParamGroup,
}
impl GpuIqlTrainer {
@@ -508,9 +528,45 @@ impl GpuIqlTrainer {
total_params,
grad_norm_blocks,
state_dim_padded,
// SP4 Task A14: Pearl C disabled until `set_pearl_c_buffers` is
// called by FusedTrainingCtx. Default group `IqlHigh`; the low-
// tau trainer overrides via `set_pearl_c_param_group`.
pearl_c_nan_flags_ptr: 0,
pearl_c_engage_buf_ptr: 0,
pearl_c_param_group: crate::cuda_pipeline::ParamGroup::IqlHigh,
})
}
/// SP4 Task A14: thread the main DQN trainer's mapped-pinned Pearl C
/// buffers down to the IQL Adam launch. Pass 0 for either pointer to
/// disable Pearl C (smoke tests / standalone runs).
pub fn set_pearl_c_buffers(&mut self, nan_flags: u64, engage_buf: u64) {
self.pearl_c_nan_flags_ptr = nan_flags;
self.pearl_c_engage_buf_ptr = engage_buf;
}
/// SP4 Task A15: parameter count accessor for the host-side Pearl C
/// engagement-rate check. Returns the total V(s) MLP parameter count.
pub fn total_params(&self) -> usize {
self.total_params
}
/// SP4 Task A14: declare which `ParamGroup` this IQL trainer's params
/// live in. `IqlHigh` (= 4) for the standard high-tau trainer, `IqlLow`
/// (= 5) for the companion low-tau trainer. Drives `engage_buf_offset
/// = group.idx() × MAX_BLOCKS_PER_ADAM` at launch.
pub fn set_pearl_c_param_group(&mut self, group: crate::cuda_pipeline::ParamGroup) {
debug_assert!(
matches!(
group,
crate::cuda_pipeline::ParamGroup::IqlHigh
| crate::cuda_pipeline::ParamGroup::IqlLow
),
"IQL Pearl C param-group must be IqlHigh or IqlLow"
);
self.pearl_c_param_group = group;
}
/// Run one IQL value network training step — fully deterministic.
///
/// Pipeline: 3 forward GEMMs + SiLU + bias_add + expectile_loss
@@ -983,6 +1039,29 @@ impl GpuIqlTrainer {
let mgn = self.config.max_grad_norm;
let t_ptr = self.t_dev_ptr;
// SP4 Task A14: Pearl C engagement-counter args. If
// `set_pearl_c_buffers` was called by FusedTrainingCtx wiring,
// the kernel records per-block engagement counts at offset
// `pearl_c_param_group * MAX_BLOCKS_PER_ADAM` (IqlHigh=4 →
// offset 1024, IqlLow=5 → offset 1280). When the host pointer
// is unset (pre-wire-up or aux test paths), pass null +
// SP4_ENGAGE_OFFSET_DISABLED so the kernel skips writeback.
let pearl_c_active = self.pearl_c_engage_buf_ptr != 0
&& self.pearl_c_nan_flags_ptr != 0;
let nan_flags_ptr: u64 = self.pearl_c_nan_flags_ptr;
let diag_slot: i32 = if pearl_c_active {
crate::cuda_pipeline::SP4_IQL_ADAM_DIAG_SLOT
} else {
-1
};
let engage_buf_ptr: u64 = self.pearl_c_engage_buf_ptr;
let engage_buf_offset: i32 = if pearl_c_active {
(self.pearl_c_param_group.idx() as i32)
* (crate::cuda_pipeline::MAX_BLOCKS_PER_ADAM as i32)
} else {
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_DISABLED
};
unsafe {
self.stream
.launch_builder(&self.adam_kernel)
@@ -1000,6 +1079,10 @@ impl GpuIqlTrainer {
.arg(&t_ptr)
.arg(&total_params_i32)
.arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp
.arg(&nan_flags_ptr) // SP4 A14: sticky engage flag store
.arg(&diag_slot) // SP4 A14: diag slot (52 = IQL)
.arg(&engage_buf_ptr) // SP4 A14: per-block engage buf
.arg(&engage_buf_offset) // SP4 A14: group offset (4 or 5)
.launch(LaunchConfig {
grid_dim: (adam_blocks as u32, 1, 1),
block_dim: (256, 1, 1),

50
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -331,6 +331,17 @@ pub struct GpuIqnHead {
/// Pre-computed target h_s2 from DQN Pass 2 (avoids redundant trunk forward).
/// Set via `set_cached_target_h_s2()`, consumed (taken) by `execute_training_pipeline`.
cached_target_h_s2_ptr: Option<u64>,
/// SP4 Task A14: Pearl C engagement-counter buffer pointers, shared
/// from the main DQN trainer's mapped-pinned buffers via
/// `set_pearl_c_buffers`. `nan_flags_buf` is the shared `[i32; 55]`
/// sticky-flag store; the IQN Adam writes `[51]=1` on overflow.
/// `clamp_engage_per_block_buf` is the shared `[i32; 2048]` per-block
/// counter buffer indexed by `engage_buf_offset = ParamGroup::Iqn.idx()
/// × MAX_BLOCKS_PER_ADAM`. Both default to 0 → kernel skips the write
/// (compatible with smoke tests / standalone `train_iqn_step_gpu`).
pearl_c_nan_flags_ptr: u64,
pearl_c_engage_buf_ptr: u64,
/// Persistent rewards buffer [B] f32
rewards_buf: CudaSlice<f32>,
/// Persistent dones buffer [B] f32
@@ -674,6 +685,8 @@ impl GpuIqnHead {
branch_actions,
target_h_s2,
cached_target_h_s2_ptr: None,
pearl_c_nan_flags_ptr: 0,
pearl_c_engage_buf_ptr: 0,
rewards_buf,
dones_buf,
save_q_online,
@@ -830,6 +843,17 @@ impl GpuIqnHead {
self.cached_target_h_s2_ptr = Some(ptr);
}
/// SP4 Task A14: thread the main DQN trainer's mapped-pinned Pearl C
/// buffers down to the IQN Adam launch. `nan_flags` is the shared
/// `[i32; SP4_NAN_FLAGS_END]` sticky-flag buffer; `engage_buf` is the
/// shared `[i32; SP4_ENGAGE_BUF_LEN]` per-block engagement-count buffer.
/// Pass 0 for either pointer to disable Pearl C for this IQN trainer
/// (smoke tests / standalone runs).
pub fn set_pearl_c_buffers(&mut self, nan_flags: u64, engage_buf: u64) {
self.pearl_c_nan_flags_ptr = nan_flags;
self.pearl_c_engage_buf_ptr = engage_buf;
}
/// Core IQN training pipeline: trunk forward → cuBLAS forward → loss → backward → Adam.
///
/// Assumes `branch_actions`, `rewards_buf`, `dones_buf` are already populated.
@@ -1475,6 +1499,21 @@ impl GpuIqnHead {
let mgn = self.config.max_grad_norm;
let t_ptr = self.t_dev_ptr;
// SP4 Task A14: Pearl C engagement-counter args. IQN params live in
// ParamGroup::Iqn (= group 3); offset = 3 × MAX_BLOCKS_PER_ADAM. The
// caller (`FusedTrainingCtx`) plumbs through `nan_flags_dev_ptr` +
// `clamp_engage_dev_ptr` from the main DQN trainer's mapped-pinned
// buffers via `set_pearl_c_buffers`.
let nan_flags_ptr = self.pearl_c_nan_flags_ptr;
let diag_slot: i32 = crate::cuda_pipeline::SP4_IQN_ADAM_DIAG_SLOT;
let engage_buf_ptr = self.pearl_c_engage_buf_ptr;
let engage_buf_offset: i32 = if engage_buf_ptr == 0 {
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_DISABLED
} else {
(crate::cuda_pipeline::ParamGroup::Iqn.idx()
* crate::cuda_pipeline::MAX_BLOCKS_PER_ADAM) as i32
};
unsafe {
effective_stream
.launch_builder(&self.adam_kernel)
@@ -1499,6 +1538,10 @@ impl GpuIqnHead {
.arg(&b2_i32)
.arg(&b3_i32)
.arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp
.arg(&nan_flags_ptr) // SP4 A14: Pearl C sticky engage flag
.arg(&diag_slot) // SP4 A14: nan_flags_buf slot
.arg(&engage_buf_ptr) // SP4 A14: per-block engage buffer
.arg(&engage_buf_offset) // SP4 A14: group offset (Iqn=3)
.launch(adam_config)
.map_err(|e| MLError::ModelError(format!("IQN adam kernel: {e}")))?;
}
@@ -1555,6 +1598,13 @@ impl GpuIqnHead {
self.total_params
}
/// SP4 Task A15: parameter count accessor for the host-side Pearl C
/// engagement-rate check. Returns `config.total_params()` cached at
/// construction.
pub fn total_params(&self) -> usize {
self.total_params
}
/// Test-only borrow of the owning CUDA stream — the same stream every
/// async DtoD copy on this head is queued against. The test issues its
/// fills and read-backs on this stream, then synchronises before

18
crates/ml/src/cuda_pipeline/gpu_tlob.rs
View File

@@ -697,6 +697,17 @@ impl GpuTlob {
let eps = 1e-8_f32;
let wd = 1e-5_f32;
let t_ptr = self.t_dev_ptr;
// SP4 Task A14: Pearl C engagement-counter args. TLOB shares
// the `attn_adam_kernel` cubin with `GpuAttention`. For Layer A
// we silently disable Pearl C bookkeeping on this trainer:
// null pointer + SP4_ENGAGE_OFFSET_DISABLED. Layer B can split
// the kernel ownership and add per-trainer engagement slots.
let nan_flags_ptr_null: u64 = 0;
let diag_slot_disabled: i32 = -1;
let engage_buf_ptr_null: u64 = 0;
let engage_off_disabled: i32 =
crate::cuda_pipeline::SP4_ENGAGE_OFFSET_DISABLED;
unsafe {
self.stream
.launch_builder(&self.adam_kernel)
@@ -714,6 +725,13 @@ impl GpuTlob {
.arg(&t_ptr)
.arg(&tp_i32)
.arg(&weight_clamp_max_abs) // SP3 Mech 9 (close-out v2): post-Adam |p| clamp
// SP4 A14: TLOB shares attn_adam_kernel with GpuAttention;
// pass null/disabled to skip per-block engagement writeback
// (Layer A — Layer B can split the kernel if needed).
.arg(&nan_flags_ptr_null)
.arg(&diag_slot_disabled)
.arg(&engage_buf_ptr_null)
.arg(&engage_off_disabled)
.launch(LaunchConfig {
grid_dim: (adam_blocks as u32, 1, 1),
block_dim: (256, 1, 1),

93
crates/ml/src/cuda_pipeline/iql_value_kernel.cu
View File

@@ -203,44 +203,83 @@ void iql_adam_kernel(
int total_params,
/* SP3 Mech 9 (close-out v2): ISV-driven |p| bound (100 × Q_ABS_REF.max(1.0))
* — same pattern wired into all five Adam kernels. 0=disabled (debug only). */
float weight_clamp_max_abs
float weight_clamp_max_abs,
/* SP4 Task A14 (2026-04-30): Pearl C engagement counter. Same warp+
* block tree-reduce pattern as `dqn_adam_update_kernel`. Caller
* passes `engage_buf_offset = group_idx × MAX_BLOCKS_PER_ADAM` (group
* = ParamGroup::IqlHigh or IqlLow depending on which IQL trainer
* launched). `engage_buf_offset < 0` skips the writeback. */
int* __restrict__ nan_flags_buf,
int diag_slot,
int* __restrict__ clamp_engage_per_block_buf,
int engage_buf_offset
)
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= total_params) return;
bool active = (i < total_params);
int local_engage = 0;
/* Gradient clipping from sum-of-squares */
float gn = sqrtf(grad_norm_sq_buf[0]);
float clip_scale = (gn > max_grad_norm && gn > 0.0f) ? (max_grad_norm / gn) : 1.0f;
float g = grads[i] * clip_scale;
if (active) {
/* Gradient clipping from sum-of-squares */
float gn = sqrtf(grad_norm_sq_buf[0]);
float clip_scale = (gn > max_grad_norm && gn > 0.0f) ? (max_grad_norm / gn) : 1.0f;
float g = grads[i] * clip_scale;
/* AdamW: decoupled weight decay */
float p = params[i];
p -= lr * weight_decay * p;
/* AdamW: decoupled weight decay */
float p = params[i];
p -= lr * weight_decay * p;
/* Adam moment updates */
float m_i = beta1 * m[i] + (1.0f - beta1) * g;
float v_i = beta2 * v[i] + (1.0f - beta2) * g * g;
/* Adam moment updates */
float m_i = beta1 * m[i] + (1.0f - beta1) * g;
float v_i = beta2 * v[i] + (1.0f - beta2) * g * g;
/* Bias correction */
int t = t_buf[0];
float m_hat = m_i / (1.0f - powf(beta1, (float)t));
float v_hat = v_i / (1.0f - powf(beta2, (float)t));
/* Bias correction */
int t = t_buf[0];
float m_hat = m_i / (1.0f - powf(beta1, (float)t));
float v_hat = v_i / (1.0f - powf(beta2, (float)t));
/* Parameter update */
p -= lr * m_hat / (sqrtf(v_hat) + eps);
/* Parameter update */
p -= lr * m_hat / (sqrtf(v_hat) + eps);
/* SP3 Mech 9: post-Adam |p| clamp. Last step before writeback —
* applies after all decay/decoupled updates so the bound holds for
* the value the next forward sees. */
if (weight_clamp_max_abs > 0.0f) {
p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
/* SP3 Mech 9 + SP4 Task A14: post-Adam |p| clamp + engagement
* counter. Last step before writeback — applies after all
* decay/decoupled updates so the bound holds for the value the
* next forward sees. */
if (weight_clamp_max_abs > 0.0f) {
if (fabsf(p) > weight_clamp_max_abs) {
if (nan_flags_buf != nullptr && diag_slot >= 0) {
nan_flags_buf[diag_slot] = 1;
}
local_engage = 1;
}
p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
}
params[i] = p;
m[i] = m_i;
v[i] = v_i;
grads[i] = 0.0f; /* zero gradient for next step */
}
params[i] = p;
m[i] = m_i;
v[i] = v_i;
grads[i] = 0.0f; /* zero gradient for next step */
/* SP4 Task A14: warp + block tree-reduce of `local_engage`. */
int e = local_engage;
for (int offset = 16; offset > 0; offset >>= 1)
e += __shfl_xor_sync(0xFFFFFFFF, e, offset);
__shared__ int warp_engage[8]; /* blockDim.x = 256 ⇒ 8 warps */
int warp_id = threadIdx.x / 32;
int warp_lane = threadIdx.x % 32;
if (warp_lane == 0) warp_engage[warp_id] = e;
__syncthreads();
if (warp_id == 0) {
int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
}
}
}
/* ------------------------------------------------------------------ */

107
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -851,6 +851,12 @@ void iqn_grad_norm_phase2(
* `apply_iqn_trunk_gradient` saw NaN at slot 26 with finite inputs because
* `online_params` here drifted via this Adam, never reached by the
* single-kernel Mech 9 in `dqn_adam_update_kernel`. Disabled when bound==0.
*
* SP4 Task A14 (2026-04-30): Pearl C engagement counter via warp+block
* tree-reduce mirroring `dqn_grad_norm_kernel` (no atomicAdd). Single
* block-leader writes `clamp_engage_per_block_buf[engage_buf_offset +
* blockIdx.x]`. `engage_buf_offset < 0` skips the writeback (defensive
* default for callers that disable Pearl C).
*/
extern "C" __global__
void iqn_adam_kernel(
@@ -870,50 +876,89 @@ void iqn_adam_kernel(
int b1_size, /* runtime: unused, for consistent interface */
int b2_size, /* runtime: unused, for consistent interface */
int b3_size, /* runtime: unused, for consistent interface */
float weight_clamp_max_abs /* SP3 Mech 9: ISV-driven |p| bound (100 × Q_ABS_REF.max(1.0)). 0=disabled. */
float weight_clamp_max_abs, /* SP3 Mech 9: ISV-driven |p| bound (100 × Q_ABS_REF.max(1.0)). 0=disabled. */
/* SP4 Task A14: Pearl C engagement counter. */
int* __restrict__ nan_flags_buf, /* sticky overflow flag */
int diag_slot, /* nan_flags_buf index for this Adam kernel */
int* __restrict__ clamp_engage_per_block_buf, /* [SP4_PARAM_GROUP_COUNT × MAX_BLOCKS_PER_ADAM] */
int engage_buf_offset /* group*MAX_BLOCKS; -1 to skip */
)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= total_params) return;
bool active = (tid < total_params);
int local_engage = 0;
float g = grads[tid];
if (active) {
float g = grads[tid];
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged. */
if (!isfinite(g)) { grads[tid] = 0.0f; return; }
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged.
* Still fall through to the block reduce so the engagement counter
* sees this thread's local_engage=0. */
if (!isfinite(g)) {
grads[tid] = 0.0f;
} else {
/* Gradient clipping: norm_sq[0] is raw sum-of-squares from grad_norm kernel */
float grad_norm_val = sqrtf(fmaxf(norm_sq[0], 0.0f));
float clip_scale = (grad_norm_val > max_grad_norm && grad_norm_val > 0.0f)
? max_grad_norm / grad_norm_val : 1.0f;
g *= clip_scale;
/* Gradient clipping: norm_sq[0] is raw sum-of-squares from grad_norm kernel */
float grad_norm_val = sqrtf(fmaxf(norm_sq[0], 0.0f));
float clip_scale = (grad_norm_val > max_grad_norm && grad_norm_val > 0.0f)
? max_grad_norm / grad_norm_val : 1.0f;
g *= clip_scale;
/* AdamW update */
float mi = beta1 * m[tid] + (1.0f - beta1) * g;
float vi = beta2 * v[tid] + (1.0f - beta2) * g * g;
m[tid] = mi;
v[tid] = vi;
/* AdamW update */
float mi = beta1 * m[tid] + (1.0f - beta1) * g;
float vi = beta2 * v[tid] + (1.0f - beta2) * g * g;
m[tid] = mi;
v[tid] = vi;
/* Bias correction */
int adam_t = adam_t_buf[0];
float m_hat = mi / (1.0f - powf(beta1, (float)adam_t));
float v_hat = vi / (1.0f - powf(beta2, (float)adam_t));
/* Bias correction */
int adam_t = adam_t_buf[0];
float m_hat = mi / (1.0f - powf(beta1, (float)adam_t));
float v_hat = vi / (1.0f - powf(beta2, (float)adam_t));
/* Parameter update with decoupled weight decay */
float p = params[tid];
p = p - lr * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * p);
/* Parameter update with decoupled weight decay */
float p = params[tid];
p = p - lr * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * p);
/* SP3 Mech 9 + SP4 Task A14: post-Adam |p| clamp +
* engagement-counter increment. Bound is host-supplied
* (100 × Q_ABS_REF.max(1.0)). One OOM below slot 48
* diagnostic threshold (1e3 × Q_ABS_REF) so the regression
* sentinel keeps 10× firing headroom. */
if (weight_clamp_max_abs > 0.0f) {
if (fabsf(p) > weight_clamp_max_abs) {
if (nan_flags_buf != nullptr && diag_slot >= 0) {
nan_flags_buf[diag_slot] = 1;
}
local_engage = 1;
}
p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
}
/* SP3 Mech 9: post-Adam |p| clamp. Bound is host-supplied
* (100 × Q_ABS_REF.max(1.0)). One OOM below slot 48 diagnostic
* threshold (1e3 × Q_ABS_REF) so the regression sentinel keeps
* 10× firing headroom. */
if (weight_clamp_max_abs > 0.0f) {
p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
params[tid] = p;
/* Zero gradient for next step */
grads[tid] = 0.0f;
}
}
params[tid] = p;
/* SP4 Task A14: warp + block tree-reduce of `local_engage`. */
int e = local_engage;
for (int offset = 16; offset > 0; offset >>= 1)
e += __shfl_xor_sync(0xFFFFFFFF, e, offset);
/* Zero gradient for next step */
grads[tid] = 0.0f;
__shared__ int warp_engage[8]; /* blockDim.x = 256 ⇒ 8 warps */
int warp_id = threadIdx.x / 32;
int warp_lane = threadIdx.x % 32;
if (warp_lane == 0) warp_engage[warp_id] = e;
__syncthreads();
if (warp_id == 0) {
int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
}
}
}
/* ── Forward-Only Kernel (Inference) ────────────────────────────────── */

7
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -65,6 +65,13 @@ pub use sp4_isv_slots::{
GRAD_CLIP_BOUND_INDEX, H_S2_BOUND_INDEX, L1_LAMBDA_TRUNK_INDEX,
SP4_PARAM_GROUP_COUNT, SP4_BRANCH_COUNT, SP4_SLOT_BASE, SP4_SLOT_END, ParamGroup,
weight_bound, adam_m_bound, adam_v_bound, wd_rate, atom_pos_bound,
// SP4 Task A14: Pearl C engagement-counter buffer constants.
MAX_BLOCKS_PER_ADAM, SP4_ENGAGE_BUF_LEN, SP4_ENGAGE_OFFSET_DISABLED,
SP4_DQN_ADAM_DIAG_SLOT, SP4_IQN_ADAM_DIAG_SLOT, SP4_IQL_ADAM_DIAG_SLOT,
SP4_ATTN_ADAM_DIAG_SLOT, SP4_CURIOSITY_ADAM_DIAG_SLOT, SP4_NAN_FLAGS_END,
SP4_ENGAGE_OFFSET_CURIOSITY_W1, SP4_ENGAGE_OFFSET_CURIOSITY_B1,
SP4_ENGAGE_OFFSET_CURIOSITY_W2, SP4_ENGAGE_OFFSET_CURIOSITY_B2,
SP4_ENGAGE_EXTRA_CURIOSITY_SUBLAUNCHES,
};
pub mod sp4_wiener_ema;
pub use sp4_wiener_ema::{

99
crates/ml/src/cuda_pipeline/sp4_isv_slots.rs
View File

@@ -48,6 +48,73 @@ pub const SP4_SLOT_BASE: usize = 131;
/// One past last slot used by SP4. Must equal `ISV_TOTAL_DIM` post-bump.
pub const SP4_SLOT_END: usize = 171;
// ── SP4 Task A14: Pearl C engagement-counter buffer layout ───────────
/// Maximum block grid size per Adam-kernel launch — the engagement
/// counter buffer reserves `MAX_BLOCKS_PER_ADAM` slots per param group,
/// indexed by `blockIdx.x`. 256 covers up to ≈65k params per group at
/// blockDim=256 threads, comfortably above all current production
/// param-group sizes (DQN trunk≈30k, DQN value≈4k, branches≈80k —
/// branch group is the only one exceeding 65k, but the launcher caps
/// `blocks` at `MAX_BLOCKS_PER_ADAM` defensively in launch code).
pub const MAX_BLOCKS_PER_ADAM: usize = 256;
/// Total length of the Pearl C engagement-counter mapped-pinned buffer
/// (`clamp_engage_per_block_buf`). The 8 main param groups occupy
/// offsets [0..2048) (8 × 256). Curiosity has 4 sub-launches (W1/b1/
/// W2/b2) sharing the same kernel; to avoid block-index collisions
/// between sub-launches, each gets its own offset slot. Sub-launch 0
/// reuses Curiosity's main group offset (group_idx=7 → 1792). Sub-
/// launches 1/2/3 occupy extra slots at offsets 2048/2304/2560 — i.e.
/// 3 extra logical "groups" beyond the canonical 8. Total length:
/// (8 + 3) × 256 = 11 × 256 = 2816 i32s. Host-side
/// `pearl_c_post_adam_engagement_check` for `ParamGroup::Curiosity`
/// sums all 4 sub-launch ranges.
pub const SP4_ENGAGE_EXTRA_CURIOSITY_SUBLAUNCHES: usize = 3;
pub const SP4_ENGAGE_BUF_LEN: usize =
(SP4_PARAM_GROUP_COUNT + SP4_ENGAGE_EXTRA_CURIOSITY_SUBLAUNCHES) * MAX_BLOCKS_PER_ADAM;
/// Curiosity sub-launch offsets in `clamp_engage_per_block_buf`.
/// Sub-launch 0 (W1) reuses the Curiosity main-group offset (7×256=1792).
/// Sub-launches 1/2/3 (b1/W2/b2) get extra slots 8/9/10 × 256
/// (= 2048/2304/2560).
pub const SP4_ENGAGE_OFFSET_CURIOSITY_W1: i32 = 7 * (MAX_BLOCKS_PER_ADAM as i32);
pub const SP4_ENGAGE_OFFSET_CURIOSITY_B1: i32 = 8 * (MAX_BLOCKS_PER_ADAM as i32);
pub const SP4_ENGAGE_OFFSET_CURIOSITY_W2: i32 = 9 * (MAX_BLOCKS_PER_ADAM as i32);
pub const SP4_ENGAGE_OFFSET_CURIOSITY_B2: i32 = 10 * (MAX_BLOCKS_PER_ADAM as i32);
/// Sentinel `engage_buf_offset` value passed by aux trainers outside
/// the SP4 8-group taxonomy (DT, OFI embed, denoise, recursive_conf,
/// sel). Causes the kernel's per-block writeback to be skipped — Pearl
/// C only tracks the canonical 8 param groups.
pub const SP4_ENGAGE_OFFSET_DISABLED: i32 = -1;
// ── SP4 Task A14: per-Adam-kernel diag_slot allocations ──────────────
// `nan_flags_buf` slots 50-54 record per-Adam-kernel sticky-engagement
// flags (idempotent writes — no race). The buffer is sized 50 today
// (see gpu_dqn_trainer.rs `nan_flags_buf` alloc); A14 grows it to 55.
/// `dqn_adam_update_kernel` (main DQN trainer + 3 post_aux helpers).
pub const SP4_DQN_ADAM_DIAG_SLOT: i32 = 50;
/// `iqn_adam_kernel` (IQN dual-head trainer).
pub const SP4_IQN_ADAM_DIAG_SLOT: i32 = 51;
/// `iql_adam_kernel` (both high- and low-tau IQL trainers share the
/// slot — identical kernel code; per-group separation is via
/// `engage_buf_offset`, not the sticky flag).
pub const SP4_IQL_ADAM_DIAG_SLOT: i32 = 52;
/// `attn_adam_kernel` (GpuAttention + GpuTlob — TLOB shares the cubin).
pub const SP4_ATTN_ADAM_DIAG_SLOT: i32 = 53;
/// `curiosity_adam_step` (4 sub-buffer launches per step).
pub const SP4_CURIOSITY_ADAM_DIAG_SLOT: i32 = 54;
/// One past the last `nan_flags_buf` index used by SP4 — drives the
/// buffer allocation size in `GpuDqnTrainer::new`.
pub const SP4_NAN_FLAGS_END: usize = 55;
/// Convenience: slot index for `WEIGHT_BOUND[group]`.
#[inline]
pub const fn weight_bound(group: usize) -> usize {
@@ -141,4 +208,36 @@ mod tests {
assert!(atom_pos_bound(SP4_BRANCH_COUNT - 1) < WEIGHT_BOUND_BASE,
"atom_pos_bound max must not alias into WEIGHT_BOUND family");
}
#[test]
fn pearl_c_engage_buf_layout() {
// 8 main groups + 3 extra curiosity sub-launch slots.
assert_eq!(MAX_BLOCKS_PER_ADAM, 256);
assert_eq!(SP4_PARAM_GROUP_COUNT, 8);
assert_eq!(SP4_ENGAGE_EXTRA_CURIOSITY_SUBLAUNCHES, 3);
assert_eq!(SP4_ENGAGE_BUF_LEN, 11 * 256);
assert_eq!(SP4_ENGAGE_BUF_LEN, 2816);
// Curiosity sub-launch offsets are distinct + contiguous starting
// from the canonical Curiosity group offset.
assert_eq!(SP4_ENGAGE_OFFSET_CURIOSITY_W1, 7 * 256);
assert_eq!(SP4_ENGAGE_OFFSET_CURIOSITY_B1, 8 * 256);
assert_eq!(SP4_ENGAGE_OFFSET_CURIOSITY_W2, 9 * 256);
assert_eq!(SP4_ENGAGE_OFFSET_CURIOSITY_B2, 10 * 256);
// Highest sub-launch end must fit within the buffer.
let highest_end = (SP4_ENGAGE_OFFSET_CURIOSITY_B2 as usize) + MAX_BLOCKS_PER_ADAM;
assert_eq!(highest_end, SP4_ENGAGE_BUF_LEN);
// Sentinel disabled value.
assert_eq!(SP4_ENGAGE_OFFSET_DISABLED, -1);
// Per-Adam-kernel diag slot indices land within the
// sticky-flag buffer (size SP4_NAN_FLAGS_END=55).
assert!((SP4_DQN_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
assert!((SP4_IQN_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
assert!((SP4_IQL_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
assert!((SP4_ATTN_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
assert!((SP4_CURIOSITY_ADAM_DIAG_SLOT as usize) < SP4_NAN_FLAGS_END);
}
}

106
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -1687,6 +1687,65 @@ impl FusedTrainingCtx {
// ── Outside graph: kernel-free host ops only ─────────────────
self.pending_vaccine_batch = None;
// SP4 Task A15: Pearl C post-Adam engagement-rate check (Layer A).
// Reads per-block engagement counts from the mapped-pinned buffer
// (zero-copy host access), reduces → engagement_rate, applies
// Pearls A+D to the rate-deficit time series. Layer A logs only;
// Layer B will mutate ISV[WEIGHT_BOUND[group]] on sustained
// over-engagement.
//
// The 5 main Adam launches captured in this step's graph-replay
// each contribute to a distinct param-group offset:
// group 0 (DqnTrunk) — covers DQN Trunk/Value/Branches
// (Layer A simplification: counted under
// group 0 only — Layer B will split).
// group 3 (Iqn) — IQN dual-head trainer.
// group 4 (IqlHigh) — IQL high-tau trainer.
// group 5 (IqlLow) — IQL low-tau trainer.
// group 6 (Attn) — 4-head attention (TLOB shares cubin
// but disables Pearl C).
// group 7 (Curiosity) — handled in `train_curiosity_gpu`,
// not here (separate caller path).
let total_params_dqn = self.trainer.total_params();
if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
crate::cuda_pipeline::ParamGroup::DqnTrunk,
total_params_dqn,
) {
tracing::warn!("SP4 Pearl C check (DQN trunk) failed (non-fatal): {e}");
}
if let Some(ref iqn) = self.gpu_iqn {
let n = iqn.total_params();
if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
crate::cuda_pipeline::ParamGroup::Iqn,
n,
) {
tracing::warn!("SP4 Pearl C check (IQN) failed (non-fatal): {e}");
}
}
let iql_high_n = self.gpu_iql.total_params();
if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
crate::cuda_pipeline::ParamGroup::IqlHigh,
iql_high_n,
) {
tracing::warn!("SP4 Pearl C check (IQL high) failed (non-fatal): {e}");
}
let iql_low_n = self.gpu_iql_low.total_params();
if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
crate::cuda_pipeline::ParamGroup::IqlLow,
iql_low_n,
) {
tracing::warn!("SP4 Pearl C check (IQL low) failed (non-fatal): {e}");
}
if let Some(ref attn) = self.gpu_attention {
let n = attn.total_params();
if let Err(e) = self.trainer.pearl_c_post_adam_engagement_check(
crate::cuda_pipeline::ParamGroup::Attn,
n,
) {
tracing::warn!("SP4 Pearl C check (Attn) failed (non-fatal): {e}");
}
}
// Adaptive IQN lambda: read loss OUTSIDE graph (pinned read, no kernel).
if let Some(ref iqn) = self.gpu_iqn {
if let Ok(loss) = iqn.read_total_loss() {
@@ -3229,6 +3288,53 @@ impl FusedTrainingCtx {
self.trainer.set_curiosity_weights(&weights.w1, &weights.b1, &weights.w2, &weights.b2);
}
/// SP4 Task A14: device pointer to the trainer's `nan_flags_buf`.
/// Used by external aux paths (`GpuExperienceCollector` for the
/// curiosity trainer) to share the Pearl C sticky-flag store.
pub(crate) fn nan_flags_buf_ptr(&self) -> u64 {
self.trainer.nan_flags_buf_ptr()
}
/// SP4 Task A14: device pointer to the trainer's
/// `clamp_engage_per_block_buf` (mapped pinned). Used by external
/// aux paths to share the Pearl C per-block engagement-counter buf.
pub(crate) fn clamp_engage_per_block_buf_dev_ptr(&self) -> u64 {
self.trainer.clamp_engage_per_block_buf_dev_ptr()
}
/// SP4 Task A14: wire the main DQN trainer's mapped-pinned Pearl C
/// buffers into all aux trainers owned by this `FusedTrainingCtx`.
/// Idempotent — safe to call multiple times. Curiosity is wired
/// separately because it lives in `GpuExperienceCollector`, not
/// here; that path goes through
/// `GpuExperienceCollector::set_curiosity_pearl_c_buffers`.
pub(crate) fn wire_aux_trainer_pearl_c_buffers(&mut self) {
let nan_flags_ptr = self.trainer.nan_flags_buf_ptr();
let engage_buf_ptr = self.trainer.clamp_engage_per_block_buf_dev_ptr();
// IQN head (group 3 = ParamGroup::Iqn).
if let Some(ref mut iqn) = self.gpu_iqn {
iqn.set_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
}
// IQL high-tau trainer (group 4 = ParamGroup::IqlHigh) — default
// group on construction; explicitly re-set for clarity.
self.gpu_iql.set_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
self.gpu_iql.set_pearl_c_param_group(
crate::cuda_pipeline::ParamGroup::IqlHigh,
);
// IQL low-tau trainer (group 5 = ParamGroup::IqlLow).
self.gpu_iql_low.set_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
self.gpu_iql_low.set_pearl_c_param_group(
crate::cuda_pipeline::ParamGroup::IqlLow,
);
// Attention head (group 6 = ParamGroup::Attn). TLOB shares the
// attn_adam_kernel cubin but disables Pearl C at the launch site
// (Layer A); we do not wire it here.
if let Some(ref mut attn) = self.gpu_attention {
attn.set_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
}
}
/// Compute Q-value statistics from raw f32 states pointer (cold path).
pub(crate) fn compute_q_stats_from_ptr(
&mut self,

12
crates/ml/src/trainers/dqn/state_reset_registry.rs
View File

@@ -527,7 +527,17 @@ impl StateResetRegistry {
RegistryEntry {
name: "sp4_clamp_engage_counters",
category: ResetCategory::FoldReset,
description: "GpuDqnTrainer.clamp_engage_per_block_buf [2048 ints = SP4_PARAM_GROUP_COUNT × MAX_BLOCKS_PER_ADAM] — SP4 Pearl C engagement counter (per-(group, Adam-launch-block) clamp-trigger count). Mapped-pinned i32; bulk write_bytes(0) at fold boundary so the new fold's engagement-rate tracking restarts from a clean 0 baseline (Task A12)",
description: "GpuDqnTrainer.clamp_engage_per_block_buf [2816 ints = (SP4_PARAM_GROUP_COUNT + 3 curiosity sub-launches) × MAX_BLOCKS_PER_ADAM = 11 × 256] — SP4 Pearl C engagement counter (per-(group, Adam-launch-block) clamp-trigger count). Mapped-pinned i32; bulk write_bytes(0) at fold boundary so the new fold's engagement-rate tracking restarts from a clean 0 baseline. Buffer length grew 2048→2816 in Task A14/A15 to give Curiosity's 4 sub-launches (W1/b1/W2/b2) distinct offset slots avoiding per-block index collisions (Task A12 + A14/A15)",
},
RegistryEntry {
name: "sp4_pearl_c_rate_deficit_state",
category: ResetCategory::FoldReset,
description: "GpuDqnTrainer.pearl_c_rate_deficit_state_buf [24 floats = SP4_PARAM_GROUP_COUNT × {sample_var, diff_var, x_lag}] — SP4 Pearl C rate-deficit Wiener-state (per param group). Mapped-pinned f32; bulk write_bytes(0) at fold boundary so the new fold's engagement-rate-deficit EMA restarts from Pearl A sentinel (`prev_x_mean == 0 AND state.x_lag == 0` → first-observation replacement). Paired with sp4_pearl_c_rate_deficit_ema — both reset together (Task A15)",
},
RegistryEntry {
name: "sp4_pearl_c_rate_deficit_ema",
category: ResetCategory::FoldReset,
description: "GpuDqnTrainer.pearl_c_rate_deficit_ema [8 floats, host-only [f32; SP4_PARAM_GROUP_COUNT]] — SP4 Pearl C rate-deficit EMA per param group (the `prev_x_mean` surrogate read by `pearls_ad_update`). Reset to all-zeros at fold boundary so the new fold's deficit tracking restarts from the Pearl A sentinel. Companion to sp4_pearl_c_rate_deficit_state (Task A15)",
},
];
Self { entries }

46
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1430,6 +1430,14 @@ impl DQNTrainer {
info!("Curiosity Q-penalty: weights wired to fused trainer (lambda={:.2})",
self.hyperparams.curiosity_q_penalty_lambda);
}
// SP4 Task A14: wire the main DQN trainer's mapped-
// pinned Pearl C buffers into all aux trainers
// (IQN, IQL hi/lo, Attn) and the curiosity trainer.
fused_ctx.wire_aux_trainer_pearl_c_buffers();
let nan_flags_ptr = fused_ctx.nan_flags_buf_ptr();
let engage_buf_ptr = fused_ctx.clamp_engage_per_block_buf_dev_ptr();
collector.set_curiosity_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
info!("SP4 Pearl C: aux trainer engagement-counter buffers wired");
}
self.gpu_experience_collector = Some(collector);
}
@@ -1774,6 +1782,20 @@ impl DQNTrainer {
) {
debug!("GPU curiosity training failed (non-fatal): {e}");
}
// SP4 Task A15: Pearl C post-Adam engagement-rate check for
// curiosity trainer (Layer A). Curiosity has 4 sub-launches
// (W1/b1/W2/b2) and `pearl_c_post_adam_engagement_check`
// sums all 4 sub-launch ranges before computing the rate.
if let Some(ref mut fused) = self.fused_ctx {
let curiosity_total = crate::cuda_pipeline::gpu_curiosity_trainer
::GpuCuriosityTrainer::total_params();
if let Err(e) = fused.trainer_mut().pearl_c_post_adam_engagement_check(
crate::cuda_pipeline::ParamGroup::Curiosity,
curiosity_total,
) {
tracing::warn!("SP4 Pearl C check (Curiosity) failed (non-fatal): {e}");
}
}
if let Some(ref stream) = self.cuda_stream {
match stream.record_event(None) {
Ok(event) => {
@@ -5841,12 +5863,34 @@ impl DQNTrainer {
}
"sp4_clamp_engage_counters" => {
// Bulk write_bytes(0) on `clamp_engage_per_block_buf`
// (2048 mapped-pinned i32). Each fold restarts Pearl C's
// (2816 mapped-pinned i32 — Task A14/A15 grew it from
// 2048 to accommodate Curiosity's 4 distinct sub-launch
// offset slots). Each fold restarts Pearl C's
// engagement-rate tracking from a clean 0 baseline.
if let Some(ref mut fused) = self.fused_ctx {
fused.trainer_mut().reset_sp4_clamp_engage_counters();
}
}
"sp4_pearl_c_rate_deficit_state" => {
// Bulk write_bytes(0) on `pearl_c_rate_deficit_state_buf`
// (24 mapped-pinned f32 = 8 groups × {sample_var, diff_var,
// x_lag}). Pearl A's first-observation sentinel requires
// `state.x_lag == 0` on the new fold's first engagement-
// rate-deficit observation; reset paired with the host-
// only `pearl_c_rate_deficit_ema` below.
if let Some(ref mut fused) = self.fused_ctx {
fused.trainer_mut().reset_sp4_pearl_c_rate_deficit_state();
}
}
"sp4_pearl_c_rate_deficit_ema" => {
// Reset the host-only `pearl_c_rate_deficit_ema [f32; 8]`
// EMA surrogate to all zeros. Companion to
// sp4_pearl_c_rate_deficit_state — both halves of Pearl A's
// sentinel contract reset together.
if let Some(ref mut fused) = self.fused_ctx {
fused.trainer_mut().reset_sp4_pearl_c_rate_deficit_ema();
}
}
"isv_gamma_dir_eff" | "isv_gamma_mag_eff" | "isv_gamma_ord_eff" | "isv_gamma_urg_eff" => {
// D.2 per-branch gamma slots. Reset to 0.0 at fold boundary;
// per_branch_gamma_update GPU kernel re-populates on next epoch-boundary launch.

2
docs/dqn-wire-up-audit.md
View File

File diff suppressed because one or more lines are too long