CUDA ML Optimization

#include <cuda_runtime.h>

__global__ void softmax_kernel(float* input, float* output, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        float max_val = input[idx];
        for (int i = 0; i < size; i++) {
            max_val = fmaxf(max_val, input[i]);
        }
        float sum = 0.0f;
        for (int i = 0; i < size; i++) {
            sum += expf(input[i] - max_val);
        }
        output[idx] = expf(input[idx] - max_val) / sum;
    }
}

void launch_softmax(float* input, float* output, int size) {
    int threads = 256;
    int blocks = (size + threads - 1) / threads;
    softmax_kernel<<<blocks, threads>>>(input, output, size);
}