feat(sp17): mean-zero identifiability in compute_expected_q
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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++;
|
||||
}
|
||||
|
||||
|
||||
@@ -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<CudaStream> {
|
||||
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<CudaStream>) -> 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<f32> = 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<f32> = (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,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user