Support generation from input embedding

tests/conftest.py

@@ -133,9 +133,11 @@ def generate(
         self,
         prompts: List[str],
         sampling_params: SamplingParams,
+        prompt_embeds: List[torch.Tensor] = None,
     ) -> List[Tuple[List[int], str]]:
         req_outputs = self.model.generate(prompts,
-                                          sampling_params=sampling_params)
+                                          sampling_params=sampling_params,
+                                          prompt_embeds=prompt_embeds)
         outputs = []
         for req_output in req_outputs:
             prompt_str = req_output.prompt
@@ -154,9 +156,12 @@ def generate_greedy(
         self,
         prompts: List[str],
         max_tokens: int,
+        prompt_embeds: List[torch.Tensor] = None,
     ) -> List[Tuple[List[int], str]]:
         greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
-        outputs = self.generate(prompts, greedy_params)
+        outputs = self.generate(prompts,
+                                greedy_params,
+                                prompt_embeds=prompt_embeds)
         return [(output_ids[0], output_str[0])
                 for output_ids, output_str in outputs]
 
@@ -165,12 +170,15 @@ def generate_beam_search(
         prompts: List[str],
         beam_width: int,
         max_tokens: int,
+        prompt_embeds: List[torch.Tensor] = None,
     ) -> List[Tuple[List[int], str]]:
         beam_search_params = SamplingParams(n=beam_width,
                                             use_beam_search=True,
                                             temperature=0.0,
                                             max_tokens=max_tokens)
-        outputs = self.generate(prompts, beam_search_params)
+        outputs = self.generate(prompts,
+                                beam_search_params,
+                                prompt_embeds=prompt_embeds)
         return outputs
 
 

tests/models/test_models.py

@@ -43,3 +43,53 @@ def test_models(
             f"Test{i}:\nHF: {hf_output_str!r}\nvLLM: {vllm_output_str!r}")
         assert hf_output_ids == vllm_output_ids, (
             f"Test{i}:\nHF: {hf_output_ids}\nvLLM: {vllm_output_ids}")
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.parametrize("max_tokens", [128])
+def test_models_from_prompt_embeds(
+    hf_runner,
+    vllm_runner,
+    example_prompts,
+    model: str,
+    dtype: str,
+    max_tokens: int,
+) -> None:
+    hf_model = hf_runner(model, dtype=dtype)
+    hf_outputs = hf_model.generate_greedy(example_prompts, max_tokens)
+
+    prompt_embeds = []
+    for prompt in example_prompts:
+        token_ids = hf_model.tokenizer(
+            prompt, return_tensors="pt").input_ids.to("cuda")
+        token_embeds = hf_model.model.get_input_embeddings()(token_ids)
+        prompt_embeds.append(token_embeds[0])
+    del hf_model
+
+    vllm_model = vllm_runner(model, dtype=dtype)
+    vllm_outputs_from_prompts = vllm_model.generate_greedy(example_prompts,
+                                                           max_tokens,
+                                                           prompt_embeds=None)
+    vllm_outputs_from_embeds = vllm_model.generate_greedy(
+        example_prompts, max_tokens, prompt_embeds=prompt_embeds)
+    del vllm_model
+
+    for i in range(len(example_prompts)):
+        prompt = example_prompts[i]
+        hf_output_str = hf_outputs[i][0]
+        vllm_output_str_from_prompts = vllm_outputs_from_prompts[i][0]
+        vllm_output_str_from_embeds = vllm_outputs_from_embeds[i][0]
+
+        assert hf_output_str == vllm_output_str_from_prompts, (
+            f"Test{i}:\n"
+            "HF: {hf_output_str!r}\n"
+            "vLLM_prompt: {vllm_output_str_from_prompts!r}")
+        assert hf_output_str == vllm_output_str_from_embeds, (
+            f"Test{i}:\n"
+            "HF: {hf_output_str}\n"
+            "vLLM_embeds: {vllm_output_str_from_embeds}")
+        assert vllm_output_str_from_prompts == vllm_output_str_from_embeds, (
+            f"Test{i}:\n"
+            "vLLM_prompt: {vllm_output_str_from_prompts}\n"
+            "vLLM_embeds: {vllm_output_str_from_embeds}")

tests/samplers/test_sampler.py

@@ -63,7 +63,7 @@ def test_sampler_all_greedy(seed: int):
                 block_tables={0: [1]},
             ))
 
-    _, _, input_metadata = worker._prepare_inputs(seq_group_metadata_list)
+    _, _, _, input_metadata = worker._prepare_inputs(seq_group_metadata_list)
     sampler_output = sampler(embedding=None,
                              hidden_states=input_tensor,
                              input_metadata=input_metadata)
@@ -96,7 +96,7 @@ def test_sampler_all_random(seed: int):
                 block_tables={0: [1]},
             ))
 
-    _, _, input_metadata = worker._prepare_inputs(seq_group_metadata_list)
+    _, _, _, input_metadata = worker._prepare_inputs(seq_group_metadata_list)
     sampler_output = sampler(embedding=None,
                              hidden_states=input_tensor,
                              input_metadata=input_metadata)
@@ -126,7 +126,7 @@ def test_sampler_all_beam(seed: int):
                 block_tables={0: [1]},
             ))
 
-    _, _, input_metadata = worker._prepare_inputs(seq_group_metadata_list)
+    _, _, _, input_metadata = worker._prepare_inputs(seq_group_metadata_list)
     sampler(embedding=None,
             hidden_states=input_tensor,
             input_metadata=input_metadata)
@@ -174,7 +174,7 @@ def test_sampler_mixed(seed: int):
                 block_tables={0: [1]},
             ))
 
-    _, _, input_metadata = worker._prepare_inputs(seq_group_metadata_list)
+    _, _, _, input_metadata = worker._prepare_inputs(seq_group_metadata_list)
     sampler_output = sampler(embedding=None,
                              hidden_states=input_tensor,
                              input_metadata=input_metadata)

vllm/engine/async_llm_engine.py

@@ -3,6 +3,7 @@
 from functools import partial
 from typing import (Any, Dict, Iterable, List, Optional, Set, Tuple, Type,
                     Union)
+import torch
 
 from vllm.config import ModelConfig
 from vllm.engine.arg_utils import AsyncEngineArgs
@@ -358,6 +359,7 @@ async def add_request(
         sampling_params: SamplingParams,
         prompt_token_ids: Optional[List[int]] = None,
         arrival_time: Optional[float] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
     ) -> AsyncStream:
         if self.log_requests:
             shortened_prompt = prompt
@@ -388,16 +390,20 @@ async def add_request(
             prompt=prompt,
             sampling_params=sampling_params,
             prompt_token_ids=prompt_token_ids,
-            arrival_time=arrival_time)
+            arrival_time=arrival_time,
+            prompt_embeds=prompt_embeds,
+        )
 
         return stream
 
     async def generate(
-            self,
-            prompt: Optional[str],
-            sampling_params: SamplingParams,
-            request_id: str,
-            prompt_token_ids: Optional[List[int]] = None) -> RequestOutput:
+        self,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        request_id: str,
+        prompt_token_ids: Optional[List[int]] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+    ) -> RequestOutput:
         """Generate outputs for a request.
 
         Generate outputs for a request. This method is a coroutine. It adds the
@@ -421,11 +427,14 @@ async def generate(
         arrival_time = time.monotonic()
 
         try:
-            stream = await self.add_request(request_id,
-                                            prompt,
-                                            sampling_params,
-                                            prompt_token_ids=prompt_token_ids,
-                                            arrival_time=arrival_time)
+            stream = await self.add_request(
+                request_id,
+                prompt,
+                sampling_params,
+                prompt_token_ids=prompt_token_ids,
+                arrival_time=arrival_time,
+                prompt_embeds=prompt_embeds,
+            )
 
             async for request_output in stream:
                 yield request_output

vllm/engine/llm_engine.py

@@ -3,6 +3,8 @@
 from functools import partial
 from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Tuple, Union
 
+import torch
+
 from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
                          SchedulerConfig)
 from vllm.core.scheduler import Scheduler, SchedulerOutputs
@@ -241,6 +243,7 @@ def add_request(
         sampling_params: SamplingParams,
         prompt_token_ids: Optional[List[int]] = None,
         arrival_time: Optional[float] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
     ) -> None:
         """Add a request to the engine's request pool.
 
@@ -250,24 +253,36 @@ def add_request(
 
         Args:
             request_id: The unique ID of the request.
-            prompt: The prompt string. Can be None if prompt_token_ids is
-                provided.
+            prompt: The prompt string. Can be None if prompt_token_ids
+                or prompt_embeds are provided.
             sampling_params: The sampling parameters for text generation.
             prompt_token_ids: The token IDs of the prompt. If None, we
                 use the tokenizer to convert the prompts to token IDs.
             arrival_time: The arrival time of the request. If None, we use
                 the current monotonic time.
+            prompt_embeds: The prompt embeddings. If set,
+                input prompt and prompt_token_ids are ignored
         """
         if arrival_time is None:
             arrival_time = time.monotonic()
-        if prompt_token_ids is None:
+
+        # If prompt_embeds is set, prompt_token_ids is filled with 0
+        if prompt_embeds is not None:
+            prompt_token_ids = [0] * prompt_embeds.size(0)
+        elif prompt_token_ids is None:
             assert prompt is not None
             prompt_token_ids = self.tokenizer.encode(prompt)
 
         # Create the sequences.
         block_size = self.cache_config.block_size
         seq_id = next(self.seq_counter)
-        seq = Sequence(seq_id, prompt, prompt_token_ids, block_size)
+        seq = Sequence(
+            seq_id,
+            prompt,
+            prompt_token_ids,
+            block_size,
+            prompt_embeds=prompt_embeds,
+        )
 
         # Create the sequence group.
         seq_group = SequenceGroup(request_id, [seq], sampling_params,
@@ -629,10 +644,17 @@ def _log_system_stats(
     def _decode_sequence(self, seq: Sequence,
                          sampling_params: SamplingParams) -> None:
         """Decodes the new token for a sequence."""
+
+        # if data has prompt embeds, all_input_ids are only output token ids
+        if seq.data.has_prompt_embeds_forwarding():
+            all_input_ids = seq.get_output_token_ids()
+        else:
+            all_input_ids = seq.get_token_ids()
+
         (new_tokens, new_output_text, prefix_offset,
          read_offset) = detokenize_incrementally(
              self.tokenizer,
-             all_input_ids=seq.get_token_ids(),
+             all_input_ids=all_input_ids,
              prev_tokens=seq.tokens,
              prefix_offset=seq.prefix_offset,
              read_offset=seq.read_offset,

vllm/entrypoints/api_server.py

@@ -5,6 +5,7 @@
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 import uvicorn
+import torch
 
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.engine.async_llm_engine import AsyncLLMEngine
@@ -23,16 +24,27 @@ async def generate(request: Request) -> Response:
 
     The request should be a JSON object with the following fields:
     - prompt: the prompt to use for the generation.
+    - prompt_embeds: the prompt embedding to use for the generation
+        instead of the prompt.
     - stream: whether to stream the results or not.
     - other fields: the sampling parameters (See `SamplingParams` for details).
     """
     request_dict = await request.json()
     prompt = request_dict.pop("prompt")
+    prompt_embeds = request_dict.pop("prompt_embeds", None)
+    if prompt_embeds is not None:
+        prompt_embeds = torch.tensor(prompt_embeds).to("cuda")
+        prompt = None
     stream = request_dict.pop("stream", False)
     sampling_params = SamplingParams(**request_dict)
     request_id = random_uuid()
 
-    results_generator = engine.generate(prompt, sampling_params, request_id)
+    results_generator = engine.generate(
+        prompt,
+        sampling_params,
+        request_id,
+        prompt_embeds=prompt_embeds,
+    )
 
     # Streaming case
     async def stream_results() -> AsyncGenerator[bytes, None]:
@@ -58,7 +70,12 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
 
     assert final_output is not None
     prompt = final_output.prompt
-    text_outputs = [prompt + output.text for output in final_output.outputs]
+    if prompt:
+        text_outputs = [
+            prompt + output.text for output in final_output.outputs
+        ]
+    else:
+        text_outputs = [output.text for output in final_output.outputs]
     ret = {"text": text_outputs}
     return JSONResponse(ret)