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, 17:15 UTC
[TPU] [Perf] Improve Memory Usage Estimation (#15671)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Mar 28, 2025 at 10:37 UTC
038bededbac67e873d3aa6155c7c05674b98db8c
1 file changed
+7
-1
vllm/v1/worker/tpu_worker.py
+7
-1
@@ -161,7 +161,13 @@ class TPUWorker:
161
# intermediate activations.
162
m = xm.get_memory_info(self.device)
163
total_memory_size = m["bytes_limit"]
164
- profiled = m["peak_bytes_used"] # Weights + intermediate activations.
164
+ current_mem = m["bytes_used"]
165
+ # Ideally we would use profiled = m["peak_bytes_used"] to
166
+ # get weights + activations. But there is memory used during
167
+ # compilation / weight loading that impacts the peak and
168
+ # there is no way to reset peak memory in XLA, So we
169
+ # use the heuristic of 2% of weights.
170
+ profiled = current_mem * 1.02
171
172
# Calculate the TPU KV cache size based on profiling.
173
usable_memory_size = int(total_memory_size *