Add batch warmup to sweep-bench (#375)

Co-authored-by: Iwan Kawrakow <iwan.kawrakow@gmail.com>
This commit is contained in:
Kawrakow
2025-05-12 07:50:26 +03:00
committed by GitHub
parent 2e585d4508
commit ceb8f513e4
3 changed files with 22 additions and 1 deletions

View File

@@ -107,7 +107,7 @@ int main(int argc, char ** argv) {
llama_batch batch = llama_batch_init(n_kv_max, 0, 1);
// warm up
{
if (params.warmup) {
llama_batch_add(batch, bos, 0, { 0 }, false);
if (!decode_helper(ctx, batch, ctx_params.n_batch)) {
@@ -115,6 +115,22 @@ int main(int argc, char ** argv) {
return 1;
}
}
if (params.batch_warmup) {
// clean up KV cache after generation
llama_kv_cache_seq_rm(ctx, 0, params.n_ubatch, -1);
// prepare batch of pp size for prompt processing performance measurement
llama_batch_clear(batch);
for (unsigned int i = 0; i < params.n_ubatch; ++i) {
llama_batch_add(batch, std::rand() % n_vocab, i, { 0 }, false);
}
if (!decode_helper(ctx, batch, ctx_params.n_ubatch)) {
LOG_TEE("%s: llama_decode() failed\n", __func__);
return 1;
}
}
llama_batch_clear(batch);
llama_kv_cache_clear(ctx);