reduce cuda sync

Signed-off-by: wang.yuqi <noooop@126.com>

Author: noooop

vllm/model_executor/layers/pooler.py

@@ -196,6 +196,13 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):
 
 def build_output(
     all_data: Union[torch.Tensor, list[torch.Tensor]], ) -> PoolerOutput:
+    # Pooling models D2H occurs only here
+    if isinstance(all_data, list):
+        all_data = [d.to("cpu", non_blocking=True) for d in all_data]
+    else:
+        all_data = all_data.to("cpu", non_blocking=True)
+    torch.cuda.synchronize()
+
     all_outputs = [PoolingSequenceGroupOutput(data) for data in all_data]
     return PoolerOutput(outputs=all_outputs)
 
@@ -706,6 +713,7 @@ def forward(
             hidden_states_lst = hidden_states
         else:
             prompt_lens = get_prompt_lens(hidden_states, pooling_metadata)
+
             hidden_states_lst = list(hidden_states.split(prompt_lens.tolist()))
 
         outputs = list[PoolingSequenceGroupOutput]()

vllm/v1/worker/gpu_input_batch.py

@@ -713,7 +713,7 @@ def pooling_metadata(self) -> PoolingMetadata:
 
         return PoolingMetadata(
             prompt_lens=torch.from_numpy(
-                self.num_prompt_tokens[:self.num_reqs]).to(self.device),
+                self.num_prompt_tokens[:self.num_reqs]),
             prompt_token_ids=self.sampling_metadata.prompt_token_ids,
             pooling_params=pooling_params,
         )

vllm/v1/worker/gpu_model_runner.py

@@ -1489,7 +1489,7 @@ def _pool(
                 raw_pooler_output, seq_lens, pooling_metadata.prompt_lens):
 
             if seq_len == prompt_len:
-                pooler_output.append(raw_output.data.cpu())
+                pooler_output.append(raw_output.data)
             else:
                 pooler_output.append(None)