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 11, 15:00 UTC
[ROCm][CI] Skip multi-GPU speculative decoding tests when insufficient GPUs available (#30527)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Dec 11, 2025 at 21:54 UTC
783644e4ac7d4ee324a1817bbe199fc5b557bc7d
1 file changed
+13
tests/v1/e2e/test_spec_decode.py
+13
@@ -16,6 +16,16 @@ from vllm.platforms import current_platform
16
MTP_SIMILARITY_RATE = 0.8
17
18
19
+def _skip_if_insufficient_gpus_for_tp(tp_size: int):
20
+ """Skip test if available GPUs < tp_size on ROCm."""
21
+ if current_platform.is_rocm():
22
+ available_gpus = torch.cuda.device_count()
23
+ if available_gpus < tp_size:
24
+ pytest.skip(
25
+ f"Test requires {tp_size} GPUs, but only {available_gpus} available"
26
+ )
27
+
28
+
29
def get_test_prompts(mm_enabled: bool):
30
prompt_types = ["repeat", "sentence"]
31
if mm_enabled:
@@ -455,6 +465,8 @@ def test_eagle_correctness(
465
m.setenv("VLLM_ROCM_USE_AITER", "1")
466
467
method, model_name, spec_model_name, tp_size = model_setup
468
+ _skip_if_insufficient_gpus_for_tp(tp_size)
469
+
470
max_model_len = 2048
471
max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len
472
@@ -525,6 +537,7 @@ def test_mtp_correctness(
537
m.setenv("VLLM_MLA_DISABLE", "1")
538
539
method, model_name, tp_size = model_setup
540
+ _skip_if_insufficient_gpus_for_tp(tp_size)
541
542
ref_llm = LLM(
543
model=model_name,