fix: introducing dedicated VisionPooler class

Signed-off-by: Sigrid Jin (Sionic AI) <sigrid@sionic.ai>

Author: Sigrid Jin (Sionic AI)

tests/models/pooling/test_jina_embeddings_v4.py

@@ -342,3 +342,40 @@ def test_vision_only_pooling(self, model):
         # embeddings should be very similar despite different text
         similarity = torch.dot(emb1, emb2).item()
         assert similarity > 0.99  # Should be nearly identical
+
+
+class TestVisionPooler:
+    """Test the VisionPooler class."""
+
+    def test_vision_pooler(self):
+        """Test that the VisionPooler correctly pools vision tokens."""
+        from vllm.config import ModelConfig
+        from vllm.model_executor.layers.pooler import VisionPooler
+        from vllm.pooling_params import PoolingParams
+        from vllm.v1.pool.metadata import PoolingMetadata
+
+        model_config = ModelConfig(model_name, task="embed")
+        model_config.hf_config.vision_start_token_id = VISION_START_TOKEN_ID
+        model_config.hf_config.vision_end_token_id = VISION_END_TOKEN_ID
+        model_config.hidden_size = 4
+
+        pooler = VisionPooler(model_config)
+
+        hidden_states = torch.randn(10, 4)
+        prompt_token_ids = torch.tensor([[
+            1, 2, VISION_START_TOKEN_ID, 4, VISION_END_TOKEN_ID, 6, 7, 8, 9, 10
+        ]])
+        prompt_lens = torch.tensor([10])
+
+        pooling_metadata = PoolingMetadata(prompt_lens=prompt_lens,
+                                           prompt_token_ids=prompt_token_ids,
+                                           pooling_params=[PoolingParams()])
+
+        output = pooler.forward(hidden_states, pooling_metadata)
+
+        vision_tokens = hidden_states[2:5]
+        expected_output = vision_tokens.mean(dim=0)
+
+        assert torch.allclose(output.outputs[0].data,
+                              expected_output,
+                              atol=1e-5)

vllm/config.py

@@ -3256,9 +3256,10 @@ def get_limit_per_prompt(self, modality: str) -> int:
 @config
 @dataclass
 class PoolerConfig:
-    """Controls the behavior of output pooling in pooling models."""
+    """Configuration for the pooler."""
 
-    pooling_type: Optional[str] = None
+    pooling_type: Optional[Literal["last", "all", "cls", "step", "mean",
+                                   "vision"]] = None
     """
     The pooling method of the pooling model. This should be a key in
     [`vllm.model_executor.layers.pooler.PoolingType`][].

vllm/model_executor/layers/pooler.py

@@ -32,6 +32,7 @@ class PoolingType(IntEnum):
     CLS = 2
     STEP = 3
     MEAN = 4
+    VISION = 5
 
 
 @dataclass(frozen=True)
@@ -91,6 +92,8 @@ def from_config_with_defaults(
 
         if pooling_type == PoolingType.STEP:
             return StepPooler.from_config(resolved_config)
+        if pooling_type == PoolingType.VISION:
+            return VisionPooler.from_config(resolved_config)
 
         return SimplePooler.from_config(resolved_config)
 
@@ -622,6 +625,86 @@ def forward(
 ClassifierFn = Callable[[torch.Tensor], torch.Tensor]
 
 
+class VisionPooler(Pooler):
+
+    @classmethod
+    def from_config(cls, model_config: ModelConfig) -> "VisionPooler":
+        return cls(model_config)
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+
+    def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
+        if task == "embed":
+            return PoolingParams(pooling_type="vision",
+                                 logits_processing_needs_token_ids=True)
+        return None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        pooling_metadata: PoolingMetadata,
+    ) -> PoolerOutput:
+        assert isinstance(pooling_metadata, V1PoolingMetadata)
+
+        pooled_outputs = []
+        for i in range(len(pooling_metadata.prompt_lens)):
+            start_pos = (pooling_metadata.prompt_token_ids[i] == self.config.
+                         hf_config.vision_start_token_id).nonzero()[-1].item()
+            end_pos = (pooling_metadata.prompt_token_ids[i] == self.config.
+                       hf_config.vision_end_token_id).nonzero()[-1].item()
+
+            seq_start = torch.cumsum(
+                torch.tensor([0] + pooling_metadata.prompt_lens.tolist()),
+                dim=0)[i]
+            seq_len = pooling_metadata.prompt_lens[i]
+
+            output = torch.empty(self.config.hidden_size,
+                                 device=hidden_states.device,
+                                 dtype=hidden_states.dtype)
+
+            grid = lambda meta: (self.config.hidden_size, )
+            mean_pool_with_position_kernel[grid](hidden_states, output,
+                                                 seq_start, seq_len,
+                                                 self.config.hidden_size,
+                                                 start_pos, end_pos + 1)
+
+            pooled_outputs.append(output)
+
+        return build_output(torch.stack(pooled_outputs))
+
+
+if HAS_TRITON:
+
+    @triton.jit
+    def mean_pool_with_position_kernel(
+        hidden_states_ptr,
+        output_ptr,
+        seq_start,
+        seq_len,
+        hidden_size,
+        pool_start,
+        pool_end,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        """Triton kernel to perform mean pooling over a specified token range."""
+        pid = tl.program_id(0)
+
+        if pid >= hidden_size:
+            return
+
+        accumulator = 0.0
+        for i in range(pool_start, pool_end):
+            hidden_val = tl.load(hidden_states_ptr +
+                                 (seq_start + i) * hidden_size + pid)
+            accumulator += hidden_val
+
+        # Store mean pooled result
+        result = accumulator / (pool_end - pool_start)
+        tl.store(output_ptr + pid, result)
+
+
 class ClassifierPooler(nn.Module):
     """A pooling layer for classification tasks.
 
@@ -709,39 +792,81 @@ def forward(
         return build_output(scores)
 
 
+class VisionPooler(Pooler):
+
+    @classmethod
+    def from_config(cls, model_config: ModelConfig) -> "VisionPooler":
+        return cls(model_config)
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+
+    def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
+        if task == "embed":
+            return PoolingParams(pooling_type="vision",
+                                 logits_processing_needs_token_ids=True)
+        return None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        pooling_metadata: PoolingMetadata,
+    ) -> PoolerOutput:
+        assert isinstance(pooling_metadata, V1PoolingMetadata)
+
+        pooled_outputs = []
+        for i in range(len(pooling_metadata.prompt_lens)):
+            start_pos = (pooling_metadata.prompt_token_ids[i] == self.config.
+                         hf_config.vision_start_token_id).nonzero()[-1].item()
+            end_pos = (pooling_metadata.prompt_token_ids[i] == self.config.
+                       hf_config.vision_end_token_id).nonzero()[-1].item()
+
+            seq_start = torch.cumsum(
+                torch.tensor([0] + pooling_metadata.prompt_lens.tolist()),
+                dim=0)[i]
+            seq_len = pooling_metadata.prompt_lens[i]
+
+            output = torch.empty(self.config.hidden_size,
+                                 device=hidden_states.device,
+                                 dtype=hidden_states.dtype)
+
+            grid = lambda meta: (self.config.hidden_size, )
+            mean_pool_with_position_kernel[grid](hidden_states, output,
+                                                 seq_start, seq_len,
+                                                 self.config.hidden_size,
+                                                 start_pos, end_pos + 1)
+
+            pooled_outputs.append(output)
+
+        return build_output(torch.stack(pooled_outputs))
+
+
 if HAS_TRITON:
 
     @triton.jit
-    def extract_vision_tokens_kernel(
+    def mean_pool_with_position_kernel(
         hidden_states_ptr,
-        token_ids_ptr,
         output_ptr,
         seq_start,
         seq_len,
         hidden_size,
-        vision_start_id: tl.constexpr,
-        vision_end_id: tl.constexpr,
+        pool_start,
+        pool_end,
         BLOCK_SIZE: tl.constexpr,
     ):
-        """Triton kernel to extract and pool vision tokens efficiently."""
+        """Triton kernel to perform mean pooling over a specified token range."""
         pid = tl.program_id(0)
 
         if pid >= hidden_size:
             return
 
-        # Find vision token range
-        vision_count = 0
         accumulator = 0.0
-
-        for i in range(seq_len):
-            token_id = tl.load(token_ids_ptr + seq_start + i)
-            if token_id >= vision_start_id and token_id <= vision_end_id:
-                hidden_val = tl.load(hidden_states_ptr +
-                                     (seq_start + i) * hidden_size + pid)
-                accumulator += hidden_val
-                vision_count += 1
+        for i in range(pool_start, pool_end):
+            hidden_val = tl.load(hidden_states_ptr +
+                                 (seq_start + i) * hidden_size + pid)
+            accumulator += hidden_val
 
         # Store mean pooled result
-        result = accumulator / vision_count if vision_count > 0 else 0.0
-
+        result = accumulator / (pool_end - pool_start)
         tl.store(output_ptr + pid, result)