From fd317426271189e12d84370c1b2155bc12b7f409 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 29 May 2026 21:34:29 +0200 Subject: [PATCH] =?UTF-8?q?fix(cuda):=20Plan=20A=20v2=20=E2=80=94=20dd049d?= =?UTF-8?q?9a4=20+=203=20kernel-only=20bug=20fixes=20(no=20Phase=202.0)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../cuda/bucket_transition_kernels.cu | 27 ++++++++++------- .../ml-alpha/cuda/compute_advantage_return.cu | 29 ++++++++++++++----- crates/ml-alpha/cuda/variable_selection.cu | 24 +++++++++++++-- 3 files changed, 58 insertions(+), 22 deletions(-) diff --git a/crates/ml-alpha/cuda/bucket_transition_kernels.cu b/crates/ml-alpha/cuda/bucket_transition_kernels.cu index 73f106680..ec8c1e3fc 100644 --- a/crates/ml-alpha/cuda/bucket_transition_kernels.cu +++ b/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(); } diff --git a/crates/ml-alpha/cuda/compute_advantage_return.cu b/crates/ml-alpha/cuda/compute_advantage_return.cu index 45d016064..e850923ed 100644 --- a/crates/ml-alpha/cuda/compute_advantage_return.cu +++ b/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; } diff --git a/crates/ml-alpha/cuda/variable_selection.cu b/crates/ml-alpha/cuda/variable_selection.cu index b86b69bfa..6c25bec86 100644 --- a/crates/ml-alpha/cuda/variable_selection.cu +++ b/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; }