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

Required reasoning process

1. Choose single-pass vs two-pass strategy. A single block can reduce up to blockDim.x elements in one pass. For inputs larger than one block, two strategies exist:

• Two-pass: launch a first kernel that reduces chunks to per-block partial results, then launch a second kernel to reduce those partials. Simple, deterministic if implemented correctly, preferred for reproducibility.
• Atomic accumulation: each block reduces its chunk and atomically combines the result into a global accumulator. Simpler launch logic, but non-deterministic for floating-point ops due to non-associativity of floating-point addition. Acceptable for max/min/int ops where atomics are exact.
• Cooperative groups grid sync: all blocks cooperate in a single kernel launch using cooperative_groups::this_grid().sync(). Requires cooperative launch (cudaLaunchCooperativeKernel) and limits grid to what fits resident on the GPU. Use only if a single-kernel solution is architecturally required.

1. Design the warp-level reduction. Use __shfl_down_sync for all warp-level reductions. The standard pattern reduces 32 lanes to 1:

`` for (int offset = 16; offset > 0; offset >>= 1) val = op(val, __shfl_down_sync(0xffffffff, val, offset)); ` The mask 0xffffffff is correct only when all 32 lanes in the warp are active. If the thread count is not a multiple of 32 (partial warp at the tail of the input), use a mask that includes only the active lanes: compute it as __ballotsync(0xffffffff, threadis_active)` or predicate the loop.

1. Design the block-level reduction. After warp reduction, each warp has a partial result in lane 0. Collect these into shared memory (one value per warp, so blockDim.x / 32 values), sync, then reduce that small array using the first warp. Total smem needed: (blockDim.x / 32) * sizeof(dtype).

1. Handle the batched / axis reduction case. For a [B, L] tensor reduced along axis=1 to [B]:

• Assign one or more thread blocks per row. If L fits in one block, assign one block per row, with grid.x = B.
• If L does not fit in one block, use a two-pass approach: first kernel writes per-block partials to a [B, numblocksper_row] intermediate tensor; second kernel reduces along axis=1 of that intermediate.
• Alternatively, assign one warp per row for small L (e.g., L = N must not read from global memory and must contribute the identity element to the reduction. Guard with `if (idx 1024, two-pass multi-block reduction per row.
• Occupancy: reduction kernels typically have low arithmetic intensity, so high occupancy (to hide memory latency) is important. Minimize smem usage (only blockDim.x / 32 elements needed for the warp-partial stage). Aim for 50-100% theoretical occupancy.

Output format

The final response must include:

1. Strategy decision: single-pass or two-pass, with justification based on input size and operator properties.
2. Kernel code: complete, compilable CUDA kernel(s). If two-pass, both kernels plus the host dispatch function.
3. Warp reduction helper: a device function for the warp-level reduction with explicit mask and operator.
4. Block reduction helper: a device function for the block-level reduction using shared memory.
5. Host dispatch: kernel launch parameters (block size, grid size), temporary buffer allocation if needed, kernel calls in the correct order.
6. Correctness notes: explicit statement of identity element used, boundary handling strategy, warp mask rationale.
7. Numerical precision notes: accumulator dtype, any loss of precision relative to full fp64 reference.
8. Known limitations: cases where this kernel will produce non-deterministic results or underperform.

Common failure modes

• Incorrect warp mask in __shfl_down_sync: using 0xffffffff when the warp is not fully active (e.g., when N is not a multiple of 32). The inactive lanes participate in the shuffle with undefined values, corrupting the result. Fix: compute the active mask with __ballot_sync or predicate the loop.
• Missing __syncthreads between smem write and read: warp leader writes to smem[warpIdx], but another warp reads smem[0] before all writes complete. This is a race condition producing non-deterministic wrong results. Fix: add __syncthreads() immediately after all warp leaders have written, before any reads.
• Atomic accumulation without identity initialization: the global accumulator retains a value from a previous kernel call or is uninitialized. Fix: zero (or set to identity) the output buffer before the kernel launch. Do not assume CUDA will zero device memory.
• Non-deterministic floating-point results: atomic adds to a shared fp32 accumulator from many blocks produce different results on different runs due to non-deterministic ordering. This can cause training non-reproducibility. Fix: use a two-pass deterministic reduction if reproducibility is required; document the non-determinism explicitly otherwise.
• Reduction over warp partial smem array using wrong size: after warp reduction, there are blockDim.x / 32 values in smem. If the first warp uses 0xffffffff as the shfl mask but there are fewer than 32 values (e.g., block size 256 → 8 warp partials), lanes 8–31 read uninitialized smem. Fix: pad smem to 32 entries with identity, or use the correctly sized mask `(1u = N) excluded from global memory reads and contributing the identity?
• [ ] For two-pass: is the partial buffer large enough (gridDim.x elements), and is it the correct dtype (accumulator dtype, not input dtype)?
• [ ] For atomic strategy: is the output buffer initialized to the identity before the kernel launch, from the host side?
• [ ] For fp16/bf16 inputs: is accumulation happening in fp32?
• [ ] Is the warp partial smem reduction using the correct mask for blockDim.x / 32 entries, not 0xffffffff?
• [ ] Has the kernel been validated against a reference CPU reduction on: N=1, N=32, N=33, N=1024, N=1025, N=large?
• [ ] Is non-determinism documented if atomics are used for floating-point ops?
• [ ] For batched reductions: does the kernel handle B=1 and variable L correctly?
• [ ] Is CUB considered and explicitly rejected with a reason, or recommended as the better choice?

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: tensormux
• Source: tensormux/kernel-skills
• License: MIT

Install and usage instructions live in the source repository linked above.