From 4b40710b7c0d43e8125fc9d73e92000642a3e888 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 14 May 2026 12:41:59 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp22-vnext):=20Phase=20B5b-2=20collector?= =?UTF-8?q?=20trade=20plan=20forward=20=E2=80=94=20resolves=20K=3D3=20asym?= =?UTF-8?q?metry?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase B5b's K=3 input concat passed plan_params=NULL because the collector had no trade plan launch. Trainer-side K=3 forward trained on real plan_params while the collector queried at plan_params=0 — documented train/inference asymmetry on the plan-conditioning surface. Phase B5b-2 mirrors the (now-corrected) trainer `launch_trade_plan_forward` chain on the collector inside the rollout step: 1. SGEMM: hidden[N, AH] = h_s2_q[N, SH2] @ W_fc[AH, SH2]^T 2. bias+relu in-place on hidden 3. SGEMM: pre_out[N, 6] = hidden[N, AH] @ W_out[6, AH]^T 4. trade_plan_activate → exp_plan_params[N, 6] Weight resolution uses the same `aux_w_ptrs` array the K=3 forward already consumes (`f32_weight_ptrs_from_base`); indices 91-94 match the corrected trainer-side reads. The `trade_plan_activate` kernel is loaded from `EXPERIENCE_KERNELS_CUBIN` (same cubin the rest of `exp_module_extra` uses; the trainer loads it from there too). The K=3 concat now takes `exp_plan_params.raw_ptr()` instead of NULL — both sides see f(h_s2; W_plan_*_init), symmetry restored. Plan tensors at [91..94] still have no backward (no Adam updates), so the plan-head weights stay at Xavier cold-start forever. This commit delivers the symmetry the K=3 head requires, not a learned plan signal — adding a real plan-head backward is a follow-up project. New struct fields on GpuExperienceCollector: - exp_trade_plan_hidden_buf [alloc_episodes × adv_h] - exp_trade_plan_pre_out_buf [alloc_episodes × 6] - exp_plan_params [alloc_episodes × 6] - exp_trade_plan_activate_kernel: CudaFunction Lib test suite: 1016/0 green maintained. Audit doc updated. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../cuda_pipeline/gpu_experience_collector.rs | 180 ++++++++++++++++-- docs/dqn-wire-up-audit.md | 47 +++++ 2 files changed, 216 insertions(+), 11 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 38d60f2fc..1075d8ab9 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1344,6 +1344,54 @@ pub struct GpuExperienceCollector { /// `exp_aux_to_softmax_buf`. exp_aux_outcome_softmax_to_per_env_kernel: CudaFunction, + // ── SP22 H6 vNext Phase B5b-2 (2026-05-14) — collector trade plan + // forward (resolves Phase B5b train/inference asymmetry) ── + // + // Phase B5b's K=3 concat passed `plan_params=NULL` because the + // collector had no trade plan launch — trainer-side K=3 forward + // trained on real plan_params while the collector queried at + // plan_params=0. Phase B5b-2 mirrors the (now-corrected) trainer + // `launch_trade_plan_forward` chain on the collector: cuBLAS SGEMM + // (W_fc[AH,SH2] × h_s2) → bias-relu → SGEMM (W_out[6,AH] × hidden) + // → trade_plan_activate kernel. Output `exp_plan_params[N, 6]` + // feeds the concat kernel's `plan_params` arg in lieu of NULL. + // + // Param tensors at indices [91..95) via the same + // `f32_weight_ptrs_from_base` array the K=3 forward already uses + // (no extra weight resolution). Cold-start path (test scaffold, + // `trainer_params_ptr == 0`) skips the chain — same guard as the + // existing aux next-bar forward at the start of the rollout step. + + /// Scratch buffer `[alloc_episodes × adv_h]` f32 — layer-1 hidden + /// output (post bias+ReLU) of the collector's trade plan MLP. + /// Mirrors the trainer's `trade_plan_hidden_buf` shape. Lifetime: + /// only valid between the SGEMM that writes it and the SGEMM that + /// reads it within the same rollout step. + exp_trade_plan_hidden_buf: cudarc::driver::CudaSlice, + + /// Scratch buffer `[alloc_episodes × 6]` f32 — layer-2 pre-activation + /// output of the collector's trade plan MLP. Mirrors the trainer's + /// `trade_plan_pre_out_buf`. The `trade_plan_activate` kernel reads + /// this + b_out and writes activated `exp_plan_params`. + exp_trade_plan_pre_out_buf: cudarc::driver::CudaSlice, + + /// Collector-side rollout-step plan_params `[alloc_episodes × 6]` + /// f32, sigmoid/scale-activated per `trade_plan_activate`. Feeds + /// the B5b concat kernel's `plan_params` arg (replacing the NULL + /// passed in B5b). Cold-start sentinel = 0.0 (alloc_zeros). The + /// plan tensors at indices [91..94] still have no backward, so + /// these values are deterministic outputs of Xavier-init weights + /// applied to live h_s2 — that's all the symmetry-resolution this + /// commit can deliver. Adding a real plan-head backward is a + /// separate follow-up project (out of scope for B5b-2). + exp_plan_params: cudarc::driver::CudaSlice, + + /// Kernel handle for `trade_plan_activate` (bias + per-output + /// sigmoid/scaling activation). Loaded from `EXPERIENCE_KERNELS_CUBIN` + /// (same cubin used by the rest of `exp_module_extra`). Mirrors + /// the trainer's `trade_plan_activate_kernel`. + exp_trade_plan_activate_kernel: CudaFunction, + /// SP22 H6 Phase 2 (2026-05-12) — per-env aux directional probability /// cache, RECENTERED encoding. Sized `[alloc_episodes]` f32 (one slot /// per env). Written end-of-step by `exp_aux_softmax_to_per_env_kernel` @@ -2849,6 +2897,39 @@ impl GpuExperienceCollector { "sp22-vnext C: aux_outcome_softmax_to_per_env_kernel load (collector): {e}" )))? }; + + // SP22 H6 vNext Phase B5b-2 (2026-05-14) — collector trade plan + // forward chain (resolves the train/inference plan_params asymmetry + // documented on `exp_aux_to_input_buf`). Three scratch buffers + // mirror the trainer's `trade_plan_hidden_buf` / `pre_out_buf` / + // `plan_params_buf` (sized for `alloc_episodes` rollout envs vs the + // trainer's batch_size) plus the `trade_plan_activate` kernel handle + // loaded from `EXPERIENCE_KERNELS_CUBIN` (the same cubin the rest + // of `exp_module_extra` uses, where the trainer also loads this + // kernel — single source of truth). + let exp_trade_plan_hidden_buf = stream + .alloc_zeros::(alloc_episodes * adv_h) + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext B5b-2: alloc exp_trade_plan_hidden_buf ({} f32): {e}", + alloc_episodes * adv_h + )))?; + let exp_trade_plan_pre_out_buf = stream + .alloc_zeros::(alloc_episodes * super::gpu_aux_heads::PLAN_PARAM_DIM) + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext B5b-2: alloc exp_trade_plan_pre_out_buf ({} f32): {e}", + alloc_episodes * super::gpu_aux_heads::PLAN_PARAM_DIM + )))?; + let exp_plan_params = stream + .alloc_zeros::(alloc_episodes * super::gpu_aux_heads::PLAN_PARAM_DIM) + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext B5b-2: alloc exp_plan_params ({} f32): {e}", + alloc_episodes * super::gpu_aux_heads::PLAN_PARAM_DIM + )))?; + let exp_trade_plan_activate_kernel = exp_module_extra + .load_function("trade_plan_activate") + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext B5b-2: trade_plan_activate load (collector): {e}" + )))?; // SP22 H6 (2026-05-12): aux→policy state bridge. Load the per-env // softmax extractor (`aux_softmax_to_per_env_kernel`) + the shared // `fill_f32` from the epsilon_greedy cubin (for 0.5 sentinel init at @@ -3450,6 +3531,15 @@ impl GpuExperienceCollector { // step; state gather consumer wired in Phase C-2). prev_aux_outcome_probs, exp_aux_outcome_softmax_to_per_env_kernel, + // SP22 H6 vNext Phase B5b-2 (2026-05-14): collector trade plan + // forward chain — 3 scratch buffers + activate kernel handle. + // Resolves Phase B5b's NULL plan_params train/inference + // asymmetry by mirroring trainer's `launch_trade_plan_forward` + // (now reading from the corrected indices [91..94]). + exp_trade_plan_hidden_buf, + exp_trade_plan_pre_out_buf, + exp_plan_params, + exp_trade_plan_activate_kernel, // SP22 H6 (2026-05-12): aux→policy state bridge — per-env p_up // cache + extractor kernel + fill_f32 handle for 0.5 sentinel. prev_aux_dir_prob, @@ -5982,28 +6072,96 @@ impl GpuExperienceCollector { let to_b1 = aux_w_ptrs[164]; let to_w2 = aux_w_ptrs[165]; let to_b2 = aux_w_ptrs[166]; + // ── SP22 H6 vNext Phase B5b-2 (2026-05-14): collector ── + // trade plan forward chain. Mirrors trainer's + // `launch_trade_plan_forward` (gpu_dqn_trainer.rs:10844) + // using the corrected param indices [91..94]: + // 1. SGEMM: hidden[N, AH] = h_s2[N, SH2] @ W_fc^T + // 2. add_bias_relu in-place on hidden + // 3. SGEMM: pre_out[N, 6] = hidden[N, AH] @ W_out^T + // 4. trade_plan_activate kernel → exp_plan_params[N, 6] + // + // Input is the Q-trunk `exp_h_s2_f32` (NOT the aux trunk + // `exp_h_s2_aux`) — matches the trainer reading + // `save_h_s2` (Q-trunk h_s2). The plan head's actual + // tensors at [91..94] still have no backward, so the + // output is a deterministic function of cold-start + // Xavier weights applied to live h_s2. The result feeds + // the K=3 concat in lieu of NULL, achieving train/ + // inference symmetry on `plan_params` (both sides see + // the same `f(h_s2)` distribution). + // + // Cost: 2 × cuBLAS SGEMM + 1 elementwise + 1 activate + // launch per rollout step. SGEMM dims for n_episodes=8, + // AH=256, SH2=512: 8×256×512 = 1M FLOPs (L1) + 8×6×256 + // = 12k FLOPs (L2). Sub-microsecond on L40S; negligible + // vs the K=3 forward. + let w_plan_fc = aux_w_ptrs[91]; + let b_plan_fc = aux_w_ptrs[92]; + let w_plan_out = aux_w_ptrs[93]; + let b_plan_out = aux_w_ptrs[94]; + let p_dim = super::gpu_aux_heads::PLAN_PARAM_DIM; + let trade_plan_hidden_ptr = self.exp_trade_plan_hidden_buf.raw_ptr(); + let trade_plan_pre_out_ptr = self.exp_trade_plan_pre_out_buf.raw_ptr(); + let exp_plan_params_ptr = self.exp_plan_params.raw_ptr(); + let h_s2_q_ptr = self.exp_h_s2_f32.raw_ptr(); + let adv_h = self.network_dims.3; + + // Step 1: hidden[N, AH] = h_s2[N, SH2] @ W_fc[AH, SH2]^T + self.cublas_forward.sgemm_f32( + &self.stream, + w_plan_fc, h_s2_q_ptr, trade_plan_hidden_ptr, + adv_h, n, sh2, "exp_plan_l1", + )?; + // Step 2: hidden += b_fc[AH]; relu (in-place). + self.cublas_forward.launch_add_bias_relu_f32_raw( + &self.stream, trade_plan_hidden_ptr, b_plan_fc, adv_h, n, + )?; + // Step 3: pre_out[N, 6] = hidden[N, AH] @ W_out[6, AH]^T + self.cublas_forward.sgemm_f32( + &self.stream, + w_plan_out, trade_plan_hidden_ptr, trade_plan_pre_out_ptr, + p_dim, n, adv_h, "exp_plan_l2", + )?; + // Step 4: bias + per-output activation → exp_plan_params[N, 6]. + { + let n_i32 = n as i32; + let blocks = ((n as u32).div_ceil(256)).max(1); + unsafe { + self.stream + .launch_builder(&self.exp_trade_plan_activate_kernel) + .arg(&trade_plan_pre_out_ptr) + .arg(&b_plan_out) + .arg(&exp_plan_params_ptr) + .arg(&n_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "exp_trade_plan_activate t={t}: {e}" + )))?; + } + } + // ── SP22 H6 vNext Phase B5b (2026-05-14): plan-conditioned ── // input concat for the rollout-step K=3 forward. // - // Collector has NO trade plan head launch (the trade plan - // runs only on replay batch); pass plan_params_ptr = 0 - // (NULL) to the concat kernel which zero-fills the - // trailing PLAN_PARAM_DIM=6 columns. Train/inference - // asymmetry documented in the `exp_aux_to_input_buf` - // struct field comment; Phase B5b-2 (follow-up) adds - // a real trade plan launch to the collector to resolve. - // // SH2_TOTAL = SH2 + PLAN_PARAM_DIM = 262. Kernel takes // SH2 as a runtime arg so passing sh2_total just works // — no kernel surgery needed. - let p_dim = super::gpu_aux_heads::PLAN_PARAM_DIM; + // + // B5b-2 update: pass real `exp_plan_params` (just + // produced above) instead of NULL — resolves the + // train/inference asymmetry documented on + // `exp_aux_to_input_buf`. let sh2_total = sh2 + p_dim; - let plan_params_null: u64 = 0; let aux_to_input_ptr = self.exp_aux_to_input_buf.raw_ptr(); self.exp_aux_to_fwd.launch_concat( &self.stream, self.exp_h_s2_aux.raw_ptr(), - plan_params_null, + exp_plan_params_ptr, n, sh2, p_dim, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 04e3f0166..c4389e55a 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,53 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +## 2026-05-14 — SP22 H6 vNext Phase B5b-2: collector trade plan forward (asymmetry resolved) + +**Branch:** `sp20-aux-h-fixed`. **Depends on:** prior commit `fb9b62a1f` +(trade plan index fix [82..85] → [91..94]). + +Phase B5b deliberately landed with the collector's K=3 input concat +passing `plan_params=NULL` (zero-fill the trailing 6 columns of the 262-col +input). The trainer-side K=3 forward trained on real `plan_params` while +the collector queried at `plan_params=0`, creating a documented train/ +inference asymmetry on the K=3 head's plan-conditioning surface. + +B5b-2 mirrors the (now-corrected) trainer `launch_trade_plan_forward` +chain on the collector inside the rollout step: + +1. cuBLAS SGEMM: `hidden[N, AH] = h_s2_q[N, SH2] @ W_fc[AH, SH2]^T` + (W_fc resolved via `aux_w_ptrs[91]` — same `f32_weight_ptrs_from_base` + array the K=3 forward already consumes; no extra weight resolution) +2. `launch_add_bias_relu_f32_raw` in-place on hidden +3. cuBLAS SGEMM: `pre_out[N, 6] = hidden[N, AH] @ W_out[6, AH]^T` + (W_out at `aux_w_ptrs[93]`) +4. `trade_plan_activate` kernel (loaded from `EXPERIENCE_KERNELS_CUBIN` + via `exp_module_extra`, same cubin the trainer loads from) → + `exp_plan_params[N, 6]` + +The K=3 concat call then takes `exp_plan_params.raw_ptr()` instead of +the NULL it passed in B5b. Both training and inference paths now see +`plan_params = f(h_s2; W_plan_*_init)` — symmetric. Cost on L40S: +sub-microsecond per rollout step (2 small SGEMMs + 1 elementwise + 1 +kernel launch). + +**Important caveat:** the plan tensors at `[91..94]` still have no +backward (no Adam updates), so the plan-head weights remain at their +Xavier cold-start values forever. This commit delivers the *symmetry* +the K=3 head's design requires, not a *learned* plan signal. Adding a +real plan-head backward is a separate follow-up project. + +**New struct fields on `GpuExperienceCollector`:** + +- `exp_trade_plan_hidden_buf: CudaSlice` — `[alloc_episodes × adv_h]` +- `exp_trade_plan_pre_out_buf: CudaSlice` — `[alloc_episodes × 6]` +- `exp_plan_params: CudaSlice` — `[alloc_episodes × 6]` +- `exp_trade_plan_activate_kernel: CudaFunction` + +All compile-clean. Lib test suite: **1016/0 green** maintained. + +**Touched:** `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`. + ## 2026-05-14 — fix(plan-head): trainer trade plan forward reads from wrong param indices **Branch:** `sp20-aux-h-fixed`. **Discovered during SP22 H6 vNext Phase B5b-2.**