feat(router): Dynamic batch sizing

Author: njhill

launcher/src/main.rs

@@ -41,6 +41,10 @@ struct Args {
     max_total_tokens: usize,
     #[clap(default_value = "32", long, env)]
     max_batch_size: usize,
+    #[clap(default_value = None, long, env)]
+    max_batch_weight: Option<usize>,
+    #[clap(default_value = None, long, env)]
+    max_prefill_weight: Option<usize>,
     #[clap(default_value = "20", long, env)]
     max_waiting_tokens: usize,
     #[clap(default_value = "3000", long, short, env)]
@@ -93,6 +97,8 @@ fn main() -> ExitCode {
         max_input_length,
         max_total_tokens,
         max_batch_size,
+        max_batch_weight,
+        max_prefill_weight,
         max_waiting_tokens,
         port,
         shard_uds_path,
@@ -392,6 +398,16 @@ fn main() -> ExitCode {
         model_id,
     ];
 
+    if let Some(max_batch_weight) = max_batch_weight {
+        argv.push("--max-batch-weight".to_string());
+        argv.push(max_batch_weight.to_string())
+    }
+
+    if let Some(max_prefill_weight) = max_prefill_weight {
+        argv.push("--max-batch-weight".to_string());
+        argv.push(max_prefill_weight.to_string())
+    }
+
     // Model optional revision
     if let Some(ref revision) = revision {
         argv.push("--revision".to_string());

router/src/infer.rs

@@ -14,6 +14,7 @@ use thiserror::Error;
 use tokio::sync::{Notify, Semaphore, TryAcquireError};
 use tokio::time::Instant;
 use tracing::{info_span, instrument, Instrument, Span};
+use crate::queue::BatchingConfig;
 
 /// Inference struct
 #[derive(Clone)]
@@ -39,19 +40,24 @@ impl Infer {
         client: ShardedClient,
         validation: Validation,
         max_batch_size: usize,
+        max_batch_weight: usize,
+        max_prefill_weight: usize,
         max_waiting_tokens: usize,
         max_concurrent_requests: usize,
     ) -> Self {
         // Infer shared state
-        let queue = Queue::new();
+        let queue = Queue::new(BatchingConfig {
+            size_limit: max_batch_size,
+            weight_limit: max_batch_weight,
+            prefill_weight_limit: max_prefill_weight,
+        });
         let shared = Arc::new(Shared {
             batching_task: Notify::new(),
         });
 
         // Spawn batching background task that contains all the inference logic
         tokio::spawn(batching_task(
             client,
-            max_batch_size,
             max_waiting_tokens,
             queue.clone(),
             shared.clone(),
@@ -99,6 +105,7 @@ impl Infer {
         // Append the request to the queue
         self.queue.append(Entry {
             request: valid_request,
+            generated_tokens: 0,
             response_tx,
             span: Span::current(),
             temp_span: None,
@@ -227,18 +234,11 @@ impl Infer {
 /// Batches requests and sends them to the inference server
 async fn batching_task(
     mut client: ShardedClient,
-    max_batch_size: usize,
+    // max_batch_size: usize,
     max_waiting_tokens: usize,
     queue: Queue,
     shared: Arc<Shared>,
 ) {
-    // Minimum batch size after which we try to add more requests
-    let limit_min_batch_size = if max_batch_size > 1 {
-        (max_batch_size / 2) as u32
-    } else {
-        0
-    };
-
     // Infinite loop
     loop {
         // Wait for a notification from the Infer struct
@@ -247,8 +247,8 @@ async fn batching_task(
         // Get the next batch from the queue
         // This batch might be smaller than the maximum batch size if there are not enough requests
         // waiting in the queue
-        while let Some((mut entries, batch, span)) = queue.next_batch(None, max_batch_size).await {
-            let mut cached_batch = prefill(&mut client, batch, &mut entries)
+        while let (_, Some((mut entries, batch, span))) = queue.next_batch(None).await {
+            let (mut cached_batch, mut some_completed) = prefill(&mut client, batch, &mut entries)
                 .instrument(span)
                 .await;
             let mut waiting_tokens = 1;
@@ -261,21 +261,16 @@ async fn batching_task(
                 let mut batches = vec![batch];
                 metrics::gauge!("tgi_batch_current_size", batch_size as f64);
 
-                // If the current batch is too small, we try to add more requests to it
-                if batch_size <= limit_min_batch_size {
-                    let min_size = match waiting_tokens {
-                        // If we didn't onboard any new requests since >= max_waiting_tokens, we try
-                        // to add a new batch even though its size might be small
-                        _ if waiting_tokens >= max_waiting_tokens => None,
-                        // Minimum size criteria
-                        _ => Some(limit_min_batch_size as usize),
-                    };
-
-                    // Try to get a new batch
-                    if let Some((mut new_entries, new_batch, span)) = queue
-                        .next_batch(min_size, max_batch_size - batch_size as usize)
-                        .await
-                    {
+                // Try to extend batch if its size reduced or enough tokens have elapsed since last one
+                if some_completed || waiting_tokens >= max_waiting_tokens {
+
+                    // Try to get a new batch - ownership of entries passed in and out
+                    let (
+                        existing_entries, new_entries
+                    ) = queue.next_batch(Some(entries)).await;
+                    entries = existing_entries.unwrap();
+
+                    if let Some((mut new_entries, new_batch, span)) = new_entries {
                         let new_batch_size = new_batch.size;
                         entries.iter_mut().for_each(|(_, entry)| {
                             // Create a new span to add the info that this entry is waiting
@@ -290,7 +285,7 @@ async fn batching_task(
                         });
 
                         // Generate one token for this new batch to have the attention past in cache
-                        let new_cached_batch = prefill(&mut client, new_batch, &mut new_entries)
+                        let (new_cached_batch, _) = prefill(&mut client, new_batch, &mut new_entries)
                             .instrument(span)
                             .await;
                         // Reset waiting counter
@@ -317,7 +312,7 @@ async fn batching_task(
                     entry.temp_span = Some(entry_batch_span);
                 });
 
-                cached_batch = decode(&mut client, batches, &mut entries)
+                (cached_batch, some_completed) = decode(&mut client, batches, &mut entries)
                     .instrument(next_batch_span)
                     .await;
                 waiting_tokens += 1;
@@ -332,24 +327,24 @@ async fn prefill(
     client: &mut ShardedClient,
     batch: Batch,
     entries: &mut IntMap<u64, Entry>,
-) -> Option<Batch> {
+) -> (Option<Batch>, bool) {
     let start_time = Instant::now();
     let batch_id = batch.id;
     metrics::increment_counter!("tgi_batch_inference_count", "method" => "prefill");
 
     match client.prefill(batch).await {
         Ok((generations, next_batch)) => {
-            send_generations(generations, entries);
+            let some_completed = send_generations(generations, entries);
             metrics::histogram!("tgi_batch_inference_duration", start_time.elapsed().as_secs_f64(), "method" => "prefill");
             metrics::increment_counter!("tgi_batch_inference_success", "method" => "prefill");
-            next_batch
+            (next_batch, some_completed)
         }
         // If we have an error, we discard the whole batch
         Err(err) => {
             let _ = client.clear_cache(Some(batch_id)).await;
             send_errors(err, entries);
             metrics::increment_counter!("tgi_batch_inference_failure", "method" => "prefill");
-            None
+            (None, true)
         }
     }
 }
@@ -359,22 +354,22 @@ async fn decode(
     client: &mut ShardedClient,
     batches: Vec<Batch>,
     entries: &mut IntMap<u64, Entry>,
-) -> Option<Batch> {
+) -> (Option<Batch>, bool) {
     let start_time = Instant::now();
     metrics::increment_counter!("tgi_batch_inference_count", "method" => "decode");
 
     match client.decode(batches).await {
         Ok((generations, next_batch)) => {
-            send_generations(generations, entries);
+            let some_completed = send_generations(generations, entries);
             metrics::histogram!("tgi_batch_inference_duration", start_time.elapsed().as_secs_f64(), "method" => "decode");
             metrics::increment_counter!("tgi_batch_inference_success", "method" => "decode");
-            next_batch
+            (next_batch, some_completed)
         }
         // If we have an error, we discard the whole batch
         Err(err) => {
             send_errors(err, entries);
             metrics::increment_counter!("tgi_batch_inference_failure", "method" => "decode");
-            None
+            (None, true)
         }
     }
 }
@@ -398,13 +393,15 @@ fn send_errors(error: ClientError, entries: &mut IntMap<u64, Entry>) {
 }
 
 /// Send one or multiple `InferStreamResponse` to Infer for all `entries`
+/// Return true if any requests completed
 #[instrument(skip_all)]
-fn send_generations(generations: Vec<Generation>, entries: &mut IntMap<u64, Entry>) {
+fn send_generations(generations: Vec<Generation>, entries: &mut IntMap<u64, Entry>) -> bool {
+    let mut some_completed = false;
     generations.into_iter().for_each(|generation| {
         // Get entry
         // We can `expect` here as the request id should always be in the entries
         let entry = entries
-            .get(&generation.request_id)
+            .get_mut(&generation.request_id)
             .expect("ID not found in entries. This is a bug.");
 
         // Create and enter a span to link this function back to the entry
@@ -445,7 +442,10 @@ fn send_generations(generations: Vec<Generation>, entries: &mut IntMap<u64, Entr
                     start: entry.batch_time.unwrap(),
                 }))
                 .unwrap_or(());
+            some_completed = true;
         } else {
+            // Increment generated token count
+            entry.generated_tokens += 1;
             // Send message
             // unwrap_or is valid here as we don't care if the receiver is gone.
             entry
@@ -454,6 +454,8 @@ fn send_generations(generations: Vec<Generation>, entries: &mut IntMap<u64, Entr
                 .unwrap_or(());
         }
     });
+
+    return some_completed;
 }
 
 #[derive(Debug)]

router/src/main.rs

@@ -33,6 +33,10 @@ struct Args {
     max_total_tokens: usize,
     #[clap(default_value = "32", long, env)]
     max_batch_size: usize,
+    #[clap(default_value = None, long, env)]
+    max_batch_weight: Option<usize>,
+    #[clap(default_value = None, long, env)]
+    max_prefill_weight: Option<usize>,
     #[clap(default_value = "20", long, env)]
     max_waiting_tokens: usize,
     #[clap(default_value = "3000", long, short, env)]
@@ -64,6 +68,8 @@ fn main() -> Result<(), std::io::Error> {
         max_input_length,
         max_total_tokens,
         max_batch_size,
+        max_batch_weight,
+        max_prefill_weight,
         max_waiting_tokens,
         port,
         master_shard_uds_path,
@@ -169,6 +175,8 @@ fn main() -> Result<(), std::io::Error> {
                 max_input_length,
                 max_total_tokens,
                 max_batch_size,
+                max_batch_weight,
+                max_prefill_weight,
                 max_waiting_tokens,
                 sharded_client,
                 tokenizer,