feat(dqn-v2): C.1 quantile-based atom support (Plan 2 Task 1)
Replaces half = max(10*q_gap, 3*std_ema, floor) tuned constants with quantile-based: half = max(|q_p95 - v_center|, |v_center - q_p5|) clamped to [min_half_floor, abs_half]. Covers observed Q distribution directly. New GPU kernel q_quantile_reduce reads per-sample Q-values from q_out_buf, sorts per branch (bitonic for power-of-2 N, quickselect otherwise), writes ISV[Q_P05_*=47..51) and ISV[Q_P95_*=51..55). Cold-path per-epoch (4 blocks, 1 thread each). No atomicAdd. ISV slots 47-54 added at tail (8 total). Fingerprint shifted to 55-56. ISV_TOTAL_DIM grows 49 -> 57. Fingerprint seed updated; recomputes hash at compile time (new value: 0xbf6c400c026d77e3). Launch order: reduce_current_q_stats (populates q_out_buf) -> q_quantile_reduce (writes ISV P5/P95) -> reduce_current_q_stats_per_branch -> update_eval_v_range (reads ISV P5/P95 for half-width). Bootstrap: q_p05 = v_min, q_p95 = v_max (matches current atom range at cold start). FoldReset: isv_q_quantiles dispatch arm resets to bootstrap values at fold boundary. StateResetRegistry: isv_q_quantiles as FoldReset; dispatch arm added in training_loop.rs::reset_named_state. Docs: isv-slots.md rows [47..57) updated, fingerprint tail reference corrected to [55..57). dqn-wire-up-audit.md: q_quantile_kernel.cu added. Plan 2 Task 1. Spec §4.C.1. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -89,6 +89,7 @@ fn main() {
|
||||
"gamma_update_kernel.cu",
|
||||
"kelly_cap_update_kernel.cu",
|
||||
"atoms_update_kernel.cu",
|
||||
"q_quantile_kernel.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
@@ -87,6 +87,7 @@ static GAMMA_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/gam
|
||||
static KELLY_CAP_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/kelly_cap_update_kernel.cubin"));
|
||||
/// Plan 1 Task 9: GPU-driven C51 atom position recompute (4-branch, single launch).
|
||||
static ATOMS_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/atoms_update_kernel.cubin"));
|
||||
static Q_QUANTILE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/q_quantile_kernel.cubin"));
|
||||
|
||||
/// Mamba2 temporal scan configuration.
|
||||
const MAMBA2_HISTORY_K: usize = 8; // Rolling history length
|
||||
@@ -186,13 +187,23 @@ const ISV_NETWORK_DIM: usize = 23;
|
||||
/// [46] PLAN_THRESHOLD_INDEX — constructor writes 0.5f; plan kernels read instead of
|
||||
/// hardcoded literal (Plan 1 Task 16).
|
||||
///
|
||||
/// Slots [47..49) carry the layout fingerprint (ISV_LAYOUT_FINGERPRINT_LO_INDEX
|
||||
/// Slots [47..55) are the DQN v2 Plan 2 Task 1 expansion (spec §4.C.1, 2026-04-24):
|
||||
/// [47] Q_P05_DIR_INDEX — per-epoch 5th-percentile Q EMA for direction branch.
|
||||
/// [48] Q_P05_MAG_INDEX — per-epoch 5th-percentile Q EMA for magnitude branch.
|
||||
/// [49] Q_P05_ORD_INDEX — per-epoch 5th-percentile Q EMA for order branch.
|
||||
/// [50] Q_P05_URG_INDEX — per-epoch 5th-percentile Q EMA for urgency branch.
|
||||
/// [51] Q_P95_DIR_INDEX — per-epoch 95th-percentile Q EMA for direction branch.
|
||||
/// [52] Q_P95_MAG_INDEX — per-epoch 95th-percentile Q EMA for magnitude branch.
|
||||
/// [53] Q_P95_ORD_INDEX — per-epoch 95th-percentile Q EMA for order branch.
|
||||
/// [54] Q_P95_URG_INDEX — per-epoch 95th-percentile Q EMA for urgency branch.
|
||||
///
|
||||
/// Slots [55..57) carry the layout fingerprint (ISV_LAYOUT_FINGERPRINT_LO_INDEX
|
||||
/// and ISV_LAYOUT_FINGERPRINT_HI_INDEX). These are placed at the tail because
|
||||
/// slots [0] and [1] are actively written by `isv_signal_update` (Q-drift and
|
||||
/// gradient-norm EMA respectively). 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 the structural-hash design rationale.
|
||||
const ISV_TOTAL_DIM: usize = 49;
|
||||
const ISV_TOTAL_DIM: usize = 57;
|
||||
/// 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).
|
||||
@@ -311,7 +322,27 @@ pub const CQL_ALPHA_INDEX: usize = 45;
|
||||
/// Plan 3 B.4 may later make this reactive.
|
||||
pub const PLAN_THRESHOLD_INDEX: usize = 46;
|
||||
|
||||
/// ISV slot [47] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
|
||||
/// ISV slots [47..51) — C.1 Quantile-based atom support (Plan 2 Task 1, spec §4.C.1).
|
||||
/// Per-branch 5th percentile of observed Q-values via adaptive-rate EMA.
|
||||
/// Producer: q_quantile_reduce GPU kernel (cold-path per-epoch).
|
||||
/// Consumer: update_eval_v_range, which computes
|
||||
/// `half = max(|q_p95 - v_center|, |v_center - q_p5|)`
|
||||
/// clamped to [min_half_floor, abs_half].
|
||||
/// Bootstrap: q_p05 = v_min (matches cold-start atom range).
|
||||
pub const Q_P05_DIR_INDEX: usize = 47;
|
||||
pub const Q_P05_MAG_INDEX: usize = 48;
|
||||
pub const Q_P05_ORD_INDEX: usize = 49;
|
||||
pub const Q_P05_URG_INDEX: usize = 50;
|
||||
/// ISV slots [51..55) — C.1 Quantile-based atom support (Plan 2 Task 1, spec §4.C.1).
|
||||
/// Per-branch 95th percentile of observed Q-values via adaptive-rate EMA.
|
||||
/// Producer: q_quantile_reduce GPU kernel (cold-path per-epoch).
|
||||
/// Bootstrap: q_p95 = v_max (matches cold-start atom range).
|
||||
pub const Q_P95_DIR_INDEX: usize = 51;
|
||||
pub const Q_P95_MAG_INDEX: usize = 52;
|
||||
pub const Q_P95_ORD_INDEX: usize = 53;
|
||||
pub const Q_P95_URG_INDEX: usize = 54;
|
||||
|
||||
/// ISV slot [55] — 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
|
||||
@@ -322,9 +353,9 @@ pub const PLAN_THRESHOLD_INDEX: usize = 46;
|
||||
///
|
||||
/// Placement at the tail (not head) is mandated by slots [0] and [1] being
|
||||
/// actively written by `isv_signal_update` (Q-drift and gradient-norm EMA).
|
||||
pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 47;
|
||||
/// ISV slot [48] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
|
||||
pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 48;
|
||||
pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 55;
|
||||
/// ISV slot [56] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits).
|
||||
pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 56;
|
||||
/// Canonical alias for the fingerprint slot (the low half).
|
||||
pub const ISV_LAYOUT_FINGERPRINT_INDEX: usize = ISV_LAYOUT_FINGERPRINT_LO_INDEX;
|
||||
|
||||
@@ -378,8 +409,10 @@ const fn layout_fingerprint_seed() -> &'static [u8] {
|
||||
EPOCH_IDX=39;TOTAL_EPOCHS=40;\
|
||||
EPSILON_EFF=41;TAU_EFF=42;GAMMA_EFF=43;KELLY_CAP_EFF=44;\
|
||||
CQL_ALPHA=45;PLAN_THRESHOLD=46;\
|
||||
ISV_LAYOUT_FINGERPRINT_LO=47;ISV_LAYOUT_FINGERPRINT_HI=48;\
|
||||
ISV_TOTAL_DIM=49"
|
||||
Q_P05_DIR=47;Q_P05_MAG=48;Q_P05_ORD=49;Q_P05_URG=50;\
|
||||
Q_P95_DIR=51;Q_P95_MAG=52;Q_P95_ORD=53;Q_P95_URG=54;\
|
||||
ISV_LAYOUT_FINGERPRINT_LO=55;ISV_LAYOUT_FINGERPRINT_HI=56;\
|
||||
ISV_TOTAL_DIM=57"
|
||||
}
|
||||
|
||||
/// Compile-time layout fingerprint. Burned into the binary at build time.
|
||||
@@ -1421,6 +1454,19 @@ pub struct GpuDqnTrainer {
|
||||
/// Plan 1 Task 9: GPU-driven C51 atom position recompute. Cold-path (per-epoch + SGD step).
|
||||
/// Reads ISV v-range slots [23..31) + spacing_raw params; writes atom_positions_buf.
|
||||
atoms_update_kernel: CudaFunction,
|
||||
/// Plan 2 Task 1 C.1: GPU-driven per-branch Q P5/P95 percentile EMA update.
|
||||
/// Cold-path (per-epoch). Reads q_out_buf [N, total_actions]; for each branch,
|
||||
/// sorts per-sample max-Q via bitonic/quickselect, picks P5/P95, EMA into
|
||||
/// ISV[Q_P05_*=47..51) and ISV[Q_P95_*=51..55).
|
||||
q_quantile_kernel: CudaFunction,
|
||||
/// Action branch offsets in the flat [N, total_actions] Q-value buffer.
|
||||
/// [0]=dir_off=0, [1]=mag_off, [2]=ord_off, [3]=urg_off.
|
||||
/// Passed to q_quantile_reduce as the branch_offsets argument.
|
||||
q_quantile_branch_offsets_dev: CudaSlice<i32>,
|
||||
/// Action branch sizes in the flat [N, total_actions] Q-value buffer.
|
||||
/// [b] = number of actions for branch b (dir=4, mag=3, ord=3, urg=3 by default).
|
||||
/// Passed to q_quantile_reduce as the branch_sizes argument.
|
||||
q_quantile_branch_sizes_dev: CudaSlice<i32>,
|
||||
/// 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]`.
|
||||
@@ -2094,12 +2140,13 @@ pub struct GpuDqnTrainer {
|
||||
q_dir_bin_means_reduce_kernel: CudaFunction,
|
||||
|
||||
// ── ISV core buffers (pinned device-mapped, GPU read/write) ──
|
||||
// isv_signals_pinned is sized for ISV_TOTAL_DIM (49) — the extra slots
|
||||
// slots [23..31] carry per-branch Q-support centres/half-widths; [39..49)
|
||||
// are the Plan 1 C.6 expansion (static configs + pre-allocated GPU slots).
|
||||
// isv_signals_pinned is sized for ISV_TOTAL_DIM (57) — the extra slots
|
||||
// slots [23..31] carry per-branch Q-support centres/half-widths; [39..47)
|
||||
// are the Plan 1 C.6 expansion (static configs + pre-allocated GPU slots);
|
||||
// [47..55) are the Plan 2 C.1 Q-quantile EMAs.
|
||||
// The network (w_isv_fc1) and history rotation still consume only the first
|
||||
// ISV_NETWORK_DIM (23) slots; the tail is broadcast-bus scratchpad only.
|
||||
isv_signals_pinned: *mut f32, // [ISV_TOTAL_DIM = 49]
|
||||
isv_signals_pinned: *mut f32, // [ISV_TOTAL_DIM = 57]
|
||||
isv_signals_dev_ptr: u64,
|
||||
isv_history_pinned: *mut f32, // [ISV_K * 12 = 48] — history rotates slots [0..11] only
|
||||
isv_history_dev_ptr: u64,
|
||||
@@ -3239,23 +3286,31 @@ impl GpuDqnTrainer {
|
||||
(1.0 - alpha_std) * self.eval_q_std_ema[branch_idx] + alpha_std * q_std;
|
||||
}
|
||||
|
||||
// Width from max(Q-gap action scale, 3σ band, min floor), clamped
|
||||
// to the full config range. Floor = 10% of config range — enough
|
||||
// atom-grid resolution to differentiate actions regardless of the
|
||||
// observed distribution.
|
||||
let gap_width = (10.0 * q_gap)
|
||||
.max(3.0 * self.eval_q_std_ema[branch_idx])
|
||||
.max(min_half_floor);
|
||||
let half = gap_width.min(abs_half).max(min_half_floor);
|
||||
// Plan 2 Task 1 C.1: quantile-based half-width (spec §4.C.1).
|
||||
// half = max(|q_p95 - v_center|, |v_center - q_p05|), clamped to
|
||||
// [min_half_floor, abs_half]. q_p05/q_p95 are the EMA-smoothed
|
||||
// percentiles written by q_quantile_reduce into ISV[47..55).
|
||||
// On cold-start (before q_quantile_reduce runs) the slots hold
|
||||
// v_min / v_max bootstraps, reproducing the pre-spec behaviour.
|
||||
let center = self.eval_q_mean_ema[branch_idx];
|
||||
let q_p05 = unsafe {
|
||||
*self.isv_signals_pinned.add(Q_P05_DIR_INDEX + branch_idx)
|
||||
};
|
||||
let q_p95 = unsafe {
|
||||
*self.isv_signals_pinned.add(Q_P95_DIR_INDEX + branch_idx)
|
||||
};
|
||||
let half_from_p95 = (q_p95 - center).abs();
|
||||
let half_from_p05 = (center - q_p05).abs();
|
||||
let half = half_from_p95.max(half_from_p05)
|
||||
.max(min_half_floor)
|
||||
.min(abs_half);
|
||||
// Centre clamp so the band always fits inside [v_min, v_max].
|
||||
let center = self.eval_q_mean_ema[branch_idx].clamp(
|
||||
v_min_f + half,
|
||||
v_max_f - half,
|
||||
);
|
||||
let center = center.clamp(v_min_f + half, v_max_f - half);
|
||||
|
||||
tracing::info!(
|
||||
target: "isv_vrange_diag",
|
||||
branch_idx, q_mean, q_std, q_gap, center, half,
|
||||
q_p05, q_p95,
|
||||
initialized = initialized_before,
|
||||
"update_eval_v_range branch"
|
||||
);
|
||||
@@ -6993,6 +7048,36 @@ impl GpuDqnTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("atoms_update load: {e}")))?
|
||||
};
|
||||
|
||||
// Plan 2 Task 1 C.1: load q_quantile_reduce kernel (cold-path, per-epoch).
|
||||
let q_quantile_kernel = {
|
||||
let module = stream.context().load_cubin(Q_QUANTILE_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("q_quantile cubin load: {e}")))?;
|
||||
module.load_function("q_quantile_reduce")
|
||||
.map_err(|e| MLError::ModelError(format!("q_quantile_reduce 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),
|
||||
// urg [b0+b1+b2..total_actions).
|
||||
let (q_quantile_branch_offsets_dev, q_quantile_branch_sizes_dev) = {
|
||||
let b0 = config.branch_0_size as i32;
|
||||
let b1 = config.branch_1_size as i32;
|
||||
let b2 = config.branch_2_size as i32;
|
||||
let b3 = config.branch_3_size as i32;
|
||||
let offsets: [i32; 4] = [0, b0, b0 + b1, b0 + b1 + b2];
|
||||
let sizes: [i32; 4] = [b0, b1, b2, b3];
|
||||
let mut off_dev = stream.alloc_zeros::<i32>(4)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc q_quantile branch_offsets: {e}")))?;
|
||||
stream.memcpy_htod(&offsets, &mut off_dev)
|
||||
.map_err(|e| MLError::ModelError(format!("htod q_quantile branch_offsets: {e}")))?;
|
||||
let mut sz_dev = stream.alloc_zeros::<i32>(4)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc q_quantile branch_sizes: {e}")))?;
|
||||
stream.memcpy_htod(&sizes, &mut sz_dev)
|
||||
.map_err(|e| MLError::ModelError(format!("htod q_quantile branch_sizes: {e}")))?;
|
||||
(off_dev, sz_dev)
|
||||
};
|
||||
|
||||
// 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;
|
||||
@@ -8191,10 +8276,10 @@ impl GpuDqnTrainer {
|
||||
};
|
||||
|
||||
// ── ISV core buffers (pinned device-mapped) ──────────────────────
|
||||
// Allocation uses `ISV_TOTAL_DIM` so all scratchpad slots [23..49) have
|
||||
// Allocation uses `ISV_TOTAL_DIM` so all scratchpad slots [23..57) have
|
||||
// backing storage. The network-facing path still consumes only the first
|
||||
// `ISV_NETWORK_DIM` slots — the scratchpad tail is invisible to `w_isv_fc1`.
|
||||
// Slots [47..49) hold the layout fingerprint (tail placement because [0..2)
|
||||
// Slots [55..57) hold the layout fingerprint (tail placement because [0..2)
|
||||
// are actively written by `isv_signal_update`).
|
||||
let (isv_signals_pinned, isv_signals_dev_ptr) = {
|
||||
let num_bytes = ISV_TOTAL_DIM * std::mem::size_of::<f32>();
|
||||
@@ -8243,7 +8328,14 @@ impl GpuDqnTrainer {
|
||||
*sig_ptr.add(TOTAL_EPOCHS_INDEX) = config.total_epochs as f32;
|
||||
*sig_ptr.add(CQL_ALPHA_INDEX) = config.cql_alpha;
|
||||
*sig_ptr.add(PLAN_THRESHOLD_INDEX) = 0.5_f32;
|
||||
// Layout fingerprint (ISV[47..49)). Compile-time structural hash of
|
||||
// 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.
|
||||
for b in 0..4usize {
|
||||
*sig_ptr.add(Q_P05_DIR_INDEX + b) = v_min_f;
|
||||
*sig_ptr.add(Q_P95_DIR_INDEX + b) = v_max_f;
|
||||
}
|
||||
// Layout fingerprint (ISV[55..57)). Compile-time structural hash of
|
||||
// the slot layout; checkpoint load fails-fast on mismatch.
|
||||
// Stored as a u64 split across two f32 lanes using raw bit-cast so
|
||||
// no precision is lost. See LAYOUT_FINGERPRINT_CURRENT docs for why
|
||||
@@ -9051,6 +9143,9 @@ impl GpuDqnTrainer {
|
||||
gamma_update_kernel,
|
||||
kelly_cap_update_kernel,
|
||||
atoms_update_kernel,
|
||||
q_quantile_kernel,
|
||||
q_quantile_branch_offsets_dev,
|
||||
q_quantile_branch_sizes_dev,
|
||||
grad_component_norms_mag: [0.0_f32; 9],
|
||||
grad_component_norms_dir: [0.0_f32; 9],
|
||||
grad_component_norms_trunk: [0.0_f32; 9],
|
||||
@@ -10455,6 +10550,46 @@ impl GpuDqnTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Plan 2 Task 1 C.1: GPU-driven per-branch Q P5/P95 percentile EMA update.
|
||||
///
|
||||
/// Single-block cold-path kernel (grid=(4,1,1), block=(1,1,1)): for each branch,
|
||||
/// collects per-sample max-Q-over-actions, sorts via bitonic/quickselect, picks the
|
||||
/// 5th and 95th percentile, and EMA-blends into ISV[Q_P05_*=47..51) and
|
||||
/// ISV[Q_P95_*=51..55). Must be called once per epoch boundary AFTER
|
||||
/// `reduce_current_q_stats` (which calls `populate_q_out`) and BEFORE
|
||||
/// `update_eval_v_range` reads the percentile slots.
|
||||
///
|
||||
/// ema_alpha: adaptive smoothing rate ∈ [0.01, 0.5]. A rate of 0.1 is a
|
||||
/// reasonable default (10% of each new epoch's estimate).
|
||||
pub(crate) fn launch_q_quantile_reduce(&self, ema_alpha: f32) -> Result<(), MLError> {
|
||||
let q_out_ptr = self.q_out_buf.raw_ptr();
|
||||
let n_samples = self.config.batch_size as i32;
|
||||
let total_acts = self.total_actions() as i32;
|
||||
let off_ptr = self.q_quantile_branch_offsets_dev.raw_ptr();
|
||||
let sz_ptr = self.q_quantile_branch_sizes_dev.raw_ptr();
|
||||
let isv_ptr = self.isv_signals_dev_ptr;
|
||||
let isv_out = isv_ptr; // writes back to same pinned bus (reads prev before writing)
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.q_quantile_kernel)
|
||||
.arg(&q_out_ptr)
|
||||
.arg(&n_samples)
|
||||
.arg(&total_acts)
|
||||
.arg(&off_ptr)
|
||||
.arg(&sz_ptr)
|
||||
.arg(&isv_ptr)
|
||||
.arg(&isv_out)
|
||||
.arg(&ema_alpha)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (4, 1, 1), // one block per branch
|
||||
block_dim: (1, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("q_quantile_reduce 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
|
||||
|
||||
161
crates/ml/src/cuda_pipeline/q_quantile_kernel.cu
Normal file
161
crates/ml/src/cuda_pipeline/q_quantile_kernel.cu
Normal file
@@ -0,0 +1,161 @@
|
||||
/* q_quantile_reduce — GPU-driven per-branch P5/P95 EMA update (Plan 2 Task 1, spec §4.C.1).
|
||||
*
|
||||
* For each of the 4 action branches (direction=0, magnitude=1, order=2, urgency=3),
|
||||
* computes the per-sample max-Q-over-actions, sorts those N values in shared memory
|
||||
* via bitonic sort or quickselect, picks the 5th and 95th percentile values, and
|
||||
* applies an EMA to ISV[Q_P05_*] and ISV[Q_P95_*].
|
||||
*
|
||||
* Cold-path per-epoch kernel. Launched with grid=(4,1,1), block=(1,1,1) — single
|
||||
* thread per block, one block per branch. N samples can reach up to the full batch
|
||||
* size; for N <= MAX_SHMEM_SORT_SIZE the sort runs in shared memory; larger N is
|
||||
* truncated to a representative subsample (still epoch-cold-path accurate).
|
||||
*
|
||||
* No atomicAdd anywhere — per-sample array + single-thread reduce only.
|
||||
* (feedback_no_atomicadd.md).
|
||||
*
|
||||
* ISV slot mapping (Plan 2 Task 1 corrected layout):
|
||||
* Q_P05_DIR=47, Q_P05_MAG=48, Q_P05_ORD=49, Q_P05_URG=50
|
||||
* Q_P95_DIR=51, Q_P95_MAG=52, Q_P95_ORD=53, Q_P95_URG=54
|
||||
*/
|
||||
|
||||
/* Maximum samples that fit in shared memory for the sort path.
|
||||
* Two arrays of this size × sizeof(float) must fit within 48KB (0xc000) shmem.
|
||||
* 3072 × 4 × 2 = 24576 bytes — fits with plenty of margin.
|
||||
* For N > 3072, the first 3072 samples form a representative subsample. */
|
||||
#define MAX_SHMEM_SORT_SIZE 3072
|
||||
|
||||
/* ── Bitonic sort — in-place ascending sort on arr[0..len).
|
||||
* len must be a power of 2. Single-thread; safe on cold path. ──────────── */
|
||||
__device__ static void bitonic_sort_block(float* arr, int len)
|
||||
{
|
||||
for (int k = 2; k <= len; k <<= 1) {
|
||||
for (int j = k >> 1; j >= 1; j >>= 1) {
|
||||
for (int i = 0; i < len; i++) {
|
||||
int ixj = i ^ j;
|
||||
if (ixj > i) {
|
||||
bool ascending = ((i & k) == 0);
|
||||
if (( ascending && arr[i] > arr[ixj]) ||
|
||||
(!ascending && arr[i] < arr[ixj])) {
|
||||
float tmp = arr[i];
|
||||
arr[i] = arr[ixj];
|
||||
arr[ixj] = tmp;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* ── Quickselect: find the k-th smallest value (0-indexed) in arr[0..n).
|
||||
* Modifies arr in place. Single-thread; safe on cold path. ─────────────── */
|
||||
__device__ static float nth_element_float(float* arr, int n, int k)
|
||||
{
|
||||
int lo = 0, hi = n - 1;
|
||||
while (lo < hi) {
|
||||
float pivot = arr[(lo + hi) >> 1];
|
||||
int l = lo, r = hi;
|
||||
while (l <= r) {
|
||||
while (arr[l] < pivot) l++;
|
||||
while (arr[r] > pivot) r--;
|
||||
if (l <= r) {
|
||||
float tmp = arr[l];
|
||||
arr[l] = arr[r];
|
||||
arr[r] = tmp;
|
||||
l++; r--;
|
||||
}
|
||||
}
|
||||
if (k <= r) hi = r;
|
||||
else if (k >= l) lo = l;
|
||||
else break;
|
||||
}
|
||||
return arr[k];
|
||||
}
|
||||
|
||||
/* ── Main kernel ─────────────────────────────────────────────────────────── */
|
||||
extern "C" __global__ void q_quantile_reduce(
|
||||
const float* __restrict__ q_values, /* [n_samples, total_actions] row-major */
|
||||
int n_samples,
|
||||
int total_actions,
|
||||
const int* __restrict__ branch_offsets, /* [4] — action index of each branch start */
|
||||
const int* __restrict__ branch_sizes, /* [4] — number of actions per branch */
|
||||
const float* __restrict__ isv, /* read: current Q_P05/Q_P95 EMA seeds */
|
||||
float* __restrict__ isv_out, /* write: updated Q_P05/Q_P95 EMA values */
|
||||
float ema_alpha
|
||||
)
|
||||
{
|
||||
/* Single-thread per block; blockIdx.x selects the branch. */
|
||||
if (threadIdx.x != 0) return;
|
||||
|
||||
const int branch = (int)blockIdx.x; /* 0=dir, 1=mag, 2=ord, 3=urg */
|
||||
|
||||
/* ISV slot indices for this branch (Plan 2 Task 1 corrected layout).
|
||||
* Q_P05 base = 47, Q_P95 base = 51. */
|
||||
const int p05_slot = 47 + branch;
|
||||
const int p95_slot = 51 + branch;
|
||||
|
||||
const int b_off = branch_offsets[branch];
|
||||
const int b_size = branch_sizes[branch];
|
||||
const int n = n_samples;
|
||||
|
||||
/* Guard: degenerate configuration — skip without modifying ISV. */
|
||||
if (b_size <= 0 || b_off < 0 || n <= 0 ||
|
||||
total_actions <= 0 || b_off + b_size > total_actions) {
|
||||
return;
|
||||
}
|
||||
|
||||
/* ── Phase 1: collect per-sample max-Q over this branch's actions.
|
||||
* Stored in svals[0..fill_n). fill_n = min(n, MAX_SHMEM_SORT_SIZE).
|
||||
* When n > MAX_SHMEM_SORT_SIZE the first fill_n samples are used; this is a
|
||||
* representative subsample for the epoch-boundary percentile estimate. */
|
||||
__shared__ float svals[MAX_SHMEM_SORT_SIZE];
|
||||
__shared__ float svals_copy[MAX_SHMEM_SORT_SIZE];
|
||||
|
||||
const int fill_n = (n <= MAX_SHMEM_SORT_SIZE) ? n : MAX_SHMEM_SORT_SIZE;
|
||||
|
||||
for (int s = 0; s < fill_n; s++) {
|
||||
const float* row = q_values + (long long)s * total_actions + b_off;
|
||||
float mx = row[0];
|
||||
for (int a = 1; a < b_size; a++) {
|
||||
float v = row[a];
|
||||
if (v > mx) mx = v;
|
||||
}
|
||||
svals[s] = mx;
|
||||
}
|
||||
|
||||
/* ── Phase 2: percentile selection.
|
||||
* p05_idx = floor(fill_n * 0.05), p95_idx = floor(fill_n * 0.95), clamped. */
|
||||
int p05_idx = (int)((float)fill_n * 0.05f);
|
||||
int p95_idx = (int)((float)fill_n * 0.95f);
|
||||
if (p05_idx < 0) p05_idx = 0;
|
||||
if (p05_idx >= fill_n) p05_idx = fill_n - 1;
|
||||
if (p95_idx < 0) p95_idx = 0;
|
||||
if (p95_idx >= fill_n) p95_idx = fill_n - 1;
|
||||
|
||||
float p05_val, p95_val;
|
||||
|
||||
/* Bitonic sort path: exact, in-place, for power-of-2 fill_n. */
|
||||
bool is_pow2 = (fill_n > 1) && ((fill_n & (fill_n - 1)) == 0);
|
||||
|
||||
if (is_pow2) {
|
||||
bitonic_sort_block(svals, fill_n);
|
||||
p05_val = svals[p05_idx];
|
||||
p95_val = svals[p95_idx];
|
||||
} else {
|
||||
/* Quickselect path (two independent passes; use svals_copy for the second
|
||||
* so svals is not clobbered between calls). */
|
||||
for (int i = 0; i < fill_n; i++) svals_copy[i] = svals[i];
|
||||
|
||||
p05_val = nth_element_float(svals, fill_n, p05_idx);
|
||||
p95_val = nth_element_float(svals_copy, fill_n, p95_idx);
|
||||
}
|
||||
|
||||
/* ── Phase 3: EMA update into ISV[p05_slot] and ISV[p95_slot].
|
||||
* Read both previous values before any write (handles isv == isv_out aliasing). */
|
||||
float alpha = fminf(fmaxf(ema_alpha, 0.01f), 0.5f);
|
||||
|
||||
float prev_p05 = isv[p05_slot];
|
||||
float prev_p95 = isv[p95_slot];
|
||||
|
||||
isv_out[p05_slot] = (1.0f - alpha) * prev_p05 + alpha * p05_val;
|
||||
isv_out[p95_slot] = (1.0f - alpha) * prev_p95 + alpha * p95_val;
|
||||
}
|
||||
@@ -3110,6 +3110,13 @@ impl FusedTrainingCtx {
|
||||
.map_err(|e| anyhow::anyhow!("launch_kelly_cap_update: {e}"))
|
||||
}
|
||||
|
||||
/// Plan 2 Task 1 C.1: GPU-driven per-branch Q P5/P95 percentile EMA update.
|
||||
/// Delegates to `GpuDqnTrainer::launch_q_quantile_reduce`. Cold-path per-epoch.
|
||||
pub(crate) fn launch_q_quantile_reduce(&self, ema_alpha: f32) -> anyhow::Result<()> {
|
||||
self.trainer.launch_q_quantile_reduce(ema_alpha)
|
||||
.map_err(|e| anyhow::anyhow!("launch_q_quantile_reduce: {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);
|
||||
|
||||
@@ -139,7 +139,7 @@ impl StateResetRegistry {
|
||||
RegistryEntry {
|
||||
name: "ISV_LAYOUT_FINGERPRINT",
|
||||
category: ResetCategory::SchemaContract,
|
||||
description: "ISV[47..49) — compile-time structural hash of slot layout; fail-fast only, no migration path",
|
||||
description: "ISV[55..57) — compile-time structural hash of slot layout; fail-fast only, no migration path",
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "isv_iql_branch_scale_floor",
|
||||
@@ -162,6 +162,11 @@ impl StateResetRegistry {
|
||||
description: "ISV[PLAN_THRESHOLD_INDEX=46] — plan-MLP activation threshold; Plan 3 B.4 may make reactive",
|
||||
},
|
||||
// ───── Fold-reset state (pre-allocated GPU-written slots) ───
|
||||
RegistryEntry {
|
||||
name: "isv_q_quantiles",
|
||||
category: ResetCategory::FoldReset,
|
||||
description: "ISV[47..55) — per-branch Q P5/P95 EMAs (Plan 2 C.1); reset to v_min/v_max bootstrap at fold boundary",
|
||||
},
|
||||
RegistryEntry {
|
||||
name: "isv_epoch_idx",
|
||||
category: ResetCategory::FoldReset,
|
||||
|
||||
@@ -1788,6 +1788,15 @@ impl DQNTrainer {
|
||||
// the only thing that updates the buffer — so it stays at
|
||||
// zero and the downstream health composition believes
|
||||
// q_gap=0 even when raw epoch q_gap peaks above 1.0.
|
||||
// Plan 2 Task 1 C.1: launch q_quantile_reduce AFTER populate_q_out
|
||||
// (done inside reduce_current_q_stats above) and BEFORE
|
||||
// update_eval_v_range reads the percentile slots from ISV.
|
||||
// EMA alpha 0.1: 10% of each new epoch's percentile estimate.
|
||||
if let Err(e) = fused.launch_q_quantile_reduce(0.1) {
|
||||
tracing::warn!(
|
||||
"launch_q_quantile_reduce at epoch-boundary failed (non-fatal): {e}"
|
||||
);
|
||||
}
|
||||
let per_branch_q_gaps = fused.get_per_branch_q_gaps();
|
||||
match fused.reduce_current_q_stats_per_branch() {
|
||||
Ok(per_branch_stats) => {
|
||||
@@ -4038,6 +4047,25 @@ impl DQNTrainer {
|
||||
}
|
||||
}
|
||||
}
|
||||
"isv_q_quantiles" => {
|
||||
// Reset Q-quantile EMA slots to v_min/v_max bootstrap at fold boundary.
|
||||
// q_p05 = v_min, q_p95 = v_max — matches cold-start atom range
|
||||
// so the first update_eval_v_range after fold start sees meaningful bounds.
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
let v_min_f = self.hyperparams.v_min as f32;
|
||||
let v_max_f = self.hyperparams.v_max as f32;
|
||||
for b in 0..4 {
|
||||
fused.trainer().write_isv_signal_at(
|
||||
crate::cuda_pipeline::gpu_dqn_trainer::Q_P05_DIR_INDEX + b,
|
||||
v_min_f,
|
||||
);
|
||||
fused.trainer().write_isv_signal_at(
|
||||
crate::cuda_pipeline::gpu_dqn_trainer::Q_P95_DIR_INDEX + b,
|
||||
v_max_f,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
"isv_epoch_idx" => {
|
||||
// Reset epoch-idx slot at fold boundary; epoch loop re-populates on next start.
|
||||
if let Some(ref fused) = self.fused_ctx {
|
||||
|
||||
@@ -54,6 +54,7 @@
|
||||
| `cuda_pipeline/kelly_cap_update_kernel.cu` | `GpuDqnTrainer::launch_kelly_cap_update` → `FusedTrainingCtx::launch_kelly_cap_update` → `training_loop.rs` epoch boundary; takes portfolio_states dev ptr + n_envs | Wired | Plan 1 Task 11 — cold-path (per-epoch) | — |
|
||||
| `trainers/dqn/monitors/atoms_monitor.rs` | Read-only observer for atoms_update kernel inputs (ISV v-range slots 23..31); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 9 (test initial-value fixup post-land; local-smoke StateResetRegistry dispatch arms + hp_f32 helper added) | — |
|
||||
| `cuda_pipeline/atoms_update_kernel.cu` | `GpuDqnTrainer::launch_atoms_update` → `FusedTrainingCtx::launch_atoms_update` → `training_loop.rs` epoch boundary; replaces per-branch `recompute_atom_positions` CPU loop | Wired | Plan 1 Task 9 — cold-path (per-epoch + SGD step) | — |
|
||||
| `cuda_pipeline/q_quantile_kernel.cu` | `GpuDqnTrainer::launch_q_quantile_reduce` → `FusedTrainingCtx::launch_q_quantile_reduce` → `training_loop.rs` epoch boundary; reads q_out_buf [N, total_actions], writes ISV[47..55) P5/P95 EMAs per branch | Wired | Plan 2 Task 1 C.1 — cold-path (per-epoch); consumer: `update_eval_v_range` reads Q_P05/Q_P95 to set quantile-based half-width | — |
|
||||
|
||||
## CUDA Pipeline — Rust Wrappers
|
||||
|
||||
|
||||
@@ -8,14 +8,14 @@ fail-fast; no migration functions exist. Backward compat is structurally
|
||||
unwritable (there is no ordered version space to pair migrations against).
|
||||
See spec §4.A.2.
|
||||
|
||||
**Tail placement rationale:** The fingerprint occupies the tail (`[47..49)`) rather
|
||||
**Tail placement rationale:** The fingerprint occupies the tail (`[55..57)`) rather
|
||||
than the head because `isv_signals[0]` and `isv_signals[1]` are actively written
|
||||
by `isv_signal_update` (Q-drift EMA and gradient-norm EMA respectively). Inserting
|
||||
at the head would displace those live signals and require shifting every upstream
|
||||
literal in `experience_kernels.cu`. The fingerprint moves to the new tail each time
|
||||
new slots are appended to the bus.
|
||||
|
||||
**Current `ISV_TOTAL_DIM`:** 49 (Plan 1 Tasks 12/15/16 + pre-allocation for Tasks 9-11/13/14). Post-full DQN v2 rollout: 72.
|
||||
**Current `ISV_TOTAL_DIM`:** 57 (Plan 1 + Plan 2 Task 1 C.1 quantile slots). Post-full DQN v2 rollout: 72.
|
||||
|
||||
| Index | Name constant | Type | Producer | Consumers | Reset-category | Notes |
|
||||
|---|---|---|---|---|---|---|
|
||||
@@ -48,6 +48,14 @@ new slots are appended to the bus.
|
||||
| [44] | `KELLY_CAP_EFF_INDEX` | f32 | GPU kelly-cap-adaptive kernel (follow-up) | experience_env_step | FoldReset | Effective Kelly cap; 0 at construction; GPU fills each step |
|
||||
| [45] | `CQL_ALPHA_INDEX` | f32 | construct (CPU static) | CQL adaptive formula in Rust | SchemaContract | CQL pessimism base coefficient; written from `config.cql_alpha`; read in `compute_cql_logit_gradients` |
|
||||
| [46] | `PLAN_THRESHOLD_INDEX` | f32 | construct (CPU static) | `experience_kernels.cu`, `backtest_plan_kernel.cu` | SchemaContract | Plan-MLP activation threshold; 0.5f at construction; read via `ISV_PLAN_THRESHOLD_IDX` in both kernels |
|
||||
| [47] | `ISV_LAYOUT_FINGERPRINT_LO_INDEX` | u32 bits (in f32) | construct | check_layout_fingerprint | SchemaContract | Low 32 bits of u64 FNV-1a structural hash. Fail-fast on mismatch — NOT a version number, no migration path. |
|
||||
| [48] | `ISV_LAYOUT_FINGERPRINT_HI_INDEX` | u32 bits (in f32) | construct | check_layout_fingerprint | SchemaContract | High 32 bits of u64 FNV-1a structural hash. |
|
||||
| [49..72) | (reserved for DQN v2) | | | | | Allocated incrementally by Plans 2-5 |
|
||||
| [47] | `Q_P05_DIR_INDEX` | f32 | q_quantile_reduce GPU kernel (cold-path per-epoch) | update_eval_v_range | FoldReset | Direction-branch 5th-percentile Q EMA. Bootstrap: v_min. |
|
||||
| [48] | `Q_P05_MAG_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Magnitude-branch 5th-percentile Q EMA. Bootstrap: v_min. |
|
||||
| [49] | `Q_P05_ORD_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Order-branch 5th-percentile Q EMA. Bootstrap: v_min. |
|
||||
| [50] | `Q_P05_URG_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Urgency-branch 5th-percentile Q EMA. Bootstrap: v_min. |
|
||||
| [51] | `Q_P95_DIR_INDEX` | f32 | q_quantile_reduce GPU kernel (cold-path per-epoch) | update_eval_v_range | FoldReset | Direction-branch 95th-percentile Q EMA. Bootstrap: v_max. |
|
||||
| [52] | `Q_P95_MAG_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Magnitude-branch 95th-percentile Q EMA. Bootstrap: v_max. |
|
||||
| [53] | `Q_P95_ORD_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Order-branch 95th-percentile Q EMA. Bootstrap: v_max. |
|
||||
| [54] | `Q_P95_URG_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Urgency-branch 95th-percentile Q EMA. Bootstrap: v_max. |
|
||||
| [55] | `ISV_LAYOUT_FINGERPRINT_LO_INDEX` | u32 bits (in f32) | construct | check_layout_fingerprint | SchemaContract | Low 32 bits of u64 FNV-1a structural hash. Fail-fast on mismatch — NOT a version number, no migration path. |
|
||||
| [56] | `ISV_LAYOUT_FINGERPRINT_HI_INDEX` | u32 bits (in f32) | construct | check_layout_fingerprint | SchemaContract | High 32 bits of u64 FNV-1a structural hash. |
|
||||
| [57..72) | (reserved for DQN v2) | | | | | Allocated incrementally by Plans 2-5 |
|
||||
|
||||
Reference in New Issue
Block a user