Merge branch 'v0.9.2-dev-pd_tbo' into 'v0.9.2-dev'

V0.9.2 dev pd tbo

See merge request dcutoolkit/deeplearing/vllm!204

vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py

@@ -215,9 +215,18 @@ class P2pNcclConnector(KVConnectorBase_V1):
 
                 inject_kv_into_layer(kv_cache_layer, kv_cache,
                                      request.slot_mapping, request.request_id)
-                tensor = self.p2p_nccl_engine.recv_store.pop(request.request_id + "#" + layer_name, None)
-                if tensor is not None:
-                    del tensor
+                tensor_id = request.request_id + "#" + layer_name
+                if tensor_id in self.p2p_nccl_engine.recv_store:
+                    tensor = self.p2p_nccl_engine.recv_store.pop(tensor_id, None)
+                    self.p2p_nccl_engine.send_request_id_to_tensor_ids.pop(
+                                request.request_id, None)
+                    self.p2p_nccl_engine.recv_request_id_to_tensor_ids.pop(
+                                request.request_id, None)
+                    addr = 0
+                    if isinstance(tensor, tuple):
+                        addr, _, _ = tensor
+                        self.p2p_nccl_engine.pool.free(addr)
+
 
     def wait_for_layer_load(self, layer_name: str) -> None:
         """Blocking until the KV for a specific layer is loaded into vLLM's

vllm/two_batch_overlap/v1/model_input_split_v1.py

@@ -287,18 +287,29 @@ def tbo_split_and_execute_model(
         attn_metadata_left = prepare_tbo_atten_metadata(runner, input_split.scheduler_output_left, input_split.req_ids_left, 0)
         attn_metadata_right = prepare_tbo_atten_metadata(runner, input_split.scheduler_output_right, input_split.req_ids_right, input_split.req_num_left)
 
-        model_output = tbo_model_executable_v1(
-            runner,
-            attn_metadata_left, 
-            attn_metadata_right, 
-            num_input_tokens_left,
-            num_input_tokens_right,
-            num_tokens_across_dp,
-            input_ids,
-            positions,
-            intermediate_tensors,
-            inputs_embeds) 
-        finished_sending, finished_recving = None, None
+        with set_forward_context(attn_metadata,
+                                runner.vllm_config,
+                                num_tokens=num_input_tokens,
+                                num_tokens_across_dp=num_tokens_across_dp,
+                                skip_cuda_graphs=True):
+            runner.maybe_setup_kv_connector(scheduler_output)
+            
+            model_output = tbo_model_executable_v1(
+                runner,
+                attn_metadata_left, 
+                attn_metadata_right, 
+                num_input_tokens_left,
+                num_input_tokens_right,
+                num_tokens_across_dp,
+                input_ids,
+                positions,
+                intermediate_tensors,
+                inputs_embeds) 
+            
+            runner.maybe_wait_for_kv_save()
+            finished_sending, finished_recving = (
+                runner.get_finished_kv_transfers(scheduler_output))
+        #finished_sending, finished_recving = None, None
     else:
         # Run the decoder.
         # Use persistent buffers for CUDA graphs.