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>
d2331f2f62