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 14, 22:35 UTC
Fix test_kv_sharing_fast_prefill flakiness (#22038)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Aug 1, 2025 at 23:55 UTC
8564dc9448ed8648088c25248313933308ae36d8
1 file changed
+30
-5
tests/v1/e2e/test_kv_sharing_fast_prefill.py
+30
-5
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import gc
4
import random
5
from typing import Optional, Union
6
@@ -10,6 +9,7 @@ import torch
9
10
from vllm import LLM, SamplingParams
11
from vllm.config import CompilationConfig, CompilationLevel
12
+from vllm.distributed import cleanup_dist_env_and_memory
13
from vllm.forward_context import get_forward_context
14
from vllm.model_executor.models.gemma3n import Gemma3nForConditionalGeneration
15
from vllm.model_executor.models.registry import ModelRegistry
@@ -18,6 +18,9 @@ from vllm.sequence import IntermediateTensors
18
19
from ...utils import fork_new_process_for_each_test
20
21
+# global seed
22
+SEED = 42
23
+
24
25
class TestGemma3nForConditionalGeneration(Gemma3nForConditionalGeneration):
26
@@ -95,8 +98,25 @@ def test_prompts():
98
return prompts
99
100
101
+def cleanup(llm: LLM, compilation_config: CompilationConfig):
102
+ # hacky: below lines are required to free up memory for the next test
103
+ # when setting VLLM_ENABLE_V1_MULTIPROCESSING=0, del llm is not sufficient
104
+ # TODO(sarckk): when enforce_eager=False, memory is not freed:
105
+ # find out why and re-enable test for enforce_eager=False case
106
+ llm_engine = llm.llm_engine.engine_core.engine_core
107
+ model_runner = llm_engine.model_executor.driver_worker.worker.model_runner
108
+ del model_runner.model
109
+ del model_runner.kv_caches
110
+ del compilation_config.static_forward_context
111
+ compilation_config.static_forward_context = {}
112
+
113
+ del llm
114
+ torch.cuda.empty_cache()
115
+ cleanup_dist_env_and_memory()
116
+
117
+
118
@fork_new_process_for_each_test
99
-@pytest.mark.parametrize("enforce_eager", [True, False])
119
+@pytest.mark.parametrize("enforce_eager", [True])
120
def test_kv_sharing_fast_prefill(
121
monkeypatch: pytest.MonkeyPatch,
122
enforce_eager: bool,
@@ -115,23 +135,28 @@ def test_kv_sharing_fast_prefill(
135
with monkeypatch.context() as m:
136
m.setenv("VLLM_USE_V1", "1")
137
138
+ # Make scheduling deterministic for reproducibility
139
+ m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
140
+
141
llm = LLM(
142
model="google/gemma-3n-E2B-it",
143
enforce_eager=enforce_eager,
144
compilation_config=compilation_config,
145
+ seed=SEED,
146
)
147
ref_responses = llm.generate(test_prompts, sampling_params)
148
125
- del llm
126
- gc.collect()
127
- torch.cuda.empty_cache()
149
+ cleanup(llm, compilation_config)
150
151
llm = LLM(model="google/gemma-3n-E2B-it",
152
enforce_eager=enforce_eager,
153
compilation_config=compilation_config,
154
+ seed=SEED,
155
kv_sharing_fast_prefill=True)
156
optimized_responses = llm.generate(test_prompts, sampling_params)
157
158
+ cleanup(llm, compilation_config)
159
+
160
misses = 0
161
162
for ref_response, optimized_response in zip(ref_responses,