From 4201a8be46cf45193c138d852fe6d725fa610559 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Tue, 18 Jul 2023 11:39:14 +0200
Subject: [PATCH] fix default value

---
 router/src/main.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/router/src/main.rs b/router/src/main.rs
index 2f793040..03e920c6 100644
--- a/router/src/main.rs
+++ b/router/src/main.rs
@@ -220,7 +220,8 @@ fn main() -> Result<(), RouterError> {
                 .map_err(RouterError::Warmup)?
             {
                 // Older models do not support automatic max-batch-total-tokens
-                None => max_batch_total_tokens.unwrap_or(16000),
+                None => max_batch_total_tokens
+                    .unwrap_or(16000.max((max_total_tokens as u32).max(max_batch_prefill_tokens))),
                 // Flash attention models return their max supported total tokens
                 Some(max_supported_batch_total_tokens) => {
                     // Warn if user added his own max-batch-total-tokens as we will ignore it