OMG I EVEN SAW THIS BEFORE AND GOT confused why we were storing twice into max_fvec0? But hmmm why do we store twice into min_fvec and zero_fvec in lines 1028-1031 above?

I spent some time last night trying to understand how this kernel works (by interrogating the LLM of my choice about it since I am below average at intuitively understanding array indexing code without lots of pictures for some reason), so I actually can answer that! Lines 1028-1031 are storing the identities for the max and sum reductions into our array of accumulators.

We're making a non-contiguous reduction vectorizable anyway by doing anarray of reductions all at once. (We can note that outer_size is basically just a batch dimension and ignore it for purposes of understanding the kernel.) We slice up the inner dimensions into chunks of length CHUNK_SIZE a; the inner loops are doing CHUNK_SIZE reductions at once. Accordingly, they have an array of CHUNK_SIZE accumulators, which is what's getting initialized in lines 1028-1031. since the CHUNK_SIZE dimension is contiguous, we can vectorize along it and get "parallelization" that way through vector arithmetic. The blocking/chunking stuff is so that the dim_size x CHUNK_SIZE "vertical panel" that we hand to each thread fits in cache, since softmax will end up reading the data 3 times (for each inner loop -- max, sum + log, and data - logsum - max).