gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(sp22-vnext): Phase B5b-2 collector trade plan forward — resolves K=3 asymmetry

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-14 12:41:59 +02:00
parent fb9b62a1f9
commit 4b40710b7c
2 changed files with 216 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

180
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -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<f32>,
/// 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<f32>,
/// 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<f32>,
/// 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::<f32>(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::<f32>(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::<f32>(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,