/*  Kernels for fused forward/backward classifier partThis fuses softmax, crossentropy, and logit gradients into a single pass, so we don't have to write unnecessary(B, T, V) tensors. Such an operation is only possible if `dloss` can be known beforehand, which doesn't seem likemuch of a restriction: In pretraining, it is just a constant 1/batch_size tensor, for fine-tuning we might zeroout the input prompt, but that is known in advance. Compile example:nvcc -O3 --use_fast_math -lcublas -lcublasLt classifier_fused.cu -o classifier_fused ./classifier_fused 1./classifier_fused 2./classifier_fused 3./classifier_fused 4*/ #include <stdio.h>#include <stdlib.h>#include <float.h>#include <cuda_runtime.h>#include <cooperative_groups.h>#include <cooperative_groups/reduce.h>#include "common.h" // todo - this file does not properly support anything but FP32// kernel 5 can be run in fp16/bf16 to test performance, but the outputs will be wrong#if defined(ENABLE_BF16)typedef __nv_bfloat16 floatX;#elif defined(ENABLE_FP16)typedef half floatX;#elsetypedef float floatX;#endiftypedef Packed128<floatX> x128; // ----------------------------------------------------------------------------// CPU code reference void softmax_forward_cpu(float* out, const float* inp, int N, int C) {    // inp is (N, C)    // out is (N, C), each row of inp will get softmaxed    for (int64_t i = 0; i < N; i++) {        const float* inp_row = inp + i * C;        float* out_row = out + i * C;         float maxval = -INFINITY;        for (int j = 0; j < C; j++) {            if (inp_row[j] > maxval) {                maxval = inp_row[j];            }        }        double sum = 0.0;        for (int j = 0; j < C; j++) {            out_row[j] = expf(inp_row[j] - maxval);            sum += out_row[j];        }        for (int j = 0; j < C; j++) {            out_row[j] /= sum;        }    }}  void crossentropy_forward_cpu(float* losses,                              const float* probs, const int* targets,                              int B, int T, int V) {    // output: losses is (B,T) of the individual losses at each position    // input: probs are (B,T,V) of the probabilities    // input: targets is (B,T) of integers giving the correct index in logits    for (int64_t bt = 0; bt < B * T; bt++) {        // loss = -log(probs[target])        const float* probs_bt = probs + bt * V;        int ix = targets[bt];        losses[bt] = -logf(probs_bt[ix]);    }} void crossentropy_softmax_backward_cpu(float* dlogits,                                       const float* dlosses, const float* probs, const int* targets,                                       int B, int T, int V) {    // backwards through both softmax and crossentropy    for (int64_t bt = 0; bt < B * T; bt++) {        float* dlogits_bt = dlogits + bt * V;        const float* probs_bt = probs + bt * V;        float dloss = dlosses[bt];        int ix = targets[bt];        for (int i = 0; i < V; i++) {            float p = probs_bt[i];            float indicator = i == ix ? 1.0f : 0.0f;            dlogits_bt[i] = (p - indicator) * dloss;        }    }} // ----------------------------------------------------// Kernel Utils // warp-level reduction for finding the maximum value__device__ float warpReduceMax(float val) {    for (int offset = 16; offset > 0; offset /= 2) {        val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset));    }    return val;} // ----------------------------------------------------------------------------// GPU kernels struct SoftmaxParams {    float Scale;    float Offset;};namespace cg = cooperative_groups;__device__ SoftmaxParams prepare_softmax(cg::thread_block_tile<32>& warp,                                         int64_t idx, const float* inp, int V, int P) {    // this warp (of 32) threads processes one row of inp, i.e. inp[idx, :] of shape (V,)    // note that inp is actually (B * T, P) but we only use the first V elements    // this function tehen calculates:    // 1) the max value to subtract for numerical stability and    // 2) the sum normalization factor    const float* x = inp + idx * P;    // thread coarsening loop, where the 32 threads serially process all V elements    // thread_rank() is in [0, 31], warp.size() is 32    float maxval = -INFINITY;    float sumval = 0.0f;    for (int i = warp.thread_rank(); i < V; i += warp.size()) {        float v = x[i];        float old_maxval = maxval;        // online softmax recurrence from "Online normalizer calculation for softmax" paper        maxval = fmaxf(maxval, v);        sumval *= expf((old_maxval - maxval));        sumval += expf(v - maxval);    }    // warp-level reduction to get the maxval across the 32 threads    float global_maxval = cg::reduce(warp, maxval, cg::greater<float>{});    // all 32 threads do a final shift of the sum considering the global max in this row    sumval *= expf((maxval - global_maxval));    // warp-level reduction to get the sumval across the 32 threads    float global_sumval = cg::reduce(warp, sumval, cg::plus<float>{});    // the final normalization factor    float norm = 1.0f / global_sumval;    return SoftmaxParams{norm, global_maxval};} __global__ void fused_classifier_kernel1(float* dlogits, float* losses,                             const float* logits, const float* dlosses, const int* targets,                             int B, int T, int V, int P) {    namespace cg = cooperative_groups;    cg::thread_block block = cg::this_thread_block();    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);    // example: B = 4, T = 1024, block_size = 128 => we'd have grid_size = 1024    // each block of 4 warps is in charge of 4 rows of the input, one warp per row    // meta_group_size is the number of warps per block (e.g. 4)    // meta_group_rank is the index of the warp in the block (e.g. 0, 1, 2, 3)    int64_t idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();    if (idx >= B * T) { // there are B * T rows in the input        return;    }    int b = idx / T;    int t = idx % T;     // calculate the offset (maxval) and scale (sumval) for the softmax    SoftmaxParams sp = prepare_softmax(warp, idx, logits, V, P);     // in each row (handled by one warp), thread 0 calculates the loss    // calculate the probability needed for the loss and update losses    if(warp.thread_rank() == 0) {        int ix = targets[b * T + t];        float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale;        losses[b * T + t] = -logf(prob);    }     // finally all threads calculate the gradients    // prob is only materialized here temporarily and in registers, never    // as a full tensor that gets written to global memory    for (int i = warp.thread_rank(); i < V; i += warp.size()) {        float prob = expf(logits[idx * P + i] - sp.Offset) * sp.Scale;        float* dlogits_bt = dlogits + b * T * P + t * P;        float dloss = dlosses[b * T + t];        int ix = targets[b * T + t];        float indicator = i == ix ? 1.0f : 0.0f;        dlogits_bt[i] = (prob - indicator) * dloss;    }}  __device__ float vec_at(const float4& vec, int index) {    return reinterpret_cast<const float*>(&vec)[index];} __device__ SoftmaxParams prepare_softmax_blockwide(cg::thread_block_tile<32>& warp,                                                   int64_t idx, const float* inp, int V, int P) {    // one row of inp, i.e. inp[idx, :] of shape (V,)    // float4 to get 128-bit loads and memory level parallelism    const float4* x_vec4 = reinterpret_cast<const float4*>(inp + idx * P);     float thread_maxval = -INFINITY;    float thread_sumval = 0.0f;    // do the loop in reverse to maximise probability of L2 cache hits    // so even small L2s get some hits on the 2nd read of the same thread    for (int i = ceil_div(V, 4) + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) {        float4 v4 = x_vec4[i];        #pragma unroll        for(int k = 0; k < 4; k++) {            if (i*4+k >= V) {  // bounds checking against real V                continue;            }            float old_maxval = thread_maxval;            thread_maxval = fmaxf(thread_maxval, vec_at(v4, k));            thread_sumval *= expf(old_maxval - thread_maxval);            thread_sumval += expf(vec_at(v4, k) - thread_maxval);        }    }     // two reductions of up to 1024 threads:    // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle)    // this results in much cleaner assembly than a multi-warp cg::reduce    __shared__ float shared_maxval[32];    __shared__ float shared_sumval[32];    int num_warps = blockDim.x / 32;    int warp_id = threadIdx.x / 32;    int lane_id = threadIdx.x % 32;     // reduce maxval within each warp    float warp_maxval = cg::reduce(warp, thread_maxval, cg::greater<float>{});    // thread 0 in each warp writes to shared memory    if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; }    __syncthreads();    // each thread now loads the maxval across previous warps    // if the thread is "out of range" of data, use -FLT_MAX as the maxval    warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX;    // now reduce the maxval among the warp threads    float block_maxval = cg::reduce(warp, warp_maxval, cg::greater<float>{});    // each thread uses maxval to scale sumval to avoid numerical instability / overflow    thread_sumval *= expf(thread_maxval - block_maxval);    // (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory    float warp_sumval = cg::reduce(warp, thread_sumval, cg::plus<float>{});    if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; }    __syncthreads();    // same strategy, now reduce sumval across warps    warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f;    float block_sumval = cg::reduce(warp, warp_sumval, cg::plus<float>{});    // return the softmax parameters    return SoftmaxParams{1.f / block_sumval, block_maxval};} // Fused forward and backward pass for classifier including softmax, and logit gradients// Writes to both probs (only for debugging) and dlogits (only for training) are optional// N.B.: We may want to reuse the logits memory for dlogits, so they should *not* be __restrict__!__global__ void fused_classifier_kernel2(float* dlogits, float* losses, float* probs,                                         const float* logits, const float* dlosses, const int* targets,                                         int B, int T, int V, int P) {    namespace cg = cooperative_groups;    cg::thread_block block = cg::this_thread_block();    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);    int64_t idx = blockIdx.x;    int ix = targets[idx];     // softmax (reading B * T * V, same logits read again below, hopefully still in cache)    SoftmaxParams sp = prepare_softmax_blockwide(warp, idx, logits, V, P);     // calculate the probability needed for the loss and update (single-threaded)    if(threadIdx.x == 0) {        float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale;        losses[idx] = -logf(prob);    }     // very sensible default for dlosses is 1/(B*T), which is the uniform loss    float dloss = dlosses != NULL ? dlosses[idx] : 1.0f / (B*T);    // calculate the gradients directly, saves bandwidth from probs during training    // but also supports writing probs for inference-only and debugging    const float4* logits_vec4 = reinterpret_cast<const float4*>(logits + idx * P);    for (int i = threadIdx.x; i < ceil_div(V, 4); i += blockDim.x) {        // this is the 2nd read of logits after the one in prepare_softmax2        // this data will never be needed again, so we reduce cache persistence        float4 v4 = __ldcs(&logits_vec4[i]);         #pragma unroll        for(int k = 0; k < 4; ++k) {            int element = i*4 + k;            float prob = expf(vec_at(v4, k) - sp.Offset) * sp.Scale;            prob = (element < V) ? prob : 0.0f; // bounds checking against real V             // this kernel is DRAM limited so cost of inner branch is ~zero            if (probs != NULL) {                probs[idx * P + element] = prob;            }            if (dlogits != NULL) {                float indicator = element == ix ? 1.0f : 0.0f;                dlogits[idx * P + element] = (prob - indicator) * dloss;            }        }    }} __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_tile<32>& warp,                                                            int64_t idx, const float* inp, int V, int P) {    // same but not float4    // one row of inp, i.e. inp[idx, :] of shape (V,)     const float* x = inp + idx * P;    float thread_maxval = -INFINITY;    float thread_sumval = 0.0f;    // do the loop in reverse to maximise probability of L2 cache hits    // so even small L2s get some hits on the 2nd read of the same thread    for (int i = V + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) {        float v = x[i];        float old_maxval = thread_maxval;        thread_maxval = fmaxf(thread_maxval, v);        thread_sumval *= expf(old_maxval - thread_maxval);        thread_sumval += expf(v - thread_maxval);    }     // two reductions of up to 1024 threads:    // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle)    // this results in much cleaner assembly than a multi-warp cg::reduce    __shared__ float shared_maxval[32];    __shared__ float shared_sumval[32];    int num_warps = blockDim.x / 32;    int warp_id = threadIdx.x / 32;    int lane_id = threadIdx.x % 32;     // reduce maxval within each warp    float warp_maxval = cg::reduce(warp, thread_maxval, cg::greater<float>{});    // thread 0 in each warp writes to shared memory    if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; }    __syncthreads();    // each thread now loads the maxval across previous warps    // if the thread is "out of range" of data, use -FLT_MAX as the maxval    warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX;    // now reduce the maxval among the warp threads    float block_maxval = cg::reduce(warp, warp_maxval, cg::greater<float>{});    // each thread uses maxval to scale sumval to avoid numerical instability / overflow    thread_sumval *= expf(thread_maxval - block_maxval);    // (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory    float warp_sumval = cg::reduce(warp, thread_sumval, cg::plus<float>{});    if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; }    __syncthreads();    // same strategy, now reduce sumval across warps    warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f;    float block_sumval = cg::reduce(warp, warp_sumval, cg::plus<float>{});    // return the softmax parameters    return SoftmaxParams{1.f / block_sumval, block_maxval};} // same as 2 but not using float4__global__ void fused_classifier_kernel3(float* dlogits, float* losses, float* probs,                                         const float* logits, const float* dlosses, const int* targets,                                         int B, int T, int V, int P) {    namespace cg = cooperative_groups;    cg::thread_block block = cg::this_thread_block();    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);    int64_t idx = blockIdx.x;    int ix = targets[idx];     // softmax (reading B * T * V, same logits read again below, hopefully still in cache)    SoftmaxParams sp = prepare_softmax_blockwide_nofloat4(warp, idx, logits, V, P);     // calculate the probability needed for the loss and update (single-threaded)    if(threadIdx.x == 0) {        float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale;        losses[idx] = -logf(prob);    }     // very sensible default for dlosses is 1/(B*T), which is the uniform loss    float dloss = dlosses != NULL ? dlosses[idx] : 1.0f / (B*T);    // calculate the gradients directly, saves bandwidth from probs during training    // but also supports writing probs for inference-only and debugging    const float* logits_vec = logits + idx * P;    for (int i = threadIdx.x; i < V; i += blockDim.x) {        // this is the 2nd read of logits after the one in prepare_softmax2        // this data will never be needed again, so we reduce cache persistence        float v = __ldcs(&logits_vec[i]);        float prob = expf(v - sp.Offset) * sp.Scale;        if (probs != NULL) {            probs[idx * P + i] = prob;        }        if (dlogits != NULL) {            float indicator = (i == ix) ? 1.0f : 0.0f;            dlogits[idx * P + i] = (prob - indicator) * dloss;        }    }} __device__ SoftmaxParams prepare_softmax_blockwide2(int64_t idx, const floatX* inp, int V, int P) {    // one row of inp, i.e. inp[idx, :] of shape (V,)     const floatX* x = inp + idx * P;    float thread_maxval = -INFINITY;    float thread_sumval = 0.0f;    // do the loop in reverse to maximise probability of L2 cache hits    // so even small L2s get some hits on the 2nd read of the same thread    for (int i = ceil_div(V, x128::size) + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) {        x128 packed_x = load128cs(x + i * x128::size); // load and do not keep in cache        for(int k = 0; k < packed_x.size; ++k) {            if (i*x128::size+k >= V) {  // bounds checking against real V                continue;            }            float v = (float)packed_x[k];            float old_maxval = thread_maxval;            thread_maxval = fmaxf(thread_maxval, v);            thread_sumval *= expf(old_maxval - thread_maxval);            thread_sumval += expf(v - thread_maxval);        }    }    // two reductions of up to 1024 threads:    // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle)    // this results in much cleaner assembly than a multi-warp cg::reduce    __shared__ float shared_maxval[32];    __shared__ float shared_sumval[32];    int num_warps = blockDim.x / 32;    int warp_id = threadIdx.x / 32;    int lane_id = threadIdx.x % 32;     // reduce maxval within each warp    float warp_maxval = warpReduceMax(thread_maxval);    // thread 0 in each warp writes to shared memory    if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; }    __syncthreads();    // each thread now loads the maxval across previous warps    // if the thread is "out of range" of data, use -FLT_MAX as the maxval    warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX;    // now reduce the maxval among the warp threads    float block_maxval = warpReduceMax(warp_maxval);    // each thread uses maxval to scale sumval to avoid numerical instability / overflow    thread_sumval *= expf(thread_maxval - block_maxval);    // (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory    float warp_sumval = warpReduceSum(thread_sumval); //cg::reduce(warp, thread_sumval, cg::plus<float>{});     if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; }    __syncthreads();    // same strategy, now reduce sumval across warps    warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f;    float block_sumval = warpReduceSum(warp_sumval); //cg::reduce(warp, thread_sumval, cg::plus<float>{});    // return the softmax parameters    return SoftmaxParams{1.f / block_sumval, block_maxval};} // same as 2 but using x128__global__ void fused_classifier_kernel4(floatX* dlogits, floatX* losses, floatX* probs,                                         const floatX* logits, const floatX* dlosses, const int* targets,                                         int B, int T, int V, int P) {    int64_t idx = blockIdx.x;    int ix = targets[idx];     // softmax (reading B * T * V, same logits read again below, hopefully still in cache)    SoftmaxParams sp = prepare_softmax_blockwide2(idx, logits, V, P);     // calculate the probability needed for the loss and update (single-threaded)    if(threadIdx.x == 0) {        float prob = expf((float)logits[idx * P + ix] - sp.Offset) * sp.Scale;        losses[idx] = -logf(prob);    }     // very sensible default for dlosses is 1/(B*T), which is the uniform loss    float dloss = dlosses != NULL ? (float)dlosses[idx] : 1.0f / (B*T);    // calculate the gradients directly, saves bandwidth from probs during training    // but also supports writing probs for inference-only and debugging    const floatX* logits_vec = logits + idx * P;    for (int i = threadIdx.x; i < ceil_div(V , x128::size); i += blockDim.x) {        // this is the 2nd read of logits after the one in prepare_softmax2        // this data will never be needed again, so we reduce cache persistence        x128 packed_logits_vec = load128cs(logits_vec + i * x128::size); // load and do not keep in cache        x128 packed_probs;        x128 packed_dlogits;        for(int k = 0; k < packed_logits_vec.size; ++k) {            int element = i*packed_logits_vec.size + k;            if (element >= V) {  // bounds checking against real V                continue;            }            float v = packed_logits_vec[k];            float prob = expf(v - sp.Offset) * sp.Scale;            packed_probs[k] = prob;            float indicator = (element == ix) ? 1.0f : 0.0f;            packed_dlogits[k] = (prob - indicator) * dloss;        }        // Note: missing .cs hint hurts our performance due to cache thrashing, fixed in kernel5        store128(dlogits + idx * P + i * packed_logits_vec.size, packed_dlogits);        if (probs != NULL) {            store128(probs + idx * P + i * packed_logits_vec.size, packed_probs);        }    }} // todo - move to common.h - or ideally somewhere it's not duplicated between train & common?// requires all 32 threads in the warp to be active, but should work for any block size// uses non-dynamic shared memory so every call increases shared memory requirements by 128 bytes// the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end// but if called inside a loop, the shared memory will be implicitly reused, so set final_sync to 1using reduction_func_t = float (*) (float);template<reduction_func_t warp_reduction>__device__ float blockReduce(float val, bool final_sync=false, float out_of_bounds=0.0f) {    // two reductions of up to 1024 threads:    // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle)    __shared__ float shared_val[32];    const int lane_id = threadIdx.x % 32;    const int warp_id = threadIdx.x / 32;    const int num_warps = blockDim.x / 32;     float warp_val = warp_reduction(val);    if (lane_id == 0) { shared_val[warp_id] = warp_val; }

Lines 1–500 of 797. Load more