feat: Input embeddings

benchmarks/benchmark_serving.py

@@ -92,17 +92,9 @@ async def get_request(
         await asyncio.sleep(interval)
 
 
-async def send_request(
-    backend: str,
-    model: str,
-    api_url: str,
-    prompt: str,
-    prompt_len: int,
-    output_len: int,
-    best_of: int,
-    use_beam_search: bool,
-    pbar: tqdm
-) -> None:
+async def send_request(backend: str, model: str, api_url: str, prompt: str,
+                       prompt_len: int, output_len: int, best_of: int,
+                       use_beam_search: bool, pbar: tqdm) -> None:
     request_start_time = time.perf_counter()
 
     headers = {"User-Agent": "Benchmark Client"}
@@ -155,7 +147,6 @@ async def send_request(
     pbar.update(1)
 
 
-
 async def benchmark(
     backend: str,
     model: str,
@@ -217,7 +208,10 @@ def main(args: argparse.Namespace):
                         type=str,
                         default="vllm",
                         choices=["vllm", "tgi"])
-    parser.add_argument("--protocol", type=str, default="http", choices=["http", "https"])
+    parser.add_argument("--protocol",
+                        type=str,
+                        default="http",
+                        choices=["http", "https"])
     parser.add_argument("--host", type=str, default="localhost")
     parser.add_argument("--port", type=int, default=8000)
     parser.add_argument("--endpoint", type=str, default="/generate")

setup.py

@@ -28,7 +28,7 @@ def _is_neuron() -> bool:
     torch_neuronx_installed = True
     try:
         subprocess.run(["neuron-ls"], capture_output=True, check=True)
-    except FileNotFoundError as e:
+    except FileNotFoundError:
         torch_neuronx_installed = False
     return torch_neuronx_installed
 
@@ -99,7 +99,8 @@ def get_hipcc_rocm_version():
 def get_neuronxcc_version():
     import sysconfig
     site_dir = sysconfig.get_paths()["purelib"]
-    version_file = os.path.join(site_dir, "neuronxcc", "version", "__init__.py")
+    version_file = os.path.join(site_dir, "neuronxcc", "version",
+                                "__init__.py")
 
     # Check if the command was executed successfully
     with open(version_file, "rt") as fp:

tests/conftest.py

@@ -61,7 +61,10 @@ def __init__(
         ).cuda()
         if tokenizer_name is None:
             tokenizer_name = model_name
-        self.tokenizer = get_tokenizer(tokenizer_name, trust_remote_code=True)
+        self.tokenizer = get_tokenizer(
+            tokenizer_name,
+            trust_remote_code=True,
+        )
 
     def generate(
         self,
@@ -178,9 +181,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
@@ -199,9 +204,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]
 
@@ -210,12 +218,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

@@ -3,6 +3,7 @@
 Run `pytest tests/models/test_models.py --forked`.
 """
 import pytest
+from vllm.sampling_params import SamplingParams
 
 MODELS = [
     "facebook/opt-125m", "meta-llama/Llama-2-7b-hf",
@@ -39,3 +40,44 @@ 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(
+    vllm_runner,
+    example_prompts,
+    model: str,
+    dtype: str,
+    max_tokens: int,
+) -> None:
+    vllm_model = vllm_runner(model, dtype=dtype)
+    tokenizer = vllm_model.model.llm_engine.tokenizer
+    input_embeddings = vllm_model.model.llm_engine.workers[
+        0].model_runner.model.get_input_embeddings()
+
+    prompt_embeds = []
+    for prompt in example_prompts:
+        token_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
+        token_embeds = input_embeddings(token_ids)
+        prompt_embeds.append(token_embeds[0])
+
+    outputs_from_prompts = vllm_model.model.generate(
+        example_prompts,
+        sampling_params=SamplingParams(temperature=0.0, max_tokens=max_tokens),
+        prompt_embeds=None)
+    outputs_from_embeds = vllm_model.model.generate(
+        None,
+        sampling_params=SamplingParams(temperature=0.0, max_tokens=max_tokens),
+        prompt_embeds=prompt_embeds,
+    )
+    del vllm_model
+
+    for output_prompt, output_embed in zip(outputs_from_prompts,
+                                           outputs_from_embeds):
+        assert output_prompt.outputs[0].token_ids == output_embed.outputs[
+            0].token_ids, (
+                f"output_prompt: {output_prompt.outputs[0].token_ids}\n",
+                f"output_embed: {output_embed.outputs[0].token_ids}",
+            )

tests/samplers/test_sampler.py

@@ -19,24 +19,29 @@ def __init__(self, vocab_size: int, fake_logits: torch.Tensor):
         self.fake_logits = fake_logits
 
     def forward(self, *args, **kwargs):
-        with patch("vllm.model_executor.layers.sampler._prune_hidden_states",
-                   lambda x, y: x), patch(
-                       "vllm.model_executor.layers.sampler._get_logits",
-                       lambda *args, **kwargs: self.fake_logits):
+        with patch(
+                "vllm.model_executor.layers.sampler._prune_hidden_states",
+                lambda x, y: x,
+        ), patch(
+                "vllm.model_executor.layers.sampler._get_logits",
+                lambda *args, **kwargs: self.fake_logits,
+        ):
             return super().forward(*args, **kwargs)
 
 
 def _prepare_test(
-    batch_size: int
+    batch_size: int,
 ) -> Tuple[torch.Tensor, torch.Tensor, MockLogitsSampler, ModelRunner]:
     vocab_size = 32000
     input_tensor = torch.rand((batch_size, 1024),
                               device="cuda",
                               dtype=torch.float16)
-    fake_logits = torch.full((batch_size, vocab_size),
-                             1e-2,
-                             device=input_tensor.device,
-                             dtype=input_tensor.dtype)
+    fake_logits = torch.full(
+        (batch_size, vocab_size),
+        1e-2,
+        device=input_tensor.device,
+        dtype=input_tensor.dtype,
+    )
     sampler = MockLogitsSampler(32000, fake_logits)
     model_runner = ModelRunner(None, None, None)
     return input_tensor, fake_logits, sampler, model_runner

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, AsyncIterator)
+import torch
 
 from vllm.config import ModelConfig
 from vllm.engine.arg_utils import AsyncEngineArgs
@@ -372,6 +373,7 @@ async def add_request(
         prompt_token_ids: Optional[List[int]] = None,
         arrival_time: Optional[float] = None,
         prefix_pos: Optional[int] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
     ) -> AsyncStream:
         if self.log_requests:
             shortened_prompt = prompt
@@ -404,7 +406,9 @@ async def add_request(
             sampling_params=sampling_params,
             prompt_token_ids=prompt_token_ids,
             arrival_time=arrival_time,
-            prefix_pos=prefix_pos)
+            prefix_pos=prefix_pos,
+            prompt_embeds=prompt_embeds,
+        )
 
         return stream
 
@@ -415,6 +419,7 @@ async def generate(
         request_id: str,
         prompt_token_ids: Optional[List[int]] = None,
         prefix_pos: Optional[int] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
     ) -> AsyncIterator[RequestOutput]:
         """Generate outputs for a request.
 
@@ -492,7 +497,8 @@ async def generate(
                                             sampling_params,
                                             prompt_token_ids=prompt_token_ids,
                                             arrival_time=arrival_time,
-                                            prefix_pos=prefix_pos)
+                                            prefix_pos=prefix_pos,
+                                            prompt_embeds=prompt_embeds)
 
             async for request_output in stream:
                 yield request_output

vllm/engine/llm_engine.py

@@ -4,6 +4,7 @@
 import time
 from typing import (TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple,
                     Union)
+import torch
 
 from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
                          SchedulerConfig)
@@ -340,6 +341,7 @@ def add_request(
         prompt_token_ids: Optional[List[int]] = None,
         arrival_time: Optional[float] = None,
         prefix_pos: Optional[int] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
     ) -> None:
         """Add a request to the engine's request pool.
 
@@ -349,8 +351,8 @@ 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.
@@ -385,17 +387,27 @@ def add_request(
             >>>    SamplingParams(temperature=0.0))
             >>> # continue the request processing
             >>> ...
+            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)
 
         # Check whether the input specifies prefix
         prefix = self.scheduler.prefix_pool.add_or_get_prefix(
@@ -835,10 +847,17 @@ def _log_system_stats(
 
     def _decode_sequence(self, seq: Sequence, prms: 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
@@ -28,20 +29,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")
     prefix_pos = request_dict.pop("prefix_pos", None)
+    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,
-                                        prefix_pos=prefix_pos)
+                                        prefix_pos=prefix_pos,
+                                        prompt_embeds=prompt_embeds)
 
     # Streaming case
     async def stream_results() -> AsyncGenerator[bytes, None]:
@@ -67,7 +75,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)