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, 13:30 UTC
[TPU][V1] Capture multimodal encoder during model compilation (#15051)
Signed-off-by: Michael Goin <mgoin64@gmail.com> Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Siyuan Liu <lsiyuan@google.com>
Nicolò Lucchesi committed
Apr 22, 2025 at 02:36 UTC
210207525e085b111a304089f4413fb5056263b4
4 files changed
+327
-48
.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
+3
-1
@@ -17,7 +17,7 @@ source /etc/environment
17
docker run --privileged --net host --shm-size=16G -it \
18
-e "HF_TOKEN=$HF_TOKEN" --name tpu-test \
19
vllm-tpu /bin/bash -c "python3 -m pip install git+https://github.com/thuml/depyf.git \
20
- && python3 -m pip install pytest tpu-info \
20
+ && python3 -m pip install pytest pytest-asyncio tpu-info \
21
&& python3 -m pip install lm_eval[api]==0.4.4 \
22
&& export VLLM_USE_V1=1 \
23
&& export VLLM_XLA_CHECK_RECOMPILATION=1 \
@@ -42,6 +42,8 @@ docker run --privileged --net host --shm-size=16G -it \
42
&& echo TEST_8 \
43
&& pytest -s -v /workspace/vllm/tests/v1/tpu/test_topk_topp_sampler.py \
44
&& echo TEST_9 \
45
+ && pytest -s -v /workspace/vllm/tests/v1/tpu/test_multimodal.py \
46
+ && echo TEST_10 \
47
&& pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py" \
48
49
tests/v1/tpu/test_multimodal.py
new
+91
@@ -0,0 +1,91 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import openai
4
+import pytest
5
+
6
+from vllm import envs
7
+from vllm.multimodal.utils import encode_image_base64, fetch_image
8
+from vllm.platforms import current_platform
9
+
10
+from ...entrypoints.openai.test_vision import TEST_IMAGE_URLS
11
+from ...utils import RemoteOpenAIServer
12
+
13
+if not envs.VLLM_USE_V1:
14
+ pytest.skip(
15
+ "Skipping V1 tests. Rerun with `VLLM_USE_V1=1` to test.",
16
+ allow_module_level=True,
17
+ )
18
+
19
+
20
+@pytest.fixture(scope="session")
21
+def base64_encoded_image() -> dict[str, str]:
22
+ return {
23
+ image_url: encode_image_base64(fetch_image(image_url))
24
+ for image_url in TEST_IMAGE_URLS
25
+ }
26
+
27
+
28
+@pytest.mark.asyncio
29
+@pytest.mark.skipif(not current_platform.is_tpu(),
30
+ reason="This test needs a TPU")
31
+@pytest.mark.parametrize("model_name", ["llava-hf/llava-1.5-7b-hf"])
32
+async def test_basic_vision(model_name: str, base64_encoded_image: dict[str,
33
+ str]):
34
+
35
+ def whats_in_this_image_msg(b64):
36
+ return [{
37
+ "role":
38
+ "user",
39
+ "content": [
40
+ {
41
+ "type": "text",
42
+ "text": "What's in this image?"
43
+ },
44
+ {
45
+ "type": "image_url",
46
+ "image_url": {
47
+ "url": f"data:image/jpeg;base64,{b64}"
48
+ },
49
+ },
50
+ ],
51
+ }]
52
+
53
+ server_args = [
54
+ "--max-model-len",
55
+ "1024",
56
+ "--max-num-seqs",
57
+ "16",
58
+ "--gpu-memory-utilization",
59
+ "0.95",
60
+ "--trust-remote-code",
61
+ "--max-num-batched-tokens",
62
+ "576",
63
+ # NOTE: max-num-batched-tokens>=mm_item_size
64
+ "--disable_chunked_mm_input",
65
+ "--chat-template",
66
+ "examples/template_llava.jinja"
67
+ ]
68
+
69
+ # Server will pre-compile on first startup (takes a long time).
70
+ with RemoteOpenAIServer(model_name, server_args,
71
+ max_wait_seconds=600) as remote_server:
72
+ client: openai.AsyncOpenAI = remote_server.get_async_client()
73
+
74
+ # Other requests now should be much faster
75
+ for image_url in TEST_IMAGE_URLS:
76
+ image_base64 = base64_encoded_image[image_url]
77
+ chat_completion_from_base64 = await client.chat.completions\
78
+ .create(
79
+ model=model_name,
80
+ messages=whats_in_this_image_msg(image_base64),
81
+ max_completion_tokens=24,
82
+ temperature=0.0)
83
+ result = chat_completion_from_base64
84
+ assert result
85
+ choice = result.choices[0]
86
+ assert choice.finish_reason == "length"
87
+
88
+ message = choice.message
89
+ message = result.choices[0].message
90
+ assert message.content is not None and len(message.content) >= 10
91
+ assert message.role == "assistant"
vllm/v1/worker/tpu_model_runner.py
+232
-46
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
import bisect
3
+import gc
4
import time
5
from typing import TYPE_CHECKING, Optional, cast
6
from unittest.mock import patch
@@ -21,7 +22,8 @@ from vllm.forward_context import set_forward_context
22
from vllm.logger import init_logger
23
from vllm.model_executor.model_loader import get_model
24
from vllm.multimodal import MULTIMODAL_REGISTRY
24
-from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
25
+from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,
26
+ PlaceholderRange)
27
from vllm.multimodal.utils import group_mm_inputs_by_modality
28
from vllm.sequence import IntermediateTensors
29
from vllm.utils import LayerBlockType, cdiv, is_pin_memory_available
@@ -37,8 +39,7 @@ from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler
39
from vllm.v1.utils import bind_kv_cache
40
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
41
40
-from .utils import (gather_mm_placeholders, sanity_check_mm_encoder_outputs,
41
- scatter_mm_placeholders)
42
+from .utils import sanity_check_mm_encoder_outputs
43
44
if TYPE_CHECKING:
45
from vllm.v1.core.sched.output import SchedulerOutput
@@ -198,7 +199,7 @@ class TPUModelRunner:
199
device="cpu")
200
self.slot_mapping_np = self.slot_mapping_cpu.numpy()
201
self.block_table_cpu = torch.zeros(
201
- (self.max_num_tokens, self.max_num_blocks_per_req),
202
+ (self.max_num_reqs, self.max_num_blocks_per_req),
203
dtype=self.input_batch.block_table.get_cpu_tensor().dtype,
204
device="cpu")
205
@@ -220,6 +221,37 @@ class TPUModelRunner:
221
self.num_reqs_paddings = _get_req_paddings(
222
min_req_size=MIN_NUM_SEQS, max_req_size=self.max_num_reqs)
223
224
+ # Get maximum number of mm items per modality (batch size).
225
+ self.max_num_mm_items_by_modality = dict()
226
+ if (self.is_multimodal_model and self.max_num_encoder_input_tokens > 0
227
+ and self.encoder_cache_size > 0):
228
+ max_tokens_by_modality_dict = (
229
+ MULTIMODAL_REGISTRY.
230
+ get_max_tokens_per_item_by_nonzero_modality(self.model_config))
231
+ for modality, max_tokens in max_tokens_by_modality_dict.items():
232
+ # Check how many items of this modality can be supported by
233
+ # the encoder budget.
234
+ encoder_budget = min(self.max_num_encoder_input_tokens,
235
+ self.encoder_cache_size)
236
+
237
+ max_num_mm_items_encoder_budget = cdiv(encoder_budget,
238
+ max_tokens)
239
+
240
+ # Check how many items of this modality can be supported by
241
+ # the decoder budget.
242
+ max_mm_items_per_req = self.mm_registry.\
243
+ get_mm_limits_per_prompt(self.model_config)[modality]
244
+
245
+ # NOTE: We do not consider max_num_batched_tokens on purpose
246
+ # because the multimodal embeddings can be generated in advance
247
+ # and chunked prefilled.
248
+ max_num_mm_items_decoder_budget = self.max_num_reqs * \
249
+ max_mm_items_per_req
250
+
251
+ max_num_mm_items = min(max_num_mm_items_encoder_budget,
252
+ max_num_mm_items_decoder_budget)
253
+ self.max_num_mm_items_by_modality[modality] = max_num_mm_items
254
+
255
def _update_num_xla_graphs(self, case_str):
256
check_comp = self.check_recompilation and not self.enforce_eager
257
if not check_comp:
@@ -606,29 +638,36 @@ class TPUModelRunner:
638
# 2. A list or tuple (length: num_items) of tensors, each of shape
639
# (feature_size, hidden_size) in case the feature size is dynamic
640
# depending on the input multimodal items.
641
+ xm.mark_step()
642
curr_group_outputs = self.model.get_multimodal_embeddings(
643
**batched_mm_inputs)
644
+ xm.mark_step()
645
646
sanity_check_mm_encoder_outputs(
647
curr_group_outputs,
648
expected_num_items=len(grouped_mm_inputs),
649
)
650
617
- for output in curr_group_outputs:
618
- encoder_outputs.append(output)
651
+ if isinstance(curr_group_outputs, torch.Tensor):
652
+ encoder_outputs.append(curr_group_outputs)
653
+ else:
654
+ assert isinstance(curr_group_outputs, (list, tuple))
655
+ for output in curr_group_outputs:
656
+ encoder_outputs.append(output)
657
658
# Cache the encoder outputs.
659
+ # NOTE (NickLucche) here we diverge from logic in other runners, as we
660
+ # assume to only have whole mm items to process. Hence we avoid the
661
+ # intrinsic dynamism that `scatter_mm_placeholders` introduces.
662
for (req_id, input_id, pos_info), output in zip(
663
req_ids_pos,
664
encoder_outputs,
665
):
666
if req_id not in self.encoder_cache:
667
self.encoder_cache[req_id] = {}
627
-
628
- self.encoder_cache[req_id][input_id] = scatter_mm_placeholders(
629
- output,
630
- is_embed=pos_info.is_embed,
631
- )
668
+ assert pos_info.is_embed is None, "Expected all positions to be"\
669
+ " contiguous and embeddings."
670
+ self.encoder_cache[req_id][input_id] = output
671
672
def _gather_mm_embeddings(
673
self,
@@ -641,6 +680,10 @@ class TPUModelRunner:
680
req_state = self.requests[req_id]
681
num_computed_tokens = req_state.num_computed_tokens
682
mm_positions = req_state.mm_positions
683
+ # TODO unroll loop and assume/enforce --disable_chunked_mm_input
684
+ # NOTE (NickLucche) here we diverge from logic in other runners, as
685
+ # we assume to only have whole mm items to process. Hence we avoid
686
+ # the intrinsic dynamism that `gather_mm_placeholders` introduces.
687
for i, pos_info in enumerate(mm_positions):
688
start_pos = pos_info.offset
689
num_encoder_tokens = pos_info.length
@@ -657,25 +700,33 @@ class TPUModelRunner:
700
# in the decoder's KV cache.
701
continue
702
660
- start_idx = max(num_computed_tokens - start_pos, 0)
661
- end_idx = min(
662
- num_computed_tokens - start_pos + num_scheduled_tokens,
663
- num_encoder_tokens)
664
- assert start_idx < end_idx
703
assert req_id in self.encoder_cache
704
assert i in self.encoder_cache[req_id]
705
+ assert pos_info.is_embed is None, "Expected all positions to"\
706
+ " be contiguous and embeddings."
707
encoder_output = self.encoder_cache[req_id][i]
668
-
669
- if (is_embed := pos_info.is_embed) is not None:
670
- is_embed = is_embed[start_idx:end_idx]
671
-
672
- mm_embeds_item = gather_mm_placeholders(
673
- encoder_output[start_idx:end_idx],
674
- is_embed=is_embed,
675
- )
676
- mm_embeds.append(mm_embeds_item)
708
+ mm_embeds.append(encoder_output)
709
return mm_embeds
710
711
+ def _get_model_inputs(self, input_ids: torch.Tensor,
712
+ mm_embeds: list[torch.Tensor]):
713
+ if self.is_multimodal_model:
714
+ # NOTE(woosuk): To unify token ids and soft tokens (vision
715
+ # embeddings), we always use embeddings (rather than token ids)
716
+ # as input to the multimodal model, even when the input is text.
717
+ if mm_embeds:
718
+ inputs_embeds = self.model.get_input_embeddings(
719
+ input_ids, mm_embeds)
720
+ else:
721
+ inputs_embeds = self.model.get_input_embeddings(input_ids)
722
+ return None, inputs_embeds
723
+ else:
724
+ # For text-only models, we use token ids as input.
725
+ # While it is possible to use embeddings as input just like the
726
+ # multimodal models, it is not desirable for performance since
727
+ # then the embedding layer is not included in the CUDA graph.
728
+ return input_ids, None
729
+
730
@torch.no_grad()
731
def execute_model(
732
self,
@@ -694,27 +745,13 @@ class TPUModelRunner:
745
mm_embeds = self._gather_mm_embeddings(scheduler_output)
746
else:
747
mm_embeds = []
697
-
748
+ xm.mark_step()
749
# Prepare inputs
750
attn_metadata, logits_indices, padded_num_reqs = self._prepare_inputs(
751
scheduler_output)
701
- if self.is_multimodal_model:
702
- # NOTE(woosuk): To unify token ids and soft tokens (vision
703
- # embeddings), we always use embeddings (rather than token ids)
704
- # as input to the multimodal model, even when the input is text.
705
- if mm_embeds:
706
- inputs_embeds = self.model.get_input_embeddings(
707
- self.input_ids, mm_embeds)
708
- else:
709
- inputs_embeds = self.model.get_input_embeddings(self.input_ids)
710
- input_ids = None
711
- else:
712
- # For text-only models, we use token ids as input.
713
- # While it is possible to use embeddings as input just like the
714
- # multimodal models, it is not desirable for performance since
715
- # then the embedding layer is not included in the CUDA graph.
716
- input_ids = self.input_ids
717
- inputs_embeds = None
752
+ input_ids, inputs_embeds = self._get_model_inputs(
753
+ self.input_ids, mm_embeds)
754
+ xm.mark_step()
755
num_reqs = self.input_batch.num_reqs
756
# Run the decoder
757
with set_forward_context(attn_metadata, self.vllm_config):
@@ -890,9 +927,70 @@ class TPUModelRunner:
927
inputs_embeds=inputs_embeds)
928
self._hidden_states_dtype = out.dtype
929
930
+ def _precompile_mm_encoder(self) -> None:
931
+ # Pre-compile MM encoder for all supported data modalities.
932
+ hf_config = self.vllm_config.model_config.hf_config
933
+ for mode, max_items_by_mode in \
934
+ self.max_num_mm_items_by_modality.items():
935
+ logger.info(
936
+ "Compiling Multimodal %s Encoder with different input"
937
+ " shapes.", mode)
938
+ start = time.perf_counter()
939
+ # No padding for MM encoder just yet.
940
+ for num_items in range(1, max_items_by_mode + 1):
941
+ logger.info(" -- mode: %s items: %d", mode, num_items)
942
+ batched_dummy_mm_inputs = self._get_mm_dummy_batch(
943
+ mode, num_items)
944
+ # Run multimodal encoder.
945
+ xm.mark_step()
946
+ mm_embeds = self.model.\
947
+ get_multimodal_embeddings(**batched_dummy_mm_inputs)
948
+ xm.mark_step()
949
+ num_patches = mm_embeds[0].shape[0]
950
+ items_size = num_patches * num_items
951
+
952
+ # NOTE (NickLucche) pre-compile `get_input_embeddings` when mm
953
+ # embeddings are present. We assume `--disable-mm-chunked`,
954
+ # hence only whole items can be scheduled. This implies we just
955
+ # need to compile when `num_items` fit the (padded) `input_ids`
956
+ for num_tokens in self.num_tokens_paddings:
957
+ if num_tokens >= items_size:
958
+ # XLA Workaround: if torch.zeros(..device) is used, XLA
959
+ # compiles a scalar+expansion op, which won't match
960
+ # the graph generated at runtime. CPU->TPU must be used
961
+ placeholders_ids = torch.zeros(num_tokens,
962
+ dtype=torch.int32,
963
+ device="cpu")
964
+ # Align placeholders and actual num mm_embeddings.
965
+ placeholders_ids[:items_size] = \
966
+ hf_config.image_token_index
967
+
968
+ placeholders_ids = placeholders_ids.to(self.device)
969
+ # Assign outputs or the graph will be cut short.
970
+ a, b = self._get_model_inputs(placeholders_ids,
971
+ [mm_embeds])
972
+ assert a is None
973
+ xm.mark_step()
974
+
975
+ # Pre-compile `get_input_embeddings` when mm_embeddings are not
976
+ # present. Chunk is only made of text, no mm_placeholders.
977
+ for num_tokens in self.num_tokens_paddings:
978
+ placeholders_ids = torch.zeros(num_tokens,
979
+ dtype=torch.int32,
980
+ device="cpu")
981
+ placeholders_ids = placeholders_ids.to(self.device)
982
+ a, b = self._get_model_inputs(placeholders_ids, [])
983
+ assert a is None
984
+ xm.mark_step()
985
+
986
+ xm.wait_device_ops()
987
+ end = time.perf_counter()
988
+ logger.info(
989
+ "Multimodal %s Encoder compilation finished in in %.2f "
990
+ "[secs].", mode, end - start)
991
+
992
def _precompile_backbone(self) -> None:
993
logger.info("Compiling the model with different input shapes.")
895
-
994
start = time.perf_counter()
995
for num_tokens in self.num_tokens_paddings:
996
logger.info(" -- num_tokens: %d", num_tokens)
@@ -962,11 +1060,70 @@ class TPUModelRunner:
1060
"""
1061
Precompile all the subgraphs with possible input shapes.
1062
"""
965
- # TODO: precompile encoder
1063
+ self._precompile_mm_encoder()
1064
self._precompile_backbone()
1065
self._precompile_select_hidden_states()
1066
self._precompile_sample_from_hidden()
1067
1068
+ def profile_run(
1069
+ self,
1070
+ num_tokens: int,
1071
+ ) -> None:
1072
+ # Profile with multimodal encoder & encoder cache.
1073
+ # TODO: handle encoder-decoder models once we support them.
1074
+ if (self.is_multimodal_model and self.max_num_encoder_input_tokens > 0
1075
+ and self.encoder_cache_size > 0):
1076
+
1077
+ # NOTE: Currently model is profiled with a single non-text
1078
+ # modality with the max possible input tokens even when
1079
+ # it supports multiple.
1080
+ dummy_data_modality, max_num_mm_items = max(
1081
+ self.max_num_mm_items_by_modality.items(), key=lambda t: t[1])
1082
+
1083
+ encoder_budget = min(self.max_num_encoder_input_tokens,
1084
+ self.encoder_cache_size)
1085
+
1086
+ logger.info(
1087
+ "Encoder cache will be initialized with a budget of %d tokens,"
1088
+ " and profiled with %s %s items of the maximum feature size.",
1089
+ encoder_budget, max_num_mm_items, dummy_data_modality)
1090
+
1091
+ # Create dummy batch of multimodal inputs.
1092
+ batched_dummy_mm_inputs = self._get_mm_dummy_batch(
1093
+ dummy_data_modality, max_num_mm_items)
1094
+
1095
+ # Run multimodal encoder.
1096
+ # Isolate encoder graph from post-processing to minimize
1097
+ # impact of recompilation until it's fixed.
1098
+ start = time.perf_counter()
1099
+ xm.mark_step()
1100
+ dummy_encoder_outputs = self.model.get_multimodal_embeddings(
1101
+ **batched_dummy_mm_inputs)
1102
+ xm.mark_step()
1103
+ xm.wait_device_ops()
1104
+ end = time.perf_counter()
1105
+ logger.info(
1106
+ "Multimodal Encoder profiling finished in in %.2f [secs].",
1107
+ end - start)
1108
+
1109
+ assert len(dummy_encoder_outputs) == max_num_mm_items, (
1110
+ "Expected dimension 0 of encoder outputs to match the number "
1111
+ f"of multimodal data items: {max_num_mm_items}, got "
1112
+ f"{len(dummy_encoder_outputs)=} instead. This is most likely "
1113
+ "due to the 'get_multimodal_embeddings' method of the model "
1114
+ "not implemented correctly.")
1115
+
1116
+ # Cache the dummy encoder outputs.
1117
+ self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
1118
+
1119
+ # Trigger compilation for general shape.
1120
+ self._dummy_run(num_tokens)
1121
+
1122
+ xm.mark_step()
1123
+ xm.wait_device_ops()
1124
+ self.encoder_cache.clear()
1125
+ gc.collect()
1126
+
1127
def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
1128
"""
1129
Initialize KV cache based on `kv_cache_config`.
@@ -1045,6 +1202,36 @@ class TPUModelRunner:
1202
def get_input_embeddings(self, *args, **kwargs):
1203
return self.model.get_input_embeddings(*args, **kwargs)
1204
1205
+ def _get_mm_dummy_batch(self, modality: str,
1206
+ batch_size: int) -> BatchedTensorInputs:
1207
+ # Dummy data for pre-compiling multimodal models.
1208
+ dummy_request_data = self.mm_registry.get_decoder_dummy_data(
1209
+ model_config=self.model_config,
1210
+ seq_len=self.max_num_tokens,
1211
+ )
1212
+ dummy_mm_data = dummy_request_data.multi_modal_data
1213
+
1214
+ # Dummy data definition in V0 may contain multiple multimodal items
1215
+ # (e.g, multiple images) for a single request, therefore here we
1216
+ # always replicate first item by max_num_mm_items times since in V1
1217
+ # they are scheduled to be processed separately.
1218
+ assert isinstance(dummy_mm_data, MultiModalKwargs), (
1219
+ "Expected dummy multimodal data to be of type "
1220
+ f"MultiModalKwargs, got {type(dummy_mm_data)=} instead. "
1221
+ "This is most likely due to the model not having a merged "
1222
+ "processor.")
1223
+
1224
+ # When models have a merged processor, their dummy data is
1225
+ # already batched `MultiModalKwargs`, therefore we take the first
1226
+ # `MultiModalKwargsItem` from the desired modality to profile on.
1227
+ dummy_mm_item = dummy_mm_data.get_item(modality=modality, item_index=0)
1228
+ dummy_mm_kwargs = MultiModalKwargs.from_items([dummy_mm_item])
1229
+
1230
+ batched_dummy_mm_inputs = MultiModalKwargs.batch([dummy_mm_kwargs] *
1231
+ batch_size)
1232
+ return MultiModalKwargs.as_kwargs(batched_dummy_mm_inputs,
1233
+ device=self.device)
1234
+
1235
1236
def _get_req_paddings(min_req_size: int, max_req_size: int) -> list[int]:
1237
logger.info("Preparing request paddings:")
@@ -1088,7 +1275,6 @@ def _get_token_paddings(min_token_size: int, max_token_size: int,
1275
if num >= max_token_size:
1276
break
1277
num *= 2
1091
-
1278
else:
1279
logger.info("Using incremental token paddings:")
1280
while num <= padding_gap:
vllm/v1/worker/tpu_worker.py
+1
-1
@@ -157,7 +157,7 @@ class TPUWorker:
157
runner_kv_caches)
158
159
# `max_num_tokens >= max_num_batched_tokens` due to padding.
160
- self.model_runner._dummy_run(self.model_runner.max_num_tokens)
160
+ self.model_runner.profile_run(self.model_runner.max_num_tokens)
161
162
# Synchronize before measuring the memory usage.
163
xm.wait_device_ops()