reduction

A sum reads every byte once and does one add per element, so the only question is whether you reach the bandwidth ceiling. The ceiling in milliseconds, four versions from an atomic per element to warp shuffles with sixteen loads in flight, and the bytes-in-flight math behind the last jump.

Normalization reads a row and writes a row, so a copy sets the ceiling and everything else is overhead you can remove. The six changes in the order a reviewer wants them, what each is worth, and the one that is a correctness fix rather than a speed fix, with the measured error that proves it.

A row of logits is half a megabyte, so the kernel is not bandwidth-bound and the cost is in how many passes you make over it. Why a full sort does far more work than the question asks, the two shapes that fit small and large k, and the four-pass radix select with a reference that matches a sort exactly.