From 7e38e46e653baa5fe38ef9ba64be94cfb34128c2 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 29 May 2026 00:30:14 +0200 Subject: [PATCH] =?UTF-8?q?fix(cuda):=20h=5Fmag=5Fper=5Fbucket=5Fkernel=20?= =?UTF-8?q?multi-warp=20reduction=20(32=E2=86=92128=20lanes)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Latent bug surfaced by f11bab542 test audit. `h_mag_per_bucket_kernel` launched with block_dim=32 (single warp) but BUCKET_DIMS terciles are [43, 43, 42] — channels 32..bdim were silently dropped from the sum, producing under-counted Controller D dead-bucket signal in production (perception.rs::h_mag_cfg). Fix: grow block_dim to 128 (smallest pow2 ≥ MAX_BUCKET_DIM=96), expand shared mem from `__shared__ float sdata[32]` to `sdata[128]`, change reduction stride from 16→32→64 (single-warp shuffle equivalent) to 64→32→16→8→4→2→1 (multi-warp block-tree). No atomicAdd, no warp divergence in tail lanes (uniform predicate tid < bdim). Changes: - bucket_transition_kernels.cu: kernel block-tree reduction over 128 lanes, shared mem 128 floats, launch comment updated - perception.rs::h_mag_cfg: block_dim 32→128, shared 32*4→128*4 - bucket_transition_kernels.rs test: launch config matches new contract + bug-surfaced comment replaced with fix-landed reference Test: tests/bucket_transition_kernels::h_mag_per_bucket_kernel_computes _mean_abs_per_bucket now passes (was deliberately failing in f11bab542 to surface this bug per feedback_no_todo_fixme — tests assert invariants not observed behavior). Co-Authored-By: Claude Opus 4.7 --- .../cuda/bucket_transition_kernels.cu | 27 +++++++++------- crates/ml-alpha/src/trainer/perception.rs | 4 +-- .../tests/bucket_transition_kernels.rs | 31 +++++-------------- 3 files changed, 26 insertions(+), 36 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/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 0213945e8..d9d71599c 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -5749,8 +5749,8 @@ impl PerceptionTrainer { h_per_k_base + ((k_seq - 1) * kb_hid_bytes) as u64; let h_mag_cfg = LaunchConfig { grid_dim: (N_HORIZONS as u32, 1, 1), - block_dim: (32, 1, 1), - shared_mem_bytes: 32 * 4, + block_dim: (128, 1, 1), + shared_mem_bytes: 128 * 4, }; { let mut args = RawArgs::new(); diff --git a/crates/ml-alpha/tests/bucket_transition_kernels.rs b/crates/ml-alpha/tests/bucket_transition_kernels.rs index 7d33bfcf2..097ae9851 100644 --- a/crates/ml-alpha/tests/bucket_transition_kernels.rs +++ b/crates/ml-alpha/tests/bucket_transition_kernels.rs @@ -647,8 +647,8 @@ fn h_mag_per_bucket_kernel_computes_mean_abs_per_bucket() -> Result<()> { let b_i32: i32 = b_sz as i32; let cfg = LaunchConfig { grid_dim: (N_HORIZONS as u32, 1, 1), - block_dim: (32, 1, 1), - shared_mem_bytes: 32 * 4, + block_dim: (128, 1, 1), + shared_mem_bytes: 128 * 4, }; let mut launch = stream.launch_builder(&func); launch @@ -663,27 +663,12 @@ fn h_mag_per_bucket_kernel_computes_mean_abs_per_bucket() -> Result<()> { stream.synchronize()?; let h_mag = download(&stream, &h_mag_d)?; - // KERNEL BUG SURFACED — DO NOT NORMALISE THIS TEST. - // - // `h_mag_per_bucket_kernel` launches with `block_dim = 32` (single - // warp; the kernel uses `__shared__ float sdata[32]` and reduces - // over 32 lanes). But this test's BUCKET_DIMS are [43, 43, 42] — - // wider than the warp. Only the first 32 channels per bucket - // contribute to the local sum, while the kernel divides by - // `bdim × B`, producing an UNDER-COUNTED mean for any bucket with - // bdim > 32. Production launches the same kernel via - // `perception.rs::h_mag_cfg = LaunchConfig { block_dim: (32, ..) }` - // against the same tercile (bdim=43/43/42) layout — so production - // is silently mis-computing Controller D's dead-bucket signal. - // - // Fix path: either grow the kernel to a multi-warp reduction - // (likely a `block_dim = 128` grid-stride over bdim), or shrink - // HIDDEN_DIM / N_HORIZONS until max(bdim) ≤ 32. This test asserts - // the IDEAL behavior (full-bucket mean); it WILL FAIL until the - // kernel is fixed. Per - // `feedback_tests_must_prove_not_lock_observations`: tests assert - // invariants, not observed behavior — surfacing latent bugs is the - // point. + // Asserts full-bucket mean: kernel now launches at block_dim=128 + // (multi-warp tree reduction) covering MAX_BUCKET_DIM=96 channels. + // Previously block_dim=32 silently dropped channels 32..bdim — fix + // landed in bucket_transition_kernels.cu h_mag_per_bucket_kernel + // (2026-05-29). Production launch site at perception.rs::h_mag_cfg + // mirrors the 128-lane config. let bucket_starts: [usize; 3] = [0, BUCKET_CUTS[0], BUCKET_CUTS[1]]; for k in 0..N_HORIZONS { let start = bucket_starts[k];