fix pd send async perfomance

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

@@ -294,7 +294,7 @@ class P2pNcclConnector(KVConnectorBase_V1):
                 slot_mapping = request.slot_mapping
                 if request.slot_mapping_device is None:
                     request.slot_mapping_device = \
-                        request.slot_mapping.pin_memory().to(device=kv_layer.device, non_blocking=True) 
+                        request.slot_mapping.pin_memory().to(device=kv_layer.device, non_blocking=True)
                 slot_mapping = request.slot_mapping_device
                 kv_cache = extract_kv_from_layer(kv_layer, slot_mapping)
                 tbo_evt = torch.cuda.Event(enable_timing=False)

vllm/v1/worker/gpu_model_runner.py

@@ -107,6 +107,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.speculative_config = vllm_config.speculative_config
         self.prompt_adapter_config = vllm_config.prompt_adapter_config
         self.observability_config = vllm_config.observability_config
+        if envs.VLLM_P2P_ASYNC:
+            self.p2p_event = torch.cuda.Event(enable_timing=False)
+            self.p2p_stream = torch.cuda.Stream()
 
         from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
         set_cpu_offload_max_bytes(
@@ -1295,7 +1298,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         scheduler_output: "SchedulerOutput",
         intermediate_tensors: Optional[IntermediateTensors] = None,
     ) -> Union[ModelRunnerOutput, IntermediateTensors]:
-
+        profile.StartTracer()
         self._update_states(scheduler_output)
         if not scheduler_output.total_num_scheduled_tokens:
             if not has_kv_transfer_group():
@@ -1381,6 +1384,34 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                                              num_tokens_across_dp, input_ids, positions,
                                              inputs_embeds, scheduler_output, intermediate_tensors,
                                              skip_cuda_graphs)
+            
+        elif envs.VLLM_P2P_ASYNC:
+            self.p2p_event.record()
+            current_stream = torch.cuda.current_stream()
+            with torch.cuda.stream(self.p2p_stream):
+                self.p2p_stream.wait_event(self.p2p_event)
+                with set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_input_tokens,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    skip_cuda_graphs=skip_cuda_graphs,
+                ):
+                    self.maybe_setup_kv_connector(scheduler_output)
+
+                    model_output = self.model(
+                        input_ids=input_ids,
+                        positions=positions,
+                        intermediate_tensors=intermediate_tensors,
+                        inputs_embeds=inputs_embeds,
+                    )
+                    
+                    self.maybe_wait_for_kv_save()
+                    finished_sending, finished_recving = (
+                        self.get_finished_kv_transfers(scheduler_output))
+                self.p2p_event.record()
+            current_stream.wait_event(self.p2p_event)
+                                    
         else:
             # Run the model.
             # Use persistent buffers for CUDA graphs.