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, 14:45 UTC
[TPU][V1] Refine tpu_model_runner to mitigate future recompilation issues (#16275)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Apr 9, 2025 at 17:51 UTC
a454748544f67a7677a2bf71ab329da0600d34a6
4 files changed
+165
-124
tests/tpu/test_compilation.py
+10
-4
@@ -44,7 +44,7 @@ def test_tpu_compilation():
44
assert generated_text.startswith(answer)
45
46
compiled_codes = sorted(
47
- glob.glob(os.path.join(temp_dir, "__transformed_code*.py")))
47
+ glob.glob(os.path.join(temp_dir, "__transformed_code*for_forward.py")))
48
49
for i, compiled_code in enumerate(compiled_codes):
50
print("{} file: {}".format(i + 1, compiled_code))
@@ -52,15 +52,21 @@ def test_tpu_compilation():
52
# We should only trigger Dynamo compilation 2 times:
53
# 1. Forward pass without kv_caches
54
# 2. Forward pass with kv_caches
55
- # Check we have 4 compiled codes
55
+ # Check we have 2 compiled codes
56
assert len(compiled_codes) == 2
57
58
kv_cache_prefix = "kv_cache"
59
attn_prefix = "ragged_paged_attention"
60
61
+ def extract_compiled_index(s):
62
+ parts = s.replace(".", "_").split("_")
63
+ numbers = [int(part) for part in parts if part.isdigit()]
64
+ return numbers[0]
65
+
66
# Check all the compilations are as expected
62
- compiled_fns = sorted(
63
- glob.glob(os.path.join(temp_dir, "__compiled_fn*Captured*.py")))
67
+ compiled_fns = sorted(glob.glob(
68
+ os.path.join(temp_dir, "__compiled_fn*Captured*.py")),
69
+ key=lambda s: extract_compiled_index(s))
70
71
for i, compiled_fn in enumerate(compiled_fns):
72
print("{} file: {}".format(i + 1, compiled_fn))
tests/v1/tpu/worker/test_tpu_model_runner.py
+19
-6
@@ -7,9 +7,9 @@ from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
7
from vllm.sampling_params import SamplingParams
8
from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
9
SchedulerOutput)
10
-from vllm.v1.worker.tpu_model_runner import (TPUModelRunner,
11
- _get_padded_token_len,
12
- _get_paddings)
10
+from vllm.v1.worker.tpu_model_runner import (
11
+ TPUModelRunner, _get_padded_num_reqs_with_upper_limit,
12
+ _get_padded_token_len, _get_req_paddings, _get_token_paddings)
13
14
# Mock torch_xla module since it may not be available in the test environments
15
torch_xla_patcher = mock.patch.dict(
@@ -296,16 +296,29 @@ def test_update_states_request_unscheduled(model_runner):
296
def test_get_paddings():
297
min_token_size, max_token_size, padding_gap = 16, 512, 64
298
expected_paddings = [16, 32, 64, 128, 192, 256, 320, 384, 448, 512]
299
- actual_paddings = _get_paddings(min_token_size, max_token_size,
300
- padding_gap)
299
+ actual_paddings = _get_token_paddings(min_token_size, max_token_size,
300
+ padding_gap)
301
assert actual_paddings == expected_paddings
302
303
304
def test_get_padded_token_len():
305
min_token_size, max_token_size, padding_gap = 16, 512, 64
306
- paddings = _get_paddings(min_token_size, max_token_size, padding_gap)
306
+ paddings = _get_token_paddings(min_token_size, max_token_size, padding_gap)
307
assert _get_padded_token_len(paddings, 1) == 16
308
assert _get_padded_token_len(paddings, 16) == 16
309
assert _get_padded_token_len(paddings, 20) == 32
310
assert _get_padded_token_len(paddings, 300) == 320
311
assert _get_padded_token_len(paddings, 512) == 512
312
+
313
+
314
+def test_get_padded_num_reqs_with_upper_limit():
315
+ assert _get_padded_num_reqs_with_upper_limit(3, 32) == 8
316
+ assert _get_padded_num_reqs_with_upper_limit(9, 32) == 16
317
+ assert _get_padded_num_reqs_with_upper_limit(19, 32) == 32
318
+ assert _get_padded_num_reqs_with_upper_limit(17, 28) == 28
319
+
320
+
321
+def test_get_req_paddings():
322
+ assert _get_req_paddings(1, 32) == [8, 16, 32]
323
+ assert _get_req_paddings(8, 32) == [8, 16, 32]
324
+ assert _get_req_paddings(8, 36) == [8, 16, 32, 36]
vllm/v1/sample/tpu/metadata.py
+35
-42
@@ -3,7 +3,6 @@ from dataclasses import dataclass, field
3
from typing import Optional
4
5
import torch
6
-import torch_xla.core.xla_model as xm
6
7
from vllm.v1.worker.gpu_input_batch import InputBatch
8
@@ -24,15 +23,15 @@ class TPUSupportedSamplingMetadata:
23
# This class exposes a more xla-friendly interface than SamplingMetadata
24
# on TPU, in particular all arguments should be traceable and no optionals
25
# are allowed, to avoid graph recompilation on Nones.
27
- temperature: torch.Tensor
26
+ temperature: torch.Tensor = None
27
29
- min_p: torch.Tensor
28
+ min_p: torch.Tensor = None
29
# Still too slow on forward_native!
30
top_k: torch.Tensor = None
31
top_p: torch.Tensor = None
32
33
# Greedy sampling flag for compiling single xla graph.
35
- all_greedy: torch.Tensor = None
34
+ all_greedy: bool = True
35
36
# Generator not supported by xla
37
generators: dict[int,
@@ -57,64 +56,58 @@ class TPUSupportedSamplingMetadata:
56
57
allowed_token_ids_mask = None
58
bad_words_token_ids = None
60
- indices_do_sample: torch.Tensor = None
59
60
@classmethod
61
def from_input_batch(
64
- cls, input_batch: InputBatch,
65
- indices_do_sample: torch.Tensor) -> "TPUSupportedSamplingMetadata":
62
+ cls,
63
+ input_batch: InputBatch,
64
+ padded_num_reqs: int,
65
+ xla_device: torch.device,
66
+ generate_params_if_all_greedy: bool = False
67
+ ) -> "TPUSupportedSamplingMetadata":
68
"""
69
Copy sampling tensors slices from `input_batch` to on device tensors.
70
71
`InputBatch._make_sampling_metadata` causes recompilation on XLA as it
72
slices dynamic shapes on device tensors. This impl moves the dynamic
71
- ops to CPU and produces tensors of fixed `padded_num_reqs` size. It
72
- also reuses the on-device persistent tensors managed in `input_batch`
73
- to reduce waste.
74
-
75
- `indices_do_sample` contains the indices to be fed to the Sampler,
76
- normally one per request, here padded to the closest pre-compiled shape
77
- We expect sampling params tensors to be padded to the same fixed shape.
78
-
79
- Eg. 3 requests, tensors padded to 4
80
- temperature: [0.7, 0.2, 0.9]=>[0.7, 0.2, 0.9, 0.0]
81
- sample indices: [4, 10, 11]=>indices_do_sample: [4, 10, 11, 0]
73
+ ops to CPU and produces tensors of fixed `padded_num_reqs` size.
74
+
75
+ Args:
76
+ input_batch: The input batch containing sampling parameters.
77
+ padded_num_reqs: The padded number of requests.
78
+ xla_device: The XLA device.
79
+ generate_params_if_all_greedy: If True, generate sampling parameters
80
+ even if all requests are greedy. this is useful for cases where
81
+ we want to pre-compile a graph with sampling parameters, even if
82
+ they are not strictly needed for greedy decoding.
83
"""
84
+ # Early return to avoid unnecessary cpu to tpu copy
85
+ if (input_batch.all_greedy is True
86
+ and generate_params_if_all_greedy is False):
87
+ return cls(all_greedy=True)
88
+
89
num_reqs = input_batch.num_reqs
84
- padded_num_reqs = len(indices_do_sample)
90
86
- def copy_slice(cpu_tensor: torch.Tensor, tpu_tensor: torch.Tensor,
87
- fill_val) -> torch.Tensor:
88
- # Copy slice from CPU to corresponding TPU pre-allocated tensor.
91
+ def fill_slice(cpu_tensor: torch.Tensor, fill_val) -> torch.Tensor:
92
# Pad value is the default one.
93
cpu_tensor[num_reqs:padded_num_reqs] = fill_val
91
- # Subtle compilation: len(tpu_tensor) must be >= `padded_num_reqs`
92
- tpu_tensor[:padded_num_reqs] = cpu_tensor[:padded_num_reqs]
94
94
- # NOTE NickLucche The sync CPU-TPU graph we produce here must be
95
- # consistent. We can't have flags to skip copies or we'll end up
96
- # recompiling.
97
- copy_slice(input_batch.temperature_cpu_tensor, input_batch.temperature,
95
+ fill_slice(input_batch.temperature_cpu_tensor,
96
DEFAULT_SAMPLING_PARAMS["temperature"])
97
# TODO Temporarily disabled until sampling options are enabled
100
- # copy_slice(input_batch.top_p_cpu_tensor, input_batch.top_p)
101
- # copy_slice(input_batch.top_k_cpu_tensor, input_batch.top_k)
102
- copy_slice(input_batch.min_p_cpu_tensor, input_batch.min_p,
98
+ # fill_slice(input_batch.top_p_cpu_tensor)
99
+ # fill_slice(input_batch.top_k_cpu_tensor)
100
+ fill_slice(input_batch.min_p_cpu_tensor,
101
DEFAULT_SAMPLING_PARAMS["min_p"])
102
105
- xm.mark_step()
106
- xm.wait_device_ops()
107
-
103
# Slice persistent device tensors to a fixed pre-compiled padded shape.
104
return cls(
110
- temperature=input_batch.temperature[:padded_num_reqs],
111
- # Scalar tensor for xla-friendly tracing.
112
- all_greedy=torch.tensor(input_batch.all_greedy,
113
- dtype=torch.bool,
114
- device=input_batch.device),
105
+ temperature=input_batch.temperature_cpu_tensor[:padded_num_reqs].
106
+ to(xla_device),
107
+ all_greedy=input_batch.all_greedy,
108
# TODO enable more and avoid returning None values
109
top_p=None, # input_batch.top_p[:padded_num_reqs],
110
top_k=None, # input_batch.top_k[:padded_num_reqs],
118
- min_p=input_batch.min_p[:padded_num_reqs],
119
- generators=input_batch.generators,
120
- indices_do_sample=indices_do_sample)
111
+ min_p=input_batch.min_p_cpu_tensor[:padded_num_reqs].to(
112
+ xla_device),
113
+ generators=input_batch.generators)
vllm/v1/worker/tpu_model_runner.py
+101
-72
@@ -32,7 +32,7 @@ from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
32
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
33
KVCacheSpec, SlidingWindowSpec)
34
from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
35
- ModelRunnerOutput, SamplerOutput)
35
+ ModelRunnerOutput)
36
from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
37
from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler
38
from vllm.v1.utils import bind_kv_cache
@@ -177,10 +177,12 @@ class TPUModelRunner:
177
# Range tensor with values [0 .. self.max_num_tokens - 1].
178
# Used to initialize positions / context_lens / seq_lens
179
self.arange_np = np.arange(self.max_num_tokens, dtype=np.int32)
180
- self.num_tokens_paddings = _get_paddings(
180
+ self.num_tokens_paddings = _get_token_paddings(
181
min_token_size=16,
182
max_token_size=self.max_num_tokens,
183
padding_gap=envs.VLLM_TPU_BUCKET_PADDING_GAP)
184
+ self.num_reqs_paddings = _get_req_paddings(
185
+ min_req_size=MIN_NUM_SEQS, max_req_size=self.max_num_reqs)
186
187
def _update_num_xla_graphs(self, case_str):
188
check_comp = self.check_recompilation and not self.enforce_eager
@@ -508,7 +510,7 @@ class TPUModelRunner:
510
# Padded to avoid recompiling when `num_reqs` varies.
511
logits_indices = self.query_start_loc_cpu[1:padded_num_reqs + 1] - 1
512
logits_indices = logits_indices.to(self.device)
511
- return attn_metadata, logits_indices
513
+ return attn_metadata, logits_indices, padded_num_reqs
514
515
def _scatter_placeholders(
516
self,
@@ -663,7 +665,8 @@ class TPUModelRunner:
665
mm_embeds = []
666
667
# Prepare inputs
666
- attn_metadata, logits_indices = self._prepare_inputs(scheduler_output)
668
+ attn_metadata, logits_indices, padded_num_reqs = self._prepare_inputs(
669
+ scheduler_output)
670
if self.is_multimodal_model:
671
# NOTE(woosuk): To unify token ids and soft tokens (vision
672
# embeddings), we always use embeddings (rather than token ids)
@@ -682,11 +685,6 @@ class TPUModelRunner:
685
input_ids = self.input_ids
686
inputs_embeds = None
687
num_reqs = self.input_batch.num_reqs
685
- # NOTE (NickLucche) here we sync with TPU: sampling params tensors
686
- # are copied to device in chunks of pre-compiled padded shape to
687
- # avoid recompilations.
688
- tpu_sampling_metadata = TPUSupportedSamplingMetadata.\
689
- from_input_batch(self.input_batch, logits_indices)
688
# Run the decoder
689
with set_forward_context(attn_metadata, self.vllm_config):
690
hidden_states = self.model(
@@ -694,6 +692,10 @@ class TPUModelRunner:
692
positions=self.position_ids,
693
inputs_embeds=inputs_embeds,
694
)
695
+ hidden_states = self.select_hidden_states(hidden_states,
696
+ logits_indices)
697
+ tpu_sampling_metadata = TPUSupportedSamplingMetadata.\
698
+ from_input_batch(self.input_batch, padded_num_reqs, self.device)
699
selected_token_ids = self.sample_from_hidden(hidden_states,
700
tpu_sampling_metadata)
701
# Remove padding on cpu and keep dynamic op outside of xla graph.
@@ -857,60 +859,78 @@ class TPUModelRunner:
859
inputs_embeds=inputs_embeds)
860
self._hidden_states_dtype = out.dtype
861
860
- def capture_model(self) -> None:
861
- """Compile the model."""
862
-
862
+ def _precompile_backbone(self) -> None:
863
logger.info("Compiling the model with different input shapes.")
864
865
start = time.perf_counter()
866
for num_tokens in self.num_tokens_paddings:
867
logger.info(" -- num_tokens: %d", num_tokens)
868
self._dummy_run(num_tokens)
869
- xm.mark_step()
869
xm.wait_device_ops()
870
end = time.perf_counter()
872
-
871
logger.info("Compilation finished in in %.2f [secs].", end - start)
874
- self._update_num_xla_graphs("model")
872
+ self._update_num_xla_graphs("model backbone")
873
876
- logger.info("Compiling sampling with different input shapes.")
874
+ def _precompile_select_hidden_states(self) -> None:
875
+ # Compile hidden state selection function for bucketed
876
+ # n_tokens x max_num_reqs. Graph is really small so this is fine.
877
+ logger.info(
878
+ "Compiling select_hidden_states with different input shapes.")
879
start = time.perf_counter()
880
hsize = self.model_config.get_hidden_size()
879
- device = self.device
880
- # Compile sampling step for different model+sampler outputs in bucketed
881
- # n_tokens x max_num_reqs. Graph is really small so this is fine.
881
for num_tokens in self.num_tokens_paddings:
883
- num_reqs_to_sample = MIN_NUM_SEQS
884
- dummy_hidden = torch.randn((num_tokens, hsize),
885
- device=device,
882
+ dummy_hidden = torch.zeros((num_tokens, hsize),
883
+ device=self.device,
884
dtype=self._hidden_states_dtype)
887
- # Compile for [8, 16, .., 128,.., `self.max_num_reqs`]
888
- while True:
889
- indices = torch.zeros(
890
- num_reqs_to_sample,
891
- dtype=torch.int32,
892
- device=device,
893
- )
894
- xm.mark_step()
895
- sampling_meta = TPUSupportedSamplingMetadata.\
896
- from_input_batch(self.input_batch, indices)
897
- logger.info(" -- num_tokens: %d, num_seqs: %d", num_tokens,
898
- num_reqs_to_sample)
899
- out = self.sample_from_hidden(dummy_hidden, sampling_meta)
900
- out = out.cpu()
901
- # Requests can't be more than tokens. But do compile for the
902
- # next bigger value in case num_tokens uses bucketed padding.
903
- if num_reqs_to_sample >= min(num_tokens, self.max_num_reqs):
904
- break
905
- # Make sure to compile the `max_num_reqs` upper-limit case
906
- num_reqs_to_sample = _get_padded_num_reqs_with_upper_limit(
907
- num_reqs_to_sample + 1, self.max_num_reqs)
885
+ torch._dynamo.mark_dynamic(dummy_hidden, 0)
886
+ for num_reqs in self.num_reqs_paddings:
887
+ indices = torch.zeros(num_reqs,
888
+ dtype=torch.int32,
889
+ device=self.device)
890
+ torch._dynamo.mark_dynamic(indices, 0)
891
+ self.select_hidden_states(dummy_hidden, indices)
892
+ logger.info(" -- num_tokens: %d", num_tokens)
893
xm.wait_device_ops()
894
end = time.perf_counter()
895
+ logger.info("Compilation finished in in %.2f [secs].", end - start)
896
+ self._update_num_xla_graphs("select_hidden_states")
897
898
+ def _precompile_sample_from_hidden(self) -> None:
899
+ logger.info("Compiling sampling with different input shapes.")
900
+ start = time.perf_counter()
901
+ hsize = self.model_config.get_hidden_size()
902
+ for num_reqs in self.num_reqs_paddings:
903
+ dummy_hidden = torch.zeros((num_reqs, hsize),
904
+ device=self.device,
905
+ dtype=self._hidden_states_dtype)
906
+ # The first dimension of dummy_hidden cannot be mark_dynamic because
907
+ # some operations in the sampler require it to be static.
908
+ for all_greedy in [False, True]:
909
+ generate_params_if_all_greedy = not all_greedy
910
+ sampling_metadata = (
911
+ TPUSupportedSamplingMetadata.from_input_batch(
912
+ self.input_batch,
913
+ num_reqs,
914
+ self.device,
915
+ generate_params_if_all_greedy,
916
+ ))
917
+ sampling_metadata.all_greedy = all_greedy
918
+ self.sample_from_hidden(dummy_hidden, sampling_metadata)
919
+ logger.info(" -- num_seqs: %d", num_reqs)
920
+ xm.wait_device_ops()
921
+ end = time.perf_counter()
922
logger.info("Compilation finished in in %.2f [secs].", end - start)
923
self._update_num_xla_graphs("sampling")
924
925
+ def capture_model(self) -> None:
926
+ """
927
+ Precompile all the subgraphs with possible input shapes.
928
+ """
929
+ # TODO: precompile encoder
930
+ self._precompile_backbone()
931
+ self._precompile_select_hidden_states()
932
+ self._precompile_sample_from_hidden()
933
+
934
def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
935
"""
936
Initialize KV cache based on `kv_cache_config`.
@@ -962,48 +982,55 @@ class TPUModelRunner:
982
compiled_model.original_code_object)
983
compiled_model.compiled_codes.clear()
984
985
+ @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
986
+ def select_hidden_states(self, hidden_states, indices_do_sample):
987
+ return hidden_states[indices_do_sample]
988
+
989
+ @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
990
def sample_from_hidden(
991
self,
967
- hidden_states: torch.Tensor,
992
+ sample_hidden_states: torch.Tensor,
993
sampling_metadata: TPUSupportedSamplingMetadata,
994
) -> torch.Tensor:
995
"""
971
- Sample with xla-friendly function. This function is to be traced
972
- separately for lighter compilation overhead.
973
- """
974
- # Tensor `sample_hidden_states` is of fixed pre-compiled size.
975
- sample_hidden_states = \
976
- hidden_states[sampling_metadata.indices_do_sample]
977
- # SamplingMetadata here for pruning output in LogitsProcessor, disabled.
996
+ Sample with xla-friendly function. This function is to be traced
997
+ separately from `forward` for lighter compilation overhead.
998
+ """
999
logits = self.model.compute_logits(sample_hidden_states, None)
979
-
980
- def sample(
981
- logits: torch.Tensor,
982
- sampling_metadata: TPUSupportedSamplingMetadata
983
- ) -> SamplerOutput:
984
- sampler_out = self.sampler(logits, sampling_metadata)
985
- return sampler_out
986
-
987
- # Optimized greedy sampling branch, tracing both paths in a single pass
988
- # NOTE all_greedy is a scalar, this is just an optimized if/else.
989
- out_tokens = torch.where(
990
- sampling_metadata.all_greedy,
991
- torch.argmax(logits, dim=-1, keepdim=True),
992
- sample(logits, sampling_metadata).sampled_token_ids)
1000
+ if sampling_metadata.all_greedy:
1001
+ out_tokens = torch.argmax(logits, dim=-1, keepdim=True)
1002
+ else:
1003
+ out_tokens = self.sampler(logits,
1004
+ sampling_metadata).sampled_token_ids
1005
return out_tokens
1006
1007
+ def get_multimodal_embeddings(self, *args, **kwargs):
1008
+ return self.model.get_multimodal_embeddings(*args, **kwargs)
1009
996
-def _get_padded_number(n: int, multiple: int) -> int:
997
- return ((n + multiple - 1) // multiple) * multiple
1010
+ def get_input_embeddings(self, *args, **kwargs):
1011
+ return self.model.get_input_embeddings(*args, **kwargs)
1012
+
1013
+
1014
+def _get_req_paddings(min_req_size: int, max_req_size: int) -> list[int]:
1015
+ logger.info("Preparing request paddings:")
1016
+ # assert min_req_size is power of 2
1017
+ assert (min_req_size & (min_req_size - 1) == 0) and min_req_size > 0
1018
+ paddings: list = []
1019
+ num = max(MIN_NUM_SEQS, min_req_size)
1020
+ while num <= max_req_size and (len(paddings) == 0 or paddings[-1] != num):
1021
+ paddings.append(num)
1022
+ logger.info(" %d", num)
1023
+ num = _get_padded_num_reqs_with_upper_limit(num + 1, max_req_size)
1024
+ return paddings
1025
1026
1000
-def _get_padded_num_reqs_with_upper_limit(x, upper_limit) -> int:
1027
+def _get_padded_num_reqs_with_upper_limit(x: int, upper_limit: int) -> int:
1028
res = MIN_NUM_SEQS if x <= MIN_NUM_SEQS else 1 << (x - 1).bit_length()
1029
return min(res, upper_limit)
1030
1031
1005
-def _get_paddings(min_token_size: int, max_token_size: int,
1006
- padding_gap: int) -> list[int]:
1032
+def _get_token_paddings(min_token_size: int, max_token_size: int,
1033
+ padding_gap: int) -> list[int]:
1034
"""Generate a list of padding size, starting from min_token_size,
1035
ending with a number that can cover max_token_size
1036
@@ -1013,18 +1040,20 @@ def _get_paddings(min_token_size: int, max_token_size: int,
1040
first increase the size to twice,
1041
then increase the padding size by padding_gap.
1042
"""
1043
+ # assert min_token_size is power of 2
1044
+ assert (min_token_size & (min_token_size - 1) == 0) and min_token_size > 0
1045
paddings = []
1046
num = min_token_size
1047
1048
if padding_gap == 0:
1020
- logger.info("Using exponential paddings:")
1049
+ logger.info("Using exponential token paddings:")
1050
while num <= max_token_size:
1051
logger.info(" %d", num)
1052
paddings.append(num)
1053
num *= 2
1054
1055
else:
1027
- logger.info("Using incremental paddings:")
1056
+ logger.info("Using incremental token paddings:")
1057
while num <= padding_gap:
1058
logger.info(" %d", num)
1059
paddings.append(num)