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>
286ea26e2a