Didn't expect, but resulted in large-degree performance loss:

One extreme case is:

It was around 1.3 TB/s and is just 468 GB/s with this PR. It seems it's due to increased register usage as vectorized domains are located inside CA positions. And this particular case resulted in using 68 registers, whereas previously used 58 registers. While the difference is not significant, it triggered to cross over an occupancy cliff, resulting in half of the original occupancy.

It's also interesting that most of the perf loss happen only with softmax benchmarks, so that probably has something to do with the structure of persistent reductions.

Since this is a non trivial heuristic decision, and it seems mostly negative as is, I'm going to close the PR for now.