From eabcf8d529542fcf3f7671be25ae235b3bfadf15 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 8 May 2026 21:13:16 +0200 Subject: [PATCH] feat(sp17): mean-zero identifiability in compute_expected_q MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Inserts per-atom mean-A reduction inside the per-branch outer loop: Q[a, z] = V[z] + (A[a, z] - mean_a A[*, z]) The post-centering atom-level identifiability holds: Σ_a A_centered[a, z] = 0 for every (sample, branch, atom). Per Wang et al. 2016 Dueling Networks: mean-zero is smoother and more stable than max-zero. Atom-level subtraction (not expectation-level) preserves distributional shape per action. Sample-local register reduction; no atomicAdd, no shared memory, no cross-thread sync. NUM_ATOMS_MAX=128 mirrors existing THOMPSON_MAX_ATOMS; device __trap() on overflow. ⚠ INTERIM STATE: c51_loss_kernel + c51_grad_kernel + mag_concat_qdir still read RAW (un-centered) advantage logits. Phase 1 Tasks 1.3-1.5 migrate them in lockstep within this branch BEFORE any L40S dispatch. A partially-migrated state is forbidden per feedback_no_partial_refactor. Plan: docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/experience_kernels.cu | 58 +++++- crates/ml/tests/sp17_dueling_oracle_tests.rs | 179 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 35 ++++ 3 files changed, 268 insertions(+), 4 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index e0229a0ba..6ddd8f898 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -4496,6 +4496,15 @@ extern "C" __global__ void portfolio_sim_kernel( * @param b3_size urgency branch size * @param v_range_buf [2] adaptive C51 z-support: [v_min, v_max] */ +/* SP17 mean-zero identifiability: per-atom mean-A reduction needs a + * compile-time-sized register array. NUM_ATOMS_MAX mirrors the existing + * THOMPSON_MAX_ATOMS = 128 ceiling (production C51 NA = 51, headroom 2.5×). + * Device-side `assert(num_atoms <= NUM_ATOMS_MAX)` triggers __trap() + * on overflow — cold-fail, never silent truncation. */ +#ifndef NUM_ATOMS_MAX +#define NUM_ATOMS_MAX 128 +#endif + extern "C" __global__ void compute_expected_q( const float* __restrict__ v_logits, const float* __restrict__ b_logits, @@ -4515,6 +4524,12 @@ extern "C" __global__ void compute_expected_q( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= N) return; + /* SP17: cold-fail if num_atoms exceeds the register-array ceiling. + * Never silent truncation; __trap() halts the launch with a clear error. */ + if (num_atoms > NUM_ATOMS_MAX) { + __trap(); + } + /* Phase 2d: per-branch (v_min, v_max, delta_z) read INSIDE the d-loop * below — direction/magnitude/order/urgency own independent ranges. */ @@ -4548,6 +4563,36 @@ extern "C" __global__ void compute_expected_q( ? atom_positions + (long long)d * num_atoms : NULL; + /* SP17 mean-zero identifiability: subtract per-atom mean of A across + * this branch's actions BEFORE softmax. The reduction is sample-local + * (one thread = one sample) and per-atom z, computed from the n_d + * action logits already in this thread's scope. No cross-thread + * reduction needed; no atomicAdd; no shared memory. + * + * Per Wang et al. 2016 Dueling Networks (mean-zero variant), this + * form is preferred over max-zero because: + * - smooth differentiable reduction (no argmax tie-breaking) + * - all actions contribute proportionally to V/A split (max-zero + * pins one action to V which creates the per-action gradient + * asymmetry the paper identifies as "less stable" in Fig. 1) + * + * Per-branch independent: dir uses dir actions only, mag uses mag + * actions only, etc. The shared V head is invariant across branches + * by construction (one V per state, all branches share). + * + * Atom-level (not expectation-level): subtract per atom z, before + * softmax. Preserves distributional shape per action (Wang & + * Rajeswar 2018 distributional dueling identifiability result). */ + float a_mean_per_atom[NUM_ATOMS_MAX]; /* register array sized to compile-time max NA */ + 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; @@ -4573,14 +4618,17 @@ extern "C" __global__ void compute_expected_q( * For num_atoms=51 the inner action loop drops by ~33% global * loads against the inner advantage logits (which dominate * because v_row may stay in L1 across the 4×n_d action loop; - * adv_a flips per action — pure global). */ + * adv_a flips per action — pure global). + * + * SP17: `adv_a[z] - a_mean_per_atom[z]` replaces `adv_a[z]`. */ float M = -1e30f; float S = 0.0f; /* sum exp(logit - M) */ float TZ = 0.0f; /* sum exp(logit - M) * z_val */ float TZ2 = 0.0f; /* sum exp(logit - M) * z_val^2 */ float TLM = 0.0f; /* sum exp(logit - M) * (logit - M) */ 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; float z_val = (atom_pos_d != NULL) ? atom_pos_d[z] : v_min + (float)z * dz; @@ -4612,10 +4660,12 @@ extern "C" __global__ void compute_expected_q( /* Pass 2: atom utilisation — count atoms with prob > threshold. * prob_z = exp(logit_z - M)/S > util_threshold * ⇔ exp(logit_z - M) > util_threshold * S - * The branch is identical to the prior 3-pass version. */ + * The branch is identical to the prior 3-pass version. + * SP17: same `adv_a[z] - a_mean_per_atom[z]` centering as Pass 1. */ float util_cutoff = util_threshold * S; 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 (expf(logit - M) > util_cutoff) sample_utilized++; } diff --git a/crates/ml/tests/sp17_dueling_oracle_tests.rs b/crates/ml/tests/sp17_dueling_oracle_tests.rs index e95b69ff1..7d5199c78 100644 --- a/crates/ml/tests/sp17_dueling_oracle_tests.rs +++ b/crates/ml/tests/sp17_dueling_oracle_tests.rs @@ -94,6 +94,10 @@ mod gpu { const SP17_V_A_MEANS_DIAG_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/v_a_means_diag_kernel.cubin")); + // Cubin handle for production experience kernels (compute_expected_q). + const SP17_EXPERIENCE_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/experience_kernels.cubin")); + fn make_test_stream() -> Arc { let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?"); ctx.default_stream() @@ -109,6 +113,16 @@ mod gpu { .expect("load v_a_means_diag_kernel function") } + fn load_compute_expected_q(stream: &Arc) -> CudaFunction { + let module = stream + .context() + .load_cubin(SP17_EXPERIENCE_CUBIN.to_vec()) + .expect("load experience_kernels cubin"); + module + .load_function("compute_expected_q") + .expect("load compute_expected_q function") + } + /// Block size matches the production launcher in `gpu_dqn_trainer.rs`. const V_A_MEANS_BLOCK: u32 = 256; /// Dynamic shared-memory bytes for the kernel: one float per thread for @@ -228,4 +242,169 @@ mod gpu { result[4], (result[4] - A_URG_VAL).abs(), ); } + + // ── B.1: compute_expected_q with mean-zero centered A matches CPU oracle ── + /// B.1: post-SP17 `compute_expected_q` GPU implementation matches the CPU + /// oracle when DUELING_CENTERED is engaged (always engaged post-SP17). + /// + /// Inputs mirror the CPU oracle (`compute_expected_q_centered_cpu_oracle`): + /// - NA = 2 atoms, atom positions [0.0, 1.0] + /// - V = [0.5, 0.5] + /// - dir branch (B0=4): A_raw = [[0.1,0.0],[0.3,0.2],[0.0,0.0],[-0.2,-0.1]] + /// - other branches (B1=B2=B3=1) carry zero advantage logits + /// + /// The kernel must reproduce per-action E[Q] with mean-zero identifiability + /// applied: `Q[a, z] = V[z] + (A[a, z] - mean_a A[*, z])` followed by + /// softmax(z) expectation against atom positions. Tolerance ε=1e-5 + /// accommodates f32 rounding from the online softmax accumulators. + #[test] + #[ignore = "requires GPU"] + fn compute_expected_q_centered_matches_cpu_oracle() { + let stream = make_test_stream(); + let kernel = load_compute_expected_q(&stream); + + // Test shapes: N=1, NA=2, B0=4 (dir), B1=B2=B3=1 (minimum padding). + const N: usize = 1; + const NA: usize = 2; + const B0: usize = 4; + const B1: usize = 1; + const B2: usize = 1; + const B3: usize = 1; + const TOTAL_ACTIONS: usize = B0 + B1 + B2 + B3; + + // ── CPU oracle inputs (mirrors compute_expected_q_centered_cpu_oracle) ── + let v: [f32; NA] = [0.5, 0.5]; + let a_raw_dir: [[f32; NA]; B0] = [ + [0.1, 0.0], + [0.3, 0.2], + [0.0, 0.0], + [-0.2, -0.1], + ]; + let atom_positions = [0.0f32, 1.0]; // shared across branches + + // ── Compute CPU oracle expected E[Q] for each dir action ─────────── + let mean_a_z = [ + (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, + ]; + let mut cpu_eq_dir = [0.0f32; B0]; + for a in 0..B0 { + let mut q_logits = [0.0f32; NA]; + for z in 0..NA { + q_logits[z] = v[z] + (a_raw_dir[a][z] - mean_a_z[z]); + } + let m = q_logits[0].max(q_logits[1]); + let e0 = (q_logits[0] - m).exp(); + let e1 = (q_logits[1] - m).exp(); + let s = e0 + e1; + cpu_eq_dir[a] = (e0 * atom_positions[0] + e1 * atom_positions[1]) / s; + } + + // ── BRANCH-MAJOR b_logits buffer ───────────────────────────────── + // Layout: [N×B0×NA | N×B1×NA | N×B2×NA | N×B3×NA] + // Branch 0 (dir): a_raw_dir, branches 1-3: zero advantage (single action each). + 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]; + // Fill dir branch: sample 0, action a, atom z → offset 0 + 0 + a*NA + z + for a in 0..B0 { + for z in 0..NA { + b_logits_host[a * NA + z] = a_raw_dir[a][z]; + } + } + // Branches 1-3 already zero-filled (zero advantage means centering yields zero too). + + // v_logits: [N, NA] = single sample with v = [0.5, 0.5] + 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. + // With explicit atom_positions provided, the kernel uses atom_positions[d * NA + z] + // and ignores v_min/dz — but we still need a valid buffer (kernel always reads it). + // Fill with (0.0, 1.0, 1.0) per branch so any code path that *did* fall back is sane. + let support_host: Vec = (0..N * 4) + .flat_map(|_| [0.0_f32, 1.0, 1.0]) + .collect(); + + // atom_positions: [4, NA] — same [0.0, 1.0] for every branch. + let mut atom_pos_host = vec![0.0f32; 4 * NA]; + for d in 0..4 { + for z in 0..NA { + atom_pos_host[d * NA + z] = atom_positions[z]; + } + } + + // ── Mapped-pinned buffers (per feedback_no_htod_htoh_only_mapped_pinned) ── + 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_values 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 atom_pos_buf = unsafe { MappedF32Buffer::new(atom_pos_host.len()) } + .expect("alloc atom_positions buf"); + atom_pos_buf.write_from_slice(&atom_pos_host); + + // Scalars. + 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 null_ptr: u64 = 0; // atom_stats_block_sums = NULL, q_variance = NULL + + // Launch: 1 block × 256 threads suffices for N=1. + 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(&null_ptr) // atom_stats_block_sums = NULL + .arg(&null_ptr) // q_variance = NULL + .arg(&atom_pos_buf.dev_ptr) + .launch(LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: 0, + }) + .expect("launch compute_expected_q"); + } + stream.synchronize().expect("sync after compute_expected_q"); + + // Read back q_values: [N=1, TOTAL_ACTIONS=7]. + let result = q_out_buf.read_all(); + assert_eq!(result.len(), N * TOTAL_ACTIONS); + + // Compare per-dir-action E[Q] against CPU oracle (action_idx 0..3 = dir). + let eps = 1e-5_f32; + for a in 0..B0 { + let gpu_eq = result[a]; + let diff = (gpu_eq - cpu_eq_dir[a]).abs(); + assert!( + diff < eps, + "dir action {a}: GPU E[Q] {gpu_eq} vs CPU oracle {} (diff {:.3e})", + cpu_eq_dir[a], diff, + ); + } + } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 6c6bdf272..621b7b643 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -90,6 +90,41 @@ SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml --test sp17_dueling_oracl Plan: `docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md` Task PP.3. +## 2026-05-08 — SP17 Task 1.2: mean-zero identifiability in `compute_expected_q` + +Inserts per-atom mean-A reduction inside `compute_expected_q`'s per-branch outer loop: +`Q[a, z] = V[z] + (A[a, z] − mean_a A[*, z])` followed by softmax(z) over atoms. +The post-centering atom-level identifiability invariant `Σ_a A_centered[a, z] = 0` +holds for every (sample, branch, atom). + +### What lands + +- **Kernel modification**: `crates/ml/src/cuda_pipeline/experience_kernels.cu::compute_expected_q` — inserted per-atom mean-A reduction in the per-branch outer loop, BEFORE the action-loop logit accumulation. `float a_mean_per_atom[NUM_ATOMS_MAX]` register array; `NUM_ATOMS_MAX=128` mirrors `THOMPSON_MAX_ATOMS`. Both inner softmax pass and atom-utilization pass now read `adv_a[z] - a_mean_per_atom[z]` instead of `adv_a[z]`. Sample-local register reduction; no atomicAdd, no shared memory, no cross-thread sync per `feedback_no_atomicadd`. Device-side `__trap()` if `num_atoms > NUM_ATOMS_MAX` — cold-fail, never silent truncation. +- **GPU oracle test** (NEW): `crates/ml/tests/sp17_dueling_oracle_tests.rs::gpu::compute_expected_q_centered_matches_cpu_oracle`. N=1, NA=2, B0=4 dir actions with `a_raw_dir = [[0.1,0.0],[0.3,0.2],[0.0,0.0],[-0.2,-0.1]]`, `v=[0.5,0.5]`, atom positions `[0.0, 1.0]`. Compares per-action GPU E[Q] against the CPU oracle (`compute_expected_q_centered_cpu_oracle`) to ε=1e-5. Mapped-pinned device buffers per `feedback_no_htod_htoh_only_mapped_pinned`. + +### Wire-up status — INTERIM + +| Consumer of `b_logits` raw advantage | Status this commit | +|---|---| +| `compute_expected_q` (this kernel) | **MIGRATED** to centered A | +| `c51_loss_kernel` | **PENDING** — Task 1.4 | +| `c51_grad_kernel` | **PENDING** — Task 1.4 | +| `mag_concat_qdir` (Thompson + mag-concat) | **PENDING** — Task 1.5 | + +⚠ **Partial-refactor state acknowledged.** Per `feedback_no_partial_refactor`, the contract for `b_logits` advantage logits has changed in `compute_expected_q` only; the other three consumers still read RAW (un-centered) advantage logits. Tasks 1.3-1.5 close the migration in lockstep within this same branch BEFORE any L40S dispatch. Atomicity is at the BRANCH level, not the COMMIT level — interim state is forbidden to escape `feat/sp17-dueling`. + +### 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 \ + compute_expected_q_centered_cpu_oracle → PASS (1/1, host) +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml --test sp17_dueling_oracle_tests --features cuda \ + compute_expected_q_centered_matches_cpu_oracle -- --ignored --nocapture → PASS (1/1, RTX 3050 Ti) +``` + +Plan: `docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md` Task 1.2. + ## 2026-05-08 — Class A audit-fix Batch 4-B: plan_threshold floor adaptive + MIN_HOLD_TEMPERATURE ISV-driven Per Class A audit-fix Batch 4-B (final 2 of 4 deferred items from P1-wiring/P1-producer). Completes the 8-commit WR-plateau intervention chain. Validation deferred to next L40S smoke.