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

Code before:
    }

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

    // register blocking: process 4 elements per thread to hide latency
    // and reduce synchronization overhead
    constexpr int num_unroll = 4;
    T             temp[num_unroll];

    for (int64_t i = 0; i < ne00; i += num_unroll * blockDim.x) {
        int64_t idx = i + tid * num_unroll;

        // thread local sequential scan
        temp[0] = (idx < ne00 ? src_row[idx] : T(0));
#pragma unroll
        for (int64_t j = 1; j < num_unroll; j++) {
            temp[j] = temp[j - 1];
            if (idx + j < ne00) {
                temp[j] += src_row[idx + j];
            } else {
                temp[j] += 0;
            }
        }

        // last emenent is sum of all values assigned to thread
        float val = (idx < ne00) ? ggml_cuda_cast<float, T>(temp[num_unroll - 1]) : 0.0f;

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

        if (lane == warp_size - 1) {
            s_warp_sums[warp] = val;
        }
        __syncthreads();

        // Exclusive scan of warp sums (warp 0 only)
        if (warp == 0) {
            float w = (tid < warps_per_block) ? s_warp_sums[tid] : 0.0f;
            float inc = warp_prefix_inclusive_sum<T, warp_size>(w);
            if (tid < warps_per_block) {
                s_warp_sums[tid] = inc - w;   // exclusive sum
            }
            if (tid == warps_per_block - 1) {
                *s_chunk_total = inc;          // total sum of this chunk
            }
        }
        __syncthreads();

        // write back results

Review comment (suggestion):
```suggestion
        // last element is sum of all values assigned to thread
```

Code after:
        return;
    }

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

    // register blocking: process 4 elements per thread to hide latency
    // and reduce synchronization overhead
    constexpr int num_unroll = 4;
    T             temp[num_unroll];

    for (int64_t i = 0; i < ne00; i += num_unroll * blockDim.x) {
        int64_t idx = i + tid * num_unroll;

        // thread local sequential scan
        temp[0] = (idx < ne00 ? src_row[idx] : T(0));
#pragma unroll
        for (int64_t j = 1; j < num_unroll; j++) {
            temp[j] = temp[j - 1];
            if (idx + j < ne00) {
                temp[j] += src_row[idx + j];
            } else {
                temp[j] += 0;
            }
        }

        // last emenent is sum of all values assigned to thread
        float val = (idx < ne00) ? ggml_cuda_cast<float, T>(temp[num_unroll - 1]) : 0.0f;

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

        if (lane == warp_size - 1) {
            s_warp_sums[warp] = val;
        }
        __syncthreads();

        // Exclusive scan of warp sums (warp 0 only)
        if (warp == 0) {
            float w = (tid < warps_per_block) ? s_warp_sums[tid] : 0.0f;
            float inc = warp_prefix_inclusive_sum<T, warp_size>(w);
            if (tid < warps_per_block) {
                s_warp_sums[tid] = inc - w;   // exclusive sum
            }
            if (tid == warps_per_block - 1) {
                *s_chunk_total = inc;          // total sum of this chunk
            }
        }
        __syncthreads();