[ROCm][CI] Ensure sliding window tests release GPU memory (#49055)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 18, 2026 at 15:44 UTC
df362b2d6d091771dbcc364b2fc96d17a78df274
1 file changed
+32
-27
tests/v1/e2e/general/test_correctness_sliding_window.py
+32
-27
@@ -33,7 +33,7 @@ model_config = {
33
@pytest.mark.parametrize("seed", [1])
34
@pytest.mark.parametrize("disable_hybrid_kv_cache_manager", [True, False])
35
def test_sliding_window_retrieval(
36
- model, batch_size, seed, disable_hybrid_kv_cache_manager
36
+ model, batch_size, seed, disable_hybrid_kv_cache_manager, vllm_runner
37
):
38
"""
39
The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then
@@ -48,34 +48,39 @@ def test_sliding_window_retrieval(
48
49
test_config = model_config[model]
50
51
- llm = LLM(
52
- model=model,
51
+ with vllm_runner(
52
+ model,
53
+ max_model_len=None,
54
+ enable_chunked_prefill=None,
55
disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,
56
enforce_eager=enforce_eager,
55
- )
56
- sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
57
-
58
- prompts, answer, indices = prep_prompts(batch_size, ln_range=test_config.ln_range)
59
-
60
- check_length(prompts, llm, test_config.sliding_window)
61
-
62
- # Fresh generation
63
- responses = llm.generate(prompts, sampling_params)
64
- check_answers(
65
- indices,
66
- answer,
67
- [response.outputs[0].text for response in responses],
68
- accept_rate=1.0,
69
- )
70
-
71
- # Re-generate with the same prompts to test prefix caching
72
- responses = llm.generate(prompts, sampling_params)
73
- check_answers(
74
- indices,
75
- answer,
76
- [response.outputs[0].text for response in responses],
77
- accept_rate=1.0,
78
- )
57
+ ) as runner:
58
+ llm = runner.get_llm()
59
+ sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
60
+
61
+ prompts, answer, indices = prep_prompts(
62
+ batch_size, ln_range=test_config.ln_range
63
+ )
64
+
65
+ check_length(prompts, llm, test_config.sliding_window)
66
+
67
+ # Fresh generation
68
+ responses = llm.generate(prompts, sampling_params)
69
+ check_answers(
70
+ indices,
71
+ answer,
72
+ [response.outputs[0].text for response in responses],
73
+ accept_rate=1.0,
74
+ )
75
+
76
+ # Re-generate with the same prompts to test prefix caching
77
+ responses = llm.generate(prompts, sampling_params)
78
+ check_answers(
79
+ indices,
80
+ answer,
81
+ [response.outputs[0].text for response in responses],
82
+ accept_rate=1.0,
83
+ )
84
85
86
def check_length(prompts: list[str], llm: LLM, sliding_window: int):