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] Fix intermediate tensor values (#13417)
Signed-off-by: Cody Yu <hao.yu.cody@gmail.com>
Cody Yu committed
Feb 17, 2025 at 13:37 UTC
6ac485a953d8bdba2c6b779207af50bd0fcc588b
2 files changed
+11
-2
vllm/sequence.py
+3
@@ -1137,6 +1137,9 @@ class IntermediateTensors:
1137
def __setitem__(self, key: str, value: torch.Tensor):
1138
self.tensors[key] = value
1139
1140
+ def items(self):
1141
+ return self.tensors.items()
1142
+
1143
def __len__(self):
1144
return len(self.tensors)
1145
vllm/v1/worker/gpu_model_runner.py
+8
-2
@@ -151,7 +151,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
151
self.positions = torch.zeros(self.max_num_tokens,
152
dtype=torch.int64,
153
device=self.device)
154
- # self.intermediate_tensors # Set after load_model
154
+ # None in the first PP rank. The rest are set after load_model.
155
+ self.intermediate_tensors: Optional[IntermediateTensors] = None
156
157
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
158
if self.uses_mrope:
@@ -922,6 +923,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
923
if get_pp_group().is_first_rank:
924
intermediate_tensors = None
925
else:
926
+ assert intermediate_tensors is not None
927
+ assert self.intermediate_tensors is not None
928
+ for k, v in intermediate_tensors.items():
929
+ self.intermediate_tensors[k][:num_input_tokens].copy_(
930
+ v[:num_input_tokens], non_blocking=True)
931
intermediate_tensors = IntermediateTensors({
932
k: v[:num_input_tokens]
933
for k, v in self.intermediate_tensors.items()
@@ -1120,7 +1126,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1126
if get_pp_group().is_first_rank:
1127
intermediate_tensors = None
1128
else:
1123
- if not hasattr(self, "intermediate_tensors"):
1129
+ if self.intermediate_tensors is None:
1130
self.intermediate_tensors = (
1131
self.model.make_empty_intermediate_tensors(
1132
batch_size=self.max_num_tokens,