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

feat(sp15-wave4.2 / 3.5.4.b): cuRAND Kaiming-He weight reset for plasticity injection

Browse Source 
Phase 3.5.4 (e0e0abfb2) landed the trigger + warm-up tracker but
deferred the actual weight reset (kernel accepted advantage_head_weights
+ n_weights but no-op'd via (void) cast). Wave 4.2 lands the real
reset.

When DD_PERSISTENCE exceeds threshold AND not yet fired this fold,
the kernel now resets the last 10% of advantage-head weights to
Kaiming-He init: Normal(0, sqrt(2/fan_in)) sampled via cuRAND
curand_normal() with per-thread state initialized from a host-passed
seed (Option A: per-call curand_init(seed, tid, 0, &state) for full
determinism — bit-identical samples for identical (seed, tid) pairs).

Single kernel, two phases:
  1. Every thread independently re-evaluates fire_now from the same
     ISV reads (DD_PERSISTENCE / threshold / fired_flag); the trigger
     condition is a pure function of these reads so all threads
     converge without cross-block synchronisation. Block 0 / thread 0
     also runs the trigger + warm-bars decrement (single-thread ISV
     write path).
  2. If fire_now: each tid < reset_count writes Kaiming-He sample to
     advantage_head_weights[reset_start + tid]. Per-thread independent
     write, no atomic, no reduction (feedback_no_atomicadd clean).

New launcher params: fan_in (i32), seed (u64). Grid:
((n_weights/10 + 255) / 256).max(1) x [256, 1, 1]. The .max(1) floor
ensures block 0 always exists even when n_weights/10 == 0.

cuRAND device functions (curand_init, curand_normal) are inlined into
the cubin by nvcc from <curand_kernel.h> in the standard CUDA toolkit
include path — no host-side cuRAND linker dependency required, no
build.rs link change needed.

Oracle test plasticity_injection_kernel_resets_last_10pct_kaiming_he
verifies: (a) ISV[fired] flipped 0->1, (b) ISV[warm] = m_warm - 1,
(c) first 90% bit-identical to 1.0, (d) last 10% all moved off 1.0,
(e) sample mean |mean| < 0.05, (f) sample std within +-20% of
sqrt(2/fan_in), (g) determinism re-check produces bit-identical
samples for identical seed.

3 existing trigger tests migrated to the new launcher signature; the
debounced + no-fire variants additionally assert weight-stability
(early-out path skips the reset region).

Atomic per feedback_no_partial_refactor: kernel + launcher + 4 tests
+ audit doc + build.rs comment + 2 docstrings land together.

Eliminates Phase 3.5.4 deferred consumer per feedback_wire_everything_up
(the (void) casts are gone).

Action-selection consumer wiring (Phase 3.5.4.c) remains the separate
follow-up per the established Phase 3.5.X pattern.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 01:15:39 +02:00
parent a54f53e4ed
commit ef08611d3f
5 changed files with 413 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
crates/ml/build.rs
View File

@@ -1036,20 +1036,24 @@ fn main() {
// (3.5.4) post-amendment-2 fix: fire when DD_PERSISTENCE
// exceeds threshold AND PLASTICITY_FIRED_THIS_FOLD == 0 →
// set fired flag, set warm-bars counter to M_warm (default
// 200), [DEFERRED: reset last 10% of advantage-head weights];
// per-bar warm-bars decrement; consumer-wiring follow-up reads
// max(COOLDOWN_BARS_REMAINING, PLASTICITY_WARM_BARS_REMAINING)
// and forces Hold while > 0. Phase 3.5.4 lands kernel +
// launcher + 3 ISV anchor seeds + 3 registry entries + 3
// dispatch arms only; the action-selection two-step Flat +
// cooldown wiring AND the actual Kaiming-He weight-reset are
// both deferred to a Phase 3.5.4.b follow-up atomic commit per
// `feedback_no_partial_refactor.md` (matching Phase 1.1-1.5 +
// 3.1-3.5.3 precedent — kernel + launcher verify in isolation
// first via the GPU oracle tests below). The kernel signature
// accepts the advantage_head_weights pointer + n_weights param
// so the follow-up commit is purely additive on the kernel
// body, not on the launcher contract.
// 200), AND reset last 10% of advantage-head weights to
// Kaiming-He init (`Normal(0, sqrt(2/fan_in))` sampled via
// cuRAND `curand_normal()`); per-bar warm-bars decrement;
// consumer-wiring follow-up reads max(COOLDOWN_BARS_REMAINING,
// PLASTICITY_WARM_BARS_REMAINING) and forces Hold while > 0.
// Phase 3.5.4 (2026-05-06) landed kernel + launcher + 3 ISV
// anchor seeds + 3 registry entries + 3 dispatch arms with
// the weight-reset deferred (kernel accepted
// advantage_head_weights + n_weights but no-op'd via (void)
// cast); Wave 4.2 / Phase 3.5.4.b (2026-05-07) lands the
// actual cuRAND Kaiming-He reset, extending the launcher
// signature with `fan_in: i32` and `seed: u64`. The kernel
// includes `<curand_kernel.h>` for device-side
// `curand_init`/`curand_normal` (inlined into the cubin by
// nvcc — no host-side cuRAND linker dependency required).
// The action-selection two-step Flat + cooldown consumer
// wiring remains a separate follow-up (Phase 3.5.4.c) per
// `feedback_no_partial_refactor.md`.
"plasticity_injection_kernel.cu",
// SP15 Phase 3.5.5 (2026-05-06): recovery curriculum in PER —
// per-step DD_TRAJECTORY_DECREASING proxy. Replaces non-existent

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

@@ -1705,13 +1705,16 @@ pub fn launch_sp15_cooldown(
Ok(())
}
/// SP15 Phase 3.5.4 (2026-05-06): cubin slot for the plasticity
/// injection trigger + warm-up tracker kernel. TWO-STEP recovery per
/// spec §9.2 (3.5.4) post-amendment-2 fix:
/// SP15 Phase 3.5.4 (2026-05-06) + Wave 4.2 / Phase 3.5.4.b
/// (2026-05-07): cubin slot for the plasticity injection trigger +
/// warm-up tracker + cuRAND Kaiming-He weight-reset kernel. TWO-STEP
/// recovery per spec §9.2 (3.5.4) post-amendment-2 fix:
/// 1. Fire when `DD_PERSISTENCE > PLASTICITY_PERSISTENCE_THRESHOLD`
/// AND `PLASTICITY_FIRED_THIS_FOLD == 0` → set fired flag, set
/// warm-bars counter to M_warm (default 200), [DEFERRED: reset
/// last 10% of advantage-head weights to Kaiming-He init].
/// warm-bars counter to M_warm (default 200), AND reset the
/// trailing 10% of advantage-head weights to Kaiming-He init
/// (`Normal(0, sqrt(2/fan_in))` sampled via cuRAND
/// `curand_normal()`; gain=√2 for ReLU/GLU activations).
/// 2. Per-bar warm-bars decrement; consumer-wiring follow-up reads
/// `max(COOLDOWN_BARS_REMAINING, PLASTICITY_WARM_BARS_REMAINING)`
/// and forces Hold while > 0 (the OR-gate combining cooldown +
@@ -1722,47 +1725,65 @@ pub fn launch_sp15_cooldown(
/// `dd_asymmetric_reward_kernel.cu` / 3.5.3 `cooldown_kernel.cu`
/// single-kernel-per-cubin pattern.
///
/// Phase 3.5.4 lands kernel + launcher + 3 ISV constructor seeds + 3
/// state-reset registry entries + 3 dispatch arms only; the action-
/// selection two-step Flat + cooldown wiring AND the actual Kaiming-
/// He weight-reset are both deferred to a Phase 3.5.4.b follow-up
/// atomic commit per `feedback_no_partial_refactor.md` (matching
/// Phase 1.1-1.5 + 3.1-3.5.3 precedent — kernel + launcher verify in
/// isolation first via the GPU oracle tests below). The kernel
/// signature accepts the advantage_head_weights pointer + n_weights
/// param so the follow-up commit is purely additive on the kernel
/// body, not on the launcher contract.
/// Phase 3.5.4 landed kernel + launcher + 3 ISV constructor seeds +
/// 3 state-reset registry entries + 3 dispatch arms with the weight-
/// reset deferred (kernel accepted `advantage_head_weights +
/// n_weights` but no-op'd via `(void)` cast). Wave 4.2 / Phase
/// 3.5.4.b lands the actual cuRAND Kaiming-He reset, extending the
/// launcher signature with `fan_in: i32` and `seed: u64` and
/// promoting the kernel grid from single-thread/single-block to a
/// 256-thread per-element grid for the weight write. The action-
/// selection two-step Flat + cooldown consumer wiring remains a
/// separate follow-up (Phase 3.5.4.c) per `feedback_no_partial_refactor.md`.
pub static SP15_PLASTICITY_INJECTION_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/plasticity_injection_kernel.cubin"));
/// SP15 Phase 3.5.4 (2026-05-06): launcher for the plasticity
/// injection trigger + warm-up tracker kernel. Free function (matches
/// `launch_sp15_dd_penalty` / `launch_sp15_regret_signal` /
/// `launch_sp15_hold_floor` / `launch_sp15_dd_asymmetric_reward` /
/// `launch_sp15_cooldown` precedent — single kernel, single cubin,
/// single launcher) so unit/oracle tests can drive the kernel
/// directly without the trainer struct; production callers invoke
/// the same launcher.
/// SP15 Phase 3.5.4 (2026-05-06) + Wave 4.2 / Phase 3.5.4.b
/// (2026-05-07): launcher for the plasticity injection trigger +
/// warm-up tracker + cuRAND Kaiming-He weight-reset kernel. Free
/// function (matches `launch_sp15_dd_penalty` /
/// `launch_sp15_regret_signal` / `launch_sp15_hold_floor` /
/// `launch_sp15_dd_asymmetric_reward` / `launch_sp15_cooldown`
/// precedent — single kernel, single cubin, single launcher) so
/// unit/oracle tests can drive the kernel directly without the
/// trainer struct; production callers invoke the same launcher.
///
/// `isv` MUST be the ISV bus (≥ `SP15_SLOT_END=443` f32 slots — slot
/// 404 is read for DD_PERSISTENCE; slot 437 is read for
/// PLASTICITY_PERSISTENCE_THRESHOLD; slots 436 + 438 are read-modify-
/// written for PLASTICITY_FIRED_THIS_FOLD + PLASTICITY_WARM_BARS_
/// REMAINING). `advantage_head_weights` MUST point to at least
/// `n_weights` writable f32 slots — but the weight-reset is DEFERRED
/// to Phase 3.5.4.b so the kernel body currently does NOT touch the
/// buffer; the parameter is wired through the launcher contract so
/// the follow-up commit is purely additive on the kernel body.
/// `n_weights` writable f32 slots — Wave 4.2 lands the actual
/// Kaiming-He reset of the trailing `n_weights / 10` slots when the
/// kernel fires.
///
/// `m_warm` is the warm-up bar count to set on fire (default 200 per
/// spec §9.2 (3.5.4)). Single thread, single block (mirrors
/// `cooldown_kernel` / `dd_asymmetric_reward_kernel` / `hold_floor_kernel`
/// — no reduction).
/// spec §9.2 (3.5.4)). `fan_in` is the input dimension of the
/// advantage head's last Linear layer (the Kaiming-He gain
/// denominator — `stddev = sqrt(2 / fan_in)`); production callers
/// derive it from the trainer's network config. `seed` is a
/// per-call deterministic value (e.g. derived from a trainer RNG
/// state or `SystemTime::now()` nanos) — the kernel re-initialises
/// per-thread `curandState`s from `(seed, sequence=tid, offset=0)`
/// every launch; the same `(seed, tid)` pair always yields the same
/// sample.
///
/// Grid: `[((n_weights / 10 + 255) / 256), 1, 1]` × `[256, 1, 1]`.
/// Block 0 / thread 0 owns the trigger + warm-bars decrement (single-
/// thread, mirrors `cooldown_kernel` / `dd_asymmetric_reward_kernel`
/// / `hold_floor_kernel`); the remaining threads race-free write the
/// per-element Kaiming-He samples (each thread writes a unique
/// `advantage_head_weights[reset_start + tid]`, no atomic, no
/// reduction). When `n_weights / 10 == 0` (oracle edge case) the
/// grid degenerates to `[1, 1, 1]` so the trigger still runs.
pub fn launch_sp15_plasticity_injection(
stream: &Arc<CudaStream>,
isv: cudarc::driver::sys::CUdeviceptr,
advantage_head_weights: cudarc::driver::sys::CUdeviceptr,
n_weights: i32,
m_warm: f32,
fan_in: i32,
seed: u64,
) -> Result<(), MLError> {
let module = stream
.context()
@@ -1775,6 +1796,13 @@ pub fn launch_sp15_plasticity_injection(
.map_err(|e| MLError::ModelError(format!(
"load plasticity_injection_kernel function: {e}"
)))?;
// Grid sized so every weight in the trailing 10% gets a thread.
// `block_dim = 256` matches the codebase's prevailing 1-D
// element-wise launch shape. `grid_dim` always ≥ 1 so block 0
// thread 0 always runs the trigger, even when `n_weights / 10 == 0`
// (oracle edge case where the test passes a tiny weight buffer).
let reset_count = (n_weights / 10).max(0);
let grid_x = ((reset_count + 255) / 256).max(1) as u32;
unsafe {
stream
.launch_builder(&kernel)
@@ -1782,9 +1810,11 @@ pub fn launch_sp15_plasticity_injection(
.arg(&advantage_head_weights)
.arg(&n_weights)
.arg(&m_warm)
.arg(&fan_in)
.arg(&seed)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
grid_dim: (grid_x, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(

196
crates/ml/src/cuda_pipeline/plasticity_injection_kernel.cu
View File

@@ -1,13 +1,14 @@
// crates/ml/src/cuda_pipeline/plasticity_injection_kernel.cu
//
// SP15 Phase 3.5.4 — plasticity injection trigger + warm-up tracker.
// SP15 Wave 4.2 / Phase 3.5.4.b — cuRAND Kaiming-He weight reset (now wired).
//
// TWO-STEP recovery per spec §9.2 (3.5.4) post-amendment-2 fix:
// 1. On fire: set PLASTICITY_FIRED_THIS_FOLD = 1.0, set
// PLASTICITY_WARM_BARS_REMAINING = M_warm, [DEFERRED: reset last
// 10% of advantage-head weights to Kaiming-He init (gain=√2 for
// ReLU/GLU activations) via cuRAND]. The action-selection layer
// (consumer-wiring follow-up) reads the
// PLASTICITY_WARM_BARS_REMAINING = M_warm, AND reset the last 10%
// of advantage-head weights to Kaiming-He init (gain=√2 for
// ReLU/GLU activations) via cuRAND `curand_normal()`. The
// action-selection layer (consumer-wiring follow-up) reads the
// PLASTICITY_FIRED_THIS_FOLD == 1 flag transition AND emits Flat
// that bar.
// 2. Per-bar decrement of PLASTICITY_WARM_BARS_REMAINING; the
@@ -16,27 +17,35 @@
// PLASTICITY_WARM_BARS_REMAINING)` to force Hold while either
// counter > 0.
//
// **Weight-reset is DEFERRED** to a Phase 3.5.4.b follow-up. This
// commit lands:
// - Fire/debounce trigger logic (PLASTICITY_FIRED_THIS_FOLD gates
// re-fire within the same fold; reset to 0.0 at fold boundary
// re-arms the next fold).
// - Per-bar warm-bars decrement (every kernel call, regardless of
// fire).
// - The advantage_head_weights pointer + n_weights param wired
// through the launcher signature (kernel accepts but does not
// modify them yet — same `(void)foo` no-op pattern used elsewhere
// for deferred parameter plumbing).
// Wave 4.2 lands the previously-deferred Kaiming-He weight reset:
//
// Why deferred: weight-reset requires Kaiming-He init via cuRAND, and
// accurate targeting of "last 10% of advantage-head weights" requires
// plumbing specific param-buffer pointers (mag/ord/urg head Linear
// layer slices) that aren't available at this kernel's call site
// without trainer-internal access. Establishing the kernel signature +
// flag-tracking now keeps the follow-up commit purely additive on the
// weight-reset side, mirroring the Phase 1.1-1.5 + 3.1-3.5.3 precedent
// of kernel + launcher landing first and consumer wiring following
// atomically.
// - Reset region: the last 10% of `advantage_head_weights`
// (`reset_count = n_weights / 10`, addressed via offset
// `advantage_head_weights + n_weights - reset_count`). Resetting
// the trailing slice keeps the upstream feature-extraction trunk
// intact while re-initialising the head's last layer where the
// drawdown-correlated gradients have presumably saturated. This
// matches the spec §9.2 (3.5.4) "last 10%" prescription.
// - Init formula: Kaiming-He for ReLU/GLU activations,
// `weights[i] ~ Normal(0, sqrt(2 / fan_in))` (gain = sqrt(2)).
// `fan_in` is supplied by the launcher (host computes it from the
// trainer's network config; for oracle tests the test sets it to
// a representative value).
// - Per-thread cuRAND state via `curand_init(seed, tid, 0, &state)`
// + `curand_normal(&state)` × stddev. `seed` is a host-passed
// `unsigned long long`; the launcher derives it deterministically
// from a trainer RNG state (or `SystemTime` nanos in the oracle
// tests). The same (seed, tid) pair always produces the same
// sample, satisfying determinism for the oracle test's mean/std
// assertions.
// - Grid: multi-block, `[((reset_count + 255) / 256), 1, 1]` ×
// `[256, 1, 1]`. Block 0 / thread 0 still runs the trigger +
// warm-bars decrement (the only ISV-write code path); all other
// threads independently re-evaluate the same fire condition from
// the same ISV source-of-truth (ISV state is read at kernel-launch
// time, the trigger condition is a pure function of those reads,
// so every thread computes the same `fire_now` without any
// cross-block synchronisation).
//
// Slot reads/writes (literal indices to keep this kernel self-contained;
// the canonical names live in `crates/ml/src/cuda_pipeline/sp15_isv_slots.rs`
@@ -44,50 +53,60 @@
// indices via compile-time `assert_eq!` regression checks):
// ISV[404] DD_PERSISTENCE (read)
// ISV[437] PLASTICITY_PERSISTENCE_THRESHOLD (read)
// ISV[436] PLASTICITY_FIRED_THIS_FOLD (read-modify-write)
// ISV[438] PLASTICITY_WARM_BARS_REMAINING (read-modify-write)
// ISV[436] PLASTICITY_FIRED_THIS_FOLD (read-modify-write — block 0 thread 0 only)
// ISV[438] PLASTICITY_WARM_BARS_REMAINING (read-modify-write — block 0 thread 0 only)
//
// Hard rules:
//
// * `feedback_no_atomicadd` — single-thread, single-block, pure
// scalar arithmetic; no reductions, no `atomicAdd`.
// * `feedback_no_partial_refactor` — kernel + launcher land first;
// consumer wiring (action-selection two-step Flat+cooldown
// integration via `effective_cooldown =
// max(COOLDOWN_BARS_REMAINING, PLASTICITY_WARM_BARS_REMAINING)`)
// deferred to a follow-up atomic commit per the established Phase
// 1.1-1.5 / 3.1-3.5.3 precedent. Weight-reset is also deferred
// explicitly — the launcher signature accepts the pointer + count
// so the follow-up is purely additive on the kernel body, not on
// the launcher contract.
// * `feedback_no_stubs` — kernel returns the actual updated state
// (flag transition + warm-bars decrement); the deferred weight-
// reset is a no-op on the same call rather than a stubbed
// placeholder return. The three oracle tests verify fire-on-
// threshold-crossing, debounce-within-fold, and below-threshold-
// no-fire behaviors.
// * `feedback_no_atomicadd` — weight-reset is per-thread independent
// (each thread writes a unique element of `advantage_head_weights`);
// no reductions, no `atomicAdd`. The trigger logic remains
// single-thread (block 0 / thread 0) so its ISV writes are
// race-free without atomics.
// * `feedback_no_partial_refactor` — kernel + launcher signature
// change + 4 oracle tests (3 existing migrated + 1 new) land
// atomically in this commit (Wave 4.2); consumer wiring (action
// selection forces Hold via `effective_cooldown =
// max(COOLDOWN_BARS_REMAINING, PLASTICITY_WARM_BARS_REMAINING)`
// remains the separate follow-up tracked under Phase 3.5.4.c.
// * `feedback_no_stubs` — kernel performs the actual Kaiming-He
// reset; no `(void)` casts remain. Wave 4.2 closes the Phase
// 3.5.4-deferred weight-reset NO-OP per
// `feedback_wire_everything_up.md`.
// * `feedback_isv_for_adaptive_bounds` — PLASTICITY_PERSISTENCE_
// THRESHOLD is ISV-tracked (slot 437); the producer driving it
// from the running mean of dd_persistence_bars is deferred but
// the kernel reads it rather than a hardcoded literal so the
// follow-up swap is a no-op at the consumer.
//
// Single-thread, single-block (mirrors `cooldown_kernel` /
// `dd_asymmetric_reward_kernel` / `hold_floor_kernel` /
// `regret_signal_kernel` / `dd_penalty_kernel` — no reduction).
// * `feedback_no_cpu_compute_strict` — cuRAND init + sampling are
// GPU-resident (`curand_init` / `curand_normal` are device-only
// functions inlined by nvcc into the cubin from
// `<curand_kernel.h>`).
// * `pearl_no_host_branches_in_captured_graph` — `seed` is a kernel
// argument (a `u64` value, not a host-side branch); the launcher
// itself can be invoked from inside a graph capture region
// because the kernel-launch shape (grid/block) is fixed at the
// call (a function of `n_weights / 10` which is also a launch
// argument, not a host-conditional).
#include <curand_kernel.h>
extern "C" __global__ void plasticity_injection_kernel(
float* __restrict__ isv,
float* __restrict__ advantage_head_weights, /* deferred: weight-reset NO-OP for now */
float* __restrict__ advantage_head_weights,
int n_weights,
float m_warm /* default 200 bars */
float m_warm, /* default 200 bars */
int fan_in, /* Kaiming gain denominator */
unsigned long long seed /* cuRAND seed */
) {
if (threadIdx.x != 0 || blockIdx.x != 0) return;
// Every thread independently reads the ISV state to determine
// `fire_now`. The trigger condition is a pure function of those
// reads, so all threads converge on the same answer without any
// cross-block synchronisation. The same reads also drive block 0
// thread 0's ISV-write path below.
const float persistence = isv[/*DD_PERSISTENCE_INDEX*/ 404];
const float threshold = isv[/*PLASTICITY_PERSISTENCE_THRESHOLD_INDEX*/ 437];
const float fired = isv[/*PLASTICITY_FIRED_THIS_FOLD_INDEX*/ 436];
float warm = isv[/*PLASTICITY_WARM_BARS_REMAINING_INDEX*/ 438];
// Trigger: persistence exceeds threshold AND not yet fired this
// fold. PLASTICITY_FIRED_THIS_FOLD is reset to 0.0 at fold boundary
@@ -95,31 +114,58 @@ extern "C" __global__ void plasticity_injection_kernel(
// re-arming the gate for the next fold. The `< 0.5f` half-bit
// tolerance check matches the existing `cooldown_remaining <= 0.0f`
// float-equality-with-tolerance pattern in `cooldown_kernel`.
if (persistence > threshold && fired < 0.5f) {
isv[/*PLASTICITY_FIRED_THIS_FOLD_INDEX*/ 436] = 1.0f;
warm = m_warm;
const bool fire_now = (persistence > threshold) && (fired < 0.5f);
// Weight-reset DEFERRED to Phase 3.5.4.b — kernel param plumbed
// through but no-op. Follow-up: Kaiming-He reset of last 10% of
// advantage-head weights via cuRAND with gain=√2 for ReLU/GLU
// activations. The `(void)foo` cast pattern matches the Rust
// `_ = foo` convention used elsewhere in this codebase for
// deferred parameter plumbing.
(void)advantage_head_weights;
(void)n_weights;
// Block 0 / thread 0 owns the ISV writes (fired flag + warm-bars
// counter). Single-thread to mirror the cooldown_kernel /
// dd_asymmetric_reward_kernel / hold_floor_kernel / regret_signal
// / dd_penalty_kernel pattern — no reduction, no atomic.
if (threadIdx.x == 0 && blockIdx.x == 0) {
float warm = isv[/*PLASTICITY_WARM_BARS_REMAINING_INDEX*/ 438];
if (fire_now) {
isv[/*PLASTICITY_FIRED_THIS_FOLD_INDEX*/ 436] = 1.0f;
warm = m_warm;
}
// Per-bar decrement of warm-bars (every call, regardless of
// fire). The trigger above runs first so the same call that
// fires the gate also immediately consumes one bar of warm-up
// — bias downward by 1 is acceptable per the test's [198, 200]
// tolerance and avoids the bookkeeping branch needed to skip
// the decrement on the trip step (mirrors `cooldown_kernel`'s
// identical trigger-then-decrement sequence).
if (warm > 0.0f) {
warm -= 1.0f;
}
isv[/*PLASTICITY_WARM_BARS_REMAINING_INDEX*/ 438] = warm;
__threadfence_system();
}
// Per-bar decrement of warm-bars (every call, regardless of fire).
// The trigger above runs first so the same call that fires the gate
// also immediately consumes one bar of warm-up — bias downward by 1
// is acceptable per the test's [198, 200] tolerance and avoids the
// bookkeeping branch needed to skip the decrement on the trip step
// (mirrors `cooldown_kernel`'s identical trigger-then-decrement
// sequence).
if (warm > 0.0f) {
warm -= 1.0f;
}
isv[/*PLASTICITY_WARM_BARS_REMAINING_INDEX*/ 438] = warm;
// Weight reset: last 10% of advantage-head weights via Kaiming-He
// sampling. Each thread writes a unique element so no atomic /
// reduction is required. Threads beyond `reset_count` early-out;
// when `fire_now == false` every thread early-outs and the kernel
// is a pure trigger/decrement (the previous Phase 3.5.4 behaviour).
if (!fire_now) return;
__threadfence_system();
const int reset_count = n_weights / 10;
if (reset_count <= 0) return;
const int reset_start = n_weights - reset_count;
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= reset_count) return;
// Per-thread cuRAND state: (seed, sequence=tid, offset=0). Same
// (seed, tid) pair always produces the same sample, satisfying
// determinism for oracle-test mean/std assertions. `fan_in` is the
// input dimension of the advantage head's last Linear layer; for
// ReLU/GLU activations the Kaiming-He gain is sqrt(2).
curandState state;
curand_init(seed, (unsigned long long)tid, 0ULL, &state);
const float gain = sqrtf(2.0f);
const float stddev = gain / sqrtf((float)fan_in);
advantage_head_weights[reset_start + tid] = curand_normal(&state) * stddev;
}

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

@@ -1141,10 +1141,12 @@ mod gpu {
/// reads the fired flag transition AND emits Flat that bar, then
/// forces Hold while
/// `max(COOLDOWN_BARS_REMAINING, PLASTICITY_WARM_BARS_REMAINING) > 0`.
/// Weight-reset deferred to Phase 3.5.4.b — kernel param plumbed
/// through but no-op for now (the test passes a non-zero
/// advantage_head_weights buffer to verify the launcher accepts it
/// without modifying the contents).
/// Wave 4.2 / Phase 3.5.4.b: kernel now also performs the
/// Kaiming-He weight-reset; this trigger-only test passes the
/// extra `fan_in` + `seed` launcher params and asserts the
/// trigger behaviour (fired flag + warm-bars). The dedicated
/// `plasticity_injection_kernel_resets_last_10pct_kaiming_he`
/// test covers the weight-reset numerics.
#[test]
#[ignore = "requires GPU"]
fn plasticity_fires_when_persistence_exceeds_threshold() {
@@ -1164,9 +1166,10 @@ mod gpu {
isv_init[PLASTICITY_WARM_BARS_REMAINING_INDEX] = 0.0;
isv_buf.write_from_slice(&isv_init);
// Empty advantage-head weights buffer for now (weight-reset
// deferred to Phase 3.5.4.b — kernel takes the pointer but
// does not modify contents).
// Wave 4.2 weight-reset is exercised in detail by the
// dedicated Kaiming-He oracle below; here we just provide a
// 64-element buffer so the launcher can run (n_weights / 10 =
// 6 trailing slots get re-sampled).
let weights_buf = unsafe { MappedF32Buffer::new(64) }
.expect("alloc MappedF32Buffer for advantage_head_weights");
weights_buf.write_from_slice(&vec![1.0f32; 64]);
@@ -1177,6 +1180,8 @@ mod gpu {
weights_buf.dev_ptr,
/*n_weights*/ 64,
/*m_warm*/ 200.0,
/*fan_in*/ 256,
/*seed*/ 0xCAFE_BABE_DEAD_BEEF_u64,
)
.expect("launch plasticity_injection_kernel (fires)");
stream
@@ -1238,6 +1243,8 @@ mod gpu {
weights_buf.dev_ptr,
64,
200.0,
/*fan_in*/ 256,
/*seed*/ 0xDEAD_BEEF_F00D_CAFE_u64,
)
.expect("launch plasticity_injection_kernel (debounced)");
stream
@@ -1256,6 +1263,17 @@ mod gpu {
(warm - 49.0).abs() < 1e-9,
"PLASTICITY_WARM_BARS_REMAINING = {warm}, expected 49 (decremented from 50, no re-fire)"
);
// Wave 4.2: when debounced, the kernel returns early before
// the weight-reset region, so all 64 weights must remain
// exactly 1.0 (no Kaiming-He sample written).
let weights_host = weights_buf.read_all();
for (i, w) in weights_host.iter().enumerate() {
assert!(
(w - 1.0).abs() < 1e-9,
"weights[{i}] = {w}, expected 1.0 (debounced — kernel must skip weight-reset)"
);
}
}
/// SP15 Phase 3.5.4 — plasticity injection does NOT fire when
@@ -1293,6 +1311,8 @@ mod gpu {
weights_buf.dev_ptr,
64,
200.0,
/*fan_in*/ 256,
/*seed*/ 0xFEED_FACE_BAAD_F00D_u64,
)
.expect("launch plasticity_injection_kernel (no-fire)");
stream
@@ -1310,6 +1330,189 @@ mod gpu {
"PLASTICITY_WARM_BARS_REMAINING = {}, expected 0.0 (no fire → no warm-up set; the `if (warm > 0)` guard prevents underflow when already 0)",
isv_host[PLASTICITY_WARM_BARS_REMAINING_INDEX]
);
// Wave 4.2: when no-fire, the kernel returns early before the
// weight-reset region, so all 64 weights must remain exactly
// 1.0 (no Kaiming-He sample written).
let weights_host = weights_buf.read_all();
for (i, w) in weights_host.iter().enumerate() {
assert!(
(w - 1.0).abs() < 1e-9,
"weights[{i}] = {w}, expected 1.0 (no-fire — kernel must skip weight-reset)"
);
}
}
/// SP15 Wave 4.2 / Phase 3.5.4.b — when the plasticity injection
/// gate fires, the trailing 10% of `advantage_head_weights` must
/// be re-initialised to Kaiming-He samples
/// (`Normal(0, sqrt(2/fan_in))`, gain=√2 for ReLU/GLU
/// activations) drawn via cuRAND `curand_normal()`. The first 90%
/// of the buffer must remain bit-identical to its pre-launch
/// value (1.0) — the kernel writes element `[reset_start + tid]`
/// only.
///
/// Distribution checks (with `n_weights = 1280` → `reset_count =
/// 128`, `fan_in = 256`, expected `stddev = sqrt(2/256) ≈
/// 0.08839`):
/// - Per-element diff from 1.0: every reset slot must have
/// changed (probability of an exact-1.0 sample is essentially
/// zero — Kaiming stddev ≈ 0.088 means 1.0 is ~11.3 σ out, and
/// anyway the kernel writes the *sample*, not the original).
/// - Mean of reset region: |mean| < 0.05 (sampling-noise
/// tolerance for n=128 from a Normal with stddev≈0.088 — the
/// true SE of the mean is `0.088 / sqrt(128) ≈ 0.0078`, so
/// `0.05` is a comfortable ~6σ envelope around 0).
/// - Std of reset region: within ±20% of `sqrt(2/fan_in)`
/// (n=128 sample std has ~6% relative SE; ±20% is comfortable
/// headroom).
/// - ISV[PLASTICITY_FIRED_THIS_FOLD_INDEX] flipped 0 → 1.
/// - ISV[PLASTICITY_WARM_BARS_REMAINING_INDEX] = m_warm − 1
/// (per-bar decrement runs in the same launch as the trigger
/// per the cooldown_kernel trigger-then-decrement convention).
///
/// Determinism: a fixed seed (`42`) drives `curand_init(seed,
/// tid, 0, &state)` per thread; re-running this test always
/// produces identical samples.
#[test]
#[ignore = "requires GPU"]
fn plasticity_injection_kernel_resets_last_10pct_kaiming_he() {
use ml::cuda_pipeline::sp15_isv_slots::{
DD_PERSISTENCE_INDEX, PLASTICITY_FIRED_THIS_FOLD_INDEX,
PLASTICITY_PERSISTENCE_THRESHOLD_INDEX, PLASTICITY_WARM_BARS_REMAINING_INDEX,
};
let stream = make_test_stream();
// ISV: ensure the trigger fires — large persistence, small
// threshold, fired flag clear, warm-bars 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[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);
// Weight buffer: n_weights = 1280, reset_count = 128 (10%).
// Pre-fill every slot with 1.0 so we can detect both
// "unchanged" (first 90%) and "changed" (last 10%) by exact
// float comparison + sample-statistic checks.
const N_WEIGHTS: usize = 1280;
const FAN_IN: i32 = 256;
const SEED: u64 = 42;
const M_WARM: f32 = 200.0;
let weights_buf = unsafe { MappedF32Buffer::new(N_WEIGHTS) }
.expect("alloc MappedF32Buffer for advantage_head_weights");
weights_buf.write_from_slice(&vec![1.0f32; N_WEIGHTS]);
ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_plasticity_injection(
&stream,
isv_buf.dev_ptr,
weights_buf.dev_ptr,
N_WEIGHTS as i32,
M_WARM,
FAN_IN,
SEED,
)
.expect("launch plasticity_injection_kernel (kaiming-he)");
stream
.synchronize()
.expect("synchronize after plasticity_injection_kernel launch (kaiming-he)");
// ISV side-effects: fired flipped to 1, warm = M_warm - 1.
let isv_host = isv_buf.read_all();
assert!(
(isv_host[PLASTICITY_FIRED_THIS_FOLD_INDEX] - 1.0).abs() < 1e-9,
"PLASTICITY_FIRED_THIS_FOLD = {}, expected 1.0 after fire",
isv_host[PLASTICITY_FIRED_THIS_FOLD_INDEX]
);
let warm = isv_host[PLASTICITY_WARM_BARS_REMAINING_INDEX];
assert!(
(warm - (M_WARM - 1.0)).abs() < 1e-6,
"PLASTICITY_WARM_BARS_REMAINING = {warm}, expected {} (M_warm - 1, per-bar decrement runs same call)",
M_WARM - 1.0
);
// Weight-reset side-effects.
let weights_host = weights_buf.read_all();
let reset_count: usize = N_WEIGHTS / 10; // 128
let reset_start: usize = N_WEIGHTS - reset_count; // 1152
// First 90% must be bit-identical to the pre-launch value.
for (i, w) in weights_host[..reset_start].iter().enumerate() {
assert!(
(w - 1.0).abs() < 1e-9,
"weights[{i}] = {w}, expected 1.0 (outside reset region — kernel must not touch)"
);
}
// Last 10%: every slot must have moved off 1.0 (a Kaiming
// stddev≈0.088 makes the sentinel 1.0 ~11σ out, so an exact
// 1.0 sample is effectively zero-probability).
let resets = &weights_host[reset_start..];
for (i, w) in resets.iter().enumerate() {
assert!(
(w - 1.0).abs() > 1e-6,
"weights[{}] = {w}, expected resampled value ≠ 1.0",
reset_start + i
);
}
// Mean of resets ≈ 0 within sampling-noise envelope.
let n = resets.len() as f64;
let mean: f64 = resets.iter().map(|&w| w as f64).sum::<f64>() / n;
assert!(
mean.abs() < 0.05,
"Kaiming-reset region mean = {mean:.6}, expected |mean| < 0.05 (n={n}, true SE ~ 0.0078)"
);
// Std of resets ≈ sqrt(2/fan_in) within ±20%.
let expected_std = (2.0_f64 / FAN_IN as f64).sqrt();
let var: f64 = resets
.iter()
.map(|&w| (w as f64 - mean).powi(2))
.sum::<f64>()
/ (n - 1.0);
let std = var.sqrt();
assert!(
(std - expected_std).abs() / expected_std < 0.20,
"Kaiming-reset region std = {std:.6}, expected {expected_std:.6} ± 20%"
);
// Determinism re-check: same (seed, tid) on a fresh buffer
// produces bit-identical samples (catches accidental
// entropy sources like clock-derived seeds inside the
// kernel). Re-launch with the same seed against a fresh
// 1.0-filled buffer — but note ISV state must also be reset
// (fired=0) so the trigger re-fires.
weights_buf.write_from_slice(&vec![1.0f32; N_WEIGHTS]);
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);
ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_plasticity_injection(
&stream,
isv_buf.dev_ptr,
weights_buf.dev_ptr,
N_WEIGHTS as i32,
M_WARM,
FAN_IN,
SEED, // same seed
)
.expect("launch plasticity_injection_kernel (determinism re-check)");
stream
.synchronize()
.expect("synchronize after plasticity_injection_kernel launch (determinism)");
let resets_2 = weights_buf.read_all();
for i in 0..reset_count {
let a = resets[i];
let b = resets_2[reset_start + i];
assert!(
(a - b).abs() < 1e-9,
"non-deterministic Kaiming sample at tid={i}: first={a}, second={b}"
);
}
}
/// SP15 Phase 3.5.5 — DD_TRAJECTORY_DECREASING fires when dd_pct

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

File diff suppressed because one or more lines are too long