fix(sp22): defensive NaN guard on state_121 reads in atom-shift kernels

Root cause: IEEE-754 rule 0 * NaN = NaN defeats W=0 safety. Two smokes
proved the bug propagates regardless of W magnitude — pattern identical
at W=[-0.5,0,+0.5,0] (train-th8pj) and W=[0,0,0,0] (train-gs4gx).

The state_121 = batch_states[i * state_dim + 121] read in 4 atom-shift
kernels can contain NaN (upstream source: aux head softmax producing
NaN at some rollout step, stored in replay buffer, sampled into
trainer's batch). 0 * NaN = NaN propagates through all atom-shift
arithmetic.

Defensive fix: guard each state_121 read with isfinite check, fallback
to 0 if NaN/inf. Applies to:
- compute_expected_q (experience_kernels.cu)
- mag_concat_qdir (experience_kernels.cu)
- quantile_q_select (experience_kernels.cu)
- c51_loss_batched (c51_loss_kernel.cu) — state_121 AND next_state_121
- c51_aux_dw_kernel (s121 AND ns121)

This unblocks Phase 3 mechanism validation. The actual state_121 NaN
source (aux head or state assembly) is to be investigated separately.

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

@@ -102,10 +102,17 @@ extern "C" __global__ void c51_aux_dw_kernel(
 
         if (a == a_d) {
             float s121 = batch_states[(long long)b * state_dim + aux_dir_prob_index];
+            /* Defensive NaN guard (2026-05-13): mirrors atom-shift kernels.
+             * State_121 NaN would make local_sum NaN, dW NaN, Adam NaN,
+             * w_aux NaN — corrupting all subsequent forward atom-shifts.
+             * Clamp NaN to 0 → no W gradient contribution from that sample
+             * (correct behavior; the gradient was undefined anyway). */
+            if (!isfinite(s121)) s121 = 0.0f;
             local_sum += inv_batch * isw * dL_dDelta_online * s121;
         }
         if (a == a_star) {
             float ns121 = next_batch_states[(long long)b * state_dim + aux_dir_prob_index];
+            if (!isfinite(ns121)) ns121 = 0.0f;
             local_sum += inv_batch * isw * dL_dDelta_target * ns121;
         }
     }

crates/ml/src/cuda_pipeline/c51_loss_kernel.cu

@@ -838,6 +838,13 @@ extern "C" __global__ void c51_loss_batched(
                                      + aux_dir_prob_index];
         next_state_121 = next_batch_states[(long long)sample_id * state_dim
                                           + aux_dir_prob_index];
+        /* Defensive NaN guard (2026-05-13): state_121 SHOULD be in
+         * [-1, +1] per the H6 spec, but IEEE-754 `0 * NaN = NaN`
+         * defeats W=0 safety if an upstream NaN propagates here. The
+         * real source (aux head or state assembly) is investigated
+         * separately. */
+        if (!isfinite(state_121))      state_121 = 0.0f;
+        if (!isfinite(next_state_121)) next_state_121 = 0.0f;
     }
 
     for (int d = 0; d < NUM_BRANCHES; d++) {

crates/ml/src/cuda_pipeline/experience_kernels.cu

@@ -5044,6 +5044,17 @@ extern "C" __global__ void compute_expected_q(
             if (d == 0 && w_aux != NULL && batch_states != NULL) {
                 float state_121 = batch_states[(long long)i * state_dim
                                              + aux_dir_prob_index];
+                /* Defensive NaN guard (2026-05-13): state_121 SHOULD be in
+                 * [-1, +1] per the H6 spec (recentered aux p_up), but if
+                 * an upstream NaN reaches here (e.g., aux head numerical
+                 * instability propagated through replay buffer), the
+                 * IEEE-754 rule `0 * NaN = NaN` makes W=0 NOT a safety
+                 * mechanism. Clamping NaN to 0 prevents propagation
+                 * through compute_expected_q while keeping the atom-shift
+                 * mathematically correct for finite state_121. The real
+                 * NaN source (likely aux head or state assembly) is to
+                 * be investigated separately. */
+                if (!isfinite(state_121)) state_121 = 0.0f;
                 aux_atom_shift = w_aux[a] * state_121;
             }
 
@@ -5565,11 +5576,14 @@ extern "C" __global__ void mag_concat_qdir(
     if (b >= B) return;
 
     /* SP22 H6 Phase 3 α atom-shift: hoist state_121 once per sample.
-     * NULL-safe per the same pattern as compute_expected_q + c51_loss. */
+     * NULL-safe per the same pattern as compute_expected_q + c51_loss.
+     * Defensive NaN guard (2026-05-13): see compute_expected_q for
+     * rationale — IEEE-754 `0 * NaN = NaN` defeats W=0 safety. */
     float state_121 = 0.0f;
     bool  aux_shift_active = (w_aux != NULL && batch_states != NULL);
     if (aux_shift_active) {
         state_121 = batch_states[(long long)b * state_dim + aux_dir_prob_index];
+        if (!isfinite(state_121)) state_121 = 0.0f;
     }
 
     /* Plan 4 Task 2c.3c.6: read trunk-output RMS once into a register.
@@ -6201,11 +6215,13 @@ extern "C" __global__ void quantile_q_select(
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= N) return;
 
-    /* SP22 H6 Phase 3 α atom-shift: hoist state_121 once. */
+    /* SP22 H6 Phase 3 α atom-shift: hoist state_121 once.
+     * Defensive NaN guard (2026-05-13): see compute_expected_q. */
     float state_121 = 0.0f;
     bool  aux_shift_active = (w_aux != NULL && batch_states != NULL);
     if (aux_shift_active) {
         state_121 = batch_states[(long long)i * state_dim + aux_dir_prob_index];
+        if (!isfinite(state_121)) state_121 = 0.0f;
     }
 
     /* Phase 2d: per-branch (v_min, delta_z) read INSIDE the d-loop below. */