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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user