parallelize part of logits processing

[kroggen]

Apply the temperature and softmax to the logits using the GPU

The argmax and sample functions were not changed

This has a drastic improvement in speed, of ~76%, when the temperature is not 0

Tested using the stories110M.bin model with an RTX 3090 with PCIE 4.0 at 24GB/s

[kroggen]

Note: the above increase in performance is with all the other PRs merged in a single file.

[kroggen]

Now I also implemented the argmax on the GPU, so it increase performance when using temperature = 0

[ankan-ban]

I knew logits processing on CPU is a bottleneck and I was planning to address this and so far I have been benchmarking with temp = 0. Thank you for the change. Will hopefully spend some time in testing/understanding it and merge soon.

[kroggen]

The argmax function can be enhanced further. The return value can be allocated once on the GPU together with the other arrays on the RunState, and released with them.

I just made it fast, but noticed that it is working as expected. The result is the same as before, with temperature 0

[kroggen]

There is a simpler way to implement parallel argmax, using the atomicMax that performs an atomic operation:

__global__ void argmax32_kernel(const float* __restrict__ v, int n, int* max_pos, float* max_val) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid < n) {
        atomicMax(max_val, v[tid]);
        __syncthreads();
        if(*max_val == v[tid]){
            *max_pos = tid;
        }
    }
}

max_val and max_pos are allocated on GPU.

In this case it can be launched with many blocks.

But I suspected that it would be slower because of the many threads stuck at the lock and just one making a step at a time (how I imagine it works, maybe I am wrong)

I did not benchmark both to compare though

[kroggen]

My suggestion: merge this PR and then later you can attempt other approaches for even higher performance