Support embedding models in V1

vllm/config.py

@@ -492,6 +492,9 @@ def _init_pooler_config(
     ) -> Optional["PoolerConfig"]:
 
         if self.runner_type == "pooling":
+            logger.warning("CUDA graph is not supported for pooling yet, "
+                           "fallback to the eager mode.")
+            self.enforce_eager = True
             user_config = override_pooler_config or PoolerConfig()
 
             base_config = get_pooling_config(self.model, self.revision)

vllm/engine/arg_utils.py

@@ -1439,12 +1439,6 @@ def _is_v1_supported_oracle(self, model_config: ModelConfig) -> bool:
                 recommend_to_remove=False)
             return False
 
-        # No Embedding Models so far.
-        if model_config.task not in ["generate"]:
-            _raise_or_fallback(feature_name=f"--task {model_config.task}",
-                               recommend_to_remove=False)
-            return False
-
         # No Mamba or Encoder-Decoder so far.
         if not model_config.is_v1_compatible:
             _raise_or_fallback(feature_name=model_config.architectures,

vllm/model_executor/layers/pooler.py

@@ -10,11 +10,15 @@
 from typing_extensions import assert_never
 
 from vllm.config import PoolerConfig
-from vllm.model_executor.pooling_metadata import (PoolingMetadata,
-                                                  PoolingTensors)
+from vllm.model_executor.pooling_metadata import (  # noqa: E501
+    PoolingMetadata as V0PoolingMetadata)
+from vllm.model_executor.pooling_metadata import PoolingTensors
 from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
 from vllm.transformers_utils.config import (
     get_cross_encoder_activation_function)
+from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata
+
+PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]
 
 
 class PoolingType(IntEnum):
@@ -78,6 +82,8 @@ def get_prompt_lens(
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
     ) -> torch.Tensor:
+        if isinstance(pooling_metadata, V1PoolingMetadata):
+            return pooling_metadata.prompt_lens
         return PoolingTensors.from_pooling_metadata(
             pooling_metadata, hidden_states.device).prompt_lens
 
@@ -181,12 +187,27 @@ def __init__(
         self.step_tag_id = step_tag_id
         self.returned_token_ids = returned_token_ids
 
+    def get_prompt_token_ids(
+        self,
+        pooling_metadata: PoolingMetadata,
+    ) -> List[torch.Tensor]:
+        if isinstance(pooling_metadata, V1PoolingMetadata):
+            return [
+                pooling_metadata.prompt_token_ids[i, :num]
+                for i, num in enumerate(pooling_metadata.prompt_lens)
+            ]
+        return [
+            seq_data_i.prompt_token_ids
+            for seq_data_i in pooling_metadata.seq_data.values()
+        ]
+
     def extract_states(
         self,
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
     ) -> Union[list[torch.Tensor], torch.Tensor]:
         prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
+        prompt_token_ids = self.get_prompt_token_ids(pooling_metadata)
 
         returned_token_ids = self.returned_token_ids
         if returned_token_ids is not None and len(returned_token_ids) > 0:
@@ -196,12 +217,11 @@ def extract_states(
 
         offset = 0
         pooled_data = list[torch.Tensor]()
-        for prompt_len, seq_data_i in zip(prompt_lens,
-                                          pooling_metadata.seq_data.values()):
+        for i, prompt_len in enumerate(prompt_lens):
             pooled_data_i = hidden_states[offset:offset + prompt_len]
             if step_tag_id is not None:
-                token_ids = torch.tensor(seq_data_i.prompt_token_ids)
-                pooled_data_i = pooled_data_i[token_ids == step_tag_id]
+                pooled_data_i = pooled_data_i[prompt_token_ids[i] ==
+                                              step_tag_id]
 
             offset += prompt_len
             pooled_data.append(pooled_data_i)
@@ -287,15 +307,24 @@ def __init__(
         self.default_activation_function = \
             get_cross_encoder_activation_function(config)
 
+    def get_prompt_lens(
+        self,
+        hidden_states: torch.Tensor,
+        pooling_metadata: PoolingMetadata,
+    ) -> torch.Tensor:
+        if isinstance(pooling_metadata, V1PoolingMetadata):
+            return pooling_metadata.prompt_lens
+        return PoolingTensors.from_pooling_metadata(
+            pooling_metadata, hidden_states.device).prompt_lens
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
     ) -> PoolerOutput:
         """Pools sentence pair scores from the hidden_states."""
 
-        prompt_lens = PoolingTensors.from_pooling_metadata(
-            pooling_metadata, hidden_states.device).prompt_lens
+        prompt_lens = self.get_prompt_lens(hidden_states, pooling_metadata)
 
         offset = 0
         pooled_data_lst = []

vllm/model_executor/models/bert.py

@@ -18,6 +18,7 @@
 from vllm.model_executor.layers.pooler import (CrossEncodingPooler, Pooler,
                                                PoolingType)
 from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.sampler import get_sampler
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -26,7 +27,7 @@
 from vllm.transformers_utils.config import (
     get_cross_encoder_activation_function)
 
-from .interfaces import SupportsCrossEncoding, SupportsV0Only
+from .interfaces import SupportsCrossEncoding
 from .utils import WeightsMapper, maybe_prefix
 
 
@@ -115,7 +116,7 @@ def forward(
         self,
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
-        for layer in self.layer:
+        for i, layer in enumerate(self.layer):
             hidden_states = layer(hidden_states)
         return hidden_states
 
@@ -323,6 +324,7 @@ def __init__(self,
                  add_pooling_layer: bool = False):
         super().__init__()
         config = vllm_config.model_config.hf_config
+        self.config = config
         self.embeddings = embedding_class(config)
         self.encoder = BertEncoder(vllm_config=vllm_config,
                                    prefix=f"{prefix}.encoder")
@@ -340,12 +342,16 @@ def forward(
             hidden_states = inputs_embeds
         else:
             attn_metadata = get_forward_context().attn_metadata
-            assert hasattr(attn_metadata, "seq_lens_tensor")
-            hidden_states = self.embeddings(
-                input_ids=input_ids,
-                seq_lens=attn_metadata.seq_lens_tensor,
-                position_ids=position_ids,
-                token_type_ids=token_type_ids)
+            seq_lens = None
+            if attn_metadata is not None:  # Can be None during warmup
+                seq_lens = getattr(attn_metadata, "seq_lens_tensor",
+                                   attn_metadata.seq_lens)
+                assert seq_lens is not None
+            hidden_states = self.embeddings(input_ids=input_ids,
+                                            seq_lens=seq_lens,
+                                            position_ids=position_ids,
+                                            token_type_ids=token_type_ids)
+
         return self.encoder(hidden_states)
 
     def load_weights(self, weights: Iterable[Tuple[str,
@@ -385,7 +391,7 @@ def load_weights(self, weights: Iterable[Tuple[str,
         return loaded_params
 
 
-class BertEmbeddingModel(nn.Module, SupportsV0Only):
+class BertEmbeddingModel(nn.Module):
     """A model that uses Bert to provide embedding functionalities.
 
    This class encapsulates the BertModel and provides an interface for
@@ -403,6 +409,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.model = self._build_model(vllm_config=vllm_config,
                                        prefix=maybe_prefix(prefix, "model"))
         self._pooler = self._build_pooler(pooler_config)
+        # TODO: Remove test scaffolding after pooling is implemented
+        self.sampler = get_sampler()
 
     def forward(
         self,
@@ -411,10 +419,11 @@ def forward(
         intermediate_tensors: Optional[IntermediateTensors] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        return self.model(input_ids=input_ids,
-                          position_ids=positions,
-                          inputs_embeds=inputs_embeds,
-                          intermediate_tensors=intermediate_tensors)
+        hidden_states = self.model(input_ids=input_ids,
+                                   position_ids=positions,
+                                   inputs_embeds=inputs_embeds,
+                                   intermediate_tensors=intermediate_tensors)
+        return hidden_states
 
     def pooler(
         self,

vllm/model_executor/models/roberta.py

@@ -80,31 +80,33 @@ def forward(
         input_shape = input_ids.size()
         inputs_embeds = self.word_embeddings(input_ids)
 
-        # Replace position ids because in RoBERTa models
-        # they have to start at padding_idx + 1 and ignore
-        # existing padding tokens
-        # References:
-        # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L133
-        # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L1669
-        pos_list = []
-        token_list = []
-        offset = 0
-        for seq_len in seq_lens:
-            pos_list.append(position_ids[offset:offset + seq_len])
-            token_list.append(input_ids[offset:offset + seq_len])
-            offset += seq_len
-
-        new_pos_list = []
-        for positions, tokens in zip(pos_list, token_list):
-            # Verify assumption that incoming position are
-            # always a sequence from 0 to N.
-            expected_pos = torch.arange(positions.size()[0],
-                                        dtype=torch.long,
-                                        device=inputs_embeds.device)
-            assert torch.equal(positions, expected_pos)
-            new_pos_list.append(
-                create_position_ids_from_input_ids(tokens, self.padding_idx))
-        position_ids = torch.cat(new_pos_list)
+        if seq_lens is not None:  # Can be None during warmup
+            # Replace position ids because in RoBERTa models
+            # they have to start at padding_idx + 1 and ignore
+            # existing padding tokens
+            # References:
+            # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L133
+            # - https://github.com/huggingface/transformers/blob/a3d69a8994d673899608a7c17fbf4f953f50474e/src/transformers/models/roberta/modeling_roberta.py#L1669
+            pos_list = []
+            token_list = []
+            offset = 0
+            for seq_len in seq_lens:
+                pos_list.append(position_ids[offset:offset + seq_len])
+                token_list.append(input_ids[offset:offset + seq_len])
+                offset += seq_len
+
+            new_pos_list = []
+            for positions, tokens in zip(pos_list, token_list):
+                # Verify assumption that incoming position are
+                # always a sequence from 0 to N.
+                expected_pos = torch.arange(positions.size()[0],
+                                            dtype=torch.long,
+                                            device=inputs_embeds.device)
+                assert torch.equal(positions, expected_pos)
+                new_pos_list.append(
+                    create_position_ids_from_input_ids(tokens,
+                                                       self.padding_idx))
+            position_ids = torch.cat(new_pos_list)
 
         # Position embeddings.
         position_embeddings = self.position_embeddings(position_ids)

vllm/outputs.py

@@ -344,10 +344,11 @@ class PoolingRequestOutput(Generic[_O]):
         finished (bool): A flag indicating whether the pooling is completed.
     """
 
-    def __init__(self, request_id: str, outputs: _O,
+    def __init__(self, request_id: str, outputs: _O, prompt: Optional[str],
                  prompt_token_ids: list[int], finished: bool):
         self.request_id = request_id
         self.prompt_token_ids = prompt_token_ids
+        self.prompt = prompt
         self.finished = finished
         self.outputs = outputs
 
@@ -359,9 +360,10 @@ def from_seq_group(seq_group: SequenceGroup) -> "PoolingRequestOutput":
         data = pooled_data.to(dtype=torch.float32, device="cpu")
         output = PoolingOutput(data)
         prompt_token_ids = seq_group.prompt_token_ids
+        prompt = seq_group.prompt
         finished = seq_group.is_finished()
 
-        return PoolingRequestOutput(seq_group.request_id, output,
+        return PoolingRequestOutput(seq_group.request_id, output, prompt,
                                     prompt_token_ids, finished)
 
     def __repr__(self):
@@ -426,6 +428,7 @@ def from_base(request_output: PoolingRequestOutput):
         return EmbeddingRequestOutput(
             request_id=request_output.request_id,
             outputs=EmbeddingOutput.from_base(request_output.outputs),
+            prompt=request_output.prompt,
             prompt_token_ids=request_output.prompt_token_ids,
             finished=request_output.finished,
         )
@@ -464,6 +467,7 @@ def from_base(request_output: PoolingRequestOutput):
         return ClassificationRequestOutput(
             request_id=request_output.request_id,
             outputs=ClassificationOutput.from_base(request_output.outputs),
+            prompt=request_output.prompt,
             prompt_token_ids=request_output.prompt_token_ids,
             finished=request_output.finished,
         )
@@ -503,6 +507,7 @@ def from_base(request_output: PoolingRequestOutput):
         return ScoringRequestOutput(
             request_id=request_output.request_id,
             outputs=ScoringOutput.from_base(request_output.outputs),
+            prompt=request_output.prompt,
             prompt_token_ids=request_output.prompt_token_ids,
             finished=request_output.finished,
         )

vllm/v1/attention/backends/flash_attn.py

@@ -198,11 +198,13 @@ def __init__(
                 f"Supported head sizes are: {support_head_sizes}. "
                 "Set VLLM_USE_V1=0 to use another attention backend.")
 
-        if attn_type != AttentionType.DECODER:
-            raise NotImplementedError("Encoder self-attention and "
-                                      "encoder/decoder cross-attention "
-                                      "are not implemented for "
+        if attn_type not in [
+                AttentionType.DECODER, AttentionType.ENCODER_ONLY
+        ]:
+            raise NotImplementedError("Encoder/decoder cross-attention "
+                                      "is not implemented for "
                                       "FlashAttentionImpl")
+        self.attn_type = attn_type
         self.vllm_flash_attn_version = get_flash_attn_version()
         if is_quantized_kv_cache(self.kv_cache_dtype) \
             and not flash_attn_supports_fp8():
@@ -265,8 +267,7 @@ def forward(
             layer._k_scale,
             layer._v_scale,
         )
-        descale_shape = (attn_metadata.query_start_loc.shape[0] - 1,
-                         key.shape[1])
+
         if self.kv_cache_dtype.startswith("fp8"):
             key_cache = key_cache.view(torch.float8_e4m3fn)
             value_cache = value_cache.view(torch.float8_e4m3fn)
@@ -280,6 +281,9 @@ def forward(
         # Compute attention and update output up to `num_actual_tokens`.
         if not attn_metadata.use_cascade:
             # Regular attention (common case).
+
+            descale_shape = (attn_metadata.query_start_loc.shape[0] - 1,
+                             key.shape[1])
             flash_attn_varlen_func(
                 q=query[:num_actual_tokens],
                 k=key_cache,
@@ -290,7 +294,7 @@ def forward(
                 seqused_k=attn_metadata.seq_lens,
                 max_seqlen_k=attn_metadata.max_seq_len,
                 softmax_scale=self.scale,
-                causal=True,
+                causal=_get_causal_option(self.attn_type),
                 alibi_slopes=self.alibi_slopes,
                 window_size=self.sliding_window,
                 block_table=attn_metadata.block_table,
@@ -483,3 +487,21 @@ def cascade_attention(
     # Merge prefix and suffix outputs, and store the result in output.
     merge_attn_states(output, prefix_output, prefix_lse, suffix_output,
                       suffix_lse)
+
+
+def _get_causal_option(attn_type: str) -> bool:
+    """
+    Determine whether the given attention type is suitable for causal 
+    attention mechanisms.
+
+    Args:
+        attn_type (AttentionType): The type of attention being evaluated
+
+    Returns:
+        bool: Returns `True` if the attention type is suitable for causal 
+        attention (i.e., not encoder, encoder-only, or encoder-decoder), 
+        otherwise returns `False`.
+    """
+    return not (attn_type == AttentionType.ENCODER
+                or attn_type == AttentionType.ENCODER_ONLY
+                or attn_type == AttentionType.ENCODER_DECODER)

vllm/v1/core/kv_cache_manager.py

@@ -126,7 +126,8 @@ def get_computed_blocks(
             self.req_to_block_hashes[request.request_id] = block_hashes
 
         self.prefix_cache_stats.requests += 1
-        if request.sampling_params.prompt_logprobs is None:
+        if request.sampling_params and \
+            request.sampling_params.prompt_logprobs is None:
             if len(block_hashes) * self.block_size == request.num_tokens:
                 # When prompt length is divisible by the block size and all
                 # blocks are cached, we need to recompute the last token. This