Use MAX_LIVE_BATCHES for prefill method

PiperOrigin-RevId: 621347176 Change-Id: I339e204fe9b9d6bde6105daea9b62da139911497
google · Apr 3, 2024 · 3ac3f3d · 3ac3f3d
1 parent 50beb83
commit 3ac3f3d
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/saxml/server/model_service_base.py b/saxml/server/model_service_base.py
@@ -1633,9 +1633,7 @@ def preprocess_rpc_tasks(
               preprocess_fn=functools.partial(
                   preprocess_rpc_tasks, prefill=True
               ),
-              # Each batch has `batch_size` requests and we should have at least
-              # `num_cache_slots` * 2 live requests
-              max_live_batches=method.num_cache_slots * 2 // method.batch_size,
+              max_live_batches=method.max_live_batches,
               batching_wait_secs=method.batching_wait_secs,
           )