Skip to content

GitLab

Unverified Commit 324a3119 authored by Yong Hoon Shin's avatar Yong Hoon Shin Committed by GitHub
Browse files

Fix test_memory_usage_no_spec (#17754) 


Signed-off-by: default avatarYong Hoon Shin <yhshin@meta.com>
parent 8a15c260
Show whitespace changes
Inline Side-by-side
Showing with 6 additions and 6 deletions
tests/spec_decode/test_memory_usage.py
View file @ 324a3119
...@@ -42,12 +42,12 @@ we can ensure we go through the _no_spec codepath for most of our engine steps. ...@@ -42,12 +42,12 @@ we can ensure we go through the _no_spec codepath for most of our engine steps.
def test_memory_usage_no_spec(): def test_memory_usage_no_spec():
previous_memory_allocated = None previous_memory_allocated = None
llm = vllm.LLM( llm = vllm.LLM(model=MAIN_MODEL,
model=MAIN_MODEL, speculative_config={
speculative_model=SPEC_MODEL, "model": SPEC_MODEL,
num_speculative_tokens=3, "num_speculative_tokens": 3,
speculative_disable_by_batch_size=SPEC_DISABLE_BATCH_SIZE, "disable_by_batch_size": SPEC_DISABLE_BATCH_SIZE,
) })
batch_sequences = set() batch_sequences = set()
engine = llm.llm_engine engine = llm.llm_engine
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment