feat: add vision pooling support for jina embeddings v4

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

Author: Sigrid Jin (Sionic AI)

vllm/config.py

@@ -3256,7 +3256,7 @@ def get_limit_per_prompt(self, modality: str) -> int:
 @config
 @dataclass
 class PoolerConfig:
-    """Configuration for the pooler."""
+    """Controls the behavior of output pooling in pooling models."""
 
     pooling_type: Optional[Literal["last", "all", "cls", "step", "mean",
                                    "vision"]] = None

vllm/model_executor/layers/pooler.py

@@ -625,56 +625,6 @@ 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
@@ -817,10 +767,12 @@ def forward(
 
         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()
+            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()),
@@ -832,41 +784,18 @@ def forward(
                                  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)
+            if HAS_TRITON:
+                mean_pool_with_position_kernel[grid](hidden_states, output,
+                                                     seq_start, seq_len,
+                                                     self.config.hidden_size,
+                                                     start_pos, end_pos + 1)
+            else:
+                # Fallback to PyTorch implementation if Triton is not available
+                vision_tokens_range = hidden_states[seq_start + start_pos : seq_start + end_pos + 1]
+                output = vision_tokens_range.mean(dim=0)
 
             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)

vllm/model_executor/models/jina_embeddings_v4.py

@@ -6,7 +6,7 @@
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.pooler import Pooler, PoolingTask
+from vllm.model_executor.layers.pooler import Pooler, PoolingTask, VisionPooler
 # yapf: disable
 from vllm.model_executor.pooling_metadata import (
     PoolingMetadata as V0PoolingMetadata)