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

feat(sp17): centered A in quantile_q_select (Commit A)

Browse Source 
Migrates `quantile_q_select` (uncertainty-driven action selection from
C51 atom CDFs) to the SP17 mean-zero advantage contract. The kernel
builds a per-(sample, branch, action) softmax over `V[z] + A[a, z]`
across THREE inner passes (max, sum_exp, CDF); all three now read
`adv_a[z] - a_mean_per_atom[z]`.

Plan flagged this as a hidden 6th consumer. Task 1.4/1.5 as authored
covered only c51_loss + c51_grad + mag_concat_qdir + Thompson; the
post-Task-1.2 audit found `quantile_q_select` reads raw advantage at
lines 5704/5709/5720 across all 4 branches and was missed entirely.

Sample-local register reduction (`float a_mean_per_atom[NUM_ATOMS_MAX]`,
NA_MAX=128 mirroring compute_expected_q); no atomicAdd, no shared mem,
no cross-thread sync per `feedback_no_atomicadd`. Device-side __trap()
on num_atoms overflow per `feedback_no_quickfixes`.

GPU oracle test: N=1, NA=3, B0=4 with V=[0,0,0] and A_raw chosen so the
per-atom mean is [0.75, 0, 0.75]. Test runs the production cubin with
iqn_readiness=0 and util_ema=1.0 and asserts each per-action q90 matches
the CPU oracle to ε=1e-5. Two structural-asymmetry assertions fail
loudly if centering regresses (action 0 q90 ≤ 0, action 3 q90 ≥ 0).
Mapped-pinned per `feedback_no_htod_htoh_only_mapped_pinned`.

Verification (RTX 3050 Ti):
  cargo check --workspace                                          → clean
  cargo test -p ml --test sp17_dueling_oracle_tests --features cuda
    -- --ignored                                                    → 3/3 PASS

⚠ INTERIM STATE: mag_concat_qdir + Thompson + aux-CQL barrier/ib still
read raw advantage. Commits B-D close them; Commit E annotates the two
already-centered c51_loss/c51_grad pre-SP17 sites. No L40S dispatch
escapes feat/sp17-dueling until every consumer is migrated 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
2026-05-08 21:33:06 +02:00
parent eabcf8d529
commit 9b6e94854f
3 changed files with 280 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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; }

213
crates/ml/tests/sp17_dueling_oracle_tests.rs
View File

@@ -407,4 +407,217 @@ mod gpu {
);
}
}
fn load_quantile_q_select(stream: &Arc<CudaStream>) -> 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<f32> = 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],
);
}
}

34
docs/dqn-wire-up-audit.md
View File

@@ -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`.