Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/cumsum.cu (CUDA)

Code before:

template<typename T, int BLOCK_SIZE>
static __global__ void cumsum_cub_kernel(
    const T* __restrict__ src,
    T* __restrict__ dst,
    const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
    const int64_t nb01, const int64_t nb02, const int64_t nb03,
    const int64_t nb1,  const int64_t nb2,  const int64_t nb3) {
#ifdef GGML_CUDA_USE_CUB
    using BlockScan = cub::BlockScan<T, BLOCK_SIZE>;

    __shared__ typename BlockScan::TempStorage temp_storage;
    __shared__ T block_carry;      // carry from previous tile

    const int tid = threadIdx.x;

    const int64_t i1 = blockIdx.x;
    const int64_t i2 = blockIdx.y;
    const int64_t i3 = blockIdx.z;

    if (i1 >= ne01 || i2 >= ne02 || i3 >= ne03) {
        return;
    }

    const T* src_row = src + i1 * nb01 + i2 * nb02 + i3 * nb03;
    T*       dst_row = dst + i1 * nb1  + i2 * nb2  + i3 * nb3;

    if (tid == 0) {
        block_carry = 0;
    }
    __syncthreads();

    for (int64_t start = 0; start < ne00; start += BLOCK_SIZE) {
        int64_t idx = start + tid;
        T x = (idx < ne00) ? src_row[idx] : T(0);

        T inclusive;
        T block_total;
        BlockScan(temp_storage).InclusiveSum(x, inclusive, block_total);

        __syncthreads();

        T final_val = inclusive + block_carry;

        // store result
        if (idx < ne00) {
            dst_row[idx] = final_val;
        }

        __syncthreads();


Review comment (suggestion):
```suggestion
    const T * src_row = src + i1 * s01 + i2 * s02 + i3 * s03;
    T *       dst_row = dst + i1 * s1  + i2 * s2  + i3 * s3;
```

Don't use the variable names for the strides in bytes for the bytes in elements, that is going to be confusing. These are the names that I have used for other kernels.

Code after:

template<typename T, int BLOCK_SIZE>
static __global__ void cumsum_cub_kernel(
        const T * __restrict__ src,
        T * __restrict__ dst,
        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
        const int64_t  s01, const int64_t  s02, const int64_t  s03,
        const int64_t   s1,  const int64_t   s2,  const int64_t   s3) {
#ifdef GGML_CUDA_USE_CUB
    using BlockScan = cub::BlockScan<T, BLOCK_SIZE>;

    __shared__ typename BlockScan::TempStorage temp_storage;
    __shared__ T block_carry;      // carry from previous tile

    const int tid = threadIdx.x;

    const int64_t i1 = blockIdx.x;
    const int64_t i2 = blockIdx.y;
    const int64_t i3 = blockIdx.z;

    if (i1 >= ne01 || i2 >= ne02 || i3 >= ne03) {
        return;
    }

    const T * src_row = src + i1 * s01 + i2 * s02 + i3 * s03;
    T *       dst_row = dst + i1 * s1  + i2 * s2  + i3 * s3;

    if (tid == 0) {
        block_carry = 0;
    }
    __syncthreads();

    for (int64_t start = 0; start < ne00; start += BLOCK_SIZE) {
        int64_t idx = start + tid;
        T x = (idx < ne00) ? src_row[idx] : T(0);

        T inclusive;
        T block_total;
        BlockScan(temp_storage).InclusiveSum(x, inclusive, block_total);

        __syncthreads();

        T final_val = inclusive + block_carry;

        // store result
        if (idx < ne00) {
            dst_row[idx] = final_val;
        }

        __syncthreads();