Repository: ggml-org/llama.cpp
Pull request: cpu: skip redudant ROPE cache updates

I was updating ROPE for Hexagon while using the CPU implementation as the reference and noticed that some threads may end up doing redundant ROPE cache updates for the rows that are mapped to other threads.
Added a bit of instrumentation and sure enough I see something like this:

<details>

```
thread-0: compute rope cache: ir 18 i3 0 i2 1
thread-0: compute rope cache: ir 19 i3 0 i2 2
thread-0: compute rope cache: ir 20 i3 0 i2 3
thread-0: compute rope cache: ir 21 i3 0 i2 4
thread-0: compute rope cache: ir 22 i3 0 i2 5
thread-0: compute rope cache: ir 23 i3 0 i2 6
thread-0: compute rope cache: ir 24 i3 0 i2 7
thread-0: compute rope cache: ir 25 i3 0 i2 8
thread-0: compute rope cache: ir 26 i3 0 i2 9
thread-0: compute rope cache: ir 27 i3 0 i2 10
thread-0: compute rope cache: ir 0 i3 0 i2 0
thread-0: use rope cache: ir 1 i3 0 i2 0 i1 0
thread-0: use rope cache: ir 2 i3 0 i2 0 i1 1
thread-0: use rope cache: ir 3 i3 0 i2 0 i1 2
thread-0: use rope cache: ir 4 i3 0 i2 0 i1 3
thread-0: use rope cache: ir 5 i3 0 i2 0 i1 4
thread-0: use rope cache: ir 6 i3 0 i2 0 i1 5
thread-0: compute rope cache: ir 7 i3 0 i2 1
thread-0: compute rope cache: ir 8 i3 0 i2 2
thread-0: compute rope cache: ir 9 i3 0 i2 3
thread-0: compute rope cache: ir 10 i3 0 i2 4
thread-0: compute rope cache: ir 11 i3 0 i2 5
thread-0: compute rope cache: ir 12 i3 0 i2 6
thread-0: compute rope cache: ir 13 i3 0 i2 7
thread-0: compute rope cache: ir 14 i3 0 i2 8
thread-0: compute rope cache: ir 15 i3 0 i2 9
thread-0: compute rope cache: ir 16 i3 0 i2 10
thread-0: compute rope cache: ir 0 i3 0 i2 0
thread-0: use rope cache: ir 1 i3 0 i2 0 i1 0
thread-0: use rope cache: ir 2 i3 0 i2 0 i1 1
thread-0: use rope cache: ir 3 i3 0 i2 0 i1 2
thread-0: use rope cache: ir 4 i3 0 i2 0 i1 3
thread-0: use rope cache: ir 5 i3 0 i2 0 i1 4
thread-0: use rope cache: ir 6 i3 0 i2 0 i1 5
thread-0: use rope cache: ir 7 i3 0 i2 0 i1 6
thread-0: use rope cache: ir 8 i3 0 i2 0 i1 7
thread-0: use rope cache: ir 9 i3 0 i2 0 i1 8
...
thread-1: compute rope cache: ir 8 i3 0 i2 1
thread-1: use rope cache: ir 9 i3 0 i2 1 i1 0
thread-1: use rope cache: ir 10 i3 0 i2 1 i1 1
thread-1: use rope cache: ir 11 i3 0 i2 1 i1 2
thread-1: use rope cache: ir 12 i3 0 i2 1 i1 3
thread-1: compute rope cache: ir 13 i3 0 i2 2
thread-1: compute rope cache: ir 14 i3 0 i2 3
thread-1: compute rope cache: ir 15 i3 0 i2 4
thread-1: compute rope cache: ir 16 i3 0 i2 5
thread-1: compute rope cache: ir 17 i3 0 i2 6
thread-1: compute rope cache: ir 18 i3 0 i2 7
thread-1: compute rope cache: ir 19 i3 0 i2 8
thread-1: compute rope cache: ir 20 i3 0 i2 9
thread-1: compute rope cache: ir 21 i3 0 i2 10
thread-1: compute rope cache: ir 0 i3 0 i2 0
thread-1: use rope cache: ir 18 i3 0 i2 0 i1 17
thread-1: use rope cache: ir 19 i3 0 i2 0 i1 18
thread-1: use rope cache: ir 20 i3 0 i2 0 i1 19
thread-1: use rope cache: ir 21 i3 0 i2 0 i1 20
thread-1: use rope cache: ir 22 i3 0 i2 0 i1 21
thread-1: use rope cache: ir 23 i3 0 i2 0 i1 22
thread-1: use rope cache: ir 24 i3 0 i2 0 i1 23

``` 

</details>

ROPE cache is updated for each i2 (seq-len) iteration but some of the iterations are mapped to other threads.
This PR adds a bit of state to compute ROPE cache only for the iterations that are actually used.

Here are some before/after numbers with llama3.2-3B-Q4_0.

```
## Ryzen AI Max+ 395
Before:
  common_perf_print: prompt eval time =     353.17 ms /   214 tokens (    1.65 ms per token,   605.93 tokens per second)
  common_perf_print:        eval time =    1631.31 ms /    83 runs   (   19.65 ms per token,    50.88 tokens per second)
After:
  common_perf_print: prompt eval time =     350.50 ms /   214 tokens (    1.64 ms per token,   610.55 tokens per second)
  common_perf_print:        eval time =    1640.84 ms /    83 runs   (   19.77 ms per token,    50.58 tokens per second)

## Snapdragon-Gen5
Before:
  common_perf_print: prompt eval time =    1091.27 ms /   205 tokens (    5.32 ms per token,   187.85 tokens per second)
  common_perf_print:        eval time =    2056.67 ms /    63 runs   (   32.65 ms per token,    30.63 tokens per second)
After:
  common_perf_print: prompt eval time =    1033.26 ms /   205 tokens (    5.04 ms per token,   198.40 tokens per second)
  common_perf_print:        eval time =    2078.37 ms /    63 runs   (   32.99 ms per token,    30.31 tokens per second)
```

ROPE cache compute is kind of expensive (lots of divs/muls/sin/cos/...) it makes sense to skip redundant updates even if the token rate bump is not huge.  

Changed files:
filename: ggml/src/ggml-cpu/ops.cpp
patch: @@ -5803,28 +5803,33 @@ static void ggml_compute_forward_rope_flt(
 
     const int32_t * pos = (const int32_t *) src1->data;
 
+    int64_t last_i2 = -1;
+
     for (int64_t i3 = 0; i3 < ne3; i3++) { // batch
         for (int64_t i2 = 0; i2 < ne2; i2++) { // seq-len
-
-            float * cache = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32)*ith;
-            if (!mrope_used) {
-                const int64_t p = pos[i2];
-                ggml_rope_cache_init(p, freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale);
-            }
-            else {
-                const int64_t p_t = pos[i2];
-                const int64_t p_h = pos[i2 + ne2];
-                const int64_t p_w = pos[i2 + ne2 * 2];
-                const int64_t p_e = pos[i2 + ne2 * 3];
-                ggml_mrope_cache_init(
-                    p_t, p_h, p_w, p_e, sections, is_imrope, is_vision,
-                    freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale);
-            }
-
             for (int64_t i1 = 0; i1 < ne1; i1++) { // attn-heads
-                if (ir++ < ir0) continue;
+                if (ir++ < ir0) continue; // skip rows mapped to other threads
                 if (ir   > ir1) break;
 
+                float * cache = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32)*ith;
+                if (last_i2 != i2) {
+                    if (!mrope_used) {
+                        const int64_t p = pos[i2];
+                        ggml_rope_cache_init(p, freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale);
+                    }
+                    else {
+                        const int64_t p_t = pos[i2];
+                        const int64_t p_h = pos[i2 + ne2];
+                        const int64_t p_w = pos[i2 + ne2 * 2];
+                        const int64_t p_e = pos[i2 + ne2 * 3];
+                        ggml_mrope_cache_init(
+                            p_t, p_h, p_w, p_e, sections, is_imrope, is_vision,
+                            freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale);
+                    }
+
+                    last_i2 = i2;
+                }
+
                 T * src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01);
                 T * dst_data  = (T *)((char *)  dst->data + i3*nb3  + i2*nb2  + i1*nb1);