Mirrors
/
hf_text-generation-inference
mirror of https://github.com/huggingface/text-generation-inference.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix(benchmarking): fix benchmarking tool

This commit is contained in:
OlivierDehaene 2023-04-26 00:54:27 +02:00
parent 45344244cf
commit 7de8a377b0
1 changed files with 17 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
benchmark/src/generation.rs
View File

@ -74,16 +74,28 @@ async fn generate_runs(
for b in batch_size { for b in batch_size {
// Warmups on batch size // Warmups on batch size
for _ in 0..warmups { for _ in 0..warmups {
let (_, decode_batch) = let (_, decode_batch) = prefill(
prefill(sequence.clone(), sequence_length, b, decode_length, &mut client).await?; sequence.clone(),
sequence_length,
b,
decode_length,
&mut client,
)
.await?;
let _ = decode(decode_batch, &mut client).await?; let _ = decode(decode_batch, &mut client).await?;
// Send warmup message // Send warmup message
run_sender.send(Ok(Message::Warmup)).await.unwrap_or(()); run_sender.send(Ok(Message::Warmup)).await.unwrap_or(());
} }
for _ in 0..n_runs { for _ in 0..n_runs {
let (prefill, decode_batch) = let (prefill, decode_batch) = prefill(
prefill(sequence.clone(), sequence_length, b, decode_length, &mut client).await?; sequence.clone(),
sequence_length,
b,
decode_length,
&mut client,
)
.await?;
// Send prefill message // Send prefill message
run_sender run_sender
.send(Ok(Message::Prefill(prefill))) .send(Ok(Message::Prefill(prefill)))
@ -143,6 +155,7 @@ async fn prefill(
id: 0, id: 0,
requests, requests,
size: batch_size, size: batch_size,
max_tokens: batch_size * (sequence_length + decode_length),
}; };
// Run prefill // Run prefill