From 286ea26e2a0c02f557237f40225cd16dfdd433b4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 13:15:33 +0200 Subject: [PATCH] perf(ml-alpha): warp-shuffle reduce in per-horizon kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cluster A/B sweep with C25 wiring showed 86 s/epoch vs 17 s/epoch baseline = ~5x regression. Root cause: per-horizon attention pool + residual head used block tree-reduce with 8 __syncthreads per K-step in a serialised K-loop, repeated H=5 times in both fwd and bwd = ~3200 barriers/step. Plus the prob_blend bwd reduce kernel ran with a single thread per block, fully serialising over K*B. Replacements: - per_horizon_attention_pool fwd/bwd: introduce block_reduce_sum / block_reduce_max helpers using intra-warp __shfl_xor_sync + cross-warp shuffle (1 syncthread per K instead of 8). Smem shrinks to [K + N_WARPS] / [2K + N_WARPS]. - per_horizon_residual_head fwd: same warp-shuffle reduce pattern. - per_horizon_prob_blend_reduce_alpha_residual: 1 thread → 1 warp per horizon, lane-strided reduction over K*B via shfl_xor_sync. Launch config updated to block_dim=(32,1,1). Tricky bug found while implementing: the cross-warp reduce in the residual head originally guarded `__shfl_xor_sync(0xffffffff, ...)` with `if (tid < PHR_N_WARPS)`, leaving 28 of 32 lanes in warp 0 outside the call. Mask 0xffffffff requires all 32 lanes to participate — divergence is UB and hung the full-pipeline smoke on Ampere/Ada. Fix: read s_warp via ternary into all 32 lanes, then shuffle inside `if (tid < 32)`. Matches the pattern used in block_reduce_sum. Verified locally on RTX 3050 (sm_86): per_horizon_attention_pool numgrad, per_horizon_residual_head numgrad, and per_horizon_full_pipeline_smoke (zero-init identity + non-zero end-to-end) all PASS. Co-Authored-By: Claude Opus 4.7 --- .../cuda/per_horizon_attention_pool.cu | 124 ++++++++++-------- .../ml-alpha/cuda/per_horizon_prob_blend.cu | 42 ++++-- .../cuda/per_horizon_residual_head.cu | 27 ++-- .../src/per_horizon_attention_pool.rs | 12 +- .../ml-alpha/src/trainer/per_horizon_state.rs | 3 +- 5 files changed, 132 insertions(+), 76 deletions(-) diff --git a/crates/ml-alpha/cuda/per_horizon_attention_pool.cu b/crates/ml-alpha/cuda/per_horizon_attention_pool.cu index 5bb5ab7ce..3de7d43df 100644 --- a/crates/ml-alpha/cuda/per_horizon_attention_pool.cu +++ b/crates/ml-alpha/cuda/per_horizon_attention_pool.cu @@ -23,12 +23,66 @@ // inside the block. Keeps grad_ln_out writes free of cross-block // races: each batch's [b, :, :] slice is touched by ONE block in bwd. // +// Reduction strategy (perf-critical): +// block_dim == HIDDEN_DIM == 128 == 4 warps. Dot products over the +// HIDDEN_DIM axis are reduced via warp-shuffle (`__shfl_xor_sync`) +// inside each warp, then a single cross-warp reduction over 4 lanes. +// This replaces ~10 `__syncthreads` per K-step with 1, ~10× fewer +// barriers on the hot inner loop. +// // Per `feedback_no_atomicadd.md`: block tree-reduce only, no atomics. #define PHA_HIDDEN_DIM 128 #define PHA_BLOCK 128 #define PHA_MAX_K 512 // safety cap; smoke uses K=16-64. #define PHA_N_HORIZONS 5 +#define PHA_N_WARPS (PHA_BLOCK / 32) // == 4 + +// Reduce a per-thread float across a 128-thread block (4 warps). +// `s_warp` must be [PHA_N_WARPS] in shared mem (caller-supplied). +// Result is broadcast to all threads in the block. +__device__ __forceinline__ float block_reduce_sum(float v, float* s_warp, int tid) { + // Intra-warp reduction. + for (int s = 16; s > 0; s >>= 1) v += __shfl_xor_sync(0xffffffff, v, s); + + const int lane = tid & 31; + const int warp = tid >> 5; + if (lane == 0) s_warp[warp] = v; + __syncthreads(); + + // Cross-warp reduction (4 lanes only). + float w = (tid < PHA_N_WARPS) ? s_warp[tid] : 0.0f; + if (tid < 32) { + for (int s = PHA_N_WARPS / 2; s > 0; s >>= 1) { + w += __shfl_xor_sync(0xffffffff, w, s); + } + } + // Re-broadcast via shared mem so all threads see the same scalar. + if (tid == 0) s_warp[0] = w; + __syncthreads(); + return s_warp[0]; +} + +// Block-wide max reduction (mirror of block_reduce_sum but with fmaxf). +__device__ __forceinline__ float block_reduce_max(float v, float* s_warp, int tid) { + for (int s = 16; s > 0; s >>= 1) { + v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, s)); + } + const int lane = tid & 31; + const int warp = tid >> 5; + if (lane == 0) s_warp[warp] = v; + __syncthreads(); + + float w = (tid < PHA_N_WARPS) ? s_warp[tid] : -INFINITY; + if (tid < 32) { + for (int s = PHA_N_WARPS / 2; s > 0; s >>= 1) { + w = fmaxf(w, __shfl_xor_sync(0xffffffff, w, s)); + } + } + if (tid == 0) s_warp[0] = w; + __syncthreads(); + return s_warp[0]; +} extern "C" __global__ void per_horizon_attention_pool_fwd( const float* __restrict__ Q_h, // [N_HORIZONS, HIDDEN_DIM] @@ -43,13 +97,11 @@ extern "C" __global__ void per_horizon_attention_pool_fwd( if (b_idx >= n_batch || tid >= PHA_BLOCK) return; extern __shared__ float smem[]; - float* s_scores = smem; // [K] - float* s_red = smem + k_seq; // [BLOCK] - float* s_context = smem + k_seq + PHA_BLOCK; // [HIDDEN_DIM] + float* s_scores = smem; // [K] + float* s_warp = smem + k_seq; // [PHA_N_WARPS] + // (no longer need a [BLOCK]-sized s_red — warp-shuffle reduce uses only N_WARPS slots) __shared__ float s_Qh[PHA_HIDDEN_DIM]; - __shared__ float s_max; - __shared__ float s_sum; const float* ln_b = ln_out + (long long)b_idx * k_seq * PHA_HIDDEN_DIM; @@ -62,16 +114,12 @@ extern "C" __global__ void per_horizon_attention_pool_fwd( __syncthreads(); // Pass 1: scores[k] = Q_h[h, :] · ln_out[b, k, :]. + // Warp-shuffle reduce → ~1 syncthreads per K instead of 8. for (int k = 0; k < k_seq; ++k) { - const float v = ln_b[k * PHA_HIDDEN_DIM + tid] * s_Qh[tid]; - s_red[tid] = v; - __syncthreads(); - for (int s = PHA_BLOCK / 2; s > 0; s >>= 1) { - if (tid < s) s_red[tid] += s_red[tid + s]; - __syncthreads(); - } - if (tid == 0) s_scores[k] = s_red[0]; - __syncthreads(); + float v = ln_b[k * PHA_HIDDEN_DIM + tid] * s_Qh[tid]; + float r = block_reduce_sum(v, s_warp, tid); + if (tid == 0) s_scores[k] = r; + __syncthreads(); // s_warp will be re-used; safe. } // Pass 2: softmax over K — max-subtract + exp + sum + divide. @@ -80,18 +128,7 @@ extern "C" __global__ void per_horizon_attention_pool_fwd( const float v = s_scores[k]; if (v > my_max) my_max = v; } - s_red[tid] = my_max; - __syncthreads(); - for (int s = PHA_BLOCK / 2; s > 0; s >>= 1) { - if (tid < s) { - const float a = s_red[tid]; - const float b = s_red[tid + s]; - s_red[tid] = (a > b) ? a : b; - } - __syncthreads(); - } - if (tid == 0) s_max = s_red[0]; - __syncthreads(); + const float s_max = block_reduce_max(my_max, s_warp, tid); float my_sum = 0.0f; for (int k = tid; k < k_seq; k += PHA_BLOCK) { @@ -99,14 +136,7 @@ extern "C" __global__ void per_horizon_attention_pool_fwd( s_scores[k] = e; my_sum += e; } - s_red[tid] = my_sum; - __syncthreads(); - for (int s = PHA_BLOCK / 2; s > 0; s >>= 1) { - if (tid < s) s_red[tid] += s_red[tid + s]; - __syncthreads(); - } - if (tid == 0) s_sum = s_red[0]; - __syncthreads(); + const float s_sum = block_reduce_sum(my_sum, s_warp, tid); // Pass 3: attn[k] = exp / sum; save; context[d] = Σ_k attn[k] * ln_out[b, k, d]. const long long attn_base = (long long)b_idx * PHA_N_HORIZONS * k_seq @@ -123,7 +153,6 @@ extern "C" __global__ void per_horizon_attention_pool_fwd( for (int k = 0; k < k_seq; ++k) { c += s_scores[k] * ln_b[k * PHA_HIDDEN_DIM + tid]; } - s_context[tid] = c; context_h[(long long)b_idx * PHA_N_HORIZONS * PHA_HIDDEN_DIM + (long long)h * PHA_HIDDEN_DIM + tid] = c; } @@ -163,12 +192,11 @@ extern "C" __global__ void per_horizon_attention_pool_bwd( extern __shared__ float smem[]; float* s_dattn = smem; // [K] float* s_dscores = smem + k_seq; // [K] - float* s_red = smem + 2 * k_seq; // [BLOCK] + float* s_warp = smem + 2 * k_seq; // [PHA_N_WARPS] __shared__ float s_Qh[PHA_HIDDEN_DIM]; __shared__ float s_attn[PHA_MAX_K]; __shared__ float s_grad_ctx[PHA_HIDDEN_DIM]; - __shared__ float s_sum_attn_dattn; const float* ln_b = ln_out + (long long)bi * k_seq * PHA_HIDDEN_DIM; float* grad_ln_b = grad_ln_out + (long long)bi * k_seq * PHA_HIDDEN_DIM; @@ -190,16 +218,11 @@ extern "C" __global__ void per_horizon_attention_pool_bwd( __syncthreads(); // Pass 1: d_attn[k] = Σ_d grad_context[d] * ln_out[b, k, d]. + // Warp-shuffle reduce per k (was 8 syncthreads per k — now 1). for (int k = 0; k < k_seq; ++k) { - const float v = (tid < PHA_HIDDEN_DIM) - ? s_grad_ctx[tid] * ln_b[k * PHA_HIDDEN_DIM + tid] : 0.0f; - s_red[tid] = v; - __syncthreads(); - for (int s = PHA_BLOCK / 2; s > 0; s >>= 1) { - if (tid < s) s_red[tid] += s_red[tid + s]; - __syncthreads(); - } - if (tid == 0) s_dattn[k] = s_red[0]; + float v = s_grad_ctx[tid] * ln_b[k * PHA_HIDDEN_DIM + tid]; + float r = block_reduce_sum(v, s_warp, tid); + if (tid == 0) s_dattn[k] = r; __syncthreads(); } @@ -208,14 +231,7 @@ extern "C" __global__ void per_horizon_attention_pool_bwd( for (int k = tid; k < k_seq; k += PHA_BLOCK) { my_sum += s_attn[k] * s_dattn[k]; } - s_red[tid] = my_sum; - __syncthreads(); - for (int s = PHA_BLOCK / 2; s > 0; s >>= 1) { - if (tid < s) s_red[tid] += s_red[tid + s]; - __syncthreads(); - } - if (tid == 0) s_sum_attn_dattn = s_red[0]; - __syncthreads(); + const float s_sum_attn_dattn = block_reduce_sum(my_sum, s_warp, tid); // Pass 3: d_scores[k] = attn[k] * (d_attn[k] - centring). for (int k = tid; k < k_seq; k += PHA_BLOCK) { diff --git a/crates/ml-alpha/cuda/per_horizon_prob_blend.cu b/crates/ml-alpha/cuda/per_horizon_prob_blend.cu index e40659f93..90068a555 100644 --- a/crates/ml-alpha/cuda/per_horizon_prob_blend.cu +++ b/crates/ml-alpha/cuda/per_horizon_prob_blend.cu @@ -63,9 +63,14 @@ extern "C" __global__ void per_horizon_prob_blend_fwd( // d_residual[b, h] = tanh(α[h]) * Σ_k d_logit_per_k[k, b, h] // d_alpha[h] = (1 - tanh²(α[h])) * Σ_{k, b} d_logit_per_k[k, b, h] * residual[b, h] // -// One block per horizon. Single-thread block (tid==0) does both -// reductions sequentially — sizes are tiny (K up to a few hundred, B -// up to 64) so a parallel reduce isn't worth the complexity. +// One block per horizon. Warp-parallel inside the block: 32 threads +// stride over the K*B index space, each maintaining its own k_sum +// per batch via warp shuffle, then a final warp shuffle for alpha_sum. +// d_residual[b, h] needs a per-batch sum (across k), so we let lane 0 +// of each warp write it after a warp reduction over the K stride. +// +// Grid: (N_HORIZONS, 1, 1). Block: (32, 1, 1) = one warp per horizon. +// Single-warp blocks → zero __syncthreads; all reduction via shuffle. extern "C" __global__ void per_horizon_prob_blend_reduce_alpha_residual( const float* __restrict__ alpha, // [N_HORIZONS] const float* __restrict__ residual, // [B, N_HORIZONS] @@ -77,17 +82,24 @@ extern "C" __global__ void per_horizon_prob_blend_reduce_alpha_residual( float* __restrict__ d_alpha // [N_HORIZONS] (written) ) { int h = blockIdx.x; - if (h >= PHB_N_HORIZONS || threadIdx.x != 0) return; + int lane = threadIdx.x; + if (h >= PHB_N_HORIZONS || lane >= 32) return; const float a = alpha[h]; const float ta = tanhf(a); const float sech2 = 1.0f - ta * ta; + // For each batch, all 32 threads cooperatively reduce over K via + // shuffle. Lane 0 writes d_residual[b, h]. alpha_sum accumulates + // across batches and is reduced once at the end. float alpha_sum = 0.0f; for (int b = 0; b < n_batch; ++b) { const float resid_bh = residual[b * PHB_N_HORIZONS + h]; + + // Per-thread strided sum over k. float k_sum = 0.0f; - for (int k = 0; k < k_seq; ++k) { + float k_resid_acc = 0.0f; + for (int k = lane; k < k_seq; k += 32) { const long long idx = (long long)k * n_batch * PHB_N_HORIZONS + (long long)b * PHB_N_HORIZONS @@ -95,10 +107,22 @@ extern "C" __global__ void per_horizon_prob_blend_reduce_alpha_residual( const float p = probs_per_k[idx]; const float gp = grad_probs[idx]; const float d_logit = gp * p * (1.0f - p); - k_sum += d_logit; - alpha_sum += d_logit * resid_bh; + k_sum += d_logit; + k_resid_acc += d_logit * resid_bh; } - d_residual[b * PHB_N_HORIZONS + h] = ta * k_sum; + // Warp-reduce both sums. + for (int s = 16; s > 0; s >>= 1) { + k_sum += __shfl_xor_sync(0xffffffff, k_sum, s); + k_resid_acc += __shfl_xor_sync(0xffffffff, k_resid_acc, s); + } + if (lane == 0) { + d_residual[b * PHB_N_HORIZONS + h] = ta * k_sum; + } + alpha_sum += k_resid_acc; // already warp-reduced above + } + + // Lane 0 already holds the full alpha_sum (added of warp-reduced values). + if (lane == 0) { + d_alpha[h] = sech2 * alpha_sum; } - d_alpha[h] = sech2 * alpha_sum; } diff --git a/crates/ml-alpha/cuda/per_horizon_residual_head.cu b/crates/ml-alpha/cuda/per_horizon_residual_head.cu index dd2f0140c..a634055f3 100644 --- a/crates/ml-alpha/cuda/per_horizon_residual_head.cu +++ b/crates/ml-alpha/cuda/per_horizon_residual_head.cu @@ -20,6 +20,7 @@ #define PHR_HIDDEN_DIM 128 #define PHR_BLOCK 128 #define PHR_N_HORIZONS 5 +#define PHR_N_WARPS (PHR_BLOCK / 32) extern "C" __global__ void per_horizon_residual_head_fwd( const float* __restrict__ context_h, // [B, N_HORIZONS, HIDDEN_DIM] @@ -32,23 +33,31 @@ extern "C" __global__ void per_horizon_residual_head_fwd( int tid = threadIdx.x; if (b >= n_batch || tid >= PHR_BLOCK) return; - __shared__ float s_red[PHR_BLOCK]; + // Warp-shuffle reduce: 1 partial per warp. + __shared__ float s_warp[PHR_N_WARPS]; for (int h = 0; h < PHR_N_HORIZONS; ++h) { // Per-thread partial: w_res[h, tid] * context_h[b, h, tid]. - const float v = (tid < PHR_HIDDEN_DIM) + float v = (tid < PHR_HIDDEN_DIM) ? w_res[h * PHR_HIDDEN_DIM + tid] * context_h[(long long)b * PHR_N_HORIZONS * PHR_HIDDEN_DIM + (long long)h * PHR_HIDDEN_DIM + tid] : 0.0f; - s_red[tid] = v; + // Intra-warp reduction via shuffle (zero barriers). + for (int s = 16; s > 0; s >>= 1) v += __shfl_xor_sync(0xffffffff, v, s); + if ((tid & 31) == 0) s_warp[tid >> 5] = v; __syncthreads(); - for (int s = PHR_BLOCK / 2; s > 0; s >>= 1) { - if (tid < s) s_red[tid] += s_red[tid + s]; - __syncthreads(); - } - if (tid == 0) { - residual_out[(long long)b * PHR_N_HORIZONS + h] = s_red[0] + bias_res[h]; + // Cross-warp reduce: WHOLE warp 0 must call __shfl_xor_sync with + // mask=0xffffffff (any divergence within the warp is UB and hangs + // on Ampere/Ada). Inactive lanes contribute 0. + float w = (tid < PHR_N_WARPS) ? s_warp[tid] : 0.0f; + if (tid < 32) { + for (int s = PHR_N_WARPS / 2; s > 0; s >>= 1) { + w += __shfl_xor_sync(0xffffffff, w, s); + } + if (tid == 0) { + residual_out[(long long)b * PHR_N_HORIZONS + h] = w + bias_res[h]; + } } __syncthreads(); } diff --git a/crates/ml-alpha/src/per_horizon_attention_pool.rs b/crates/ml-alpha/src/per_horizon_attention_pool.rs index e1b3345a3..c3f3cb4ff 100644 --- a/crates/ml-alpha/src/per_horizon_attention_pool.rs +++ b/crates/ml-alpha/src/per_horizon_attention_pool.rs @@ -19,6 +19,8 @@ use std::sync::Arc; pub const PHA_HIDDEN_DIM: usize = 128; pub const PHA_BLOCK: usize = 128; pub const PHA_N_HORIZONS: usize = 5; +/// Warp-shuffle reduce footprint inside the block: 1 partial per warp. +pub const PHA_N_WARPS: usize = PHA_BLOCK / 32; const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/per_horizon_attention_pool.cubin")); @@ -63,8 +65,10 @@ impl PerHorizonAttentionPool { context_h: &mut CudaSlice, attn_h: &mut CudaSlice, ) -> Result<()> { - let smem_bytes = ((k_seq as usize + PHA_BLOCK + PHA_HIDDEN_DIM) * std::mem::size_of::()) - as u32; + // Smem layout: s_scores[K] + s_warp[N_WARPS]. Warp-shuffle reduce + // removed the prior [BLOCK] + [HIDDEN_DIM] regions. + let smem_bytes = + ((k_seq as usize + PHA_N_WARPS) * std::mem::size_of::()) as u32; let cfg = LaunchConfig { grid_dim: (n_batch as u32, 1, 1), block_dim: (PHA_BLOCK as u32, 1, 1), @@ -102,8 +106,10 @@ impl PerHorizonAttentionPool { grad_qh_scratch: &mut CudaSlice, grad_ln_out: &mut CudaSlice, ) -> Result<()> { + // Smem layout: s_dattn[K] + s_dscores[K] + s_warp[N_WARPS]. + // Warp-shuffle reduce removed the prior [BLOCK] region. let smem_bytes = - ((2 * k_seq as usize + PHA_BLOCK) * std::mem::size_of::()) as u32; + ((2 * k_seq as usize + PHA_N_WARPS) * std::mem::size_of::()) as u32; let cfg = LaunchConfig { grid_dim: (n_batch as u32, 1, 1), block_dim: (PHA_BLOCK as u32, 1, 1), diff --git a/crates/ml-alpha/src/trainer/per_horizon_state.rs b/crates/ml-alpha/src/trainer/per_horizon_state.rs index 78f471882..e8cee4109 100644 --- a/crates/ml-alpha/src/trainer/per_horizon_state.rs +++ b/crates/ml-alpha/src/trainer/per_horizon_state.rs @@ -270,9 +270,10 @@ impl PerHorizonTrainState { // 1. Reduce: compute d_logit_baseline = grad_probs * p * (1-p) // inline + reduce over k for d_residual, over k+b for d_alpha. + // One warp per horizon: 32 threads cooperatively reduce via shuffle. let cfg_per_h = LaunchConfig { grid_dim: (N_HORIZONS as u32, 1, 1), - block_dim: (1, 1, 1), + block_dim: (32, 1, 1), shared_mem_bytes: 0, }; let mut launch = self.stream.launch_builder(&self.prob_blend_reduce_fn);