diff --git a/crates/ml-alpha/cuda/nan_scan.cu b/crates/ml-alpha/cuda/nan_scan.cu index 2bfe38637..1623e6a14 100644 --- a/crates/ml-alpha/cuda/nan_scan.cu +++ b/crates/ml-alpha/cuda/nan_scan.cu @@ -269,6 +269,9 @@ extern "C" __global__ void nan_scan( case 44: label = "mamba2_l1_fwd_b_proj_pre_scan"; break; case 45: label = "mamba2_l1_fwd_x_pre_scan"; break; case 46: label = "mamba2_l1_fwd_h_s2_pre_scan"; break; + case 47: label = "train_vsn_out"; break; + case 48: label = "step_entry_w_in_W"; break; + case 49: label = "step_entry_w_in_b"; break; default: break; } const int step = (int)isv[RL_STEP_COUNTER_ISV_INDEX]; diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 913476b80..aad1cd442 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -5200,6 +5200,21 @@ impl PerceptionTrainer { (w.raw_ptr(), w.len()) }; self.nan_scan(41, l1_w_a_ptr, l1_w_a_len)?; + // 2026-05-29: w_in weights at step entry were uncovered by 41 + // (which only covers w_a). PROJ verdict identified the L1 w_in + // gemm path as a NaN producer — scan its weights here to + // disambiguate "w_in weights NaN at step entry" vs + // "w_in weights clean but gemm produces NaN from clean inputs". + let (l1_w_in_w_ptr, l1_w_in_w_len) = { + let w = &self.trunk.mamba2_l1().w_in.weight; + (w.raw_ptr(), w.len()) + }; + self.nan_scan(48, l1_w_in_w_ptr, l1_w_in_w_len)?; + let (l1_w_in_b_ptr, l1_w_in_b_len) = { + let b = &self.trunk.mamba2_l1().w_in.bias; + (b.raw_ptr(), b.len()) + }; + self.nan_scan(49, l1_w_in_b_ptr, l1_w_in_b_len)?; let (l2_w_a_ptr, l2_w_a_len) = { let w = &self.trunk.mamba2_l2().w_a.weight; (w.raw_ptr(), w.len()) @@ -5307,6 +5322,12 @@ impl PerceptionTrainer { .map_err(|e| anyhow::anyhow!("variable_selection_fwd: {:?}", e))?; } } + // 2026-05-29: scan VSN output in the TRAIN path (label 6 only + // covers the forward-graph dispatch path, not this train path). + // PROJ verdict from prior coder identified Mamba2 L1 projections + // as NaN producers — disambiguate "VSN output NaN feeds + // gemms" vs "VSN out clean, gemms still NaN". + self.nan_scan(47, self.vsn_out_d.cuda_data().raw_ptr(), self.vsn_out_d.numel())?; // ── 2. Mamba2 stack-1 forward — writes into mamba2_fwd_scratch. // Consumes vsn_out_d (gated features). in_dim = FEATURE_DIM (40).