diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 7591f16b3..b66d1e191 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -821,6 +821,18 @@ fn main() { // separate from the K=2 sibling for diagnostic isolation, sparse- // label semantic clarity, and future per-class weighting headroom. "aux_trade_outcome_loss_reduce_kernel.cu", + // SP22 H6 vNext Phase A5 (2026-05-13): backward kernel for the + // K=3 trade-outcome aux head. Mirrors `aux_next_bar_backward` + // (K=2 sibling): per-sample partials (dW1, db1, dW2, db2, + // dh_s2_aux), d_logits = (softmax − one_hot)/B_valid with mask + // handling, ELU backward via post-activation identity. Sparse- + // label arithmetic amplifies per-trade-close gradient magnitude + // (smaller B_valid → larger inv_B); Phase E's Adam may need + // class-weighted CE or per-group LR tuning. Reuses the existing + // `aux_param_grad_reduce` kernel from aux_heads_kernel.cu for + // the per-sample → final-grad batch-reduce step. Dead code + // until Phase B wireup lands the Rust launcher chain. + "aux_trade_outcome_backward_kernel.cu", // SP22 H6 (2026-05-12): per-env p_up extractor that copies // `aux_softmax[env, 1]` → `prev_aux_dir_prob[env]` after each aux // forward. The cache buffer is read by `experience_state_gather` diff --git a/crates/ml/src/cuda_pipeline/aux_trade_outcome_backward_kernel.cu b/crates/ml/src/cuda_pipeline/aux_trade_outcome_backward_kernel.cu new file mode 100644 index 000000000..25d3979b0 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/aux_trade_outcome_backward_kernel.cu @@ -0,0 +1,252 @@ +// crates/ml/src/cuda_pipeline/aux_trade_outcome_backward_kernel.cu +// +// SP22 H6 vNext Phase A5 (2026-05-13) — backward kernel for the K=3 +// trade-outcome aux head. Closes the forward → loss → backward chain +// for the new aux head (A3 forward, A4 loss-reduce, A5 backward). +// +// Mirrors `aux_next_bar_backward` (K=2 sibling) line-for-line because +// the gradient flow doesn't depend on K — it's just +// d_logits[b, kc] = (softmax[b, kc] − one_hot(labels[b], kc)) / B_valid +// propagated through the same Linear → ELU → Linear chain. Masked rows +// (`labels[b] == -1`) zero the K-vector, which zeroes the per-sample +// partials for that row (the chain rule's multiplicative zero collapses +// everything downstream). +// +// Per-sample partials (caller reduces along batch dim, mirroring +// aux_next_bar_backward's contract): +// db2_partial[b, kc] = d_logits[b, kc] +// dW2_partial[b, kc, k] = d_logits[b, kc] × h_post[b, k] +// db1_partial[b, k] = d_h_pre[b, k] +// dW1_partial[b, k, j] = d_h_pre[b, k] × h_s2_aux[b, j] +// dh_s2_aux_out[b, j] = sum_k d_h_pre[b, k] × w1[k, j] +// +// Where: +// 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]) +// +// Sparse-label arithmetic +// ─────────────────────── +// Trade-close events are ~1-5% of bars, so `B_valid` from the +// `aux_trade_outcome_loss_reduce` kernel is typically small relative to +// the nominal batch size. `inv_B = 1 / B_valid` is therefore LARGER than +// the K=2 sibling's `inv_B = 1 / (~B)` — per-trade-close gradients have +// proportionally higher magnitude. This is correct credit assignment +// (each valid trade-close speaks louder because it's rarer signal) but +// Phase E's Adam β1/β2/ε per-group hyperparams may need to compensate +// to keep the head's update magnitude in a similar range to the K=2 +// sibling. +// +// SP14 Phase C.5b separation +// ────────────────────────── +// Reads `h_s2_aux` (aux trunk output) NOT Q's `h_s2`. Writes into +// `dh_s2_aux_out` which SAXPYs into the aux trunk's `dh_s2_aux_accum`, +// NOT Q's `bw_d_h_s2`. The aux trunk's backward kernel terminates at +// the encoder boundary (no `dx_in` output param), so Q's encoder is +// structurally protected from aux contamination per +// `feedback_no_partial_refactor` and SP14 Phase C.5b's structural +// stop-gradient design. +// +// Phase B forward-reference note +// ────────────────────────────── +// Phase B will switch the forward kernel's input to concat +// `[h_s2_aux (256) || plan_params (6)] = 262-dim`. This backward kernel +// will need a corresponding update: either (a) split `dh_s2_aux_out` so +// the trailing 6 elements feed back to plan_params source (or use +// stop-grad on those elements), or (b) bump SH2 from 256 → 262 and let +// the downstream wiring handle the split. Phase B lands the decision. +// For Phase A5 (this commit) the signature stays K=3-only with the +// original SH2 contract — the kernel is dead code at this stage. +// +// Discipline +// ────────── +// - No atomicAdd per `feedback_no_atomicadd.md` — per-sample partials +// emitted; a separate `aux_param_grad_reduce` pass aggregates along +// the batch dim (existing kernel in aux_heads_kernel.cu, reusable). +// - GPU-only per `feedback_cpu_is_read_only.md`. +// - ELU backward via the post-activation identity: f'(x) = 1 if x > 0, +// else 1 + f(x). The forward saved `h_post`, and `f'(x) = (h_post > 0) +// ? 1 : 1 + h_post` recovers the derivative without re-evaluating x_pre. +// +// Block: AUX_BLOCK threads. Grid: (B, 1, 1). +// Shared memory: (2 * H + K) floats — d_h_pre cache + h_post cache + +// d_logits cache for Step 1-2's reads. + +#include +#include + +#ifndef AUX_HIDDEN_DIM +#define AUX_HIDDEN_DIM 128 +#endif + +#ifndef AUX_BLOCK +#define AUX_BLOCK 256 +#endif + +/* ELU backward derivative from POST-activation y = ELU(x). + * x > 0 ⇒ y > 0 ⇒ f'(x) = 1 + * x ≤ 0 ⇒ y ≤ 0 ⇒ f'(x) = exp(x) = 1 + y + * Distinguishing on `y > 0` recovers the partition without re-evaluating + * x_pre — matches the K=2 sibling's helper. */ +__device__ __forceinline__ float aux_to_elu_bwd_from_post(float y) { + return (y > 0.0f) ? 1.0f : (1.0f + y); +} + +/* ===================================================================== + * aux_trade_outcome_backward — backward for the K=3 trade-outcome head. + * + * Step 0 (thread 0): compute d_logits[b, :] from softmax + label. + * if label ∈ {-1, <0, ≥K}: zero the K-vector (masked row → no grad) + * else: d_logits[kc] = inv_B × (softmax[kc] − one_hot(label, kc)) + * + * Step 1 (stride-H): cache h_post and compute d_h_pre per hidden unit. + * d_h_post[k] = sum_kc d_logits[kc] × w2[kc, k] + * d_h_pre[k] = d_h_post[k] × elu_bwd(h_post[k]) + * + * Step 2 (parallel writes): per-sample param-grad partials. + * db2_partial[b, kc] = d_logits[kc] + * dW2_partial[b, kc, k] = d_logits[kc] × h_post[k] + * db1_partial[b, k] = d_h_pre[k] + * dW1_partial[b, k, j] = d_h_pre[k] × h_s2_aux[b, j] + * + * Step 3 (stride-SH2): dh_s2_aux[b, j] = sum_k d_h_pre[k] × w1[k, j]. + * ===================================================================== */ +extern "C" __global__ void aux_trade_outcome_backward( + /* Forward inputs (re-read for dW). */ + const float* __restrict__ h_s2_aux, /* [B, SH2] (aux trunk output) */ + const float* __restrict__ w1, /* [H, SH2] row-major */ + const float* __restrict__ w2, /* [K=3, H] row-major */ + const float* __restrict__ hidden_post, /* [B, H] saved h_post */ + /* Loss inputs (K=3 sparse CE). */ + const float* __restrict__ softmax, /* [B, K=3] saved softmax tile */ + const int* __restrict__ labels, /* [B] i32 in {-1, 0, 1, 2} */ + const float* __restrict__ valid_count, /* [1] B_valid from loss kernel */ + int B, + int SH2, + int K, + /* Per-sample partial outputs (caller reduces along batch dim via + * the existing `aux_param_grad_reduce` kernel in aux_heads_kernel.cu + * — no separate reduce kernel needed; the reducer is K-generic). */ + float* __restrict__ dW1_partial, /* [B, H, SH2] flat */ + float* __restrict__ db1_partial, /* [B, H] */ + float* __restrict__ dW2_partial, /* [B, K=3, H] flat */ + float* __restrict__ db2_partial, /* [B, K=3] */ + /* Aux trunk-output gradient: SAXPYs into the aux trunk's + * `dh_s2_aux_accum` (NOT Q's `bw_d_h_s2`). Encoder stop-grad + * enforced structurally by aux_trunk_backward. */ + float* __restrict__ dh_s2_aux_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=3] */ + + /* Read the label + B_valid scalar once per block. Masked rows + * (label == -1) zero the K-vector → all downstream partials become + * zero (chain rule's multiplicative zero). B_valid floor at 1.0f + * mirrors the loss kernel's `fmaxf(valid, 1)` divisor — an all-skip + * batch produces zero gradients across the board (no NaN). */ + const int tgt = labels[b]; + const float B_v = fmaxf(valid_count[0], 1.0f); + const float inv_B = 1.0f / B_v; + + /* Step 0: compute d_logits[b, :] from softmax + label. Single thread + * — K=3 fanout is small enough that a tree reduce wastes more cycles + * in shfl bookkeeping than a serial pass (mirrors K=2 sibling at + * line 550 of aux_heads_kernel.cu). + * + * Sparse-label arithmetic: most bars have `tgt == -1` and zero the + * K-vector, producing a no-op partial write downstream. The valid- + * row branch uses `inv_B = 1 / B_valid` where B_valid is small (~1-5% + * of nominal batch); this amplifies per-trade-close gradient + * magnitude proportionally. */ + if (tid == 0) { + if (tgt == -1 || tgt < 0 || tgt >= K) { + for (int kc = 0; kc < K; ++kc) { + sh_dlogits[kc] = 0.0f; + } + } else { + const float* p_row = softmax + (size_t)b * K; + for (int kc = 0; kc < K; ++kc) { + const float onehot = (kc == tgt) ? 1.0f : 0.0f; + sh_dlogits[kc] = inv_B * (p_row[kc] - onehot); + } + } + } + __syncthreads(); + + /* Step 1: cache h_post and compute d_h_pre per hidden unit. + * 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_to_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_aux[b, j]. */ + { + const float* h_row = h_s2_aux + (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_aux[b, j] = sum_k sh_dh_pre[k] × w1[k, j]. + * SAXPYs into the aux trunk's `dh_s2_aux_accum` at the call site — + * the aux trunk's backward kernel reads this as its upstream + * gradient. Encoder stop-grad is enforced STRUCTURALLY by + * `aux_trunk_backward` (no `dx_in` output param), so Q's encoder + * is protected from aux contamination regardless of what gradient + * magnitude this kernel emits. + * + * Masked rows (label == -1) propagate zero d_h_pre through this + * sum → dh_s2_aux row is zero for skipped samples (the trunk + * SAXPY adds zero — no spurious gradient on bars without trade + * close). */ + { + float* dh_row = dh_s2_aux_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; + } + } +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 96f0fd7f6..2fbca3275 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17722,3 +17722,45 @@ Third foundation commit for the SP22 H6 vNext trade-outcome aux head. Lands the - Phase F: validation smoke Cargo check clean. Ready for Phase A5. + +#### Phase A5 — aux_trade_outcome backward kernel (2026-05-14) + +Fourth and final foundation commit for the SP22 H6 vNext trade-outcome aux head. Closes the forward → loss → backward chain. + +**NEW kernel `aux_trade_outcome_backward_kernel.cu`** — K=3 backward producing per-sample partials. +- Mirrors `aux_next_bar_backward` line-for-line (gradient flow is K-independent): `d_logits[b, kc] = (softmax[b, kc] − one_hot(labels[b], kc)) / B_valid` propagated through `Linear → ELU → Linear` chain via the standard softmax-CE derivative. +- Per-sample partials (caller reduces along batch dim via the existing `aux_param_grad_reduce` kernel — reusable, K-generic): `dW1_partial [B, H, SH2]`, `db1_partial [B, H]`, `dW2_partial [B, K=3, H]`, `db2_partial [B, K=3]`, `dh_s2_aux_out [B, SH2]`. +- Mask handling: `labels[b] == -1` zeros the K-vector → all downstream partials zero (chain rule's multiplicative zero). All-skip batch produces `valid_count=0` → `inv_B = 1/fmaxf(0, 1) = 1.0` but `d_logits=0` for every row → zero gradients across the board, no NaN. +- Sparse-label gradient amplification: `B_valid` is typically ~1-5% of nominal batch (trade-close events are rare), so `inv_B = 1/B_valid` is much larger than the K=2 sibling's `inv_B = 1/(~B)`. Per-trade-close gradients have proportionally higher magnitude — correct credit assignment (rare signal speaks louder) but Phase E's Adam may need class-weighted CE or per-group LR tuning to keep update magnitude comparable to the next-bar head. +- ELU backward via post-activation identity: `f'(x) = (h_post > 0) ? 1 : 1 + h_post` — recovers the derivative without re-evaluating `x_pre`. +- SP14 Phase C.5b separation: reads `h_s2_aux` (aux trunk output), writes `dh_s2_aux_out` which SAXPYs into `dh_s2_aux_accum`. Q's encoder structurally protected (aux_trunk_backward has no `dx_in` output). +- Cubin: `aux_trade_outcome_backward_kernel.cubin` (24.8 KB). +- Phase A5 (this commit) **dead code** — no Rust launcher. Phase B lands the full launcher chain. + +**Phase A series complete** (A1+A2+A3+A4+A5): +- A1: NEW ISV slots (none needed — reuses padding slots 121-123) — already done in spec +- A2: `trade_outcome_label_kernel.cu` (label producer, K=3 with mask=-1) +- A3: `aux_trade_outcome_forward_kernel.cu` + save-for-backward buffers + experience_env_step wireup +- A4: `aux_trade_outcome_loss_reduce_kernel.cu` (sparse CE, B_valid saved) +- A5: `aux_trade_outcome_backward_kernel.cu` (per-sample partials) + +All 5 foundation kernels compile, register in build.rs, produce cubins. The contract chain composes end-to-end on the GPU side: + +``` +trade_close_per_sample + pnl_vs_target/stop_at_close ──→ trade_outcome_label_kernel ──→ labels [B] i32 +aux trunk output h_s2_aux ──→ aux_trade_outcome_forward ──→ softmax [B, 3], hidden_post [B, H], logits [B, 3] +softmax + labels ──→ aux_trade_outcome_loss_reduce ──→ loss [1], valid_count [1] +softmax + labels + valid_count + hidden_post + h_s2_aux + w1 + w2 ──→ aux_trade_outcome_backward ──→ dW1_p, db1_p, dW2_p, db2_p, dh_s2_aux_out +dW*_partial + db*_partial ──→ aux_param_grad_reduce (existing, K-generic) ──→ final dW, db +``` + +**Remaining vNext work**: +- Phase B: Rust launcher chain + 262-dim input concat (`h_s2_aux` || `plan_params`) +- Phase C: 3-slot state assembly (state[121..124] = p_Profit, p_Stop, p_Timeout) +- Phase D: 12-weight W atom-shift (4 actions × 3 outcomes) +- Phase E: dW + Adam for W[4, 3] +- Phase F: validation smoke + +Phase B is the next major commit and the largest scope shift (Rust-side struct fields for aux head W1/W2/b1/b2 weights, Adam state, hidden_post / logits / softmax / dW*_partial buffers; gpu_aux_heads.rs gets a parallel `AuxTradeOutcomeForwardOps` + `AuxTradeOutcomeBackwardOps` struct; Phase B wires the chain into `gpu_experience_collector.rs::collect_experiences_gpu` immediately after the K=2 head's launches). + +Cargo check clean. Phase A complete. Ready for Phase B.