fix(sp22): W-read guards across ALL atom-shift kernels

Verify smoke train-t5885 partial result: dir CLEAN, mag NaN. The
isfinite(W) guard added in compute_expected_q didn't extend to the
other atom-shift consumers. Adam-corrupted W propagates through them
via 0*NaN=NaN.

Adds isfinite(w_aux[a]) guards to all atom-shift kernels:
- mag_concat_qdir
- quantile_q_select
- c51_loss_kernel (eq_per_action shift + effective_reward W[a0]/W[best_next_a])

Also strengthens c51_aux_dw_kernel: guard sp/isw/dz/gamma/done (not
just dz<1e-7 which doesn't catch NaN). Any NaN/Inf input -> skip
sample (no dW contribution).

Defense-in-depth complete: NaN cannot propagate through any
atom-shift path regardless of source.

Cargo check clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu

@@ -93,11 +93,18 @@ extern "C" __global__ void c51_aux_dw_kernel(
         float sp = aux_proj_logdiff_dir[b];
         float isw = fminf(is_weights[b], 10.0f);
         float dz = per_sample_support[b * 12 + 0 * 3 + 2];
-        if (dz < 1e-7f) continue;  /* degenerate dir support: no gradient (matches forward skip) */
-
-        float dL_dDelta_online = sp / dz;
+        /* DEFENSIVE NaN guards (2026-05-13): the dz check `dz < 1e-7` does
+         * NOT catch NaN (NaN comparison is always false), so NaN dz would
+         * slip through and produce NaN dW. Bisect cut #1 proved this
+         * kernel + Adam are the NaN source. Guard ALL upstream-derived
+         * inputs: sp / is_weights / dz / gamma / done. Any NaN/Inf in
+         * any input → skip sample (no contribution to dW). */
+        if (!isfinite(sp) || !isfinite(isw) || !isfinite(dz) || dz < 1e-7f) continue;
         float gamma_eff = gamma_buf[b];
         float done = dones[b];
+        if (!isfinite(gamma_eff) || !isfinite(done)) continue;
+
+        float dL_dDelta_online = sp / dz;
         float dL_dDelta_target = -gamma_eff * (1.0f - done) * dL_dDelta_online;
 
         if (a == a_d) {

crates/ml/src/cuda_pipeline/c51_loss_kernel.cu

@@ -1088,7 +1088,11 @@ extern "C" __global__ void c51_loss_batched(
                  * the same amount (linearity of conditional expectation).
                  */
                 if (d == 0 && aux_shift_active) {
-                    cvar += w_aux[a] * next_state_121;
+                    /* DEFENSIVE NaN guard (2026-05-13): W can carry NaN from
+                     * Adam corruption; 0 * NaN = NaN propagates. */
+                    float w_a = w_aux[a];
+                    if (!isfinite(w_a)) w_a = 0.0f;
+                    cvar += w_a * next_state_121;
                 }
                 eq_per_action[a] = cvar;
             }
@@ -1233,8 +1237,11 @@ extern "C" __global__ void c51_loss_batched(
          * branches and when aux_shift_active == false. */
         float effective_reward = reward;
         if (d == 0 && aux_shift_active) {
-            float aux_shift_online = w_aux[a0]          * state_121;
-            float aux_shift_target = w_aux[best_next_a] * next_state_121;
+            /* DEFENSIVE NaN guard (2026-05-13): W can carry NaN. */
+            float w_ad   = w_aux[a0];          if (!isfinite(w_ad))   w_ad = 0.0f;
+            float w_star = w_aux[best_next_a]; if (!isfinite(w_star)) w_star = 0.0f;
+            float aux_shift_online = w_ad   * state_121;
+            float aux_shift_target = w_star * next_state_121;
             effective_reward = reward
                              + gamma_eff * aux_shift_target * (1.0f - done)
                              - aux_shift_online;

crates/ml/src/cuda_pipeline/experience_kernels.cu

@@ -5660,9 +5660,14 @@ extern "C" __global__ void mag_concat_qdir(
         /* SP22 H6 Phase 3 α atom-shift: per-action shift = W[a]*state_121.
          * Effective atom positions for this action: z_n + W[a]*state_121.
          * Since Σ_n prob_n = 1, the expected Q shifts by exactly the
-         * scalar W[a]*state_121 (linearity of expectation). */
+         * scalar W[a]*state_121 (linearity of expectation).
+         * DEFENSIVE NaN guard (2026-05-13): same pattern as
+         * compute_expected_q — W can carry NaN from Adam corruption, and
+         * 0 * NaN = NaN propagates. Verify finite. */
         if (aux_shift_active) {
-            eq += w_aux[a] * state_121;
+            float w_a = w_aux[a];
+            if (!isfinite(w_a)) w_a = 0.0f;
+            eq += w_a * state_121;
         }
         eq_local[a] = eq;
     }
@@ -6311,9 +6316,13 @@ extern "C" __global__ void quantile_q_select(
             /* SP22 H6 Phase 3 α atom-shift: dir branch (d==0) only.
              * Per-action shift = W[a]*state_121 uniformly shifts q10/q50/q90
              * (they're all positions on the shifted support), so the blend
-             * shifts by exactly W[a]*state_121. */
+             * shifts by exactly W[a]*state_121.
+             * DEFENSIVE NaN guard (2026-05-13): same pattern as
+             * compute_expected_q + mag_concat_qdir. */
             if (d == 0 && aux_shift_active) {
-                q_blended += w_aux[a] * state_121;
+                float w_a = w_aux[a];
+                if (!isfinite(w_a)) w_a = 0.0f;
+                q_blended += w_a * state_121;
             }
 
             q_select[(long long)i * total_actions + action_idx] = q_blended;

docs/dqn-wire-up-audit.md

@@ -17428,3 +17428,22 @@ Together these provide belt-and-suspenders against NaN propagation:
 The underlying NaN source in Step 8/11 is still unknown — but the defensive guards prevent it from corrupting forward Q values. Smoke verification next.
 
 Cargo check clean.
+
+#### Comprehensive W-read guards across all atom-shift kernels (2026-05-13)
+
+**Verification smoke `train-t5885` @ commit `57cf5c1e0` partial result**:
+- `q_var_per_branch [dir=0 mag=NaN ord=0 urg=0]` — dir CLEAN, mag NaN
+- `q_by_action [hold=-0 long=-0 short=-0 flat=-0]` — dir CLEAN
+- `v_a_means [v=-0.11 a_dir=-0.003 a_mag=NaN a_ord=0.004 a_urg=0.003]` — a_mag NaN
+
+**Diagnostic**: defensive W-guard in `compute_expected_q` fixed dir branch. But `mag_concat_qdir` (which feeds mag SGEMM input via the dir-Q concat) does NOT have the W guard yet. NaN W there propagates: W=NaN * state_121=0 = NaN → eq_local NaN → mag_concat_buf NaN → mag SGEMM forward NaN → a_mag NaN.
+
+**This commit**: applies the SAME isfinite guard on `w_aux[a]` reads across ALL atom-shift consumer kernels:
+- `mag_concat_qdir` (experience_kernels.cu): guard W in `eq += W[a]*state_121`
+- `quantile_q_select` (experience_kernels.cu): guard W in `q_blended += W[a]*state_121`
+- `c51_loss_kernel`: guard W in `cvar += W[a]*next_state_121` AND in `effective_reward` substitution (both W[a0] and W[best_next_a])
+- `c51_aux_dw_kernel`: extended guards to cover sp/isw/dz/gamma/done — any NaN/Inf in any input → skip sample (the dz < 1e-7 check did NOT catch NaN since NaN comparisons are always false)
+
+Together with the previously-added guards on `state_121` and on `total` (dW write), Adam inputs (`g, lr, m, v`), and `W` in `compute_expected_q`, this provides defense-in-depth across the entire atom-shift path. No NaN can propagate through any of the 4 forward kernels OR the dW/Adam backward.
+
+Cargo check clean.