write-cuda-reduction-kernel

Skill: Write CUDA Reduction Kernel

Purpose

Guide the agent through designing and implementing a correct, efficient CUDA reduction kernel for a given operator (sum, max, min, or custom binary associative op), covering warp-level primitives, block-level reduction, multi-block strategies, and when to use CUB instead.

Use this when

• You need a reduction over a 1D array, a specific axis of a multi-dimensional tensor, or a segmented reduction with irregular segment sizes
• The reduction operator is non-standard (e.g., log-sum-exp, online variance update, argmax with index tracking) and is not directly supported by CUB or Thrust
• You need to fuse the reduction with a preceding or following per-element transformation and cannot afford the extra memory round-trip
• You are implementing a custom training loop component (e.g., gradient norm, loss reduction) where you need exact control over accumulation order or precision

Do not use this when

• The reduction is a standard sum/min/max/count over a contiguous array: use cub::DeviceReduce — it handles multi-block staging, SM-specific tuning, and dtype variants correctly and will outperform a first-attempt custom kernel
• The input is large (> 1M elements) and throughput is the only concern: CUB's DeviceReduce uses a highly tuned multi-block algorithm with kernel fusion
• You need segmented reductions over fixed-size segments: use cub::DeviceSegmentedReduce
• The reduction is over a batch of small vectors and you just need row-wise sums: a simple warp-per-row kernel may suffice; use that pattern instead of a full multi-block reduction

Inputs the agent should gather first

• Reduction operator: sum, max, min, product, logical AND/OR, argmax (value + index pair), custom binary op — the op must be associative; commutativity affects atomics strategy but is not strictly required
• Input dtype: fp32, fp16, bf16, int32, int64, uint8; whether mixed precision (e.g., fp16 input, fp32 accumulator) is needed
• Input shape: total element count; whether it is a 1D flat reduction or a reduction along an axis of a multi-dimensional tensor (e.g., reduce axis=1 of a [B, L] tensor → output shape [B])
• Memory layout: contiguous or strided input; stride values for the reduction axis and non-reduction axes
• Numerical precision requirements: is fp32 accumulation required for fp16 inputs, or is fp16 accumulation acceptable? Is the result expected to be deterministic across runs?
• Output: scalar output (single value), or one output per non-reduced dimension (batched reduction)
• Hardware target: SM architecture, for warp size (always 32 on current NVIDIA hardware), and to choose between atomics vs two-pass strategies