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

Code before:
        if (idx < ne00) {
            dst_row[idx] = static_cast<T>(final_val);
        }
        __syncthreads();

        // Update carry for next chunk
        if (tid == 0) {
            *s_carry += *s_chunk_total;
        }
        __syncthreads();
    }
}

template<typename T>
static void cumsum_cuda(
    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,
    cudaStream_t stream) {

    const size_t type_size = sizeof(T);
    bool use_cub = false;
#ifdef GGML_CUDA_USE_CUB
    // Check if we can use CUB (data must be contiguous along innermost dimension)
    const bool is_contiguous = (nb00 == type_size) && (nb0 == type_size);

    if (is_contiguous) {
        use_cub = true;
    }
#endif // GGML_CUDA_USE_CUB
    dim3 grid_dims(ne01, ne02, ne03);
    const int num_warps = (ne00 + WARP_SIZE - 1) / WARP_SIZE;
    int block_size = num_warps * WARP_SIZE;
    block_size = std::min(block_size, CUDA_CUMSUM_BLOCK_SIZE);
    dim3 block_dims(block_size, 1, 1);
    const int warps_per_block = block_size / WARP_SIZE;
    const size_t shmem_size = (block_size + warps_per_block + 2) * sizeof(float);

    if (use_cub) {
        cumsum_cub_kernel<T, CUDA_CUMSUM_BLOCK_SIZE><<<grid_dims, CUDA_CUMSUM_BLOCK_SIZE, 0, stream>>>(
            src, dst,
            ne00, ne01, ne02, ne03,
            nb01 / type_size, nb02 / type_size, nb03 / type_size,
            nb1 / type_size,  nb2 / type_size,  nb3 / type_size
        );
    } else {
        cumsum_kernel<<<grid_dims, block_dims, shmem_size, stream>>>(
            src, dst,
            ne00, ne01, ne02, ne03,
            nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,

Review comment (suggestion):
I think you can use `ggml_is_contigous`, however that would need to be done from the calling function

Code after:
        if (idx < ne00) {
            dst_row[idx] = ggml_cuda_cast<T, float>(final_val);
        }
        __syncthreads();

        // Update carry for next chunk
        if (tid == 0) {
            *s_carry += *s_chunk_total;
        }
        __syncthreads();
    }
}

template<typename T>
static void cumsum_cuda(
        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,
        cudaStream_t stream) {

    const size_t type_size = sizeof(T);
    bool use_cub = false;
#ifdef GGML_CUDA_USE_CUB
    // Check if we can use CUB (data must be contiguous along innermost dimension)
    const bool is_contiguous = (nb00 == type_size) && (nb0 == type_size);

    if (is_contiguous) {
        use_cub = true;
    }
#endif // GGML_CUDA_USE_CUB
    dim3 grid_dims(ne01, ne02, ne03);
    const auto &info = ggml_cuda_info().devices[ggml_cuda_get_device()];
    const int warp_size = info.warp_size;
    const int num_warps = (ne00 + warp_size - 1) / warp_size;
    int block_size = num_warps * warp_size;
    block_size = std::min(block_size, CUDA_CUMSUM_BLOCK_SIZE);
    dim3 block_dims(block_size, 1, 1);
    const int warps_per_block = block_size / warp_size;
    const size_t shmem_size = (block_size + warps_per_block + 2) * sizeof(float);

    if (use_cub) {
        cumsum_cub_kernel<T, CUDA_CUMSUM_BLOCK_SIZE><<<grid_dims, CUDA_CUMSUM_BLOCK_SIZE, 0, stream>>>(
            src, dst,
            ne00, ne01, ne02, ne03,
            nb01 / type_size, nb02 / type_size, nb03 / type_size,
            nb1 / type_size,  nb2 / type_size,  nb3 / type_size
        );
    } else {
        cumsum_kernel<<<grid_dims, block_dims, shmem_size, stream>>>(
            src, dst,