[P/D][Feature]添加pd分离p实例对tbo的支持

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.