update another test for overlap

Signed-off-by: Izzy Putterman <[email protected]>

Author: IzzyPutterman

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -1203,12 +1203,12 @@ def _executor_loop_overlap(self):
                     self._process_previous_batch()
                     self.previous_batch: Optional[BatchState] = None
 
-                if self.enable_iter_perf_stats:
-                    iter_stats.inflight_batching_stats.num_ctx_tokens = self.model_engine.iter_states[
-                        'num_ctx_tokens']
-
                 if scheduled_batch.batch_size > 0:
 
+                    if self.enable_iter_perf_stats:
+                        iter_stats.inflight_batching_stats.num_ctx_tokens = self.model_engine.iter_states[
+                            'num_ctx_tokens']
+
                     self.previous_batch = BatchState(
                         sample_state=sample_state,
                         iter_start_time=iter_start_time,

tests/integration/defs/llmapi/test_llm_api_connector.py

@@ -87,8 +87,7 @@ def test_connector_simple(enforce_single_worker, model_with_connector,
     assert len(scheduler.update_state_after_alloc.call_args.args[1]) == 1
 
     # With the overlap scheduler, we generate one extra token.
-    assert scheduler.build_connector_meta.call_count == NUM_TOKENS + int(
-        use_overlap_scheduler)
+    assert scheduler.build_connector_meta.call_count == NUM_TOKENS
 
     # We should have a single `SchedulerOutput` per forward pass.
     for i, call in enumerate(scheduler.build_connector_meta.call_args_list):
@@ -108,8 +107,7 @@ def test_connector_simple(enforce_single_worker, model_with_connector,
             assert len(scheduler_output.cached_requests[0].new_tokens) == 1
 
     # We call `start_load_kv` once at the beginning of each forward pass.
-    assert worker.start_load_kv.call_count == NUM_TOKENS + int(
-        use_overlap_scheduler)
+    assert worker.start_load_kv.call_count == NUM_TOKENS
 
     # Only called once when the request is received.
     assert scheduler.get_num_new_matched_tokens.call_count == 1
@@ -118,19 +116,16 @@ def test_connector_simple(enforce_single_worker, model_with_connector,
                      for call in worker.wait_for_layer_load.call_args_list) + 1
 
     # Called num_layers * num_forward_passes times.
-    assert worker.wait_for_layer_load.call_count == num_layers * (
-        NUM_TOKENS + int(use_overlap_scheduler))
-    assert worker.save_kv_layer.call_count == num_layers * (
-        NUM_TOKENS + int(use_overlap_scheduler))
+    assert worker.wait_for_layer_load.call_count == num_layers * (NUM_TOKENS)
+    assert worker.save_kv_layer.call_count == num_layers * (NUM_TOKENS)
 
     for i, call in enumerate(worker.wait_for_layer_load.call_args_list):
         assert call.args[0] == i % num_layers
 
     for i, call in enumerate(worker.save_kv_layer.call_args_list):
         assert call.args[0] == i % num_layers
 
-    assert worker.wait_for_save.call_count == NUM_TOKENS + int(
-        use_overlap_scheduler)
+    assert worker.wait_for_save.call_count == NUM_TOKENS
 
     assert scheduler.request_finished.call_count == 1
 
@@ -238,8 +233,7 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector,
             NUM_INPUT_TOKENS / BLOCK_SIZE)
 
     # Additional token when using the overlap scheduler.
-    assert scheduler.build_connector_meta.call_count == NUM_TOKENS + int(
-        use_overlap_scheduler)
+    assert scheduler.build_connector_meta.call_count == NUM_TOKENS
 
     for i, call in enumerate(scheduler.build_connector_meta.call_args_list):
         sched_output = call.args[0]

tests/unittest/llmapi/test_llm_pytorch.py

@@ -174,7 +174,10 @@ def test_llm_reward_model():
 
 
 def test_llm_perf_metrics():
-    llm = LLM(model=llama_model_path, kv_cache_config=global_kvcache_config)
+    disable_overlap_scheduler = False
+    llm = LLM(model=llama_model_path,
+              kv_cache_config=global_kvcache_config,
+              disable_overlap_scheduler=disable_overlap_scheduler)
     sampling_params = SamplingParams(max_tokens=10, return_perf_metrics=True)
     outputs = llm.generate(prompts, sampling_params)
     assert outputs[0].outputs[0].request_perf_metrics is not None
@@ -194,7 +197,8 @@ def test_llm_perf_metrics():
     assert kv_cache_metrics.kv_cache_hit_rate == 0
 
     assert perf_metrics.first_iter is not None
-    assert perf_metrics.iter - perf_metrics.first_iter == sampling_params.max_tokens - 1
+    assert perf_metrics.iter - perf_metrics.first_iter == sampling_params.max_tokens - (
+        1 if disable_overlap_scheduler else 2)
     assert perf_metrics.last_iter == perf_metrics.iter