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

feat(sp15-p3.5.4.c): production caller + OR-gate consumer for plasticity injection — closes 3.5.4 end-to-end

Browse Source 
Wave 4.2 (ef08611d3) landed cuRAND Kaiming-He weight reset as orphan-
with-tests-only. Phase 3.5.4.c creates the production caller and the
action-selection consumer atomically per spec lines 4346-4448.

Production caller in gpu_experience_collector.rs step 4-pre (BEFORE
experience_action_select):
  - fan_in = cfg.adv_h = 128 (CORRECTED from Wave 4.2's audit-doc
    error claiming 6528 = adv_h × num_atoms; that's total weight
    count of one branch's projection, not per-unit input dim. Per
    feedback_trust_code_not_docs the audit-doc error is also fixed
    inline in this commit.)
  - n_weights = branch_0_size × num_atoms × adv_h = 4 × 51 × 128 =
    26112 (default config). Resets last 10% = 2611 directional
    advantage-head weights via cuRAND curand_normal × sqrt(2/128) ≈
    0.125 stddev (was wrongly documented as 0.0175).
  - Branch: w_b0out (directional) only — plasticity is about
    escaping stuck-in-flat regimes; targeted at directional choice.
  - Seed: mix_seed(Phase-3.5.4.c-unique base) ^ t per-step
    deterministic source; same (FOXHUNT_SEED, t) always yields the
    same Kaiming-He samples per kernel thread.
  - Trainer exposes target via new pub fn sp15_w_b0out_target()
    returning (dev_ptr, n_weights, fan_in); collector consumes via
    new set_sp15_plasticity_target(); training_loop wires the two.

OR-gate consumer in experience_action_select:
  - New kernel arg float plasticity_m_warm threaded into the
    cooldown-mask code path.
  - Wave 1.B's cooldown_active = (cooldown_remaining > 0) extended
    to cooldown_active = (cooldown_remaining > 0) ||
                         (plasticity_warm_remaining > 0).
  - Flat-on-fire-bar detection — option (a), warm ≥ m_warm − 1.5f
    per the trigger-then-decrement convention. The kernel sees
    warm = m_warm − 1 on the fire bar; subsequent warm-up bars
    observe warm ≤ m_warm − 2. New if (plasticity_fire_bar) branch
    BEFORE the existing else if (cooldown_active) branch — the
    fire-bar gets DIR_FLAT, subsequent warm bars get DIR_HOLD via
    the OR-gate cooldown.

Two-step recovery semantics:
  - Bar T (fire): plasticity launches → ISV[436]=1, ISV[438]=
    m_warm then decrements to m_warm−1. action_select detects
    fire-bar → dir_idx = DIR_FLAT.
  - Bars T+1..T+M_warm−1 (warm): action_select OR-gate forces
    dir_idx = DIR_HOLD.
  - Bar T+M_warm: warm transitions to 0, OR-gate inactive, normal
    Thompson/argmax resumes.

When the production caller is unwired (test scaffold path):
plasticity_m_warm passes 0.0f, the kernel's fire-bar predicate is
dead (warm == 0 too), and the OR-gate degenerates to cooldown-only
— bit-identical to the pre-3.5.4.c behaviour. Existing 3 SP15
3.5.b/3.5.3.b action_select oracle tests pass unchanged at
m_warm = 0.0f.

New behavioral oracle test plasticity_fires_force_flat_then_cooldown_
holds verifies the full sequence on RTX 3050 Ti with M_warm = 5
(shrunk from production's 200 for test runtime): bar 0 → DIR_FLAT,
bars 1-3 → DIR_HOLD, bar 4 (warm boundary) → DIR_LONG. ISV side-
effects (fired flips 0→1 then debounces, warm decrements with
underflow guard) verified bar-by-bar.

Eval/backtest path passes m_warm = 0.0f because plasticity is a
training-only mechanism (during deterministic eval the weights must
remain frozen at their checkpoint values).

Atomic per feedback_no_partial_refactor: kernel + caller + consumer +
trainer plumbing + 4 launcher-call-site updates (1 production
collector + 1 eval-path + 2 test scaffolds) + new behavioral test +
audit-doc fan_in inline correction + new audit-doc entry all in this
commit. No fallback. SP15 Phase 3.5 recovery-dynamics chain is now
end-to-end production-wired (3.5.2 + 3.5.3 + 3.5.4 + 3.5.4.c +
3.5.5 + 3.5.5.b all firing).

Verified: cargo check -p ml --features cuda clean; cargo check -p ml
--features cuda --tests clean; CUDA_COMPUTE_CAP=86 cargo test -p ml
--test sp15_phase1_oracle_tests --features cuda -- --ignored
--nocapture 34 of 34 SP15 oracle tests green (5 plasticity + 3
action_select + 3 cooldown + 23 others); cargo test -p ml --features
cuda --lib HOLDS the Wave 4.3 baseline (946 pass / 13 fail) on RTX
3050 Ti — no new regressions.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 08:58:17 +02:00
parent 19f6cce510
commit 483cef454c
9 changed files with 664 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -1061,7 +1061,34 @@ extern "C" __global__ void experience_action_select(
const float* __restrict__ b_logits_dir, /* [N, b0_size, n_atoms] direction-branch RAW C51 logits */
const float* __restrict__ per_sample_support, /* [N, b0_size, 3] stride-12 per-direction [v_min, v_max, delta_z] */
const float* __restrict__ atom_positions, /* [b0_size, n_atoms] adaptive non-linear positions; NULL = linear */
int n_atoms /* C51 atom count */
int n_atoms, /* C51 atom count */
/* SP15 Phase 3.5.4.c (2026-05-07) — plasticity-injection consumer
* args. The plasticity-injection kernel runs BEFORE this kernel on
* the same stream and writes ISV[PLASTICITY_FIRED_THIS_FOLD=436] +
* ISV[PLASTICITY_WARM_BARS_REMAINING=438]. This kernel reads those
* slots to:
* (a) detect the FIRE BAR — the trigger writes warm = m_warm
* then immediately decrements to (m_warm − 1). On the fire
* bar this kernel observes warm ≥ (m_warm − 1.5f). On that
* bar dir_idx is forced to DIR_FLAT (close any open
* position — step 1 of the spec §9.2 (3.5.4) two-step
* recovery).
* (b) extend the SP15 Phase 3.5.3.b cooldown OR-gate from
* (cooldown_remaining > 0) to
* (max(cooldown_remaining, plasticity_warm_remaining) > 0);
* on bars T+1..T+M (post-fire warm-up) plasticity_warm_
* remaining > 0 and dir_idx is forced to DIR_HOLD (step 2 of
* the two-step). The pre-existing 3.5.3.b code path that
* reads ISV[435] is preserved bit-identically when
* plasticity is unwired (m_warm == 0): warm stays 0, so the
* OR-gate degenerates to (cooldown_remaining > 0).
*
* `m_warm` is the warm-up bar count the plasticity-injection
* launcher passes to its trigger (default 200 per spec). Pass
* 0.0f to disable the fire-bar detection (test path / production
* code path that hasn't wired the plasticity launch yet — when
* m_warm == 0, warm == 0 too, so neither branch fires). */
float plasticity_m_warm
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
@@ -1296,6 +1323,15 @@ extern "C" __global__ void experience_action_select(
float q_eff_dir[4];
float hold_floor = 0.0f;
int cooldown_active = 0;
/* SP15 Phase 3.5.4.c (2026-05-07) — plasticity-injection two-step
* recovery state. The trigger kernel writes ISV[438] = m_warm on
* the fire bar (then immediately decrements by 1 in the same
* launch), and decrements per bar thereafter. This kernel reads
* the warm counter and m_warm to discriminate the fire bar from
* subsequent warm-up bars; the fire-bar gets DIR_FLAT (step 1),
* subsequent warm-up bars get DIR_HOLD via the OR-gate (step 2). */
int plasticity_fire_bar = 0;
int plasticity_warm_active = 0;
if (isv_signals_ptr != NULL && b0_size <= 4) {
/* Direction-policy entropy. Stable softmax over e_dir[0..b0_size). */
float ed_max = e_dir[0];
@@ -1325,7 +1361,40 @@ extern "C" __global__ void experience_action_select(
/* Cooldown counter — ISV[COOLDOWN_BARS_REMAINING=435]. f32
* representation of an integer counter; > 0 means cooldown
* still active. */
cooldown_active = (isv_signals_ptr[/*COOLDOWN_BARS_REMAINING_INDEX*/ 435] > 0.0f);
const float cooldown_remaining =
isv_signals_ptr[/*COOLDOWN_BARS_REMAINING_INDEX*/ 435];
/* SP15 Phase 3.5.4.c — plasticity warm counter +
* fire-bar detection. The trigger kernel sets warm = m_warm
* on fire and decrements once in the same launch, so the
* fire bar observes warm == m_warm − 1 (use a 0.5f tolerance
* for the f32 equality compare per the cooldown_kernel
* trigger-then-decrement convention). When plasticity is
* unwired (m_warm == 0), warm == 0 too — both predicates are
* false and behaviour reduces to the pre-3.5.4.c (cooldown-
* only) path bit-identically.
*
* Per pearl_no_host_branches_in_captured_graph: the kernel-
* side detection avoids any host read of ISV between the
* plasticity launch and this action_select launch (both run
* outside the exp_fwd_graph capture region — exp_fwd_graph
* ends at the cuBLAS forward in gpu_experience_collector.rs
* before step 4 — but we keep the detection on-GPU
* regardless to preserve graph-friendly ordering invariants
* for any future graph-capture extension). */
const float plasticity_warm_remaining =
isv_signals_ptr[/*PLASTICITY_WARM_BARS_REMAINING_INDEX*/ 438];
plasticity_warm_active = (plasticity_warm_remaining > 0.0f);
plasticity_fire_bar = (plasticity_m_warm > 0.0f) &&
(plasticity_warm_remaining
>= plasticity_m_warm - 1.5f);
/* OR-gate per spec §9.2 (3.5.4) post-amendment-2: the cooldown
* mask fires when EITHER counter is active. Pre-3.5.4.c
* behaviour preserved when plasticity is unwired (warm == 0
* → only the cooldown branch can flip the gate). */
cooldown_active = (cooldown_remaining > 0.0f)
|| plasticity_warm_active;
}
for (int d = 0; d < b0_size; d++) {
q_eff_dir[d] = e_dir[d];
@@ -1353,10 +1422,32 @@ extern "C" __global__ void experience_action_select(
if (!(thompson_temp > 0.5f)) thompson_temp = 0.5f;
if (thompson_temp > 2.0f) thompson_temp = 2.0f;
if (cooldown_active && DIR_HOLD < b0_size) {
if (plasticity_fire_bar && DIR_FLAT < b0_size) {
/* SP15 Phase 3.5.4.c — fire-bar Flat injection (step 1 of the
* spec §9.2 (3.5.4) two-step recovery). On the bar where the
* plasticity-injection trigger fired this fold (detected via
* `plasticity_warm_remaining ≥ m_warm − 1.5f` per the
* trigger-then-decrement convention) we force dir_idx =
* DIR_FLAT to close any open position. The OR-gate cooldown
* below would have forced DIR_HOLD on the same bar (warm > 0);
* the fire-bar branch supersedes it because Flat-this-bar is
* the directed exit, not Hold-this-bar.
*
* Per spec the warm counter then runs out over bars
* T+1..T+M during which the OR-gate forces DIR_HOLD (the
* `else if (cooldown_active)` branch below). DIR_FLAT = 3
* by `state_layout.cuh`. Skips Pass 2 + RNG draws same as
* the cooldown branch — reproducibility-at-fixed-seed
* preserved because plasticity_fire_bar is a deterministic
* function of the kernel's ISV reads. */
dir_idx = DIR_FLAT;
} else if (cooldown_active && DIR_HOLD < b0_size) {
/* SP15 Phase 3.5.3.b — cooldown mask: when
* `ISV[COOLDOWN_BARS_REMAINING=435] > 0`, force the picked
* direction to Hold per spec §9.2 (3.5.3) post-amendment-2 fix.
* SP15 Phase 3.5.4.c extends the predicate to the OR of
* COOLDOWN_BARS_REMAINING and PLASTICITY_WARM_BARS_REMAINING
* so the post-fire warm-up window also forces Hold.
* Hard short-circuit (skip Pass 2 entirely) — sidesteps the
* temperature-blend numerics where a finite-sentinel-on-non-Hold
* approach would still let a pure-Thompson (τ=1) sample dominate

11
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

@@ -1994,6 +1994,17 @@ impl GpuBacktestEvaluator {
.arg(&per_sample_support_ptr) // per_sample_support: [N, 4, 3] stride-12 mapped pinned (trainer-owned)
.arg(&atom_positions_ptr) // atom_positions: [b0, n_atoms] adaptive C51 positions (trainer-owned)
.arg(&n_atoms_i32) // n_atoms: C51 atom count (i32 to match `int` ABI)
// SP15 Phase 3.5.4.c (2026-05-07): plasticity_m_warm —
// pass 0.0f during eval/backtest. Plasticity injection
// is a TRAINING-only mechanism (it resets advantage-head
// weights to break out of stuck-in-flat regimes during
// experience collection); during deterministic eval the
// weights must remain frozen at their checkpoint values.
// m_warm = 0.0f makes the action_select fire-bar
// predicate dead, and PLASTICITY_WARM_BARS_REMAINING
// stays at its 0 sentinel through eval (no producer
// running) so the OR-gate degenerates to cooldown-only.
.arg(&0.0f32)
.launch(launch_cfg)
.map_err(|e| MLError::ModelError(format!(
"chunked action_select chunk_start={chunk_start}: {e}"

37
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -10157,6 +10157,43 @@ impl GpuDqnTrainer {
self.target_params_buf.raw_ptr()
}
/// SP15 Phase 3.5.4.c (2026-05-07): return the
/// `(w_b0out_dev_ptr, n_weights, fan_in)` triple for the
/// directional advantage-head last-Linear weight tensor (param-
/// buffer tensor index 19). The collector's
/// `set_sp15_plasticity_target()` consumes this triple to wire
/// the per-step `launch_sp15_plasticity_injection` against the
/// directional branch only.
///
/// Layout (per `compute_param_sizes` and the row/col-dim pairs
/// at `gpu_dqn_trainer.rs:29946`):
/// `w_b0out: [branch_0_size × num_atoms, adv_h]` row-major
/// (output_dim rows × input_dim columns).
/// `n_weights = branch_0_size × num_atoms × adv_h`
/// (default config: 4 × 51 × 128 = 26112).
/// `fan_in = adv_h` (= input dim of the per-unit weight
/// vector — Kaiming-He gain `sqrt(2/fan_in)` is per-output-
/// unit variance; default config: 128).
///
/// The Wave 4.2 audit doc previously documented `fan_in =
/// adv_h × num_atoms = 6528`. That value is the total weight
/// count of one row of branch outputs (`num_atoms` atoms per
/// direction × `adv_h` inputs), not the input dim per unit; the
/// Kaiming-He variance derivation uses input dim per unit, which
/// is `adv_h`. The Phase 3.5.4.c audit-doc inline correction
/// restores the right value (`128`) and the right stddev
/// (`sqrt(2/128) ≈ 0.125`).
pub fn sp15_w_b0out_target(&self) -> (u64, i32, i32) {
let param_sizes = compute_param_sizes(&self.config);
let byte_offset = padded_byte_offset(&param_sizes, 19);
let dev_ptr = self.params_buf.raw_ptr() + byte_offset;
let n_weights = (self.config.branch_0_size
* self.config.num_atoms
* self.config.adv_h) as i32;
let fan_in = self.config.adv_h as i32;
(dev_ptr, n_weights, fan_in)
}
/// Backward gradient w.r.t. h_s2 (trunk activation) — f32.
pub fn bw_d_h_s2_buf(&self) -> &CudaSlice<f32> {
&self.bw_d_h_s2

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

@@ -983,6 +983,48 @@ pub struct GpuExperienceCollector {
/// pattern.
sp15_dd_trajectory_prev_dd_dev_ptr: u64,
/// SP15 Phase 3.5.4.c (2026-05-07): trainer's `params_buf`
/// device pointer offset to `w_b0out` (directional advantage-head
/// last-Linear weights, tensor index 19 per `compute_param_sizes`).
/// Set via `set_sp15_plasticity_target()` after trainer
/// construction. 0 = NULL → the per-step plasticity-injection
/// launch in step 4-pre is skipped AND `m_warm` is passed as 0.0f
/// to `experience_action_select` so the OR-gate consumer
/// degenerates to cooldown-only (bit-identical pre-3.5.4.c
/// behaviour). Production sets this in `training_loop.rs`
/// immediately after `set_sp15_dd_trajectory_prev_dd_ptr` —
/// mirrors the same wiring pattern.
///
/// Plasticity is injected into the directional branch only —
/// per spec §9.2 (3.5.4) the recovery is about escaping stuck-
/// in-flat regimes (a directional decision problem), so resetting
/// the trailing 10% of the directional output projection is
/// targeted at the exact failure mode. mag/order/urgency branches
/// are untouched — their conditioning on direction makes them
/// naturally re-train once direction starts choosing again.
sp15_w_b0out_dev_ptr: u64,
/// SP15 Phase 3.5.4.c: total weight count for `w_b0out`
/// (= `branch_0_size × num_atoms × adv_h`; default config
/// 4 × 51 × 128 = 26112). Computed once at wiring time from
/// the trainer's network config; passed as the `n_weights`
/// arg to `launch_sp15_plasticity_injection` so the kernel's
/// `reset_count = n_weights / 10` (~2611 default config)
/// addresses the trailing-10%-slice within the directional
/// output projection.
sp15_w_b0out_n_weights: i32,
/// SP15 Phase 3.5.4.c: Kaiming-He fan_in for `w_b0out`. Per
/// `compute_param_sizes` line 4074 the layout is
/// `[branch_0_size × num_atoms, adv_h]` (row-major output_dim ×
/// input_dim) so fan_in = adv_h (default config 128). The Wave
/// 4.2 audit doc previously documented fan_in =
/// `adv_h × num_atoms = 6528` — that was incorrect (total
/// weight count, not input dim per unit) and the Phase 3.5.4.c
/// audit-doc inline correction restores the right value. With
/// the correct fan_in the Kaiming-He stddev is
/// `sqrt(2 / 128) ≈ 0.125` (vs the audit doc's earlier
/// erroneous `sqrt(2/6528) ≈ 0.0175`).
sp15_w_b0out_fan_in: i32,
/// SP4 Layer A Task A13.5 (2026-05-01); GPU-Pearls refactor (2026-05-01)
/// — cross-boundary Pearls A+D wiring for
/// `launch_reward_component_ema_inplace`. The producer kernel writes
@@ -1999,6 +2041,9 @@ impl GpuExperienceCollector {
isv_signals_dev_ptr: 0, // NULL until trainer sets it
sp15_alpha_warm_count_dev_ptr: 0, // NULL until trainer wires it (SP15 Wave 2)
sp15_dd_trajectory_prev_dd_dev_ptr: 0, // NULL until trainer wires it (SP15 Wave 4.3 / 3.5.5.b)
sp15_w_b0out_dev_ptr: 0, // NULL until trainer wires it (SP15 Phase 3.5.4.c)
sp15_w_b0out_n_weights: 0,
sp15_w_b0out_fan_in: 0,
// SP4 Task A13.5: Pearls A+D buffers for reward_component_ema.
// All NULL until FusedTrainingCtx wires them post-collector
// construction (mirrors A14/A15 Pearl C wiring path). 2026-05-01
@@ -2100,6 +2145,42 @@ impl GpuExperienceCollector {
self.sp15_dd_trajectory_prev_dd_dev_ptr = dev_ptr;
}
/// SP15 Phase 3.5.4.c (2026-05-07): set the trainer's `params_buf`
/// device pointer offset to `w_b0out` (directional advantage-head
/// last-Linear weights, tensor index 19 per `compute_param_sizes`)
/// + the matching `n_weights` (= `branch_0_size × num_atoms ×
/// adv_h`; default config 4 × 51 × 128 = 26112) + Kaiming-He
/// `fan_in` (= `adv_h`; default 128). When all three are non-zero
/// alongside `isv_signals_dev_ptr`, the per-step
/// `launch_sp15_plasticity_injection` runs in step 4-pre (BEFORE
/// `experience_action_select`); the trigger writes
/// `ISV[PLASTICITY_FIRED_THIS_FOLD=436]` + `ISV[PLASTICITY_WARM_BARS_
/// REMAINING=438]`, the action_select kernel's OR-gate consumer
/// reads both slots + `m_warm` (defaulted 200.0 per spec §9.2
/// (3.5.4)) to force `dir_idx = DIR_FLAT` on the fire bar and
/// `dir_idx = DIR_HOLD` on subsequent warm-up bars. Pass 0 for any
/// of the three to disable (test scaffold path; the OR-gate
/// degenerates to cooldown-only — bit-identical to the
/// pre-3.5.4.c path).
///
/// The trainer derives `n_weights` and `fan_in` from
/// `compute_param_sizes(&self.config)[19]` and `self.config.adv_h`
/// respectively (per the param-buffer layout at
/// `gpu_dqn_trainer.rs:4074` and the Kaiming-He fan-dim layout at
/// `gpu_dqn_trainer.rs:29946` — the row-dim `branch_0_size ×
/// num_atoms` is the OUTPUT count, the col-dim `adv_h` is the
/// INPUT count = fan_in for the per-unit Kaiming-He variance).
pub fn set_sp15_plasticity_target(
&mut self,
w_b0out_dev_ptr: u64,
n_weights: i32,
fan_in: i32,
) {
self.sp15_w_b0out_dev_ptr = w_b0out_dev_ptr;
self.sp15_w_b0out_n_weights = n_weights;
self.sp15_w_b0out_fan_in = fan_in;
}
/// Set trade plan params device pointer for plan activation/enforcement.
/// The env_step kernel reads plan params to copy into portfolio state on
/// Flat→Positioned transitions. Pass 0 to disable (NULL = no plan).
@@ -4199,6 +4280,70 @@ impl GpuExperienceCollector {
let na_i32 = self.num_atoms as i32;
let support_dev_ptr = self.per_sample_support_buf.dev_ptr;
let null_atom_positions = 0u64; // NULL = use linear atom spacing
// ── 4-pre. SP15 Phase 3.5.4.c (2026-05-07): plasticity-injection
// production launch — runs BEFORE experience_action_select on the
// same stream so the trigger's ISV writes (slots 436 / 438) are
// serialized-visible to the action_select kernel's reads. Closes
// out the Wave 4.2 (`ef08611d3`) orphan-with-tests-only landing
// per `feedback_wire_everything_up`.
//
// Gated on `sp15_w_b0out_dev_ptr != 0`: the trainer wires this
// dev_ptr via `set_sp15_plasticity_target` after construction
// (mirrors `set_sp15_alpha_warm_count_ptr` / `set_sp15_dd_
// trajectory_prev_dd_ptr`). Until wired, the launch is skipped
// and the ISV slots stay at their constructor sentinels (warm = 0,
// fired = 0); the action_select kernel's OR-gate consumer then
// degenerates to cooldown-only (bit-identical to the
// pre-3.5.4.c path).
//
// `plasticity_m_warm_arg` is the warm-up bar count (default 200
// per spec §9.2 (3.5.4)) — the trigger writes it to ISV[438] on
// fire, action_select reads it back to detect the fire bar
// (warm ≥ m_warm − 1.5f per the kernel's trigger-then-decrement
// convention). Pass 0.0f when plasticity is unwired so the fire-
// bar predicate is dead in the kernel.
//
// `plasticity_seed` derives from the global FOXHUNT_SEED
// (`mix_seed`) XOR'd with the per-step timestep `t` so the
// cuRAND init in the trigger kernel produces deterministic-but-
// distinct samples each fire (multi-fold runs that fire on
// different timesteps see different Kaiming-He samples). Same
// (FOXHUNT_SEED, t) pair always yields the same reset.
//
// Out of the exp_fwd_graph capture region (which ends at the
// cuBLAS forward earlier in this loop body), so per
// `pearl_no_host_branches_in_captured_graph` the host-side
// gating compiles cleanly.
let plasticity_m_warm_arg: f32 =
if self.sp15_w_b0out_dev_ptr != 0 && self.isv_signals_dev_ptr != 0 {
const M_WARM: f32 = 200.0; // spec §9.2 (3.5.4) default
let plasticity_seed = crate::cuda_pipeline::mix_seed(
/* SP15 Phase 3.5.4.c historic base — arbitrary
* 64-bit constant unique to this launch site so
* multi-launcher runs (alpha_split / cooldown /
* etc.) get uncorrelated mix outputs even when
* FOXHUNT_SEED defaults to 42. */
0xC0DE_F00D_5915_3F4Cu64
) ^ (t as u64);
crate::cuda_pipeline::gpu_dqn_trainer::launch_sp15_plasticity_injection(
&self.stream,
self.isv_signals_dev_ptr,
self.sp15_w_b0out_dev_ptr,
self.sp15_w_b0out_n_weights,
M_WARM,
self.sp15_w_b0out_fan_in,
plasticity_seed,
).map_err(|e| MLError::ModelError(format!(
"sp15 plasticity_injection t={t}: {e}"
)))?;
M_WARM
} else {
// Plasticity unwired — pass m_warm = 0.0 so the
// action_select fire-bar branch is inactive.
0.0f32
};
unsafe {
self.stream
.launch_builder(&self.action_select_kernel)
@@ -4233,6 +4378,19 @@ impl GpuExperienceCollector {
.arg(&support_dev_ptr) // per_sample_support: [N, 4, 3] stride-12 mapped pinned
.arg(&null_atom_positions) // atom_positions: NULL → linear support
.arg(&na_i32) // n_atoms: C51 atom count (i32 to match `int` ABI)
// SP15 Phase 3.5.4.c (2026-05-07): plasticity-injection
// m_warm. When the per-step plasticity launch above
// wired this collector's `sp15_w_b0out_dev_ptr`, the
// m_warm passed to the trigger flows into the
// consumer here so the action_select kernel can
// detect the fire bar (warm ≥ m_warm − 1.5f) and
// force DIR_FLAT, then DIR_HOLD on subsequent
// warm-up bars via the OR-gate. When unwired
// (`sp15_w_b0out_dev_ptr == 0`), m_warm = 0.0f
// disables the fire-bar branch (warm == 0 too →
// OR-gate degenerates to cooldown-only — bit-
// identical to the pre-3.5.4.c path).
.arg(&plasticity_m_warm_arg)
.launch(launch_cfg)
.map_err(|e| MLError::ModelError(format!(
"experience_action_select t={t}: {e}"

3
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -1369,6 +1369,9 @@ mod tests {
.arg(&support_buf) // per_sample_support [batch, b0, 3]
.arg(&null_ptr) // atom_positions = NULL → linear support
.arg(&n_atoms) // n_atoms
// SP15 Phase 3.5.4.c: plasticity_m_warm — 0.0f for this
// diagnostic test (no plasticity wiring in this test scaffold).
.arg(&0.0f32)
.launch(cfg)
.expect("launch experience_action_select");
}

3
crates/ml/src/trainers/dqn/distributional_q_tests.rs
View File

@@ -1025,6 +1025,9 @@ impl ProdActionSelectFixture {
.arg(&self.per_sample_support)
.arg(&null_ptr) // atom_positions = NULL → linear support
.arg(&self.n_atoms)
// SP15 Phase 3.5.4.c: plasticity_m_warm — 0.0f for this
// distributional Q test (no plasticity wiring needed).
.arg(&0.0f32)
.launch(cfg)
.expect("launch experience_action_select");
}

26
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -893,6 +893,32 @@ impl DQNTrainer {
}
}
// SP15 Phase 3.5.4.c (2026-05-07): wire the trainer's
// `params_buf` device pointer offset to `w_b0out` (directional
// advantage-head last-Linear weights, tensor index 19 per
// `compute_param_sizes`) + `n_weights` (= `branch_0_size ×
// num_atoms × adv_h`; default config 4 × 51 × 128 = 26112) +
// Kaiming-He `fan_in` (= `adv_h`; default 128) into the
// collector so the per-step `launch_sp15_plasticity_injection`
// can run inside `collect_experiences_gpu` BEFORE
// `experience_action_select`. Closes out the Wave 4.2
// (`ef08611d3`) orphan-with-tests-only landing per
// `feedback_wire_everything_up`. The plasticity trigger writes
// ISV[PLASTICITY_FIRED_THIS_FOLD=436] +
// ISV[PLASTICITY_WARM_BARS_REMAINING=438] which the
// action_select kernel then reads to force DIR_FLAT on the
// fire bar and DIR_HOLD on subsequent warm-up bars (OR-gate
// extension of the SP15 Phase 3.5.3.b cooldown mask).
// Mirrors the `set_sp15_dd_trajectory_prev_dd_ptr` wiring
// pattern.
if let Some(ref fused) = self.fused_ctx {
let (w_ptr, n_weights, fan_in) =
fused.trainer().sp15_w_b0out_target();
if let Some(ref mut collector) = self.gpu_experience_collector {
collector.set_sp15_plasticity_target(w_ptr, n_weights, fan_in);
}
}
// SP15 Wave 4.3 / Phase 3.5.5.b (2026-05-07): wire the
// trainer's ISV bus dev_ptr into the GPU PER replay buffer
// so `per_insert_pa` can read

329
crates/ml/tests/sp15_phase1_oracle_tests.rs
View File

@@ -2123,6 +2123,84 @@ mod gpu {
.arg(per_sample_support)
.arg(&null_ptr) // atom_positions = NULL → linear support
.arg(&n_atoms)
// SP15 Phase 3.5.4.c (2026-05-07): plasticity_m_warm.
// Default 0.0f → fire-bar predicate dead. The new
// `plasticity_fires_force_flat_then_cooldown_holds` test
// overrides this via `launch_action_select_for_test_with_m_warm`.
.arg(&0.0f32)
.launch(cfg)
.expect("launch experience_action_select");
}
stream.synchronize().expect("sync after action_select");
stream.clone_dtoh(actions_buf).expect("readback actions")
}
/// SP15 Phase 3.5.4.c (2026-05-07): variant of `launch_action_select_for_test`
/// that takes an explicit `m_warm` so the new behavioral test
/// `plasticity_fires_force_flat_then_cooldown_holds` can drive the
/// kernel's fire-bar / OR-gate consumer paths. All other args
/// match the m_warm-defaulted helper.
#[allow(clippy::too_many_arguments)]
fn launch_action_select_for_test_with_m_warm(
stream: &Arc<CudaStream>,
batch: usize,
n_atoms: i32,
q_values: &CudaSlice<f32>,
b_logits_dir: &CudaSlice<f32>,
per_sample_support: &CudaSlice<f32>,
actions_buf: &mut CudaSlice<i32>,
q_gaps_buf: &mut CudaSlice<f32>,
intent_buf: &mut CudaSlice<i32>,
conv_buf: &mut CudaSlice<f32>,
mag_conv_buf: &mut CudaSlice<f32>,
isv_dev_ptr: u64,
m_warm: f32,
) -> Vec<i32> {
let kernel = load_action_select(stream);
let blocks = ((batch as u32).div_ceil(256)).max(1);
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
let n_i32: i32 = batch as i32;
let null_ptr: u64 = 0;
let eps_start: f32 = 0.0;
let eps_end: f32 = 0.0;
unsafe {
stream
.launch_builder(kernel)
.arg(q_values)
.arg(&mut *actions_buf)
.arg(&mut *q_gaps_buf)
.arg(&eps_start)
.arg(&eps_end)
.arg(&0_i32) // current_epoch
.arg(&1_i32) // total_epochs
.arg(&n_i32)
.arg(&AS_B0)
.arg(&AS_B1)
.arg(&AS_B2)
.arg(&AS_B3)
.arg(&0.0_f32) // q_gap_threshold
.arg(&null_ptr)
.arg(&0_i32)
.arg(&0.0_f32)
.arg(&1.0_f32)
.arg(&1.0_f32)
.arg(&1.0_f32)
.arg(&0_i32)
.arg(&null_ptr)
.arg(&isv_dev_ptr)
.arg(&0_i32)
.arg(&mut *intent_buf)
.arg(&mut *conv_buf)
.arg(&mut *mag_conv_buf)
.arg(b_logits_dir)
.arg(per_sample_support)
.arg(&null_ptr)
.arg(&n_atoms)
.arg(&m_warm)
.launch(cfg)
.expect("launch experience_action_select");
}
@@ -2345,6 +2423,257 @@ mod gpu {
);
}
/// SP15 Phase 3.5.4.c (2026-05-07) — full plasticity-injection two-step
/// recovery sequence end-to-end. Sets up the ISV state so the trigger
/// fires (DD_PERSISTENCE > threshold, fired flag clear, warm = 0),
/// then drives the per-step launch sequence
/// `launch_sp15_plasticity_injection` → `experience_action_select`
/// across `M_warm + 2` bars. Asserts:
///
/// - **Bar 0 (fire bar)**: dir_picked == DIR_FLAT — the trigger sets
/// warm = M_warm then decrements to (M_warm − 1), so action_select
/// observes warm ≥ M_warm − 1.5f and forces DIR_FLAT (step 1 of
/// the spec §9.2 (3.5.4) two-step). e_dir is set to strongly
/// prefer Long (gap = +1.0); without the fire-bar branch the
/// argmax would land on Long, so observing DIR_FLAT proves the
/// fire-bar consumer fired.
///
/// - **Bars 1..M_warm − 1 (warm-up)**: dir_picked == DIR_HOLD on
/// every bar. The trigger debounces (already fired this fold)
/// but continues to decrement warm; the action_select OR-gate
/// forces Hold while warm > 0 (step 2 of the two-step). Bar 1
/// observes warm = M_warm − 2; bar M_warm − 1 observes warm = 1.
///
/// - **Bar M_warm (warm expired)**: warm = 0 (the trigger's
/// decrement hit floor on the previous bar). action_select OR-gate
/// inactive (cooldown also 0); standard Thompson + hold_floor
/// argmax resumes; with Long strongly preferred and entropy near
/// 0 → hold_floor ≈ 0 → dir_picked == DIR_LONG.
///
/// Uses M_warm = 5 (down from production's 200) to keep the test fast
/// while preserving the invariant. The `fan_in` / `n_weights` choices
/// don't matter for the consumer-side assertions (the kernel writes
/// the trailing 10% of advantage_head_weights in addition to its ISV
/// updates, but this test doesn't read those weights). Driven via
/// `launch_action_select_for_test_with_m_warm` so the consumer sees
/// the same `m_warm` the trigger wrote.
#[test]
#[ignore = "requires GPU"]
fn plasticity_fires_force_flat_then_cooldown_holds() {
use ml::cuda_pipeline::sp15_isv_slots::{
COOLDOWN_BARS_REMAINING_INDEX, DD_PERSISTENCE_INDEX,
HOLD_FLOOR_ALPHA_INDEX, HOLD_FLOOR_EPS0_INDEX, HOLD_FLOOR_K_INDEX,
PLASTICITY_FIRED_THIS_FOLD_INDEX,
PLASTICITY_PERSISTENCE_THRESHOLD_INDEX,
PLASTICITY_WARM_BARS_REMAINING_INDEX,
};
let stream = make_test_stream();
let batch: usize = 1;
let n_atoms: i32 = 51;
let v_min: f32 = -1.0;
let v_max: f32 = 1.0;
let delta_z = (v_max - v_min) / (n_atoms as f32 - 1.0);
// Long strongly preferred → without fire-bar / OR-gate, argmax = Long.
let target_v = [0.0_f32, 0.0_f32, 1.0_f32, 0.0_f32];
let b_logits_host = fill_peaked_logits_action_select(
batch, n_atoms, v_min, delta_z, &target_v, /*peak_logit*/ 50.0,
);
let support_host = fill_linear_support_action_select(batch, v_min, v_max, delta_z);
let q_values_host = vec![0.0_f32; batch * AS_Q_STRIDE];
// ISV bus pre-fire: persistence way above threshold, fired clear,
// warm 0, cooldown 0. hold_floor params identical to the cooldown
// tests above so any DIR_HOLD outcome is forced by the OR-gate, not
// by hold_floor (entropy near 0 with one-hot logits → hold_floor ≈ 0).
let isv_buf = unsafe { MappedF32Buffer::new(ISV_LEN) }
.expect("alloc MappedF32Buffer for isv bus");
let mut isv_init = vec![0.0f32; ISV_LEN];
isv_init[339] = 0.5; // ISV_EVAL_THOMPSON_TEMP_IDX (floor)
isv_init[HOLD_FLOOR_ALPHA_INDEX] = 0.5;
isv_init[HOLD_FLOOR_K_INDEX] = 10.0;
isv_init[HOLD_FLOOR_EPS0_INDEX] = 1.0;
isv_init[COOLDOWN_BARS_REMAINING_INDEX] = 0.0;
isv_init[DD_PERSISTENCE_INDEX] = 1_000.0;
isv_init[PLASTICITY_PERSISTENCE_THRESHOLD_INDEX] = 1.0;
isv_init[PLASTICITY_FIRED_THIS_FOLD_INDEX] = 0.0;
isv_init[PLASTICITY_WARM_BARS_REMAINING_INDEX] = 0.0;
isv_buf.write_from_slice(&isv_init);
// advantage_head_weights buffer for the trigger's Kaiming-He reset.
// Uses the same 64-elem skeleton the existing plasticity oracle
// tests use; n_weights / 10 = 6 → 6 trailing slots get rewritten.
const N_WEIGHTS: i32 = 64;
const FAN_IN: i32 = 128; // matches default config adv_h
const M_WARM: f32 = 5.0; // shrunk from production 200.0 for test runtime
let weights_buf = unsafe { MappedF32Buffer::new(N_WEIGHTS as usize) }
.expect("alloc MappedF32Buffer for advantage_head_weights");
weights_buf.write_from_slice(&vec![1.0f32; N_WEIGHTS as usize]);
// Per-step GPU buffers (allocated once, reused across bars).
let q_values = stream.clone_htod(&q_values_host).expect("upload q_values");
let b_logits_dir = stream.clone_htod(&b_logits_host).expect("upload logits");
let per_sample_support = stream
.clone_htod(&support_host)
.expect("upload support");
let mut actions_buf = stream.alloc_zeros::<i32>(batch).expect("alloc actions");
let mut q_gaps_buf = stream.alloc_zeros::<f32>(batch).expect("alloc q_gaps");
let mut intent_buf = stream.alloc_zeros::<i32>(batch).expect("alloc intent");
let mut conv_buf = stream.alloc_zeros::<f32>(batch).expect("alloc conv");
let mut mag_conv = stream.alloc_zeros::<f32>(batch).expect("alloc mag_conv");
let m_warm_i = M_WARM as usize;
// ── Bar 0: fire bar → expect DIR_FLAT ──────────────────────────────
ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_plasticity_injection(
&stream,
isv_buf.dev_ptr,
weights_buf.dev_ptr,
N_WEIGHTS,
M_WARM,
FAN_IN,
/*seed*/ 0xFEED_BEEF_CAFE_BABEu64,
).expect("launch plasticity (fire bar)");
let actions = launch_action_select_for_test_with_m_warm(
&stream, batch, n_atoms,
&q_values, &b_logits_dir, &per_sample_support,
&mut actions_buf, &mut q_gaps_buf, &mut intent_buf,
&mut conv_buf, &mut mag_conv,
isv_buf.dev_ptr, M_WARM,
);
let dir_picked = decode_dir(actions[0]);
// DIR_FLAT = 3 per state_layout.cuh.
assert_eq!(
dir_picked, 3,
"fire bar must force DIR_FLAT (=3); got dir = {dir_picked} \
(expected 3, even though Long had +1.0 advantage). \
ISV[fired]={}, ISV[warm]={}",
isv_buf.read_all()[PLASTICITY_FIRED_THIS_FOLD_INDEX],
isv_buf.read_all()[PLASTICITY_WARM_BARS_REMAINING_INDEX],
);
// Sanity-check the ISV state after the fire bar:
// fired flipped 0 → 1, warm = M_warm − 1 (per the kernel's
// trigger-then-decrement convention).
let isv_after_fire = isv_buf.read_all();
assert!(
(isv_after_fire[PLASTICITY_FIRED_THIS_FOLD_INDEX] - 1.0).abs() < 1e-9,
"after fire-bar PLASTICITY_FIRED_THIS_FOLD must be 1.0; got {}",
isv_after_fire[PLASTICITY_FIRED_THIS_FOLD_INDEX]
);
let warm_after_fire = isv_after_fire[PLASTICITY_WARM_BARS_REMAINING_INDEX];
assert!(
(warm_after_fire - (M_WARM - 1.0)).abs() < 1e-4,
"after fire-bar PLASTICITY_WARM_BARS_REMAINING must be M_warm − 1 = {}; got {}",
M_WARM - 1.0, warm_after_fire
);
// ── Bars 1..M_warm − 1 (warm-up) → expect DIR_HOLD ──────────────────
// Trigger continues to run on every bar (debounced — fired = 1, no
// re-fire), per-bar warm decrement runs unconditionally; consumer
// sees warm decreasing from M_warm − 2 down to 1; OR-gate forces Hold.
//
// Trace (M_warm = 5):
// Bar 0 (fire): trigger sets warm=5 → decrements to 4. action_select
// reads warm=4 ≥ 3.5 → DIR_FLAT. (asserted above)
// Bar 1: trigger debounces, warm 4→3. action_select reads warm=3 → Hold.
// Bar 2: warm 3→2 → Hold.
// Bar 3: warm 2→1 → Hold.
// Bar 4: warm 1→0. action_select reads warm=0 — OR-gate inactive →
// normal Thompson (NOT Hold). This bar is the "post-warm"
// assertion below, not part of the warm-up loop.
//
// So the warm-up loop runs for bars [1, 2, 3] = bars 1..M_warm − 1
// = bars 1..4 in iter terms (3 iterations). The previous formulation
// `1..m_warm_i` overran by one bar.
for bar in 1..(m_warm_i - 1) {
ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_plasticity_injection(
&stream,
isv_buf.dev_ptr,
weights_buf.dev_ptr,
N_WEIGHTS,
M_WARM,
FAN_IN,
/*seed*/ 0xFEED_BEEF_CAFE_BABEu64 ^ (bar as u64),
).expect("launch plasticity (warm-up)");
let actions = launch_action_select_for_test_with_m_warm(
&stream, batch, n_atoms,
&q_values, &b_logits_dir, &per_sample_support,
&mut actions_buf, &mut q_gaps_buf, &mut intent_buf,
&mut conv_buf, &mut mag_conv,
isv_buf.dev_ptr, M_WARM,
);
let dir_picked = decode_dir(actions[0]);
assert_eq!(
dir_picked, AS_DIR_HOLD,
"warm-up bar {bar} must force DIR_HOLD (={AS_DIR_HOLD}); \
got dir = {dir_picked}. ISV[warm]={}",
isv_buf.read_all()[PLASTICITY_WARM_BARS_REMAINING_INDEX]
);
}
// After (M_warm − 2) warm-up bars the warm counter is now 1 — the
// very next call's per-bar decrement takes it to 0 (the post-warm
// boundary). Sanity-check that we're at the boundary so the post-
// warm assertion below tests the ACTUAL transition, not some later
// already-decremented state.
let isv_at_boundary = isv_buf.read_all();
let warm_at_boundary = isv_at_boundary[PLASTICITY_WARM_BARS_REMAINING_INDEX];
assert!(
(warm_at_boundary - 1.0).abs() < 1e-4,
"after warm-up loop PLASTICITY_WARM_BARS_REMAINING must be 1 (= boundary); got {}",
warm_at_boundary
);
// ── Bar M_warm − 1 (warm transitions to 0) → normal action selection resumes ─
// The trigger still runs (debounced — `fired` stays 1 for the
// remainder of the fold) and decrements warm 1 → 0. action_select
// reads warm = 0: OR-gate inactive (cooldown also 0), fire-bar
// predicate false (warm=0 < 3.5). With Long strongly preferred and
// hold_floor ≈ 0 (one-hot logits → low entropy), Thompson argmax
// returns DIR_LONG.
ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_plasticity_injection(
&stream,
isv_buf.dev_ptr,
weights_buf.dev_ptr,
N_WEIGHTS,
M_WARM,
FAN_IN,
/*seed*/ 0xFEED_BEEF_CAFE_BABEu64 ^ (m_warm_i as u64 + 1),
).expect("launch plasticity (post-warm)");
let actions = launch_action_select_for_test_with_m_warm(
&stream, batch, n_atoms,
&q_values, &b_logits_dir, &per_sample_support,
&mut actions_buf, &mut q_gaps_buf, &mut intent_buf,
&mut conv_buf, &mut mag_conv,
isv_buf.dev_ptr, M_WARM,
);
let dir_picked = decode_dir(actions[0]);
assert_eq!(
dir_picked, AS_DIR_LONG,
"post-warm bar must follow normal Thompson argmax → DIR_LONG; \
got dir = {dir_picked}. ISV[warm]={}, ISV[fired]={}",
isv_buf.read_all()[PLASTICITY_WARM_BARS_REMAINING_INDEX],
isv_buf.read_all()[PLASTICITY_FIRED_THIS_FOLD_INDEX]
);
// Post-test ISV state: fired stays 1 (debounced for remainder of
// fold — fold-reset registry arms re-arm to 0 on next fold
// boundary), warm at 0 (decrement guard prevents underflow).
let isv_final = isv_buf.read_all();
assert!(
(isv_final[PLASTICITY_FIRED_THIS_FOLD_INDEX] - 1.0).abs() < 1e-9,
"post-test PLASTICITY_FIRED_THIS_FOLD must remain 1 (debounced); got {}",
isv_final[PLASTICITY_FIRED_THIS_FOLD_INDEX]
);
assert!(
isv_final[PLASTICITY_WARM_BARS_REMAINING_INDEX].abs() < 1e-4,
"post-test PLASTICITY_WARM_BARS_REMAINING must be 0; got {}",
isv_final[PLASTICITY_WARM_BARS_REMAINING_INDEX]
);
}
// ====================================================================
// SP15 Wave 2 (2026-05-06) — fused post-SP11 reward-axis composer
// oracle tests for `compute_sp15_final_reward_kernel`.

4
docs/dqn-wire-up-audit.md
View File

File diff suppressed because one or more lines are too long