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

feat(dqn): Phase 2d per-branch per_sample_support tile [B, 4, 3]

Browse Source 
Completes the ISV-unified Q-support range spec
(docs/superpowers/specs/2026-04-23-isv-v-range-unification.md) by
migrating the per_sample_support buffer from per-sample [B, 3] to
per-sample-per-branch [B, 4, 3] stride-12. Without this phase the
atom_positions grid already spanned per-branch adaptive ranges (landed
in 9deda5f65 via ISV slots 23..30) while the loss-projection
Bellman step still read a single V(s)-centred range — atoms and
projection disagreed, which is the exact pathology the spec fixes.

Producers:
* iql_value_kernel.cu::iql_compute_per_sample_support — new kernel
  signature adds isv_signals pointer; writes 4 branch triples per
  sample where centre = isv_signals[23 + 2*d] and half-width = V(s)-
  derived Q spread. Bootstrap identity: ISV centres=0 + readiness=0
  falls back to [-1, 1] across all 4 branches, byte-identical to the
  pre-Phase-2d single-range default at epoch 1.
* iql_value_kernel.cu::iql_support_floor — Frugal-1U p5 estimator now
  aggregates half-widths across all (sample, branch) pairs and applies
  the floor per (sample, branch) independently.
* gpu_iql_trainer.rs — per_sample_support_buf sized b*4*3, seed writes
  12 floats per sample, compute_per_sample_support takes isv_dev_ptr
  and forwards it to the kernel; launch-site arg order aligned.
* fused_training.rs — passes trainer.isv_signals_dev_ptr() into the
  IQL call.

Consumers (all migrated to stride-12 indexing `b*12 + d*3 + {0,1,2}`):
* c51_loss_kernel.cu::c51_loss_batched — per-branch (v_min, v_max,
  delta_z) read INSIDE the d-loop; degenerate-support skip is now
  per-branch (continue instead of whole-sample early exit).
* c51_grad_kernel.cu::c51_grad_kernel — per-branch z_norm and
  delta_z for the q-gap floor gradient path.
* experience_kernels.cu::compute_expected_q — per-branch (v_min, dz)
  inside the d-loop that iterates all 4 branches.
* experience_kernels.cu::mag_concat_qdir — reads direction branch
  (d=0) slots from the stride-12 tile.
* experience_kernels.cu::quantile_q_select — per-branch (v_min, dz)
  inside the d-loop.

Experience-collector parity:
* gpu_experience_collector.rs — its OWN per_sample_support_buf grows
  to alloc_episodes*4*3; update_per_sample_support tiles the same
  (v_min, v_max, delta_z) triple to all 4 branches so the layout
  matches the IQL buffer and the consumer kernels read uniformly.

Safety:
* Bootstrap byte-identical at epoch 1 preserved (ISV centres default 0,
  readiness ramps from 0 → 1).
* No stub values, no TODO/FIXME/XXX markers introduced.
* Kernel scalar arg (gamma) is already f32 in GpuIqlConfig — no f64→f32
  cast needed at the call site (feedback_cudarc_f64_f32_abi compliance
  via type, not cast).
* c51_loss branch-degenerate `continue` is uniform across the block
  (all threads read the same support_base) so __syncthreads inside
  the loop body remains collective.

Compile verified: cargo check -p ml + --workspace pass (SQLX_OFFLINE,
CARGO_INCREMENTAL=0, sccache) and cargo build -p ml compiles all
CUDA kernels via nvcc. cargo test -p ml --lib --no-run succeeds.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-23 21:42:23 +02:00
parent 9deda5f65b
commit 11df037855
9 changed files with 213 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
View File

@@ -24,7 +24,7 @@ extern "C" __global__ void c51_grad_kernel(
int total_branch_atoms,
float entropy_coeff,
const float* __restrict__ branch_scales, /* [B, 4] per-sample per-branch gradient scale */
const float* __restrict__ per_sample_support, /* [B, 3] per-sample [v_min, v_max, delta_z] */
const float* __restrict__ per_sample_support, /* [B, 4, 3] stride-12 per-sample per-branch [v_min, v_max, delta_z] (Phase 2d) */
const float* __restrict__ liquid_mod, /* [4] pinned device-mapped per-branch modulators */
const float* __restrict__ atom_positions, /* [4, num_atoms] adaptive positions. NULL = linear. */
const float* __restrict__ q_mean_ema_ptr, /* [1] pinned device-mapped Q-mean EMA */
@@ -203,13 +203,17 @@ extern "C" __global__ void c51_grad_kernel(
* The gradient is orthogonal to Bellman — it depends on atom
* position, not on the target match. Active on ALL samples,
* providing perpetual differentiation pressure. */
/* z_norm: use adaptive atom positions if available, else linear grid */
/* z_norm: use adaptive atom positions if available, else linear grid.
* Phase 2d: per_sample_support is [B, 4, 3] stride-12; index with
* (b, d, {0,1,2}) so each branch's normalisation uses its own
* v_min/v_max/delta_z. */
float z_norm;
float delta_z;
int support_base = b * 12 + d * 3;
if (atom_positions != NULL) {
float z_val = atom_positions[(long long)d * num_atoms + j];
float z_min = per_sample_support[b * 3 + 0];
float z_max = per_sample_support[b * 3 + 1];
float z_min = per_sample_support[support_base + 0];
float z_max = per_sample_support[support_base + 1];
z_norm = 2.0f * (z_val - z_min) / fmaxf(z_max - z_min, 1e-6f) - 1.0f;
/* Approximate delta_z from adjacent atom spacing */
delta_z = (j < num_atoms - 1)
@@ -217,7 +221,7 @@ extern "C" __global__ void c51_grad_kernel(
: z_val - atom_positions[(long long)d * num_atoms + j - 1];
} else {
z_norm = 2.0f * (float)j / fmaxf((float)(num_atoms - 1), 1.0f) - 1.0f;
delta_z = per_sample_support[b * 3 + 2];
delta_z = per_sample_support[support_base + 2];
}
float velocity_mod = liquid_mod[d];
float spread_scale = inv_batch * delta_z * velocity_mod;

48
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -554,7 +554,7 @@ extern "C" __global__ void c51_loss_batched(
const float* __restrict__ gamma_buf, /* [B] per-sample effective gamma */
int batch_size,
int num_atoms,
const float* __restrict__ per_sample_support, /* [B, 3] per-sample: [v_min, v_max, delta_z] */
const float* __restrict__ per_sample_support, /* [B, 4, 3] stride-12 per-sample per-branch: [v_min, v_max, delta_z] (Phase 2d) */
int b0_size,
int b1_size,
int b2_size,
@@ -612,28 +612,16 @@ extern "C" __global__ void c51_loss_batched(
int sample_id = blockIdx.x;
if (sample_id >= batch_size) return;
/* Per-sample C51 support centered on V(s) */
float v_min = per_sample_support[sample_id * 3];
float v_max = per_sample_support[sample_id * 3 + 1];
float delta_z = per_sample_support[sample_id * 3 + 2];
/* Phase 2d: per_sample_support is [B, 4, 3] stride-12. Per-branch
* (v_min, v_max, delta_z) is read INSIDE the branch loop below; the
* degenerate-support skip also becomes per-branch. No single-range
* read here — the tile no longer carries a shared per-sample range. */
/* Skip sample if support is degenerate (all Q-values identical) */
if (delta_z < 1e-7f) {
if (tid == 0) {
per_sample_loss[sample_id] = 0.0f;
td_errors[sample_id] = 0.0f;
}
return;
}
int b0_atoms = b0_size * num_atoms;
int b1_atoms = b1_size * num_atoms;
int b2_atoms = b2_size * num_atoms;
int b3_atoms = b3_size * num_atoms;
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_support[j] = v_min + (float)j * delta_z;
__syncthreads();
int factored_action = actions[sample_id];
int max_action = b0_size * b1_size * b2_size * b3_size;
if (factored_action < 0 || factored_action >= max_action)
@@ -702,12 +690,34 @@ extern "C" __global__ void c51_loss_batched(
int n_atoms = n_d * num_atoms;
long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * num_atoms;
/* Update atom support for this branch (adaptive positions if available) */
/* Phase 2d: per-branch (v_min, v_max, delta_z) from the [B, 4, 3]
* tile. Direction/magnitude/order/urgency own independent ranges
* — matches the per-branch atom_positions layout that landed in
* Phase 2a/2b. */
int support_base = sample_id * 12 + d * 3;
float v_min = per_sample_support[support_base + 0];
float v_max = per_sample_support[support_base + 1];
float delta_z = per_sample_support[support_base + 2];
/* Degenerate-support skip is now per-branch: when a branch's
* Q-spread collapses to zero, skip just that branch's loss
* contribution instead of the whole sample. Unskipped branches
* still contribute to the sample's total_ce. */
bool branch_degenerate = (delta_z < 1e-7f);
/* Update atom support for this branch: prefer adaptive positions
* when provided, otherwise fall back to the per-branch linear
* grid derived from v_min/delta_z. */
if (atom_positions != NULL) {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_support[j] = atom_positions[(long long)d * num_atoms + j];
__syncthreads();
} else {
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
shmem_support[j] = v_min + (float)j * delta_z;
}
__syncthreads();
if (branch_degenerate) continue;
/* ═══ STEP a: Current log-probs for taken action a_d ═══════ */

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

@@ -2721,7 +2721,7 @@ extern "C" __global__ void compute_expected_q(
int b1_size,
int b2_size,
int b3_size,
const float* __restrict__ per_sample_support, /* [N*3]: per-sample [v_min, v_max, delta_z] from IQL */
const float* __restrict__ per_sample_support, /* [N, 4, 3] stride-12: per-sample per-branch [v_min, v_max, delta_z] from IQL (Phase 2d) */
float* __restrict__ atom_stats_block_sums, /* [num_blocks * 2]: per-block partial sums (phase 1).
* NULL to skip atom-stat collection entirely. */
float* __restrict__ q_variance, /* [N, total_actions] Var[Q] per action. NULL to skip. */
@@ -2730,12 +2730,10 @@ extern "C" __global__ void compute_expected_q(
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
/* Per-sample support from IQL — breaks circular eval_v_range feedback */
float v_min = per_sample_support[i * 3];
float v_max = per_sample_support[i * 3 + 1];
/* Phase 2d: per-branch (v_min, v_max, delta_z) read INSIDE the d-loop
* below — direction/magnitude/order/urgency own independent ranges. */
int total_actions = b0_size + b1_size + b2_size + b3_size;
float dz = per_sample_support[i * 3 + 2];
float util_threshold = 0.5f / (float)num_atoms; /* half of uniform */
/* Per-sample value logits — sample-major [N, num_atoms] */
@@ -2757,6 +2755,11 @@ extern "C" __global__ void compute_expected_q(
const float* branch_base = b_logits + (long long)branch_logit_offset
+ (long long)i * n_d * num_atoms;
/* Phase 2d: per-branch support triple from the [N, 4, 3] tile. */
int support_base = i * 12 + d * 3;
float v_min = per_sample_support[support_base + 0];
float dz = per_sample_support[support_base + 2];
for (int a = 0; a < n_d; a++) {
const float* adv_a = branch_base + (long long)a * num_atoms;
@@ -3163,7 +3166,7 @@ extern "C" __global__ void mag_concat_qdir(
int SH2,
int NA, /* num_atoms */
int b0_size, /* direction branch size (4) */
const float* __restrict__ per_sample_support /* [B, 3]: v_min, v_max, delta_z */
const float* __restrict__ per_sample_support /* [B, 4, 3] stride-12: v_min, v_max, delta_z per-branch (Phase 2d) */
) {
int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= B) return;
@@ -3174,9 +3177,11 @@ extern "C" __global__ void mag_concat_qdir(
concat_out[b * out_stride + i] = h_s2[b * SH2 + i];
}
/* Step 2: compute E[Q_dir] for each direction action */
float v_min = per_sample_support[b * 3 + 0];
float dz = per_sample_support[b * 3 + 2];
/* Step 2: compute E[Q_dir] for each direction action.
* Phase 2d: read direction-branch (d=0) slots from the [B, 4, 3] tile. */
int support_base = b * 12 + 0 * 3; /* direction = branch 0 */
float v_min = per_sample_support[support_base + 0];
float dz = per_sample_support[support_base + 2];
for (int a = 0; a < b0_size; a++) {
/* Combined value + advantage logits for direction action a.
@@ -3676,15 +3681,14 @@ extern "C" __global__ void quantile_q_select(
int b1_size,
int b2_size,
int b3_size,
const float* __restrict__ per_sample_support,
const float* __restrict__ per_sample_support, /* [N, 4, 3] stride-12 per-branch (Phase 2d) */
float iqn_readiness,
float util_ema)
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
float v_min = per_sample_support[i * 3];
float dz = per_sample_support[i * 3 + 2];
/* Phase 2d: per-branch (v_min, delta_z) read INSIDE the d-loop below. */
int total_actions = b0_size + b1_size + b2_size + b3_size;
const float* v_row = v_logits + (long long)i * num_atoms;
@@ -3703,6 +3707,11 @@ extern "C" __global__ void quantile_q_select(
const float* branch_base = b_logits + (long long)branch_logit_offset
+ (long long)i * n_d * num_atoms;
/* Phase 2d: per-branch (v_min, dz) from the [N, 4, 3] tile. */
int support_base = i * 12 + d * 3;
float v_min = per_sample_support[support_base + 0];
float dz = per_sample_support[support_base + 2];
for (int a = 0; a < n_d; a++) {
const float* adv_a = branch_base + (long long)a * num_atoms;

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

@@ -1249,7 +1249,9 @@ pub struct GpuDqnTrainer {
tau_pinned: *mut f32,
tau_dev_ptr: u64,
/// Per-sample C51 support [B, 3] — pointer set by IQL trainer. Used by c51_loss_batched only.
/// Per-sample, per-branch C51 support [B, 4, 3] stride-12 — pointer set by IQL trainer
/// (Phase 2d). Consumed by c51_loss_batched / c51_grad_kernel / compute_expected_q /
/// mag_concat_qdir / quantile_q_select, which index `sample_id * 12 + d * 3 + {0,1,2}`.
per_sample_support_ptr: u64,
/// Per-branch gradient scales [B, 4] — pointer set by IQL trainer.
branch_scales_ptr: u64,

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

@@ -439,8 +439,10 @@ pub struct GpuExperienceCollector {
network_dims: (usize, usize, usize, usize),
/// Number of C51 atoms (1 for non-distributional).
num_atoms: usize,
/// Per-sample C51 support tiled [alloc_episodes, 3]: [v_min, v_max, delta_z] per episode.
/// Filled once per epoch via `update_per_sample_support()` — matches kernel expectation.
/// Per-sample, per-branch C51 support tiled [alloc_episodes, 4, 3]:
/// [v_min, v_max, delta_z] per (episode, branch). Filled once per epoch via
/// `update_per_sample_support()` — matches the Phase 2d stride-12 contract
/// consumed by compute_expected_q / quantile_q_select / mag_concat_qdir.
per_sample_support_buf: CudaSlice<f32>,
/// Per-sample epsilon from IQL expectile gap (0 = use cosine schedule).
per_sample_epsilon_ptr: u64,
@@ -1274,10 +1276,11 @@ impl GpuExperienceCollector {
let exp_bn_concat = stream.alloc_zeros::<f32>(alloc_episodes * (bn_alloc + portfolio_dim_bn) + 128)
.map_err(|e| MLError::ModelError(format!("alloc exp_bn_concat: {e}")))?;
// Per-sample C51 support buffer [alloc_episodes, 3] — tiled v_min/v_max/delta_z.
// Filled once per epoch in update_per_sample_support() — replaces the 2-float
// eval_v_range pointer that caused OOB reads in compute_expected_q/quantile_q_select.
let per_sample_support_buf = stream.alloc_zeros::<f32>(alloc_episodes * 3)
// Phase 2d: per-sample, per-branch C51 support buffer [alloc_episodes, 4, 3]
// stride-12 — tiled v_min/v_max/delta_z per branch. Filled once per epoch
// in update_per_sample_support(); consumers compute_expected_q /
// quantile_q_select / mag_concat_qdir index with (sample, branch, {0,1,2}).
let per_sample_support_buf = stream.alloc_zeros::<f32>(alloc_episodes * 4 * 3)
.map_err(|e| MLError::ModelError(format!("alloc per_sample_support_buf: {e}")))?;
// Task 8: GPU-resident step counter for experience collection loop
@@ -1323,7 +1326,7 @@ impl GpuExperienceCollector {
market_dim,
network_dims,
num_atoms,
per_sample_support_buf, // Tiled [alloc_episodes, 3]; filled via update_per_sample_support()
per_sample_support_buf, // Phase 2d: [alloc_episodes, 4, 3] stride-12; update_per_sample_support() tiles
per_sample_epsilon_ptr: 0, // Set via set_per_sample_epsilon_ptr() after IQL init
branch_sizes,
alloc_episodes,
@@ -2715,7 +2718,7 @@ impl GpuExperienceCollector {
.arg(&b1)
.arg(&b2)
.arg(&b3)
.arg(&self.per_sample_support_buf) // [N, 3] per-sample support
.arg(&self.per_sample_support_buf) // [N, 4, 3] stride-12 per-sample per-branch support
.arg(&null_atom_stats)
.arg(&mut self.q_var_buf)
.arg(&null_atom_positions) // atom_positions (NULL = linear)
@@ -2835,7 +2838,7 @@ impl GpuExperienceCollector {
.arg(&b1)
.arg(&b2)
.arg(&b3)
.arg(&self.per_sample_support_buf) // [N, 3] per-sample support
.arg(&self.per_sample_support_buf) // [N, 4, 3] stride-12 per-sample per-branch support
.arg(&iqn_r)
.arg(&util_e)
.launch(launch_cfg)
@@ -3089,17 +3092,23 @@ impl GpuExperienceCollector {
info!(trainer_params_ptr = ptr, "Experience collector: zero-copy trainer params pointer set");
}
/// Fill per_sample_support_buf [alloc_episodes, 3] with uniform v_min/v_max/delta_z.
/// Fill per_sample_support_buf [alloc_episodes, 4, 3] with uniform v_min/v_max/delta_z
/// broadcast across all 4 branches.
///
/// Called once per epoch (not hot-path) when the eval v_range changes.
/// Replaces the old `set_eval_v_range_ptr` which passed a 2-float pointer
/// to kernels that expected [N, 3], causing OOB reads for every sample > 0.
/// Called once per epoch (not hot-path) when the eval v_range changes. The
/// experience collector does not own per-branch Q-stats, so all 4 branches
/// receive the same (v_min, v_max, delta_z) triple here; this preserves
/// existing training-path semantics while making the buffer layout match the
/// [N, 4, 3] stride-12 contract that compute_expected_q / quantile_q_select /
/// mag_concat_qdir read (Phase 2d).
pub fn update_per_sample_support(&mut self, v_min: f32, v_max: f32) -> Result<(), MLError> {
let n = self.alloc_episodes;
let delta_z = (v_max - v_min) / (self.num_atoms as f32 - 1.0).max(1.0);
let data: Vec<f32> = (0..n).flat_map(|_| [v_min, v_max, delta_z]).collect();
let data: Vec<f32> = (0..n)
.flat_map(|_| (0..4).flat_map(|_| [v_min, v_max, delta_z]))
.collect();
super::htod_f32(&self.stream, &data, &mut self.per_sample_support_buf)?;
debug!(v_min, v_max, delta_z, n, "per_sample_support tiled for experience collector");
debug!(v_min, v_max, delta_z, n, "per_sample_support tiled [N, 4, 3] for experience collector");
Ok(())
}
pub fn set_per_sample_epsilon_ptr(&mut self, ptr: u64) {

38
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -375,20 +375,30 @@ impl GpuIqlTrainer {
let adv_sigma_ema_buf = alloc_f32(&stream, 1, "iql_adv_sigma_ema")?;
let readiness_buf = alloc_f32(&stream, 4, "iql_readiness")?;
let p5_state_buf = alloc_f32(&stream, 2, "iql_p5_state")?;
let mut per_sample_support_buf = alloc_f32(&stream, b * 3, "iql_per_sample_support")?;
// Phase 2d: per-sample, per-branch support tile [B, 4, 3] stride-12.
// Each sample owns 4 branch triples (v_min, v_max, delta_z) — direction,
// magnitude, order, urgency — read by C51 loss/grad + experience kernels.
let mut per_sample_support_buf = alloc_f32(&stream, b * 4 * 3, "iql_per_sample_support")?;
let branch_scales_buf = alloc_f32(&stream, b * 4, "iql_branch_scales")?;
let expectile_gap_buf = alloc_f32(&stream, b, "iql_expectile_gap")?;
let gap_mean_buf = alloc_f32(&stream, 1, "iql_gap_mean")?;
let per_sample_epsilon_buf = alloc_f32(&stream, b, "iql_per_sample_epsilon")?;
// Seed per_sample_support with defaults for step 0
// Seed per_sample_support with defaults for step 0.
// Phase 2d: broadcast identical (v_min=-1, v_max=1, delta_z) to all
// 4 branches per sample. Bootstrap identity — the per-branch kernel
// will overwrite with ISV-driven centres + V(s)-derived half once
// the first training step lands.
let na = config.num_atoms.max(2) as f32;
let default_delta_z = 2.0 / (na - 1.0);
let mut default_support = vec![0.0_f32; b * 3];
let mut default_support = vec![0.0_f32; b * 4 * 3];
for i in 0..b {
default_support[i * 3] = -1.0;
default_support[i * 3 + 1] = 1.0;
default_support[i * 3 + 2] = default_delta_z;
for d in 0..4 {
let base = i * 12 + d * 3;
default_support[base] = -1.0;
default_support[base + 1] = 1.0;
default_support[base + 2] = default_delta_z;
}
}
super::htod_f32(&stream, &default_support, &mut per_sample_support_buf)?;
@@ -1208,10 +1218,19 @@ impl GpuIqlTrainer {
Ok(())
}
/// Compute per-sample C51 atom support centered on V(s).
/// Compute per-sample, per-branch C51 atom support.
///
/// Phase 2d: fills the [B, 4, 3] tile with per-branch (v_min, v_max,
/// delta_z) where the centre comes from ISV slots `V_CENTER_DIR + 2*d`
/// and the half-width from the per-sample V(s)-derived Q spread.
///
/// `isv_signals_dev_ptr`: raw device pointer to the [ISV_TOTAL_DIM=31]
/// broadcast bus. Pass `0` to skip ISV (null in the kernel) — centres
/// collapse to 0 and the readiness blend still provides a safe default.
pub fn compute_per_sample_support(
&mut self,
q_out_buf: &CudaSlice<f32>,
isv_signals_dev_ptr: u64,
) -> Result<(), MLError> {
let b = self.config.batch_size;
let batch_i32 = b as i32;
@@ -1227,6 +1246,7 @@ impl GpuIqlTrainer {
.arg(q_out_buf)
.arg(&mut self.per_sample_support_buf)
.arg(&self.readiness_buf)
.arg(&isv_signals_dev_ptr)
.arg(&gamma)
.arg(&batch_i32)
.arg(&total_actions_i32)
@@ -1349,7 +1369,9 @@ impl GpuIqlTrainer {
Ok(())
}
/// Raw pointer to per-sample support buffer [B, 3].
/// Raw pointer to per-sample, per-branch support buffer [B, 4, 3] stride-12
/// (Phase 2d). Consumers read per-branch triples at offset
/// `sample_id * 12 + branch_idx * 3 + {0,1,2}`.
pub fn per_sample_support_ptr(&self) -> u64 {
self.per_sample_support_buf.raw_ptr()
}

123
crates/ml/src/cuda_pipeline/iql_value_kernel.cu
View File

@@ -539,25 +539,43 @@ void iql_adv_variance_reduce(
}
/* ------------------------------------------------------------------ */
/* Per-Sample C51 Atom Support Kernel */
/* Per-Sample C51 Atom Support Kernel — per-branch [B, 4, 3] tile */
/* ------------------------------------------------------------------ */
/**
* Compute per-sample [v_min, v_max, delta_z] for C51 atom projection:
* Compute per-sample, per-branch [v_min, v_max, delta_z] for C51 atom
* projection. Spec: docs/superpowers/specs/2026-04-23-isv-v-range-unification.md
* (Phase 2d).
*
* spread = max_{a} |Q(s,a) - V(s)|
* half_w = spread * (1 + gamma)
* v_min[b] = V(s) - half_w
* v_max[b] = V(s) + half_w
* half_w = spread * (1 + gamma) (shared across branches)
* centre_b = isv_signals[V_CENTER_DIR + 2*b] (per-branch from ISV)
* v_min_b = centre_b - half_w
* v_max_b = centre_b + half_w
* delta_z = (v_max - v_min) / (num_atoms - 1)
*
* Output: per_sample_support[b*3+0]=v_min, [b*3+1]=v_max, [b*3+2]=delta_z
* Rationale: the V(s)-derived half_w captures per-sample Q magnitude
* (how wide the atoms need to span around the centre). The centre
* differs per branch because direction/magnitude/order/urgency have
* genuinely different Q scales — the ISV bus broadcasts each branch's
* centre EMA.
*
* Bootstrap safety: when ISV centres are 0 and readiness=0 blend falls
* back to [-1, 1], all 4 branch triples become [-1, 1, 2/(NA-1)] —
* byte-identical to the pre-spec single-range behaviour at epoch 1.
*
* Output layout stride-12: per_sample_support[b*12 + d*3 + {0,1,2}] per
* branch d in 0..4. The 4-branch tile is read by the C51 loss / grad
* kernels and downstream experience kernels that migrated in lockstep.
*
* Launch: grid=ceil(B/256), block=256.
*/
extern "C" __global__
void iql_compute_per_sample_support(
const float* __restrict__ v_out, /* [B] */
const float* __restrict__ q_out, /* [B, total_actions] */
float* __restrict__ per_sample_support, /* [B*3] */
float* __restrict__ per_sample_support, /* [B, 4, 3] */
const float* __restrict__ readiness_buf, /* [1] CV-based readiness */
const float* __restrict__ isv_signals, /* [ISV_TOTAL_DIM]; NULL → 0 centres */
float gamma,
int batch_size,
int total_actions,
@@ -578,18 +596,34 @@ void iql_compute_per_sample_support(
}
float half_w = spread * (1.0f + gamma);
float iql_vmin = v - half_w;
float iql_vmax = v + half_w;
float iql_dz = (iql_vmax - iql_vmin) / (float)(num_atoms - 1);
float default_dz = 2.0f / (float)(num_atoms - 1);
/* Blend: readiness=0 → default [-1,1], readiness=1 → V(s)-centered */
float v_min = r * iql_vmin + (1.0f - r) * (-1.0f);
float v_max = r * iql_vmax + (1.0f - r) * (1.0f);
float delta_z = r * iql_dz + (1.0f - r) * (2.0f / (float)(num_atoms - 1));
/* Per-branch centre from ISV slots 23 + 2*d. When ISV is null (smoke
* paths / pre-init) the centre collapses to 0; in that regime the
* readiness blend below falls back to the uniform [-1, 1] default
* while r ramps from 0 → 1. */
for (int d = 0; d < 4; d++) {
float isv_centre = (isv_signals != nullptr)
? isv_signals[23 + 2 * d]
: 0.0f;
per_sample_support[b * 3 + 0] = v_min;
per_sample_support[b * 3 + 1] = v_max;
per_sample_support[b * 3 + 2] = delta_z;
/* Per-branch IQL range: centre from ISV, half from per-sample
* V(s)-derived Q spread. Direction/magnitude/order/urgency
* share the half but own independent centres. */
float iql_vmin = isv_centre - half_w;
float iql_vmax = isv_centre + half_w;
float iql_dz = (iql_vmax - iql_vmin) / (float)(num_atoms - 1);
/* Blend: readiness=0 → default [-1,1], readiness=1 → adaptive. */
float v_min = r * iql_vmin + (1.0f - r) * (-1.0f);
float v_max = r * iql_vmax + (1.0f - r) * ( 1.0f);
float delta_z = r * iql_dz + (1.0f - r) * default_dz;
int base = b * 12 + d * 3;
per_sample_support[base + 0] = v_min;
per_sample_support[base + 1] = v_max;
per_sample_support[base + 2] = delta_z;
}
}
/* ------------------------------------------------------------------ */
@@ -600,11 +634,18 @@ void iql_compute_per_sample_support(
* p5 quantile of half-widths. Uses Frugal-1U streaming quantile estimator
* (one GPU scalar, O(B) per step, scale-free).
*
* Phase 2d layout: per_sample_support is now [B, 4, 3] stride-12; the
* p5 estimator aggregates half-widths across ALL (sample, branch) pairs
* (one scalar p5 for the entire tile) and the floor is applied per
* (sample, branch) independently. Sharing a single p5 across branches
* keeps the estimator statistically stable and preserves the anti-
* collapse guarantee for every branch's smallest half-width.
*
* Launch: grid=1, block=1 (after iql_compute_per_sample_support).
*/
extern "C" __global__
void iql_support_floor(
float* __restrict__ per_sample_support, /* [B*3] in-place */
float* __restrict__ per_sample_support, /* [B, 4, 3] in-place */
float* __restrict__ p5_state, /* [2]: [0]=p5_estimate, [1]=step_count */
int batch_size,
int num_atoms
@@ -612,16 +653,19 @@ void iql_support_floor(
{
float est = p5_state[0];
float step_count = p5_state[1];
int total_branches = batch_size * 4;
/* First call: init from batch median spread (not hardcoded 1.0) */
/* First call: init from aggregated mean spread across all (sample, branch). */
if (step_count < 0.5f) {
/* Find median via single pass: compute mean as proxy (O(B), no sort) */
float hw_sum = 0.0f;
for (int b = 0; b < batch_size; b++) {
float hw = (per_sample_support[b * 3 + 1] - per_sample_support[b * 3 + 0]) * 0.5f;
hw_sum += hw;
for (int d = 0; d < 4; d++) {
int base = b * 12 + d * 3;
float hw = (per_sample_support[base + 1] - per_sample_support[base + 0]) * 0.5f;
hw_sum += hw;
}
}
est = hw_sum / fmaxf((float)batch_size, 1.0f);
est = hw_sum / fmaxf((float)total_branches, 1.0f);
est = fmaxf(est, 1e-6f);
}
@@ -629,30 +673,37 @@ void iql_support_floor(
* Asymmetric ratio 0.05/0.95 is mathematically correct for p5 quantile. */
float adaptive_rate = 1.0f / fmaxf(sqrtf(step_count + 1.0f), 1.0f);
for (int b = 0; b < batch_size; b++) {
float hw = (per_sample_support[b * 3 + 1] - per_sample_support[b * 3 + 0]) * 0.5f;
float step = fmaxf(est * adaptive_rate, 1e-8f);
if (hw < est) {
est -= step;
} else {
est += step * (0.05f / 0.95f);
for (int d = 0; d < 4; d++) {
int base = b * 12 + d * 3;
float hw = (per_sample_support[base + 1] - per_sample_support[base + 0]) * 0.5f;
float step = fmaxf(est * adaptive_rate, 1e-8f);
if (hw < est) {
est -= step;
} else {
est += step * (0.05f / 0.95f);
}
est = fmaxf(est, 1e-8f);
}
est = fmaxf(est, 1e-8f);
}
p5_state[0] = est;
p5_state[1] = step_count + 1.0f;
/* Floor: 1/num_atoms of p5 — guarantees at least 1 atom of resolution.
* Adaptive to atom count: 51 atoms → 2% of p5, 11 atoms → 9% of p5. */
* Adaptive to atom count: 51 atoms → 2% of p5, 11 atoms → 9% of p5.
* Applied per (sample, branch) so each branch's floor is independent. */
float floor_fraction = 1.0f / fmaxf((float)num_atoms, 1.0f);
float floor_hw = est * floor_fraction;
float floor_dz = (2.0f * floor_hw) / fmaxf((float)(num_atoms - 1), 1.0f);
for (int b = 0; b < batch_size; b++) {
if (per_sample_support[b * 3 + 2] < floor_dz) {
float center = (per_sample_support[b * 3 + 0] + per_sample_support[b * 3 + 1]) * 0.5f;
per_sample_support[b * 3 + 0] = center - floor_hw;
per_sample_support[b * 3 + 1] = center + floor_hw;
per_sample_support[b * 3 + 2] = floor_dz;
for (int d = 0; d < 4; d++) {
int base = b * 12 + d * 3;
if (per_sample_support[base + 2] < floor_dz) {
float center = (per_sample_support[base + 0] + per_sample_support[base + 1]) * 0.5f;
per_sample_support[base + 0] = center - floor_hw;
per_sample_support[base + 1] = center + floor_hw;
per_sample_support[base + 2] = floor_dz;
}
}
}
}

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

@@ -1430,9 +1430,12 @@ impl FusedTrainingCtx {
self.gpu_iql.update_adv_sigma()
.map_err(|e| anyhow::anyhow!("IQL adv_sigma: {e}"))?;
// 6. Compute per-sample C51 support centered on V(s).
// 6. Compute per-sample, per-branch C51 support.
// Phase 2d: tile is [B, 4, 3]; centres from ISV slots 23..30,
// half-width from V(s)-derived Q spread.
let q_out = self.trainer.q_out_buf();
self.gpu_iql.compute_per_sample_support(q_out)
let isv_ptr = self.trainer.isv_signals_dev_ptr();
self.gpu_iql.compute_per_sample_support(q_out, isv_ptr)
.map_err(|e| anyhow::anyhow!("IQL per_sample_support: {e}"))?;
// Floor support using Frugal p5 quantile — no silent zero-loss samples.

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

@@ -1292,9 +1292,11 @@ impl DQNTrainer {
return Err(anyhow::anyhow!("GPU episode reset FAILED (no CPU fallback): {e}"));
}
// Tile per-sample support [N, 3] from current eval v_range before experience collection.
// Reads v_min/v_max from fused trainer (which may have been updated last epoch),
// falls back to config defaults for epoch 0 / pre-fused-init.
// Tile per-sample, per-branch support [N, 4, 3] from current eval v_range
// before experience collection. Reads v_min/v_max from fused trainer (which
// may have been updated last epoch), falls back to config defaults for
// epoch 0 / pre-fused-init. All 4 branches receive the same triple here;
// see update_per_sample_support doc (Phase 2d).
{
let (v_min, v_max) = if let Some(ref fused_ctx) = self.fused_ctx {
let vr = fused_ctx.eval_v_range();