non_blocking torch.zeros in _build_encoder_only_attn_metadata

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

Author: noooop

vllm/model_executor/layers/pooler.py

@@ -617,8 +617,10 @@ def extract_states(
         self,
         hidden_states: Union[torch.Tensor, list[torch.Tensor]],
         pooling_metadata: PoolingMetadata,
+        num_scheduled_tokens: torch.Tensor,
     ) -> Union[list[torch.Tensor], torch.Tensor]:
-        pooled_data_lst = self.pooling(hidden_states, pooling_metadata)
+        pooled_data_lst = self.pooling(hidden_states, pooling_metadata,
+                                       num_scheduled_tokens)
         prompt_token_ids = get_prompt_token_ids(pooling_metadata)
 
         pooled_data = list[torch.Tensor]()
@@ -652,7 +654,8 @@ def forward(
         pooling_metadata: PoolingMetadata,
         num_scheduled_tokens: torch.Tensor,
     ) -> PoolerOutput:
-        pooled_data = self.extract_states(hidden_states, pooling_metadata)
+        pooled_data = self.extract_states(hidden_states, pooling_metadata,
+                                          num_scheduled_tokens)
         pooled_data = self.head(pooled_data, pooling_metadata)
         return build_output(pooled_data)
 

vllm/v1/worker/gpu_model_runner.py

@@ -3324,10 +3324,14 @@ def _build_encoder_only_attn_metadata(
 
         dummy_block_table = torch.zeros((num_reqs, 1),
                                         dtype=torch.int32,
-                                        device=self.device)
+                                        pin_memory=self.pin_memory,
+                                        device="cpu").to(self.device,
+                                                         non_blocking=True)
         dummy_slot_mapping = torch.zeros((total_num_scheduled_tokens, ),
                                          dtype=torch.int32,
-                                         device=self.device)
+                                         pin_memory=self.pin_memory,
+                                         device="cpu").to(self.device,
+                                                          non_blocking=True)
 
         group_metadata = dict[str, tuple[CommonAttentionMetadata, Any]]()