gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(ml-alpha): warp-shuffle reduce in per-horizon kernels

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 13:15:33 +02:00
parent 83546b5c37
commit 286ea26e2a
5 changed files with 132 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

124
crates/ml-alpha/cuda/per_horizon_attention_pool.cu
View File

@@ -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) {

42
crates/ml-alpha/cuda/per_horizon_prob_blend.cu
View File

@@ -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;
}

27
crates/ml-alpha/cuda/per_horizon_residual_head.cu
View File

@@ -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();
}

12
crates/ml-alpha/src/per_horizon_attention_pool.rs
View File

@@ -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<f32>,
attn_h: &mut CudaSlice<f32>,
) -> Result<()> {
let smem_bytes = ((k_seq as usize + PHA_BLOCK + PHA_HIDDEN_DIM) * std::mem::size_of::<f32>())
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::<f32>()) 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<f32>,
grad_ln_out: &mut CudaSlice<f32>,
) -> 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::<f32>()) as u32;
((2 * k_seq as usize + PHA_N_WARPS) * std::mem::size_of::<f32>()) as u32;
let cfg = LaunchConfig {
grid_dim: (n_batch as u32, 1, 1),
block_dim: (PHA_BLOCK as u32, 1, 1),

3
crates/ml-alpha/src/trainer/per_horizon_state.rs
View File

@@ -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);