[Core] generate from input embeds

tests/conftest.py

@@ -27,8 +27,9 @@
 from vllm.distributed import (cleanup_dist_env_and_memory,
                               init_distributed_environment,
                               initialize_model_parallel)
-from vllm.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
-                         to_enc_dec_tuple_list, zip_enc_dec_prompts)
+from vllm.inputs import (EmbedsPrompt, ExplicitEncoderDecoderPrompt,
+                         TextPrompt, to_enc_dec_tuple_list,
+                         zip_enc_dec_prompts)
 from vllm.logger import init_logger
 from vllm.outputs import RequestOutput
 from vllm.platforms import current_platform
@@ -654,21 +655,26 @@ def __init__(
 
     def get_inputs(
         self,
-        prompts: List[str],
+        prompts_or_prompt_embeds: Union[List[str], List[torch.Tensor]],
         images: Optional[PromptImageInput] = None,
         videos: Optional[PromptVideoInput] = None,
         audios: Optional[PromptAudioInput] = None,
-    ) -> List[TextPrompt]:
+    ) -> List[Union[TextPrompt, EmbedsPrompt]]:
         if images is not None:
-            assert len(prompts) == len(images)
+            assert len(prompts_or_prompt_embeds) == len(images)
 
         if videos is not None:
-            assert len(prompts) == len(videos)
+            assert len(prompts_or_prompt_embeds) == len(videos)
 
         if audios is not None:
-            assert len(prompts) == len(audios)
+            assert len(prompts_or_prompt_embeds) == len(audios)
+
+        inputs = [
+            EmbedsPrompt(prompt_embeds=prompt) if isinstance(
+                prompt, torch.Tensor) else TextPrompt(prompt=prompt)
+            for prompt in prompts_or_prompt_embeds
+        ]
 
-        inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
         if images is not None:
             for i, image in enumerate(images):
                 if image is not None:
@@ -696,13 +702,13 @@ def classify(self, prompts: List[str]) -> List[str]:
 
     def generate(
         self,
-        prompts: List[str],
+        prompts_or_prompt_embeds: Union[List[str], List[torch.Tensor]],
         sampling_params: SamplingParams,
         images: Optional[PromptImageInput] = None,
         videos: Optional[PromptVideoInput] = None,
         audios: Optional[PromptAudioInput] = None,
     ) -> List[Tuple[List[List[int]], List[str]]]:
-        inputs = self.get_inputs(prompts,
+        inputs = self.get_inputs(prompts_or_prompt_embeds,
                                  images=images,
                                  videos=videos,
                                  audios=audios)
@@ -720,7 +726,7 @@ def generate(
                 output_str = sample.text
                 output_ids = list(sample.token_ids)
                 req_sample_output_ids.append(prompt_ids + output_ids)
-                req_sample_output_strs.append(prompt_str + output_str)
+                req_sample_output_strs.append((prompt_str or "") + output_str)
             outputs.append((req_sample_output_ids, req_sample_output_strs))
         return outputs
 
@@ -785,14 +791,14 @@ def generate_encoder_decoder_w_logprobs(
 
     def generate_greedy(
         self,
-        prompts: List[str],
+        prompts_or_prompt_embeds: Union[List[str], List[torch.Tensor]],
         max_tokens: int,
         images: Optional[PromptImageInput] = None,
         videos: Optional[PromptVideoInput] = None,
         audios: Optional[PromptAudioInput] = None,
     ) -> List[Tuple[List[int], str]]:
         greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
-        outputs = self.generate(prompts,
+        outputs = self.generate(prompts_or_prompt_embeds,
                                 greedy_params,
                                 images=images,
                                 videos=videos,

tests/models/decoder_only/language/test_models.py

@@ -54,9 +54,22 @@ def test_models(
         hf_outputs = hf_model.generate_greedy_logprobs_limit(
             example_prompts, max_tokens, num_logprobs)
 
+        prompt_embeds = []
+        prompt_token_ids = []
+        for prompt in example_prompts:
+            token_ids = hf_model.tokenizer(prompt,
+                                           return_tensors="pt").input_ids.to(
+                                               hf_model.model.device)
+            prompt_token_ids.append(token_ids)
+            prompt_embeds.append(
+                hf_model.model.get_input_embeddings()(token_ids).squeeze(0))
+
     with vllm_runner(model, dtype=dtype) as vllm_model:
         vllm_outputs = vllm_model.generate_greedy_logprobs(
             example_prompts, max_tokens, num_logprobs)
+        vllm_outputs_from_embeds = vllm_model.generate_greedy_logprobs(
+            prompt_embeds, max_tokens, num_logprobs)
+
         # This test is for verifying whether the model's extra_repr
         # can be printed correctly.
         print(vllm_model.model.llm_engine.model_executor.driver_worker.
@@ -68,3 +81,10 @@ def test_models(
         name_0="hf",
         name_1="vllm",
     )
+
+    check_logprobs_close(
+        outputs_0_lst=vllm_outputs,
+        outputs_1_lst=vllm_outputs_from_embeds,
+        name_0="vllm",
+        name_1="vllm_from_embeds",
+    )

tests/worker/test_model_runner.py

@@ -1,3 +1,4 @@
+import random
 from typing import List
 
 import pytest
@@ -22,8 +23,9 @@ def _create_model_runner(model: str, *args, **kwargs) -> ModelRunner:
     return model_runner
 
 
-@pytest.mark.parametrize("batch_size", list(range(1, 257)))
-def test_prepare_prompt(batch_size):
+@pytest.mark.parametrize("batch_size", list(range(1, 257, 3)))
+@pytest.mark.parametrize("prompt_embeds_ratio", (0.0, 0.5, 1.0))
+def test_prepare_prompt(batch_size, prompt_embeds_ratio):
     model_runner = _create_model_runner(
         "facebook/opt-125m",
         max_num_batched_tokens=100000,
@@ -34,11 +36,19 @@ def test_prepare_prompt(batch_size):
     seq_lens: List[int] = []
     seq_group_metadata_list: List[SequenceGroupMetadata] = []
     block_tables = {0: [1]}
+    input_embeds_len = 0
     for i in range(batch_size):
         # make sure all tokens fit into one block
         seq_len = i % (model_runner.block_size - 1) + 1
         seq_lens.append(seq_len)
-        seq_data = SequenceData.from_seqs(range(seq_len))
+        if random.random() < prompt_embeds_ratio:
+            seq_data = SequenceData.from_seqs(
+                range(seq_len),
+                prompt_embeds=torch.rand(seq_len, 10),
+            )
+            input_embeds_len += seq_len
+        else:
+            seq_data = SequenceData.from_seqs(range(seq_len))
         seq_group_metadata = SequenceGroupMetadata(
             request_id=f"test_{i}",
             is_prompt=True,
@@ -59,6 +69,8 @@ def test_prepare_prompt(batch_size):
         seq_group_metadata_list)
     input_tokens = model_input.input_tokens
     input_positions = model_input.input_positions
+    input_embeds = model_input.input_embeds
+    input_embeds_masks = model_input.input_embeds_masks
     attn_metadata = model_input.attn_metadata
     return_seq_lens = model_input.seq_lens
     slot_mapping = attn_metadata.slot_mapping
@@ -112,7 +124,12 @@ def test_prepare_prompt(batch_size):
 
     assert len(input_tokens) == sum(seq_lens)
     assert len(input_positions) == sum(seq_lens)
-    torch.testing.assert_close(input_tokens, input_positions)
+    assert len(input_embeds_masks) == sum(seq_lens)
+    if input_embeds_len == 0:
+        torch.testing.assert_close(input_tokens, input_positions)
+        assert input_embeds is None
+    else:
+        assert len(input_embeds) == input_embeds_len
 
     sampling_metadata = SamplingMetadata.prepare(
         seq_group_metadata_list,
@@ -136,8 +153,9 @@ def test_prepare_prompt(batch_size):
     torch.testing.assert_close(actual, expected)
 
 
-@pytest.mark.parametrize("batch_size", list(range(1, 257)))
-def test_prepare_decode_cuda_graph(batch_size):
+@pytest.mark.parametrize("batch_size", list(range(1, 257, 3)))
+@pytest.mark.parametrize("prompt_embeds_ratio", (0.0, 0.5, 1.0))
+def test_prepare_decode_cuda_graph(batch_size, prompt_embeds_ratio):
     model_runner = _create_model_runner(
         "facebook/opt-125m",
         seed=0,
@@ -151,11 +169,19 @@ def test_prepare_decode_cuda_graph(batch_size):
     context_lens: List[int] = []
     seq_group_metadata_list: List[SequenceGroupMetadata] = []
     # Assume each seq group finishes prefill.
+    input_embeds_len = 0
     for i in range(batch_size):
         # make sure all tokens fit into one block
         context_len = i % (model_runner.block_size - 1) + 1
         context_lens.append(context_len)
-        seq_data = SequenceData.from_seqs(range(context_len))
+        if random.random() < prompt_embeds_ratio:
+            seq_data = SequenceData.from_seqs(
+                [],
+                prompt_embeds=torch.rand(context_len, 10),
+            )
+            input_embeds_len += context_len
+        else:
+            seq_data = SequenceData.from_seqs(range(context_len))
         seq_data.update_num_computed_tokens(context_len)
         # Append one token ID since prefill is finished.
         seq_data.append_token_id(1, 0)
@@ -171,9 +197,13 @@ def test_prepare_decode_cuda_graph(batch_size):
 
     model_input = model_runner._prepare_model_input_tensors(
         seq_group_metadata_list)
-    input_tokens, input_positions, attn_metadata, slot_mapping = (
-        model_input.input_tokens, model_input.input_positions,
-        model_input.attn_metadata, model_input.attn_metadata.slot_mapping)
+    (input_tokens, input_positions, input_embeds, input_embeds_masks,
+     attn_metadata,
+     slot_mapping) = (model_input.input_tokens, model_input.input_positions,
+                      model_input.input_embeds, model_input.input_embeds_masks,
+                      model_input.attn_metadata,
+                      model_input.attn_metadata.slot_mapping)
+
     assert len(slot_mapping) == len(input_tokens)
 
     expected_bs = _get_graph_batch_size(len(seq_group_metadata_list))
@@ -226,6 +256,8 @@ def test_prepare_decode_cuda_graph(batch_size):
     assert len(input_tokens) == expected_bs
     assert len(input_positions) == expected_bs
     torch.allclose(input_tokens, input_positions)
+    assert input_embeds is None
+    assert input_embeds_masks is None
 
     # Verify Sampling
     expected_selected_token_indices = []
@@ -256,14 +288,19 @@ def test_empty_seq_group():
     seq_group_metadata_list: List[SequenceGroupMetadata] = []
     model_input = model_runner._prepare_model_input_tensors(
         seq_group_metadata_list)
-    input_tokens, input_positions, attn_metadata = (
-        model_input.input_tokens,
-        model_input.input_positions,
-        model_input.attn_metadata,
-    )
+    (input_tokens, input_positions, input_embeds, input_embeds_masks,
+     attn_metadata) = (
+         model_input.input_tokens,
+         model_input.input_positions,
+         model_input.input_embeds,
+         model_input.input_embeds_masks,
+         model_input.attn_metadata,
+     )
     assert input_tokens is None
     assert input_positions is None
     assert attn_metadata is None
+    assert input_embeds is None
+    assert input_embeds_masks is None
 
     model_input = model_runner._prepare_model_input_tensors(
         seq_group_metadata_list)
@@ -289,9 +326,11 @@ def distributed_init():
     ensure_model_parallel_initialized(1, 1)
 
 
-@pytest.mark.parametrize("batch_size", list(range(2, 128)))
+@pytest.mark.parametrize("batch_size", list(range(2, 128, 3)))
 @pytest.mark.parametrize("enforce_eager", [True, False])
-def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
+@pytest.mark.parametrize('prompt_embeds_ratio', [0.0, 0.5, 1.0])
+def test_hybrid_batches(batch_size, enforce_eager, prompt_embeds_ratio,
+                        distributed_init):
     model_runner = _create_model_runner(
         "facebook/opt-125m",
         seed=0,
@@ -310,11 +349,19 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
     block_tables = {0: [1]}
     prefill_batch_size = batch_size // 2
     decode_batch_size = batch_size - prefill_batch_size
+    input_embeds_len = 0
     for i in range(prefill_batch_size):
         # make sure all tokens fit into one block
         seq_len = i % (model_runner.block_size - 1) + 1
         seq_lens.append(seq_len)
-        seq_data = SequenceData.from_seqs(range(seq_len))
+        if random.random() < prompt_embeds_ratio:
+            seq_data = SequenceData.from_seqs(
+                [],
+                prompt_embeds=torch.rand(seq_len, 10),
+            )
+            input_embeds_len += seq_len
+        else:
+            seq_data = SequenceData.from_seqs(range(seq_len))
         seq_group_metadata = SequenceGroupMetadata(
             request_id=f"test_{i}",
             is_prompt=True,
@@ -330,7 +377,13 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
     for i in range(prefill_batch_size, batch_size):
         # make sure all tokens fit into one block
         context_len = i % (model_runner.block_size - 1) + 1
-        seq_data = SequenceData.from_seqs(range(context_len))
+        if random.random() < prompt_embeds_ratio:
+            seq_data = SequenceData.from_seqs(
+                [],
+                prompt_embeds=torch.rand(context_len, 10),
+            ),
+        else:
+            seq_data = SequenceData.from_seqs(range(context_len))
         seq_data.append_token_id(1, 0)
         seq_data.update_num_computed_tokens(context_len)
         seq_group_metadata = SequenceGroupMetadata(
@@ -345,11 +398,14 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
         decode_metadata_list.append(seq_group_metadata)
 
     model_input = model_runner.prepare_model_input(seq_group_metadata_list)
-    (input_tokens, input_positions, attn_metadata) = (
-        model_input.input_tokens,
-        model_input.input_positions,
-        model_input.attn_metadata,
-    )
+    (input_tokens, input_positions, attn_metadata, input_embeds,
+     input_embeds_masks) = (
+         model_input.input_tokens,
+         model_input.input_positions,
+         model_input.attn_metadata,
+         model_input.input_embeds,
+         model_input.input_embeds_masks,
+     )
 
     prefill_meta_actual = attn_metadata.prefill_metadata
     decode_meta_actual = attn_metadata.decode_metadata
@@ -359,6 +415,11 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
     assert attn_metadata.num_prefills == prefill_batch_size
     assert attn_metadata.num_decode_tokens == decode_batch_size
     assert attn_metadata.num_prefill_tokens == sum(seq_lens)
+    assert len(input_embeds_masks) == sum(seq_lens)
+    if input_embeds_len == 0:
+        assert input_embeds is None
+    else:
+        assert len(input_embeds) == input_embeds_len
 
     # Verify attn metadata is consistent. We don't need to test individual
     # values here because they are tested above.