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

feat(dqn-v2): C.6 Task 14 — epsilon GPU kernel + CPU monitor (exploration)

Browse Source 
epsilon_update kernel computes effective epsilon from ISV[EPOCH_IDX=39,
TOTAL_EPOCHS=40, LEARNING_HEALTH=12] with cosine schedule (eps_start ->
eps_end) plus health-coupled boost (up to +0.1 at collapse). Single-thread
cold-path kernel writes ISV[EPSILON_EFF_INDEX=41].

EpsilonMonitor is a read-only observer exposing eps_eff, epoch_idx,
total_ep, health, progress, fire_rate in DiagSnapshot.

Wires epsilon kernel launch at epoch boundary alongside tau kernel. Consumer
migration: log_training_config ISV-adaptive epsilon block now reads
ISV[EPSILON_EFF_INDEX] instead of computing `base_floor * (0.5 + volatility)`.

Tests: 3 monitor unit tests pass. cargo check -p ml at 8-warning baseline.

Plan 1 Task 14. Spec §4.C.6 (2026-04-24 revision).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-24 17:53:30 +02:00
parent 1032be3897
commit d346e03d48
9 changed files with 211 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml/build.rs
View File

@@ -85,6 +85,7 @@ fn main() {
"branch_grad_balance_kernel.cu",
"backtest_plan_kernel.cu",
"tau_update_kernel.cu",
"epsilon_update_kernel.cu",
];
// ALL kernels get common header (BF16 types + wrappers)

27
crates/ml/src/cuda_pipeline/epsilon_update_kernel.cu Normal file
View File

@@ -0,0 +1,27 @@
/* epsilon_update — effective exploration epsilon, GPU-driven (Plan 1 Task 14).
* Reads: ISV[EPOCH_IDX_INDEX=39], ISV[TOTAL_EPOCHS_INDEX=40],
* ISV[LEARNING_HEALTH_INDEX=12]
* Writes: ISV[EPSILON_EFF_INDEX=41]
* Cold path (per-epoch boundary). Single-thread kernel.
*
* Formula: cosine schedule (eps_start -> eps_end) over training, plus
* a volatility boost of up to +0.1 when health=0 (full collapse).
*/
extern "C" __global__ void epsilon_update(
const float* __restrict__ isv,
float* __restrict__ isv_out,
float eps_start,
float eps_end
) {
if (threadIdx.x != 0 || blockIdx.x != 0) return;
float epoch = isv[39];
float total_ep = fmaxf(isv[40], 1.0f);
float health = fminf(fmaxf(isv[12], 0.0f), 1.0f);
float progress = fminf(fmaxf(epoch / total_ep, 0.0f), 1.0f);
float cos_f = 0.5f * (1.0f + cosf(3.14159265358979f * progress));
float eps_cos = eps_end + (eps_start - eps_end) * cos_f;
float boost = (1.0f - health) * 0.1f; /* up to +0.1 when health=0 */
float eps_eff = eps_cos + boost;
eps_eff = fminf(fmaxf(eps_eff, 0.0f), 1.0f);
isv_out[41] = eps_eff;
}

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

@@ -79,6 +79,8 @@ static GRAD_DECOMP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grad
static BRANCH_GRAD_BALANCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/branch_grad_balance_kernel.cubin"));
/// Plan 1 Task 13: GPU-driven Polyak-EMA tau coefficient.
static TAU_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/tau_update_kernel.cubin"));
/// Plan 1 Task 14: GPU-driven effective exploration epsilon.
static EPSILON_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/epsilon_update_kernel.cubin"));
/// Mamba2 temporal scan configuration.
const MAMBA2_HISTORY_K: usize = 8; // Rolling history length
@@ -1402,6 +1404,9 @@ pub struct GpuDqnTrainer {
/// Plan 1 Task 13: GPU-driven Polyak-EMA tau. Cold-path (per-epoch).
/// Reads ISV[39, 40, 12]; writes ISV[TAU_EFF_INDEX=42].
tau_update_kernel: CudaFunction,
/// Plan 1 Task 14: GPU-driven effective epsilon. Cold-path (per-epoch).
/// Reads ISV[39, 40, 12]; writes ISV[EPSILON_EFF_INDEX=41].
epsilon_update_kernel: CudaFunction,
/// Host-side cached per-component norms populated at epoch boundary
/// from the pinned result slot. Index order matches `grad_mag_*_ratio`
/// accessors: `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`.
@@ -7042,6 +7047,14 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("tau_update load: {e}")))?
};
// Plan 1 Task 14: load epsilon_update kernel (cold-path, per-epoch).
let epsilon_update_kernel = {
let module = stream.context().load_cubin(EPSILON_UPDATE_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("epsilon_update cubin load: {e}")))?;
module.load_function("epsilon_update")
.map_err(|e| MLError::ModelError(format!("epsilon_update load: {e}")))?
};
// eval_v_range — pinned device-mapped (zero-copy write from CPU).
let eval_v_range_pinned: *mut f32 = unsafe {
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
@@ -9099,6 +9112,7 @@ impl GpuDqnTrainer {
branch_grad_scales_dev,
branch_slice_max_len,
tau_update_kernel,
epsilon_update_kernel,
grad_component_norms_mag: [0.0_f32; 9],
grad_component_norms_dir: [0.0_f32; 9],
grad_component_norms_trunk: [0.0_f32; 9],
@@ -10433,6 +10447,32 @@ impl GpuDqnTrainer {
Ok(())
}
/// Plan 1 Task 14: GPU-driven epsilon update.
///
/// Single-thread cold-path kernel: reads ISV[EPOCH_IDX=39, TOTAL_EPOCHS=40,
/// LEARNING_HEALTH=12]; writes ISV[EPSILON_EFF_INDEX=41] with cosine schedule
/// + health-coupled boost. Must be called once per epoch boundary BEFORE
/// any epsilon consumer reads ISV[EPSILON_EFF_INDEX].
pub(crate) fn launch_epsilon_update(&self, eps_start: f32, eps_end: f32) -> Result<(), MLError> {
let isv_ptr = self.isv_signals_dev_ptr;
let isv_out = isv_ptr; // kernel writes to the same ISV bus (offset 41)
unsafe {
self.stream
.launch_builder(&self.epsilon_update_kernel)
.arg(&isv_ptr)
.arg(&isv_out)
.arg(&eps_start)
.arg(&eps_end)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("epsilon_update launch: {e}")))?;
}
Ok(())
}
/// Read back the per-branch scales applied on the most recent
/// `launch_branch_grad_balance` call — `[dir, mag, ord, urg]`.
/// Values in `[0, 1]`; `1.0` means the branch passed through

8
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -3104,6 +3104,14 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("launch_tau_update: {e}"))
}
/// Plan 1 Task 14: delegate GPU epsilon update to trainer.
pub(crate) fn launch_epsilon_update(&self, eps_start: f32, eps_end: f32)
-> anyhow::Result<()>
{
self.trainer.launch_epsilon_update(eps_start, eps_end)
.map_err(|e| anyhow::anyhow!("launch_epsilon_update: {e}"))
}
/// G5: Set the EMA threshold for epistemic variance gating.
pub(crate) fn set_var_ema(&self, val: f32) {
self.trainer.set_var_ema(val);

113
crates/ml/src/trainers/dqn/monitors/epsilon_monitor.rs Normal file
View File

@@ -0,0 +1,113 @@
//! CPU-side observer for the GPU-driven effective exploration epsilon.
//!
//! The `epsilon_update` CUDA kernel reads ISV[EPOCH_IDX_INDEX=39],
//! ISV[TOTAL_EPOCHS_INDEX=40], and ISV[LEARNING_HEALTH_INDEX=12] to compute
//! a cosine-annealed epsilon with a health-coupled boost, writing the result
//! to ISV[EPSILON_EFF_INDEX=41]. This monitor reads that slot for HEALTH_DIAG
//! emission and fire-rate tracking.
//!
//! The monitor's `read()` returns `isv[EPSILON_EFF_INDEX]`. `diagnose()`
//! exposes all four input slots plus the effective epsilon and fire-rate.
use crate::cuda_pipeline::gpu_dqn_trainer::{
EPOCH_IDX_INDEX, TOTAL_EPOCHS_INDEX, LEARNING_HEALTH_INDEX, EPSILON_EFF_INDEX,
};
use crate::trainers::dqn::adaptive_monitor::{
AdaptiveMonitor, DiagSnapshot, FireRateStats, IsvBus,
};
pub(crate) struct EpsilonMonitor {
last_eps: f32,
fire_rate: FireRateStats,
}
impl EpsilonMonitor {
pub(crate) fn new() -> Self {
Self { last_eps: 0.0, fire_rate: FireRateStats::default() }
}
}
impl AdaptiveMonitor for EpsilonMonitor {
fn read(&self, isv: &IsvBus<'_>) -> f32 {
isv.read(EPSILON_EFF_INDEX)
}
fn diagnose(&self, isv: &IsvBus<'_>) -> DiagSnapshot {
let eps_eff = isv.read(EPSILON_EFF_INDEX);
let epoch = isv.read(EPOCH_IDX_INDEX);
let total_ep = isv.read(TOTAL_EPOCHS_INDEX);
let health = isv.read(LEARNING_HEALTH_INDEX);
let progress = if total_ep > 0.0 { (epoch / total_ep).clamp(0.0, 1.0) } else { 0.0 };
DiagSnapshot::new()
.with("eps_eff", eps_eff)
.with("epoch_idx", epoch)
.with("total_ep", total_ep)
.with("health", health)
.with("progress", progress)
.with("fire_rate", self.fire_rate.fire_rate())
}
fn observe(&mut self, current: f32) {
let fired = (current - self.last_eps).abs() > 1e-6;
self.fire_rate.record_fire(fired);
self.last_eps = current;
}
fn fire_rate(&self) -> &FireRateStats {
&self.fire_rate
}
fn name(&self) -> &'static str {
"epsilon"
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_isv(epoch: f32, total: f32, health: f32, eps_eff: f32) -> Vec<f32> {
let mut v = vec![0.0f32; 49];
v[EPOCH_IDX_INDEX] = epoch;
v[TOTAL_EPOCHS_INDEX] = total;
v[LEARNING_HEALTH_INDEX] = health;
v[EPSILON_EFF_INDEX] = eps_eff;
v
}
#[test]
fn read_returns_epsilon_eff_slot() {
let mut buf = make_isv(5.0, 100.0, 0.8, 0.15);
let bus = IsvBus::new(&mut buf);
let mon = EpsilonMonitor::new();
let v = mon.read(&bus);
assert!((v - 0.15).abs() < 1e-7, "expected eps_eff=0.15, got {v}");
}
#[test]
fn diagnose_exposes_all_fields() {
let mut buf = make_isv(20.0, 100.0, 0.7, 0.12);
let bus = IsvBus::new(&mut buf);
let mon = EpsilonMonitor::new();
let snap = mon.diagnose(&bus);
assert!(snap.fields.contains_key("eps_eff"));
assert!(snap.fields.contains_key("epoch_idx"));
assert!(snap.fields.contains_key("total_ep"));
assert!(snap.fields.contains_key("health"));
assert!(snap.fields.contains_key("progress"));
assert!(snap.fields.contains_key("fire_rate"));
let progress = snap.fields["progress"];
assert!((progress - 0.2).abs() < 1e-6, "expected progress=0.2, got {progress}");
}
#[test]
fn observe_tracks_fire_rate_on_change() {
let mut mon = EpsilonMonitor::new();
mon.observe(0.20);
mon.observe(0.20);
mon.observe(0.15);
let rate = mon.fire_rate().fire_rate();
assert!((rate - 2.0 / 3.0).abs() < 1e-6,
"expected 2/3 fire rate, got {rate}");
}
}

1
crates/ml/src/trainers/dqn/monitors/mod.rs
View File

@@ -5,5 +5,6 @@
//! read those slots for HEALTH_DIAG emission and `controller_activity`
//! smoke-test fire-rate tracking. No CPU-side adaptive computation.
pub(crate) mod epsilon_monitor;
pub(crate) mod grad_balancer_monitor;
pub(crate) mod tau_monitor;

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

@@ -21,7 +21,7 @@ use cudarc::driver::CudaSlice;
use common::metrics::{questdb_sink, training_metrics};
use tracing::{debug, info, warn};
use crate::cuda_pipeline::gpu_dqn_trainer::EPOCH_IDX_INDEX;
use crate::cuda_pipeline::gpu_dqn_trainer::{EPOCH_IDX_INDEX, EPSILON_EFF_INDEX};
use crate::dqn::logging::{log_epoch_start, log_epoch_end, log_training_progress};
use crate::dqn::target_update::convergence_half_life;
use crate::evaluation::metrics::calculate_var_cvar;
@@ -258,6 +258,15 @@ impl DQNTrainer {
}
}
// Plan 1 Task 14: launch GPU epsilon update kernel at epoch boundary.
if let Some(ref fused) = self.fused_ctx {
let eps_start = self.hyperparams.epsilon_start as f32;
let eps_end = self.hyperparams.epsilon_end as f32;
if let Err(e) = fused.launch_epsilon_update(eps_start, eps_end) {
tracing::warn!(epoch, "launch_epsilon_update failed (non-fatal): {e}");
}
}
// Collect pending async validation from PREVIOUS epoch
if let Some(val) = self.pending_val_loss.take() {
// Store for this epoch's logging (one-epoch delay is acceptable)
@@ -928,17 +937,20 @@ impl DQNTrainer {
}
{
// ISV-adaptive exploration: read regime signals, modulate epsilon + sigma.
// Volatile regime → explore more (higher epsilon, sigma).
// Stable regime → exploit (lower epsilon, sigma).
// Plan 1 Task 14: read GPU-computed epsilon from ISV[EPSILON_EFF_INDEX].
// epsilon_update kernel ran at epoch boundary with cosine schedule +
// health-coupled boost; no CPU-side adaptive computation here.
let (regime_stab, volatility) = if let Some(ref fused) = self.fused_ctx {
fused.read_isv_regime()
} else {
(0.5, 0.5)
};
let base_floor = self.hyperparams.noisy_epsilon_floor.unwrap_or(0.05) as f32;
let adaptive_epsilon = base_floor * (0.5 + volatility);
let adaptive_epsilon = if let Some(ref fused) = self.fused_ctx {
fused.trainer().read_isv_signal_at(EPSILON_EFF_INDEX)
} else {
self.hyperparams.noisy_epsilon_floor.unwrap_or(0.05) as f32
};
let mut agent = self.agent.write().await;
agent.set_epsilon(adaptive_epsilon as f64);

1
docs/dqn-gpu-hot-path-audit.md
View File

@@ -68,6 +68,7 @@
| `training_loop.rs:1003` | `clone_htod_f32` → targets device buffer | COLD-PATH | Data loading: DBN targets uploaded before step loop | OK |
| `training_loop.rs:1007` | `clone_htod_f32` → features device buffer | COLD-PATH | Data loading: feature matrix uploaded before step loop | OK |
| `tau_update_kernel.cu` (Plan 1 Task 13) | single-thread kernel: reads ISV[39,40,12]; writes ISV[42] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; zero hot-path overhead | OK |
| `epsilon_update_kernel.cu` (Plan 1 Task 14) | single-thread kernel: reads ISV[39,40,12]; writes ISV[41] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; zero hot-path overhead | OK |
## Summary

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

@@ -46,6 +46,8 @@
| `trainers/dqn/monitors/grad_balancer_monitor.rs` | Read-only observer for grad_balance_isv_update kernel output (ISV slots 31..35); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 17 | — |
| `trainers/dqn/monitors/tau_monitor.rs` | Read-only observer for tau_update kernel output (ISV slot 42); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 13 | — |
| `cuda_pipeline/tau_update_kernel.cu` | `GpuDqnTrainer::launch_tau_update` → `FusedTrainingCtx::launch_tau_update` → `training_loop.rs` epoch boundary | Wired | Plan 1 Task 13 — cold-path (per-epoch) | — |
| `trainers/dqn/monitors/epsilon_monitor.rs` | Read-only observer for epsilon_update kernel output (ISV slot 41); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 14 | — |
| `cuda_pipeline/epsilon_update_kernel.cu` | `GpuDqnTrainer::launch_epsilon_update` → `FusedTrainingCtx::launch_epsilon_update` → `training_loop.rs` epoch boundary | Wired | Plan 1 Task 14 — cold-path (per-epoch) | — |
## CUDA Pipeline — Rust Wrappers