Repository: ggml-org/llama.cpp
Pull request: batched-bench : add "separate text gen" mode

Adding `-tgs` to `llama-batched-bench` would make it decode the sequences separately, one by one:

```bash
# no -tgs
0123 0123 0123 ...

# -tgs
0 0 0 ... 1 1 1 ... 2 2 2 ... 3 3 3 ...
```

This is useful for benchmarking the performance of the unified KV cache where it's important to detect and skip masked regions in the KQ mask.

Example with the Metal backend:

```bash
# unified KV cache with up to 4 sequences, running one by one
llama-batched-bench -m ../models/gemma-3-4b-it/ggml-model-f16.gguf -c 33792 -npp 8192 -ntg 32 -npl 1,2,4 -kvu -tgs

# the cache looks like this
#
#                        prompt processing ends here v
# 000...[8192 tokens]...000111...111222...222333...333000...[32 tokens]...000111...111222...222333...333
#.                             text generation starts ^
```

With the -INF block optimizations in the FA kernels:

```
main: n_kv_max = 33792, n_batch = 2048, n_ubatch = 512, flash_attn = -1, is_pp_shared = 0, is_tg_separate = 1, n_gpu_layers = -1, n_threads = 16, n_threads_batch = 16

|    PP |     TG |    B |   N_KV |   T_PP s | S_PP t/s |   T_TG s | S_TG t/s |      T s |    S t/s |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|  8192 |     32 |    1 |   8224 |    3.101 |  2641.88 |    0.478 |    66.95 |    3.579 |  2298.00 |
|  8192 |     32 |    2 |  16448 |    6.091 |  2689.76 |    0.971 |    65.90 |    7.062 |  2328.95 |
|  8192 |     32 |    4 |  32896 |   12.373 |  2648.43 |    1.965 |    65.15 |   14.337 |  2294.45 |
```

Disabling the -INF block optimizations in the FA kernels:

<details>
<summary>patch</summary>

```diff
diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal
index cea535ade..6c249fb56 100644
--- a/ggml/src/ggml-metal/ggml-metal.metal
+++ b/ggml/src/ggml-metal/ggml-metal.metal
@@ -4633,7 +4633,7 @@ kernel void kernel_flash_attn_ext_blk(
     const int32_t nblk0 = ((args.ne30 + C - 1)/C);
 
     if (tiisg == 0) {
-        dst[((i3*args.ne32 + i2)*nblk1 + i1)*nblk0 + i0] = res;
+        dst[((i3*args.ne32 + i2)*nblk1 + i1)*nblk0 + i0] = 1;
     }
 }
 
@@ -5660,7 +5660,7 @@ void kernel_flash_attn_ext_vec_impl(
             }
 
             // skip -INF blocks
-            if (simd_max(sm[tiisg]) == -INFINITY) {
+            if (false) {
                 continue;
             }
 
```
</details>

```
main: n_kv_max = 33792, n_batch = 2048, n_ubatch = 512, flash_attn = -1, is_pp_shared = 0, is_tg_separate = 1, n_gpu_layers = -1, n_threads = 16, n_threads_batch = 16

|    PP |     TG |    B |   N_KV |   T_PP s | S_PP t/s |   T_TG s | S_TG t/s |      T s |    S t/s |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|  8192 |     32 |    1 |   8224 |    3.528 |  2321.91 |    0.500 |    64.05 |    4.028 |  2041.84 |
|  8192 |     32 |    2 |  16448 |    7.393 |  2216.28 |    1.027 |    62.30 |    8.420 |  1953.47 |
|  8192 |     32 |    4 |  32896 |   16.157 |  2028.06 |    2.159 |    59.30 |   18.316 |  1796.04 |
```

Observe that both `pp` and `tg` perf is worse and it's amplified with more sequences in the cache.

Changed files:
filename: common/arg.cpp
patch: @@ -2253,6 +2253,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.is_pp_shared = true;
         }
     ).set_examples({LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL}));
+    add_opt(common_arg(
+        {"-tgs"},
+        string_format("is the text generation separated across the different sequences (default: %s)", params.is_tg_separate ? "true" : "false"),
+        [](common_params & params) {
+            params.is_tg_separate = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL}));
     add_opt(common_arg(
         {"-npp"}, "n0,n1,...",
         "number of prompt tokens",

filename: common/common.h
patch: @@ -460,7 +460,8 @@ struct common_params {
     float slot_prompt_similarity = 0.1f;
 
     // batched-bench params
-    bool is_pp_shared = false;
+    bool is_pp_shared   = false;
+    bool is_tg_separate = false;
 
     std::vector<int32_t> n_pp;
     std::vector<int32_t> n_tg;

filename: tools/batched-bench/batched-bench.cpp
patch: @@ -23,7 +23,8 @@ int main(int argc, char ** argv) {
 
     common_init();
 
-    int is_pp_shared = params.is_pp_shared;
+    int is_pp_shared   = params.is_pp_shared;
+    int is_tg_separate = params.is_tg_separate;
 
     std::vector<int> n_pp = params.n_pp;
     std::vector<int> n_tg = params.n_tg;
@@ -72,8 +73,8 @@ int main(int argc, char ** argv) {
 
     // decode in batches of ctx_params.n_batch tokens
     auto decode_helper = [](llama_context * ctx, llama_batch & batch, int32_t n_batch, bool synchronize) {
-        for (int32_t i = 0; i < (int32_t) batch.n_tokens; i += n_batch) {
-            const int32_t n_tokens = std::min(n_batch, (int32_t) (batch.n_tokens - i));
+        for (int32_t i = 0; i < batch.n_tokens; i += n_batch) {
+            const int32_t n_tokens = std::min(n_batch, batch.n_tokens - i);
 
             llama_batch batch_view = {
                 n_tokens,
@@ -113,7 +114,7 @@ int main(int argc, char ** argv) {
 
     if (!params.batched_bench_output_jsonl) {
         LOG("\n");
-        LOG("%s: n_kv_max = %d, n_batch = %d, n_ubatch = %d, flash_attn = %d, is_pp_shared = %d, n_gpu_layers = %d, n_threads = %u, n_threads_batch = %u\n", __func__, n_kv_max, params.n_batch, params.n_ubatch, int(params.flash_attn_type), params.is_pp_shared, params.n_gpu_layers, ctx_params.n_threads, ctx_params.n_threads_batch);
+        LOG("%s: n_kv_max = %d, n_batch = %d, n_ubatch = %d, flash_attn = %d, is_pp_shared = %d, is_tg_separate = %d, n_gpu_layers = %d, n_threads = %u, n_threads_batch = %u\n", __func__, n_kv_max, params.n_batch, params.n_ubatch, int(params.flash_attn_type), is_pp_shared, is_tg_separate, params.n_gpu_layers, ctx_params.n_threads, ctx_params.n_threads_batch);
         LOG("\n");
         LOG("|%6s | %6s | %4s | %6s | %8s | %8s | %8s | %8s | %8s | %8s |\n", "PP", "TG", "B", "N_KV", "T_PP s", "S_PP t/s", "T_TG s", "S_TG t/s", "T s", "S t/s");
         LOG("|%6s-|-%6s-|-%4s-|-%6s-|-%8s-|-%8s-|-%8s-|-%8s-|-%8s-|-%8s-|\n", "------", "------", "----", "------", "--------", "--------", "--------", "--------", "--------", "--------");
@@ -172,16 +173,35 @@ int main(int argc, char ** argv) {
 
                 const auto t_tg_start = ggml_time_us();
 
-                for (int i = 0; i < tg; ++i) {
-                    common_batch_clear(batch);
-
+                if (is_tg_separate) {
+                    // decode pattern:
+                    // 0 0 0 ... 1 1 1 ... 2 2 2 ... 3 3 3 ...
                     for (int j = 0; j < pl; ++j) {
-                        common_batch_add(batch, get_token_rand(), pp + i, { j }, true);
+                        for (int i = 0; i < tg; ++i) {
+                            common_batch_clear(batch);
+
+                            common_batch_add(batch, get_token_rand(), pp + i, { j }, true);
+
+                            if (!decode_helper(ctx, batch, ctx_params.n_batch, true)) {
+                                LOG_ERR("%s: llama_decode() failed\n", __func__);
+                                return 1;
+                            }
+                        }
                     }
+                } else {
+                    // decode pattern:
+                    // 0123 0123 0123 ...
+                    for (int i = 0; i < tg; ++i) {
+                        common_batch_clear(batch);
 
-                    if (!decode_helper(ctx, batch, ctx_params.n_batch, true)) {
-                        LOG_ERR("%s: llama_decode() failed\n", __func__);
-                        return 1;
+                        for (int j = 0; j < pl; ++j) {
+                            common_batch_add(batch, get_token_rand(), pp + i, { j }, true);
+                        }
+
+                        if (!decode_helper(ctx, batch, ctx_params.n_batch, true)) {
+                            LOG_ERR("%s: llama_decode() failed\n", __func__);
+                            return 1;
+                        }
                     }
                 }