Use as many batch operations as possible

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

Author: noooop

vllm/model_executor/layers/pooler.py

@@ -469,6 +469,10 @@ def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
 
         pooling_params = get_pooling_params(pooling_metadata)
 
+        if isinstance(pooled_data, list):
+            pooled_data = torch.stack(pooled_data)
+        # pooled_data shape: [batchsize, embedding_dimension]
+
         # for matryoshka representation
         dimensions_list = [
             pooling_param.dimensions for pooling_param in pooling_params
@@ -659,6 +663,10 @@ def forward(
     ) -> PoolerOutput:
         pooled_data = self.pooling(hidden_states, pooling_metadata)
 
+        if isinstance(pooled_data, list):
+            pooled_data = torch.stack(pooled_data)
+        # pooled_data shape: [batchsize, hidden_size]
+
         if self.classifier is not None:
             # apply classifier once on the full batch if possible
             if isinstance(pooled_data, torch.Tensor):