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, 23:30 UTC
[V1][PP] Cache Intermediate Tensors (#13353)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Feb 16, 2025 at 10:02 UTC
e18227b04a7ef082e55380e143ae9e56f1dc6f86
1 file changed
+28
-8
vllm/v1/worker/gpu_model_runner.py
+28
-8
@@ -2,7 +2,7 @@
2
3
import gc
4
import time
5
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast
5
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union, cast
6
7
import numpy as np
8
import torch
@@ -149,6 +149,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
149
self.positions = torch.zeros(self.max_num_tokens,
150
dtype=torch.int64,
151
device=self.device)
152
+ # self.intermediate_tensors # Set after load_model
153
154
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
155
if self.uses_mrope:
@@ -869,7 +870,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
870
self,
871
scheduler_output: "SchedulerOutput",
872
intermediate_tensors: Optional[IntermediateTensors] = None,
872
- ) -> ModelRunnerOutput:
873
+ ) -> Union[ModelRunnerOutput, torch.Tensor]:
874
batch_changed = self._update_states(scheduler_output)
875
876
if self.is_multimodal_model:
@@ -919,6 +920,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
920
else:
921
positions = self.positions[:num_input_tokens]
922
923
+ if get_pp_group().is_first_rank:
924
+ intermediate_tensors = None
925
+ else:
926
+ intermediate_tensors = IntermediateTensors({
927
+ k: v[:num_input_tokens]
928
+ for k, v in self.intermediate_tensors.items()
929
+ })
930
+
931
# Run the decoder.
932
# Use persistent buffers for CUDA graphs.
933
with set_forward_context(attn_metadata, self.vllm_config):
@@ -931,7 +940,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
940
inputs_embeds=inputs_embeds,
941
)
942
if not get_pp_group().is_last_rank:
943
+ # For mid-pipeline stages, return the hidden states.
944
return hidden_states
945
+
946
hidden_states = hidden_states[:num_scheduled_tokens]
947
sample_hidden_states = hidden_states[logits_indices]
948
logits = self.model.compute_logits(sample_hidden_states, None)
@@ -1118,12 +1129,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1129
positions = self.mrope_positions[:, :num_tokens]
1130
else:
1131
positions = self.positions[:num_tokens]
1121
- intermediate_tensors = None
1122
- if not get_pp_group().is_first_rank:
1123
- intermediate_tensors = self.model.make_empty_intermediate_tensors(
1124
- batch_size=num_tokens,
1125
- dtype=self.model_config.dtype,
1126
- device=self.device)
1132
+
1133
+ if get_pp_group().is_first_rank:
1134
+ intermediate_tensors = None
1135
+ else:
1136
+ if not hasattr(self, "intermediate_tensors"):
1137
+ self.intermediate_tensors = (
1138
+ self.model.make_empty_intermediate_tensors(
1139
+ batch_size=self.max_num_tokens,
1140
+ dtype=self.model_config.dtype,
1141
+ device=self.device))
1142
+ intermediate_tensors = IntermediateTensors({
1143
+ k: v[:num_tokens]
1144
+ for k, v in self.intermediate_tensors.items()
1145
+ })
1146
+
1147
with set_forward_context(None, self.vllm_config):
1148
hidden_states = model(
1149
input_ids=input_ids,