diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 6ddd8f898..94d25ab08 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5685,6 +5685,12 @@ extern "C" __global__ void quantile_q_select( * Prevents entropy regularization from making quantile blend meaningless. */ float util_factor = fminf(fmaxf(util_ema / 0.5f, 0.0f), 1.0f); + /* SP17: cold-fail if num_atoms exceeds the register-array ceiling. Same + * NUM_ATOMS_MAX=128 contract as compute_expected_q above. */ + if (num_atoms > NUM_ATOMS_MAX) { + __trap(); + } + for (int d = 0; d < 4; d++) { int n_d = branch_sizes[d]; const float* branch_base = b_logits + (long long)branch_logit_offset @@ -5695,21 +5701,42 @@ extern "C" __global__ void quantile_q_select( float v_min = per_sample_support[support_base + 0]; float dz = per_sample_support[support_base + 2]; + /* SP17 mean-zero identifiability: per-atom mean of A across this + * branch's actions. Sample-local register reduction (one thread = + * one sample), no atomicAdd / shared mem / cross-thread sync per + * `feedback_no_atomicadd`. Identical contract to compute_expected_q + * — every advantage read in this kernel goes through + * `adv_a[z] - a_mean_per_atom[z]`. The three inner softmax passes + * (max → sum_exp → CDF) all migrate atomically per + * `feedback_no_partial_refactor`. */ + float a_mean_per_atom[NUM_ATOMS_MAX]; + float inv_n_d = 1.0f / (float)n_d; + for (int z = 0; z < num_atoms; z++) { + float s = 0.0f; + for (int a = 0; a < n_d; a++) { + s += branch_base[(long long)a * num_atoms + z]; + } + a_mean_per_atom[z] = s * inv_n_d; + } + for (int a = 0; a < n_d; a++) { const float* adv_a = branch_base + (long long)a * num_atoms; - /* Softmax over atoms */ + /* SP17: Pass 1 — max over centered logits (V[z] + A[a,z] − mean_a A[*,z]). */ float max_logit = -1e30f; for (int z = 0; z < num_atoms; z++) { - float logit = v_row[z] + adv_a[z]; + float a_centered = adv_a[z] - a_mean_per_atom[z]; /* SP17 */ + float logit = v_row[z] + a_centered; if (logit > max_logit) max_logit = logit; } + /* SP17: Pass 2 — same centering. */ float sum_exp = 0.0f; for (int z = 0; z < num_atoms; z++) { - sum_exp += expf((v_row[z] + adv_a[z]) - max_logit); + float a_centered = adv_a[z] - a_mean_per_atom[z]; /* SP17 */ + sum_exp += expf((v_row[z] + a_centered) - max_logit); } - /* Build CDF and extract quantiles */ + /* SP17: Pass 3 — CDF over centered probs. */ float cdf = 0.0f; float q10 = v_min; float q50 = v_min; @@ -5717,7 +5744,8 @@ extern "C" __global__ void quantile_q_select( int got10 = 0, got50 = 0, got90 = 0; for (int z = 0; z < num_atoms; z++) { - float prob = expf((v_row[z] + adv_a[z]) - max_logit) / sum_exp; + float a_centered = adv_a[z] - a_mean_per_atom[z]; /* SP17 */ + float prob = expf((v_row[z] + a_centered) - max_logit) / sum_exp; cdf += prob; float z_val = v_min + (float)z * dz; if (!got10 && cdf >= 0.10f) { q10 = z_val; got10 = 1; } diff --git a/crates/ml/tests/sp17_dueling_oracle_tests.rs b/crates/ml/tests/sp17_dueling_oracle_tests.rs index 7d5199c78..13e15d168 100644 --- a/crates/ml/tests/sp17_dueling_oracle_tests.rs +++ b/crates/ml/tests/sp17_dueling_oracle_tests.rs @@ -407,4 +407,217 @@ mod gpu { ); } } + + fn load_quantile_q_select(stream: &Arc) -> CudaFunction { + let module = stream + .context() + .load_cubin(SP17_EXPERIENCE_CUBIN.to_vec()) + .expect("load experience_kernels cubin"); + module + .load_function("quantile_q_select") + .expect("load quantile_q_select function") + } + + // ── B.2: quantile_q_select with mean-zero centered A matches CPU oracle ── + /// SP17 Commit A: post-migration `quantile_q_select` reads centered advantage + /// across all three inner softmax passes (max → sum_exp → CDF). Synthetic + /// inputs where the per-atom mean A is non-zero so centering changes the + /// emitted softmax and therefore the CDF and the q10/q50/q90 quantiles. + /// + /// Configuration: + /// - NA = 3 atoms at v_min=−1.0, dz=1.0 ⇒ atom positions [−1.0, 0.0, 1.0]. + /// - V = [0.0, 0.0, 0.0] (V row neutral so A drives the distribution). + /// - dir branch (B0=4) advantage logits constructed so the centered logits + /// for action 0 concentrate mass on atom 2 (z=+1) and the centered + /// logits for action 3 concentrate mass on atom 0 (z=−1). Other branches + /// carry zero advantage (single action each). + /// - `iqn_readiness=0` (alpha=0 → optimistic Q90), `util_ema=1.0` + /// (full quantile blend, util_factor=1). + /// + /// The CPU oracle computes the post-centering CDF for each action and + /// extracts q90 directly (alpha=0); the GPU result must match to ε=1e-5. + /// Pre-centering the answer would diverge because the raw softmax at each + /// action sees `V + A_raw` (a different, non-zero-mean distribution) and + /// produces different quantile boundaries. + #[test] + #[ignore = "requires GPU"] + fn quantile_q_select_centered_matches_cpu_oracle() { + let stream = make_test_stream(); + let kernel = load_quantile_q_select(&stream); + + const N: usize = 1; + const NA: usize = 3; + const B0: usize = 4; + const B1: usize = 1; + const B2: usize = 1; + const B3: usize = 1; + const TOTAL_ACTIONS: usize = B0 + B1 + B2 + B3; + + // V row neutral so centered A drives the entire distribution. + let v: [f32; NA] = [0.0, 0.0, 0.0]; + // Direction-branch advantage logits. The mean over a per atom z is: + // mean_z0 = (3 + 0 + 0 + 0)/4 = 0.75 + // mean_z1 = 0 + // mean_z2 = (0 + 0 + 0 + 3)/4 = 0.75 + // After centering, action 0 has logits [+2.25, 0, -0.75] (mass at z=0/1) + // and action 3 has logits [-0.75, 0, +2.25] (mass at z=1/2). Distinct + // CDFs ⇒ distinct quantile triples ⇒ a real test that centering ran. + let a_raw_dir: [[f32; NA]; B0] = [ + [3.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.0, 3.0], + ]; + + let dir_len = N * B0 * NA; + let mag_len = N * B1 * NA; + let ord_len = N * B2 * NA; + let urg_len = N * B3 * NA; + let total_b_len = dir_len + mag_len + ord_len + urg_len; + let mut b_logits_host = vec![0.0_f32; total_b_len]; + for a in 0..B0 { + for z in 0..NA { + b_logits_host[a * NA + z] = a_raw_dir[a][z]; + } + } + + let v_logits_host: Vec = v.iter().copied().collect(); + + // per_sample_support: [N, 4, 3] stride-12 = (v_min, v_max, delta_z) per branch. + // Branch 0 uses v_min=-1, v_max=1, dz=1 → linear atom positions [-1, 0, +1]. + // Other branches share the same support so quantile arithmetic stays sane. + let mut support_host = vec![0.0_f32; N * 4 * 3]; + for d in 0..4 { + support_host[d * 3 + 0] = -1.0; // v_min + support_host[d * 3 + 1] = 1.0; // v_max + support_host[d * 3 + 2] = 1.0; // delta_z + } + + // ── CPU oracle: centered softmax per action, extract q90 (alpha=0). ── + // util_factor = 1.0 (util_ema=1, threshold 0.5 → util_factor=1) so + // q_blended = quantile_blend = (1-alpha)*q90 + alpha*q10 = q90. + let mean_a_per_atom = [ + (a_raw_dir[0][0] + a_raw_dir[1][0] + a_raw_dir[2][0] + a_raw_dir[3][0]) / 4.0_f32, + (a_raw_dir[0][1] + a_raw_dir[1][1] + a_raw_dir[2][1] + a_raw_dir[3][1]) / 4.0_f32, + (a_raw_dir[0][2] + a_raw_dir[1][2] + a_raw_dir[2][2] + a_raw_dir[3][2]) / 4.0_f32, + ]; + let mut cpu_q_select = [0.0f32; B0]; + for a in 0..B0 { + // Stable softmax over centered logits (V + A_centered). + let mut logits = [0.0f32; NA]; + for z in 0..NA { + logits[z] = v[z] + (a_raw_dir[a][z] - mean_a_per_atom[z]); + } + let m = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let mut sum_exp = 0.0_f32; + for z in 0..NA { + sum_exp += (logits[z] - m).exp(); + } + // Build CDF and pick q10/q50/q90 with the same first-cross logic + // the kernel uses. + let v_min = -1.0_f32; + let dz = 1.0_f32; + let mut cdf = 0.0_f32; + let (mut q10, mut q50, mut q90) = (v_min, v_min, v_min); + let (mut g10, mut g50, mut g90) = (false, false, false); + for z in 0..NA { + let p = (logits[z] - m).exp() / sum_exp; + cdf += p; + let z_val = v_min + (z as f32) * dz; + if !g10 && cdf >= 0.10 { q10 = z_val; g10 = true; } + if !g50 && cdf >= 0.50 { q50 = z_val; g50 = true; } + if !g90 && cdf >= 0.90 { q90 = z_val; g90 = true; } + } + // alpha=0 (optimistic), util_factor=1 → q_blended = q90. + let _ = (q10, q50); + cpu_q_select[a] = q90; + } + + // ── Mapped-pinned device buffers ─────────────────────────────────── + let v_buf = unsafe { MappedF32Buffer::new(v_logits_host.len()) } + .expect("alloc v_logits buf"); + v_buf.write_from_slice(&v_logits_host); + let b_buf = unsafe { MappedF32Buffer::new(total_b_len) } + .expect("alloc b_logits buf"); + b_buf.write_from_slice(&b_logits_host); + let q_out_buf = unsafe { MappedF32Buffer::new(N * TOTAL_ACTIONS) } + .expect("alloc q_select buf"); + q_out_buf.write_from_slice(&vec![0.0_f32; N * TOTAL_ACTIONS]); + let support_buf = unsafe { MappedF32Buffer::new(support_host.len()) } + .expect("alloc support buf"); + support_buf.write_from_slice(&support_host); + + let n_i32: i32 = N as i32; + let na_i32: i32 = NA as i32; + let b0_i32: i32 = B0 as i32; + let b1_i32: i32 = B1 as i32; + let b2_i32: i32 = B2 as i32; + let b3_i32: i32 = B3 as i32; + let iqn_r: f32 = 0.0; // alpha = 0 ⇒ q90 selection + let util_e: f32 = 1.0; // util_factor = 1 ⇒ no q50 fallback + + let block_dim: u32 = 256; + let grid_dim: u32 = ((N as u32 + block_dim - 1) / block_dim).max(1); + + unsafe { + stream + .launch_builder(&kernel) + .arg(&v_buf.dev_ptr) + .arg(&b_buf.dev_ptr) + .arg(&q_out_buf.dev_ptr) + .arg(&n_i32) + .arg(&na_i32) + .arg(&b0_i32) + .arg(&b1_i32) + .arg(&b2_i32) + .arg(&b3_i32) + .arg(&support_buf.dev_ptr) + .arg(&iqn_r) + .arg(&util_e) + .launch(LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: 0, + }) + .expect("launch quantile_q_select"); + } + stream.synchronize().expect("sync after quantile_q_select"); + + let result = q_out_buf.read_all(); + assert_eq!(result.len(), N * TOTAL_ACTIONS); + + let eps = 1e-5_f32; + for a in 0..B0 { + let gpu = result[a]; + let diff = (gpu - cpu_q_select[a]).abs(); + assert!( + diff < eps, + "dir action {a}: GPU q_select {gpu} vs CPU oracle {} (diff {:.3e})", + cpu_q_select[a], diff, + ); + } + + // Also verify the centered-vs-raw distinction: had we NOT centered the + // advantage, action 1 (logits all-zero raw) would yield a uniform + // softmax over [-1, 0, 1] giving q90=+1.0; with centering, action 1 + // sees logits [−0.75, 0, −0.75] (post-centering) which still yields + // q90=+1 because the CDF at z=2 still hits 1.0 — but action 0's + // q90 must be ≤ 0 because mass concentrates at the negative-side + // atom under raw, vs the positive-side atom under centered. We + // assert this asymmetry to fail-loudly if centering regresses. + // (Action 0 raw-centered logits before V add: A_centered[0] = + // [3 - 0.75, 0 - 0, 0 - 0.75] = [+2.25, 0, -0.75]. Softmax peaks + // at z=0 ⇒ CDF crosses 0.10 immediately at z=0 (q10=-1), still + // reaches 0.90 by z=1 because most mass is at z=0. So q90 should be 0.0.) + assert!( + result[0] <= 0.0 + 1e-5, + "centering regression: action 0 q90 expected ≤ 0.0 (centered mass at z=0), got {}", + result[0], + ); + assert!( + result[3] >= 0.0 - 1e-5, + "centering regression: action 3 q90 expected ≥ 0.0 (centered mass at z=2), got {}", + result[3], + ); + } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 621b7b643..91c3e897c 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,40 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +## 2026-05-08 — SP17 Commit A: mean-zero identifiability in `quantile_q_select` + +Migrates the second of five raw-advantage consumers identified by the post-Task 1.2 audit. `quantile_q_select` (uncertainty-driven action selection from C51 atom CDFs) builds a per-(sample, branch, action) softmax over `V[z] + A[a, z]` across **three** inner passes — Pass 1 finds the max logit, Pass 2 normalises `sum_exp`, Pass 3 walks the CDF to extract `q10/q50/q90`. All three passes now read `adv_a[z] - a_mean_per_atom[z]` per the SP17 mean-zero contract. The plan flagged this as a hidden 6th consumer (Task 1.4/1.5 as authored covered only `c51_loss_kernel` + `c51_grad_kernel` + `mag_concat_qdir` + Thompson; the audit found that `quantile_q_select` reads raw advantage at lines 5704/5709/5720 across all 4 branches and was missed entirely). + +### What lands + +- **Kernel modification**: `crates/ml/src/cuda_pipeline/experience_kernels.cu::quantile_q_select` — added per-atom mean-A reduction inside the per-branch outer loop (BEFORE the inner 3-pass softmax/CDF block). All three passes now read `adv_a[z] - a_mean_per_atom[z]`. Sample-local register reduction (one thread per sample), `float a_mean_per_atom[NUM_ATOMS_MAX]` register array sized to the same `NUM_ATOMS_MAX=128` ceiling as `compute_expected_q`. Device-side `__trap()` if `num_atoms > NUM_ATOMS_MAX` — cold-fail per `feedback_no_atomicadd` + `feedback_no_quickfixes`. +- **GPU oracle test** (NEW): `crates/ml/tests/sp17_dueling_oracle_tests.rs::gpu::quantile_q_select_centered_matches_cpu_oracle`. N=1, NA=3, B0=4 with `a_raw_dir = [[3,0,0], [0,0,0], [0,0,0], [0,0,3]]` and `V = [0,0,0]`. The non-zero per-atom mean `[0.75, 0, 0.75]` makes raw and centered distributions diverge measurably. Test launches the production cubin with `iqn_readiness=0` (alpha→q90) and `util_ema=1.0` (no q50 fallback); each per-action q90 must match the CPU oracle to ε=1e-5. Two extra structural assertions: action 0 q90 ≤ 0 (centered mass at the negative-side atom) and action 3 q90 ≥ 0 (centered mass at the positive-side atom) — both fail loudly if a future regression strips centering. Mapped-pinned per `feedback_no_htod_htoh_only_mapped_pinned`. + +### Wire-up status — INTERIM + +| Consumer of `b_logits` raw advantage | Status this commit | +|---|---| +| `compute_expected_q` (Task 1.2) | MIGRATED | +| `quantile_q_select` (this commit) | **MIGRATED** | +| `c51_loss_batched` | already mean-zero (annotation pending — Commit E) | +| `c51_grad_kernel` | already mean-zero Jacobian (annotation pending — Commit E) | +| `mag_concat_qdir` | PENDING — Commit B | +| Thompson direction-select (`experience_action_select`) | PENDING — Commit C | +| `barrier_gradient_direction` (aux-CQL) | PENDING — Commit D | +| `ib_gradient_direction` (aux-CQL) | PENDING — Commit D | + +⚠ **Partial-refactor state acknowledged.** Per `feedback_no_partial_refactor`, the contract gap that Task 1.2 opened is still partially open; Commits B–D close the remaining four production consumers + Commit E annotates the two pre-SP17 already-centered c51_loss / c51_grad sites. Atomicity is at the BRANCH level — no L40S dispatch escapes `feat/sp17-dueling` while any consumer reads raw advantage. + +### Verification + +``` +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check --workspace → clean +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml --test sp17_dueling_oracle_tests --features cuda \ + -- --ignored --nocapture → 3 PASSED (RTX 3050 Ti) +``` + +Plan: `docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md` Task 1.4 (extended — plan undercounted consumers; audit by prior agent surfaced this). + ## 2026-05-08 — SP17 Task PP.2: dueling-Q identifiability ISV slots [474..483) Allocates the 9 ISV slots that drive the SP17 dueling-Q identifiability diagnostic chain. **Additive infrastructure only — no consumer wiring in this commit.** The Phase 1 mean-centering producer (atomic across `compute_expected_q` + `c51_loss` + `c51_grad` + `mag_concat_qdir`) lands in the next 4 commits per `feedback_no_partial_refactor`.