ctx_pp2_gen_pp1_draft

Signed-off-by: Patrice Castonguay <[email protected]>

Author: pcastonguay

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

@@ -814,7 +814,19 @@ void CacheFormatter::unformat(TransferSession& session)
     if (selfConfig.getModelConfig().mNbKvHeadsPerLayer.size() != destConfig.getModelConfig().mNbKvHeadsPerLayer.size())
     {
         TLLM_LOG_WARNING("CacheFormatter::inquireSupport: only support same number of layers");
-        return false;
+        TLLM_LOG_WARNING("self: %d dest %d", selfConfig.getModelConfig().mNbKvHeadsPerLayer.size(),
+            destConfig.getModelConfig().mNbKvHeadsPerLayer.size());
+
+        auto selfTotalLayers = selfConfig.getModelConfig().mNbKvHeadsPerLayer.size()
+            * selfConfig.getParallelConfig().mPipelineParallelism;
+        auto destTotalLayers = destConfig.getModelConfig().mNbKvHeadsPerLayer.size()
+            * destConfig.getParallelConfig().mPipelineParallelism;
+        if (selfTotalLayers != destTotalLayers)
+        {
+            TLLM_LOG_WARNING("CacheFormatter::inquireSupport: incompatible total layer counts: self=%d, dest=%d",
+                static_cast<int>(selfTotalLayers), static_cast<int>(destTotalLayers));
+            return false;
+        }
     }
     int selfNumLayers = selfConfig.getModelConfig().mNbKvHeadsPerLayer.size();
     int selfPPSize = selfConfig.getParallelConfig().mPipelineParallelism;

examples/disaggregated/disagg_config.yaml

@@ -1,13 +1,13 @@
 hostname: localhost
 port: 8000
-model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+model: /home/scratch.trt_llm_data/llm-models/llama-3.1-model/Meta-Llama-3.1-8B
 free_gpu_memory_fraction: 0.25
 backend: "pytorch"
 disable_overlap_scheduler: True
 context_servers:
   num_instances: 1
   tensor_parallel_size: 1
-  pipeline_parallel_size: 1
+  pipeline_parallel_size: 2
   kv_cache_config:
     free_gpu_memory_fraction: 0.2
   cache_transceiver_config:

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -122,6 +122,7 @@ class BatchState:
 @dataclasses.dataclass
 class BatchStatePP(BatchState):
     microbatch_id: int = -1
+    scheduled_ctx_reqs: list[LlmRequest] = None
 
 
 class PyExecutor:
@@ -656,6 +657,9 @@ def _executor_loop_pp(self):
                 if self.should_stop_processing:
                     break
 
+                if self.kv_cache_transceiver:
+                    self._check_disagg_gen_transfer_status()
+
                 if self.enable_iter_perf_stats:
                     iter_stats = self._get_init_iter_stats(
                         len(new_requests),
@@ -664,9 +668,28 @@ def _executor_loop_pp(self):
 
                 self._pad_attention_dp_dummy_request()
 
-                scheduled_batch, _, _ = self._schedule()
+                scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule(
+                )
+
+                if self.kv_cache_transceiver:
+
+                    # For requests that are fitting disagg gen init, also prepare resources for KV cache manager
+                    self._prepare_disagg_gen_init(
+                        fitting_disagg_gen_init_requests)
+
+                    if num_fitting_reqs == 0 and not fitting_disagg_gen_init_requests:
+                        logger.warning(
+                            "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache"
+                        )
+                        self.kv_cache_transceiver.check_context_transfer_status(
+                            1)
+                else:
+                    assert scheduled_batch.batch_size > 0, (
+                        "fail to schedule any pending request, "
+                        "probably run out of resource.")
 
                 self.num_scheduled_requests = scheduled_batch.batch_size
+
                 logger.debug(
                     f'has {len(self.active_requests)} active_request, '
                     f'scheduled {len(scheduled_batch.context_requests)} context requests and '
@@ -688,8 +711,28 @@ def _executor_loop_pp(self):
                     self.micro_batches[microbatch_id] = None
                 else:
                     self._add_inflight_ids(scheduled_batch)
+
+                    if self.kv_cache_transceiver:
+                        # For generation requests which have completed KV cache transfer
+                        self._prepare_disagg_gen_transmission_complete(
+                            scheduled_batch)
+
                     self.resource_manager.prepare_resources(scheduled_batch)
 
+                    # 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
+                    # for disaggregated serving. For non-disaggregated serving,
+                    # the generation requests always have batch_idx.
+                    scheduled_batch.generation_requests = sorted(  # stable sort
+                        scheduled_batch.generation_requests,
+                        key=lambda req: int(req.py_batch_idx is not None),
+                    )
+
+                    if self.kv_cache_transceiver:
+                        # Return the first token to the client
+                        self._handle_first_token_response(scheduled_batch)
+
                     # Stage 1: Async forward (all ranks) and decoding pass (last rank only)
                     if not self.dist.is_last_pp_rank:
                         sample_state = self._forward_step_inter_pp(
@@ -720,6 +763,7 @@ def _executor_loop_pp(self):
                         iter_start_time=iter_start_time,
                         iter_stats=iter_stats,
                         microbatch_id=microbatch_id,
+                        scheduled_ctx_reqs=scheduled_batch.context_requests,
                     )
 
                     self.micro_batches[microbatch_id] = batch_state
@@ -784,6 +828,12 @@ def _executor_loop_pp(self):
                 if previous_batch is not None:
                     with torch.cuda.nvtx.range("_handle_previous_batch_pp"):
                         self._update_requests(previous_batch.sample_state)
+
+                        if self.kv_cache_transceiver and previous_batch.scheduled_ctx_reqs:
+                            ctx_transmission_reqs = self._send_disagg_ctx_cache(
+                                previous_batch.scheduled_ctx_reqs
+                            ) if self.kv_cache_transceiver else []
+
                         self._handle_canceled_requests()
                         finished_requests = self._handle_responses()
                         previous_scheduled_batch = previous_batch.sample_state.scheduled_requests
@@ -792,6 +842,9 @@ def _executor_loop_pp(self):
                         self._remove_inflight_ids(previous_scheduled_batch)
                     self.micro_batches[prev_microbatch_id] = None
 
+                if self.kv_cache_transceiver and self.ctx_in_transmission_requests:
+                    self._terminate_ctx_finished_requests()
+
                 # march forward in microbatch slots
                 microbatch_id = (microbatch_id + 1) % self.num_micro_batches