llm.c icon indicating copy to clipboard operation
llm.c copied to clipboard

Online Softmax is wrong

Open NoSavedDATA opened this issue 4 months ago • 0 comments

You should change

__global__ void softmax_forward_online_kernel8(float* out, const float* inp, int N, int C) {

    const int warpsPerBlock = blockDim.x / warpSize;
    int tid = threadIdx.x;

    if (tid >= C) {
        return;
    }

    int warpId = tid / warpSize;
    int laneId = tid % warpSize;
    // one warp one row
    int row = blockIdx.x * warpsPerBlock + warpId;

    if (row >= N) {
        return;
    }

Into

    const int warpsPerBlock = blockDim.x / warpSize;
    int tid = threadIdx.x;


    int warpId = tid / warpSize;
    int laneId = tid % warpSize;
    // one warp one row
    int row = blockIdx.x * warpsPerBlock + warpId;

    if (laneId >= C) {
        return;
    }

    if (row >= N) {
        return;
    }

NoSavedDATA avatar Oct 04 '24 02:10 NoSavedDATA