@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cc2a77d7f1
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 12:15 UTC

[Core] [Bugfix] Add Input Embeddings (#15428)

Signed-off-by: Andrew Sansom <andrew@protopia.ai> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: 临景 <linjing.yx@alibaba-inc.com> Co-authored-by: Bryce1010 <bryceyx@gmail.com> Co-authored-by: Nan2018 <nan@protopia.ai> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>

Andrew Sansom committed May 2, 2025 at 03:06 UTC cc2a77d7f1dfa94fe07150ae3a637604d49aa32c
22 files changed +691 -113
tests/conftest.py
+10 -8
@@ -787,7 +787,7 @@ class VllmRunner:
787
788 def get_inputs(
789 self,
790 - prompts: list[str],
790 + prompts: Union[list[str], list[torch.Tensor]],
791 images: Optional[PromptImageInput] = None,
792 videos: Optional[PromptVideoInput] = None,
793 audios: Optional[PromptAudioInput] = None,
@@ -809,16 +809,18 @@ class VllmRunner:
809 if audios is not None and (audio := audios[i]) is not None:
810 multi_modal_data["audio"] = audio
811
812 - inputs.append(
813 - TextPrompt(prompt=prompt,
814 - multi_modal_data=multi_modal_data
815 - if multi_modal_data else None))
812 + text_prompt_kwargs = {
813 + ("prompt" if isinstance(prompt, str) else "prompt_embeds"):
814 + prompt,
815 + "multi_modal_data": multi_modal_data or None
816 + }
817 + inputs.append(TextPrompt(**text_prompt_kwargs))
818
819 return inputs
820
821 def generate(
822 self,
821 - prompts: list[str],
823 + prompts: Union[list[str], list[torch.Tensor]],
824 sampling_params: SamplingParams,
825 images: Optional[PromptImageInput] = None,
826 videos: Optional[PromptVideoInput] = None,
@@ -844,7 +846,7 @@ class VllmRunner:
846 output_str = sample.text
847 output_ids = list(sample.token_ids)
848 req_sample_output_ids.append(prompt_ids + output_ids)
847 - req_sample_output_strs.append(prompt_str + output_str)
849 + req_sample_output_strs.append((prompt_str or "") + output_str)
850 outputs.append((req_sample_output_ids, req_sample_output_strs))
851 return outputs
852
@@ -911,7 +913,7 @@ class VllmRunner:
913
914 def generate_greedy(
915 self,
914 - prompts: list[str],
916 + prompts: Union[list[str], list[torch.Tensor]],
917 max_tokens: int,
918 images: Optional[PromptImageInput] = None,
919 videos: Optional[PromptVideoInput] = None,
tests/core/test_scheduler.py
+73 -1
@@ -2,16 +2,18 @@
2
3 import time
4 from collections import deque
5 +from typing import Optional
6 from unittest.mock import MagicMock
7
8 import pytest # noqa
9 +import torch
10 from torch import Use # noqa
11
12 from vllm.config import CacheConfig, LoRAConfig, SchedulerConfig
13 from vllm.core.interfaces import AllocStatus
14 from vllm.core.scheduler import Scheduler, SchedulingBudget
15 from vllm.lora.request import LoRARequest
14 -from vllm.sequence import SequenceGroup
16 +from vllm.sequence import SequenceGroup, SequenceStatus
17
18 from .utils import (append_new_token, append_new_token_seq,
19 append_new_token_seq_group, create_dummy_prompt,
@@ -968,3 +970,73 @@ def test_no_multiple_partial_prefills_with_chunked_prefill_and_prefix_caching(
970 ), "A partial prefix of C (4 tokens) should be prefilled, with the "
971 "remaining tokens fit into 3 token budget (4-1 from the seqA). It will "
972 "then be rounded down to 2 tokens on block size, thus 6 tokens in total."
973 +
974 +
975 +def test_no_batches_mixed_with_prompt_tokens_and_prompt_embeds():
976 + """
977 + Test that the scheduler does not schedule batches with prompt tokens and
978 + prompt embeddings co-mingled.
979 + """
980 + block_size = 2
981 + max_seq_group = 3
982 + scheduler = initialize_scheduler(
983 + block_size=block_size,
984 + num_cpu_blocks=16,
985 + num_gpu_blocks=16,
986 + max_num_seqs=max_seq_group,
987 + max_model_len=100,
988 + enable_prefix_caching=True,
989 + )
990 +
991 + # the odd indexed inputs should be passed in via embeddings,
992 + # evens via token_ids
993 + seq_length = 7
994 + embedding_size = 5
995 + num_seqs = 11
996 + seq_tokens: list[list[int]] = []
997 + seq_embeds: list[Optional[torch.Tensor]] = []
998 + for i in range(num_seqs):
999 + if i % 2:
1000 + seq_tokens.append(list(range(seq_length)))
1001 + seq_embeds.append(None)
1002 + else:
1003 + seq_tokens.append([0] * seq_length)
1004 + seq_embeds.append(torch.rand(embedding_size))
1005 +
1006 + seq_and_seq_groups = [
1007 + create_dummy_prompt(f"{i}",
1008 + prompt_tokens=seq_tokens[i],
1009 + prompt_embeds=seq_embeds[i],
1010 + block_size=block_size)
1011 + for i in range(len(seq_tokens))
1012 + ]
1013 +
1014 + for _, seq_group in seq_and_seq_groups:
1015 + scheduler.add_seq_group(seq_group)
1016 +
1017 + while not all(seq.is_finished() for seq, _ in seq_and_seq_groups):
1018 + unfinished_seq_groups = [
1019 + seq_group for _, seq_group in seq_and_seq_groups
1020 + if not seq_group.is_finished()
1021 + ]
1022 + _, out = schedule_and_update_computed_tokens(scheduler)
1023 + assert len(out.scheduled_seq_groups) > 0
1024 + batch_is_prompt_embeds = out.scheduled_seq_groups[
1025 + 0].seq_group.uses_prompt_embeds()
1026 + expected_scheduled_seq_groups = [
1027 + seq_group for seq_group in unfinished_seq_groups
1028 + if seq_group.uses_prompt_embeds() == batch_is_prompt_embeds
1029 + ]
1030 +
1031 + # We should have as many scheduled groups as possible, without mixing
1032 + assert len(out.scheduled_seq_groups) == min(
1033 + max_seq_group, len(expected_scheduled_seq_groups))
1034 + assert all(scheduled_seq_group.seq_group.uses_prompt_embeds() ==
1035 + batch_is_prompt_embeds
1036 + for scheduled_seq_group in out.scheduled_seq_groups)
1037 +
1038 + # Finish the scheduled groups
1039 + for scheduled_seq_group in out.scheduled_seq_groups:
1040 + for seq in scheduled_seq_group.seq_group.seqs:
1041 + seq.status = SequenceStatus.FINISHED_STOPPED
1042 + scheduler.free_finished_seq_groups()
tests/core/utils.py
+9 -2
@@ -5,9 +5,11 @@ from collections import defaultdict
5 from collections.abc import Sequence as GenericSequence
6 from typing import Any, Optional
7
8 +import torch
9 +
10 from vllm import SamplingParams
11 from vllm.core.scheduler import Scheduler, SchedulerOutputs
10 -from vllm.inputs import EncoderDecoderInputs, token_inputs
12 +from vllm.inputs import EncoderDecoderInputs, embeds_inputs, token_inputs
13 from vllm.lora.request import LoRARequest
14 from vllm.sequence import (Logprob, Sequence, SequenceGroup,
15 SequenceGroupMetadata)
@@ -19,6 +21,7 @@ def create_dummy_prompt(
21 block_size: Optional[int] = None,
22 lora_request: Optional[LoRARequest] = None,
23 prompt_tokens: Optional[list[int]] = None,
24 + prompt_embeds: Optional[torch.Tensor] = None,
25 min_tokens: int = 0,
26 max_tokens: int = 16,
27 ) -> tuple[Sequence, SequenceGroup]:
@@ -31,9 +34,13 @@ def create_dummy_prompt(
34 prompt_tokens = list(range(prompt_length))
35
36 prompt_str = " ".join([str(t) for t in prompt_tokens])
37 + inputs = token_inputs(
38 + prompt_token_ids=prompt_tokens,
39 + prompt=prompt_str) if prompt_embeds is None else embeds_inputs(
40 + prompt_embeds=prompt_embeds)
41 prompt = Sequence(
42 int(request_id),
36 - inputs=token_inputs(prompt_tokens, prompt=prompt_str),
43 + inputs=inputs,
44 block_size=block_size,
45 )
46 seq_group = SequenceGroup(
tests/models/language/generation/test_common.py
+26
@@ -1,4 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 +import os
3 +from typing import Optional
4 +
5 import pytest
6 import torch
7
@@ -110,6 +113,18 @@ def test_models(hf_runner, vllm_runner, example_prompts, model: str,
113 hf_outputs = hf_model.generate_greedy_logprobs_limit(
114 example_prompts, max_tokens, num_logprobs)
115
116 + prompt_embeds: Optional[list[torch.Tensor]] = [] if os.getenv(
117 + "VLLM_USE_V1") == "0" else None
118 + prompt_token_ids = []
119 + for prompt in example_prompts:
120 + token_ids = hf_model.tokenizer(prompt,
121 + return_tensors="pt").input_ids.to(
122 + hf_model.model.device)
123 + prompt_token_ids.append(token_ids)
124 + if prompt_embeds is not None:
125 + prompt_embeds.append(hf_model.model.get_input_embeddings()(
126 + token_ids).squeeze(0))
127 +
128 with vllm_runner(
129 model,
130 tokenizer_name=model_info.tokenizer or model,
@@ -119,6 +134,9 @@ def test_models(hf_runner, vllm_runner, example_prompts, model: str,
134 ) as vllm_model:
135 vllm_outputs = vllm_model.generate_greedy_logprobs(
136 example_prompts, max_tokens, num_logprobs)
137 + if prompt_embeds is not None:
138 + vllm_outputs_from_embeds = vllm_model.generate_greedy_logprobs(
139 + prompt_embeds, max_tokens, num_logprobs)
140
141 check_logprobs_close(
142 outputs_0_lst=hf_outputs,
@@ -126,6 +144,14 @@ def test_models(hf_runner, vllm_runner, example_prompts, model: str,
144 name_0="hf",
145 name_1="vllm",
146 )
147 + if prompt_embeds is not None:
148 + check_logprobs_close(
149 + outputs_0_lst=vllm_outputs,
150 + outputs_1_lst=vllm_outputs_from_embeds,
151 + name_0="vllm",
152 + name_1="vllm_from_embeds",
153 + )
154 +
155 if use_rocm_aiter:
156 # this is to ensure that vllm engine
157 # has deallocated the memory before running the next
tests/worker/test_model_runner.py
+109 -34
@@ -31,8 +31,13 @@ def test_deepseek_mla_attn_backend_module():
31 assert model_runner.attn_backend.__name__ == "TritonMLABackend"
32
33
34 -@pytest.mark.parametrize("batch_size", list(range(1, 257)))
35 -def test_prepare_prompt(batch_size):
34 +@pytest.mark.parametrize("batch_size", list(range(1, 257, 3)))
35 +@pytest.mark.parametrize("use_prompt_embeds", [True, False])
36 +def test_prepare_prompt(batch_size, use_prompt_embeds, monkeypatch):
37 + if use_prompt_embeds:
38 + # Prompt Embeddings is only currently supported on V0
39 + monkeypatch.setenv("VLLM_USE_V1", "0")
40 +
41 model_runner = _create_model_runner(
42 "facebook/opt-125m",
43 max_num_batched_tokens=100000,
@@ -43,11 +48,20 @@ def test_prepare_prompt(batch_size):
48 seq_lens: list[int] = []
49 seq_group_metadata_list: list[SequenceGroupMetadata] = []
50 block_tables = {0: [1]}
51 + expected_input_embeds_len = 0
52 for i in range(batch_size):
53 # make sure all tokens fit into one block
54 seq_len = i % (model_runner.block_size - 1) + 1
55 seq_lens.append(seq_len)
50 - seq_data = SequenceData.from_seqs(range(seq_len))
56 + if use_prompt_embeds:
57 + seq_data = SequenceData.from_seqs(
58 + prompt_token_ids=[0] * seq_len,
59 + prompt_embeds=torch.rand(seq_len, 10),
60 + )
61 + expected_input_embeds_len += seq_len
62 + else:
63 + seq_data = SequenceData.from_seqs(prompt_token_ids=range(seq_len))
64 +
65 seq_group_metadata = SequenceGroupMetadata(
66 request_id=f"test_{i}",
67 is_prompt=True,
@@ -68,6 +82,7 @@ def test_prepare_prompt(batch_size):
82 seq_group_metadata_list)
83 input_tokens = model_input.input_tokens
84 input_positions = model_input.input_positions
85 + input_embeds = model_input.inputs_embeds
86 attn_metadata = model_input.attn_metadata
87 return_seq_lens = model_input.seq_lens
88 slot_mapping = attn_metadata.slot_mapping
@@ -121,7 +136,11 @@ def test_prepare_prompt(batch_size):
136
137 assert len(input_tokens) == sum(seq_lens)
138 assert len(input_positions) == sum(seq_lens)
124 - torch.testing.assert_close(input_tokens, input_positions)
139 + if expected_input_embeds_len == 0:
140 + torch.testing.assert_close(input_tokens, input_positions)
141 + assert input_embeds is None
142 + else:
143 + assert len(input_embeds) == expected_input_embeds_len
144
145 sampling_metadata = SamplingMetadata.prepare(
146 seq_group_metadata_list,
@@ -145,8 +164,13 @@ def test_prepare_prompt(batch_size):
164 torch.testing.assert_close(actual, expected)
165
166
148 -@pytest.mark.parametrize("batch_size", list(range(1, 257)))
149 -def test_prepare_decode_cuda_graph(batch_size):
167 +@pytest.mark.parametrize("batch_size", list(range(1, 257, 3)))
168 +@pytest.mark.parametrize("use_prompt_embeds", [True, False])
169 +def test_prepare_decode_cuda_graph(batch_size, use_prompt_embeds, monkeypatch):
170 + if use_prompt_embeds:
171 + # Prompt Embeddings is only currently supported on V0
172 + monkeypatch.setenv("VLLM_USE_V1", "0")
173 +
174 model_runner = _create_model_runner(
175 "facebook/opt-125m",
176 seed=0,
@@ -164,10 +188,19 @@ def test_prepare_decode_cuda_graph(batch_size):
188 # make sure all tokens fit into one block
189 context_len = i % (model_runner.block_size - 1) + 1
190 context_lens.append(context_len)
167 - seq_data = SequenceData.from_seqs(range(context_len))
191 + if use_prompt_embeds:
192 + seq_data = SequenceData.from_seqs(
193 + prompt_token_ids=[0] * context_len,
194 + prompt_embeds=torch.rand(context_len, 10),
195 + )
196 + output_embed = torch.rand(10)
197 + else:
198 + seq_data = SequenceData.from_seqs(
199 + prompt_token_ids=range(context_len))
200 + output_embed = None
201 seq_data.update_num_computed_tokens(context_len)
202 # Append one token ID since prefill is finished.
170 - seq_data.append_token_id(1, 0)
203 + seq_data.append_token_id(1, 0, output_embed)
204 seq_group_metadata = SequenceGroupMetadata(
205 request_id=f"test_{i}",
206 is_prompt=False,
@@ -180,9 +213,12 @@ def test_prepare_decode_cuda_graph(batch_size):
213
214 model_input = model_runner._prepare_model_input_tensors(
215 seq_group_metadata_list)
183 - input_tokens, input_positions, attn_metadata, slot_mapping = (
184 - model_input.input_tokens, model_input.input_positions,
185 - model_input.attn_metadata, model_input.attn_metadata.slot_mapping)
216 + input_tokens = model_input.input_tokens
217 + input_positions = model_input.input_positions
218 + input_embeds = model_input.inputs_embeds
219 + attn_metadata = model_input.attn_metadata
220 + slot_mapping = attn_metadata.slot_mapping
221 +
222 assert len(slot_mapping) == len(input_tokens)
223
224 expected_bs = model_runner.vllm_config.pad_for_cudagraph(
@@ -227,7 +263,7 @@ def test_prepare_decode_cuda_graph(batch_size):
263 # block table's first index corresponds to each batch, meaning in
264 # decoding it is each token.
265 assert attn_metadata.block_tables.shape[0] == len(input_tokens)
230 - # Block table's second dim correspondsd to each token's block number.
266 + # Block table's second dim corresponds to each token's block number.
267 # It is padded up to
268 assert attn_metadata.block_tables.shape[1] == (
269 model_runner.get_max_block_per_batch())
@@ -235,7 +271,12 @@ def test_prepare_decode_cuda_graph(batch_size):
271
272 assert len(input_tokens) == expected_bs
273 assert len(input_positions) == expected_bs
238 - torch.allclose(input_tokens, input_positions)
274 + if use_prompt_embeds:
275 + expected_input_embeds_length = start_loc[-1]
276 + assert len(input_embeds) == expected_input_embeds_length
277 + assert expected_input_embeds_length <= expected_bs
278 + else:
279 + assert input_embeds is None
280
281 # Verify Sampling
282 expected_selected_token_indices = []
@@ -266,25 +307,27 @@ def test_empty_seq_group():
307 seq_group_metadata_list: list[SequenceGroupMetadata] = []
308 model_input = model_runner._prepare_model_input_tensors(
309 seq_group_metadata_list)
269 - input_tokens, input_positions, attn_metadata = (
270 - model_input.input_tokens,
271 - model_input.input_positions,
272 - model_input.attn_metadata,
273 - )
310 +
311 + input_tokens = model_input.input_tokens
312 + input_positions = model_input.input_positions
313 + attn_metadata = model_input.attn_metadata
314 +
315 assert input_tokens is None
316 assert input_positions is None
317 assert attn_metadata is None
318
319 model_input = model_runner._prepare_model_input_tensors(
320 seq_group_metadata_list)
280 - (input_tokens, input_positions, attn_metadata, return_seq_lens) = (
281 - model_input.input_tokens,
282 - model_input.input_positions,
283 - model_input.attn_metadata,
284 - model_input.seq_lens,
285 - )
321 +
322 + input_tokens = model_input.input_tokens
323 + input_positions = model_input.input_positions
324 + input_embeds = model_input.inputs_embeds
325 + attn_metadata = model_input.attn_metadata
326 + return_seq_lens = model_input.seq_lens
327 +
328 assert input_tokens is None
329 assert input_positions is None
330 + assert input_embeds is None
331 assert attn_metadata is None
332 assert return_seq_lens is None
333
@@ -299,9 +342,15 @@ def distributed_init():
342 ensure_model_parallel_initialized(1, 1)
343
344
302 -@pytest.mark.parametrize("batch_size", list(range(2, 128)))
345 +@pytest.mark.parametrize("batch_size", list(range(2, 128, 3)))
346 @pytest.mark.parametrize("enforce_eager", [True, False])
304 -def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
347 +@pytest.mark.parametrize('use_prompt_embeds', [True, False])
348 +def test_hybrid_batches(batch_size, enforce_eager, use_prompt_embeds,
349 + distributed_init, monkeypatch):
350 + if use_prompt_embeds:
351 + # Prompt Embeddings is only currently supported on V0
352 + monkeypatch.setenv("VLLM_USE_V1", "0")
353 +
354 model_runner = _create_model_runner(
355 "facebook/opt-125m",
356 seed=0,
@@ -320,11 +369,20 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
369 block_tables = {0: [1]}
370 prefill_batch_size = batch_size // 2
371 decode_batch_size = batch_size - prefill_batch_size
372 + expected_input_embeds_len = 0
373 for i in range(prefill_batch_size):
374 # make sure all tokens fit into one block
375 seq_len = i % (model_runner.block_size - 1) + 1
376 seq_lens.append(seq_len)
327 - seq_data = SequenceData.from_seqs(range(seq_len))
377 + if use_prompt_embeds:
378 + seq_data = SequenceData.from_seqs(
379 + prompt_token_ids=[0] * seq_len,
380 + prompt_embeds=torch.rand(seq_len, 10),
381 + )
382 + expected_input_embeds_len += seq_len
383 + else:
384 + seq_data = SequenceData.from_seqs(
385 + prompt_token_ids=range(seq_len), )
386 seq_group_metadata = SequenceGroupMetadata(
387 request_id=f"test_{i}",
388 is_prompt=True,
@@ -340,8 +398,21 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
398 for i in range(prefill_batch_size, batch_size):
399 # make sure all tokens fit into one block
400 context_len = i % (model_runner.block_size - 1) + 1
343 - seq_data = SequenceData.from_seqs(range(context_len))
344 - seq_data.append_token_id(1, 0)
401 + if use_prompt_embeds:
402 + seq_data = SequenceData.from_seqs(
403 + prompt_token_ids=[0] * context_len,
404 + prompt_embeds=torch.rand(context_len, 10),
405 + )
406 + output_embed = torch.rand(10)
407 + # This also iterates the expected input_embeds, because the model
408 + # needs both the input and output embeddings passed into together
409 + expected_input_embeds_len += 1
410 + else:
411 + seq_data = SequenceData.from_seqs(
412 + prompt_token_ids=range(context_len), )
413 + output_embed = None
414 + assert len(seq_data.prompt_token_ids) == context_len
415 + seq_data.append_token_id(1, 0, output_embed)
416 seq_data.update_num_computed_tokens(context_len)
417 seq_group_metadata = SequenceGroupMetadata(
418 request_id=f"test_{i}",
@@ -355,11 +426,11 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
426 decode_metadata_list.append(seq_group_metadata)
427
428 model_input = model_runner.prepare_model_input(seq_group_metadata_list)
358 - (input_tokens, input_positions, attn_metadata) = (
359 - model_input.input_tokens,
360 - model_input.input_positions,
361 - model_input.attn_metadata,
362 - )
429 +
430 + input_tokens = model_input.input_tokens
431 + input_positions = model_input.input_positions
432 + input_embeds = model_input.inputs_embeds
433 + attn_metadata = model_input.attn_metadata
434
435 prefill_meta_actual = attn_metadata.prefill_metadata
436 decode_meta_actual = attn_metadata.decode_metadata
@@ -369,6 +440,10 @@ def test_hybrid_batches(batch_size, enforce_eager, distributed_init):
440 assert attn_metadata.num_prefills == prefill_batch_size
441 assert attn_metadata.num_decode_tokens == decode_batch_size
442 assert attn_metadata.num_prefill_tokens == sum(seq_lens)
443 + if expected_input_embeds_len == 0:
444 + assert input_embeds is None
445 + else:
446 + assert len(input_embeds) == expected_input_embeds_len
447
448 # Verify attn metadata is consistent. We don't need to test individual
449 # values here because they are tested above.
vllm/attention/backends/flashinfer.py
+11 -3
@@ -367,9 +367,17 @@ class FlashInferState(AttentionState):
367 # scheduled while CUDA graph mode is enabled. We don't run graph in that
368 # case.
369 if use_cuda_graph and is_decode:
370 - batch_size = model_input.input_tokens.shape[0]
371 - state = (self.runner.graph_runners[model_input.virtual_engine]
372 - [batch_size].attn_state)
370 + if model_input.inputs_embeds is None:
371 + batch_size = model_input.input_tokens.shape[0]
372 + state = (
373 + self.runner.graph_runners[model_input.virtual_engine][(
374 + batch_size, False)].attn_state)
375 + else:
376 + batch_size = model_input.inputs_embeds.shape[0]
377 + state = (
378 + self.runner.graph_runners[model_input.virtual_engine][(
379 + batch_size, True)].attn_state)
380 +
381 model_input.attn_metadata.prefill_wrapper = state._get_prefill_wrapper(
382 )
383 model_input.attn_metadata.decode_wrapper = state._get_decode_wrapper()
vllm/core/scheduler.py
+32
@@ -1071,6 +1071,7 @@ class Scheduler:
1071 )
1072 ignored_seq_groups: List[SequenceGroup] = []
1073 seq_groups: List[ScheduledSequenceGroup] = []
1074 + using_prompt_embeds: bool = False
1075
1076 waiting_queue = self.waiting
1077
@@ -1138,6 +1139,15 @@ class Scheduler:
1139 waiting_queue.popleft()
1140 continue
1141
1142 + # We cannot mix sequence groups that use prompt embeds and
1143 + # those that do not.
1144 + if len(seq_groups) == 0:
1145 + using_prompt_embeds = seq_group.uses_prompt_embeds()
1146 + if using_prompt_embeds != seq_group.uses_prompt_embeds():
1147 + leftover_waiting_sequences.appendleft(seq_group)
1148 + waiting_queue.popleft()
1149 + continue
1150 +
1151 lora_int_id = 0
1152 if self.lora_enabled:
1153 lora_int_id = seq_group.lora_int_id
@@ -1295,17 +1305,39 @@ class Scheduler:
1305
1306 # Merge lists
1307 num_prefill_groups = len(prefills.seq_groups)
1308 + ignored_seq_groups_for_embeds = list[SequenceGroup]()
1309 if num_prefill_groups > 0:
1310 scheduled_seq_groups = prefills.seq_groups
1311 scheduled_seq_groups.extend(running_scheduled.decode_seq_groups)
1312 + ignored_seq_groups_for_embeds.clear()
1313 else:
1314 scheduled_seq_groups = running_scheduled.decode_seq_groups
1315 + if len(scheduled_seq_groups) > 0:
1316 + using_prompt_embeds = scheduled_seq_groups[
1317 + 0].seq_group.uses_prompt_embeds()
1318 + ignored_seq_groups_for_embeds.clear()
1319 + indices_ignored = list[int]()
1320 + for i, schedule_seq_group in enumerate(scheduled_seq_groups):
1321 + if using_prompt_embeds !=\
1322 + schedule_seq_group.seq_group.uses_prompt_embeds():
1323 + ignored_seq_groups_for_embeds.append(
1324 + schedule_seq_group.seq_group)
1325 + indices_ignored.append(i)
1326 + if len(ignored_seq_groups_for_embeds) > 0:
1327 + scheduled_seq_groups = [
1328 + group for i, group in enumerate(scheduled_seq_groups)
1329 + if i not in indices_ignored
1330 + ]
1331 + else:
1332 + ignored_seq_groups_for_embeds.clear()
1333 +
1334 scheduled_seq_groups.extend(swapped_in.decode_seq_groups)
1335
1336 blocks_to_copy = running_scheduled.blocks_to_copy
1337 blocks_to_copy.extend(swapped_in.blocks_to_copy)
1338
1339 ignored_seq_groups = prefills.ignored_seq_groups
1340 + ignored_seq_groups.extend(ignored_seq_groups_for_embeds)
1341 ignored_seq_groups.extend(swapped_in.infeasible_seq_groups)
1342
1343 return SchedulerOutputs(
vllm/engine/async_llm_engine.py
+8
@@ -489,6 +489,14 @@ class _AsyncLLMEngine(LLMEngine):
489 if arrival_time is None:
490 arrival_time = time.time()
491
492 + if (isinstance(prompt, dict)
493 + and prompt.get("prompt_embeds", None) is not None
494 + and not prompt.get("prompt_token_ids", None)):
495 + # We use the -2 dimension (instead of 0) in case a batched input
496 + # of batch size 1 is passed in.
497 + prompt["prompt_token_ids"] = [0
498 + ] * prompt["prompt_embeds"].shape[-2]
499 +
500 if self.tokenizer is not None:
501 tokenizer = await self.get_tokenizer_async(lora_request)
502 self._validate_token_prompt(prompt, tokenizer=tokenizer)
vllm/engine/llm_engine.py
+13 -3
@@ -753,6 +753,12 @@ class LLMEngine:
753 if arrival_time is None:
754 arrival_time = time.time()
755
756 + if (isinstance(prompt, dict)
757 + and prompt.get("prompt_embeds", None) is not None
758 + and not prompt.get("prompt_token_ids", None)):
759 + seq_len = prompt["prompt_embeds"].shape[0]
760 + prompt["prompt_token_ids"] = [0] * seq_len
761 +
762 if self.tokenizer is not None:
763 self._validate_token_prompt(
764 prompt,
@@ -1267,11 +1273,13 @@ class LLMEngine:
1273 if self.scheduler_config.is_multi_step:
1274 is_prefill_append = seq.data.get_num_uncomputed_tokens(
1275 ) == 0
1270 - seq.append_token_id(sample.output_token, sample.logprobs)
1276 + seq.append_token_id(sample.output_token, sample.logprobs,
1277 + sample.output_embed)
1278 if not is_prefill_append:
1279 seq_group.update_num_computed_tokens(1)
1280 else:
1274 - seq.append_token_id(sample.output_token, sample.logprobs)
1281 + seq.append_token_id(sample.output_token, sample.logprobs,
1282 + sample.output_embed)
1283
1284 def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]:
1285 """Performs one decoding iteration and returns newly generated results.
@@ -2032,10 +2040,12 @@ class LLMEngine:
2040 tokenizer = (None if self.tokenizer is None else
2041 self.tokenizer.get_lora_tokenizer(lora_request))
2042
2035 - prompt_ids = prompt_inputs["prompt_token_ids"]
2043 + prompt_ids = prompt_inputs.get("prompt_token_ids", [])
2044 if not prompt_ids:
2045 if prompt_type == "encoder" and model_config.is_multimodal_model:
2046 pass # Mllama may have empty encoder inputs for text-only data
2047 + if prompt_inputs["type"] == "embeds":
2048 + pass
2049 else:
2050 raise ValueError(f"The {prompt_type} prompt cannot be empty")
2051
vllm/engine/output_processor/multi_step.py
+4 -2
@@ -167,6 +167,7 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
167 sampling_params: SamplingParams) -> None:
168 output_token_ids = [sample.output_token for sample in valid_samples]
169 output_logprobs = [sample.logprobs for sample in valid_samples]
170 + output_embeds = [sample.output_embed for sample in valid_samples]
171
172 # Truncate to max_tokens if necessary.
173 remaining_tokens = sampling_params.max_tokens - (seq.get_output_len() +
@@ -190,11 +191,12 @@ class MultiStepOutputProcessor(SequenceGroupOutputProcessor):
191 is_prefill_sampled_token = seq.data.get_num_uncomputed_tokens() == 0
192 # Incrementally append tokens to the sequence, as if we had only one new
193 # token.
193 - for output_token_id, output_logprob in zip(output_token_ids,
194 - output_logprobs):
194 + for output_token_id, output_logprob, output_embed in zip(
195 + output_token_ids, output_logprobs, output_embeds):
196 seq.append_token_id(
197 token_id=output_token_id,
198 logprobs=output_logprob,
199 + token_embed=output_embed,
200 )
201
202 if is_prefill_sampled_token:
vllm/engine/output_processor/single_step.py
+2 -1
@@ -119,7 +119,8 @@ class SingleStepOutputProcessor(SequenceGroupOutputProcessor):
119 sample = outputs.samples[0]
120 seq = seq_group.first_seq
121 if not is_async:
122 - seq.append_token_id(sample.output_token, sample.logprobs)
122 + seq.append_token_id(sample.output_token, sample.logprobs,
123 + sample.output_embed)
124 if sampling_params.detokenize and self.detokenizer:
125 new_char_count = self.detokenizer.decode_sequence_inplace(
126 seq, sampling_params)
vllm/inputs/__init__.py
+4 -2
@@ -1,9 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from .data import (DecoderOnlyInputs, EncoderDecoderInputs,
3 +from .data import (DecoderOnlyInputs, EmbedsInputs, EncoderDecoderInputs,
4 ExplicitEncoderDecoderPrompt, ProcessorInputs, PromptType,
5 SingletonInputs, SingletonPrompt, TextPrompt, TokenInputs,
6 - TokensPrompt, build_explicit_enc_dec_prompt,
6 + TokensPrompt, build_explicit_enc_dec_prompt, embeds_inputs,
7 to_enc_dec_tuple_list, token_inputs, zip_enc_dec_prompts)
8 from .registry import (DummyData, InputContext, InputProcessingContext,
9 InputRegistry)
@@ -21,7 +21,9 @@ __all__ = [
21 "SingletonPrompt",
22 "ExplicitEncoderDecoderPrompt",
23 "TokenInputs",
24 + "EmbedsInputs",
25 "token_inputs",
26 + "embeds_inputs",
27 "DecoderOnlyInputs",
28 "EncoderDecoderInputs",
29 "ProcessorInputs",
vllm/inputs/data.py
+33 -3
@@ -2,6 +2,7 @@
2 from collections.abc import Iterable
3 from typing import TYPE_CHECKING, Any, Generic, Literal, Optional, Union, cast
4
5 +import torch
6 from typing_extensions import NotRequired, TypedDict, TypeVar
7
8 if TYPE_CHECKING:
@@ -63,12 +64,20 @@ class TokensPrompt(TypedDict):
64 """
65
66
66 -SingletonPrompt = Union[str, TextPrompt, TokensPrompt]
67 +class EmbedsPrompt(TypedDict):
68 + """Schema for a prompt provided via token embeddings."""
69 +
70 + prompt_embeds: torch.Tensor
71 + """The embeddings of the prompt."""
72 +
73 +
74 +SingletonPrompt = Union[str, TextPrompt, TokensPrompt, EmbedsPrompt]
75 """
76 Set of possible schemas for a single prompt:
77
78 - A text prompt (:class:`str` or :class:`TextPrompt`)
79 - A tokenized prompt (:class:`TokensPrompt`)
80 +- An embeddings prompt (:class:`EmbedsPrompt`)
81
82 Note that "singleton" is as opposed to a data structure
83 which encapsulates multiple prompts, i.e. of the sort
@@ -129,6 +138,7 @@ both decoder-only and encoder/decoder input types:
138
139 - A text prompt (:class:`str` or :class:`TextPrompt`)
140 - A tokenized prompt (:class:`TokensPrompt`)
141 +- An embeddings prompt (:class:`EmbedsPrompt`)
142 - A single data structure containing both an encoder and a decoder prompt
143 (:class:`ExplicitEncoderDecoderPrompt`)
144 """
@@ -176,7 +186,27 @@ def token_inputs(
186 return inputs
187
188
179 -DecoderOnlyInputs = Union[TokenInputs, "MultiModalInputs"]
189 +class EmbedsInputs(TypedDict):
190 + """Represents embeddings-based inputs."""
191 +
192 + type: Literal["embeds"]
193 + """The type of inputs."""
194 +
195 + prompt_embeds: torch.Tensor
196 + """The embeddings of the prompt."""
197 +
198 +
199 +def embeds_inputs(prompt_embeds: torch.Tensor) -> EmbedsInputs:
200 + """Construct :class:`EmbedsInputs` from optional values."""
201 + inputs = EmbedsInputs(
202 + type="embeds",
203 + prompt_embeds=prompt_embeds,
204 + )
205 +
206 + return inputs
207 +
208 +
209 +DecoderOnlyInputs = Union[TokenInputs, EmbedsInputs, "MultiModalInputs"]
210 """
211 The inputs in :class:`~vllm.LLMEngine` before they are
212 passed to the model executor.
@@ -198,7 +228,7 @@ class EncoderDecoderInputs(TypedDict):
228 """The inputs for the decoder portion."""
229
230
201 -SingletonInputs = Union[TokenInputs, "MultiModalInputs"]
231 +SingletonInputs = Union[TokenInputs, EmbedsInputs, "MultiModalInputs"]
232 """
233 A processed :class:`SingletonPrompt` which can be passed to
234 :class:`vllm.sequence.Sequence`.
vllm/inputs/parse.py
+48 -8
@@ -6,8 +6,9 @@ from typing_extensions import TypeIs
6
7 from vllm.utils import is_list_of
8
9 -from .data import (ExplicitEncoderDecoderPrompt, ProcessorInputs, PromptType,
10 - SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt)
9 +from .data import (EmbedsInputs, EmbedsPrompt, ExplicitEncoderDecoderPrompt,
10 + ProcessorInputs, PromptType, SingletonInputs,
11 + SingletonPrompt, TextPrompt, TokensPrompt)
12
13
14 class ParsedText(TypedDict):
@@ -84,30 +85,69 @@ class ParsedTokensPrompt(TypedDict):
85 content: TokensPrompt
86
87
88 +class ParsedEmbedsPrompt(TypedDict):
89 + type: Literal['embeds']
90 + content: EmbedsPrompt
91 +
92 +
93 +@overload
94 +def parse_singleton_prompt(prompt: str) -> ParsedStrPrompt:
95 + ...
96 +
97 +
98 +@overload
99 +def parse_singleton_prompt(prompt: TextPrompt) -> ParsedTextPrompt:
100 + ...
101 +
102 +
103 +@overload
104 +def parse_singleton_prompt(prompt: TokensPrompt) -> ParsedTokensPrompt:
105 + ...
106 +
107 +
108 +@overload
109 +def parse_singleton_prompt(prompt: EmbedsPrompt) -> ParsedEmbedsPrompt:
110 + ...
111 +
112 +
113 def parse_singleton_prompt(
114 prompt: SingletonPrompt,
89 -) -> Union[ParsedStrPrompt, ParsedTextPrompt, ParsedTokensPrompt]:
115 +) -> Union[ParsedStrPrompt, ParsedTextPrompt, ParsedTokensPrompt,
116 + ParsedEmbedsPrompt]:
117 if isinstance(prompt, str):
118 return ParsedStrPrompt(type="str", content=prompt)
119 elif isinstance(prompt, dict):
93 - if "prompt_token_ids" in prompt:
94 - return ParsedTokensPrompt(type="tokens",
95 - content=prompt) # type: ignore
120 + # Type ignores are because mypy does not correctly infer the TypedDicts
121 + # Pyright does succeed.
122 + if "prompt_embeds" in prompt:
123 + return ParsedEmbedsPrompt(
124 + type="embeds", content=prompt) # type: ignore[typeddict-item]
125 + elif "prompt_token_ids" in prompt:
126 + return ParsedTokensPrompt(
127 + type="tokens", content=prompt) # type: ignore[typeddict-item]
128 elif "prompt" in prompt:
129 return ParsedTextPrompt(type="text", content=prompt)
98 -
99 - raise TypeError("inputs must be a string, TextPrompt, or TokensPrompt")
130 + raise TypeError(
131 + "inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt")
132
133
134 def is_token_prompt(prompt: PromptType) -> TypeIs[TokensPrompt]:
135 return isinstance(prompt, dict) and "prompt_token_ids" in prompt
136
137
138 +def is_embeds_prompt(prompt: PromptType) -> TypeIs[EmbedsPrompt]:
139 + return isinstance(prompt, dict) and "prompt_embeds" in prompt
140 +
141 +
142 def is_explicit_encoder_decoder_prompt(
143 prompt: PromptType) -> TypeIs[ExplicitEncoderDecoderPrompt]:
144 return isinstance(prompt, dict) and "encoder_prompt" in prompt
145
146
147 +def is_embeds_inputs(inputs: SingletonInputs) -> TypeIs[EmbedsInputs]:
148 + return isinstance(inputs, dict) and inputs["type"] == "embeds"
149 +
150 +
151 def split_enc_dec_inputs(
152 inputs: ProcessorInputs,
153 ) -> tuple[Optional[SingletonInputs], SingletonInputs]:
vllm/inputs/preprocess.py
+63 -8
@@ -6,6 +6,7 @@ from typing import Any, Optional, Union, cast
6
7 from typing_extensions import assert_never
8
9 +from vllm import envs
10 from vllm.config import ModelConfig
11 from vllm.logger import init_logger
12 from vllm.lora.request import LoRARequest
@@ -15,9 +16,11 @@ from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalEncDecInputs,
16 from vllm.prompt_adapter.request import PromptAdapterRequest
17 from vllm.transformers_utils.tokenizer_group import TokenizerGroup
18
18 -from .data import (DecoderOnlyInputs, EncoderDecoderInputs, ProcessorInputs,
19 - PromptType, SingletonInputs, SingletonPrompt, token_inputs)
20 -from .parse import (ParsedStrPrompt, ParsedTextPrompt, ParsedTokensPrompt,
19 +from .data import (DecoderOnlyInputs, EmbedsInputs, EncoderDecoderInputs,
20 + ProcessorInputs, PromptType, SingletonInputs,
21 + SingletonPrompt, TokenInputs, embeds_inputs, token_inputs)
22 +from .parse import (ParsedEmbedsPrompt, ParsedStrPrompt, ParsedTextPrompt,
23 + ParsedTokensPrompt, is_embeds_inputs,
24 is_explicit_encoder_decoder_prompt, parse_singleton_prompt)
25
26 logger = init_logger(__name__)
@@ -328,6 +331,10 @@ class InputPreprocessor:
331 * :class:`SingletonInputs` instance
332 """
333 parsed = parse_singleton_prompt(prompt)
334 +
335 + if parsed["type"] == "embeds":
336 + return self._process_prompt_embeds(parsed)
337 +
338 prompt_text, prompt_token_ids, token_type_ids, cache_salt = \
339 self._get_prompt_data(parsed)
340
@@ -359,6 +366,8 @@ class InputPreprocessor:
366 cache_salt=cache_salt,
367 )
368
369 + assert_never(parsed)
370 +
371 async def _prompt_to_llm_inputs_async(
372 self,
373 prompt: SingletonPrompt,
@@ -369,6 +378,9 @@ class InputPreprocessor:
378 """Async version of :meth:`_extract_prompt_components`."""
379 parsed = parse_singleton_prompt(prompt)
380
381 + if parsed["type"] == "embeds":
382 + return self._process_prompt_embeds(parsed)
383 +
384 prompt_text, prompt_token_ids, token_type_ids, cache_salt = \
385 self._get_prompt_data(parsed)
386
@@ -399,10 +411,34 @@ class InputPreprocessor:
411 cache_salt=cache_salt,
412 )
413
414 + def _process_prompt_embeds(self,
415 + parsed: ParsedEmbedsPrompt) -> EmbedsInputs:
416 + if envs.VLLM_USE_V1:
417 + raise ValueError("prompt_embeds is only available in V0.")
418 +
419 + prompt_embeds_content = parsed["content"]
420 +
421 + prompt_embeds = prompt_embeds_content["prompt_embeds"]
422 +
423 + # prompt_embeds must be (seq_len, hidden_size), but if the user
424 + # passes in a batch of size 1, i.e. (1, seq_len, hidden_size),
425 + # we can unambiguously process the intent by squeezing the batch
426 + # dimension.
427 + if prompt_embeds.ndim == 3 and prompt_embeds.shape[0] == 1:
428 + prompt_embeds = prompt_embeds.squeeze(dim=0)
429 +
430 + if prompt_embeds.ndim != 2:
431 + raise ValueError(
432 + "prompt_embeds must be of shape (seq_len, hidden_size).")
433 +
434 + return embeds_inputs(prompt_embeds=prompt_embeds)
435 +
436 + assert_never(parsed)
437 +
438 def _build_enc_dec_llm_inputs(
439 self,
404 - encoder_inputs: SingletonInputs,
405 - decoder_inputs: Optional[SingletonInputs],
440 + encoder_inputs: Union[TokenInputs, MultiModalInputs],
441 + decoder_inputs: Optional[Union[TokenInputs, MultiModalInputs]],
442 ) -> EncoderDecoderInputs:
443 if (encoder_inputs["type"] == "token"
444 or encoder_inputs["type"] == "multimodal"):
@@ -410,6 +446,9 @@ class InputPreprocessor:
446 else:
447 assert_never(encoder_inputs) # type: ignore[arg-type]
448
449 + # Mypy does not correctly infer that EmbedsInputs is impossible
450 + assert "prompt_token_ids" in encoder_inputs
451 +
452 if decoder_inputs is None:
453 if self.model_config.hf_config.model_type == "whisper":
454 # For Whisper models, the text prompt should go to the decoder.
@@ -441,7 +480,8 @@ class InputPreprocessor:
480 def _separate_enc_dec_inputs_from_mm_processor_outputs(
481 self,
482 inputs: SingletonInputs,
444 - decoder_inputs_to_override: Optional[SingletonInputs] = None,
483 + decoder_inputs_to_override: Optional[Union[TokenInputs,
484 + MultiModalInputs]] = None,
485 ) -> tuple[SingletonInputs, SingletonInputs]:
486 """
487 For encoder/decoder models only:
@@ -540,6 +580,8 @@ class InputPreprocessor:
580 # For multimodal model, override decoder prompt from processor
581 # with explicit decoder prompt.
582 if self.model_config.is_multimodal_model:
583 + assert decoder_inputs is None or not is_embeds_inputs(
584 + decoder_inputs)
585 encoder_inputs, decoder_inputs = (
586 self._separate_enc_dec_inputs_from_mm_processor_outputs(
587 encoder_inputs, decoder_inputs))
@@ -555,9 +597,12 @@ class InputPreprocessor:
597 inputs))
598 else:
599 encoder_inputs = inputs
558 -
600 decoder_inputs = None
601
602 + # Mypy does not do type inference well with TypedDicts with Literal
603 + # values.
604 + assert not is_embeds_inputs(encoder_inputs)
605 + assert decoder_inputs is None or not is_embeds_inputs(decoder_inputs)
606 return self._build_enc_dec_llm_inputs(encoder_inputs, decoder_inputs)
607
608 async def _process_encoder_decoder_prompt_async(
@@ -590,6 +635,8 @@ class InputPreprocessor:
635 # For multimodal model, override decoder prompt from processor
636 # with explicit decoder prompt.
637 if self.model_config.is_multimodal_model:
638 + assert decoder_inputs is None or not is_embeds_inputs(
639 + decoder_inputs)
640 encoder_inputs, decoder_inputs = (
641 self._separate_enc_dec_inputs_from_mm_processor_outputs(
642 encoder_inputs, decoder_inputs))
@@ -605,9 +652,12 @@ class InputPreprocessor:
652 inputs))
653 else:
654 encoder_inputs = inputs
608 -
655 decoder_inputs = None
656
657 + # Mypy does not do type inference well with TypedDicts with Literal
658 + # values.
659 + assert not is_embeds_inputs(encoder_inputs)
660 + assert decoder_inputs is None or not is_embeds_inputs(decoder_inputs)
661 return self._build_enc_dec_llm_inputs(encoder_inputs, decoder_inputs)
662
663 def _build_decoder_only_llm_inputs(
@@ -617,10 +667,15 @@ class InputPreprocessor:
667 ) -> DecoderOnlyInputs:
668 if (prompt_inputs["type"] == "token"
669 or prompt_inputs["type"] == "multimodal"):
670 + # Mypy does not do type inference well with typedicts and Literal
671 + # values
672 + assert not is_embeds_inputs(prompt_inputs)
673 prompt_inputs["prompt_token_ids"] = self._apply_prompt_adapter(
674 prompt_inputs["prompt_token_ids"],
675 prompt_adapter_request=prompt_adapter_request,
676 )
677 + elif (prompt_inputs["type"] == "embeds"):
678 + pass
679 else:
680 assert_never(prompt_inputs) # type: ignore[arg-type]
681
vllm/model_executor/layers/sampler.py
+6 -1
@@ -110,6 +110,11 @@ class SamplerOutput(
110 # 'broadcasted' to all other PP ranks for next step.
111 sampled_token_ids_cpu: Optional[torch.Tensor] = None
112
113 + # On-device tensor containing the sampled token embeddings (embeddings
114 + # corresponding to the sampled token ids). Used when prompt embeddings are
115 + # specified in lieu of prompt token ids or text.
116 + sampled_token_embeds: Optional[torch.Tensor] = None
117 +
118 # Spec decode metrics populated by workers.
119 spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None
120
@@ -183,7 +188,7 @@ class Sampler(nn.Module):
188
189 # Whether or not the SamplerOutput should have on-device tensors
190 # containing the sampled token ids and probabilities. This is used by
186 - # speculative decoding.
191 + # speculative decoding and when prompt embeddings are specified.
192 self.include_gpu_probs_tensor = False
193 self.should_modify_greedy_probs_inplace = False
194
vllm/sequence.py
+84 -7
@@ -166,6 +166,9 @@ class SequenceData(msgspec.Struct,
166 _output_token_ids: array = msgspec.field(
167 default_factory=lambda: array(VLLM_TOKEN_ID_ARRAY_TYPE, []))
168
169 + _prompt_embeds: Optional[torch.Tensor] = None
170 + _output_embeds: Optional[torch.Tensor] = None
171 +
172 ### The below fields should not be passed as an argument ###
173 _cumulative_logprob: float = 0.0
174 _prompt_token_ids_tuple: tuple[int,
@@ -176,6 +179,7 @@ class SequenceData(msgspec.Struct,
179 _num_cached_tokens: int = 0
180 _stage: SequenceStage = SequenceStage.PREFILL
181 _cached_all_token_ids: list[int] = msgspec.field(default_factory=list)
182 + _cached_all_token_embeds: Optional[torch.Tensor] = None
183
184 # It is used to get delta input. It is reset when `get_delta_and_reset`
185 # is called.
@@ -208,6 +212,8 @@ class SequenceData(msgspec.Struct,
212 def from_seqs(
213 prompt_token_ids: GenericSequence[int],
214 output_token_ids: Optional[GenericSequence[int]] = None,
215 + *,
216 + prompt_embeds: Optional[torch.Tensor] = None,
217 ) -> "SequenceData":
218 """
219 Construct a :class:`SequenceData` instance from prompt and output
@@ -217,13 +223,15 @@ class SequenceData(msgspec.Struct,
223 prompt_token_ids)
224
225 if output_token_ids is None:
220 - return SequenceData(prompt_token_ids_arr)
226 + return SequenceData(prompt_token_ids_arr,
227 + _prompt_embeds=prompt_embeds)
228
229 output_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE,
230 output_token_ids)
231
232 return SequenceData(prompt_token_ids_arr,
226 - _output_token_ids=output_token_ids_arr)
233 + _output_token_ids=output_token_ids_arr,
234 + _prompt_embeds=prompt_embeds)
235
236 def __post_init__(self) -> None:
237 assert self._prompt_token_ids.typecode == "l"
@@ -231,6 +239,8 @@ class SequenceData(msgspec.Struct,
239 self._prompt_token_ids_tuple: tuple[int, ...] = tuple(
240 self._prompt_token_ids)
241 self._update_cached_all_tokens()
242 + if self._prompt_embeds is not None:
243 + self._update_cached_all_token_embeds()
244
245 def _update_cached_all_tokens(self):
246 assert isinstance(self._prompt_token_ids, array)
@@ -238,6 +248,13 @@ class SequenceData(msgspec.Struct,
248 self._cached_all_token_ids: list[int] = list(self._prompt_token_ids +
249 self._output_token_ids)
250
251 + def _update_cached_all_token_embeds(self):
252 + assert isinstance(self._prompt_embeds, torch.Tensor)
253 + self._cached_all_token_embeds: torch.Tensor = self._prompt_embeds
254 + if self._output_embeds is not None:
255 + self._cached_all_token_embeds = torch.cat(
256 + (self._cached_all_token_embeds, self._output_embeds), dim=0)
257 +
258 @property
259 def cumulative_logprob(self) -> float:
260 return self._cumulative_logprob
@@ -270,6 +287,15 @@ class SequenceData(msgspec.Struct,
287 new_output_token_ids)
288 self._update_cached_all_tokens()
289
290 + @property
291 + def output_embeds(self) -> Optional[torch.Tensor]:
292 + return self._output_embeds
293 +
294 + @output_embeds.setter
295 + def output_embeds(self, new_output_token_embeds: torch.Tensor) -> None:
296 + self._output_token_embeds = new_output_token_embeds
297 + self._update_cached_all_token_embeds()
298 +
299 @property
300 def output_token_ids_array(self) -> array:
301 """Return the prompt token ids in array type.
@@ -280,6 +306,15 @@ class SequenceData(msgspec.Struct,
306 assert isinstance(self._output_token_ids, array)
307 return self._output_token_ids
308
309 + @property
310 + def prompt_embeds(self) -> Optional[torch.Tensor]:
311 + return self._prompt_embeds
312 +
313 + @prompt_embeds.setter
314 + def prompt_embeds(self, prompt_embeds: torch.Tensor) -> None:
315 + self._prompt_embeds = prompt_embeds
316 + self._update_cached_all_token_embeds()
317 +
318 @property
319 def mrope_position_delta(self) -> Optional[int]:
320 return self._mrope_position_delta
@@ -288,11 +323,28 @@ class SequenceData(msgspec.Struct,
323 def mrope_position_delta(self, new_mrope_position_delta):
324 self._mrope_position_delta = new_mrope_position_delta
325
291 - def append_token_id(self, token_id: int, logprob: float) -> None:
326 + def append_token_id(self,
327 + token_id: int,
328 + logprob: float,
329 + token_embed: Optional[torch.Tensor] = None) -> None:
330 self._output_token_ids.append(token_id)
331 self._new_appended_tokens.append(token_id)
332 self._cached_all_token_ids.append(token_id)
333 self._cumulative_logprob += logprob
334 + if token_embed is not None:
335 + # Do not pass in with batch or sequence dimensions
336 + assert token_embed.ndim == 1
337 + token_embed = token_embed.detach().cpu().unsqueeze(0)
338 + if self._output_embeds is None:
339 + self._output_embeds = token_embed
340 + else:
341 + self._output_embeds = torch.cat(
342 + (self._output_embeds, token_embed), dim=0)
343 + assert self._cached_all_token_embeds is not None
344 + self._cached_all_token_embeds = torch.cat(
345 + (self._cached_all_token_embeds,
346 + token_embed.to(device=self._cached_all_token_embeds.device)),
347 + dim=0)
348
349 def get_len(self) -> int:
350 return len(self._output_token_ids) + len(self._prompt_token_ids)
@@ -306,6 +358,9 @@ class SequenceData(msgspec.Struct,
358 def get_token_ids(self) -> list[int]:
359 return self._cached_all_token_ids
360
361 + def get_token_embeddings(self) -> Optional[torch.Tensor]:
362 + return self._cached_all_token_embeds
363 +
364 def get_prefix_token_ids(
365 self, num_tokens: int
366 ) -> tuple[tuple[int, ...], Optional[tuple[int, ...]]]:
@@ -387,6 +442,8 @@ class SequenceData(msgspec.Struct,
442 def __repr__(self) -> str:
443 return (f"SequenceData("
444 f"prompt_token_ids={self._prompt_token_ids}, "
445 + f"prompt_embeds.shape="
446 + f"{getattr(self._prompt_embeds, 'shape', None)}, "
447 f"output_token_ids={self.output_token_ids}, "
448 f"cumulative_logprob={self.cumulative_logprob}, "
449 f"get_num_computed_tokens={self.get_num_computed_tokens()})")
@@ -425,7 +482,10 @@ class Sequence:
482 self.lora_request = lora_request
483 self.prompt_adapter_request = prompt_adapter_request
484
428 - self.data = SequenceData.from_seqs(self.prompt_token_ids)
485 + self.data = SequenceData.from_seqs(
486 + self.prompt_token_ids,
487 + prompt_embeds=self.inputs["prompt_embeds"]
488 + if self.inputs["type"] == "embeds" else None)
489 self.output_logprobs: SampleLogprobs = []
490 self.output_text = ""
491
@@ -448,14 +508,20 @@ class Sequence:
508
509 @property
510 def prompt(self) -> Optional[str]:
511 + if self.inputs["type"] == "embeds":
512 + return None
513 return self.inputs.get("prompt")
514
515 @property
516 def prompt_token_ids(self) -> list[int]:
517 + if self.inputs["type"] == "embeds":
518 + return [0] * len(self.inputs["prompt_embeds"])
519 return self.inputs["prompt_token_ids"]
520
521 @property
522 def token_type_ids(self) -> list[int]:
523 + if self.inputs["type"] == "embeds":
524 + return []
525 return self.inputs.get("token_type_ids", [])
526
527 @property
@@ -554,11 +620,14 @@ class Sequence:
620 """Reset the sequence states for recomputation."""
621 self.data.reset_state_for_recompute()
622
557 - def append_token_id(self, token_id: int, logprobs: dict[int,
558 - Logprob]) -> None:
623 + def append_token_id(self,
624 + token_id: int,
625 + logprobs: dict[int, Logprob],
626 + token_embed: Optional[torch.Tensor] = None) -> None:
627 assert token_id in logprobs
628 self.output_logprobs.append(logprobs)
561 - self.data.append_token_id(token_id, logprobs[token_id].logprob)
629 + self.data.append_token_id(token_id, logprobs[token_id].logprob,
630 + token_embed)
631
632 def get_len(self) -> int:
633 return self.data.get_len()
@@ -889,6 +958,10 @@ class SequenceGroup:
958 f"sampling_params={self.sampling_params}, "
959 f"num_seqs={len(self.seqs)})")
960
961 + def uses_prompt_embeds(self) -> bool:
962 + """Returns True if the sequence group uses input embeds."""
963 + return any(seq.data.prompt_embeds is not None for seq in self.seqs)
964 +
965
966 class SequenceGroupMetadataDelta(
967 msgspec.Struct,
@@ -1043,10 +1116,14 @@ class SequenceOutput(
1116 parent_seq_id: int
1117 output_token: int
1118 logprobs: dict[int, Logprob]
1119 + output_embed: Optional[torch.Tensor] = None
1120
1121 def __repr__(self) -> str:
1122 + output_embed_shape = \
1123 + self.output_embed.shape if self.output_embed is not None else None
1124 return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
1125 f"output_token={self.output_token}, "
1126 + f"output_embed.shape={output_embed_shape}"
1127 f"logprobs={self.logprobs})")
1128
1129 def __eq__(self, other: object) -> bool:
vllm/spec_decode/draft_model_runner.py
+14 -3
@@ -201,6 +201,9 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
201 if self.prompt_adapter_config is not None:
202 raise ValueError("TP1DraftModelRunner has no support for "
203 "prompt_adapter_config")
204 + if model_input.inputs_embeds is not None:
205 + raise ValueError("TP1DraftModelRunner has no support for "
206 + "inputs_embeds")
207 if model_input.multi_modal_kwargs:
208 raise ValueError(
209 "TP1DraftModelRunner has no support for multi_modal_kwargs"
@@ -242,9 +245,16 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
245
246 # Get model
247 if use_cuda_graph:
245 - graph_batch_size = model_input.input_tokens.shape[0]
246 - model_executable = (self.graph_runners[model_input.virtual_engine]
247 - [graph_batch_size])
248 + if model_input.inputs_embeds is None:
249 + graph_batch_size = model_input.input_tokens.shape[0]
250 + model_executable = (
251 + self.graph_runners[model_input.virtual_engine][(
252 + graph_batch_size, False)])
253 + else:
254 + graph_batch_size = model_input.inputs_embeds.shape[0]
255 + model_executable = (
256 + self.graph_runners[model_input.virtual_engine][(
257 + graph_batch_size, True)])
258
259 if previous_hidden_states is not None:
260 hidden_states = torch.cat([
@@ -281,6 +291,7 @@ class TP1DraftModelRunner(ModelRunnerWrapperBase):
291 self.vllm_config):
292 hidden_states = model_executable(
293 input_ids=model_input.input_tokens,
294 + inputs_embeds=None,
295 positions=model_input.input_positions,
296 intermediate_tensors=intermediate_tensors,
297 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
vllm/spec_decode/multi_step_worker.py
+2 -1
@@ -282,7 +282,8 @@ class MultiStepWorker(ProposerWorkerBase, DelegateWorkerBase):
282 else:
283 count += 1
284
285 - seq.append_token_id(token_id, token_logprob.logprob)
285 + seq.append_token_id(token_id, token_logprob.logprob,
286 + seq_output.output_embed)
287 seq.update_num_computed_tokens(1)
288
289 @staticmethod
vllm/worker/enc_dec_model_runner.py
+13 -4
@@ -49,6 +49,7 @@ class EncoderDecoderModelInput(ModelInputForGPUWithSamplingMetadata):
49 def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
50 tensor_dict = {
51 "input_tokens": self.input_tokens,
52 + "inputs_embeds": self.inputs_embeds,
53 "input_positions": self.input_positions,
54 "encoder_input_tokens": self.encoder_input_tokens,
55 "encoder_input_positions": self.encoder_input_positions,
@@ -172,10 +173,17 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):
173 if (model_input.attn_metadata is not None
174 and model_input.attn_metadata.prefill_metadata is None
175 and model_input.attn_metadata.decode_metadata.use_cuda_graph):
175 - assert model_input.input_tokens is not None
176 - graph_batch_size = model_input.input_tokens.shape[0]
177 - model_executable = self.graph_runners[
178 - model_input.virtual_engine][graph_batch_size]
176 + if model_input.inputs_embeds is None:
177 + assert model_input.input_tokens is not None
178 + graph_batch_size = model_input.input_tokens.shape[0]
179 + model_executable = (
180 + self.graph_runners[model_input.virtual_engine][(
181 + graph_batch_size, False)])
182 + else:
183 + graph_batch_size = model_input.inputs_embeds.shape[0]
184 + model_executable = (
185 + self.graph_runners[model_input.virtual_engine][(
186 + graph_batch_size, True)])
187 else:
188 model_executable = self.model
189
@@ -189,6 +197,7 @@ class EncoderDecoderModelRunner(GPUModelRunnerBase[EncoderDecoderModelInput]):
197 model_input.virtual_engine):
198 hidden_or_intermediate_states = model_executable(
199 input_ids=model_input.input_tokens,
200 + inputs_embeds=model_input.inputs_embeds,
201 positions=model_input.input_positions,
202 encoder_input_ids=model_input.encoder_input_tokens,
203 encoder_positions=model_input.encoder_input_positions,
vllm/worker/model_runner.py
+116 -18
@@ -35,7 +35,8 @@ from vllm.lora.request import LoRARequest
35 from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
36 from vllm.model_executor import SamplingMetadata, SamplingMetadataCache
37 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
38 -from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
38 +from vllm.model_executor.layers.sampler import (Sampler, SamplerOutput,
39 + get_sampler)
40 from vllm.model_executor.model_loader import get_model
41 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
42 from vllm.model_executor.models import supports_lora, supports_multimodal
@@ -85,6 +86,7 @@ class ModelInputForGPU(ModelRunnerInputBase):
86 additional fields.
87 """
88 input_tokens: Optional[torch.Tensor] = None
89 + inputs_embeds: Optional[torch.Tensor] = None
90 input_positions: Optional[torch.Tensor] = None
91 token_types: Optional[torch.Tensor] = None
92 seq_lens: Optional[List[int]] = None
@@ -105,6 +107,7 @@ class ModelInputForGPU(ModelRunnerInputBase):
107 def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
108 tensor_dict = {
109 "input_tokens": self.input_tokens,
110 + "inputs_embeds": self.inputs_embeds,
111 "input_positions": self.input_positions,
112 "lora_requests": self.lora_requests,
113 "lora_mapping": self.lora_mapping,
@@ -155,6 +158,7 @@ class ModelInputForGPUWithSamplingMetadata(ModelInputForGPU):
158 def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
159 tensor_dict = {
160 "input_tokens": self.input_tokens,
161 + "inputs_embeds": self.inputs_embeds,
162 "input_positions": self.input_positions,
163 "lora_requests": self.lora_requests,
164 "lora_mapping": self.lora_mapping,
@@ -194,6 +198,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
198
199 def simple_reinit(self):
200 self.input_tokens[0].clear() # type: ignore
201 + self.inputs_embeds = None # type: ignore
202 self.input_positions[0].clear() # type: ignore
203 self.token_types[0].clear() # type: ignore
204 self.mrope_input_positions = None # type: ignore
@@ -221,6 +226,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
226
227 # Input tokens and positions.
228 input_tokens: Optional[List[List[int]]] = None,
229 + inputs_embeds: Optional[torch.Tensor] = None,
230 input_positions: Optional[List[List[int]]] = None,
231 token_types: Optional[List[List[int]]] = None,
232 mrope_input_positions: Optional[List[List[List[int]]]] = None,
@@ -282,6 +288,8 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
288 for seq_id in range(len(self.seq_ids)):
289 self.input_tokens[seq_id].clear()
290
291 + self.inputs_embeds = inputs_embeds
292 +
293 if input_positions:
294 self.input_positions = input_positions
295 else:
@@ -356,6 +364,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
364
365 else:
366 self.input_tokens = input_tokens or []
367 + self.inputs_embeds = inputs_embeds
368 self.input_positions = input_positions or []
369 self.token_types = token_types or []
370 self.mrope_input_positions = mrope_input_positions or None
@@ -401,6 +410,26 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
410 self.lora_index_mapping = []
411 self.lora_prompt_mapping = []
412
413 + def __repr__(self) -> str:
414 + return (f"InterDataForSeqGroup("
415 + f"request_id={self.request_id}, "
416 + f"seq_ids={self.seq_ids}, "
417 + f"is_prompt={self.is_prompt}, "
418 + f"block_tables={self.block_tables}, "
419 + f"computed_block_nums={self.computed_block_nums}, "
420 + f"n_seqs={self.n_seqs}, "
421 + f"input_tokens={self.input_tokens}, "
422 + f"inputs_embeds.shape="
423 + f"{getattr(self.inputs_embeds, 'shape', None)}, "
424 + f"input_positions={self.input_positions}, "
425 + f"token_types={self.token_types}, "
426 + f"mrope_input_positions={self.mrope_input_positions}, "
427 + f"seq_lens={self.seq_lens}, "
428 + f"orig_seq_lens={self.orig_seq_lens}, "
429 + f"query_lens={self.query_lens}, "
430 + f"context_lens={self.context_lens}, "
431 + f"multi_modal_kwargs={self.multi_modal_kwargs}")
432 +
433 def gen_inter_data_builder(self, num_seqs: int):
434 return lambda: ModelInputForGPUBuilder.InterDataForSeqGroup(
435 request_id="",
@@ -511,13 +540,21 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
540 context_len = seq_data.get_num_computed_tokens()
541
542 # Compute tokens.
514 - tokens = seq_data.get_token_ids()[context_len:seq_len]
543 + if seq_data.prompt_embeds is None:
544 + tokens = seq_data.get_token_ids()[context_len:seq_len]
545 + prompt_embeds = None
546 + else:
547 + tokens = [0] * (seq_len - context_len)
548 + prompt_embeds = seq_data.get_token_embeddings(
549 + )[context_len:seq_len]
550 +
551 token_types = seq_group_metadata.token_type_ids
552
553 inter_data.seq_lens[seq_idx] = seq_len
554 inter_data.orig_seq_lens[seq_idx] = seq_len
555 inter_data.context_lens[seq_idx] = context_len
556 inter_data.input_tokens[seq_idx].extend(tokens)
557 + inter_data.inputs_embeds = prompt_embeds
558 inter_data.input_positions[seq_idx].extend(range(context_len, seq_len))
559 inter_data.token_types[seq_idx].extend(
560 token_types if token_types else [])
@@ -822,15 +859,29 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
859 create on-device tensors.
860 """
861 # Combine and flatten intermediate data.
825 - input_tokens = []
826 - token_types = []
862 + input_tokens = list[int]()
863 + inputs_embeds_lst = list[torch.Tensor]()
864 + token_types = list[int]()
865 for inter_data in self.inter_data_list:
866 for cur_input_tokens in inter_data.input_tokens:
867 input_tokens.extend(cur_input_tokens)
868 for cur_token_types in inter_data.token_types:
869 token_types.extend(cur_token_types)
870 + if inter_data.inputs_embeds is not None:
871 + inputs_embeds_lst.append(
872 + inter_data.inputs_embeds.to(
873 + dtype=self.runner.model_config.dtype,
874 + device=self.runner.device))
875 + inputs_embeds: Optional[torch.Tensor]
876 + if len(inputs_embeds_lst) == 0:
877 + inputs_embeds = None
878 + else:
879 + inputs_embeds = torch.cat(inputs_embeds_lst, dim=0).to(
880 + dtype=self.runner.model_config.dtype,
881 + device=self.runner.device)
882 + assert len(inputs_embeds) == len(input_tokens)
883
833 - if not input_tokens:
884 + if not input_tokens and inputs_embeds is None:
885 # This may happen when all prefill requests hit
886 # prefix caching and there is no decode request.
887 return self.model_input_cls()
@@ -980,6 +1031,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
1031
1032 return self.model_input_cls(
1033 input_tokens=input_tokens_tensor,
1034 + inputs_embeds=inputs_embeds,
1035 input_positions=input_positions_tensor,
1036 token_types=token_types_tensor,
1037 attn_metadata=attn_metadata,
@@ -1029,7 +1081,8 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1081 self.max_batchsize_to_capture = \
1082 self.vllm_config.compilation_config.max_capture_size
1083
1032 - self.graph_runners: List[Dict[int, CUDAGraphRunner]] = [
1084 + #
1085 + self.graph_runners: List[Dict[Tuple[int, bool], CUDAGraphRunner]] = [
1086 {} for _ in range(self.parallel_config.pipeline_parallel_size)
1087 ]
1088 self.graph_memory_pool: Optional[Tuple[
@@ -1466,6 +1519,10 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1519 input_positions = torch.zeros(max_batch_size,
1520 dtype=torch.long,
1521 device=self.device)
1522 + inputs_embeds = torch.zeros(
1523 + (max_batch_size, self.model_config.get_hidden_size()),
1524 + dtype=self.model_config.dtype,
1525 + device=self.device)
1526 if self.model_config.uses_mrope:
1527 input_positions = torch.tile(input_positions,
1528 (3, 1)).cuda(device=self.device)
@@ -1503,15 +1560,22 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1560 # memory usage of CUDA graph.
1561 for virtual_engine in range(
1562 self.parallel_config.pipeline_parallel_size):
1506 - # Only rank 0 should print progress bar during capture
1507 - cudagraph_capture_sizes = (tqdm(
1508 - self.vllm_config.compilation_config.
1563 + # We need to not only iterate over batch sizes, but also whether
1564 + # to use inputs_embeds or not, hence we use the cartesian
1565 + # product.
1566 + cudagraph_capture_sizes = self.vllm_config.compilation_config\
1567 + .cudagraph_capture_sizes
1568 + cudagraph_inputs_embeds = (True, False)
1569 + compilation_cases = itertools.product(
1570 cudagraph_capture_sizes,
1510 - desc="Capturing CUDA graph shapes",
1511 - ) if get_tensor_model_parallel_rank() == 0 else
1512 - self.vllm_config.compilation_config.
1513 - cudagraph_capture_sizes)
1514 - for batch_size in cudagraph_capture_sizes:
1571 + cudagraph_inputs_embeds,
1572 + )
1573 + # Only rank 0 should print progress bar during capture
1574 + if get_tensor_model_parallel_rank() == 0:
1575 + compilation_cases = tqdm(
1576 + list(compilation_cases),
1577 + desc="Capturing CUDA graph shapes")
1578 + for batch_size, use_inputs_embeds in compilation_cases:
1579 attn_metadata = (
1580 self.attn_state.graph_capture_get_metadata_for_batch(
1581 batch_size,
@@ -1542,6 +1606,9 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1606 capture_inputs = {
1607 "input_ids":
1608 input_tokens[:batch_size],
1609 + "inputs_embeds":
1610 + inputs_embeds[:batch_size]
1611 + if use_inputs_embeds else None,
1612 "positions":
1613 input_positions[..., :batch_size],
1614 "intermediate_inputs":
@@ -1578,8 +1645,8 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1645 virtual_engine):
1646 graph_runner.capture(**capture_inputs)
1647 self.graph_memory_pool = graph_runner.graph.pool()
1581 - self.graph_runners[virtual_engine][batch_size] = (
1582 - graph_runner)
1648 + self.graph_runners[virtual_engine][(
1649 + batch_size, use_inputs_embeds)] = graph_runner
1650
1651 if self.lora_config:
1652 self._remove_dummy_loras()
@@ -1711,8 +1778,9 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1778 if prefill_meta is None and decode_meta.use_cuda_graph:
1779 assert model_input.input_tokens is not None
1780 graph_batch_size = model_input.input_tokens.shape[0]
1714 - model_executable = self.graph_runners[virtual_engine][
1715 - graph_batch_size]
1781 + use_inputs_embeds = model_input.inputs_embeds is not None
1782 + model_executable = self.graph_runners[virtual_engine][(
1783 + graph_batch_size, use_inputs_embeds)]
1784 if previous_hidden_states is not None:
1785 previous_hidden_states = torch.cat([
1786 previous_hidden_states,
@@ -1763,6 +1831,7 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1831 self.vllm_config, virtual_engine):
1832 hidden_or_intermediate_states = model_executable(
1833 input_ids=model_input.input_tokens,
1834 + inputs_embeds=model_input.inputs_embeds,
1835 positions=model_input.input_positions,
1836 intermediate_tensors=intermediate_tensors,
1837 **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
@@ -1817,6 +1886,11 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1886 model_input.async_callback()
1887
1888 # Sample the next token.
1889 + assert isinstance(self.sampler, Sampler)
1890 + orig_include_gpu_probs_tensor = self.sampler.include_gpu_probs_tensor
1891 + if model_input.inputs_embeds is not None:
1892 + self.sampler.include_gpu_probs_tensor = True
1893 +
1894 output: SamplerOutput = self.sampler(
1895 logits=logits,
1896 sampling_metadata=model_input.sampling_metadata,
@@ -1838,6 +1912,18 @@ class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):
1912 output.model_forward_time = (orig_model_forward_time +
1913 model_forward_time)
1914
1915 + if model_input.inputs_embeds is not None:
1916 + self.sampler.include_gpu_probs_tensor = \
1917 + orig_include_gpu_probs_tensor
1918 + if output.sampled_token_ids is not None:
1919 + output.sampled_token_embeds = self.model.get_input_embeddings(
1920 + output.sampled_token_ids.squeeze(1))
1921 +
1922 + for token_embed, sequence_group_output in zip(
1923 + output.sampled_token_embeds, output.outputs):
1924 + assert len(sequence_group_output.samples) == 1
1925 + sequence_group_output.samples[0].output_embed = token_embed
1926 +
1927 if self.return_hidden_states:
1928 # we only need to pass hidden states of most recent token
1929 assert model_input.sampling_metadata is not None
@@ -1931,6 +2017,7 @@ class CUDAGraphRunner(nn.Module):
2017 def capture(
2018 self,
2019 input_ids: torch.Tensor,
2020 + inputs_embeds: Optional[torch.Tensor],
2021 positions: torch.Tensor,
2022 intermediate_inputs: Optional[IntermediateTensors],
2023 kv_caches: List[torch.Tensor],
@@ -1947,6 +2034,7 @@ class CUDAGraphRunner(nn.Module):
2034 for _ in range(_NUM_WARMUP_ITERS):
2035 self.model(
2036 input_ids=input_ids,
2037 + inputs_embeds=inputs_embeds,
2038 positions=positions,
2039 intermediate_tensors=intermediate_inputs,
2040 **kwargs,
@@ -1959,6 +2047,9 @@ class CUDAGraphRunner(nn.Module):
2047 with torch.cuda.graph(self._graph, pool=memory_pool, stream=stream):
2048 output_hidden_or_intermediate_states = self.model(
2049 input_ids=input_ids,
2050 + **({
2051 + "inputs_embeds": inputs_embeds,
2052 + } if inputs_embeds is not None else {}),
2053 positions=positions,
2054 intermediate_tensors=intermediate_inputs,
2055 **kwargs,
@@ -1986,6 +2077,9 @@ class CUDAGraphRunner(nn.Module):
2077 self.input_buffers = {
2078 "input_ids":
2079 input_ids,
2080 + **({
2081 + "inputs_embeds": inputs_embeds,
2082 + } if inputs_embeds is not None else {}),
2083 "positions":
2084 positions,
2085 "kv_caches":
@@ -2006,6 +2100,7 @@ class CUDAGraphRunner(nn.Module):
2100 def forward(
2101 self,
2102 input_ids: torch.Tensor,
2103 + inputs_embeds: Optional[torch.Tensor],
2104 positions: torch.Tensor,
2105 intermediate_tensors: Optional[IntermediateTensors],
2106 **kwargs,
@@ -2020,6 +2115,9 @@ class CUDAGraphRunner(nn.Module):
2115 # so the shape is not padded, we need to copy partial only
2116 self.input_buffers["positions"][:positions.shape[0]].copy_(
2117 positions, non_blocking=True)
2118 + if inputs_embeds is not None:
2119 + self.input_buffers["inputs_embeds"][:inputs_embeds.shape[0]].copy_(
2120 + inputs_embeds, non_blocking=True)
2121
2122 if self.backend_name != "NO_ATTENTION":
2123 self.input_buffers["slot_mapping"].copy_(
vllm/worker/pooling_model_runner.py
+11 -4
@@ -84,10 +84,17 @@ class PoolingModelRunner(
84 # explore how to leverage it.
85 if (prefill_meta is None and decode_meta is not None
86 and decode_meta.use_cuda_graph):
87 - assert model_input.input_tokens is not None
88 - graph_batch_size = model_input.input_tokens.shape[0]
89 - model_executable = self.graph_runners[virtual_engine][
90 - graph_batch_size]
87 + if model_input.inputs_embeds is None:
88 + assert model_input.input_tokens is not None
89 + graph_batch_size = model_input.input_tokens.shape[0]
90 + model_executable = (
91 + self.graph_runners[model_input.virtual_engine][(
92 + graph_batch_size, False)])
93 + else:
94 + graph_batch_size = model_input.inputs_embeds.shape[0]
95 + model_executable = (
96 + self.graph_runners[model_input.virtual_engine][(
97 + graph_batch_size, True)])
98 else:
99 model_executable = self.model
100