Fully overlap model execution (#134)

Dependent on vllm-project/vllm#23569

---------

Signed-off-by: Tianmu Li <[email protected]>
Co-authored-by: Chendi.Xue <[email protected]>

Author: tianmu-li

tests/full_tests/ci_gsm8k_tests.sh

@@ -140,6 +140,18 @@ if [ $? -ne 0 ]; then
 fi
 echo "Test with granite-8b passed"
 
+# used to check asynchronous scheduling
+echo "Testing GSM8K on ganite-8b with async scheduling"
+echo VLLM_CONTIGUOUS_PA=False VLLM_SKIP_WARMUP=True PT_HPU_LAZY_MODE=1 VLLM_USE_V1=1 ASYNC_SCHEDULING=1 \
+pytest -v -s vllm-gaudi/tests/models/language/generation/test_common.py --model_card_path vllm-gaudi/tests/full_tests/model_cards/granite-8b.yaml
+VLLM_CONTIGUOUS_PA=False VLLM_SKIP_WARMUP=True PT_HPU_LAZY_MODE=1 VLLM_USE_V1=1 ASYNC_SCHEDULING=1 \
+pytest -v -s vllm-gaudi/tests/models/language/generation/test_common.py --model_card_path vllm-gaudi/tests/full_tests/model_cards/granite-8b.yaml
+if [ $? -ne 0 ]; then
+    echo "Error: Test failed for granite-8b + async_scheduling" >&2
+    exit -1
+fi
+echo "Test with granite-8b + async_scheduling passed"
+
 # used to check MLA + MOE
 echo "Testing GSM8K on deepseek v2 lite"
 # deepseek-R1

tests/models/language/generation/test_common.py

@@ -23,9 +23,11 @@ def launch_lm_eval(eval_config):
     enforce_eager = os.environ.get('ENFORCE_EAGER', 'False').lower() in ['true', '1']
     kv_cache_dtype = os.environ.get('KV_CACHE_DTYPE', None)
     task = eval_config.get('tasks', 'gsm8k')
+    async_scheduling = os.environ.get('ASYNC_SCHEDULING', 'False').lower() in ['true', '1']
     model_args = {
         'pretrained': eval_config['model_name'],
         'tensor_parallel_size': tp_size,
+        'async_scheduling': async_scheduling,
         'enforce_eager': enforce_eager,
         'enable_prefix_caching': enable_apc,
         'add_bos_token': True,

vllm_gaudi/v1/worker/hpu_input_batch.py

@@ -215,6 +215,11 @@ def __init__(
         self.sampling_metadata = self._make_sampling_metadata()
         self.pooling_params: dict[str, PoolingParams] = {}
 
+        # Cached reference to the GPU tensor of previously sampled tokens
+        self.prev_sampled_token_ids: Optional[torch.Tensor] = None
+        self.prev_sampled_token_ids_invalid_indices: Optional[set[int]] = None
+        self.prev_req_id_to_index: Optional[dict[str, int]] = None
+
         self.req_type: dict[str, str] = {}
 
     @property