diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 94d25ab08..2365b76c4 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5111,28 +5111,60 @@ extern "C" __global__ void mag_concat_qdir( float v_min = per_sample_support[support_base + 0]; float dz = per_sample_support[support_base + 2]; - /* Pass 1: softmax → E[Q] per direction action; stash for the RMS-match pass. */ + /* SP17 mean-zero identifiability: per-atom mean of A across the b0_size + * direction actions, computed sample-local from registers — same + * contract as compute_expected_q + quantile_q_select. The dueling-Q + * decomposition `Q = V + (A − mean_a A)` MUST be applied here too, + * otherwise the magnitude branch trunk-input sees a different dueling + * convention than the c51_loss/c51_grad backward path and the trunk + * learns a mixed-contract mapping. NUM_ATOMS_MAX guards the register + * array size (NA <= 128 — production uses 51). + * + * `MAG_CONCAT_MAX_DIR=4` already bounds the action loop above; we + * still iterate `b0_size` for ABI flexibility but the production + * caller passes b0_size=4. The atom dimension uses the kernel's NA + * parameter (no branch-major NUM_ATOMS_MAX gate added because dz/v_min + * already scale linearly in NA). */ + if (NA > NUM_ATOMS_MAX) { + __trap(); + } + float a_mean_per_atom[NUM_ATOMS_MAX]; + float inv_b0 = 1.0f / (float)b0_size; + for (int z = 0; z < NA; z++) { + float s = 0.0f; + for (int a = 0; a < b0_size; a++) { + s += dir_logits_b[(long long)b * b0_size * NA + (long long)a * NA + z]; + } + a_mean_per_atom[z] = s * inv_b0; + } + + /* Pass 1: softmax → E[Q] per direction action; stash for the RMS-match pass. + * SP17: every read of `dir_logits_b[..]` flows through `- a_mean_per_atom[z]` + * (atomic across all 3 passes — Pass 1 max, Pass 2 sum_exp, Pass 3 prob/z). */ float eq_local[MAG_CONCAT_MAX_DIR]; for (int a = 0; a < b0_size; a++) { - /* Combined value + advantage logits for direction action a. + /* Combined value + centered-advantage logits for direction action a. * Branch logits layout: BRANCH-MAJOR [N × b0 × NA], so action a * for sample b is at: dir_logits_b[b * b0_size * NA + a * NA + z] */ float max_logit = -1e30f; for (int z = 0; z < NA; z++) { - float logit = dir_logits_v[b * NA + z] - + dir_logits_b[b * b0_size * NA + a * NA + z]; + float a_centered = dir_logits_b[b * b0_size * NA + a * NA + z] + - a_mean_per_atom[z]; /* SP17 */ + float logit = dir_logits_v[b * NA + z] + a_centered; if (logit > max_logit) max_logit = logit; } float sum_exp = 0.0f; for (int z = 0; z < NA; z++) { - float logit = dir_logits_v[b * NA + z] - + dir_logits_b[b * b0_size * NA + a * NA + z]; + float a_centered = dir_logits_b[b * b0_size * NA + a * NA + z] + - a_mean_per_atom[z]; /* SP17 */ + float logit = dir_logits_v[b * NA + z] + a_centered; sum_exp += expf(logit - max_logit); } float eq = 0.0f; for (int z = 0; z < NA; z++) { - float logit = dir_logits_v[b * NA + z] - + dir_logits_b[b * b0_size * NA + a * NA + z]; + float a_centered = dir_logits_b[b * b0_size * NA + a * NA + z] + - a_mean_per_atom[z]; /* SP17 */ + float logit = dir_logits_v[b * NA + z] + a_centered; float prob = expf(logit - max_logit) / sum_exp; float z_val = v_min + (float)z * dz; eq += prob * z_val; diff --git a/crates/ml/tests/sp17_dueling_oracle_tests.rs b/crates/ml/tests/sp17_dueling_oracle_tests.rs index 13e15d168..cbfeb0f84 100644 --- a/crates/ml/tests/sp17_dueling_oracle_tests.rs +++ b/crates/ml/tests/sp17_dueling_oracle_tests.rs @@ -418,6 +418,16 @@ mod gpu { .expect("load quantile_q_select function") } + fn load_mag_concat_qdir(stream: &Arc) -> CudaFunction { + let module = stream + .context() + .load_cubin(SP17_EXPERIENCE_CUBIN.to_vec()) + .expect("load experience_kernels cubin"); + module + .load_function("mag_concat_qdir") + .expect("load mag_concat_qdir 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 @@ -620,4 +630,193 @@ mod gpu { result[3], ); } + + // ── B.3: mag_concat_qdir with mean-zero centered A matches CPU oracle ── + /// SP17 Commit B: post-migration `mag_concat_qdir` reads centered direction + /// advantage across all three inner softmax passes (max → sum_exp → E[Q]). + /// The kernel concatenates `h_s2` with the per-action E[Q_dir], scaled to + /// match the trunk-output RMS via `ISV[H_S2_RMS_EMA_INDEX=96]`. Centering + /// ensures the magnitude branch's input matches the dueling decomposition + /// the c51_loss/c51_grad backward path is already using. + /// + /// Configuration: + /// - B=1, SH2=2 (tiny trunk activation), NA=3, b0_size=4 + /// - h_s2 = [1.0, 2.0] — copied verbatim into concat[0..SH2] + /// - V row = [0,0,0] + /// - Direction advantage: same `[[3,0,0], [0,0,0], [0,0,0], [0,0,3]]` + /// pattern as the quantile_q_select test so per-atom mean A is non-zero + /// (`[0.75, 0, 0.75]`). + /// - per_sample_support[0..3] = (v_min=-1, v_max=1, dz=1). + /// - ISV slot 96 = 1.0 (cold-start neutral) → trunk RMS target 1.0. + /// + /// The CPU oracle computes the four centered E[Q_dir], then applies the + /// q_rms scaling `(1.0 / q_rms)` so that the per-sample RMS over the four + /// scaled values equals 1.0. The GPU result must match to ε=1e-5. + #[test] + #[ignore = "requires GPU"] + fn mag_concat_qdir_centered_matches_cpu_oracle() { + let stream = make_test_stream(); + let kernel = load_mag_concat_qdir(&stream); + + const B: usize = 1; + const SH2: usize = 2; + const NA: usize = 3; + const B0: usize = 4; + const ISV_DIM: usize = 483; // matches sp14_isv_slots.rs ISV_TOTAL_DIM + const H_S2_RMS_EMA_INDEX: usize = 96; + + let h_s2_host: Vec = vec![1.0, 2.0]; + let v_logits_host: Vec = vec![0.0; NA]; + 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 mut b_logits_host = vec![0.0_f32; B * B0 * NA]; + for a in 0..B0 { + for z in 0..NA { + b_logits_host[a * NA + z] = a_raw_dir[a][z]; + } + } + + // per_sample_support [B, 4, 3]: only branch 0 (direction) consumed. + let mut support_host = vec![0.0_f32; B * 4 * 3]; + for d in 0..4 { + support_host[d * 3 + 0] = -1.0; + support_host[d * 3 + 1] = 1.0; + support_host[d * 3 + 2] = 1.0; + } + + // ISV bus: cold-start neutral RMS target = 1.0. + let mut isv_host = vec![0.0_f32; ISV_DIM]; + isv_host[H_S2_RMS_EMA_INDEX] = 1.0; + + // ── CPU oracle: centered softmax → E[Q_dir], then RMS scale to 1.0. ── + let mean_a = [ + (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 v_min = -1.0_f32; + let dz = 1.0_f32; + let mut eq_centered = [0.0_f32; B0]; + for a in 0..B0 { + let mut logits = [0.0f32; NA]; + for z in 0..NA { + logits[z] = 0.0 + (a_raw_dir[a][z] - mean_a[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(); + } + let mut eq = 0.0_f32; + for z in 0..NA { + let p = (logits[z] - m).exp() / sum_exp; + let z_val = v_min + (z as f32) * dz; + eq += p * z_val; + } + eq_centered[a] = eq; + } + let q_sum_sq: f32 = eq_centered.iter().map(|x| x * x).sum(); + let q_rms = (q_sum_sq / B0 as f32).sqrt(); + // Scale: q_rms > 1e-6 → (h_s2_rms_ema / q_rms) where h_s2_rms_ema = 1.0 + let scale = if q_rms > 1e-6 { 1.0 / q_rms } else { 1.0 / dz.max(1e-6) }; + let mut cpu_concat_tail = [0.0_f32; B0]; + for a in 0..B0 { + cpu_concat_tail[a] = eq_centered[a] * scale; + } + + // ── Mapped-pinned device buffers ─────────────────────────────────── + let h_s2_buf = unsafe { MappedF32Buffer::new(B * SH2) } + .expect("alloc h_s2 buf"); + h_s2_buf.write_from_slice(&h_s2_host); + let v_buf = unsafe { MappedF32Buffer::new(B * NA) } + .expect("alloc v_logits buf"); + v_buf.write_from_slice(&v_logits_host); + let b_buf = unsafe { MappedF32Buffer::new(B * B0 * NA) } + .expect("alloc b_logits buf"); + b_buf.write_from_slice(&b_logits_host); + let out_buf = unsafe { MappedF32Buffer::new(B * (SH2 + B0)) } + .expect("alloc concat_out buf"); + out_buf.write_from_slice(&vec![0.0_f32; B * (SH2 + B0)]); + let support_buf = unsafe { MappedF32Buffer::new(support_host.len()) } + .expect("alloc support buf"); + support_buf.write_from_slice(&support_host); + let isv_buf = unsafe { MappedF32Buffer::new(ISV_DIM) } + .expect("alloc isv buf"); + isv_buf.write_from_slice(&isv_host); + + let b_i32: i32 = B as i32; + let sh2_i32: i32 = SH2 as i32; + let na_i32: i32 = NA as i32; + let b0_i32: i32 = B0 as i32; + let isv_idx_i32: i32 = H_S2_RMS_EMA_INDEX as i32; + + let block_dim: u32 = 256; + let grid_dim: u32 = ((B as u32 + block_dim - 1) / block_dim).max(1); + + unsafe { + stream + .launch_builder(&kernel) + .arg(&h_s2_buf.dev_ptr) + .arg(&v_buf.dev_ptr) + .arg(&b_buf.dev_ptr) + .arg(&out_buf.dev_ptr) + .arg(&b_i32) + .arg(&sh2_i32) + .arg(&na_i32) + .arg(&b0_i32) + .arg(&support_buf.dev_ptr) + .arg(&isv_buf.dev_ptr) + .arg(&isv_idx_i32) + .launch(LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: 0, + }) + .expect("launch mag_concat_qdir"); + } + stream.synchronize().expect("sync after mag_concat_qdir"); + + let result = out_buf.read_all(); + assert_eq!(result.len(), B * (SH2 + B0)); + + // h_s2 prefix copied verbatim. + let eps = 1e-5_f32; + for i in 0..SH2 { + let diff = (result[i] - h_s2_host[i]).abs(); + assert!( + diff < eps, + "h_s2 prefix slot {i}: GPU {} vs host {} (diff {:.3e})", + result[i], h_s2_host[i], diff, + ); + } + // E[Q_dir] tail (post-centering, post-RMS scale). + for a in 0..B0 { + let gpu = result[SH2 + a]; + let cpu = cpu_concat_tail[a]; + let diff = (gpu - cpu).abs(); + assert!( + diff < eps, + "Q_dir tail action {a}: GPU {} vs CPU oracle {} (diff {:.3e})", + gpu, cpu, diff, + ); + } + // Centering structural check: post-centering action 0 has E[Q] < 0 + // (centered logits [+2.25, 0, -0.75] put mass on z=0 ⇒ negative E[Q]) + // and action 3 has E[Q] > 0 (mirror). After RMS scale the sign + // structure is preserved. + assert!( + result[SH2 + 0] < 0.0 - 1e-5, + "centering regression: tail[0] expected < 0 (centered mass at z=0), got {}", + result[SH2 + 0], + ); + assert!( + result[SH2 + 3] > 0.0 + 1e-5, + "centering regression: tail[3] expected > 0 (centered mass at z=2), got {}", + result[SH2 + 3], + ); + } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 91c3e897c..f147081d3 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,38 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +## 2026-05-08 — SP17 Commit B: mean-zero identifiability in `mag_concat_qdir` + +Migrates the magnitude-branch's direction-Q conditioning input to the SP17 mean-zero contract. `mag_concat_qdir` builds a per-(sample, action) softmax over `V[z] + A_dir[a, z]` across **three** inner passes (max → sum_exp → E[Q]) and concatenates the four E[Q_dir] values onto `h_s2` as the magnitude-branch input. The post-Plan-4 `q_rms` adaptive scale (ISV[H_S2_RMS_EMA_INDEX=96]) is preserved unchanged — only the per-atom mean-A subtraction inside the softmax inputs is added. If the magnitude branch saw raw E[Q_dir] while c51_loss already saw centered, the trunk would learn a mixed-contract mapping per `feedback_no_partial_refactor`. + +### What lands + +- **Kernel modification**: `crates/ml/src/cuda_pipeline/experience_kernels.cu::mag_concat_qdir` — added per-atom mean-A reduction over the b0_size direction actions (sample-local register array `float a_mean_per_atom[NUM_ATOMS_MAX]`). All three inner passes now read `dir_logits_b[..] - a_mean_per_atom[z]`. NA cold-fail via `__trap()` mirrors the contract from `compute_expected_q` and `quantile_q_select`. The `q_rms` adaptive scale path is left intact; centering changes the per-action E[Q_dir] values that feed it but does not alter the scale formula itself. No change to bw path because backward into the direction logits flows through c51_grad (already centered). +- **GPU oracle test** (NEW): `crates/ml/tests/sp17_dueling_oracle_tests.rs::gpu::mag_concat_qdir_centered_matches_cpu_oracle`. B=1, SH2=2, NA=3, B0=4 with the same `[[3,0,0],…,[0,0,3]]` synthetic pattern. The test verifies (a) `h_s2[0..SH2]` copied verbatim into the concat prefix, (b) the 4 tail slots match the CPU oracle's `centered_E[Q_dir] × (h_s2_rms_ema / q_rms)` to ε=1e-5, and (c) tail[0] < 0 / tail[3] > 0 sign asymmetry — fail-loud regression detector. ISV[96] = 1.0 (cold-start neutral) so the RMS-scale becomes `1 / q_rms`. 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` (Commit A) | MIGRATED | +| `mag_concat_qdir` (this commit) | **MIGRATED** | +| `c51_loss_batched` | already mean-zero (annotation pending — Commit E) | +| `c51_grad_kernel` | already mean-zero Jacobian (annotation pending — Commit E) | +| 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 | + +### 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 → 4 PASSED (RTX 3050 Ti) +``` + +Plan: `docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md` Task 1.5 (split: mag_concat is Commit B, Thompson is Commit C). + ## 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).