Repository: ggml-org/llama.cpp
Pull request: cpu: introduce chunking for repack matmuls and enable matmul-id chunking

Similar to #16829 and tested in tandem.

A very simple dynamic chunking mechanism for repack matmuls. Helps on platforms with significant performance difference between the CPU cores, and helps distribute the work better under load in general. 
I tested on M4 Pro and a few Snapdragons but it should work on all platforms.

See the details below. 
I included a trace with instrumented matmuls that shows how threads threads endup processing chunks.

<details>

```
## M4 Pro

Before (no other load)
| model                  |       size |     params | backend    | ngl | threads | fa | dev   |   test |            t/s |
| ---------------------- | ---------: | ---------: | ---------- | --: | ------: | -: | ----- | -----: |  ------------: |
| gpt-oss 20B MXFP4 MoE  |  11.27 GiB |    20.91 B | Metal      |   0 |       6 |  1 | none  |  pp256 |   75.67 ± 0.43 |
| gpt-oss 20B MXFP4 MoE  |  11.27 GiB |    20.91 B | Metal      |   0 |       6 |  1 | none  |   tg64 |   56.13 ± 0.26 |

| model                  |       size |     params | backend    | ngl | threads | fa | dev   |   test |            t/s |
| ---------------------- | ---------: | ---------: | ---------- | --: | ------: | -: | ----- | -----: | -------------: |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       2 |  1 | none  |  pp256 |  100.81 ± 2.22 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       2 |  1 | none  |   tg64 |   37.27 ± 1.06 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       4 |  1 | none  |  pp256 |  198.34 ± 0.21 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       4 |  1 | none  |   tg64 |   67.88 ± 0.63 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       6 |  1 | none  |  pp256 |  275.03 ± 8.60 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       6 |  1 | none  |   tg64 |   92.09 ± 1.40 |

After (no other load)
| model                  |       size |     params | backend    | ngl | threads | fa | dev   |   test |            t/s |
| ---------------------- | ---------: | ---------: | ---------- | --: | ------: | -: | ----- | -----: |  ------------: |
| gpt-oss 20B MXFP4 MoE  |  11.27 GiB |    20.91 B | Metal      |   0 |       6 |  1 | none  |  pp256 |   76.57 ± 0.15 |
| gpt-oss 20B MXFP4 MoE  |  11.27 GiB |    20.91 B | Metal      |   0 |       6 |  1 | none  |   tg64 |   55.66 ± 0.46 |

| model                  |       size |     params | backend    | ngl | threads | fa | dev   |   test |            t/s |
| ---------------------  | ---------: | ---------: | ---------- | --: | ------: | -: | ----- | -----: | -------------: |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       2 |  1 | none  |  pp256 |  105.01 ± 0.33 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       2 |  1 | none  |   tg64 |   38.63 ± 0.10 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       4 |  1 | none  |  pp256 |  198.66 ± 0.19 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       4 |  1 | none  |   tg64 |   67.40 ± 0.29 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       6 |  1 | none  |  pp256 |  290.01 ± 1.31 |
| llama 3B Q4_0          |   1.78 GiB |     3.21 B | Metal      |   0 |       6 |  1 | none  |   tg64 |   89.92 ± 0.09 |


Chunking in action (no load)
thread-5: matmul ffn_up-0 nchunks 4 usec 7219
thread-3: matmul ffn_up-0 nchunks 4 usec 7221
thread-2: matmul ffn_up-0 nchunks 4 usec 7232
thread-1: matmul ffn_up-0 nchunks 4 usec 7247
thread-0: matmul ffn_up-0 nchunks 4 usec 7259
thread-4: matmul ffn_up-0 nchunks 4 usec 7260
thread-3: matmul ffn_out-0 nchunks 4 usec 7402
thread-1: matmul ffn_out-0 nchunks 4 usec 7423
thread-2: matmul ffn_out-0 nchunks 4 usec 7425
thread-4: matmul ffn_out-0 nchunks 4 usec 7402
thread-0: matmul ffn_out-0 nchunks 4 usec 7411
thread-5: matmul ffn_out-0 nchunks 4 usec 7402

Chunking in action (heavy other load)
thread-3: matmul ffn_up-6 nchunks 3 usec 8080
thread-1: matmul ffn_up-6 nchunks 5 usec 9055
thread-4: matmul ffn_up-6 nchunks 5 usec 9070
thread-5: matmul ffn_up-6 nchunks 5 usec 9428
thread-2: matmul ffn_up-6 nchunks 3 usec 9502
thread-0: matmul ffn_up-6 nchunks 3 usec 10552
thread-3: matmul ffn_out-6 nchunks 4 usec 8556
thread-0: matmul ffn_out-6 nchunks 3 usec 8612
thread-4: matmul ffn_out-6 nchunks 4 usec 8809
thread-1: matmul ffn_out-6 nchunks 5 usec 9275
thread-5: matmul ffn_out-6 nchunks 5 usec 9750
thread-2: matmul ffn_out-6 nchunks 3 usec 9963


## Snapdragon 8-Elite Gen5

## LLama3.2 1B Q4_0
  llama_model_loader: - type  f32:   34 tensors
  llama_model_loader: - type q4_0:  112 tensors
  llama_model_loader: - type q6_K:    1 tensors

Before
| model          |       size |     params | backend    | ngl | threads | fa | dev   |   test |            t/s |
| -------------- | ---------: | ---------: | ---------- | --: | ------: | -: | ----- | -----: | -------------: |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       2 |  1 | none  |  pp128 |  384.94 ± 9.15 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       2 |  1 | none  |   tg64 |   65.17 ± 1.49 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       4 |  1 | none  |  pp128 |  351.52 ± 0.28 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       4 |  1 | none  |   tg64 |   71.00 ± 1.49 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       6 |  1 | none  |  pp128 |  512.93 ± 1.78 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       6 |  1 | none  |   tg64 |   77.26 ± 1.29 |


After
| model          |       size |     params | backend    | ngl | threads | fa | dev   |    test |            t/s |
| -------------- | ---------: | ---------: | ---------- | --: | ------: | -: | ----- | ------: |--------------: |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       2 |  1 | none  |   pp128 |  395.65 ± 7.81 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       2 |  1 | none  |    tg64 |   64.40 ± 0.85 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       4 |  1 | none  |   pp128 |  459.51 ± 1.04 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       4 |  1 | none  |    tg64 |   73.62 ± 0.67 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       6 |  1 | none  |   pp128 |  669.03 ± 1.41 |
| llama 1B Q4_0  | 727.75 MiB |     1.24 B | CPU        |   0 |       6 |  1 | none  |    tg64 |   79.75 ± 0.56 |


## Llama3.2 3B Q4_0
  llama_model_loader: - type  f32:   58 tensors
  llama_model_loader: - type q4_0:  196 tensors
  llama_model_loader: - type q6_K:    1 tensors

Before
| model           |       size |     params | backend    | ngl | threads | fa | dev   |  test |             t/s |
| --------------- | ---------: | ---------: | ---------- | --: | ------: | -: | ----- | ----: | --------------: |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       2 |  1 | none  | pp128 |   127.73 ± 2.43 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       2 |  1 | none  |  tg64 |    27.91 ± 0.61 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       4 |  1 | none  | pp128 |   122.97 ± 0.02 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       4 |  1 | none  |  tg64 |    29.72 ± 1.09 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       6 |  1 | none  | pp128 |  159.59 ± 14.06 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       6 |  1 | none  |  tg64 |    30.33 ± 0.60 |


After
| model           |       size |     params | backend    | ngl | threads | fa | dev   |  test |             t/s |
| --------------- | ---------: | ---------: | ---------- | --: | ------: | -: | ----- | ----: | --------------: |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       2 |  1 | none  | pp128 |   128.16 ± 2.09 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       2 |  1 | none  |  tg64 |    27.46 ± 0.47 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       4 |  1 | none  | pp128 |   161.89 ± 0.30 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       4 |  1 | none  |  tg64 |    30.07 ± 0.65 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       6 |  1 | none  | pp128 |   227.02 ± 7.26 |
| llama 3B Q4_0   |   1.78 GiB |     3.21 B | CPU        |   0 |       6 |  1 | none  |  tg64 |    32.16 ± 0.58 |


## Llama3.2 1B chunking in action

thread-0: matmul ffn_up-11 nchunks 7 usec 143
thread-5: matmul ffn_up-11 nchunks 8 usec 147
thread-3: matmul ffn_up-11 nchunks 2 usec 150
thread-1: matmul ffn_up-11 nchunks 2 usec 150
thread-2: matmul ffn_up-11 nchunks 2 usec 152
thread-4: matmul ffn_up-11 nchunks 3 usec 158
thread-0: matmul ffn_out-11 nchunks 7 usec 124
thread-1: matmul ffn_out-11 nchunks 2 usec 125
thread-5: matmul ffn_out-11 nchunks 8 usec 128
thread-4: matmul ffn_out-11 nchunks 2 usec 129
thread-2: matmul ffn_out-11 nchunks 2 usec 139
thread-3: matmul ffn_out-11 nchunks 3 usec 150


## Galaxy S25+ (Snapdragon 8-Elite Gen4)

## LLama3.2 1B chunking in action

thread-2: matmul ffn_up-11 nchunks 6 usec 147
thread-3: matmul ffn_up-11 nchunks 3 usec 150
thread-0: matmul ffn_up-11 nchunks 6 usec 147
thread-5: matmul ffn_up-11 nchunks 3 usec 150
thread-1: matmul ffn_up-11 nchunks 3 usec 147
thread-4: matmul ffn_up-11 nchunks 3 usec 152
thread-4: matmul ffn_out-11 nchunks 3 usec 136
thread-2: matmul ffn_out-11 nchunks 6 usec 142
thread-5: matmul ffn_out-11 nchunks 3 usec 146
thread-1: matmul ffn_out-11 nchunks 3 usec 136
thread-0: matmul ffn_out-11 nchunks 6 usec 144
thread-3: matmul ffn_out-11 nchunks 3 usec 136

```

</details>




Changed files:
filename: ggml/src/ggml-cpu/ggml-cpu.c
patch: @@ -1613,13 +1613,8 @@ static void ggml_compute_forward_mul_mat_id(
             chunk_size = 64;
         }
 
-#if defined(__aarch64__)
-        // disable for ARM
-        const bool disable_chunking = true;
-#else
         // disable for NUMA
         const bool disable_chunking = ggml_is_numa();
-#endif // defined(__aarch64__)
 
         int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size;
         int64_t nchunk1 = (nr1 + chunk_size - 1) / chunk_size;

filename: ggml/src/ggml-cpu/repack.cpp
patch: @@ -1600,6 +1600,32 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
         return false;
     }
 
+    void forward_mul_mat_one_chunk(ggml_compute_params * params, ggml_tensor * op, int64_t src0_start, int64_t src0_end) {
+        const ggml_tensor * src0 = op->src[0];
+        const ggml_tensor * src1 = op->src[1];
+        ggml_tensor *       dst  = op;
+
+        GGML_TENSOR_BINARY_OP_LOCALS
+
+        const void * src1_wdata      = params->wdata;
+        const size_t src1_col_stride = ggml_row_size(PARAM_TYPE, ne10);
+
+        // If there are more than three rows in src1, use gemm; otherwise, use gemv.
+        if (ne11 > 3) {
+            gemm<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(ne00,
+                    (float *) ((char *) dst->data) + src0_start, ne01,
+                    (const char *) src0->data + src0_start * nb01,
+                    (const char *) src1_wdata, ne11 - ne11 % 4, src0_end - src0_start);
+        }
+        for (int iter = ne11 - ne11 % 4; iter < ne11; iter++) {
+            gemv<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(ne00,
+                    (float *) ((char *) dst->data + (iter * nb1)) + src0_start, ne01,
+                    (const char *) src0->data + src0_start * nb01,
+                    (const char *) src1_wdata + (src1_col_stride * iter), 1,
+                    src0_end - src0_start);
+        }
+    }
+
     void forward_mul_mat(ggml_compute_params * params, ggml_tensor * op) {
         const ggml_tensor * src0 = op->src[0];
         const ggml_tensor * src1 = op->src[1];
@@ -1643,31 +1669,41 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS, ggml_type PAR
             from_float((float *) ((char *) src1->data + i11 * nb11), (void *) (wdata + i11 * nbw1), ne10);
         }
 
-        ggml_barrier(params->threadpool);
+        // disable for NUMA
+        const bool disable_chunking = ggml_is_numa();
 
-        const void * src1_wdata      = params->wdata;
-        const size_t src1_col_stride = ggml_row_size(PARAM_TYPE, ne10);
-        int64_t      src0_start      = (ith * ne01) / nth;
-        int64_t      src0_end        = ((ith + 1) * ne01) / nth;
-        src0_start = (src0_start % NB_COLS) ? src0_start + NB_COLS - (src0_start % NB_COLS) : src0_start;
-        src0_end   = (src0_end   % NB_COLS) ? src0_end   + NB_COLS - (src0_end   % NB_COLS) : src0_end;
-        if (src0_start >= src0_end) {
-            return;
+        // 4x chunks per thread
+        int64_t nr = ggml_nrows(op->src[0]);
+        int nth_scaled = nth * 4;
+        int64_t chunk_size = (nr + nth_scaled - 1) / nth_scaled;
+        int64_t nchunk     = (nr + chunk_size - 1) / chunk_size;
+
+        if (nth == 1 || nchunk < nth || disable_chunking) {
+            nchunk = nth;
         }
 
-        // If there are more than three rows in src1, use gemm; otherwise, use gemv.
-        if (ne11 > 3) {
-            gemm<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(ne00,
-                    (float *) ((char *) dst->data) + src0_start, ne01,
-                    (const char *) src0->data + src0_start * nb01,
-                    (const char *) src1_wdata, ne11 - ne11 % 4, src0_end - src0_start);
+        if (ith == 0) {
+            // Every thread starts at ith, so the first unprocessed chunk is nth.  This save a bit of coordination right at the start.
+            ggml_threadpool_chunk_set(params->threadpool, nth);
         }
-        for (int iter = ne11 - ne11 % 4; iter < ne11; iter++) {
-            gemv<BLOC_TYPE, INTER_SIZE, NB_COLS, PARAM_TYPE>(ne00,
-                    (float *) ((char *) dst->data + (iter * nb1)) + src0_start, ne01,
-                    (const char *) src0->data + src0_start * nb01,
-                    (const char *) src1_wdata + (src1_col_stride * iter), 1,
-                    src0_end - src0_start);
+
+        ggml_barrier(params->threadpool);
+
+        // The first chunk comes from our thread_id, the rest will get auto-assigned.
+        int current_chunk = ith;
+
+        while (current_chunk < nchunk) {
+            int64_t src0_start = (current_chunk * ne01) / nchunk;
+            int64_t src0_end   = ((current_chunk + 1) * ne01) / nchunk;
+            src0_start = (src0_start % NB_COLS) ? src0_start + NB_COLS - (src0_start % NB_COLS) : src0_start;
+            src0_end   = (src0_end   % NB_COLS) ? src0_end   + NB_COLS - (src0_end   % NB_COLS) : src0_end;
+            if (src0_start >= src0_end) {
+                break;
+            }
+
+            forward_mul_mat_one_chunk(params, dst, src0_start, src0_end);
+
+            current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1);
         }
     }