From 3416f98c587a45b34161af08b905781a4c758af0 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sun, 9 Aug 2026 10:56:05 +0800
Subject: [PATCH] fix: wire benchmark cache selection

---
 scripts/tools/benchmark.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py
index f80afbd..0ff09f0 100644
--- a/scripts/tools/benchmark.py
+++ b/scripts/tools/benchmark.py
@@ -68,6 +68,15 @@ class GenerationBenchmark:
         self.cuda_graph = cuda_graph
 
     def _make_pool(self, batch_size: int, max_seq_len: int) -> PagePool:
+        if self.cache_type == "contiguous":
+            n_tokens = None
+        elif self.cache_type == "paged":
+            # Keep the total token capacity equal to contiguous mode while
+            # routing allocation through the shared paged pool.
+            n_tokens = batch_size * max_seq_len
+        else:
+            raise ValueError(f"unsupported cache type: {self.cache_type}")
+
         return PagePool(
             n_layers=self.config.num_hidden_layers,
             n_kv_heads=self.config.num_key_value_heads,
@@ -77,7 +86,7 @@ class GenerationBenchmark:
             device=self.device,
             dtype=self.dtype,
             page_size=1,
-            n_tokens=None,
+            n_tokens=n_tokens,
         )
 
     @staticmethod