Adds MEAN pooling type

Signed-off-by: Flavia Beo <[email protected]>

vllm/model_executor/layers/pooler.py

@@ -14,7 +14,6 @@ class PoolingType(IntEnum):
     ALL = 1
     CLS = 2
     MEAN = 3
-    MAX = 4
 
 
 class PoolingConfig():
@@ -49,7 +48,7 @@ class Pooler(nn.Module):
     3. Returns structured results as `PoolerOutput`.
 
     Attributes:
-        pooling_type: The type of pooling to use (LAST, ALL, CLS).
+        pooling_type: The type of pooling to use (LAST, ALL, CLS, MEAN).
         normalize: Whether to normalize the pooled data.
     """
 
@@ -83,6 +82,17 @@ def forward(
             for prompt_len in prompt_lens:
                 pooled_data.append(hidden_states[offset:offset + prompt_len])
                 offset += prompt_len
+        elif self.pooling_type == PoolingType.MEAN:
+            # Calculate mean pooling
+            cumsum = torch.cumsum(hidden_states, dim=0)
+            start_indices = torch.cat([
+                torch.tensor([0], device=hidden_states.device),
+                torch.cumsum(prompt_lens[:-1], dim=0)
+            ])
+            end_indices = torch.cumsum(prompt_lens, dim=0)
+            pooled_data = (
+                cumsum[end_indices - 1] - cumsum[start_indices] +
+                hidden_states[start_indices]) / prompt_lens.unsqueeze(1)
         else:
             raise ValueError(f"Invalid pooling type: {self.pooling_type}")