ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: size KV pool from prompt/gen args in benchmark

Browse Source 
- Drop hardcoded CACHE_MAX_SEQ=2048 which overflowed at long prompts
- Size prefill pool to prompt_length and decode pool to prompt+5+gen*num_trials
- Unblocks decode/prefill benchmark at prompt 4096+ (was KV cache index OOB)
This commit is contained in:
ViperEkura
2026-08-02 16:25:27 +08:00
parent b1b65a657e
commit 8447f88f61
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/tools/benchmark.py
+7 -5
View File
@@ -13,7 +13,6 @@ from astrai.model import AutoModel
_DTYPES = ["bfloat16", "float16", "float32"] _DTYPES = ["bfloat16", "float16", "float32"]
_CACHES = ["contiguous", "paged"] _CACHES = ["contiguous", "paged"]
_BACKENDS = ["cuda", "torch_native"] _BACKENDS = ["cuda", "torch_native"]
CACHE_MAX_SEQ = 2048
_BACKEND_MAP = { _BACKEND_MAP = {
"cuda": ATTN_BACKEND.CUDA, "cuda": ATTN_BACKEND.CUDA,
@@ -56,13 +55,13 @@ class GenerationBenchmark:
self.config = config self.config = config
self.backend = backend self.backend = backend
def _make_pool(self, batch_size: int) -> PagePool: def _make_pool(self, batch_size: int, max_seq_len: int) -> PagePool:
return PagePool( return PagePool(
n_layers=self.config.num_hidden_layers, n_layers=self.config.num_hidden_layers,
n_kv_heads=self.config.num_key_value_heads, n_kv_heads=self.config.num_key_value_heads,
head_dim=self.config.hidden_size // self.config.num_attention_heads, head_dim=self.config.hidden_size // self.config.num_attention_heads,
max_batch_size=batch_size, max_batch_size=batch_size,
max_seq_len=CACHE_MAX_SEQ, max_seq_len=max_seq_len,
device=self.device, device=self.device,
dtype=self.dtype, dtype=self.dtype,
page_size=1, page_size=1,
@@ -128,7 +127,7 @@ class GenerationBenchmark:
) -> BenchmarkResult: ) -> BenchmarkResult:
import time import time
pool = self._make_pool(batch_size) pool = self._make_pool(batch_size, prompt_length)
task_ids = [f"bench_prefill_{i}" for i in range(batch_size)] task_ids = [f"bench_prefill_{i}" for i in range(batch_size)]
for tid in task_ids: for tid in task_ids:
pool.task_alloc(tid, list(range(prompt_length))) pool.task_alloc(tid, list(range(prompt_length)))
@@ -189,7 +188,10 @@ class GenerationBenchmark:
) -> BenchmarkResult: ) -> BenchmarkResult:
import time import time
pool = self._make_pool(batch_size) # Decode grows seq_len monotonically up to prompt + 5 + gen*num_trials
# (warmup 5 steps, then one step per trial), so size the pool to cover it.
max_seq_len = prompt_length + 5 + gen_length * num_trials
pool = self._make_pool(batch_size, max_seq_len)
task_ids = self._run_prefill(pool, batch_size, prompt_length) task_ids = self._run_prefill(pool, batch_size, prompt_length)
for i in range(5): for i in range(5):