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

Code before:
    }
#else
    NO_DEVICE_CODE;
#endif
}

// Fallback kernel implementation (original)
template<typename T>
static __global__ void cumsum_kernel(
    const T * src, T * dst,
    const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
    const int64_t nb00, const int64_t nb01, const int64_t nb02, const int64_t nb03,
    const int64_t nb0,  const int64_t nb1,  const int64_t nb2,  const int64_t nb3) {

    GGML_UNUSED_VARS(nb00, nb0);

    const int tid = threadIdx.x;
    const int lane = tid & (WARP_SIZE - 1);
    const int warp = tid / WARP_SIZE;
    const int warps_per_block = blockDim.x / WARP_SIZE;

    extern __shared__ float smem[];
    float* s_vals = smem;
    float* s_warp_sums = smem + blockDim.x;
    float* s_carry = smem + blockDim.x + warps_per_block;
    float* s_chunk_total = s_carry + 1;

    // Initialize carry
    if (tid == 0) {
        *s_carry = 0.0f;
    }
    __syncthreads();

    const int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;
    const int64_t i1 = blockIdx.x;
    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
        return;
    }

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

    for (int64_t start = 0; start < ne00; start += blockDim.x) {
        int64_t idx = start + tid;
        float val = (idx < ne00) ? ggml_cuda_cast<float, T>(src_row[idx]) : 0.0f;

        // 1. Warp inclusive scan
        val = warp_prefix_inclusive_sum(val);
        s_vals[tid] = val;


Review comment (suggestion):
```suggestion
    float * s_vals = smem;
    float * s_warp_sums = smem + blockDim.x;
    float * s_carry = smem + blockDim.x + warps_per_block;
    float * s_chunk_total = s_carry + 1;
```

Code after:
    }
#else
    NO_DEVICE_CODE;
#endif // GGML_CUDA_USE_CUB
}

// Fallback kernel implementation (original)
template<typename T>
static __global__ void cumsum_kernel(
        const T * src, T * dst,
        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
        const int64_t  s00, const int64_t  s01, const int64_t  s02, const int64_t  s03,
        const int64_t   s0, const int64_t   s1, const int64_t   s2, const int64_t   s3) {

    GGML_UNUSED_VARS(s00, s0);

    const int tid = threadIdx.x;
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    const int lane = tid % warp_size;
    const int warp = tid / warp_size;
    const int warps_per_block = blockDim.x / warp_size;

    extern __shared__ float smem[];
    float * s_vals = smem;
    float * s_warp_sums = smem + blockDim.x;
    float * s_carry = smem + blockDim.x + warps_per_block;
    float * s_chunk_total = s_carry + 1;

    // Initialize carry
    if (tid == 0) {
        *s_carry = 0.0f;
    }
    __syncthreads();

    const int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;
    const int64_t i1 = blockIdx.x;
    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
        return;
    }

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

    for (int64_t start = 0; start < ne00; start += blockDim.x) {
        int64_t idx = start + tid;
        float val = (idx < ne00) ? ggml_cuda_cast<float, T>(src_row[idx]) : 0.0f;

        // 1. Warp inclusive scan
        val = warp_prefix_inclusive_sum<T, warp_size>(val);
        s_vals[tid] = val;