fix(cuda): Plan A v2 — dd049d9a4 + 3 kernel-only bug fixes (no Phase 2.0)

Plan A v1 (commit 2d68ef5d8 = revert Phase 2.1 on top of 104fe81ca)
failed at cluster (alpha-rl-4sjzw): entropy collapsed to 0.69 by step 3000.
The Phase 2.0 V envelope clamp (c52282fb4) — kept in v1 — was likely
the culprit: it bounds V_pred to a tight envelope, biasing advantages
and breaking PPO gradient flow.

Plan A v2: start from dd049d9a4 (proven working at cluster wr=0.57
+$6.3M @ step 15000 today via alpha-rl-lpbp8) and apply ONLY the
kernel-only bug fixes that don't affect training dynamics:

- variable_selection.cu: VSN stride 40→56 fix (104fe81ca) — prevents
  step-4 NaN from reading wrong-stride window_tensor
- bucket_transition_kernels.cu: h_mag_per_bucket multi-warp 32→128
  (7e38e46e6) — correct warp-shuffle reduction for HIDDEN_DIM=128
- compute_advantage_return.cu: branch-gate done flag (a6acc25ec) —
  done's terminal-state semantics applied via gate, resolves step-4 NaN

NOT applied (intentionally):
- c52282fb4 Phase 2.0 V envelope clamp — biases V regression target
- b4aadff75 V envelope ±200→±10 — extends Phase 2.0
- 10d4614fb atomicAdd removal — kernel non-determinism is real but
  intrusive (Rust changes), dd049d9a4 works without this fix
- db4d9a16f Phase 2.1 dueling — broken decomposition per spec analysis
- 72672c9e7+ Phase 2.2/2.3/H/P1+P2+P3 — built on broken Phase 2.1

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

crates/ml-alpha/cuda/bucket_transition_kernels.cu

@@ -392,10 +392,14 @@ extern "C" __global__ void heads_lr_multiplier_scale_kernel(
 // `channels_in_bucket[bucket][i]` populated at transition by
 // `channels_in_bucket_kernel`.
 //
-// Launch: grid = (N_HORIZONS, 1, 1), block = (32, 1, 1). Each block
+// Launch: grid = (N_HORIZONS, 1, 1), block = (128, 1, 1). Each block
 // reduces one bucket. Per `feedback_no_atomicadd`, reduction is
 // block-tree on shared memory (no atomicAdd).
 //
+// Block size = 128 covers MAX_BUCKET_DIM=96 (smallest pow2 ≥ 96).
+// Previous block_dim=32 silently dropped channels 32..bdim when bdim>32
+// — caught by `tests/bucket_transition_kernels.rs` (block_dim=43 case).
+//
 // Input `h_state` is `[B × HIDDEN_DIM]` (ORIGINAL channel layout). Each
 // block reads its bucket's `bucket_dim_k[bucket]` channels (via
 // `channels_in_bucket[bucket][0..bdim]`) per sample (across all B
@@ -412,20 +416,20 @@ extern "C" __global__ void h_mag_per_bucket_kernel(
     if (bucket >= N_HORIZONS) return;
     int bdim = (int)bucket_dim_k[bucket];
 
-    // sdata sized for a single-warp reduction (32 lanes).
-    __shared__ float sdata[32];
+    // sdata sized for the 128-lane block reduction.
+    __shared__ float sdata[128];
     int tid = threadIdx.x;
 
     // Each thread sums |h| over its bucket-local index (tid), looking up
     // the ORIGINAL channel via channels_in_bucket. Threads with tid >= bdim
-    // idle — uniform predicate, no warp divergence inside [0, 32).
+    // contribute 0. With bdim ∈ [1, MAX_BUCKET_DIM=96] and block_dim=128,
+    // tail lanes [bdim, 128) are always inactive — uniform predicate.
     float local_sum = 0.0f;
     if (tid < bdim) {
         unsigned int c = channels_in_bucket[bucket * MAX_BUCKET_DIM + tid];
-        // Defensive: sentinel slot OR out-of-range channel index should
-        // not contribute. Predicate is uniform across the warp because all
-        // active threads (tid < bdim) hold valid entries by construction
-        // of channels_in_bucket_kernel.
+        // Defensive: out-of-range channel index should not contribute.
+        // Predicate is uniform across active lanes (tid < bdim) — all hold
+        // valid entries by construction of channels_in_bucket_kernel.
         if (c < (unsigned int)HIDDEN_DIM) {
             for (int b = 0; b < B; ++b) {
                 float v = h_state[b * HIDDEN_DIM + c];
@@ -433,11 +437,12 @@ extern "C" __global__ void h_mag_per_bucket_kernel(
             }
         }
     }
-    sdata[tid] = (tid < 32) ? local_sum : 0.0f;
+    sdata[tid] = local_sum;
     __syncthreads();
 
-    // Block-tree reduction over 32 lanes (single warp). No atomicAdd.
-    for (int s = 16; s > 0; s >>= 1) {
+    // Block-tree reduction over 128 lanes (multi-warp). No atomicAdd.
+    // Stages: 64→32→16→8→4→2→1. Each stage halves the active lane count.
+    for (int s = 64; s > 0; s >>= 1) {
         if (tid < s) sdata[tid] += sdata[tid + s];
         __syncthreads();
     }

crates/ml-alpha/cuda/compute_advantage_return.cu

@@ -19,13 +19,26 @@ extern "C" __global__ void compute_advantage_return(
     const int b = blockIdx.x * blockDim.x + threadIdx.x;
     if (b >= b_size) return;
 
-    const float gamma = isv[RL_GAMMA_INDEX];
-    const float r     = rewards[b];
-    const float done  = dones[b];
-    const float vt    = v_t[b];
-    const float vtp1  = v_tp1[b];
+    const float gamma  = isv[RL_GAMMA_INDEX];
+    const float r      = rewards[b];
+    const float is_done = (dones[b] > 0.5f);
+    const float vt     = v_t[b];
+    const float vtp1   = v_tp1[b];
 
-    const float ret = r + gamma * (1.0f - done) * vtp1;
-    returns[b]    = ret;
-    advantages[b] = done * (ret - vt);
+    // 2026-05-29: branch-gate (not multiplication-gate) the V_tp1 term
+    // and the (ret - vt) advantage. Multiplication-gating fails on IEEE
+    // `0 * Inf = NaN`: if any batch's encoder produces a non-finite V
+    // prediction early in training (random init outliers), the prior
+    // `done * (ret - vt)` multiplication propagated NaN to ALL non-done
+    // batches' advantages, which then broke compute_advantage_rms (sum
+    // of A² → NaN) and PPO (A/RMS → NaN, ratio×A → NaN, l_pi → NaN).
+    // Branch-gating ensures non-finite vtp1/vt are NEVER mixed into a
+    // non-done batch's advantage. Validated by the smoke bisect at
+    // 10d4614fb (deterministic NaN at step 4 with l_v=6.329 stable —
+    // proving V REGRESSION is fine while V FORWARD has at least one
+    // non-finite output that the 0*Inf trick was promoting to NaN
+    // everywhere). Per `pearl_atomicadd_masks_v_instability`.
+    const float ret = is_done ? r : (r + gamma * vtp1);
+    returns[b]      = ret;
+    advantages[b]   = is_done ? (ret - vt) : 0.0f;
 }

crates/ml-alpha/cuda/variable_selection.cu

@@ -26,6 +26,24 @@
 #define VSN_FEATURE_DIM 40
 #define VSN_BLOCK       64    // round up to warp-multiple; threads i >= FEATURE_DIM idle.
 
+// 2026-05-29 stride-mismatch fix.
+// VSN's input buffer (window_tensor_d) is allocated [B, K, ENCODER_INPUT_DIM=56]
+// by perception.rs (snap features [0..40) + per-batch broadcast context
+// [40..56) written by rl_encoder_context_broadcast). VSN only processes the
+// first VSN_FEATURE_DIM=40 features per row (snap features), but the input
+// rows are spaced 56 floats apart, not 40. The kernel originally indexed x
+// with stride VSN_FEATURE_DIM=40 — correct ONLY for row 0; every subsequent
+// row read mixed broadcast-context + snap features across the [B, K, 56]
+// row boundaries. Symptoms: intermittent step-4 NaN as accumulating trade
+// context magnitudes overflowed VSN's softmax via the bleed.
+//
+// Fix: use VSN_X_ROW_STRIDE=56 for reading x in both forward and backward.
+// Output (gates, y) and gradient outputs (grad_W, grad_b, grad_x) remain at
+// VSN_FEATURE_DIM=40 because the downstream consumers (Mamba2 L1 with
+// in_dim=40) read at compact 40-stride. grad_x is unused downstream (see
+// `vsn_grad_x_d` audit — write-only), so its stride doesn't matter.
+#define VSN_X_ROW_STRIDE 56   // = ENCODER_INPUT_DIM in heads.rs / perception.rs
+
 extern "C" __global__ void variable_selection_fwd(
     const float* __restrict__ W_vsn,         // [FEATURE_DIM, FEATURE_DIM]
     const float* __restrict__ b_vsn,         // [FEATURE_DIM]
@@ -38,7 +56,7 @@ extern "C" __global__ void variable_selection_fwd(
     int tid = threadIdx.x;
     if (row >= n_rows) return;
 
-    const float* x_row = x + (long long)row * VSN_FEATURE_DIM;
+    const float* x_row = x + (long long)row * VSN_X_ROW_STRIDE;
 
     // Shared mem: gate_logit + max-reduce scratch + sum-reduce scratch.
     __shared__ float s_logit[VSN_FEATURE_DIM];
@@ -170,7 +188,7 @@ extern "C" __global__ void variable_selection_bwd(
     float dy_i = 0.0f;
     if (tid < VSN_FEATURE_DIM) {
         gates_i = gates[(long long)row * VSN_FEATURE_DIM + tid];
-        x_i = x[(long long)row * VSN_FEATURE_DIM + tid];
+        x_i = x[(long long)row * VSN_X_ROW_STRIDE + tid];
         dy_i = grad_y[(long long)row * VSN_FEATURE_DIM + tid];
         s_gates[tid] = gates_i;
         s_dgates[tid] = dy_i * x_i;
@@ -201,7 +219,7 @@ extern "C" __global__ void variable_selection_bwd(
         const float dl_t = s_dlogit[tid];
         #pragma unroll
         for (int j = 0; j < VSN_FEATURE_DIM; ++j) {
-            const float xj = x[(long long)row * VSN_FEATURE_DIM + j];
+            const float xj = x[(long long)row * VSN_X_ROW_STRIDE + j];
             grad_W_vsn_scratch[row_FF + (long long)tid * VSN_FEATURE_DIM + j]
                 += dl_t * xj;
         }