add support for the overlap scheduler + little refactoring

Author: jthomson04

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -276,7 +276,9 @@ def estimate_max_tokens(self, py_executor: PyExecutor) -> None:
         executor_config.kv_cache_config.max_tokens = kv_cache_max_tokens
 
     def _create_kv_cache_manager(
-            self, model_engine: PyTorchModelEngine) -> KVCacheManager:
+            self,
+            model_engine: PyTorchModelEngine,
+            for_estimation: bool = False) -> KVCacheManager:
         executor_config = self._executor_config
         mapping = self._mapping
         assert model_engine.model.model_config.is_generation, "Only construct KV cache for generation models."
@@ -317,15 +319,16 @@ def _create_kv_cache_manager(
                 dtype=kv_cache_dtype,
                 spec_config=spec_config,
                 max_beam_width=executor_config.max_beam_width,
-                kv_connector_manager=self._kv_connector_manager,
+                kv_connector_manager=self._kv_connector_manager
+                if not for_estimation else None,
             )
         elif is_nemotron_hybrid(config):
             if executor_config.max_beam_width > 1:
                 raise ValueError(
                     "MambaHybridCacheManager + beam search is not supported yet."
                 )
 
-            if self._kv_connector_manager is not None:
+            if not for_estimation and self._kv_connector_manager is not None:
                 raise ValueError(
                     "Connector manager is not supported for MambaHybridCacheManager."
                 )
@@ -387,25 +390,29 @@ def _create_kv_cache_manager(
                 max_num_tokens=executor_config.max_num_tokens,
                 model_config=binding_model_config,
                 max_beam_width=executor_config.max_beam_width,
-                kv_connector_manager=self._kv_connector_manager,
+                kv_connector_manager=self._kv_connector_manager
+                if not for_estimation else None,
             )
         # KVCacheManager (Non-draft) modifies the max_seq_len field, update it to executor_config
         if model_engine.kv_cache_manager_key == ResourceManagerType.KV_CACHE_MANAGER:
             executor_config.max_seq_len = kv_cache_manager.max_seq_len
 
         return kv_cache_manager
 
-    def build_managers(self, resources: Dict) -> None:
+    def build_managers(self,
+                       resources: Dict,
+                       for_estimation: bool = False) -> None:
         """Construct KV caches for model and draft model (if applicable)."""
-        kv_cache_manager = self._create_kv_cache_manager(self._model_engine)
+        kv_cache_manager = self._create_kv_cache_manager(
+            self._model_engine, for_estimation)
 
-        if self._kv_connector_manager is not None and self._draft_model_engine is not None:
+        if not for_estimation and self._kv_connector_manager is not None and self._draft_model_engine is not None:
             raise ValueError(
                 "Connector manager is not supported for draft model.")
 
         draft_kv_cache_manager = self._create_kv_cache_manager(
-            self._draft_model_engine
-        ) if self._draft_model_engine is not None else None
+            self._draft_model_engine,
+            for_estimation) if self._draft_model_engine is not None else None
 
         resources[ResourceManagerType.KV_CACHE_MANAGER] = kv_cache_manager
         resources[

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -274,11 +274,6 @@ def _maybe_init_kv_connector_manager(self):
                     "KV Cache Connector is not supported with pipeline parallelism."
                 )
 
-            if not self.disable_overlap_scheduler:
-                raise NotImplementedError(
-                    "KV Cache Connector is not supported with overlap scheduler."
-                )
-
             # TODO: This does NOT support pipeline parallel.
             layer_kv_tensors = {
                 layer_idx: self.kv_cache_manager.get_buffers(layer_idx)
@@ -948,6 +943,19 @@ def _execute_guided_decoder(self, scheduled_batch, logits):
             self.guided_decoder.build(scheduled_batch)
             self.guided_decoder.execute(scheduled_batch, logits)
 
+    def _execute_kv_connector(self, scheduled_batch):
+        if self.kv_connector_manager:
+            self.kv_connector_manager.take_scheduled_requests_pending_load(
+                scheduled_batch)
+            self.kv_connector_manager.handle_metadata()
+            self.kv_connector_manager.worker.start_load_kv()
+
+    def _terminate_async_save_requests(self):
+        if self.kv_connector_manager:
+            reqs_to_terminate = self.kv_connector_manager.get_finished()
+            for req in reqs_to_terminate:
+                self.resource_manager.free_resources(req)
+
     def _executor_loop(self):
         torch.cuda.set_device(self.device_id)
         with self._profiler() as profile_step:
@@ -976,14 +984,9 @@ def _executor_loop(self):
 
                         # Return the first token to the client
                         self._handle_first_token_response(scheduled_batch)
-                    scheduled_batch.is_warmup = self.is_warmup
                     self.resource_manager.prepare_resources(scheduled_batch)
 
-                    if self.kv_connector_manager:
-                        self.kv_connector_manager.take_scheduled_requests_pending_load(
-                            scheduled_batch)
-                        self.kv_connector_manager.handle_metadata()
-                        self.kv_connector_manager.worker.start_load_kv()
+                    self._execute_kv_connector(scheduled_batch)
 
                 if scheduled_batch.batch_size > 0 or (
                         self.enable_attention_dp and self.dist.tp_size > 1):
@@ -1017,10 +1020,8 @@ def _executor_loop(self):
 
                 if self.kv_cache_transceiver and self.ctx_in_transmission_requests:
                     self._terminate_ctx_finished_requests()
-                elif self.kv_connector_manager:
-                    reqs_to_terminate = self.kv_connector_manager.get_finished()
-                    for req in reqs_to_terminate:
-                        self.resource_manager.free_resources(req)
+
+                self._terminate_async_save_requests()
 
                 if self.enable_iter_perf_stats:
                     iter_stats.inflight_batching_stats.num_ctx_tokens = self.model_engine.iter_states[
@@ -1086,9 +1087,12 @@ def _executor_loop_overlap(self):
                         # For generation requests which have completed KV cache transfer
                         self._prepare_disagg_gen_transmission_complete(
                             scheduled_batch)
-
                     self.resource_manager.prepare_resources(scheduled_batch)
 
+                    self._execute_kv_connector(scheduled_batch)
+
+                if scheduled_batch.batch_size > 0:
+
                     # The generation requests that are do not have batch_idx,
                     # needs to be in front of the batch due to the assumptions
                     # made in model_engine.py::_forward_step. This is only important
@@ -1141,6 +1145,8 @@ def _executor_loop_overlap(self):
                 if self.kv_cache_transceiver and self.ctx_in_transmission_requests:
                     self._terminate_ctx_finished_requests()
 
+                self._terminate_async_save_requests()
+
     def _process_previous_batch(self):
         if self.kv_cache_transceiver and self.previous_batch.ctx_transmission_reqs:
             for req in self.previous_batch.ctx_transmission_reqs:

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -392,7 +392,7 @@ def create_py_executor(
         with mem_monitor.observe_creation_stage(
                 _ExecutorCreationStage.INIT_KV_CACHE
                 if estimating_kv_cache else _ExecutorCreationStage.KV_CACHE):
-            kv_cache_creator.build_managers(resources)
+            kv_cache_creator.build_managers(resources, estimating_kv_cache)
 
     # Resource managers for speculative decoding
     # For user-specified drafters, use extra_resource_managers in PyTorchBackend config
@@ -443,7 +443,7 @@ def create_py_executor(
             # create_kv_cache_manager above, which caps executor_config.max_seq_len. Restoring
             # the original value before creating the final KV cache.
             executor_config.max_seq_len = max_seq_len
-            kv_cache_creator.build_managers(resources)
+            kv_cache_creator.build_managers(resources, False)
 
             for eng in [model_engine, draft_model_engine]:
                 if eng is None:

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -377,8 +377,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                         req.py_request_id,
                         seq_len + (len(req.query_id) if self.mapping.cp_rank
                                    == self.mapping.cp_size - 1 else 0),
-                        req_beam_width, req, self.kv_connector_manager
-                        if not scheduled_batch.is_warmup else None)
+                        req_beam_width, req, self.kv_connector_manager)
             else:
                 # TODO(jthomson04): This is begging for a mega refactor, and can likely be significantly simplified.
                 # In add sequence, the connector API's get_num_new_matched_tokens is called.
@@ -388,10 +387,9 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                 # When that happens, the request will go through this same code path, but with is_kv_cache_connector_async_onboard set to True.
                 # Because of this, we need to filter this case out to avoid adding the same sequence twice.
                 if req.is_first_context_chunk and not req.is_kv_cache_connector_async_onboard:
-                    self.impl.add_sequence(
-                        req.py_request_id, req.prompt_len, req_beam_width, req,
-                        self.kv_connector_manager
-                        if not scheduled_batch.is_warmup else None)
+                    self.impl.add_sequence(req.py_request_id, req.prompt_len,
+                                           req_beam_width, req,
+                                           self.kv_connector_manager)
                     for _ in range(self.num_extra_kv_tokens):
                         self.impl.add_token(req.py_request_id)
                     for _ in range(get_draft_token_length(req)):

tensorrt_llm/_torch/pyexecutor/scheduler.py

@@ -21,7 +21,6 @@ def __init__(self):
         self.context_requests: RequestList = []
         self.generation_requests: RequestList = []
         self.paused_requests: RequestList = []
-        self.is_warmup: bool = False
 
     @property
     def is_generation_only(self) -> bool:

tests/integration/defs/llmapi/test_llm_api_connector.py

@@ -71,15 +71,16 @@ def model_fn(*args, **kwargs):
 
 
 @pytest.mark.threadleak(enabled=False)
-def test_connector_simple(model_with_connector):
+@pytest.mark.parametrize("use_overlap_scheduler", [True, False])
+def test_connector_simple(model_with_connector, use_overlap_scheduler):
     NUM_TOKENS = 8
 
     model_fn, scheduler, worker = model_with_connector
 
     model = model_fn(
         model="Qwen/Qwen2-0.5B",
         backend="pytorch",
-        disable_overlap_scheduler=True,
+        disable_overlap_scheduler=not use_overlap_scheduler,
         cuda_graph_config=None,
         kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1))
 
@@ -93,7 +94,9 @@ def test_connector_simple(model_with_connector):
 
     model.generate(["Hello, world"], sampling_params)
 
-    assert scheduler.build_connector_meta.call_count == NUM_TOKENS
+    # With the overlap scheduler, we generate one extra token.
+    assert scheduler.build_connector_meta.call_count == NUM_TOKENS + int(
+        use_overlap_scheduler)
 
     # We should have a single `SchedulerOutput` per forward pass.
     for i, call in enumerate(scheduler.build_connector_meta.call_args_list):
@@ -105,7 +108,8 @@ def test_connector_simple(model_with_connector):
             assert len(scheduler_output.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
+    assert worker.start_load_kv.call_count == NUM_TOKENS + int(
+        use_overlap_scheduler)
 
     # Only called once when the request is received.
     assert scheduler.get_num_new_matched_tokens.call_count == 1
@@ -114,31 +118,36 @@ def test_connector_simple(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
-    assert worker.save_kv_layer.call_count == num_layers * NUM_TOKENS
+    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))
 
     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
+    assert worker.wait_for_save.call_count == NUM_TOKENS + int(
+        use_overlap_scheduler)
 
     assert scheduler.request_finished.call_count == 1
-    assert worker.get_finished.call_count == NUM_TOKENS
+    assert worker.get_finished.call_count == NUM_TOKENS + int(
+        use_overlap_scheduler)
 
 
 @pytest.mark.threadleak(enabled=False)
-def test_connector_async_onboard(model_with_connector):
+@pytest.mark.parametrize("use_overlap_scheduler", [True, False])
+def test_connector_async_onboard(model_with_connector, use_overlap_scheduler):
     NUM_TOKENS = 8
 
     model_fn, scheduler, worker = model_with_connector
 
     model = model_fn(
         model="Qwen/Qwen2-0.5B",
         backend="pytorch",
-        disable_overlap_scheduler=True,
+        disable_overlap_scheduler=not use_overlap_scheduler,
         cuda_graph_config=None,
         kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1))
 
@@ -153,23 +162,25 @@ def test_connector_async_onboard(model_with_connector):
         "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua."
     ], SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True))
 
-    # Once for the initial poll, then once for each token.
-    assert worker.get_finished.call_count == NUM_TOKENS + 1
+    # Once for the initial poll, then once for each token. One extra token when using the overlap scheduler.
+    assert worker.get_finished.call_count == NUM_TOKENS + 1 + int(
+        use_overlap_scheduler)
 
     # In the first iteration, there should be a single request id provided.
     assert len(worker.get_finished.call_args_list[0].args[1]) == 1
 
 
 @pytest.mark.threadleak(enabled=False)
-def test_connector_async_save(model_with_connector):
+@pytest.mark.parametrize("use_overlap_scheduler", [True, False])
+def test_connector_async_save(model_with_connector, use_overlap_scheduler):
     NUM_TOKENS = 8
 
     model_fn, scheduler, worker = model_with_connector
 
     model = model_fn(
         model="Qwen/Qwen2-0.5B",
         backend="pytorch",
-        disable_overlap_scheduler=True,
+        disable_overlap_scheduler=not use_overlap_scheduler,
         cuda_graph_config=None,
         kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1))
 
@@ -188,12 +199,13 @@ def test_connector_async_save(model_with_connector):
 
     assert scheduler.request_finished.call_count == 1
 
-    # On the last call to get_finished, we should be providing the async saving request.
-    assert worker.get_finished.call_count == NUM_TOKENS
+    # On the last call to get_finished, we should be providing the async saving request. One extra token when using the overlap scheduler.
+    assert worker.get_finished.call_count == NUM_TOKENS + int(
+        use_overlap_scheduler)
 
-    for i in range(NUM_TOKENS):
-        args = worker.get_finished.call_args_list[i].args
-        if i != NUM_TOKENS - 1:
+    for i, call in enumerate(worker.get_finished.call_args_list):
+        args = call.args
+        if i != len(worker.get_finished.call_args_list) - 1:
             assert args == ([], [])
         else:
             assert len(args[0]) == 1
@@ -202,7 +214,9 @@ def test_connector_async_save(model_with_connector):
 
 
 @pytest.mark.threadleak(enabled=False)
-def test_connector_scheduler_output(model_with_connector):
+@pytest.mark.parametrize("use_overlap_scheduler", [True, False])
+def test_connector_scheduler_output(model_with_connector,
+                                    use_overlap_scheduler):
     NUM_INPUT_TOKENS = 48
     NUM_TOKENS = 32
     BLOCK_SIZE = 32
@@ -212,7 +226,7 @@ def test_connector_scheduler_output(model_with_connector):
     model = model_fn(
         model="Qwen/Qwen2-0.5B",
         backend="pytorch",
-        disable_overlap_scheduler=True,
+        disable_overlap_scheduler=not use_overlap_scheduler,
         cuda_graph_config=None,
         kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1))
 
@@ -226,7 +240,9 @@ def test_connector_scheduler_output(model_with_connector):
 
     model.generate([0] * NUM_INPUT_TOKENS, sampling_params)
 
-    assert scheduler.build_connector_meta.call_count == NUM_TOKENS
+    # Additional token when using the overlap scheduler.
+    assert scheduler.build_connector_meta.call_count == NUM_TOKENS + int(
+        use_overlap_scheduler)
 
     for i, call in enumerate(scheduler.build_connector_meta.call_args_list):
         sched_output = call.args[0]
@@ -241,7 +257,8 @@ def test_connector_scheduler_output(model_with_connector):
         else:
             assert len(request.new_tokens) == 1
 
-            if request.computed_position % BLOCK_SIZE == 0:
+            if (request.computed_position +
+                    int(use_overlap_scheduler)) % BLOCK_SIZE == 0:
                 assert len(request.new_block_ids) == 1
             else:
                 assert request.new_block_ids == []
@@ -257,7 +274,9 @@ def test_connector_scheduler_output(model_with_connector):
 
 
 @pytest.mark.threadleak(enabled=False)
-def test_connector_scheduler_output_chunked_context(model_with_connector):
+@pytest.mark.parametrize("use_overlap_scheduler", [True, False])
+def test_connector_scheduler_output_chunked_context(model_with_connector,
+                                                    use_overlap_scheduler):
     model_fn, scheduler, worker = model_with_connector
 
     CHUNK_SIZE = 128
@@ -266,7 +285,7 @@ def test_connector_scheduler_output_chunked_context(model_with_connector):
     model = model_fn(
         model="Qwen/Qwen2-0.5B",
         backend="pytorch",
-        disable_overlap_scheduler=True,
+        disable_overlap_scheduler=not use_overlap_scheduler,
         cuda_graph_config=None,
         kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.1),
         enable_chunked_prefill=True,