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

feat(dqn-v2): Plan 4 Task 6 Commit A — aux heads scaffolding (kernels + params [119..127) + 2 ISV slots, additive)

Browse Source 
Multi-task auxiliary heads (E.6) — additive scaffolding ONLY. NO
production callers in this commit; Commit B wires forward/backward +
training-loop loss accumulation.

Two `Linear(SH2 → 32) → ELU → Linear(32 → K)` MLPs branch off h_s2:
- Next-bar return regression head (K=1, MSE loss).
- 5-class regime classification head (K=5, cross-entropy loss).

Lands in this commit (additive, no behavioural change):

1. Two CUDA kernel files (kernel count 62 → 64):
   - `aux_heads_kernel.cu` — 8 entry points (per-head forward, backward,
     loss-reduce; shared regime label-builder + per-tensor batch-reduce).
     Single-block-per-sample, shmem-tree reductions, no atomicAdd.
   - `aux_heads_loss_ema_kernel.cu` — single-thread single-block ISV
     producer mirroring h_s2_rms_ema_kernel.cu's footprint.
2. Rust orchestrator `cuda_pipeline/gpu_aux_heads.rs` —
   `AuxHeadsForwardOps` + `AuxHeadsBackwardOps` mirror gpu_grn.rs's
   raw-u64-pointer ABI for graph-capture safety. NO callers.
3. Param tensors at [119..127) (NUM_WEIGHT_TENSORS 119 → 127):
   aux_nb_{w1,b1,w2,b2}, aux_rg_{w1,b1,w2,b2}. Xavier on weights,
   zero on biases. Adam SAXPY iterates 0..total_params (count-driven),
   so the new range gets covered automatically; with no producer for
   grad_buf[119..127) in this commit, Adam keeps params at Xavier
   init (intended dormant state).
4. Two new ISV slots tail-appended after VSN:
   AUX_NEXT_BAR_MSE_EMA_INDEX=113, AUX_REGIME_CE_EMA_INDEX=114.
   Producer: aux_heads_loss_ema_update. Cold-start 0.0; FoldReset → 0.0.
   Diagnostic only; no GPU consumer kernel reads slots [113..115).
5. Fingerprint pair shifted [111,112] → [115,116]; ISV_TOTAL_DIM
   113 → 117. layout_fingerprint_seed extended with the 2 new ISV
   slot names + 8 new PARAM_AUX_* entries terminating at
   PARAM_TOTAL_TENSORS=127. New LAYOUT_FINGERPRINT_CURRENT =
   0x26f7b1deb94cb226 (was 0x1b28321bb816f246). Checkpoint break by
   intent — fail-fast at constructor load on mismatch.
6. Two new FoldReset entries in state_reset_registry.rs
   (isv_aux_next_bar_mse_ema, isv_aux_regime_ce_ema) WITH matching
   dispatch arm in training_loop.rs::reset_named_state (the wire
   VSN-rc2 missed and chain-final fix made dispatch-mandatory).
7. Polyak EMA target sync NOT extended for aux heads — design
   decision: aux heads are online-only supervised heads, not used in
   Bellman bootstrapping. target_params_buf[119..127) initialised by
   xavier_init_params_buf (same Xavier values as online) and never
   updated; verified no consumer reads target_params_buf past
   padded_byte_offset(119).

Constraints honoured: GPU-only (every reduction GPU-side, zero DtoH);
no atomicAdd (shmem-tree reductions throughout backward + final
batch-dim collapse); no stubs (every Rust function is real
implementation; "no callers yet" applies only to orchestrator wrappers
waiting on Commit B); partial-refactor invariant honoured
(NUM_WEIGHT_TENSORS, Adam range, xavier init, fingerprint seed,
ISV_TOTAL_DIM, FoldReset entries + dispatch arm all migrate together);
no `// ok:` band-aids; no tuned constants.

cargo check --workspace clean at 11 warnings (workspace baseline
preserved). Smoke deferred to Commit B (no behaviour change in this
commit — aux head params are dormant Xavier weights).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-26 08:43:38 +02:00
parent 31e0f219a5
commit 5478e7c821
9 changed files with 1398 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
crates/ml/build.rs
View File

@@ -153,6 +153,19 @@ fn main() {
// EMA-updates 6 ISV slots. Producer-only; consumer-side ISV slot
// allocation lands in 1B-ii.
"vsn_mask_ema_kernel.cu",
// Plan 4 Task 6 Commit A: multi-task auxiliary heads (E.6).
// Two `Linear(SH2 → 32) → ELU → Linear(32 → K)` MLPs branching off
// h_s2 — next-bar return regression (K=1, MSE) + 5-class regime
// classification (K=5, cross-entropy). Forward + backward + loss-
// reduce + label-builder + per-tensor batch-reduce kernels. Module
// is additive — ZERO production callers in this commit; Commit B
// wires the forward/backward + training-loop loss accumulation.
"aux_heads_kernel.cu",
// Plan 4 Task 6 Commit A: aux-head loss EMA producer (single-thread
// single-block kernel mirroring h_s2_rms_ema_kernel's footprint).
// Producer-only; ISV[AUX_NEXT_BAR_MSE_EMA_INDEX] +
// ISV[AUX_REGIME_CE_EMA_INDEX] consumer wires in Commit B.
"aux_heads_loss_ema_kernel.cu",
];
// ALL kernels get common header (BF16 types + wrappers)

627
crates/ml/src/cuda_pipeline/aux_heads_kernel.cu Normal file
View File

@@ -0,0 +1,627 @@
/*
* aux_heads_kernel.cu — Plan 4 Task 6 Commit A.
*
* Multi-task auxiliary heads (E.6): two small MLPs branching off the
* trunk's `h_s2 [B, SH2]` post-GRN activation:
*
* (1) Next-bar return regression head:
* pred = Linear_2(ELU(Linear_1(h_s2))) where K_out = 1
* Loss = MSE against `next_close_pct` derived from `next_states`.
*
* (2) Five-class regime classification head:
* logits = Linear_2(ELU(Linear_1(h_s2))) where K_out = 5
* Loss = cross-entropy against a discretised regime label
* (`floor(states[regime_score_idx] * 5)`, clamped to [0, 4]).
*
* Both heads share the `Linear(SH2 → 32) → ELU → Linear(32 → K)` shape
* with `H = AUX_HIDDEN_DIM = 32`.
*
* This commit lands the kernels + Rust orchestrator + params/ISV/registry
* scaffolding only — there are NO production callers in this commit.
* Commit B wires the forward/backward + training-loop loss accumulation.
*
* Pearls applied:
* - feedback_no_atomicadd.md — per-block shmem-tree reductions, no
* atomicAdd anywhere (param-grad accumulation is two-phase: per-block
* partial → final reduce, mirroring `grn_layernorm_backward_dgamma_dbeta`).
* - pearl_cold_path_no_exception_to_gpu_drives.md — every reduction
* stays GPU-side; no DtoH / host roundtrips.
* - feedback_no_feature_flags.md — no enable_/use_ booleans; the only
* axis of behavioural variation is K_out (1 vs 5), which is a
* structural property of the head, not a toggle.
*
* Layout convention: row-major `[B, …]` for all activation / gradient
* tensors. `w1 [H, SH2]` row-major; `w2 [K, H]` row-major. Matches the
* VSN per-group MLP convention from `vsn_feature_selection_kernel.cu`.
*
* ELU activation: f(x) = x if x > 0 else (exp(x) - 1).
* Backward: f'(x) = 1 if x > 0 else exp(x) = 1 + f(x).
*
* Cross-entropy: numerically-stable form — subtract max(logits) before exp.
* Loss[b] = -log( exp(logits[b, label[b]] - max) / sum_k exp(logits[b, k] - max) )
* d_logits[b, k] = (softmax[b, k] - one_hot(label[b], k))
*
* Mean-over-batch loss reduction: dW / db both divide by B in the
* accumulation reduce so the kernel-internal scale matches `1/B * sum_b ∂L/∂W`.
*/
#include <cuda_runtime.h>
#include <math_constants.h>
/* AUX_HIDDEN_DIM kept as a compile-time literal so kernel launch configs
* don't need to thread `H` through. Both heads share the same hidden width. */
#ifndef AUX_HIDDEN_DIM
#define AUX_HIDDEN_DIM 32
#endif
/* Block dim used by reduction-shaped launches (per-tensor backward, label
* builders). Kept identical to `vsn_feature_selection_kernel.cu`'s 256
* threads/block convention. */
#ifndef AUX_BLOCK
#define AUX_BLOCK 256
#endif
/* ELU helpers — branchless on x sign via `x > 0.f`. */
__device__ __forceinline__ float aux_elu_fwd(float x) {
return (x > 0.0f) ? x : (expf(x) - 1.0f);
}
/* Backward derivative from POST-activation y = ELU(x). For x > 0,
* y == x and y > 0; for x <= 0, y = exp(x) - 1 ∈ (-1, 0] so 1 + y = exp(x).
* Distinguishing on `y > 0` recovers the same partition. */
__device__ __forceinline__ float aux_elu_bwd_from_post(float y) {
return (y > 0.0f) ? 1.0f : (1.0f + y);
}
/* =====================================================================
* aux_next_bar_forward — Linear → ELU → Linear with K=1 output.
*
* Per sample b (one block per sample, AUX_BLOCK threads):
* 1. h[k] = ELU( b1[k] + sum_j w1[k, j] * h_s2[b, j] ) for k ∈ [0, H)
* 2. pred[b] = b2 + sum_k w2[0, k] * h[k]
*
* Hidden activations `h [B, H]` are SAVED for backward (consumed as the
* ELU post-activation buffer; backward derives `f'(x_pre)` from `h_post`
* via the standard ELU backward identity).
*
* Block: AUX_BLOCK threads. Grid: (B, 1, 1). Shared memory: H floats
* (h cache for the second linear's reduction).
* ===================================================================== */
extern "C" __global__ void aux_next_bar_forward(
const float* __restrict__ h_s2, /* [B, SH2] row-major */
const float* __restrict__ w1, /* [H, SH2] row-major */
const float* __restrict__ b1, /* [H] */
const float* __restrict__ w2, /* [1, H] (== [H] flat) */
const float* __restrict__ b2, /* [1] */
int B,
int SH2,
float* __restrict__ hidden_out, /* [B, H] OUT (saved for backward) */
float* __restrict__ pred_out /* [B, 1] OUT */
) {
const int b = blockIdx.x;
if (b >= B) return;
const int tid = threadIdx.x;
const int H = AUX_HIDDEN_DIM;
extern __shared__ float smem[]; /* h cache [H] */
/* Step 1: per-hidden-unit Linear_1 + ELU. Each thread computes one or
* more hidden lanes via stride loop. SH2 is small enough that each
* thread doing a serial dot over SH2 is cheaper than block-cooperative
* reduce per lane. */
for (int k = tid; k < H; k += blockDim.x) {
float acc = b1[k];
const float* w_row = w1 + (size_t)k * SH2;
const float* h_row = h_s2 + (size_t)b * SH2;
for (int j = 0; j < SH2; ++j) {
acc += w_row[j] * h_row[j];
}
const float h_post = aux_elu_fwd(acc);
smem[k] = h_post;
hidden_out[(size_t)b * H + k] = h_post;
}
__syncthreads();
/* Step 2: scalar Linear_2. One thread does the H-element dot. H is
* small (32) — a tree reduce wastes more cycles in shfl bookkeeping
* than the serial pass. */
if (tid == 0) {
float acc = b2[0];
for (int k = 0; k < H; ++k) {
acc += w2[k] * smem[k];
}
pred_out[b] = acc;
}
}
/* =====================================================================
* aux_regime_forward — Linear → ELU → Linear with K=5 output.
*
* Same architecture as next_bar but with K_out = 5.
*
* Per sample b: same Step 1 as next_bar (cached in shmem); Step 2 computes
* 5 logits via thread 0 (5 × H = 160 multiplies — well under the cost of
* a 5-way fanout block reduce).
*
* Block: AUX_BLOCK threads. Grid: (B, 1, 1). Shared memory: H floats.
* ===================================================================== */
extern "C" __global__ void aux_regime_forward(
const float* __restrict__ h_s2, /* [B, SH2] row-major */
const float* __restrict__ w1, /* [H, SH2] row-major */
const float* __restrict__ b1, /* [H] */
const float* __restrict__ w2, /* [K, H] row-major (K = 5) */
const float* __restrict__ b2, /* [K] */
int B,
int SH2,
int K,
float* __restrict__ hidden_out, /* [B, H] OUT (saved for backward) */
float* __restrict__ logits_out /* [B, K] OUT */
) {
const int b = blockIdx.x;
if (b >= B) return;
const int tid = threadIdx.x;
const int H = AUX_HIDDEN_DIM;
extern __shared__ float smem[]; /* h cache [H] */
for (int k = tid; k < H; k += blockDim.x) {
float acc = b1[k];
const float* w_row = w1 + (size_t)k * SH2;
const float* h_row = h_s2 + (size_t)b * SH2;
for (int j = 0; j < SH2; ++j) {
acc += w_row[j] * h_row[j];
}
const float h_post = aux_elu_fwd(acc);
smem[k] = h_post;
hidden_out[(size_t)b * H + k] = h_post;
}
__syncthreads();
/* Step 2: K-way Linear_2. Thread 0 emits all K logits (K is at most 5
* in this commit's design). */
if (tid == 0) {
for (int kc = 0; kc < K; ++kc) {
float acc = b2[kc];
const float* w2_row = w2 + (size_t)kc * H;
for (int k = 0; k < H; ++k) {
acc += w2_row[k] * smem[k];
}
logits_out[(size_t)b * K + kc] = acc;
}
}
}
/* =====================================================================
* aux_next_bar_loss_reduce — single-block shmem-reduce mean MSE.
*
* Computes loss = (1/B) * sum_b (pred[b] - label[b])^2 into loss_out[0].
* Single-block (AUX_BLOCK threads) shmem-tree reduction; no atomicAdd.
*
* Block: AUX_BLOCK threads. Grid: (1, 1, 1). Shared mem: AUX_BLOCK floats.
* ===================================================================== */
extern "C" __global__ void aux_next_bar_loss_reduce(
const float* __restrict__ pred, /* [B] (== [B, 1] flat) */
const float* __restrict__ label, /* [B] */
int B,
float* __restrict__ loss_out /* [1] */
) {
if (blockIdx.x != 0) return;
extern __shared__ float smem[];
const int tid = threadIdx.x;
const int block = blockDim.x;
float local = 0.0f;
for (int i = tid; i < B; i += block) {
const float d = pred[i] - label[i];
local += d * d;
}
smem[tid] = local;
__syncthreads();
for (int s = block / 2; s > 0; s >>= 1) {
if (tid < s) smem[tid] += smem[tid + s];
__syncthreads();
}
/* B > 0 invariant — caller guarantees positive batch. NaN propagation
* preferred over masking per the same rationale documented in
* h_s2_rms_ema_kernel.cu and vsn_mask_ema_kernel.cu. */
if (tid == 0) {
loss_out[0] = smem[0] / (float)B;
}
}
/* =====================================================================
* aux_regime_loss_reduce — single-block shmem-reduce mean cross-entropy.
*
* For each sample b, computes per-sample CE = -log(softmax(logits)[label])
* using the numerically-stable max-shift form. Reduces both the loss and
* the count of samples whose argmax(logits) == label, writing scalar
* means into loss_out[0] and correct_out[0].
*
* Block: AUX_BLOCK threads. Grid: (1, 1, 1).
* Shared memory: 2 * AUX_BLOCK floats (loss + correct partial reductions).
* ===================================================================== */
extern "C" __global__ void aux_regime_loss_reduce(
const float* __restrict__ logits, /* [B, K] row-major */
const unsigned char* __restrict__ label, /* [B] uint8 in [0, K) */
int B,
int K,
float* __restrict__ loss_out, /* [1] */
float* __restrict__ correct_out /* [1] mean accuracy in [0, 1] */
) {
if (blockIdx.x != 0) return;
extern __shared__ float smem[];
float* sh_loss = smem;
float* sh_correct = smem + blockDim.x;
const int tid = threadIdx.x;
const int block = blockDim.x;
float local_loss = 0.0f;
float local_correct = 0.0f;
for (int i = tid; i < B; i += block) {
const float* row = logits + (size_t)i * K;
/* Numerically-stable softmax: subtract max before exp. */
float lmax = row[0];
int amax = 0;
for (int k = 1; k < K; ++k) {
if (row[k] > lmax) { lmax = row[k]; amax = k; }
}
float sum_e = 0.0f;
for (int k = 0; k < K; ++k) {
sum_e += expf(row[k] - lmax);
}
const int tgt = (int)label[i];
const float ce = (lmax - row[tgt]) + logf(sum_e);
local_loss += ce;
local_correct += (amax == tgt) ? 1.0f : 0.0f;
}
sh_loss[tid] = local_loss;
sh_correct[tid] = local_correct;
__syncthreads();
for (int s = block / 2; s > 0; s >>= 1) {
if (tid < s) {
sh_loss[tid] += sh_loss[tid + s];
sh_correct[tid] += sh_correct[tid + s];
}
__syncthreads();
}
if (tid == 0) {
loss_out[0] = sh_loss[0] / (float)B;
correct_out[0] = sh_correct[0] / (float)B;
}
}
/* =====================================================================
* aux_regime_label_from_states — discretise states[:, regime_score_idx]
* into uint8 labels in [0, K).
*
* Reads `states[b, regime_score_idx]` (a float in [0, 1] — the Plan 4
* spec's regime stability score / regime cluster score), maps to an
* integer bin via floor(s * K) clamped to [0, K-1].
*
* Block: AUX_BLOCK threads. Grid: ceil(B / AUX_BLOCK) blocks. One thread
* per sample. No atomicAdd (each thread writes a unique label slot).
* ===================================================================== */
extern "C" __global__ void aux_regime_label_from_states(
const float* __restrict__ states, /* [B, F] row-major */
int B,
int F,
int regime_score_idx,
int K,
unsigned char* __restrict__ label_out /* [B] uint8 */
) {
const int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B) return;
const float s = states[(size_t)b * F + regime_score_idx];
/* floor(s * K) with K in {1..255}; clamp to [0, K-1]. The clamp covers
* both the s == 1.0 boundary (floor(K) == K out-of-range) and any
* out-of-spec negative values. */
int bin = (int)floorf(s * (float)K);
if (bin < 0) bin = 0;
if (bin > K - 1) bin = K - 1;
label_out[b] = (unsigned char)bin;
}
/* =====================================================================
* aux_next_bar_backward — full backward for the next-bar regression head.
*
* Forward was:
* h_pre[b, k] = b1[k] + sum_j w1[k, j] * h_s2[b, j]
* h_post[b, k] = ELU(h_pre[b, k])
* pred[b] = b2[0] + sum_k w2[0, k] * h_post[b, k]
*
* Loss = (1/B) * sum_b (pred[b] - label[b])^2.
*
* Mean-over-batch backward:
* d_pred[b] = (2 / B) * (pred[b] - label[b])
* d_b2[0] = sum_b d_pred[b]
* d_w2[0, k] = sum_b d_pred[b] * h_post[b, k]
* d_h_post[b, k] = d_pred[b] * w2[0, k]
* d_h_pre[b, k] = d_h_post[b, k] * elu_bwd(h_post[b, k])
* d_b1[k] = sum_b d_h_pre[b, k]
* d_w1[k, j] = sum_b d_h_pre[b, k] * h_s2[b, j]
* d_h_s2[b, j] = sum_k d_h_pre[b, k] * w1[k, j]
*
* Implementation: one block per sample (B blocks), AUX_BLOCK threads.
* Each block computes its sample's `d_h_pre` into shared memory, then:
* - writes its sample's contribution to `d_h_s2` row b (no contention),
* - writes per-block partials `dW1_partial[b, k, j]` and
* `dW2_partial[b, 0, k]` and per-block bias partials.
*
* To stay atomicAdd-free, we DO NOT do block-cooperative param-grad
* accumulation. Instead the caller runs a follow-up REDUCE pass over the
* `[B, …]` partials. This kernel emits the per-sample partials as
* contiguous tensors (size O(B * H * SH2 + B * H + B * H + B * 1)) and a
* separate `aux_param_grad_reduce_*` kernel finalises them.
*
* For Commit A (no callers), the partial-emission contract is the
* artefact that Commit B's training-loop wire-up consumes; the reduce
* step lands together with the wire-up since the partials → final-reduce
* pattern needs both endpoints to be exercised at once.
*
* The dh_s2 output is `[B, SH2]`: the trunk's h_s2 backward path will SUM
* this into the existing main-path d_h_s2 (no atomicAdd, just a SAXPY in
* the caller).
*
* Block: AUX_BLOCK threads. Grid: (B, 1, 1).
* Shared memory: 2 * H floats (d_h_pre cache + ELU-derivative cache).
* ===================================================================== */
extern "C" __global__ void aux_next_bar_backward(
/* Forward inputs (re-read for dW). */
const float* __restrict__ h_s2, /* [B, SH2] */
const float* __restrict__ w1, /* [H, SH2] */
const float* __restrict__ w2, /* [H] flat */
const float* __restrict__ hidden_post, /* [B, H] saved h_post */
/* Loss inputs. */
const float* __restrict__ pred, /* [B] */
const float* __restrict__ label, /* [B] */
int B,
int SH2,
/* Partial outputs (per-sample). Caller reduces along batch dim. */
float* __restrict__ dW1_partial, /* [B, H, SH2] flat */
float* __restrict__ db1_partial, /* [B, H] */
float* __restrict__ dW2_partial, /* [B, H] */
float* __restrict__ db2_partial, /* [B] */
/* Trunk gradient (contiguous in B, no batch-reduce needed — caller
* accumulates into trunk's main-path d_h_s2 via SAXPY). */
float* __restrict__ dh_s2_out /* [B, SH2] */
) {
const int b = blockIdx.x;
if (b >= B) return;
const int tid = threadIdx.x;
const int H = AUX_HIDDEN_DIM;
extern __shared__ float smem[];
float* sh_dh_pre = smem; /* [H] */
float* sh_h_post = smem + H; /* [H] */
/* Mean-over-batch derivative scale: dL/dpred[b] = (2/B) * (pred[b] - label[b]). */
const float d_pred_b = (2.0f / (float)B) * (pred[b] - label[b]);
/* Step 1: cache h_post and compute d_h_pre per hidden unit. */
for (int k = tid; k < H; k += blockDim.x) {
const float h_post = hidden_post[(size_t)b * H + k];
sh_h_post[k] = h_post;
const float d_h_post = d_pred_b * w2[k];
sh_dh_pre[k] = d_h_post * aux_elu_bwd_from_post(h_post);
}
__syncthreads();
/* Step 2: per-sample partials for {db2, dW2, db1, dW1}.
* db2_partial[b] = d_pred_b
* dW2_partial[b, k] = d_pred_b * h_post[b, k]
* db1_partial[b, k] = sh_dh_pre[k]
* dW1_partial[b, k, j] = sh_dh_pre[k] * h_s2[b, j]
*/
if (tid == 0) {
db2_partial[b] = d_pred_b;
}
for (int k = tid; k < H; k += blockDim.x) {
const float dh = sh_dh_pre[k];
dW2_partial[(size_t)b * H + k] = d_pred_b * sh_h_post[k];
db1_partial[(size_t)b * H + k] = dh;
}
/* dW1 partial — per-(k, j) write. Stride loop over (H * SH2). */
{
const float* h_row = h_s2 + (size_t)b * SH2;
float* dW1_b = dW1_partial + (size_t)b * H * SH2;
const int total = H * SH2;
for (int idx = tid; idx < total; idx += blockDim.x) {
const int k = idx / SH2;
const int j = idx - k * SH2;
dW1_b[idx] = sh_dh_pre[k] * h_row[j];
}
}
/* Step 3: dh_s2[b, j] = sum_k sh_dh_pre[k] * w1[k, j].
* Per-sample, per-feature serial reduction over H lanes — H is small
* (32) and SH2 is moderate; stride loop over j. */
{
float* dh_row = dh_s2_out + (size_t)b * SH2;
for (int j = tid; j < SH2; j += blockDim.x) {
float acc = 0.0f;
for (int k = 0; k < H; ++k) {
acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j];
}
dh_row[j] = acc;
}
}
}
/* =====================================================================
* aux_regime_backward — full backward for the 5-class regime CE head.
*
* Forward was:
* logits[b, kc] = b2[kc] + sum_k w2[kc, k] * h_post[b, k]
* p[b, kc] = softmax(logits[b, :])[kc]
* CE[b] = -log(p[b, label[b]])
*
* Mean-over-batch loss `L = (1/B) * sum_b CE[b]` ⇒
* d_logits[b, kc] = (1/B) * (p[b, kc] - one_hot(label[b], kc))
*
* The remaining backward (dW2, db2, dW1, db1, dh_s2) mirrors the structure
* of `aux_next_bar_backward` but with a K-vector `d_logits` instead of a
* scalar `d_pred`.
*
* Block: AUX_BLOCK threads. Grid: (B, 1, 1).
* Shared memory: 2 * H + K floats (d_h_pre cache, h_post cache, d_logits cache).
* ===================================================================== */
extern "C" __global__ void aux_regime_backward(
/* Forward inputs. */
const float* __restrict__ h_s2, /* [B, SH2] */
const float* __restrict__ w1, /* [H, SH2] */
const float* __restrict__ w2, /* [K, H] row-major */
const float* __restrict__ hidden_post, /* [B, H] saved h_post */
const float* __restrict__ logits, /* [B, K] saved (re-softmaxed here) */
/* Loss inputs. */
const unsigned char* __restrict__ label, /* [B] uint8 in [0, K) */
int B,
int SH2,
int K,
/* Partial outputs (per-sample). */
float* __restrict__ dW1_partial, /* [B, H, SH2] */
float* __restrict__ db1_partial, /* [B, H] */
float* __restrict__ dW2_partial, /* [B, K, H] */
float* __restrict__ db2_partial, /* [B, K] */
/* Trunk gradient. */
float* __restrict__ dh_s2_out /* [B, SH2] */
) {
const int b = blockIdx.x;
if (b >= B) return;
const int tid = threadIdx.x;
const int H = AUX_HIDDEN_DIM;
extern __shared__ float smem[];
float* sh_dh_pre = smem; /* [H] */
float* sh_h_post = smem + H; /* [H] */
float* sh_dlogits = smem + 2 * H; /* [K] */
/* Step 0: compute d_logits[b, :] from the saved logits + label. Single
* thread (K is small — at most 5). Mean-over-batch scaling: divide by B. */
if (tid == 0) {
const float* row = logits + (size_t)b * K;
float lmax = row[0];
for (int k = 1; k < K; ++k) if (row[k] > lmax) lmax = row[k];
float sum_e = 0.0f;
for (int k = 0; k < K; ++k) sum_e += expf(row[k] - lmax);
const float inv_sum = 1.0f / sum_e;
const int tgt = (int)label[b];
const float inv_B = 1.0f / (float)B;
for (int k = 0; k < K; ++k) {
const float p = expf(row[k] - lmax) * inv_sum;
const float onehot = (k == tgt) ? 1.0f : 0.0f;
sh_dlogits[k] = inv_B * (p - onehot);
}
}
__syncthreads();
/* Step 1: cache h_post and compute d_h_pre.
* d_h_post[b, k] = sum_kc d_logits[b, kc] * w2[kc, k]
* d_h_pre[b, k] = d_h_post[b, k] * elu_bwd(h_post[b, k])
*/
for (int k = tid; k < H; k += blockDim.x) {
const float h_post = hidden_post[(size_t)b * H + k];
sh_h_post[k] = h_post;
float d_h_post = 0.0f;
for (int kc = 0; kc < K; ++kc) {
d_h_post += sh_dlogits[kc] * w2[(size_t)kc * H + k];
}
sh_dh_pre[k] = d_h_post * aux_elu_bwd_from_post(h_post);
}
__syncthreads();
/* Step 2: per-sample param-grad partials. */
if (tid == 0) {
for (int kc = 0; kc < K; ++kc) {
db2_partial[(size_t)b * K + kc] = sh_dlogits[kc];
}
}
/* dW2_partial[b, kc, k] = sh_dlogits[kc] * sh_h_post[k]. */
{
const int total = K * H;
for (int idx = tid; idx < total; idx += blockDim.x) {
const int kc = idx / H;
const int k = idx - kc * H;
dW2_partial[(size_t)b * K * H + idx] = sh_dlogits[kc] * sh_h_post[k];
}
}
for (int k = tid; k < H; k += blockDim.x) {
db1_partial[(size_t)b * H + k] = sh_dh_pre[k];
}
/* dW1_partial[b, k, j] = sh_dh_pre[k] * h_s2[b, j]. */
{
const float* h_row = h_s2 + (size_t)b * SH2;
float* dW1_b = dW1_partial + (size_t)b * H * SH2;
const int total = H * SH2;
for (int idx = tid; idx < total; idx += blockDim.x) {
const int k = idx / SH2;
const int j = idx - k * SH2;
dW1_b[idx] = sh_dh_pre[k] * h_row[j];
}
}
/* Step 3: dh_s2[b, j] = sum_k sh_dh_pre[k] * w1[k, j]. */
{
float* dh_row = dh_s2_out + (size_t)b * SH2;
for (int j = tid; j < SH2; j += blockDim.x) {
float acc = 0.0f;
for (int k = 0; k < H; ++k) {
acc += sh_dh_pre[k] * w1[(size_t)k * SH2 + j];
}
dh_row[j] = acc;
}
}
}
/* =====================================================================
* aux_param_grad_reduce — final-reduce kernel for the per-sample partials
* emitted by both backward kernels.
*
* Aggregates `partials [B, P]` along the batch dimension into
* `final [P]`. Single-block 256-thread (or grid-stride per-element)
* shmem-tree reduction; no atomicAdd. Each block handles one column of
* the partials (one element of the final output).
*
* Grid: (P, 1, 1) blocks. Block: AUX_BLOCK threads. Shared mem: AUX_BLOCK floats.
* ===================================================================== */
extern "C" __global__ void aux_param_grad_reduce(
const float* __restrict__ partials, /* [B, P] row-major */
int B,
int P,
float* __restrict__ final_out /* [P] */
) {
const int p = blockIdx.x;
if (p >= P) return;
const int tid = threadIdx.x;
const int block = blockDim.x;
extern __shared__ float smem[];
float local = 0.0f;
for (int i = tid; i < B; i += block) {
local += partials[(size_t)i * P + p];
}
smem[tid] = local;
__syncthreads();
for (int s = block / 2; s > 0; s >>= 1) {
if (tid < s) smem[tid] += smem[tid + s];
__syncthreads();
}
if (tid == 0) {
final_out[p] = smem[0];
}
}

50
crates/ml/src/cuda_pipeline/aux_heads_loss_ema_kernel.cu Normal file
View File

@@ -0,0 +1,50 @@
/* aux_heads_loss_ema — GPU-driven EMA of the two aux-head loss scalars
* into ISV[AUX_NEXT_BAR_MSE_EMA_INDEX] and ISV[AUX_REGIME_CE_EMA_INDEX].
*
* Plan 4 Task 6 Commit A. Producer-only kernel — no consumer reads either
* slot in this commit. Commit B wires consumers (HEALTH_DIAG emission +
* aux-loss attribution monitor).
*
* Mirrors `h_s2_rms_ema_kernel.cu`'s shape — single thread, single block,
* cold-path-cadence (one launch per training step alongside the other
* Plan 4 ISV producers in `training_loop.rs`).
*
* The two scalar input buffers are produced by `aux_next_bar_loss_reduce`
* and `aux_regime_loss_reduce` (in `aux_heads_kernel.cu`); they are
* always already a scalar mean — this kernel just performs the EMA blend
* with the prior ISV value.
*
* GPU-drives-CPU-reads (spec §4.C.6) and pearl_cold_path_no_exception_to
* _gpu_drives.md compliant. No atomicAdd (per feedback_no_atomicadd.md);
* the EMA itself is two reads + two writes by a single thread.
*
* Cost: 2 global reads + 2 global writes per launch. Comfortably under
* any kernel-launch overhead — fusing with `h_s2_rms_ema_update` would
* not measurably help.
*/
extern "C" __global__ void aux_heads_loss_ema_update(
const float* __restrict__ next_bar_loss_scalar, /* [1] from aux_next_bar_loss_reduce */
const float* __restrict__ regime_loss_scalar, /* [1] from aux_regime_loss_reduce */
float* __restrict__ isv, /* ISV bus [ISV_TOTAL_DIM] */
int isv_aux_next_bar_mse_index, /* AUX_NEXT_BAR_MSE_EMA_INDEX = 113 */
int isv_aux_regime_ce_index, /* AUX_REGIME_CE_EMA_INDEX = 114 */
float ema_alpha /* per-batch EMA rate (0,1) */
) {
/* Single-block, single-thread contract — match h_s2_rms_ema's grid guard.
* The work is two scalar EMAs; spawning a 256-thread block would waste
* every thread except thread 0. */
if (blockIdx.x != 0) return;
if (threadIdx.x != 0) return;
const float nb_loss = next_bar_loss_scalar[0];
const float rg_loss = regime_loss_scalar[0];
const int nb_slot = isv_aux_next_bar_mse_index;
const int rg_slot = isv_aux_regime_ce_index;
const float prev_nb = isv[nb_slot];
const float prev_rg = isv[rg_slot];
isv[nb_slot] = (1.0f - ema_alpha) * prev_nb + ema_alpha * nb_loss;
isv[rg_slot] = (1.0f - ema_alpha) * prev_rg + ema_alpha * rg_loss;
}

535
crates/ml/src/cuda_pipeline/gpu_aux_heads.rs Normal file
View File

@@ -0,0 +1,535 @@
#![allow(unsafe_code)] // Required for CUDA kernel launches.
//! `GpuAuxHeads` — Plan 4 Task 6 Commit A.
//!
//! Multi-task auxiliary heads orchestrator (E.6). Owns the per-head
//! `CudaFunction` handles for both forward / backward / loss-reduce kernels
//! plus the regime-label builder, and exposes thin `pub(crate)` wrappers
//! that mirror the kernel ABIs in row-major device-pointer form.
//!
//! NO PRODUCTION CALLERS in this commit — Commit B's training-loop wire-up
//! invokes these methods. The orchestrator is built as additive scaffolding
//! per the same convention used for `gpu_grn.rs` (1B-i / 2c.2 pattern):
//! the kernels and orchestrator land first, the wire-up commit lights them
//! up.
//!
//! # Architecture
//!
//! Both heads share the `Linear(SH2 → 32) → ELU → Linear(32 → K)` shape:
//!
//! - Next-bar regression head: K = 1, MSE loss.
//! - Regime classification head: K = 5, cross-entropy loss.
//!
//! Param tensors live in the trainer's flat `params_buf` at indices
//! `[119..127)` (8 tensors total — w1/b1/w2/b2 per head). See
//! `compute_param_sizes()` in `gpu_dqn_trainer.rs` for the exact layout.
//!
//! # Pearls applied
//!
//! - `feedback_no_atomicadd.md` — backward kernels emit per-sample partials;
//! the trailing `aux_param_grad_reduce` kernel collapses along the batch
//! dim with shmem-tree reduction (no atomicAdd anywhere).
//! - `pearl_cold_path_no_exception_to_gpu_drives.md` — every reduction is
//! GPU-side; the orchestrator does zero CPU/DtoH work.
//! - `feedback_wire_everything_up.md` — this orchestrator is intentionally
//! built as a pre-wire scaffold for Commit B, mirroring the same
//! commit-by-commit ordering used for VSN (1B-ii orchestrator → 1B-iii
//! wire-up). The Commit B wire-up retires the orphan-by-design status.
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
use crate::MLError;
// ── Cubin embedding (kernels live in aux_heads_kernel.cu / aux_heads_loss_ema_kernel.cu) ──
use super::gpu_dqn_trainer::{AUX_HEADS_CUBIN, AUX_HEADS_LOSS_EMA_CUBIN};
/// Hidden dim shared by both auxiliary heads (Linear_1 output width).
/// Kept in sync with the `AUX_HIDDEN_DIM` macro at the top of
/// `aux_heads_kernel.cu`.
pub(crate) const AUX_HIDDEN_DIM: usize = 32;
/// Output cardinality of the next-bar regression head (scalar prediction).
pub(crate) const AUX_NEXT_BAR_K: usize = 1;
/// Output cardinality of the regime classification head (5-way softmax).
pub(crate) const AUX_REGIME_K: usize = 5;
/// Block dim used by the per-sample forward / backward kernels.
const AUX_BLOCK: u32 = 256;
/// Forward orchestrator — owns the 5 forward / loss-reduce / label-builder
/// kernel handles for both auxiliary heads.
///
/// Dropping this struct releases the underlying `CudaFunction` handles via
/// cudarc's RAII; nothing else needs explicit teardown.
#[allow(missing_debug_implementations)]
pub(crate) struct AuxHeadsForwardOps {
next_bar_forward_kernel: CudaFunction,
regime_forward_kernel: CudaFunction,
next_bar_loss_reduce_kernel: CudaFunction,
regime_loss_reduce_kernel: CudaFunction,
regime_label_kernel: CudaFunction,
/// Aux loss EMA producer kernel (single thread, single block; mirrors
/// `h_s2_rms_ema_update`'s shape). Loaded from `AUX_HEADS_LOSS_EMA_CUBIN`.
loss_ema_kernel: CudaFunction,
}
impl AuxHeadsForwardOps {
/// Load all forward / reduce / EMA kernel handles from the precompiled
/// cubins. Mirrors `GrnBlock::new`'s loader pattern.
pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let context = stream.context();
let aux_module = context
.load_cubin(AUX_HEADS_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("aux_heads cubin load: {e}")))?;
let next_bar_forward_kernel = aux_module
.load_function("aux_next_bar_forward")
.map_err(|e| MLError::ModelError(format!("aux_next_bar_forward load: {e}")))?;
let regime_forward_kernel = aux_module
.load_function("aux_regime_forward")
.map_err(|e| MLError::ModelError(format!("aux_regime_forward load: {e}")))?;
let next_bar_loss_reduce_kernel = aux_module
.load_function("aux_next_bar_loss_reduce")
.map_err(|e| MLError::ModelError(format!("aux_next_bar_loss_reduce load: {e}")))?;
let regime_loss_reduce_kernel = aux_module
.load_function("aux_regime_loss_reduce")
.map_err(|e| MLError::ModelError(format!("aux_regime_loss_reduce load: {e}")))?;
let regime_label_kernel = aux_module
.load_function("aux_regime_label_from_states")
.map_err(|e| MLError::ModelError(format!("aux_regime_label_from_states load: {e}")))?;
let ema_module = context
.load_cubin(AUX_HEADS_LOSS_EMA_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema cubin load: {e}")))?;
let loss_ema_kernel = ema_module
.load_function("aux_heads_loss_ema_update")
.map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema_update load: {e}")))?;
Ok(Self {
next_bar_forward_kernel,
regime_forward_kernel,
next_bar_loss_reduce_kernel,
regime_loss_reduce_kernel,
regime_label_kernel,
loss_ema_kernel,
})
}
/// Launch `aux_next_bar_forward`: `Linear(h_s2) → ELU → Linear → pred`.
///
/// Caller-owned buffers (raw `u64` device pointers for graph-capture
/// safety; mirrors `gpu_grn.rs::forward_raw`):
/// * `h_s2_ptr` — `[B, SH2]` row-major.
/// * `w1_ptr`/`b1_ptr` — `[H, SH2]` / `[H]`. Slice from `params_buf[119]`/`[120]`.
/// * `w2_ptr`/`b2_ptr` — `[1, H]` / `[1]`. Slice from `params_buf[121]`/`[122]`.
/// * `hidden_out_ptr` — `[B, H]` SAVED post-ELU buffer (consumed by backward).
/// * `pred_out_ptr` — `[B, 1]` scalar prediction per sample.
///
/// Block: `AUX_BLOCK` threads. Grid: `B` blocks. Shared mem: `H` floats.
#[allow(clippy::too_many_arguments)]
pub(crate) fn forward_next_bar(
&self,
stream: &Arc<CudaStream>,
h_s2_ptr: u64,
w1_ptr: u64,
b1_ptr: u64,
w2_ptr: u64,
b2_ptr: u64,
b: usize,
sh2: usize,
hidden_out_ptr: u64,
pred_out_ptr: u64,
) -> Result<(), MLError> {
let b_i32 = b as i32;
let sh2_i32 = sh2 as i32;
let smem_bytes = (AUX_HIDDEN_DIM as u32) * std::mem::size_of::<f32>() as u32;
unsafe {
stream
.launch_builder(&self.next_bar_forward_kernel)
.arg(&h_s2_ptr)
.arg(&w1_ptr)
.arg(&b1_ptr)
.arg(&w2_ptr)
.arg(&b2_ptr)
.arg(&b_i32)
.arg(&sh2_i32)
.arg(&hidden_out_ptr)
.arg(&pred_out_ptr)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (AUX_BLOCK, 1, 1),
shared_mem_bytes: smem_bytes,
})
.map_err(|e| MLError::ModelError(format!("aux_next_bar_forward: {e}")))?;
}
Ok(())
}
/// Launch `aux_regime_forward`: `Linear(h_s2) → ELU → Linear → logits[K]`.
///
/// `K` is `AUX_REGIME_K = 5` in this commit's design; passed as runtime
/// arg for kernel-signature stability.
#[allow(clippy::too_many_arguments)]
pub(crate) fn forward_regime(
&self,
stream: &Arc<CudaStream>,
h_s2_ptr: u64,
w1_ptr: u64,
b1_ptr: u64,
w2_ptr: u64,
b2_ptr: u64,
b: usize,
sh2: usize,
k: usize,
hidden_out_ptr: u64,
logits_out_ptr: u64,
) -> Result<(), MLError> {
let b_i32 = b as i32;
let sh2_i32 = sh2 as i32;
let k_i32 = k as i32;
let smem_bytes = (AUX_HIDDEN_DIM as u32) * std::mem::size_of::<f32>() as u32;
unsafe {
stream
.launch_builder(&self.regime_forward_kernel)
.arg(&h_s2_ptr)
.arg(&w1_ptr)
.arg(&b1_ptr)
.arg(&w2_ptr)
.arg(&b2_ptr)
.arg(&b_i32)
.arg(&sh2_i32)
.arg(&k_i32)
.arg(&hidden_out_ptr)
.arg(&logits_out_ptr)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (AUX_BLOCK, 1, 1),
shared_mem_bytes: smem_bytes,
})
.map_err(|e| MLError::ModelError(format!("aux_regime_forward: {e}")))?;
}
Ok(())
}
/// Launch `aux_next_bar_loss_reduce`: scalar mean MSE into `loss_out[1]`.
pub(crate) fn next_bar_loss_reduce(
&self,
stream: &Arc<CudaStream>,
pred_ptr: u64,
label_ptr: u64,
b: usize,
loss_out_ptr: u64,
) -> Result<(), MLError> {
let b_i32 = b as i32;
let smem_bytes = AUX_BLOCK * std::mem::size_of::<f32>() as u32;
unsafe {
stream
.launch_builder(&self.next_bar_loss_reduce_kernel)
.arg(&pred_ptr)
.arg(&label_ptr)
.arg(&b_i32)
.arg(&loss_out_ptr)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (AUX_BLOCK, 1, 1),
shared_mem_bytes: smem_bytes,
})
.map_err(|e| MLError::ModelError(format!("aux_next_bar_loss_reduce: {e}")))?;
}
Ok(())
}
/// Launch `aux_regime_loss_reduce`: scalar mean CE + scalar mean accuracy.
pub(crate) fn regime_loss_reduce(
&self,
stream: &Arc<CudaStream>,
logits_ptr: u64,
label_ptr: u64,
b: usize,
k: usize,
loss_out_ptr: u64,
correct_out_ptr: u64,
) -> Result<(), MLError> {
let b_i32 = b as i32;
let k_i32 = k as i32;
// Two parallel partial-reduction strips: loss + accuracy. See
// `aux_regime_loss_reduce` shmem layout in the kernel.
let smem_bytes = 2 * AUX_BLOCK * std::mem::size_of::<f32>() as u32;
unsafe {
stream
.launch_builder(&self.regime_loss_reduce_kernel)
.arg(&logits_ptr)
.arg(&label_ptr)
.arg(&b_i32)
.arg(&k_i32)
.arg(&loss_out_ptr)
.arg(&correct_out_ptr)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (AUX_BLOCK, 1, 1),
shared_mem_bytes: smem_bytes,
})
.map_err(|e| MLError::ModelError(format!("aux_regime_loss_reduce: {e}")))?;
}
Ok(())
}
/// Launch `aux_regime_label_from_states`: discretise the regime score
/// channel of the state vector into `[B] uint8` regime labels in `[0, K)`.
///
/// `regime_score_idx` is a feature offset into the row-major `[B, F]`
/// state buffer (typically `ml_core::state_layout::REGIME_STABILITY_INDEX`,
/// resolved by the caller).
#[allow(clippy::too_many_arguments)]
pub(crate) fn regime_label_from_states(
&self,
stream: &Arc<CudaStream>,
states_ptr: u64,
b: usize,
f: usize,
regime_score_idx: usize,
k: usize,
label_out_ptr: u64,
) -> Result<(), MLError> {
let b_i32 = b as i32;
let f_i32 = f as i32;
let r_i32 = regime_score_idx as i32;
let k_i32 = k as i32;
let blocks = ((b as u32) + AUX_BLOCK - 1) / AUX_BLOCK;
unsafe {
stream
.launch_builder(&self.regime_label_kernel)
.arg(&states_ptr)
.arg(&b_i32)
.arg(&f_i32)
.arg(&r_i32)
.arg(&k_i32)
.arg(&label_out_ptr)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (AUX_BLOCK, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("aux_regime_label_from_states: {e}")))?;
}
Ok(())
}
/// Launch `aux_heads_loss_ema_update`: cold-path-cadence EMA of both
/// scalar loss buffers into `ISV[isv_aux_next_bar_mse_index]` and
/// `ISV[isv_aux_regime_ce_index]`. Single-thread, single-block — mirrors
/// `h_s2_rms_ema_update`'s footprint.
#[allow(clippy::too_many_arguments)]
pub(crate) fn launch_loss_ema(
&self,
stream: &Arc<CudaStream>,
next_bar_loss_scalar_ptr: u64,
regime_loss_scalar_ptr: u64,
isv_dev_ptr: u64,
isv_aux_next_bar_mse_index: i32,
isv_aux_regime_ce_index: i32,
ema_alpha: f32,
) -> Result<(), MLError> {
unsafe {
stream
.launch_builder(&self.loss_ema_kernel)
.arg(&next_bar_loss_scalar_ptr)
.arg(&regime_loss_scalar_ptr)
.arg(&isv_dev_ptr)
.arg(&isv_aux_next_bar_mse_index)
.arg(&isv_aux_regime_ce_index)
.arg(&ema_alpha)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("aux_heads_loss_ema_update: {e}")))?;
}
Ok(())
}
}
/// Backward orchestrator — owns the 3 backward / final-reduce kernel
/// handles. Held separate from the forward ops because Commit B's wire-up
/// will dispatch forward + backward at distinct points in the training
/// step (forward runs alongside the main encoder forward; backward runs
/// after the main loss backward and SAXPY-accumulates `dh_s2` into the
/// trunk's main-path gradient).
#[allow(missing_debug_implementations)]
pub(crate) struct AuxHeadsBackwardOps {
next_bar_backward_kernel: CudaFunction,
regime_backward_kernel: CudaFunction,
/// Final-reduce kernel for the per-sample param-grad partials emitted
/// by both backward kernels. `aux_param_grad_reduce(partials [B, P], P,
/// final [P])`.
param_grad_reduce_kernel: CudaFunction,
}
impl AuxHeadsBackwardOps {
/// Load the 3 backward kernel handles from the precompiled aux-heads
/// cubin.
pub(crate) fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let context = stream.context();
let aux_module = context
.load_cubin(AUX_HEADS_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("aux_heads cubin load (bwd): {e}")))?;
let next_bar_backward_kernel = aux_module
.load_function("aux_next_bar_backward")
.map_err(|e| MLError::ModelError(format!("aux_next_bar_backward load: {e}")))?;
let regime_backward_kernel = aux_module
.load_function("aux_regime_backward")
.map_err(|e| MLError::ModelError(format!("aux_regime_backward load: {e}")))?;
let param_grad_reduce_kernel = aux_module
.load_function("aux_param_grad_reduce")
.map_err(|e| MLError::ModelError(format!("aux_param_grad_reduce load: {e}")))?;
Ok(Self {
next_bar_backward_kernel,
regime_backward_kernel,
param_grad_reduce_kernel,
})
}
/// Launch `aux_next_bar_backward` — emits per-sample param-grad
/// partials + per-sample `dh_s2 [B, SH2]`. Caller follows up with
/// `param_grad_reduce` for each of {`dW1`, `db1`, `dW2`, `db2`} to
/// collapse along the batch dim.
#[allow(clippy::too_many_arguments)]
pub(crate) fn backward_next_bar(
&self,
stream: &Arc<CudaStream>,
h_s2_ptr: u64,
w1_ptr: u64,
w2_ptr: u64,
hidden_post_ptr: u64,
pred_ptr: u64,
label_ptr: u64,
b: usize,
sh2: usize,
dw1_partial_ptr: u64,
db1_partial_ptr: u64,
dw2_partial_ptr: u64,
db2_partial_ptr: u64,
dh_s2_out_ptr: u64,
) -> Result<(), MLError> {
let b_i32 = b as i32;
let sh2_i32 = sh2 as i32;
let smem_bytes = (2 * AUX_HIDDEN_DIM as u32) * std::mem::size_of::<f32>() as u32;
unsafe {
stream
.launch_builder(&self.next_bar_backward_kernel)
.arg(&h_s2_ptr)
.arg(&w1_ptr)
.arg(&w2_ptr)
.arg(&hidden_post_ptr)
.arg(&pred_ptr)
.arg(&label_ptr)
.arg(&b_i32)
.arg(&sh2_i32)
.arg(&dw1_partial_ptr)
.arg(&db1_partial_ptr)
.arg(&dw2_partial_ptr)
.arg(&db2_partial_ptr)
.arg(&dh_s2_out_ptr)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (AUX_BLOCK, 1, 1),
shared_mem_bytes: smem_bytes,
})
.map_err(|e| MLError::ModelError(format!("aux_next_bar_backward: {e}")))?;
}
Ok(())
}
/// Launch `aux_regime_backward` — emits per-sample param-grad partials
/// + per-sample `dh_s2 [B, SH2]`. Caller follows up with
/// `param_grad_reduce` for each of {`dW1`, `db1`, `dW2`, `db2`}.
#[allow(clippy::too_many_arguments)]
pub(crate) fn backward_regime(
&self,
stream: &Arc<CudaStream>,
h_s2_ptr: u64,
w1_ptr: u64,
w2_ptr: u64,
hidden_post_ptr: u64,
logits_ptr: u64,
label_ptr: u64,
b: usize,
sh2: usize,
k: usize,
dw1_partial_ptr: u64,
db1_partial_ptr: u64,
dw2_partial_ptr: u64,
db2_partial_ptr: u64,
dh_s2_out_ptr: u64,
) -> Result<(), MLError> {
let b_i32 = b as i32;
let sh2_i32 = sh2 as i32;
let k_i32 = k as i32;
let smem_bytes =
((2 * AUX_HIDDEN_DIM as u32) + (k as u32)) * std::mem::size_of::<f32>() as u32;
unsafe {
stream
.launch_builder(&self.regime_backward_kernel)
.arg(&h_s2_ptr)
.arg(&w1_ptr)
.arg(&w2_ptr)
.arg(&hidden_post_ptr)
.arg(&logits_ptr)
.arg(&label_ptr)
.arg(&b_i32)
.arg(&sh2_i32)
.arg(&k_i32)
.arg(&dw1_partial_ptr)
.arg(&db1_partial_ptr)
.arg(&dw2_partial_ptr)
.arg(&db2_partial_ptr)
.arg(&dh_s2_out_ptr)
.launch(LaunchConfig {
grid_dim: (b as u32, 1, 1),
block_dim: (AUX_BLOCK, 1, 1),
shared_mem_bytes: smem_bytes,
})
.map_err(|e| MLError::ModelError(format!("aux_regime_backward: {e}")))?;
}
Ok(())
}
/// Launch `aux_param_grad_reduce`: collapse `partials [B, P]` along
/// the batch dim into `final_out [P]`. One block per output element
/// (`P` blocks); each block does a shmem-tree reduction over `B`
/// samples. No atomicAdd.
pub(crate) fn param_grad_reduce(
&self,
stream: &Arc<CudaStream>,
partials_ptr: u64,
b: usize,
p: usize,
final_out_ptr: u64,
) -> Result<(), MLError> {
let b_i32 = b as i32;
let p_i32 = p as i32;
let smem_bytes = AUX_BLOCK * std::mem::size_of::<f32>() as u32;
unsafe {
stream
.launch_builder(&self.param_grad_reduce_kernel)
.arg(&partials_ptr)
.arg(&b_i32)
.arg(&p_i32)
.arg(&final_out_ptr)
.launch(LaunchConfig {
grid_dim: (p as u32, 1, 1),
block_dim: (AUX_BLOCK, 1, 1),
shared_mem_bytes: smem_bytes,
})
.map_err(|e| MLError::ModelError(format!("aux_param_grad_reduce: {e}")))?;
}
Ok(())
}
}

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

@@ -155,6 +155,27 @@ pub(crate) static VSN_FEATURE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(en
/// 105..111).
pub(crate) static VSN_MASK_EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/vsn_mask_ema_kernel.cubin"));
/// Plan 4 Task 6 Commit A: aux-heads kernels (forward + backward + loss-
/// reduce + label-builder + per-tensor batch-reduce). 7 entry points:
/// `aux_next_bar_forward`, `aux_regime_forward`, `aux_next_bar_loss_reduce`,
/// `aux_regime_loss_reduce`, `aux_regime_label_from_states`,
/// `aux_next_bar_backward`, `aux_regime_backward`, `aux_param_grad_reduce`.
/// Loaded by `gpu_aux_heads::{AuxHeadsForwardOps, AuxHeadsBackwardOps}`.
/// `pub(crate)` so the orchestrator module in
/// `cuda_pipeline/gpu_aux_heads.rs` can `include_bytes!()` the same cubin
/// without a duplicate copy. Module is **additive** in this commit — ZERO
/// production callers; Commit B wires the forward/backward and training-
/// loop loss accumulation.
#[allow(dead_code)]
pub(crate) static AUX_HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_heads_kernel.cubin"));
/// Plan 4 Task 6 Commit A: aux-heads loss EMA producer (single-thread,
/// single-block kernel mirroring `h_s2_rms_ema_kernel.cu`). Producer for
/// ISV[AUX_NEXT_BAR_MSE_EMA_INDEX] + ISV[AUX_REGIME_CE_EMA_INDEX]; consumer
/// wires in Commit B (HEALTH_DIAG aux-loss diagnostic line).
#[allow(dead_code)]
pub(crate) static AUX_HEADS_LOSS_EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_heads_loss_ema_kernel.cubin"));
/// Plan 4 Task 1B-ii: per-group VSN MLP hidden dim. `Linear_1[g]` is
/// `[VSN_HIDDEN_DIM, group_dim_g]`, `Linear_2[g]` is `[1, VSN_HIDDEN_DIM]`.
/// 16 chosen as a tractable expansion that fits comfortably against the
@@ -381,13 +402,25 @@ const ISV_NETWORK_DIM: usize = 23;
/// [108] VSN_MASK_GROUP_3_EMA_INDEX — `mtf` group, same producer/contract.
/// [109] VSN_MASK_GROUP_4_EMA_INDEX — `portfolio` group, same producer/contract.
/// [110] VSN_MASK_GROUP_5_EMA_INDEX — `plan_isv` group, same producer/contract.
/// [111] ISV_LAYOUT_FINGERPRINT_LO_INDEX — low 32 bits of layout fingerprint
/// (shifted 103→111 in 1B-ii).
/// [112] ISV_LAYOUT_FINGERPRINT_HI_INDEX — high 32 bits of layout fingerprint
/// (shifted 104→112 in 1B-ii).
/// [113] AUX_NEXT_BAR_MSE_EMA_INDEX — Plan 4 Task 6 Commit A — EMA of the
/// next-bar regression head's mean-batch MSE (α=0.05). Producer:
/// `aux_heads_loss_ema_update` GPU kernel (single-thread, single-block;
/// mirrors `h_s2_rms_ema_update`'s shape). Launch site wires in
/// Commit B alongside `launch_h_s2_rms_ema`. Cold-start 0.0 (no aux
/// forwards have happened yet). FoldReset → 0.0. Diagnostic only —
/// no consumer kernel reads slot [113] in either Commit A or B
/// (Commit B's HEALTH_DIAG emission is a CPU-side text consumer
/// only, not a kernel input).
/// [114] AUX_REGIME_CE_EMA_INDEX — Plan 4 Task 6 Commit A — EMA of the
/// regime classification head's mean-batch cross-entropy
/// (α=0.05). Same producer/contract as [113]; same FoldReset → 0.0.
/// [115] ISV_LAYOUT_FINGERPRINT_LO_INDEX — low 32 bits of layout fingerprint
/// (shifted 111→115 in Plan 4 Task 6 Commit A).
/// [116] ISV_LAYOUT_FINGERPRINT_HI_INDEX — high 32 bits of layout fingerprint
/// (shifted 112→116 in Plan 4 Task 6 Commit A).
/// Written by the constructor; checked at checkpoint load. Fail-fast only — no migration
/// path exists. See spec §4.A.2 and `LAYOUT_FINGERPRINT_CURRENT` for structural-hash rationale.
const ISV_TOTAL_DIM: usize = 113;
const ISV_TOTAL_DIM: usize = 117;
/// Legacy alias preserved for call sites that haven't been audited for the
/// network-vs-total split. New code should pick `ISV_NETWORK_DIM` (for weight
/// tensor sizing) or `ISV_TOTAL_DIM` (for the broadcast bus buffer).
@@ -704,7 +737,32 @@ pub const VSN_MASK_GROUP_3_EMA_INDEX: usize = 108;
pub const VSN_MASK_GROUP_4_EMA_INDEX: usize = 109;
pub const VSN_MASK_GROUP_5_EMA_INDEX: usize = 110;
/// ISV slot [111] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
/// ISV slots [113..115) — Plan 4 Task 6 Commit A: aux-head loss EMAs.
///
/// Two slots: `[113] AUX_NEXT_BAR_MSE_EMA_INDEX` carries the EMA of the
/// next-bar regression head's batch-mean MSE; `[114] AUX_REGIME_CE_EMA_INDEX`
/// carries the EMA of the regime classification head's batch-mean
/// cross-entropy. Producer: `aux_heads_loss_ema_update` GPU kernel
/// (single-thread, single-block; mirrors `h_s2_rms_ema_update`'s footprint
/// — the work is two scalar EMAs, no reduction needed because the upstream
/// loss-reduce kernels emit scalars). EMA α matches the rest of the
/// Plan 4 producer family (cold-path-cadence, one launch per training
/// step alongside `h_s2_rms_ema_update`).
///
/// Cold-start: 0.0 (no aux forwards have happened yet — Commit B wires the
/// forwards). FoldReset → 0.0.
///
/// Producer-only after Commit B — the only diagnostic consumer is
/// HEALTH_DIAG's `aux[next_bar_mse=… regime_ce=…]` line (CPU-side text
/// emission, not a kernel input). No GPU kernel reads slots [113..115)
/// in either commit. Deliberately tail-appended after Plan 4 Task 1B-ii's
/// VSN slots so the ISV layout grows monotonically per
/// `feedback_no_partial_refactor.md` (every consumer of `ISV_TOTAL_DIM`
/// or the fingerprint shifts together in this commit).
pub const AUX_NEXT_BAR_MSE_EMA_INDEX: usize = 113;
pub const AUX_REGIME_CE_EMA_INDEX: usize = 114;
/// ISV slot [115] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
///
/// Design note: this is NOT a version number. There is no ordered version space.
/// The fingerprint is a structural hash that automatically changes whenever any
@@ -721,14 +779,17 @@ pub const VSN_MASK_GROUP_5_EMA_INDEX: usize = 110;
/// EMA replacing legacy CPU-DtoH path), 94→97 in Plan 4 Task 2c.3c.5
/// (H_S2_RMS_EMA producer slot), 97→103 in Plan 4 Task 3 E.3 (IQN
/// fixed-τ multi-quantile head, +4 diagnostic slots), 103→111 in
/// Plan 4 Task 1B-ii (VSN per-group mask EMA, +6 slots).
pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 111;
/// ISV slot [112] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
/// Plan 4 Task 1B-ii (VSN per-group mask EMA, +6 slots), 111→115 in
/// Plan 4 Task 6 Commit A (aux-head loss EMA, +2 slots — gap at [112]
/// preserved to keep the new slots bin-aligned with the existing
/// gap-anchored cadence).
pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 115;
/// ISV slot [116] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
/// Shifted 62→70 in Plan 3 Task 1, 70→74 in Plan 3 Task 3, 74→77 in Plan 3 Task 4 B.4,
/// 77→81 in Plan 3 Task 7 C.3, 81→86 in Plan 3 Task 8 B.3, 86→91 in Plan 4 Task 5,
/// 91→95 in Plan 4 follow-up, 95→98 in Plan 4 Task 2c.3c.5, 98→104 in Plan 4
/// Task 3 E.3, 104→112 in Plan 4 Task 1B-ii.
pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 112;
/// Task 3 E.3, 104→112 in Plan 4 Task 1B-ii, 112→116 in Plan 4 Task 6 Commit A.
pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 116;
/// Canonical alias for the fingerprint slot (the low half).
pub const ISV_LAYOUT_FINGERPRINT_INDEX: usize = ISV_LAYOUT_FINGERPRINT_LO_INDEX;
@@ -810,8 +871,9 @@ const fn layout_fingerprint_seed() -> &'static [u8] {
IQN_Q_P05_EMA=99;IQN_Q_P25_EMA=100;IQN_Q_P75_EMA=101;IQN_Q_P95_EMA=102;\
VSN_MASK_GROUP_0_EMA=105;VSN_MASK_GROUP_1_EMA=106;VSN_MASK_GROUP_2_EMA=107;\
VSN_MASK_GROUP_3_EMA=108;VSN_MASK_GROUP_4_EMA=109;VSN_MASK_GROUP_5_EMA=110;\
ISV_LAYOUT_FINGERPRINT_LO=111;ISV_LAYOUT_FINGERPRINT_HI=112;\
ISV_TOTAL_DIM=113;\
AUX_NEXT_BAR_MSE_EMA=113;AUX_REGIME_CE_EMA=114;\
ISV_LAYOUT_FINGERPRINT_LO=115;ISV_LAYOUT_FINGERPRINT_HI=116;\
ISV_TOTAL_DIM=117;\
PARAM_W_A_H_S1=0;PARAM_B_A_H_S1=1;PARAM_W_B_H_S1=2;PARAM_B_B_H_S1=3;\
PARAM_W_RESIDUAL_H_S1=4;PARAM_GAMMA_H_S1=5;PARAM_BETA_H_S1=6;\
PARAM_W_A_H_S2=7;PARAM_B_A_H_S2=8;PARAM_W_B_H_S2=9;PARAM_B_B_H_S2=10;\
@@ -858,7 +920,9 @@ const fn layout_fingerprint_seed() -> &'static [u8] {
PARAM_VSN_W1_G3=107;PARAM_VSN_B1_G3=108;PARAM_VSN_W2_G3=109;PARAM_VSN_B2_G3=110;\
PARAM_VSN_W1_G4=111;PARAM_VSN_B1_G4=112;PARAM_VSN_W2_G4=113;PARAM_VSN_B2_G4=114;\
PARAM_VSN_W1_G5=115;PARAM_VSN_B1_G5=116;PARAM_VSN_W2_G5=117;PARAM_VSN_B2_G5=118;\
PARAM_TOTAL_TENSORS=119"
PARAM_AUX_NB_W1=119;PARAM_AUX_NB_B1=120;PARAM_AUX_NB_W2=121;PARAM_AUX_NB_B2=122;\
PARAM_AUX_RG_W1=123;PARAM_AUX_RG_B1=124;PARAM_AUX_RG_W2=125;PARAM_AUX_RG_B2=126;\
PARAM_TOTAL_TENSORS=127"
}
/// Compile-time layout fingerprint. Burned into the binary at build time.
@@ -1258,7 +1322,17 @@ impl Default for CausalInterventionConfig {
/// `VSN_HIDDEN_DIM = 16`) → indices [95..119). Producer-only params in
/// this commit; the `vsn_softmax_and_gate_forward` consumer kernel and
/// the cuBLAS `Linear_1`/`Linear_2` wire-in land in 1B-iii.
pub(crate) const NUM_WEIGHT_TENSORS: usize = 119;
/// + 8 Plan 4 Task 6 Commit A aux-head MLP params (2 heads × 4 tensors —
/// `aux_nb_{w1,b1,w2,b2}` for the next-bar regression head with K=1,
/// `aux_rg_{w1,b1,w2,b2}` for the regime classification head with K=5,
/// shared hidden dim `AUX_HIDDEN_DIM = 32`) → indices [119..127).
/// Producer-only params in this commit; the `aux_*_forward` consumer
/// kernels and the training-loop loss accumulation wire-in land in
/// Commit B. Adam SAXPY iterates `0..NUM_WEIGHT_TENSORS` so these
/// tensors will receive zero-gradient SAXPYs each step until Commit B
/// activates the backward — Adam keeps them at Xavier init meanwhile,
/// which is the intended dormant state.
pub(crate) const NUM_WEIGHT_TENSORS: usize = 127;
/// Compute the size (element count) of each weight tensor.
///
@@ -1429,9 +1503,32 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT
idx += 4;
g += 1;
}
debug_assert!(idx == NUM_WEIGHT_TENSORS,
"compute_param_sizes: VSN tail fill ended at {} but NUM_WEIGHT_TENSORS={}",
idx, NUM_WEIGHT_TENSORS);
debug_assert!(idx == 119,
"compute_param_sizes: VSN tail fill ended at {} but expected 119 before aux-heads",
idx);
// ── Plan 4 Task 6 Commit A: aux-head MLP params (8 tensors) ───
// Two heads: next-bar regression (K=1) + regime classification (K=5).
// Each head is `Linear(SH2 → AUX_HIDDEN_DIM=32) → ELU → Linear(32 → K)`.
// Per-head total = 32*SH2 + 32 + K*32 + K. Sum over both heads:
// 2*32*SH2 + 2*32 + (1+5)*32 + (1+5) = 2*32*SH2 + 64 + 192 + 6
// = 2*32*SH2 + 262 floats.
// No production callers in this commit — params consumed by Commit B
// (forward orchestrator + training-loop loss accumulation + backward
// dispatcher).
let sh2 = cfg.shared_h2;
sizes[119] = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM * sh2; // [119] aux_nb_w1 [32, SH2]
sizes[120] = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; // [120] aux_nb_b1 [32]
sizes[121] = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; // [121] aux_nb_w2 [1, 32]
sizes[122] = 1; // [122] aux_nb_b2 [1]
sizes[123] = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM * sh2; // [123] aux_rg_w1 [32, SH2]
sizes[124] = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; // [124] aux_rg_b1 [32]
sizes[125] = crate::cuda_pipeline::gpu_aux_heads::AUX_REGIME_K
* crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM; // [125] aux_rg_w2 [5, 32]
sizes[126] = crate::cuda_pipeline::gpu_aux_heads::AUX_REGIME_K; // [126] aux_rg_b2 [5]
debug_assert!(NUM_WEIGHT_TENSORS == 127,
"compute_param_sizes: NUM_WEIGHT_TENSORS expected 127 after aux-heads, got {}",
NUM_WEIGHT_TENSORS);
sizes
}
@@ -16986,6 +17083,29 @@ impl GpuDqnTrainer {
}
}
// ── Plan 4 Task 6 Commit A: aux-head fan dims (8 tensors at [119..127)) ─
// Both heads: `Linear(SH2 → AUX_HIDDEN_DIM=32) → ELU → Linear(32 → K)`.
// Xavier on weight matrices (uses standard `sqrt(6/(fan_in+fan_out))`);
// zero on biases. Cold-start: predictions and logits ≈ 0 → next-bar
// pred ≈ 0 (matches the typical centred next-bar return), regime
// logits ≈ uniform softmax (no class preference).
// Indices: aux_nb_w1=119, aux_nb_b1=120, aux_nb_w2=121, aux_nb_b2=122,
// aux_rg_w1=123, aux_rg_b1=124, aux_rg_w2=125, aux_rg_b2=126.
{
let h_dim = crate::cuda_pipeline::gpu_aux_heads::AUX_HIDDEN_DIM;
let k_rg = crate::cuda_pipeline::gpu_aux_heads::AUX_REGIME_K;
// Next-bar head (K=1).
fan_dims[119] = (h_dim, cfg.shared_h2); // aux_nb_w1 [32, SH2] (Xavier)
fan_dims[120] = (0, 0); // aux_nb_b1 (zero)
fan_dims[121] = (1, h_dim); // aux_nb_w2 [1, 32] (Xavier)
fan_dims[122] = (0, 0); // aux_nb_b2 (zero)
// Regime head (K=5).
fan_dims[123] = (h_dim, cfg.shared_h2); // aux_rg_w1 [32, SH2] (Xavier)
fan_dims[124] = (0, 0); // aux_rg_b1 (zero)
fan_dims[125] = (k_rg, h_dim); // aux_rg_w2 [5, 32] (Xavier)
fan_dims[126] = (0, 0); // aux_rg_b2 (zero)
}
// Build flat host buffer: Xavier init for weights, zeros for biases + padding.
let cutlass_tile_pad = 32 * cfg.adv_h.max(cfg.value_h);
let buf_len = total + cutlass_tile_pad;

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

@@ -41,6 +41,7 @@ pub mod gpu_iqn_head;
pub mod gpu_attention;
pub mod gpu_tlob;
pub mod gpu_grn;
pub mod gpu_aux_heads;
pub mod decision_transformer;
pub mod learning_health;
pub mod q_snapshot;

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

@@ -331,6 +331,17 @@ impl StateResetRegistry {
category: ResetCategory::FoldReset,
description: "ISV[VSN_MASK_GROUP_5_EMA_INDEX=110] — `plan_isv` group importance mask EMA; same producer/contract as g0; cold-start 1/6",
},
// ───── Plan 4 Task 6 Commit A: aux-head loss EMAs ─────────
RegistryEntry {
name: "isv_aux_next_bar_mse_ema",
category: ResetCategory::FoldReset,
description: "ISV[AUX_NEXT_BAR_MSE_EMA_INDEX=113] — next-bar regression head batch-mean MSE EMA (α=0.05); GPU aux_heads_loss_ema_update kernel fills (Plan 4 Task 6 Commit A; producer kernel landed in Commit A, launch site wires in Commit B alongside launch_h_s2_rms_ema). Cold-start 0.0 (no aux forwards have happened yet) reapplied at fold boundary. Diagnostic only — no GPU consumer kernel reads slot 113 in either commit",
},
RegistryEntry {
name: "isv_aux_regime_ce_ema",
category: ResetCategory::FoldReset,
description: "ISV[AUX_REGIME_CE_EMA_INDEX=114] — regime classification head batch-mean cross-entropy EMA; same producer/contract as next_bar_mse_ema; cold-start 0.0",
},
];
Self { entries }
}

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

@@ -4578,6 +4578,27 @@ impl DQNTrainer {
fused.trainer().write_isv_signal_at(idx, uniform);
}
}
// Plan 4 Task 6 Commit A: aux-head loss EMA producer slots
// (next-bar regression MSE + regime classification CE). Reset to
// 0.0 at fold boundary so the first `aux_heads_loss_ema_update`
// fire on the new fold EMAs the measured per-head loss toward 0
// rather than carrying a stale fold-N estimate. Diagnostic only —
// no GPU consumer kernel reads slots [113..115) in either Commit
// A or B (Commit B's HEALTH_DIAG line is a CPU-side text
// emitter, not a kernel input). Producer kernel landed in
// Commit A; launch site wires in Commit B.
"isv_aux_next_bar_mse_ema" | "isv_aux_regime_ce_ema" => {
if let Some(ref fused) = self.fused_ctx {
let idx = match name {
"isv_aux_next_bar_mse_ema"
=> crate::cuda_pipeline::gpu_dqn_trainer::AUX_NEXT_BAR_MSE_EMA_INDEX,
"isv_aux_regime_ce_ema"
=> crate::cuda_pipeline::gpu_dqn_trainer::AUX_REGIME_CE_EMA_INDEX,
_ => unreachable!(),
};
fused.trainer().write_isv_signal_at(idx, 0.0);
}
}
"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