Repository: ggml-org/llama.cpp
Pull request: Kimi Linear chunk size = 16

*Make sure to read the [contributing guidelines](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md) before submitting a PR*

~15% gain in pp. 

Also, CUDA0 compute buffer of IQ3_M at 192k context reduces from 1693.51MB to 991.29MB on 3090. This allows context to further increase to 272k. 

Also, tried to implement the same thing for gated delta net. However, both pp speed lost and CUDA0 compute buffer lost is observed. Anyway, I included the code for GDN if someone would like to explore whether it is possible to speed up GDN via block.

<details>
<summary>orig, pp 849t/s, tg 61t/s</summary>

```
| model                          |       size |     params | backend    | ngl | n_batch |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | ------: | --------------: | -------------------: |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |      64 |   pp512 @ d8192 |        504.32 ± 1.84 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |      64 |    tg32 @ d8192 |         61.20 ± 1.85 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     128 |   pp512 @ d8192 |        522.16 ± 1.00 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     128 |    tg32 @ d8192 |         61.25 ± 1.71 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     256 |   pp512 @ d8192 |        690.77 ± 2.34 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     256 |    tg32 @ d8192 |         61.07 ± 1.95 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     512 |   pp512 @ d8192 |        849.02 ± 6.51 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     512 |    tg32 @ d8192 |         60.70 ± 1.63 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |    1024 |   pp512 @ d8192 |        844.93 ± 6.80 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |    1024 |    tg32 @ d8192 |         60.74 ± 1.64 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |    2048 |   pp512 @ d8192 |        844.30 ± 7.04 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |    2048 |    tg32 @ d8192 |         60.49 ± 1.42 |

build: 23cccea2 (8102)
```
</details>

<details>
<summary>block, pp 974t/s, tg 61t/s</summary>

```
| model                          |       size |     params | backend    | ngl | n_batch |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | ------: | --------------: | -------------------: |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |      64 |   pp512 @ d8192 |        519.94 ± 1.99 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |      64 |    tg32 @ d8192 |         62.09 ± 1.99 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     128 |   pp512 @ d8192 |        556.54 ± 1.10 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     128 |    tg32 @ d8192 |         61.78 ± 1.71 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     256 |   pp512 @ d8192 |        762.60 ± 0.48 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     256 |    tg32 @ d8192 |         61.53 ± 1.98 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     512 |   pp512 @ d8192 |       973.75 ± 12.92 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |     512 |    tg32 @ d8192 |         61.44 ± 1.80 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |    1024 |   pp512 @ d8192 |       969.33 ± 12.86 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |    1024 |    tg32 @ d8192 |         61.35 ± 1.62 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |    2048 |   pp512 @ d8192 |       968.47 ± 14.16 |
| kimi-linear 48B.A3B Q2_K - Medium |  16.78 GiB |    49.12 B | CUDA       |  99 |    2048 |    tg32 @ d8192 |         61.28 ± 1.60 |

build: 23cccea2 (8102)
```
</details>

<details>
<summary>block implementation for gated delta net</summary>

```
        const int64_t HB = H_k * n_seqs;
        kb = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, CS, CS, n_chunks, HB);
        kb = ggml_clamp(ctx0, kb, 0.0f, 0.0f);
        kq = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, CS, CS, n_chunks, HB);
        kq = ggml_clamp(ctx0, kq, 0.0f, 0.0f);
        const int64_t block_size = BLOCK_SIZE;
        const int64_t n_blocks = CHUNK_SIZE / BLOCK_SIZE;

        ggml_tensor * k_block[n_blocks];
        ggml_tensor * q_block[n_blocks];
        ggml_tensor * g_block[n_blocks];
        ggml_tensor * g_block_bc[n_blocks];
        for (int64_t j = 0; j < n_blocks; ++j) {
            int64_t j_start = j * block_size;

            // k_i_block: [S, block_size, C, HB]
            k_block[j] = ggml_view_4d(ctx0, k,
                S_k, block_size, n_chunks, HB,
                k->nb[1], k->nb[2], k->nb[3],
                j_start * k->nb[1]);

            // q_i_block: [S, block_size, C, HB]
            q_block[j] = ggml_view_4d(ctx0, q,
                    S_k, block_size, n_chunks, HB,
                    q->nb[1], q->nb[2], q->nb[3],
                    j_start * q->nb[1]);

            // g_j_block: [S, block_size, C, HB]
            g_block[j] = ggml_cont(ctx0, ggml_view_4d(ctx0, g_cs,
                block_size, 1, n_chunks, HB,
                g_cs->nb[1], g_cs->nb[2], g_cs->nb[3],
                j_start * g_cs->nb[0]));

            g_block_bc[j] = ggml_reshape_4d(ctx0, g_block[j], 1, block_size, n_chunks, HB);
            g_block_bc[j] = ggml_repeat_4d(ctx0, g_block_bc[j], block_size, block_size, n_chunks, HB);
        }

        for (int64_t j = 0; j < n_blocks; ++j) {
            int64_t j_start = j * block_size;

            for (int64_t i = 0; i <= j; ++i) {
                int64_t i_start = i * block_size;

                ggml_tensor * decay_mask = ggml_sub(ctx0, g_block_bc[j], g_block[i]);
                cb(decay_mask, "decay_mask", il);

                // Apply diag_mask only at diagnoal blocks
                if (i == j) {
                    decay_mask = ggml_tri(ctx0, decay_mask, GGML_TRI_TYPE_LOWER_DIAG);
                }
                decay_mask = ggml_exp(ctx0, decay_mask);

                ggml_tensor * Akk_block = ggml_mul_mat(ctx0, k_block[i], k_block[j]);
                ggml_tensor * Aqk_block = ggml_mul_mat(ctx0, q_block[i], k_block[j]);
                Akk_block = ggml_mul(ctx0, Akk_block, decay_mask);
                Aqk_block = ggml_mul(ctx0, Aqk_block, decay_mask);

                if (i == j) {
                    Aqk_block = ggml_tri(ctx0, Aqk_block, GGML_TRI_TYPE_LOWER_DIAG);
                    Akk_block = ggml_tri(ctx0, Akk_block, GGML_TRI_TYPE_LOWER);
                }

                // Accumulate into Akk and Aqk at position [j_start:j_end, i_start:i_end]
                kb = ggml_set(ctx0, kb, Akk_block,
                    kb->nb[1], kb->nb[2], kb->nb[3],
                    i_start * kb->nb[0] + j_start * kb->nb[1]);
                kq = ggml_set(ctx0, kq, Aqk_block,
                    kq->nb[1], kq->nb[2], kq->nb[3],
            }
        }
        kb = ggml_mul(ctx0, kb, b);
        cb(kq, "kq", il);
```
</details>

Changed files:
filename: src/models/delta-net-base.cpp
patch: @@ -1,7 +1,5 @@
 #include "models.h"
 
-#define CHUNK_SIZE 64
-
 // utility to get one slice from the third dimension
 // input dim:  [x, y, c, b]
 // output dim: [x, y, 1, b]
@@ -57,7 +55,7 @@ std::pair<ggml_tensor *, ggml_tensor *> llm_build_delta_net_base::build_delta_ne
     g = ggml_permute(ctx0, g, 0, 2, 1, 3); // [g_0, n_tokens, H_v, n_seqs]
     b = ggml_permute(ctx0, b, 0, 2, 1, 3); // [  1, n_tokens, H_v, n_seqs]
 
-    const int CS = CHUNK_SIZE;
+    const int CS = kda ? 16 : 64; // chunk size
 
     const int pad = (CS - n_tokens % CS) % CS;
     const int n_chunks = (n_tokens + pad) / CS;