[ET-VK] softmax_buffer normalises each row with 4 threads and aggregates serially in every thread

[msluszniak]

🐛 Describe the bug

softmax_buffer.glsl hardcodes

#define NWORKERS 4
#define MAX_NTHREADS 16

and pick_softmax_gwg() matches it with lwg_extents[reduce_dim] = 4u, so four threads normalise a softmax row however long it is.

On all-MiniLM-L6-v2 at its published 254-token shape, Adreno 840, that is six dispatches at 633 us each, global {1,256,12} local {4,1,1}: 21.2% of GPU time (3.80 ms of 17.96 ms), second only to the linear layers.

There is a second problem behind it. The aggregation after each barrier is a serial walk of all NWORKERS shared entries run by every thread, not just thread 0:

T max_val = shared_max[0];
for (int i = 1; i < NWORKERS; ++i) {
  max_val = max(max_val, shared_max[i]);
}

so simply widening the group makes that walk proportionally more expensive and recovers only a quarter of the time. Both passes (max and sum) need a shared-memory tree.

Fix in #22349: all-MiniLM-L6-v2 @254 goes 18.30 -> 15.53 ms on an Adreno 840 (-15.1%), of which the tree is the larger part; softmax itself 3.80 -> 2.91 ms. Output matches the XNNPACK build to cosine 0.9999975 and is 60/60 bit-identical.

The texture-path softmax.glsl has the same constants and is not covered by that PR.

Versions

main @ 60cb889

[Gasoonjia]

@SS-JIA mind take a look?