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

feat(sp17): centered A in mag_concat_qdir (Commit B)

Browse Source 
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]); all three now read `dir_logits_b[..] - a_mean_per_atom[z]`.

Without this migration the magnitude trunk-input would see raw E[Q_dir]
while c51_loss/c51_grad backward consumes centered A — mixed contract
ruled out by `feedback_no_partial_refactor`.

The Plan-4 q_rms adaptive scale (ISV[96]) is preserved unchanged;
centering changes the per-action E[Q_dir] values that feed it but the
scale formula itself is structurally orthogonal. Backward into the
direction logits flows through c51_grad (already mean-zero) so no bw
change is needed in this kernel.

GPU oracle test: B=1, SH2=2, NA=3, B0=4 with the same synthetic where
mean_a = [0.75, 0, 0.75] is non-zero. Asserts:
  (a) h_s2[0..SH2] copied verbatim into concat prefix
  (b) tail slots = centered_E[Q_dir] × (1.0 / q_rms) to ε=1e-5
  (c) tail[0] < 0 / tail[3] > 0 sign asymmetry (regression detector)

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

⚠ INTERIM STATE: Thompson direction-select + aux-CQL barrier/ib still
read raw advantage. Commits C-D close them; Commit E annotates the two
already-centered c51_loss/c51_grad pre-SP17 sites.

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:38:03 +02:00
parent 9b6e94854f
commit 4802879494
3 changed files with 271 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

@@ -418,6 +418,16 @@ mod gpu {
.expect("load quantile_q_select function")
}
fn load_mag_concat_qdir(stream: &Arc<CudaStream>) -> 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<f32> = vec![1.0, 2.0];
let v_logits_host: Vec<f32> = 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],
);
}
}

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

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