[CB] refactoring warmup for batch size 1 (#347)

### [CB] refactoring warmup for batch size 1

From
#312 (comment)
there is a request for a nicer integration of batch size 1 support
during warmup. Most of the code is already on main, thus this PR.

Signed-off-by: Yannick Schnider <[email protected]>

Author: yannicks1

vllm_spyre/v1/worker/spyre_model_runner.py

@@ -290,8 +290,13 @@ def load_model(self, prompt_lens: Iterable[int],
         )
 
     def build_input_batch(self) -> SamplingInputBatch:
+        # Fix for batch size 1: set input batch to fit 2 requests for warmup,
+        # and reset input batch to fit max_num_seqs requests after warmup
+        min_seqs_required = 2 if self.warmup_mode else 1
+
         return SamplingInputBatch(
-            max_num_reqs=self.scheduler_config.max_num_seqs,
+            max_num_reqs=max(min_seqs_required,
+                             self.scheduler_config.max_num_seqs),
             max_model_len=self.model_config.max_model_len,
             device=self.device,
             pin_memory=self.pin_memory,
@@ -802,7 +807,10 @@ def __init__(
             vocab_size=vllm_config.model_config.get_vocab_size(),
         )
 
-    def finish_warmup(self) -> None:
+    def complete_warmup(self) -> None:
+        super().complete_warmup()
+        # Fix for batch size 1: need to update the input_batch after the warmup
+        self.input_batch = self.build_input_batch()
         # get the number or pages from the actual Spyre card after the warmup
         # and set it accordingly in the model runner and the kv cache size
         n_blocks_avail = self._get_num_blocks_available()

vllm_spyre/v1/worker/spyre_worker.py

@@ -29,7 +29,6 @@
 import vllm_spyre.perf_metrics as perf_metrics
 from vllm_spyre.model_executor.model_loader import spyre_setup
 from vllm_spyre.platform import SpyrePlatform
-from vllm_spyre.v1.worker.spyre_input_batch import SamplingInputBatch
 from vllm_spyre.v1.worker.spyre_model_runner import (
     ContinuousBatchingSpyreModelRunner, SpyrePoolingModelRunner,
     StaticBatchingSpyreModelRunner, SupportedTask)
@@ -110,6 +109,9 @@ def compile_or_warm_up_model(self) -> None:
                 prompt_len, num_decode_tokens, batch_size)
             self._warmup_spyre_fixed_size(prompt_len, num_decode_tokens,
                                           self.restricted_tokens, batch_size)
+
+        self.model_runner.complete_warmup()
+
         all_warmup_end_t = time.time()
         all_warmup_total_t = all_warmup_end_t - all_warmup_start_t
         self.perf_metrics.log("total warmup time", all_warmup_total_t)
@@ -119,7 +121,6 @@ def compile_or_warm_up_model(self) -> None:
             "[WARMUP] All %d prompt/decode/batchsize-shape "
             "combinations finished in %.3fs", num_shape_combinations,
             all_warmup_total_t)
-        self.model_runner.complete_warmup()
 
     def check_health(self) -> None:
         """Basic health check (override for device-specific checks)."""
@@ -339,18 +340,6 @@ def _warmup_spyre_dynamic_size(self, special_token_ids):
         prompt_len = 42
         num_decode_tokens = 2
 
-        # Fix for batch size 1: set input batch to fit 2 requests for warmup
-        if model_runner.vllm_config.scheduler_config.max_num_seqs == 1:
-            model_runner.input_batch = SamplingInputBatch(
-                max_num_reqs=2,
-                max_model_len=model_runner.vllm_config.model_config.
-                max_model_len,
-                device=model_runner.device,
-                pin_memory=model_runner.pin_memory,
-                vocab_size=model_runner.vllm_config.model_config.
-                get_vocab_size(),
-            )
-
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
             0, len(valid_token_ids_tensor), (batch_size + 1, prompt_len))]
@@ -398,20 +387,7 @@ def _warmup_spyre_dynamic_size(self, special_token_ids):
         self.execute_model(scheduler_output)
         self._cleanup_model_runner(request=[add_dummy_request])
 
-        # Fix for batch size 1: reset input batch to fit max_num_seqs requests
-        if model_runner.vllm_config.scheduler_config.max_num_seqs == 1:
-            model_runner.input_batch = SamplingInputBatch(
-                max_num_reqs=model_runner.vllm_config.scheduler_config.
-                max_num_seqs,
-                max_model_len=model_runner.vllm_config.model_config.
-                max_model_len,
-                device=model_runner.device,
-                pin_memory=model_runner.pin_memory,
-                vocab_size=model_runner.vllm_config.model_config.
-                get_vocab_size(),
-            )
-
-        model_runner.finish_warmup()
+        model_runner.complete_warmup()
 
         warmup_end_t = time.time()
         warmup_total_t = warmup_end_t - warmup_start_t