Repository: ggml-org/llama.cpp
Pull request: ggml : optimize cuda cumsum fallback (~2.5x speedup vs CUB)

### Description
While analyzing the `cumsum` kernel (critical for sampling in newer architectures like Qwen-2/3 and SSMs), I noticed the fallback implementation was a standard naive block scan with heavy synchronization overhead.

This PR implements **Register Blocking** (Thread Coarsening) with `unroll = 4`. By processing 4 elements per thread in registers before hitting shared memory, we significantly reduce the ratio of `__syncthreads()` calls to data processed.

### Benchmarks
I do not have the VRAM to run the full Qwen-3 next (50GB) model end-to-end, so I wrote a standalone benchmark to test the kernel latency directly on an **RTX 4050 Laptop GPU**.

**Test Case:** N = 152,000 (Approx Qwen Vocab Size)

| Implementation | Time (ms) | Relative Speed |
| :--- | :--- | :--- |
| Old Fallback (Naive) | ~1.018 ms | 1.0x |
| Current CUB Wrapper | ~1.00 ms | 1.018x |
| **New Fallback (This PR)** | **0.40 ms** | **2.54x** |

*(Note: The new kernel is ~2.5x faster than the CUB library wrapper itself due to better instruction-level parallelism and reduced barrier stalls.)*

### Motivation & Next Steps
Currently, `ggml` prefers CUB if available. However, given that this custom kernel outperforms the CUB wrapper significantly on my hardware, it might be worth considering:
1.  Merging this as the optimized fallback (safe option, what this PR does).
2.  Eventually replacing the CUB path entirely with this kernel to reduce dependencies and improve performance across the board.

I would appreciate if a maintainer with access to larger hardware/models could verify the end-to-end impact on token generation speed for `cumsum`-heavy models.

---
### Proposal for CUB Removal
Since this kernel appears to outperform the CUB wrapper significantly, **I am happy to update this PR to remove the CUB dependency entirely** if the maintainers agree.

This would:
1. Simplify `cumsum.cu` (removing the CUB include and wrapper logic).
2. Reduce build times/dependencies.
3. Ensure the faster path is used by default on all CUDA devices.

Let me know if you would like me to proceed with that!

Changed files:
filename: ggml/src/ggml-cuda/cumsum.cu
patch: @@ -69,7 +69,7 @@ static __global__ void cumsum_cub_kernel(
 #endif // GGML_CUDA_USE_CUB
 }
 
-// Fallback kernel implementation (original)
+// Fallback kernel implementation
 template<typename T>
 static __global__ void cumsum_kernel(
         const T * src, T * dst,
@@ -86,10 +86,10 @@ static __global__ void cumsum_kernel(
     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;
+    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) {
@@ -107,21 +107,39 @@ static __global__ void cumsum_kernel(
     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;
+    // 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;
 
-        // 1. Warp inclusive scan
+        // Warp inclusive scan
         val = warp_prefix_inclusive_sum<T, warp_size>(val);
         s_vals[tid] = val;
 
-        // Store warp total
         if (lane == warp_size - 1) {
             s_warp_sums[warp] = val;
         }
         __syncthreads();
 
-        // 2. Exclusive scan of warp sums (warp 0 only)
+        // 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);
@@ -134,12 +152,17 @@ static __global__ void cumsum_kernel(
         }
         __syncthreads();
 
+        // write back results
         float carry = *s_carry;
-        float final_val = s_vals[tid] + s_warp_sums[warp] + carry;
-        if (idx < ne00) {
-            dst_row[idx] = ggml_cuda_cast<T, float>(final_val);
+        // calculate sum offset for this thread
+        float final_val_offset = s_vals[tid] + s_warp_sums[warp] + carry - temp[num_unroll - 1];
+
+#pragma unroll
+        for (int32_t j = 0; j < num_unroll; j++) {
+            if (idx + j < ne00) {
+                dst_row[idx + j] = temp[j] + ggml_cuda_cast<T, float>(final_val_offset);
+            }
         }
-        __syncthreads();
 
         // Update carry for next chunk
         if (tid == 0) {