feat(sp22-vnext): Phase B5b — full plan-conditioning integration

The K=3 head's input is now [h_s2_aux (256) || plan_params (6)] = 262-
dim, matching the spec's intended architecture. Implements the input
concat in both the trainer's replay-batch path and the collector's
rollout-step path, with appropriate handling of the backward-side
stride mismatch.

NEW kernel strided_row_saxpy_kernel.cu: row-truncating SAXPY that
accumulates first n_cols_copy columns per row of src [B, src_cols]
into dst [B, dst_cols] scaled by alpha. Handles stride mismatch
(src_cols != dst_cols). Needed because backward emits dh_s2_aux_to_buf
[B, 262] but dh_s2_aux_accum [B, 256] only consumes first 256 cols.

PLAN_PARAM_DIM = 6 constant in gpu_aux_heads.rs.

Ops struct updates:
- AuxTradeOutcomeForwardOps gains concat_kernel + launch_concat()
- AuxTradeOutcomeBackwardOps gains strided_saxpy_kernel +
  launch_strided_row_saxpy()
- Both load new cubins in new()

Weight tensor resize:
- sizes[163] = H × (SH2 + PLAN_PARAM_DIM) = 128 × 262 = 33,536 floats
- fan_dims[163] = (H, SH2 + PLAN_PARAM_DIM)

Trainer changes:
- New aux_to_input_buf [B × 262] field
- aux_dh_s2_to_buf resized to [B × 262]
- aux_partial_to_w1 resized to [B × H × 262]
- max_aux_tensor_len bumped for param_grad_final scratch

Trainer forward (aux_heads_forward):
- Concat h_s2_aux + plan_params_buf → aux_to_input_buf
- forward() with SH2_TOTAL=262

Trainer backward (aux_heads_backward):
- backward() with SH2_TOTAL=262; reads aux_to_input_buf
- saxpy_f32_kernel SAXPY for dh_s2_aux REPLACED by
  launch_strided_row_saxpy: copies only first SH2=256 cols per row;
  trailing 6 cols (plan_params gradient) are dropped — STOP-GRAD on
  trade plan head from aux loss.

Collector forward (rollout):
- New exp_aux_to_input_buf [N × 262] field
- Concat with plan_params_ptr = 0 (NULL) → zero-fill trailing 6 cols
- forward() with SH2_TOTAL=262

Train/inference asymmetry (documented):
- Trainer: real plan_params from trade plan head output
- Collector: zeros (no trade plan launch in collector)
- The head is trained on real plan-conditional outcomes but queried
  at rollout time with plan_params=0. Phase B5b-2 follow-up would
  add a trade plan launch to the collector to resolve. Deferred —
  current state is functional, head still receives plan signal
  during training.

Stop-grad on plan_params: backward writes full [B, 262] gradient, but
strided SAXPY only copies first 256 cols. Trade plan head weights
NOT trained by K=3 aux loss in this commit.

Verification:
- cargo check -p ml clean.
- cargo test -p ml --lib → 1016/0 green.

Phase D next: 12-weight W atom-shift (4 actions × 3 outcomes).

Audit: docs/dqn-wire-up-audit.md Phase B5b section.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-14 10:53:23 +02:00
parent 53462a28d9
commit d2331f2f62
6 changed files with 487 additions and 47 deletions

View File

@@ -841,6 +841,14 @@ fn main() {
// this commit — Phase B5b allocates buffers + wires launches +
// flips SH2_TOTAL=262 in forward/backward.
"aux_to_input_concat_kernel.cu",
// SP22 H6 vNext Phase B5b (2026-05-14): row-truncating SAXPY for
// backward dh_s2_aux accumulate. Adds first n_cols_copy columns
// per row of [B, src_cols] into [B, dst_cols] with row stride
// mismatch handling. Needed when B5b's plan-conditioned backward
// emits dh_s2_aux_to_buf [B, 262] but the aux trunk's accumulator
// dh_s2_aux_accum [B, 256] only consumes first 256 cols (plan_
// params gradient is stop-grad).
"strided_row_saxpy_kernel.cu",
// SP22 H6 vNext Phase C (2026-05-14): per-env 3-slot softmax
// extractor for the K=3 trade-outcome aux head. Copies
// aux_outcome_softmax [n_envs, K=3] → prev_aux_outcome_probs