docs(sp17): annotate c51_loss/c51_grad SP17 compliance (Commit E)

Audit-only commit per user design call DD9. The plan's Task 1.4 incorrectly claimed c51_loss_batched + c51_grad_kernel needed migration; both kernels have been mean-zero centered since commit 56373f094. Audit findings: c51_loss_kernel.cu::c51_loss_batched (line 760-806): - Forward dueling reduction at lines 765-779 already computes `a_mean = (1/n_d) Σ shmem_adv[a, j]` and applies `centered = shmem_adv[a_d, j] - a_mean` — IS the SP17 mean-zero projection. Annotated `/* SP17 mean-zero */`. - Counterfactual magnitude path at lines 788-805: same pattern. - Spectral decoupling at lines 812-827: same pattern. - Bellman target argmax at lines 848-855: shmem_proj[j] IS a_mean_per_atom. - Per-d=1 magnitude std normalization (line 770-778) is ORTHOGONAL to SP17 mean-zero — a separate variance-control concern for the magnitude branch's tighter Q distribution. Annotated `/* SP17: per-d=1 mag std normalization (orthogonal) */`. c51_grad_kernel.cu::c51_grad_kernel (line 287-291): - `dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A)` at line 288 IS the SP17 mean-zero Jacobian: J[a, a_d] = δ(a, a_d) − 1/n_d. The chain rule for centered logit `A_taken − mean_a A` w.r.t. A[a'] produces exactly this Kronecker-delta-minus-uniform pattern. - Per-d=1 magnitude std grad multiplier (line 290) — same orthogonal concern as c51_loss. Zero code-path change. Both kernels behave bit-identically to pre-E. Wire-up status (FINAL — every consumer SP17-compliant): compute_expected_q (Task 1.2) MIGRATED quantile_q_select (Commit A) MIGRATED mag_concat_qdir (Commit B) MIGRATED Thompson direction-select (Commit C) MIGRATED + V wired in barrier_gradient_direction (Commit D) MIGRATED ib_gradient_direction (Commit D) MIGRATED c51_loss_batched (this commit) ALREADY-COMPLIANT, ANNOTATED c51_grad_kernel (this commit) ALREADY-COMPLIANT, ANNOTATED After this commit, EVERY advantage-logit read in cuda_pipeline goes through the SP17 mean-zero contract. No un-centered raw-advantage reads remain in production kernels. Verification (RTX 3050 Ti): cargo check --workspace → clean cargo test sp17_dueling_oracle_tests --features cuda -- --ignored → 6/6 PASS grep adv_a[ kernels (excluding centered/comments) → empty grep dir_logits_b[ kernels (excluding centered/comments) → empty Plan: docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-08 22:03:17 +02:00
parent ffa6fda868
commit 6f53d676fb
3 changed files with 82 additions and 4 deletions
--- a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
+++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
@@ -285,9 +285,19 @@ extern "C" __global__ void c51_grad_kernel(
        }

        for (int a = 0; a < A_d; a++) {
+            /* SP17 annotation (Commit E, 2026-05-08): this `dueling_grad` IS
+             * the SP17 mean-zero Jacobian — `J[a, a_d] = δ(a, a_d) − 1/n_d`
+             * (where `inv_A = 1.0 / A_d`). For the action `a` that was taken
+             * in the forward pass, the per-action centered logit is
+             * `A_taken − mean_a A`; the chain rule gradient w.r.t. logit at
+             * action a' is `∂(A_taken − mean_a A)/∂A[a']` = δ(a', taken) −
+             * 1/n_d. This site has carried the correct mean-zero Jacobian
+             * since 56373f094 — no code-path change in this commit.
+             *
+             * SP17: this IS the mean-zero Jacobian — Δ(a,a_d) - 1/n_d. */
            float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
            float grad_val = branch_scale * d_combined * dueling_grad;
-            if (d == 1) grad_val *= inv_a_std;
+            if (d == 1) grad_val *= inv_a_std;   /* SP17: per-d=1 mag std normalization (orthogonal) */

            float spread_grad;
            if (a == a_d) {
--- a/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
+++ b/crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
@@ -760,13 +760,29 @@ extern "C" __global__ void c51_loss_batched(
        /* Dueling: Q[j] = V[j] + A'[a_d,j] where A' is centered (all branches)
         * and standardized to unit variance (magnitude branch d==1 only).
         * Standardization prevents winner-take-all Q-value collapse where Small
-         * dominates due to lower PnL variance → tighter C51 distributions. */
+         * dominates due to lower PnL variance → tighter C51 distributions.
+         *
+         * SP17 annotation (Commit E, 2026-05-08): this `centered = shmem_adv[a_d, j]
+         * − a_mean` IS the SP17 mean-zero identifiability projection — same
+         * contract that compute_expected_q + quantile_q_select + mag_concat_qdir
+         * + Thompson + barrier_gradient + ib_gradient now also enforce. The
+         * `Σ_a (centered_a) = 0` invariant holds at every (sample, branch, atom).
+         * No code-path change in this commit; this site has been mean-zero
+         * since 56373f094.
+         *
+         * SP17: per-d=1 magnitude std normalization preserved (orthogonal to
+         * SP17 mean-zero). The `centered /= a_std` is a SEPARATE concern —
+         * the magnitude branch's PnL variance is lower than direction's by
+         * construction (Kelly cap shrinks per-action realised PnL spread),
+         * so the C51 distribution shape is tighter and a winner-take-all
+         * collapse is otherwise easy to fall into. The std-normalization
+         * exists at this site only (d==1) and is NOT part of SP17 scope. */
        for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * num_atoms + j];
            a_mean /= (float)n_d;
-            float centered = shmem_adv[a_d * num_atoms + j] - a_mean;
+            float centered = shmem_adv[a_d * num_atoms + j] - a_mean;   /* SP17 mean-zero */
            if (d == 1) {
                float sq_sum = 0.0f;
                for (int a = 0; a < n_d; a++) {
@@ -774,7 +790,7 @@ extern "C" __global__ void c51_loss_batched(
                    sq_sum += diff * diff;
                }
                float a_std = sqrtf(sq_sum / (float)n_d + 1e-12f);
-                centered /= (a_std + 1e-6f);
+                centered /= (a_std + 1e-6f);   /* SP17: per-d=1 mag std normalization (orthogonal) */
            }
            shmem_lp[j] = shmem_val[j] + centered;
        }
--- a/docs/dqn-wire-up-audit.md
+++ b/docs/dqn-wire-up-audit.md
@@ -2,6 +2,58 @@

 **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.

+## 2026-05-08 — SP17 Commit E: c51_loss_batched + c51_grad_kernel SP17-compliance annotations (audit-only, no code-path change)
+
+User design call DD9: audit `c51_loss_batched` + `c51_grad_kernel` for SP17-contract compliance and annotate the existing already-centered code with explicit SP17 markers. The plan's Task 1.4 incorrectly claimed these kernels needed migration — they have been mean-zero centered since commit `56373f094`. Verifying this audit closes the contract gap that the plan opened by undercounting consumers.
+
+### Audit findings
+
+**c51_loss_kernel.cu::c51_loss_batched** (line 760-806):
+- Forward pass at lines 765-779 already computes `a_mean = (1/n_d) Σ shmem_adv[a, j]` and applies `centered = shmem_adv[a_d, j] − a_mean`. **CONFIRMED SP17-compliant.**
+- Counterfactual magnitude path at lines 788-805 (Gem 3+) uses the same `a_mean − centered` reduction. **CONFIRMED SP17-compliant.**
+- Spectral decoupling at lines 812-827 iterates over all actions with the same `a_mean − centered` pattern. **CONFIRMED SP17-compliant.**
+- Bellman target argmax recompute at lines 848-855 uses `a_mean = (1/n_d) Σ shmem_adv[a, j]` then `shmem_proj[j] = a_mean`, and the next softmax (line 862) does `shmem_lp[j] = shmem_val[j] + shmem_adv[a, j] - shmem_proj[j]`. **CONFIRMED SP17-compliant** (the `shmem_proj` IS the `a_mean_per_atom`).
+- Per-d=1 magnitude std normalization (line 770-778) is **orthogonal** to SP17 mean-zero — a separate variance-control concern for the magnitude branch's tighter Q distribution. Annotated with `// SP17: per-d=1 mag std normalization (orthogonal)` to disambiguate.
+
+**c51_grad_kernel.cu::c51_grad_kernel** (line 287-291):
+- The `dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A)` formula at line 288 IS the SP17 mean-zero Jacobian: `J[a, a_d] = δ(a, a_d) − 1/n_d`. The chain rule for the centered logit `A_taken − mean_a A` w.r.t. `A[a']` produces exactly this Kronecker-delta-minus-uniform pattern. **CONFIRMED SP17-compliant.** Annotated with `// SP17: this IS the mean-zero Jacobian`.
+- Per-d=1 magnitude std grad multiplier at line 290 — same orthogonal variance-control concern as c51_loss. Annotated with `// SP17: per-d=1 mag std normalization (orthogonal)`.
+
+### What lands
+
+- **Annotations only** — zero code-path changes. Both kernels behave bit-identically to pre-Commit-E.
+- `crates/ml/src/cuda_pipeline/c51_loss_kernel.cu:769-787` — added SP17 comment block above the dueling reduction, marked `centered` with `/* SP17 mean-zero */`, marked `centered /= (a_std + 1e-6f)` with `/* SP17: per-d=1 mag std normalization (orthogonal) */`.
+- `crates/ml/src/cuda_pipeline/c51_grad_kernel.cu:287-290` — added SP17 comment block above `dueling_grad`, marked the std multiplier with the same orthogonal-concern annotation.
+
+### Wire-up status — FINAL
+
+| Consumer of `b_logits` raw advantage | Status this commit |
+|---|---|
+| `compute_expected_q` (Task 1.2) | MIGRATED (SP17 mean-zero) |
+| `quantile_q_select` (Commit A) | MIGRATED (SP17 mean-zero) |
+| `mag_concat_qdir` (Commit B) | MIGRATED (SP17 mean-zero) |
+| Thompson direction-select (Commit C) | MIGRATED (SP17 mean-zero + V wired in) |
+| `barrier_gradient_direction` (Commit D) | MIGRATED (SP17 mean-zero) |
+| `ib_gradient_direction` (Commit D) | MIGRATED (SP17 mean-zero) |
+| `c51_loss_batched` (this commit) | **ALREADY-COMPLIANT, ANNOTATED** |
+| `c51_grad_kernel` (this commit) | **ALREADY-COMPLIANT, ANNOTATED** |
+
+After this commit, **every advantage-logit read in the DQN cuda_pipeline goes through the SP17 mean-zero contract**. No un-centered raw-advantage reads remain in production kernels.
+
+### 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                                                                                             → 6 PASSED (RTX 3050 Ti)
+
+# Wire-up audit grep — should produce ONLY annotations / non-kernel false positives:
+grep -rn "adv_a\\[" crates/ml/src/cuda_pipeline/*.cu | grep -vE "(centered|a_mean_per_atom|//|\\*)"   → empty
+grep -rn "dir_logits_b\\[" crates/ml/src/cuda_pipeline/*.cu | grep -vE "(centered|a_mean_per_atom|//|\\*)"   → empty
+```
+
+Plan: `docs/superpowers/plans/2026-05-08-sp17-dueling-q-network.md` Task 1.4 (DD9 amendment — plan claimed migration needed; audit confirmed already-compliant since 56373f094).
+
 ## 2026-05-08 — SP17 Commit D: aux-CQL `barrier_gradient_direction` + `ib_gradient_direction` migration

 User design call DD11: migrate the two aux-CQL gradient kernels living in `c51_loss_kernel.cu` (NOT `experience_kernels.cu` — the audit's location ref needed correction). Both kernels read raw advantage to compute per-direction-action E[Q] for the barrier (max gap below `min_req`) and IB (variance below `min_var`) gradient pushes. The pre-SP17 code at `barrier_gradient_direction:1212` carried this comment: