feat(dqn-v2): Plan 4 Task 2c.3c.5 — append H_S2_RMS_EMA ISV slot + producer kernel
ISV bus tail-append: - [96] H_S2_RMS_EMA_INDEX — per-batch RMS(save_h_s2) EMA (alpha=0.05) - [97] ISV_LAYOUT_FINGERPRINT_LO_INDEX (shifted from 94) - [98] ISV_LAYOUT_FINGERPRINT_HI_INDEX (shifted from 95) - ISV_TOTAL_DIM 96 -> 99 - New LAYOUT_FINGERPRINT_CURRENT = 0x3e21acecd922e540 (was 0xcf3a24b0a1f70057) Producer kernel `h_s2_rms_ema_kernel.cu`: - Single-block reduction (256 threads, shmem-tree, no atomicAdd) of the online trunk's `save_h_s2 [B, SH2]` post-GRN activation - RMS = sqrt(sum_sq / (B*SH2)); EMA into ISV[96] with alpha=0.05 - Launched once per training step alongside reward_component_ema and the other Plan 3/4 ISV producers in training_loop.rs Constructor cold-start writes ISV[96]=1.0 (neutral RMS). StateResetRegistry: H_S2_RMS_EMA registered as FoldReset -> 1.0. `launch_h_s2_rms_ema` and the kernel both treat the unconditional constructor allocation as an invariant — debug_assert! on the host side, no defensive (N>0) ternary in the kernel — same proper-resolution pattern as the 2c.3c.4 followup. Producer-only commit. 2c.3c.6 wires the consumer in mag_concat_qdir's adaptive-scale path so the magnitude-branch decoder sees a scale- invariant residual stack regardless of GRN drift across training. Smoke deferred — producer-only with zero consumers, runtime behaviour structurally unchanged from 2c.3c.4 (which validated the GRN backward chain end-to-end with all 3 fold checkpoints). cargo check + cargo build (new h_s2_rms_ema_kernel.cubin compiles clean under nvcc sm_80) is the standing validation; meaningful multi_fold_convergence re-run lands with 2c.3c.6 when the consumer wires up. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -131,6 +131,12 @@ fn main() {
|
||||
// Module is additive — no production callers in this commit; Task 2c.3+4
|
||||
// wires it into the trunk encoder.
|
||||
"grn_kernel.cu",
|
||||
// Plan 4 Task 2c.3c.5: per-batch RMS EMA of `save_h_s2` into ISV[96].
|
||||
// Single-block shmem-reduce kernel (256 threads, no atomicAdd) launched
|
||||
// alongside `reward_component_ema` from `training_loop.rs`. Producer-only
|
||||
// in this commit — 2c.3c.6 wires the consumer in `mag_concat_qdir`'s
|
||||
// adaptive-scale path.
|
||||
"h_s2_rms_ema_kernel.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
@@ -126,6 +126,12 @@ pub(crate) static TARGET_DRIFT_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_D
|
||||
/// production callers added in 2c.3+4.
|
||||
pub(crate) static GRN_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grn_kernel.cubin"));
|
||||
|
||||
/// Plan 4 Task 2c.3c.5: per-batch RMS EMA of `save_h_s2` into ISV[96].
|
||||
/// Single-block (256 threads, shmem-reduce, no atomicAdd) kernel launched
|
||||
/// alongside `reward_component_ema` from `training_loop.rs`. Producer-only —
|
||||
/// 2c.3c.6 wires the consumer in `mag_concat_qdir`'s adaptive-scale path.
|
||||
static H_S2_RMS_EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/h_s2_rms_ema_kernel.cubin"));
|
||||
|
||||
/// Mamba2 temporal scan configuration.
|
||||
const MAMBA2_HISTORY_K: usize = 8; // Rolling history length
|
||||
const MAMBA2_STATE_DIM: usize = 16; // SSM state dimension
|
||||
@@ -289,11 +295,14 @@ const ISV_NETWORK_DIM: usize = 23;
|
||||
/// replaces legacy CPU-DtoH `per_branch_target_drift()`. Written by
|
||||
/// `target_drift_ema_update` GPU kernel. Magnitude branch.
|
||||
/// [93] TARGET_DRIFT_DIR_EMA_INDEX — same kernel/contract; direction branch.
|
||||
/// [94] ISV_LAYOUT_FINGERPRINT_LO_INDEX — low 32 bits of layout fingerprint (shifted from 90).
|
||||
/// [95] ISV_LAYOUT_FINGERPRINT_HI_INDEX — high 32 bits of layout fingerprint (shifted from 91).
|
||||
/// [96] H_S2_RMS_EMA_INDEX — EMA of RMS(save_h_s2) per batch. Producer-only
|
||||
/// in this commit (no consumers). Wired in 2c.3c.6 to drive the
|
||||
/// adaptive scale of `mag_concat_qdir`'s residual stack.
|
||||
/// [97] ISV_LAYOUT_FINGERPRINT_LO_INDEX — low 32 bits of layout fingerprint (shifted from 94).
|
||||
/// [98] ISV_LAYOUT_FINGERPRINT_HI_INDEX — high 32 bits of layout fingerprint (shifted from 95).
|
||||
/// Written by the constructor; checked at checkpoint load. Fail-fast only — no migration
|
||||
/// path exists. See spec §4.A.2 and `LAYOUT_FINGERPRINT_CURRENT` for structural-hash rationale.
|
||||
const ISV_TOTAL_DIM: usize = 96;
|
||||
const ISV_TOTAL_DIM: usize = 99;
|
||||
/// Legacy alias preserved for call sites that haven't been audited for the
|
||||
/// network-vs-total split. New code should pick `ISV_NETWORK_DIM` (for weight
|
||||
/// tensor sizing) or `ISV_TOTAL_DIM` (for the broadcast bus buffer).
|
||||
@@ -547,7 +556,21 @@ pub const TARGET_DRIFT_MAG_EMA_INDEX: usize = 92;
|
||||
/// ISV slot [93] — same producer/contract as [92]; direction branch.
|
||||
pub const TARGET_DRIFT_DIR_EMA_INDEX: usize = 93;
|
||||
|
||||
/// ISV slot [94] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
|
||||
/// ISV slot [96] — Plan 4 Task 2c.3c.5: EMA of RMS(save_h_s2) per batch.
|
||||
/// Producer: `h_s2_rms_ema_update` GPU kernel (single-block, 256 threads,
|
||||
/// shmem-reduction, no atomicAdd). Cold-path-cadence — launched once per
|
||||
/// experience-collection epoch alongside the other Plan 3/Plan 4 ISV producers
|
||||
/// in `training_loop.rs`. EMA α = 0.05 (≈ 13-batch half-life). Cold-start
|
||||
/// 1.0 (neutral RMS so first kernel fire EMAs measured RMS toward 1.0
|
||||
/// rather than 0). FoldReset → 1.0.
|
||||
///
|
||||
/// Producer-only in this commit. 2c.3c.6 wires the consumer in
|
||||
/// `mag_concat_qdir`'s adaptive-scale path (residual stack normalised by
|
||||
/// the trunk-output RMS so the magnitude-branch decoder sees a scale-
|
||||
/// invariant input regardless of the GRN trunk's drift).
|
||||
pub const H_S2_RMS_EMA_INDEX: usize = 96;
|
||||
|
||||
/// ISV slot [97] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
|
||||
///
|
||||
/// Design note: this is NOT a version number. There is no ordered version space.
|
||||
/// The fingerprint is a structural hash that automatically changes whenever any
|
||||
@@ -560,14 +583,15 @@ pub const TARGET_DRIFT_DIR_EMA_INDEX: usize = 93;
|
||||
/// actively written by `isv_signal_update` (Q-drift and gradient-norm EMA).
|
||||
/// Shifted 61→69 in Plan 3 Task 1, 69→73 in Plan 3 Task 3, 73→76 in
|
||||
/// Plan 3 Task 4 B.4, 76→80 in Plan 3 Task 7 C.3, 80→85 in Plan 3 Task 8 B.3,
|
||||
/// 85→90 in Plan 4 Task 5 Mode A, then 90→94 in Plan 4 follow-up (target-drift
|
||||
/// EMA replacing legacy CPU-DtoH path).
|
||||
pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 94;
|
||||
/// ISV slot [95] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
|
||||
/// 85→90 in Plan 4 Task 5 Mode A, 90→94 in Plan 4 follow-up (target-drift
|
||||
/// EMA replacing legacy CPU-DtoH path), then 94→97 in Plan 4 Task 2c.3c.5
|
||||
/// (H_S2_RMS_EMA producer slot).
|
||||
pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 97;
|
||||
/// ISV slot [98] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
|
||||
/// Shifted 62→70 in Plan 3 Task 1, 70→74 in Plan 3 Task 3, 74→77 in Plan 3 Task 4 B.4,
|
||||
/// 77→81 in Plan 3 Task 7 C.3, 81→86 in Plan 3 Task 8 B.3, 86→91 in Plan 4 Task 5,
|
||||
/// then 91→95 in Plan 4 follow-up.
|
||||
pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 95;
|
||||
/// 91→95 in Plan 4 follow-up, then 95→98 in Plan 4 Task 2c.3c.5.
|
||||
pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 98;
|
||||
/// Canonical alias for the fingerprint slot (the low half).
|
||||
pub const ISV_LAYOUT_FINGERPRINT_INDEX: usize = ISV_LAYOUT_FINGERPRINT_LO_INDEX;
|
||||
|
||||
@@ -645,8 +669,9 @@ const fn layout_fingerprint_seed() -> &'static [u8] {
|
||||
SEED_STEPS_TARGET=82;SEED_STEPS_DONE=83;SEED_FRAC_EMA=84;\
|
||||
VSN_MAG_EMA=87;VSN_DIR_EMA=88;MAMBA2_RETENTION_EMA=89;\
|
||||
TARGET_DRIFT_MAG_EMA=92;TARGET_DRIFT_DIR_EMA=93;\
|
||||
ISV_LAYOUT_FINGERPRINT_LO=94;ISV_LAYOUT_FINGERPRINT_HI=95;\
|
||||
ISV_TOTAL_DIM=96;\
|
||||
H_S2_RMS_EMA=96;\
|
||||
ISV_LAYOUT_FINGERPRINT_LO=97;ISV_LAYOUT_FINGERPRINT_HI=98;\
|
||||
ISV_TOTAL_DIM=99;\
|
||||
PARAM_W_A_H_S1=0;PARAM_B_A_H_S1=1;PARAM_W_B_H_S1=2;PARAM_B_B_H_S1=3;\
|
||||
PARAM_W_RESIDUAL_H_S1=4;PARAM_GAMMA_H_S1=5;PARAM_BETA_H_S1=6;\
|
||||
PARAM_W_A_H_S2=7;PARAM_B_A_H_S2=8;PARAM_W_B_H_S2=9;PARAM_B_B_H_S2=10;\
|
||||
@@ -2234,6 +2259,15 @@ pub struct GpuDqnTrainer {
|
||||
/// Cold-path: one launch per training step. Loaded from reward_component_ema_kernel.cubin.
|
||||
reward_component_ema_kernel: CudaFunction,
|
||||
|
||||
// ── Plan 4 Task 2c.3c.5: H_S2 RMS EMA producer ────────────────────
|
||||
/// Single-block kernel (256 threads, shmem-reduction, no atomicAdd). Reduces
|
||||
/// RMS = sqrt(sum_sq / (B*SH2)) over `save_h_s2 [B, SH2]` and EMA-updates
|
||||
/// ISV[H_S2_RMS_EMA_INDEX=96] with α = 0.05. Cold-path-cadence (one launch
|
||||
/// per training step, alongside `reward_component_ema`). Producer-only —
|
||||
/// 2c.3c.6 wires the consumer in `mag_concat_qdir`'s adaptive-scale path.
|
||||
/// Loaded from `h_s2_rms_ema_kernel.cubin`.
|
||||
h_s2_rms_ema_kernel: CudaFunction,
|
||||
|
||||
// ── Q-mean drift correction (Component 8) ──────────────────────────
|
||||
/// Phase 1 kernel: computes global mean of Q-values into q_mean_scratch (pinned).
|
||||
q_mean_reduce_kernel: CudaFunction,
|
||||
@@ -6674,6 +6708,49 @@ impl GpuDqnTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Plan 4 Task 2c.3c.5: launch `h_s2_rms_ema_update` (single-block,
|
||||
/// 256-thread shmem-reduction kernel). Reduces RMS = sqrt(sum_sq / (B*SH2))
|
||||
/// over the trainer's `save_h_s2 [B, SH2]` buffer (the online trunk's
|
||||
/// post-GRN activation, valid after the most recent forward) and EMAs the
|
||||
/// result into ISV[H_S2_RMS_EMA_INDEX=96].
|
||||
///
|
||||
/// Producer-only — 2c.3c.6 wires the consumer in `mag_concat_qdir`'s
|
||||
/// adaptive-scale path. Cold-path-cadence: launched once per training step
|
||||
/// alongside `launch_reward_component_ema`. No atomicAdd, no DtoH.
|
||||
pub fn launch_h_s2_rms_ema(&self, ema_alpha: f32) -> Result<(), MLError> {
|
||||
// `isv_signals_dev_ptr` is set by the constructor via
|
||||
// `cuMemHostGetDevicePointer_v2` on a `cuMemAllocHost_v2`-backed pinned
|
||||
// buffer (with its own `assert_eq!` on success) and is never reassigned,
|
||||
// so non-zero is an invariant. Silent-skip on `== 0` would mask a
|
||||
// misconfiguration; `debug_assert!` flags it in dev builds and release
|
||||
// inherits the loud kernel-launch error rather than a no-op.
|
||||
debug_assert!(self.isv_signals_dev_ptr != 0,
|
||||
"launch_h_s2_rms_ema: isv_signals_dev_ptr must be allocated by constructor");
|
||||
let h_s2_ptr = self.save_h_s2.raw_ptr();
|
||||
let b_i = self.config.batch_size as i32;
|
||||
let sh2_i = self.config.shared_h2 as i32;
|
||||
let isv_idx = H_S2_RMS_EMA_INDEX as i32;
|
||||
let isv_dev_ptr = self.isv_signals_dev_ptr;
|
||||
const BLOCK_DIM: u32 = 256;
|
||||
let smem_bytes = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.h_s2_rms_ema_kernel)
|
||||
.arg(&h_s2_ptr)
|
||||
.arg(&b_i)
|
||||
.arg(&sh2_i)
|
||||
.arg(&isv_dev_ptr)
|
||||
.arg(&isv_idx)
|
||||
.arg(&ema_alpha)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (BLOCK_DIM, 1, 1),
|
||||
shared_mem_bytes: smem_bytes,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("h_s2_rms_ema_update: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Apply spectral normalization to all 10 weight matrices (trunk + 8 heads).
|
||||
///
|
||||
/// One step of power iteration per call (standard practice — single step
|
||||
@@ -7345,6 +7422,16 @@ impl GpuDqnTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("reward_component_ema load: {e}")))?
|
||||
};
|
||||
|
||||
// Plan 4 Task 2c.3c.5: load h_s2_rms_ema kernel (cold-path, per-step).
|
||||
// Single-block 256-thread shmem-reduction kernel; no atomicAdd. Producer
|
||||
// for ISV[H_S2_RMS_EMA_INDEX=96]; consumer wired in 2c.3c.6.
|
||||
let h_s2_rms_ema_kernel = {
|
||||
let module = stream.context().load_cubin(H_S2_RMS_EMA_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("h_s2_rms_ema cubin load: {e}")))?;
|
||||
module.load_function("h_s2_rms_ema_update")
|
||||
.map_err(|e| MLError::ModelError(format!("h_s2_rms_ema_update load: {e}")))?
|
||||
};
|
||||
|
||||
// Plan 2 Task 1 C.1: allocate branch-action offset/size device buffers for q_quantile_reduce.
|
||||
// Branch layout in q_out_buf: [dir | mag | ord | urg] row-major per sample.
|
||||
// dir occupies actions [0..b0), mag [b0..b0+b1), ord [b0+b1..b0+b1+b2),
|
||||
@@ -8699,6 +8786,13 @@ impl GpuDqnTrainer {
|
||||
// Plan 4 follow-up: target-drift EMA cold-start.
|
||||
*sig_ptr.add(TARGET_DRIFT_MAG_EMA_INDEX) = 0.0_f32;
|
||||
*sig_ptr.add(TARGET_DRIFT_DIR_EMA_INDEX) = 0.0_f32;
|
||||
/* Plan 4 Task 2c.3c.5: H_S2 RMS EMA cold-start. Set to 1.0
|
||||
* (neutral RMS) so the first kernel fire's α=0.05 EMA pulls
|
||||
* the slot toward the measured RMS rather than collapsing
|
||||
* toward 0. Half-life ≈ 13 batches (log(0.5)/log(0.95)).
|
||||
* Producer-only in this commit; 2c.3c.6 wires the consumer
|
||||
* in `mag_concat_qdir`'s adaptive-scale path. FoldReset → 1.0. */
|
||||
*sig_ptr.add(H_S2_RMS_EMA_INDEX) = 1.0_f32;
|
||||
// Plan 2 C.1 Q-quantile bootstrap (2026-04-24):
|
||||
// q_p05 = v_min, q_p95 = v_max — matches the cold-start atom range
|
||||
// so the first update_eval_v_range call reads meaningful bounds.
|
||||
@@ -9714,6 +9808,7 @@ impl GpuDqnTrainer {
|
||||
popart_var,
|
||||
popart_count,
|
||||
reward_component_ema_kernel,
|
||||
h_s2_rms_ema_kernel,
|
||||
q_mean_reduce_kernel,
|
||||
q_mean_subtract_kernel,
|
||||
q_mean_scratch_pinned,
|
||||
|
||||
69
crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu
Normal file
69
crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu
Normal file
@@ -0,0 +1,69 @@
|
||||
/* h_s2_rms_ema — GPU-driven RMS(save_h_s2) EMA into ISV[96].
|
||||
*
|
||||
* Plan 4 Task 2c.3c.5. GPU-drives-CPU-reads (spec §4.C.6) and
|
||||
* pearl_cold_path_no_exception_to_gpu_drives.md compliant.
|
||||
*
|
||||
* Reduces RMS = sqrt(sum_sq / N) over the flattened `save_h_s2 [B, SH2]`
|
||||
* buffer (the online trunk's post-GRN activation) and EMA-updates
|
||||
* ISV[H_S2_RMS_EMA_INDEX] with rate `ema_alpha`.
|
||||
*
|
||||
* Single-block kernel (256 threads, shmem-reduction). No atomicAdd
|
||||
* (per `feedback_no_atomicadd.md`); the in-block tree reduction is
|
||||
* standard `s >>= 1` with `__syncthreads()`. Only thread 0 writes the
|
||||
* EMA result.
|
||||
*
|
||||
* Producer-only — no consumer reads ISV[96] in this commit. 2c.3c.6
|
||||
* wires the consumer in `mag_concat_qdir`'s adaptive-scale path so the
|
||||
* residual stack is normalised by the trunk-output RMS regardless of
|
||||
* GRN drift across training.
|
||||
*
|
||||
* Cost: one launch per training step (cold-path cadence — same site as
|
||||
* `reward_component_ema`). 256 threads × ceil(B*SH2/256) strided loads
|
||||
* + log2(256)=8 reduce steps + 1 global write.
|
||||
*/
|
||||
|
||||
extern "C" __global__ void h_s2_rms_ema_update(
|
||||
const float* __restrict__ h_s2, /* [B, SH2] row-major */
|
||||
int B,
|
||||
int SH2,
|
||||
float* __restrict__ isv, /* ISV bus [ISV_TOTAL_DIM] */
|
||||
int isv_h_s2_rms_ema_index, /* H_S2_RMS_EMA_INDEX = 96 */
|
||||
float ema_alpha /* per-batch EMA rate (0,1) */
|
||||
) {
|
||||
/* Single-block contract — match reward_component_ema_kernel's grid guard. */
|
||||
if (blockIdx.x != 0) return;
|
||||
|
||||
extern __shared__ float smem[];
|
||||
const int tid = (int)threadIdx.x;
|
||||
const int block = (int)blockDim.x;
|
||||
const int N = B * SH2;
|
||||
|
||||
/* Per-thread sum of squares over a strided slice. */
|
||||
float local_sumsq = 0.0f;
|
||||
for (int i = tid; i < N; i += block) {
|
||||
const float v = h_s2[i];
|
||||
local_sumsq += v * v;
|
||||
}
|
||||
smem[tid] = local_sumsq;
|
||||
__syncthreads();
|
||||
|
||||
/* In-block tree reduction (no atomicAdd). */
|
||||
for (int s = block / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) smem[tid] += smem[tid + s];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
/* Thread 0 writes the EMA. The host caller's trainer config requires
|
||||
* batch_size > 0 and shared_h2 > 0, so N = B*SH2 > 0 is invariant.
|
||||
* No defensive `(N > 0) ? N : 1` ternary here — masking the divide
|
||||
* would silently EMA toward 0 if the invariant ever broke; instead the
|
||||
* NaN that division-by-zero produces propagates to ISV[96] and
|
||||
* downstream consumers (mag_concat_qdir in 2c.3c.6) where it is
|
||||
* loudly visible. */
|
||||
if (tid == 0) {
|
||||
const float rms = sqrtf(smem[0] / (float)N);
|
||||
const int slot = isv_h_s2_rms_ema_index;
|
||||
const float prev = isv[slot];
|
||||
isv[slot] = (1.0f - ema_alpha) * prev + ema_alpha * rms;
|
||||
}
|
||||
}
|
||||
@@ -273,6 +273,12 @@ impl StateResetRegistry {
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[MAMBA2_RETENTION_EMA_INDEX=89] — Mamba2 state-transition magnitude EMA, retention proxy; GPU attention_focus_ema_update kernel fills (Plan 4 E.5 Mode A)",
|
||||
},
|
||||
// ───── Plan 4 Task 2c.3c.5: trunk-RMS EMA producer ─────────────
|
||||
RegistryEntry {
|
||||
name: "isv_h_s2_rms_ema",
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[H_S2_RMS_EMA_INDEX=96] — per-batch RMS(save_h_s2) EMA (α=0.05); GPU h_s2_rms_ema_update kernel fills; producer-only in 2c.3c.5, consumer wired in 2c.3c.6 (mag_concat_qdir adaptive-scale). Cold-start 1.0 (neutral RMS) reapplied at fold boundary",
|
||||
},
|
||||
];
|
||||
Self { entries }
|
||||
}
|
||||
|
||||
@@ -2803,6 +2803,19 @@ impl DQNTrainer {
|
||||
tracing::warn!("C.5 cql_alpha_seed_update launch failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
// Plan 4 Task 2c.3c.5: per-step H_S2 RMS EMA into ISV[96].
|
||||
// Single-block 256-thread shmem-reduction kernel reads the
|
||||
// online trunk's `save_h_s2 [B, SH2]` (populated by the most
|
||||
// recent forward) and EMAs sqrt(sum_sq / N) into ISV[96] at
|
||||
// α=0.05. Producer-only — 2c.3c.6 wires the consumer in
|
||||
// `mag_concat_qdir`'s adaptive-scale path. Same launch
|
||||
// cadence as the Plan 3 producers above.
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
if let Err(e) = fused.trainer().launch_h_s2_rms_ema(ema_alpha) {
|
||||
tracing::warn!("Plan 4 2c.3c.5 h_s2_rms_ema launch failed: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// B.2 Plan 3 Task 3: freeze TRADE_TARGET_RATE at the configured
|
||||
@@ -4470,6 +4483,18 @@ impl DQNTrainer {
|
||||
fused.trainer().write_isv_signal_at(idx, 0.0);
|
||||
}
|
||||
}
|
||||
// Plan 4 Task 2c.3c.5: trunk-RMS EMA producer slot. Reset to
|
||||
// 1.0 (cold-start neutral RMS) at fold boundary so the first
|
||||
// h_s2_rms_ema_update fire on the new fold EMAs the measured
|
||||
// RMS toward 1.0 rather than collapsing toward 0.
|
||||
"isv_h_s2_rms_ema" => {
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
fused.trainer().write_isv_signal_at(
|
||||
crate::cuda_pipeline::gpu_dqn_trainer::H_S2_RMS_EMA_INDEX,
|
||||
1.0,
|
||||
);
|
||||
}
|
||||
}
|
||||
"isv_gamma_dir_eff" | "isv_gamma_mag_eff" | "isv_gamma_ord_eff" | "isv_gamma_urg_eff" => {
|
||||
// D.2 per-branch gamma slots. Reset to 0.0 at fold boundary;
|
||||
// per_branch_gamma_update GPU kernel re-populates on next epoch-boundary launch.
|
||||
|
||||
@@ -291,6 +291,8 @@ Plan 4 Task 2c.3c.2 (2026-04-24): backward infrastructure additions — two help
|
||||
|
||||
Plan 4 Task 2c.3c.3 (2026-04-24): GRN trunk backward scratch buffer allocation on `GpuDqnTrainer`. 8 new `CudaSlice<f32>` device buffers added (4 per GRN block × 2 blocks for h_s1 + h_s2): `bw_grn_h_s2_d_pre_ln` [B, SH2], `bw_grn_h_s2_d_linear_b_out` [B, 2*SH2], `bw_grn_h_s2_d_elu_out` [B, SH2], `bw_grn_h_s2_d_linear_a_out` [B, SH2], plus the four h_s1 mirrors at [B, SH1] (and 2*SH1 for the linear_b_out scratch — GLU's pre-activation has 2× hidden width because value-path + gate_pre are concatenated). The `d_pre_LN` buffer aliases both `d_glu_out` and `d_residual` (pure pass-through under `GrnBlock::backward_raw_phase1`'s pointer aliasing convention from 2c.3c.1); `d_linear_b_out` carries the phase1 → caller's Linear_b backward GEMM gradient; `d_elu_out` carries Linear_b backward's output into phase2; `d_linear_a_out` carries phase2's output to the caller's Linear_a backward GEMM. Allocated alongside the existing `alloc_backward_scratch` call in the trainer constructor (mirrors that helper's `stream.alloc_zeros::<f32>(size).map_err(...)?` pattern). Total footprint at SH1=SH2=256, B=512: 8 × 5 × 256 × 4B = 5.0 MB per fold (negligible against existing ~32 MB per-branch cuBLAS workspaces). Each field is `#[allow(dead_code)]` until 2c.3c.4's `apply_grn_trunk_backward_raw` helper consumes them — the comment `wired by Task 2c.3c.4` annotates each suppression. **Additive only — ZERO production callers in this commit**; the three backward panic gates (`backward_full`, `apply_iqn_trunk_gradient`, `apply_ensemble_diversity_backward`) remain in place until 2c.3c.4. Adam state for the new GRN tensors is auto-allocated since `m_buf`/`v_buf` are sized to `total_params + cutlass_tile_pad` which already includes the 13 GRN tensors after 2c.3a. cargo check clean at 11 warnings (baseline preserved); cargo build compiles 80 cubins (unchanged — no new kernel). +57 LOC. No new module / kernel / ISV slot / Orphan row.
|
||||
|
||||
Plan 4 Task 2c.3c.5 (2026-04-25): producer-only ISV slot append for the trunk-RMS adaptive scale. New ISV slot `H_S2_RMS_EMA_INDEX=96` tail-appended between the prior tail-data block and the layout-fingerprint pair; `ISV_LAYOUT_FINGERPRINT_LO_INDEX` shifted 94→97, `ISV_LAYOUT_FINGERPRINT_HI_INDEX` shifted 95→98, `ISV_TOTAL_DIM` 96→99. New CUDA kernel `h_s2_rms_ema_kernel.cu` registered in `build.rs::kernels_with_common` (kernel count 80 → 81): single-block 256-thread shmem-tree reduction (no atomicAdd per `feedback_no_atomicadd.md`) computing `RMS = sqrt(sum_sq / (B*SH2))` over the trainer's `save_h_s2 [B, SH2]` buffer (online trunk's post-GRN activation), then EMA-updating ISV[96] with α=0.05 (≈13-batch half-life). Cubin loaded in `GpuDqnTrainer::new` and stored in a new `h_s2_rms_ema_kernel: CudaFunction` field; `pub fn launch_h_s2_rms_ema(&self, ema_alpha)` mirrors `launch_reward_component_ema`'s style and reads `save_h_s2.raw_ptr()` + `config.shared_h2` + `isv_signals_dev_ptr` from the trainer (no plumbing through the experience collector — the buffer lives on the trainer). Launched from `training_loop.rs` immediately after the existing per-step ISV producer block (`reward_component_ema` + `trade_attempt_rate_ema` + `plan_threshold_update` + `seed_step_counter` + `cql_alpha_seed_update`), at the same launch cadence — once per `collect_experiences_gpu` epoch. Constructor cold-start writes ISV[96]=1.0 (neutral RMS) so the first kernel fire EMAs measured RMS toward 1.0 rather than collapsing toward 0; `StateResetRegistry::isv_h_s2_rms_ema` registered as `FoldReset` with the same 1.0 reapplication at fold boundary (dispatch arm added to `training_loop.rs::reset_named_state`). `layout_fingerprint_seed()` extended with `H_S2_RMS_EMA=96;ISV_LAYOUT_FINGERPRINT_LO=97;ISV_LAYOUT_FINGERPRINT_HI=98;ISV_TOTAL_DIM=99;` — new `LAYOUT_FINGERPRINT_CURRENT = 0x3e21acecd922e540` (was `0xcf3a24b0a1f70057`). **Producer-only — ZERO consumers in this commit**. 2c.3c.6 wires the consumer in `mag_concat_qdir`'s adaptive-scale path so the magnitude-branch decoder's residual stack is normalised by the trunk-output RMS regardless of GRN drift across training. cargo check clean at 11 warnings (baseline preserved); cargo build compiles 81 cubins (was 80). +109 / -10 LOC across `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (slot constants + docstring + CUBIN ref + field + load + launcher + cold-start + fingerprint seed entry), `crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu` (new file, 60 LOC), `crates/ml/build.rs` (1 entry), `crates/ml/src/trainers/dqn/state_reset_registry.rs` (1 entry), `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (per-step launch + FoldReset dispatch arm). Smoke deferred — 2c.3c.5 is producer-only with zero consumers, so runtime behaviour is structurally unchanged from 2c.3c.4 (which passed `multi_fold_convergence` in 599.73s with all 3 fold checkpoints written, train Sharpes 7.52 / 60.94 / 10.40 and val_Sharpes 0.51 / 0.64 / 0.82). The new producer's only runtime effects are (a) writing ISV[96] each step — read by nobody yet — and (b) bumping `LAYOUT_FINGERPRINT_CURRENT`, with no checkpoint-load path exercised by the smoke (it starts from scratch). cargo check + cargo build (with the new `h_s2_rms_ema_kernel.cubin` compiling clean under nvcc sm_80) is the standing validation; the meaningful `multi_fold_convergence` re-run lands with 2c.3c.6 when the consumer wires up in `mag_concat_qdir`'s adaptive-scale path. No new Orphan row — the kernel is wired to a real launch in this commit (cold-start + per-step EMA), classified Wired-Producer-Only until 2c.3c.6 promotes it to Wired-Producer+Consumer. 0 panic gates added/removed.
|
||||
|
||||
Plan 4 Task 2c.3c.4 followup (2026-04-25): stub-return defensive guards replaced with proper invariants + signal escalation in `gpu_dqn_trainer.rs`. Three classes of cleanup, prompted by the pre-commit stub-return heuristic flagging code paths exposed by the 2c.3c.4 commit's surrounding edits. (1) **Unreachable null/bounds guards → `debug_assert!` + deletion.** `read_isv_signal_at`, `read_atom_utilization`, `compute_q_spectral_gap` each guarded `is_null()` on `isv_signals_pinned` / `q_readback_pinned`, but both pointers are unconditionally allocated by the constructor (`cuMemAllocHost_v2` + `assert_eq!`) and never reassigned after construction outside Drop — the guards masked the contract instead of enforcing it. Same for `read_isv_signal_at`'s `index >= ISV_TOTAL_DIM` early-return (every caller passes a named ISV-slot constant) and `compute_q_spectral_gap`'s `n_cols == 0` early-return (`total_actions()` is the sum of branch sizes which config requires non-zero). All four converted to `debug_assert!(...)` so dev/test catches invariant violations loudly; release builds inherit the unsafe-deref crash on the broken precondition rather than a silent stub return. (2) **Silent training-instability mask → warn-log + collapse sentinel.** `compute_q_spectral_gap` previously returned 1.0 (= "balanced/healthy" per the downstream `NormalizedComponents::from_raw` calibration) on the first non-finite Q-value, masking NaN/Inf gradient corruption with the same value a healthy network produces. New behaviour: scan all samples, flag non-finite, then emit `tracing::warn!` and return 100.0 (the same collapse sentinel emitted when `lambda_2 < 1e-12` or `sigma2 < 1e-6`). Operators see the failure in logs alongside HEALTH_DIAG rather than chasing a phantom balanced spectral gap. (3) **Genuine domain encodings annotated** with `// ok: domain encoding` + explanatory comment: `lambda_2 < 1e-12 → 100.0` and `sigma2 < 1e-6 → 100.0` (rank-1 collapse — sigma_1/sigma_2 diverges, calibrated collapse sentinel), `power_iteration_largest`'s `norm < 1e-20 → 0.0` (M·v vanishes ⇒ eigenvalue is zero). `power_iteration_largest`'s `n == 0 → 0.0` was unreachable from the spectral_gap caller; replaced with `debug_assert!(n > 0)` at the helper entry. No behaviour change in production runs: the unreachable guards weren't firing and the new invariants match the constructor's contract; the non-finite Q-value path strictly improves observability. cargo check clean at 11 warnings (baseline preserved). +57 / -14 LOC, all in `gpu_dqn_trainer.rs`. No new module / kernel / ISV slot / Orphan row. Tiny followup tightens the remaining `.unwrap()` in `compute_q_spectral_gap` to `.expect("q_sample_history just received push_back; back() cannot be None")` documenting the invariant that the immediately-prior `push_back` makes the `back()` return Some-by-construction.
|
||||
|
||||
Plan 4 Task 2c.3c.4 (2026-04-25): GRN backward chain wired through all 3 panic-gated trunk-backward sites — smoke validates. Three panic gates removed: `BatchedBackward::backward_full`, `GpuDqnTrainer::apply_iqn_trunk_gradient`, `GpuDqnTrainer::apply_ensemble_diversity_backward`. New helper `BatchedForward::encoder_backward_chain` orchestrates the full GRN trunk backward composition (h_s2 → h_s1) symmetrically with `encoder_forward_only`: `backward_raw_phase1` (LN_dx + LN_dgamma/dbeta no-atomicAdd reduction + GLU_bwd) → cuBLAS Linear_b dW + dX → `backward_raw_phase2` (ELU_bwd) → cuBLAS Linear_a dW + dX → for h_s1 only: Linear_residual dW (no_bias_lda) + dX accumulation when bottleneck active → for h_s2 only: `saxpy_inplace(alpha=1.0)` accumulates `d_pre_ln_h_s2` into `d_h_s1` (identity residual gradient). Used by 4 callers writing to 3 different gradient targets: main backward (writes to `grad_buf` for the 13 GRN trunk tensors at indices [0..13)), CQL backward (writes to `cql_grad_scratch` — same layout as grad_buf), and the two auxiliary paths (`apply_iqn_trunk_gradient` and `apply_ensemble_diversity_backward` write to `iqn_trunk_m` then SAXPY into `grad_buf` with their respective scale factors). `iqn_trunk_m` grown from legacy 4-tensor element count (`sh1*sd + sh1 + sh2*sh1 + sh2`) to `padded_byte_offset(¶m_sizes, 13) / sizeof(f32)` so the layout matches `grad_buf`'s padded byte offsets exactly for the trunk portion — element-wise SAXPY across that span lands every per-tensor gradient at the correct location without needing per-tensor offset computation. `apply_ensemble_diversity_backward` additionally fixes value-head weight indices: `w_v1` 4→13, `w_v2` 6→15 (legacy indices were stale post-2c.3a — caught by the panic gate before they could write garbage). `BatchedBackward::launch_dw_only_no_bias_lda` added because Linear_residual (h_s1, no bias) needs the padded states stride that plain `launch_dw_only_no_bias` doesn't accept; `launch_dw_only` would dereference NULL inside the bias-grad kernel. ReLU masks on `bw_d_h_s2` and `bw_d_h_s1` are gone — the GRN trunk ends in LayerNorm which has no truncation derivative; the value-head FC retains its ReLU mask (value head still uses ReLU for its hidden activation). `launch_cublas_backward_to` migrated from `&self` to `&mut self` because the GRN backward needs to scribble per-block partial-reduction scratch inside `grn_h_s{1,2}_online`; the borrow split lets the trainer hand `&mut BatchedForward` + `&BatchedBackward` to `encoder_backward_chain` simultaneously. Smoke (`cargo test … multi_fold_convergence --ignored`, 599.73s, 3 folds × 5 epochs): all 3 `dqn_fold{N}_best.safetensors` checkpoints written; per-fold best train Sharpe 7.52 / 60.94 / 10.40 at epochs 2 / 4 / 2; per-fold val_Sharpe peaks 0.51 / 0.64 / 0.82; fold 2 epoch 5 hit PF=5.29 with +988% return — gradients flow cleanly through the GRN composition (no NaN, no Inf, no policy collapse, no fold-boundary explosion). Audit row #11 (IQN target trunk orphan) was already retired by 2c.3b; this commit retires the 3 backward panic gates that 2c.3a introduced. cargo check clean at 11 warnings (baseline preserved). +582 / -320 LOC across `batched_backward.rs` (+154/-154 — panic-gate removal + `launch_dw_only_no_bias_lda`), `batched_forward.rs` (+284/-0 — `encoder_backward_chain`), and `gpu_dqn_trainer.rs` (+221/-243 — 3 trunk-backward call sites rewritten + `iqn_trunk_m` grown + 4 main/CQL/IQN/ensemble backward call sites updated). No new module / kernel / ISV slot. Closes 3 of the 6 panic gates introduced by 2c.3a (3 forward gates were retired by 2c.3b); panic-gate count: 6 → 0.
|
||||
|
||||
Reference in New Issue
Block a user