fix main

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

Author: yannicks1

vllm_spyre/v1/worker/spyre_worker.py

@@ -29,7 +29,7 @@
 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 InputBatch
+from vllm_spyre.v1.worker.spyre_input_batch import SamplingInputBatch
 from vllm_spyre.v1.worker.spyre_model_runner import (
     ContinuousBatchingSpyreModelRunner, SpyrePoolingModelRunner,
     StaticBatchingSpyreModelRunner, SupportedTask)
@@ -329,7 +329,7 @@ def _warmup_spyre_dynamic_size(self, special_token_ids):
 
         # 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 = InputBatch(
+            model_runner.input_batch = SamplingInputBatch(
                 max_num_reqs=2,
                 max_model_len=model_runner.vllm_config.model_config.
                 max_model_len,
@@ -388,7 +388,7 @@ def _warmup_spyre_dynamic_size(self, special_token_ids):
 
         # 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 = InputBatch(
+            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.