From bbd52c3aa7ee40bb9b6c7d65f38e0b09d0452452 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 14 May 2026 13:12:34 +0200 Subject: [PATCH] feat(sp22-vnext): FoldReset registry entries for B5b/C collector buffers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two collector-side device buffers used by the K=3 trade-outcome head were missing FoldReset registry coverage: - prev_aux_outcome_probs [alloc_episodes × 3] TRUE stale-read risk. Producer writes end-of-step; consumer (experience_state_gather) reads start-of-next-step into state[121..124). Without FoldReset the new fold's step-0 state gather would inject the previous fold's last-step softmax probs into the first batch's state slots. - exp_aux_to_input_buf [alloc_episodes × 262] Cleanliness-only. Concat kernel overwrites all 262 columns every step before the K=3 forward reads them, so no steady-state stale- read risk. Registered for parity with the rest of the K=3 pipeline + to satisfy feedback_registry_entries_need_dispatch_ arms (the pin test asserts every registry entry has a matching dispatch arm in reset_named_state). Both fields promoted to pub(crate) on GpuExperienceCollector so reset_named_state can reach them. Matching dispatch arms added with the standard memset_zeros pattern (is_win_per_env / hold_baseline_ buffer style). Tests: All 10 state_reset_registry tests pass, including the critical every_fold_and_soft_reset_entry_has_dispatch_arm pin test that walks the dispatch body and validates parity with registry entries. Full lib suite 1015/1 (the failing test is the pre-existing test_dqn_checkpoint_round_trip NoisyLinear flake — pred1/pred2 sign mismatch surfacing ~30-50% of full-suite runs, documented in project_sp22_h6_vnext_resume memory as unrelated to this work). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../cuda_pipeline/gpu_experience_collector.rs | 12 +++++- .../src/trainers/dqn/state_reset_registry.rs | 19 ++++++++++ .../src/trainers/dqn/trainer/training_loop.rs | 27 ++++++++++++++ docs/dqn-wire-up-audit.md | 37 +++++++++++++++++++ 4 files changed, 93 insertions(+), 2 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 1075d8ab9..efa8f9f3a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1316,7 +1316,11 @@ pub struct GpuExperienceCollector { /// to next-step state[121..124)). Phase B5b-2 (follow-up) adds a /// trade plan head launch to the collector to populate real /// plan_params at rollout time, resolving the asymmetry. - exp_aux_to_input_buf: cudarc::driver::CudaSlice, + /// + /// `pub(crate)` so the trainer's `reset_named_state` can reach it + /// at fold-boundary reset (per the matching `RegistryEntry` in + /// `state_reset_registry.rs::exp_aux_to_input_buf`). + pub(crate) exp_aux_to_input_buf: cudarc::driver::CudaSlice, /// SP22 H6 vNext Phase C (2026-05-14) — per-env K=3 trade-outcome /// softmax cache. Sized `[alloc_episodes × AUX_OUTCOME_K=3]` f32, @@ -1336,7 +1340,11 @@ pub struct GpuExperienceCollector { /// Cold-start sentinel = 0.0 across all 3 slots. The state gather /// reads "no Profit/Stop/Timeout signal" on step 0, then this /// kernel writes real K=3 softmax probs for step 1+ state assembly. - prev_aux_outcome_probs: cudarc::driver::CudaSlice, + /// + /// `pub(crate)` so the trainer's `reset_named_state` can reach it + /// at fold-boundary reset (per the matching `RegistryEntry` in + /// `state_reset_registry.rs::prev_aux_outcome_probs`). + pub(crate) prev_aux_outcome_probs: cudarc::driver::CudaSlice, /// SP22 H6 vNext Phase C (2026-05-14) — kernel handle for /// `aux_outcome_softmax_to_per_env_kernel`. Loaded from the Phase /// C cubin. Launched on the collector's stream right after the diff --git a/crates/ml/src/trainers/dqn/state_reset_registry.rs b/crates/ml/src/trainers/dqn/state_reset_registry.rs index e3bbf3f2e..42663740e 100644 --- a/crates/ml/src/trainers/dqn/state_reset_registry.rs +++ b/crates/ml/src/trainers/dqn/state_reset_registry.rs @@ -2188,6 +2188,25 @@ impl StateResetRegistry { category: ResetCategory::FoldReset, description: "GpuExperienceCollector.pnl_vs_stop_at_close_per_env [alloc_episodes] f32 device-resident scratch — SP22 H6 vNext Phase A3 (2026-05-13) per-env trade-close P&L ratio vs `plan_stop × equity`. Companion to `pnl_vs_target_at_close_per_env` — same segment_complete producer site, same FoldReset semantics, same NULL-tolerant kernel arg pattern. Value semantics: `plan_isv[PLAN_ISV_PNL_VS_STOP]` (the local plan_isv value already computed at line ~877 from `-f_unrealized_pnl / (plan_stop × f_equity + 1e-6f)` clamped [-2, +2]). Consumer (eventual): `trade_outcome_label_kernel.cu`'s `pnl_vs_stop_at_close` arg — `≥ 1.0` predicate classifies as `Stop` (label 1). When neither pvt nor pvs ≥ 1.0 the kernel emits `Timeout` (label 2). FoldReset sentinel 0.0; reset path: `cudarc::driver::CudaStream::memset_zeros`.", }, + // SP22 H6 vNext Phase C cache + Phase B5b concat input — both + // collector-side device buffers used by the K=3 trade-outcome + // head's rollout-step launch. Stale-read risk differs (true + // for `prev_aux_outcome_probs` which feeds state[121..124) at + // step start; cleanliness-only for `exp_aux_to_input_buf` + // which is written-before-read every step). Both registered + // for FoldReset to satisfy `feedback_registry_entries_need_ + // dispatch_arms` discipline + match the per-fold reset + // ergonomics of the rest of the K=3 head pipeline. + RegistryEntry { + name: "prev_aux_outcome_probs", + category: ResetCategory::FoldReset, + description: "GpuExperienceCollector.prev_aux_outcome_probs [alloc_episodes × AUX_OUTCOME_K=3] f32 device-resident cache — SP22 H6 vNext Phase C (2026-05-14) per-env K=3 trade-outcome softmax cache. Producer: `aux_outcome_softmax_to_per_env_kernel` (end-of-step, reads `exp_aux_to_softmax_buf`). Consumer: `experience_state_gather` start-of-next-step, populates `state[AUX_OUTCOME_PROFIT_INDEX..AUX_OUTCOME_TIMEOUT_INDEX+1) = [121, 122, 123]`. FoldReset sentinel 0.0 across all `alloc_episodes × 3` slots — TRUE stale-read risk if previous fold's softmax probs leak into the new fold's step-0 state gather (would inject prior-fold's K=3 distribution into the new fold's first state slots, polluting the first batch's training signal). Reset path: `cudarc::driver::CudaStream::memset_zeros` on the `CudaSlice` (device-resident, no host buffer to fill). Per `pearl_first_observation_bootstrap`: the 0.0 cold-start reads as 'no Profit/Stop/Timeout signal' on step 0; the per-env producer then writes real K=3 softmax probs starting step 1.", + }, + RegistryEntry { + name: "exp_aux_to_input_buf", + category: ResetCategory::FoldReset, + description: "GpuExperienceCollector.exp_aux_to_input_buf [alloc_episodes × (SH2 + PLAN_PARAM_DIM) = N × 262] f32 device-resident input buffer — SP22 H6 vNext Phase B5b (2026-05-14) plan-conditioned input for the rollout-step K=3 forward. Producer: `aux_to_input_concat_kernel` (every rollout step, writes ALL 262 columns from `exp_h_s2_aux[N, 256]` + `exp_plan_params[N, 6]`). Consumer: `aux_trade_outcome_forward` (same step, reads as `[N, 262]` GEMM input). Stale-read risk: NONE in steady state (every step's concat overwrites the full buffer before the forward reads it). FoldReset is included for registry-discipline parity with the rest of the K=3 pipeline (per `feedback_registry_entries_need_dispatch_arms`) and to guarantee a clean baseline at fold boundary if a future change ever turns this from write-every-step into a conditional write. Sentinel 0.0; reset path: `cudarc::driver::CudaStream::memset_zeros`.", + }, ]; Self { entries } } diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index cc9c2e57e..2d9c33429 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -10125,6 +10125,33 @@ impl DQNTrainer { )))?; } } + // SP22 H6 vNext Phase C cache + Phase B5b concat input — both + // collector-side device buffers feeding the K=3 trade-outcome + // head. `prev_aux_outcome_probs` has a TRUE stale-read risk + // (state gather reads it on step 0 of new fold) — must be + // zeroed. `exp_aux_to_input_buf` is written-before-read every + // step (concat overwrites all 262 cols), so its reset is + // cleanliness-only — but the registry entry requires a + // matching dispatch arm per `feedback_registry_entries_need_ + // dispatch_arms`. Same `memset_zeros` pattern as the rest. + "prev_aux_outcome_probs" => { + if let Some(ref mut collector) = self.gpu_experience_collector { + let stream = collector.stream().clone(); + stream.memset_zeros(&mut collector.prev_aux_outcome_probs) + .map_err(|e| crate::MLError::ModelError(format!( + "prev_aux_outcome_probs memset_zeros: {e}" + )))?; + } + } + "exp_aux_to_input_buf" => { + if let Some(ref mut collector) = self.gpu_experience_collector { + let stream = collector.stream().clone(); + stream.memset_zeros(&mut collector.exp_aux_to_input_buf) + .map_err(|e| crate::MLError::ModelError(format!( + "exp_aux_to_input_buf memset_zeros: {e}" + )))?; + } + } // SP20 Phase 3 Task 3.2 (2026-05-10): per-env Hold // opportunity-cost circular buffer (Component 2). Same // device-resident `CudaSlice` reset pattern as diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index c4389e55a..d4b9d6958 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,43 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +## 2026-05-14 — SP22 H6 vNext: FoldReset registry entries for B5b/C collector buffers + +**Branch:** `sp20-aux-h-fixed`. **Builds on:** `4b40710b7` (B5b-2 collector +trade plan launch). + +Two collector-side device buffers used by the K=3 trade-outcome head +were missing FoldReset registry coverage: + +- `prev_aux_outcome_probs` `[alloc_episodes × 3]` — **TRUE stale-read + risk.** Producer (`aux_outcome_softmax_to_per_env_kernel`) writes + end-of-step; consumer (`experience_state_gather`) reads start-of- + next-step into `state[121..124)`. Without FoldReset the new fold's + step-0 state gather would inject the previous fold's last-step + softmax probs into the first batch's state slots. +- `exp_aux_to_input_buf` `[alloc_episodes × 262]` — cleanliness-only. + Producer (`aux_to_input_concat_kernel`) overwrites all 262 columns + every step before the K=3 forward reads them, so no stale-read risk + in steady state. Registered for parity with the rest of the K=3 + pipeline + to satisfy `feedback_registry_entries_need_dispatch_arms`. + +Both fields promoted to `pub(crate)` on `GpuExperienceCollector` so +`reset_named_state` (in `training_loop.rs`) can reach them. Matching +match arms added with the standard `memset_zeros` dispatch pattern +(`is_win_per_env` / `hold_baseline_buffer` style). + +**Tests:** All 10 `state_reset_registry` tests pass, including the +critical `every_fold_and_soft_reset_entry_has_dispatch_arm` pin test +that walks the dispatch body and validates parity with registry +entries. Lib suite 1015/1 (the failing test is the pre-existing +`test_dqn_checkpoint_round_trip` NoisyLinear flake — surfaces ~30-50% +of full-suite runs, unrelated to this commit). + +**Touched:** `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` +(field visibility + doc-comment), `crates/ml/src/trainers/dqn/state_ +reset_registry.rs` (2 new entries), `crates/ml/src/trainers/dqn/ +trainer/training_loop.rs` (2 new dispatch arms). + ## 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`