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, 12:15 UTC
[BugFix] fix speculative decoding memory leak when speculation is disabled (#15506)
Signed-off-by: Noah Yoshida <noahcy117@gmail.com>
Noah Yoshida committed
Apr 30, 2025 at 23:28 UTC
13cf6b62367b1bae260cb58b5158060c5f6852bc
2 files changed
+91
tests/spec_decode/test_memory_usage.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""This docstring details important information on the testing methodology.
3
+
4
+This test verifies that memory usage remains constant (or never grows) when
5
+we enable / disable speculation via --speculative-disable-by-batch-size.
6
+
7
+There are a lot of things we try to keep track of between batches of requests
8
+and if certain tensors are not freed from memory, can result in CUDA ooms.
9
+
10
+This is particularly relevant for production situations where speculation might
11
+be enabled during off hours, but disabled once traffic peaks during the workday.
12
+Since traffic will stay high for a long period of time, verifying we do not
13
+increase our memory usage over time is essential to prevent possible CUDA ooms.
14
+"""
15
+
16
+import torch
17
+
18
+import vllm
19
+from tests.core.utils import create_dummy_prompt
20
+from vllm.sequence import SequenceGroup
21
+
22
+ITERATIONS = 100
23
+MAIN_MODEL = "JackFram/llama-68m"
24
+
25
+# speculative model
26
+SPEC_MODEL = "abhigoyal/vllm-medusa-llama-68m-random"
27
+
28
+BATCH_SIZE = 5
29
+SPEC_DISABLE_BATCH_SIZE = 2
30
+
31
+
32
+def add_seq_group_to_engine(engine: vllm.LLMEngine, seq_group: SequenceGroup):
33
+ scheduler = engine.scheduler[0]
34
+ scheduler.add_seq_group(seq_group)
35
+
36
+
37
+"""
38
+Since we are using a batch size greater than the disabled batch size,
39
+we can ensure we go through the _no_spec codepath for most of our engine steps.
40
+"""
41
+
42
+
43
+def test_memory_usage_no_spec():
44
+ previous_memory_allocated = None
45
+ llm = vllm.LLM(
46
+ model=MAIN_MODEL,
47
+ speculative_model=SPEC_MODEL,
48
+ num_speculative_tokens=3,
49
+ speculative_disable_by_batch_size=SPEC_DISABLE_BATCH_SIZE,
50
+ )
51
+
52
+ batch_sequences = set()
53
+ engine = llm.llm_engine
54
+
55
+ for i in range(ITERATIONS):
56
+ seq, seq_group = create_dummy_prompt(request_id=str(i),
57
+ prompt_length=10,
58
+ min_tokens=10,
59
+ max_tokens=10)
60
+
61
+ add_seq_group_to_engine(engine, seq_group)
62
+
63
+ batch_sequences.add(seq)
64
+ engine.step()
65
+ for seq in list(batch_sequences):
66
+ if seq.is_finished():
67
+ batch_sequences.remove(seq)
68
+
69
+ # If we aren't at our batch size yet, continue
70
+ if len(batch_sequences) <= BATCH_SIZE:
71
+ continue
72
+
73
+ # Otherwise, loop until at least one request is done
74
+ while not any(seq.is_finished() for seq in batch_sequences):
75
+ engine.step()
76
+
77
+ # Remove it from the set
78
+ for seq in list(batch_sequences):
79
+ if seq.is_finished():
80
+ batch_sequences.remove(seq)
81
+
82
+ # At this point, we are always at the case where we have finished
83
+ # processing some number of requests from the batch after running
84
+ # several _no_spec executions. The memory should not have
85
+ # increased between the previous time this was recorded and the
86
+ # current time.
87
+ if previous_memory_allocated is None:
88
+ previous_memory_allocated = torch.cuda.memory_allocated()
89
+ else:
90
+ assert previous_memory_allocated == torch.cuda.memory_allocated()
vllm/spec_decode/spec_decode_worker.py
+1
@@ -695,6 +695,7 @@ class SpecDecodeWorker(LoRANotSupportedWorkerBase):
695
seq_group_meta_with_hidden):
696
self.previous_hidden_states.update(hidden_states,
697
seq_group_meta_with_hidden)
698
+ self.previous_hidden_states.prune(seq_group_meta_with_hidden)
699
700
if not skip_proposer:
701
# We prepare the prefill hidden states here so that there no