huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 12:54:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge commit 'refs/pull/2352/head' of github.com:huggingface/text-generation-inference into pr-2352-ci-branch

Browse Source
This commit is contained in:
drbh 2024-08-08 16:27:00 +00:00
commit 6497ae61e2
2 changed files with 12 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
router/src/infer/v2/queue.rs
View File

@ -205,6 +205,13 @@ impl State {
} }
} }
if let Some(max_size) = max_size {
if max_size == 0 {
tracing::debug!("No capacity");
return None;
}
}
// Pad prefill_token_budget to be a multiple of block size // Pad prefill_token_budget to be a multiple of block size
let prefill_token_budget = let prefill_token_budget =
((prefill_token_budget + self.block_size - 1) / self.block_size) * self.block_size; ((prefill_token_budget + self.block_size - 1) / self.block_size) * self.block_size;
@ -297,7 +304,7 @@ impl State {
batch_entries.insert(id, entry); batch_entries.insert(id, entry);
// Check if max_size // Check if max_size
if Some(batch_requests.len()) == max_size { if Some(batch_requests.len()) >= max_size {
break; break;
} }
} }

5
router/src/infer/v2/scheduler.rs
View File

@ -161,7 +161,10 @@ pub(crate) async fn batching_task(
}; };
let token_budget = max_batch_total_tokens.saturating_sub(batch_max_tokens); let token_budget = max_batch_total_tokens.saturating_sub(batch_max_tokens);
let max_size = max_batch_size.map(|max_size| max_size - batch_size as usize); let max_size = max_batch_size.map(|max_size| {
if batch_size as usize > max_size { 0 } else { max_size - batch_size as usize }
});
// Try to get a new batch // Try to get a new batch
if let Some((mut new_entries, new_batch, span)) = queue if let Some((mut new_entries, new_batch, span)) = queue