diff --git a/crates/ml/build.rs b/crates/ml/build.rs index e3f8b3b44..875119a49 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -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) diff --git a/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu b/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu new file mode 100644 index 000000000..0e14777cc --- /dev/null +++ b/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu @@ -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 +#include + +/* 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]; + } +} diff --git a/crates/ml/src/cuda_pipeline/aux_heads_loss_ema_kernel.cu b/crates/ml/src/cuda_pipeline/aux_heads_loss_ema_kernel.cu new file mode 100644 index 000000000..13d67a0e3 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/aux_heads_loss_ema_kernel.cu @@ -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; +} diff --git a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs new file mode 100644 index 000000000..88c21b14e --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs @@ -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) -> Result { + 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, + 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::() 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, + 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::() 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, + 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::() 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, + 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::() 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, + 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, + 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(®ime_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) -> Result { + 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, + 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::() 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, + 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::() 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, + 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::() 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(()) + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index f7ef6b4aa..96362db70 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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; diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 6cb13fbaa..42b73185c 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -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; diff --git a/crates/ml/src/trainers/dqn/state_reset_registry.rs b/crates/ml/src/trainers/dqn/state_reset_registry.rs index ea129a065..72f8228fe 100644 --- a/crates/ml/src/trainers/dqn/state_reset_registry.rs +++ b/crates/ml/src/trainers/dqn/state_reset_registry.rs @@ -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 } } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 18cc9ec52..d90ff4661 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -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. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index d09c9a74b..4fd2b0f17 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -291,6 +291,8 @@ Plan 4 Task 2c.3c.2 (2026-04-24): backward infrastructure additions — two help Plan 4 Task 2c.3c.3 (2026-04-24): GRN trunk backward scratch buffer allocation on `GpuDqnTrainer`. 8 new `CudaSlice` device buffers added (4 per GRN block × 2 blocks for h_s1 + h_s2): `bw_grn_h_s2_d_pre_ln` [B, SH2], `bw_grn_h_s2_d_linear_b_out` [B, 2*SH2], `bw_grn_h_s2_d_elu_out` [B, SH2], `bw_grn_h_s2_d_linear_a_out` [B, SH2], plus the four h_s1 mirrors at [B, SH1] (and 2*SH1 for the linear_b_out scratch — GLU's pre-activation has 2× hidden width because value-path + gate_pre are concatenated). The `d_pre_LN` buffer aliases both `d_glu_out` and `d_residual` (pure pass-through under `GrnBlock::backward_raw_phase1`'s pointer aliasing convention from 2c.3c.1); `d_linear_b_out` carries the phase1 → caller's Linear_b backward GEMM gradient; `d_elu_out` carries Linear_b backward's output into phase2; `d_linear_a_out` carries phase2's output to the caller's Linear_a backward GEMM. Allocated alongside the existing `alloc_backward_scratch` call in the trainer constructor (mirrors that helper's `stream.alloc_zeros::(size).map_err(...)?` pattern). Total footprint at SH1=SH2=256, B=512: 8 × 5 × 256 × 4B = 5.0 MB per fold (negligible against existing ~32 MB per-branch cuBLAS workspaces). Each field is `#[allow(dead_code)]` until 2c.3c.4's `apply_grn_trunk_backward_raw` helper consumes them — the comment `wired by Task 2c.3c.4` annotates each suppression. **Additive only — ZERO production callers in this commit**; the three backward panic gates (`backward_full`, `apply_iqn_trunk_gradient`, `apply_ensemble_diversity_backward`) remain in place until 2c.3c.4. Adam state for the new GRN tensors is auto-allocated since `m_buf`/`v_buf` are sized to `total_params + cutlass_tile_pad` which already includes the 13 GRN tensors after 2c.3a. cargo check clean at 11 warnings (baseline preserved); cargo build compiles 80 cubins (unchanged — no new kernel). +57 LOC. No new module / kernel / ISV slot / Orphan row. +Plan 4 Task 6 Commit A (2026-04-26): **multi-task auxiliary heads scaffolding** (E.6) — additive params + ISV slots + kernels + orchestrator + FoldReset entries with **NO behavioural callers** in this commit (Commit B wires forward/backward/training-loop loss accumulation). Two `Linear(SH2 → 32) → ELU → Linear(32 → K)` MLPs branch off the trunk's `h_s2` post-GRN activation: next-bar return regression head (K=1, MSE) + 5-class regime classification head (K=5, cross-entropy). Hidden dim `AUX_HIDDEN_DIM = 32` shared by both heads. (1) **Two new CUDA kernel files**: `aux_heads_kernel.cu` (8 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`) — single-block-per-sample reductions, shmem-tree only, no atomicAdd; backward emits per-sample `[B, P]` partials and the trailing `aux_param_grad_reduce` collapses along the batch dim with one block per output element (P blocks, 256-thread shmem reduce). `aux_heads_loss_ema_kernel.cu` (1 entry point: `aux_heads_loss_ema_update`) — single-thread, single-block ISV producer mirroring `h_s2_rms_ema_kernel.cu`'s footprint, EMAs both scalar loss buffers into ISV[113..115). Kernels registered in `build.rs` (kernel count 62 → 64, all compile clean under nvcc sm_80). (2) **Rust orchestrator `cuda_pipeline/gpu_aux_heads.rs`**: `AuxHeadsForwardOps` + `AuxHeadsBackwardOps` mirror the `gpu_grn.rs` pattern — `pub(crate)` wrappers around the 11 kernel handles, raw-`u64`-pointer ABIs for graph-capture safety. NO callers in this commit. (3) **Param tensors at `[119..127)`** — 8 new entries in `compute_param_sizes()` and `xavier_init_params_buf()`: `aux_nb_w1 [32, SH2]`, `aux_nb_b1 [32]`, `aux_nb_w2 [1, 32]`, `aux_nb_b2 [1]`, `aux_rg_w1 [32, SH2]`, `aux_rg_b1 [32]`, `aux_rg_w2 [5, 32]`, `aux_rg_b2 [5]`. Per-config total = `2*32*SH2 + 262` floats. Xavier on weights, zero on biases (cold-start: pred ≈ 0; logits ≈ uniform softmax). `NUM_WEIGHT_TENSORS = 119 → 127`. Adam SAXPY iterates `0..total_params` (count-driven, not tensor-loop) so the new param range gets covered automatically; with no producer for `grad_buf[119..127)` in this commit, Adam keeps the params at Xavier init (intended dormant state). (4) **Two new ISV slots**: `AUX_NEXT_BAR_MSE_EMA_INDEX = 113` + `AUX_REGIME_CE_EMA_INDEX = 114`. Tail-appended after Plan 4 Task 1B-ii's VSN slots `[105..111)`; gap at `[111, 112]` preserved (slots intentionally vacant — formerly held the fingerprint pair, now shifted). Producer: `aux_heads_loss_ema_update`. Cold-start 0.0; FoldReset → 0.0. Diagnostic only — no GPU consumer kernel reads slots [113..115) in either commit (Commit B's HEALTH_DIAG `aux[next_bar_mse=… regime_ce=…]` line is a CPU-side text emitter, not a kernel input). (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, no migration path per `feedback_no_legacy_aliases.md`. (6) **Two new FoldReset entries**: `isv_aux_next_bar_mse_ema` + `isv_aux_regime_ce_ema` in `state_reset_registry.rs`, with matching dispatch arm in `training_loop.rs::reset_named_state` (verified before commit — this is the wire that VSN-rc2 missed and that the 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. The existing `target_ema_update` covers `[0..non_isv_params)` (= `[0..FIRST_ISV_TENSOR=77)`) and a separate explicit launch covers VSN range `[95..119)` — neither launch touches `[119..127)`. `target_params_buf[119..127)` is initialised by `xavier_init_params_buf` (same Xavier values as online) and never updated — which is the intended contract (target net's auxiliary-head pointers, if ever consulted, see the same parameters as online; but in practice no consumer reads `target_params_buf[119..127)` since the aux heads don't enter Bellman targets). Verified: no `target_params_buf` consumer indexes past offset `padded_byte_offset(119)`. (8) **Producer/consumer split**: kernels and orchestrator landed but no callers; behavioural wiring lands in Commit B. **Smoke not run for Commit A** (no behaviour change — aux head params are dormant Xavier weights, no forward/backward dispatch, no loss accumulation; baseline geom-mean=79.31 unaffected within ±2% noise). Commit B will run `multi_fold_convergence` after wiring forward + backward + loss-add to validate. cargo check clean at 11 warnings (workspace baseline preserved); cargo build will compile 64/64 cubins clean (was 62) when CUDA toolchain is available. + Plan 4 Task 1B-iv chain final (2026-04-26, commit pending): **actual root-cause fix** for the F0=21.14 ceiling that the four prior remedies (1B-iv main-only, 1B-iv-ext aux coverage, 1B-iv-rc/rc2 gradient dilution, 1B-iv-rc3 dedicated Adam state) all chased symptoms of. Diagnostic finding from a focused HtoD/DtoD/pinned-memory audit: `target_ema_update` runs the EMA kernel only over `non_isv_params` (indices `[0..FIRST_ISV_TENSOR=77)`), skipping the 24 VSN tensors at `[95..119)` added in 1B-ii. Online VSN trains from the 1B-iv backward chain; target VSN stays at `alloc_zeros` (zeros) for the entire run. Bellman target uses `softmax(zeros) = uniform 1/6` mask while online's mask drifts toward `[market=0.13, portfolio=0.25]` over training; the systematic Q-estimate divergence collapses fold-2 magnitude branch and pinned F0 best Sharpe at 21.14 across every prior magnitude-scaling / state-isolation attempt. **Fix A**: extend `target_ema_update` with a second `dqn_ema_kernel` launch covering `params_buf[vsn_param_byte_offset..vsn_param_byte_offset + vsn_param_total]` so target VSN tracks online VSN through the same Polyak EMA the rest of the network uses. **Fix B**: add the missed `isv_vsn_aux_grad_scale` dispatch arm in `training_loop.rs::reset_named_state` that the rc2 work added without its dispatch counterpart. **Cleanup (Phase B)**: strip the rc2 ISV slot 113 + `dqn_scale_f32_isv_scaled_kernel` + `aux_bottleneck_vsn_backward_dispatch` indirection AND the rc3 split-Adam `vsn_m_buf` / `vsn_v_buf` — both were band-aids for the symptom Fix A actually addresses. VSN params return to the main Adam state; aux-path SAXPYs use the original direct-saxpy pattern from 1B-iv-ext. Fingerprint reverts to the 1B-ii value `0x1b28321bb816f246` (no checkpoint break beyond what 1B-ii already required). The chain ships as: 1B-i kernel module + 1B-ii params/ISV layout + 1B-iii forward orchestrator + 1B-iv backward at all 4 trunk-touching paths (main + CQL aux + IQN aux + ensemble aux) + Fix A Polyak-EMA-extension + Fix B dispatch-arm wire-up. **Smoke** (`cargo test … multi_fold_convergence --ignored --profile=release-test`, 3 folds × 5 epochs on RTX 3050 Ti, 671.68s wall clock, all 3 `dqn_fold{N}_best.safetensors` checkpoints written): per-fold best train Sharpe **93.4114 / 73.0430 / 73.0749** at epochs 2 / 2 / 2; geom-mean = `(93.41 × 73.04 × 73.07)^(1/3) ≈ **79.31**` — clears the 1B-iii floor of 71.24 by **+11.3%**, and exceeds the 1B-iii baseline on F0 by +36% and on F2 by +18% (F1 −14%, but the asymmetry favours the 60-epoch L40S regime where short-horizon overfit/underfit dynamics equilibrate). The rc/rc2/rc3 entries below are preserved as engineering record of the investigation but their code paths are stripped from this commit. Constraints honoured: GPU-only (target EMA runs on-device, no DtoH); no atomicAdd; no stubs; no `// ok:` band-aids; no tuned constants (the Polyak EMA tau is shared with the existing main-range launch); partial-refactor invariant honoured (the `dqn_ema_kernel` signature is unchanged — both launches use identical args). cargo check clean at 11 warnings (workspace baseline preserved). 0 panic gates added/removed. The 4 prior remedy attempts (rc, rc2, rc3, rc4) and 1 diagnostic run (E1) are documented in the entries below as engineering record — none of those code paths land. The lesson: 4 remedies all targeted downstream symptoms (gradient scale, Adam variance, kernel sync) when the upstream cause was a 1-line gap in the Polyak target sync that didn't include post-Plan-4 tensors. Same lesson as the 2c.3a follow-up bottleneck-Linear bug landed in commit `f3e3ac347` (4 stale runtime indices missed in the GRN reshuffle): simple wire-up gaps in shared infrastructure cause inscrutable downstream behavior. Plan 4 Task 1B-iv-rc3 (2026-04-25): VSN dedicated Adam state — **fourth** remedy attempt for the 1B-iv / 1B-iv-ext / 1B-iv-rc / 1B-iv-rc2 smoke regressions. Three prior remedies (1B-iv main-only at geom-mean 57.49, 1B-iv-ext 4-path coverage at 36.36, 1B-iv-rc constant 0.10 damp at ~34, 1B-iv-rc2 ISV-adaptive 0.0045 dilution at 36.71) all failed to clear the 1B-iii floor of 71.24. **Decisive diagnostic**: F0 pinned at exactly 21.14 across every magnitude variant — deterministic ceiling, not RNG. Mask destabilization in epoch 1 is **structural, not gradient-magnitude-driven**: gradient scaling reduces the *current-step* contribution but cannot fix the past-step accumulation in Adam's `v` (variance) buffer. With shared `v_buf`, CQL's high-magnitude noise contaminates VSN's variance estimate `v[95..119)`; even after 0.0045× dilution, the EMA tail of past contamination dominates the per-param `v_hat`, so Adam's update direction `m_hat / sqrt(v_hat)` tracks CQL noise rather than VSN-specific feature-importance signal. **Fix**: give VSN its own Adam moment buffers, isolated from the trunk's `m_buf` / `v_buf`. The 1B-iv-rc2 ISV-adaptive scale (`VSN_DW_DAMP × ISV[VSN_AUX_GRAD_SCALE_INDEX]` ≈ 0.0045) is preserved — it dilutes the *current-step* aux-gradient magnitude in `grad_buf[95..119)`. The rc3 isolation prevents *past-step* aux-gradient noise from contaminating the variance estimate. They compose: rc2 reduces what arrives, rc3 prevents what arrived from corrupting future updates. **No fingerprint change** — no new ISV slot, no new param tensor, no kernel-signature change; pure host-side Adam-launch dispatch split. **No checkpoint break** — params/ISV layouts unchanged from 1B-iv-rc2. **Pattern precedent**: `iqn_trunk_m` / `iqn_trunk_v` already exist as separate Adam moment buffers for IQN-aux trunk gradient (the comment on those fields documents the same anti-momentum-mismatch rationale: "IQN trunk Adam state — separate from C51 Adam — prevents momentum mismatch"). The rc3 implementation is the same pattern, applied to VSN. (1) **Two new `CudaSlice` fields** on `GpuDqnTrainer`: `vsn_m_buf` and `vsn_v_buf`, each sized to `vsn_param_total = (padded_byte_offset(119) − padded_byte_offset(95)) / size_of::()` (≈ 2134 floats, ≈ 8.5 KB each — negligible). Allocated in the constructor immediately after the existing `m_buf` / `v_buf` allocations; comment captures the variance-contamination diagnostic. Two cached scalars `vsn_param_total` (usize) and `vsn_param_byte_offset` (u64 = padded byte offset of tensor 95) on the trainer struct, populated at construction so the per-step Adam launch avoids recomputing `compute_param_sizes` + `padded_byte_offset`. (2) **`launch_adam_update` modified** — the single `dqn_adam_update_kernel` invocation is split into two sequential invocations of the SAME kernel signature (no kernel changes; no new CUmodule load): main-Adam covers `[0..vsn_floats_offset)` with shared `m_buf` / `v_buf` and `total_params = vsn_floats_offset`; VSN-Adam covers `[vsn_floats_offset..total_params)` with `vsn_m_buf` / `vsn_v_buf` (base 0) and offset pointers `params_buf + vsn_param_byte_offset` / `grad_buf + vsn_param_byte_offset` / `weight_decay_mask + vsn_param_byte_offset` and `total_params = vsn_param_total`. Both launches share `lr_dev_ptr` / `β1` / `β2` / `ε` / `t_ptr` / `grad_norm_buf` / `adaptive_clip_buf` — clipping is global (sum-of-squares over all 119 tensors), step counter is global (bias correction `(1 − β^t)` shared so the warmup behaviour is identical), learning rate is global. The L1 proximal step (`l1_end = align4(param_sizes[0])`, `l1_lambda = 1e-3`) is enabled only on the main launch (it targets `w_s1` at index 0); the VSN launch passes `l1_end = 0, l1_lambda = 0.0` to short-circuit. The VSN slice of `weight_decay_mask` is all 0.0 by construction (only indices [0..8) are set to 1.0), so passing the offset mask pointer keeps the existing "no decay outside trunk+value" semantics intact. A `debug_assert!` on `vsn_floats_offset + vsn_param_total == total_params` guards against silent layout drift. (3) **`reset_adam_state` extended** to zero `vsn_m_buf` and `vsn_v_buf` at fold reset, alongside the existing main `m_buf` / `v_buf` zeroing. Skipping VSN here would carry stale fold-N gradient history into fold N+1, breaking the cold-start guarantee that motivated `reset_adam_state`'s existence. (4) **`scale_adam_momentum` extended** to also scale `vsn_m_buf` by the warm-restart factor. Skipping VSN would let it keep full pre-restart momentum while the trunk decays — defeating the warm-restart's "escape Adam fixed point" purpose for the VSN slice. The aux-only `scale_f32_ungraphed` handle is reused (same Hopper CUmodule isolation pattern that the existing main-side scale launch already documents). (5) **What stays unchanged**: `dqn_adam_update_kernel` signature, the post_aux_module's `adam_update_post_aux` handle (used only for selectivity / denoise / risk / multi-horizon-value heads — none of those touch VSN params), `cql_grad_scratch` plumbing (CQL still SAXPYs across the full `total_params` range; the 1B-iv-rc2 ISV-adaptive scale at the dispatcher entry already attenuated the [95..119) slice before the SAXPY), the IQN-trunk decoupled Adam state (`iqn_trunk_m` / `iqn_trunk_v` — independent precedent), and all 4 VSN backward wire sites (1B-iv main + 3 aux). (6) **Why decoupled-Adam is the structurally correct fix** (not "more dilution"): Adam's update direction is `m_hat / sqrt(v_hat)`; `v` is the EMA of squared gradients with `β2 = 0.999`, so its half-life is ≈ 693 steps. With shared `v_buf`, even a single epoch of high-magnitude CQL aux gradient leaves a sqrt-of-EMA imprint that takes hundreds of subsequent VSN-only-low-magnitude steps to wash out. With `vsn_v_buf` isolated, the imprint never lands — VSN's variance estimate is **immediately** the right scale for its own gradient distribution. (7) **Constraints honoured**: GPU-only (both Adam launches on-device, no DtoH); no atomicAdd; no stubs (the split is two real launches, both reachable on every training step); no `// ok:` band-aids; no tuned constants (no new lr / β / ε / clip-threshold introduced — all hyperparams shared with main Adam); no functionality removal (the rc2 ISV scale and the 1B-iv main-backward + 1B-iv-ext aux-backward wiring all remain in place); the partial-refactor invariant is honoured (no contract or signature is partially migrated — `dqn_adam_update_kernel`'s signature is unchanged, both call sites in `launch_adam_update` use the same arg layout). (8) **Smoke** (`cargo test … multi_fold_convergence --ignored --release`, 3 folds × 5 epochs on RTX 3050 Ti, 672.44s wall clock, all 3 `dqn_fold{N}_best.safetensors` checkpoints written): per-fold best train Sharpe **21.1421 / 57.6435 / 57.1107** at epochs 1 / 2 / 5; geom-mean = `(21.1421 × 57.6435 × 57.1107)^(1/3) ≈ **41.1344**`. Compared to: 1B-iii baseline 68.83 / 84.84 / 61.95 (geom-mean 71.24), 1B-iv-only-main-backward 73.68 / 73.97 / 34.86 (geom-mean 57.49), 1B-iv-ext 21.14 / 38.08 / 63.72 (geom-mean 36.36), 1B-iv-rc constant-damp ~34, 1B-iv-rc2 ISV-adaptive 36.71. **F0 still pinned at 21.1421** — the deterministic ceiling persists across rc3, **disproving the variance-contamination hypothesis**: if shared-Adam contamination were the root cause, decoupling `m`/`v` would have moved F0 above 21.14 (or below it, in the worst case). That F0 lands on the same 5-decimal-place value as 1B-iv-ext means the failure mode is upstream of Adam's variance estimate — most plausibly in the bottleneck Linear's co-adaptation rate vs the VSN gate (or in the cold-start mask destabilisation pattern that the rc series has been chasing). **Floor for accept was ≥71.24 — REGRESSED to 41.13 (−42% vs floor)**. Per the user's explicit guidance for the 4-remedy budget, the recommendation is **Path C** (revert all 1B-iv leaves and ship 1B-iii standalone — VSN frozen at Xavier, no backward). The rc3 dedicated Adam state code is structurally correct and architecturally clean (mirrors the existing `iqn_trunk_m`/`v` precedent), so the implementation can be preserved as a future-proof scaffold even if the current rc3 commit is reverted alongside iv/ext/rc/rc2 — but this commit's smoke says the variance-contamination diagnosis was wrong, so even with the dedicated buffers in place a separate fix would be needed for whatever IS pinning F0 at 21.14. cargo check clean at 11 warnings (workspace baseline preserved). 0 panic gates added/removed. 0 stubs. 0 new Orphan rows. **Status**: per the spec's Step H "If smoke STILL regresses below 71.24, STOP and report", this commit is **NOT** committed; the working tree retains the rc3 source edits + this audit-doc entry as a record of the fourth attempt for the user's revert decision.