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 1, 22:30 UTC
[ROCm][CI] Stabilize memory-release in the Hybrid model generation tests (#44046)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jun 4, 2026 at 09:34 UTC
6f68ca3e9174d5007425c44f17ea3c8b110ea4d4
1 file changed
+76
-41
tests/models/language/generation/test_hybrid.py
+76
-41
@@ -2,11 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable
5
+from contextlib import contextmanager, nullcontext
6
7
import pytest
8
9
from tests.models.registry import HF_EXAMPLE_MODELS
9
-from tests.utils import multi_gpu_test
10
+from tests.utils import multi_gpu_test, wait_for_gpu_memory_to_clear
11
from vllm import LLM
12
from vllm.engine.arg_utils import EngineArgs
13
from vllm.platforms import current_platform
@@ -404,6 +405,30 @@ def _get_vllm_runner_params(
405
}
406
407
408
+def _wait_for_rocm_memory_to_settle() -> None:
409
+ if not current_platform.is_rocm():
410
+ return
411
+
412
+ num_gpus = current_platform.device_count()
413
+ if num_gpus == 0:
414
+ return
415
+
416
+ wait_for_gpu_memory_to_clear(
417
+ devices=list(range(num_gpus)),
418
+ threshold_ratio=0.01,
419
+ timeout_s=120,
420
+ )
421
+
422
+
423
+@contextmanager
424
+def _owned_vLLM_runner(vllm_runner, kwargs):
425
+ try:
426
+ with vllm_runner(**kwargs) as runner:
427
+ yield runner
428
+ finally:
429
+ _wait_for_rocm_memory_to_settle()
430
+
431
+
432
def _get_vLLM_output(
433
vllm_runner,
434
kwargs,
@@ -413,17 +438,21 @@ def _get_vLLM_output(
438
num_repetitions=1,
439
vllm_model=None,
440
):
416
- outs = []
417
- if vllm_model is None:
418
- vllm_model = vllm_runner(**kwargs)
419
- for _ in range(num_repetitions):
420
- if num_logprobs < 0:
421
- vllm_output = vllm_model.generate_greedy(prompts, max_tokens)
422
- else:
423
- vllm_output = vllm_model.generate_greedy_logprobs(
424
- prompts, max_tokens, num_logprobs
425
- )
426
- outs.append(vllm_output)
441
+ runner_context = (
442
+ _owned_vLLM_runner(vllm_runner, kwargs)
443
+ if vllm_model is None
444
+ else nullcontext(vllm_model)
445
+ )
446
+ with runner_context as runner:
447
+ outs = []
448
+ for _ in range(num_repetitions):
449
+ if num_logprobs < 0:
450
+ vllm_output = runner.generate_greedy(prompts, max_tokens)
451
+ else:
452
+ vllm_output = runner.generate_greedy_logprobs(
453
+ prompts, max_tokens, num_logprobs
454
+ )
455
+ outs.append(vllm_output)
456
457
return outs, vllm_model
458
@@ -772,38 +801,44 @@ def test_apc_multiple_prompts_partial_cached_outputs(
801
802
# Cache only part of all the prompts
803
vllm_runner_kwargs["enable_prefix_caching"] = True
775
- vllm_outputs_partial_cache, vllm_model = _get_vLLM_output(
776
- vllm_runner, vllm_runner_kwargs, generated_prompts[:3], max_tokens, num_logprobs
777
- )
778
-
779
- compare_operator(
780
- outputs_0_lst=vllm_outputs_no_cache[0][:3],
781
- outputs_1_lst=vllm_outputs_partial_cache[0],
782
- name_0="vllm_no_cache",
783
- name_1="vllm_partial_cache",
784
- )
785
-
786
- vllm_outputs_cache_rep, _ = _get_vLLM_output(
787
- vllm_runner,
788
- vllm_runner_kwargs,
789
- generated_prompts,
790
- max_tokens,
791
- num_logprobs,
792
- n_repetitions,
793
- vllm_model=vllm_model,
794
- )
795
-
796
- for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
797
- # In the first repetition, the caches are filled
798
- # In the second repetition, these caches are reused
804
+ with _owned_vLLM_runner(vllm_runner, vllm_runner_kwargs) as vllm_model:
805
+ vllm_outputs_partial_cache, _ = _get_vLLM_output(
806
+ vllm_runner,
807
+ vllm_runner_kwargs,
808
+ generated_prompts[:3],
809
+ max_tokens,
810
+ num_logprobs,
811
+ vllm_model=vllm_model,
812
+ )
813
814
compare_operator(
801
- outputs_0_lst=vllm_outputs_no_cache[0],
802
- outputs_1_lst=vllm_outputs_cache_itn,
815
+ outputs_0_lst=vllm_outputs_no_cache[0][:3],
816
+ outputs_1_lst=vllm_outputs_partial_cache[0],
817
name_0="vllm_no_cache",
804
- name_1=f"vllm_cache_it_{r_idx + 1}",
818
+ name_1="vllm_partial_cache",
819
+ )
820
+
821
+ vllm_outputs_cache_rep, _ = _get_vLLM_output(
822
+ vllm_runner,
823
+ vllm_runner_kwargs,
824
+ generated_prompts,
825
+ max_tokens,
826
+ num_logprobs,
827
+ n_repetitions,
828
+ vllm_model=vllm_model,
829
)
830
831
+ for r_idx, vllm_outputs_cache_itn in enumerate(vllm_outputs_cache_rep):
832
+ # In the first repetition, the caches are filled
833
+ # In the second repetition, these caches are reused
834
+
835
+ compare_operator(
836
+ outputs_0_lst=vllm_outputs_no_cache[0],
837
+ outputs_1_lst=vllm_outputs_cache_itn,
838
+ name_0="vllm_no_cache",
839
+ name_1=f"vllm_cache_it_{r_idx + 1}",
840
+ )
841
+
842
843
# Test that outputs match whether prefix caching is enabled or not for mamba.
844
@pytest.mark.parametrize("model", ["tiiuae/falcon-mamba-7b"])
@@ -826,7 +861,7 @@ def test_same_mamba_output_apc_on_vs_off(
861
862
# No prefix caching
863
kwargs_no_apc = {**base_kwargs, "enable_prefix_caching": False}
829
- with vllm_runner(**kwargs_no_apc) as vllm_model:
864
+ with _owned_vLLM_runner(vllm_runner, kwargs_no_apc) as vllm_model:
865
outputs_no_apc, _ = _get_vLLM_output(
866
vllm_runner,
867
kwargs_no_apc,
@@ -841,7 +876,7 @@ def test_same_mamba_output_apc_on_vs_off(
876
"enable_prefix_caching": True,
877
"mamba_block_size": 16,
878
}
844
- with vllm_runner(**kwargs_with_apc) as vllm_model:
879
+ with _owned_vLLM_runner(vllm_runner, kwargs_with_apc) as vllm_model:
880
outputs_with_apc, _ = _get_vLLM_output(
881
vllm_runner,
882
kwargs_with_apc,