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

Enabling custom block size schedule.

Browse Source
This commit is contained in:
Nicolas Patry 2024-05-30 05:17:00 +00:00
parent cf59593454
commit 13caf958eb
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
router/src/infer.rs
View File

@ -77,6 +77,9 @@ impl Infer {
false false
}; };
let block_size = if flashdecoding { 256 } else { 32 }; let block_size = if flashdecoding { 256 } else { 32 };
let block_size = std::env::var("BLOCK_SIZE")
.map(|b| b.parse().unwrap_or(block_size))
.unwrap_or(block_size);
let queue = Queue::new(requires_padding, block_size, window_size, speculate); let queue = Queue::new(requires_padding, block_size, window_size, speculate);
let shared = Arc::new(Shared { let shared = Arc::new(Shared {
batching_task: Notify::new(), batching_task: Notify::new(),