Skip to content

Commit 41a55cc

Browse files
sarckkminpeter
authored andcommitted
Fix test_memory_usage_no_spec (vllm-project#17754)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com> Signed-off-by: minpeter <kali2005611@gmail.com>
1 parent 54b04b8 commit 41a55cc

File tree

1 file changed

+6
-6
lines changed

1 file changed

+6
-6
lines changed

tests/spec_decode/test_memory_usage.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -42,12 +42,12 @@ def add_seq_group_to_engine(engine: vllm.LLMEngine, seq_group: SequenceGroup):
4242

4343
def test_memory_usage_no_spec():
4444
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-
)
45+
llm = vllm.LLM(model=MAIN_MODEL,
46+
speculative_config={
47+
"model": SPEC_MODEL,
48+
"num_speculative_tokens": 3,
49+
"disable_by_batch_size": SPEC_DISABLE_BATCH_SIZE,
50+
})
5151

5252
batch_sequences = set()
5353
engine = llm.llm_engine

0 commit comments

Comments
 (0)