@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3f5a1e1733

[ROCm][CI] Expand basic correctness target suites (#46573)

Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Matt <156021403+mawong-amd@users.noreply.github.com>

Andreas Karatzas committed Jun 24, 2026 at 23:18 UTC 3f5a1e1733200760169ff31ebe60a271072b199e

4 files changed +174 -60

.buildkite/test-amd.yaml

+4 -4

@@ -132,7 +132,7 @@ steps:

132 - vllm/_aiter_ops.py

133 - vllm/platforms/rocm.py

134 commands:

135 - - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

135 + - TARGET_TEST_SUITE=MI250 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

136 - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'

137 - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'

138 - pytest models/language -v -s -m 'distributed(num_gpus=2)'

@@ -646,7 +646,7 @@ steps:

646 commands:

647 - export VLLM_WORKER_MULTIPROC_METHOD=spawn

648 - pytest -v -s basic_correctness/test_mem.py

649 - - pytest -v -s basic_correctness/test_basic_correctness.py

649 + - VLLM_TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py

650 - pytest -v -s basic_correctness/test_cpu_offload.py

651

652 - label: Distributed Model Tests (2 GPUs) # TBD

@@ -668,7 +668,7 @@ steps:

668 - tests/model_executor/model_loader/test_sharded_state_loader.py

669 - tests/models/

670 commands:

671 - - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

671 + - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

672 - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'

673 - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'

674 - pytest models/language -v -s -m 'distributed(num_gpus=2)'

@@ -887,7 +887,7 @@ steps:

887 commands:

888 - pytest -v -s distributed/test_custom_all_reduce.py

889 - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py

890 - - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

890 + - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

891 - pytest -v -s -x lora/test_mixtral.py

892

893 - label: Distributed Torchrun + Examples (4 GPUs) # TBD

tests/basic_correctness/test_basic_correctness.py

+71 -12

@@ -29,7 +29,53 @@ MODELS = [

29 "meta-llama/Llama-3.2-1B-Instruct",

30 ]

32 -TARGET_TEST_SUITE = os.environ.get("TARGET_TEST_SUITE", "L4")

32 +TARGET_TEST_SUITE_ENV = "VLLM_TARGET_TEST_SUITE"

33 +LEGACY_TARGET_TEST_SUITE_ENV = "TARGET_TEST_SUITE"

34 +

35 +GENERIC_DISTRIBUTED_TEST_SUITES = ("L4", "MI250", "MI300", "MI325", "MI355")

36 +ALL_DISTRIBUTED_TEST_SUITES = (*GENERIC_DISTRIBUTED_TEST_SUITES, "A100")

37 +

38 +

39 +def _default_target_test_suite() -> str:

40 + if not current_platform.is_rocm():

41 + return "L4"

42 +

43 + try:

44 + device_name = current_platform.get_device_name().upper()

45 + except Exception:

46 + device_name = ""

47 +

48 + if "MI355" in device_name:

49 + return "MI355"

50 + if "MI300" in device_name:

51 + return "MI300"

52 + if "MI325" in device_name:

53 + return "MI325"

54 + if "MI250" in device_name:

55 + return "MI250"

56 +

57 + try:

58 + from vllm.platforms import rocm as rocm_platform

59 +

60 + if rocm_platform.on_gfx950():

61 + return "MI355"

62 + if rocm_platform.on_gfx942():

63 + return "MI300"

64 + except Exception:

65 + pass

66 +

67 + return "MI250"

68 +

69 +

70 +def _resolve_target_test_suite() -> str:

71 + for env_name in (TARGET_TEST_SUITE_ENV, LEGACY_TARGET_TEST_SUITE_ENV):

72 + value = os.environ.get(env_name, "").strip().upper()

73 + if value:

74 + return value

75 + return _default_target_test_suite()

76 +

77 +

78 +TARGET_TEST_SUITE = _resolve_target_test_suite()

81 def test_vllm_gc_ed():

@@ -131,14 +177,27 @@ def test_models(

177

178 @multi_gpu_test(num_gpus=2)

179 @pytest.mark.parametrize(

134 - "model, distributed_executor_backend, attention_backend, test_suite, extra_env",

180 + (

181 + "model, distributed_executor_backend, attention_backend, "

182 + "target_test_suites, extra_env"

183 + ),

184 [

136 - ("facebook/opt-125m", "ray", "", "L4", {}),

137 - ("facebook/opt-125m", "mp", "", "L4", {}),

138 - ("meta-llama/Llama-3.2-1B-Instruct", "ray", "", "L4", {}),

139 - ("meta-llama/Llama-3.2-1B-Instruct", "mp", "", "L4", {}),

140 - ("facebook/opt-125m", "ray", "", "A100", {}),

141 - ("facebook/opt-125m", "mp", "", "A100", {}),

185 + ("facebook/opt-125m", "ray", "", ALL_DISTRIBUTED_TEST_SUITES, {}),

186 + ("facebook/opt-125m", "mp", "", ALL_DISTRIBUTED_TEST_SUITES, {}),

187 + (

188 + "meta-llama/Llama-3.2-1B-Instruct",

189 + "ray",

190 + "",

191 + GENERIC_DISTRIBUTED_TEST_SUITES,

192 + {},

193 + ),

194 + (

195 + "meta-llama/Llama-3.2-1B-Instruct",

196 + "mp",

197 + "",

198 + GENERIC_DISTRIBUTED_TEST_SUITES,

199 + {},

200 + ),

201 ],

202 )

203 @pytest.mark.parametrize("enable_prompt_embeds", [True, False])

@@ -150,19 +209,19 @@ def test_models_distributed(

209 model: str,

210 distributed_executor_backend: str,

211 attention_backend: str,

153 - test_suite: str,

212 + target_test_suites: tuple[str, ...],

213 extra_env: dict[str, str],

214 enable_prompt_embeds: bool,

215 ) -> None:

157 - if test_suite != TARGET_TEST_SUITE:

158 - pytest.skip(f"Skip test for {test_suite}")

216 + if TARGET_TEST_SUITE and TARGET_TEST_SUITE not in target_test_suites:

217 + pytest.skip(f"Skip test for {TARGET_TEST_SUITE}")

218

219 with monkeypatch.context() as monkeypatch_context:

220 if (

221 model == "meta-llama/Llama-3.2-1B-Instruct"

222 and distributed_executor_backend == "ray"

223 and attention_backend == ""

165 - and test_suite == "L4"

224 + and TARGET_TEST_SUITE == "L4"

225 and enable_prompt_embeds

226 ): # noqa

227 pytest.skip("enable_prompt_embeds does not work with ray compiled dag.")

tests/conftest.py

+24 -1

@@ -64,6 +64,7 @@ from vllm.logprobs import Logprob

64 from vllm.multimodal.media import MediaWithBytes

65 from vllm.multimodal.utils import fetch_image

66 from vllm.outputs import RequestOutput

67 +from vllm.platforms import current_platform

68 from vllm.sampling_params import BeamSearchParams

69 from vllm.transformers_utils.utils import maybe_model_redirect

70 from vllm.utils.collection_utils import is_list_of

@@ -852,6 +853,24 @@ class HfRunner:

853 return self.model.predict(prompts, *args, convert_to_tensor=True, **kwargs)

854

855 def __enter__(self):

856 + if current_platform.is_rocm():

857 + # Record starting memory usage stats on ROCm so that we can wait for

858 + # memory to roughly settle back below these levels on shutdown. This is

859 + # helpful in cases where the HfRunner is initialized after significant GPU

860 + # memory is already occupied, e.g. in

861 + # tests/basic_correctness/test_basic_correctness.py::test_models_distributed

862 + from tests.utils import (

863 + get_physical_device_indices,

864 + record_gpu_memory_usage_stats,

865 + )

866 +

867 + if (device_count := current_platform.device_count()) > 0:

868 + devices = get_physical_device_indices(devices=list(range(device_count)))

869 + mem_usage_stats = record_gpu_memory_usage_stats(devices=devices)

870 + self.threshold_ratios = {

871 + device: 0.05 + mem_used / mem_tot

872 + for device, (mem_used, mem_tot) in mem_usage_stats.items()

873 + }

874 return self

875

876 def __exit__(self, exc_type, exc_value, traceback):

@@ -861,7 +880,11 @@ class HfRunner:

880 cleanup_dist_env_and_memory()

881 # ROCm frees VRAM lazily; wait so a runner started right after this HF

882 # model exits does not OOM on its startup memory guard.

864 - wait_for_rocm_memory_to_settle()

883 + wait_for_rocm_memory_to_settle(

884 + threshold_ratio=getattr(self, "threshold_ratios", None)

885 + )

886 + if hasattr(self, "threshold_ratios"):

887 + del self.threshold_ratios

888

889

890 @pytest.fixture(scope="session")

tests/utils.py

+75 -43

@@ -1464,7 +1464,7 @@ def error_on_warning(category: type[Warning] = Warning):

1464 yield

1465

1466

1467 -def get_physical_device_indices(devices):

1467 +def get_physical_device_indices(devices: list[int]):

1468 visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")

1469 if visible_devices is None:

1470 return devices

@@ -1475,82 +1475,114 @@ def get_physical_device_indices(devices):

1475

1476

1477 @_nvml()

1478 +def record_gpu_memory_usage_stats(

1479 + *,

1480 + devices: list[int],

1481 +) -> dict[int, tuple[float, float]]:

1482 + output: dict[int, tuple[float, float]] = {}

1483 + for device in devices:

1484 + if current_platform.is_rocm():

1485 + dev_handle = amdsmi_get_processor_handles()[device]

1486 + mem_info = amdsmi_get_gpu_vram_usage(dev_handle)

1487 + gb_used = mem_info["vram_used"] / 2**10

1488 + gb_total = mem_info["vram_total"] / 2**10

1489 + else:

1490 + dev_handle = nvmlDeviceGetHandleByIndex(device)

1491 + mem_info = nvmlDeviceGetMemoryInfo(dev_handle)

1492 + gb_used = mem_info.used / 2**30

1493 + gb_total = mem_info.total / 2**30

1494 + output[device] = (gb_used, gb_total)

1495 + return output

1496 +

1497 +

1498 def wait_for_gpu_memory_to_clear(

1499 *,

1500 devices: list[int],

1481 - threshold_bytes: int | None = None,

1482 - threshold_ratio: float | None = None,

1501 + threshold_bytes: int | dict[int, int] | None = None,

1502 + threshold_ratio: float | dict[int, float] | None = None,

1503 timeout_s: float = 120,

1504 ) -> None:

1505 assert threshold_bytes is not None or threshold_ratio is not None

1486 - if (

1487 - current_platform.is_rocm()

1488 - and threshold_ratio is not None

1489 - and threshold_ratio < 0.05

1490 - ):

1506 + devices = get_physical_device_indices(devices)

1507 + if isinstance(threshold_bytes, int):

1508 + threshold_bytes = {device: threshold_bytes for device in devices}

1509 + elif isinstance(threshold_bytes, dict):

1510 + assert threshold_bytes.keys() == set(devices)

1511 + if isinstance(threshold_ratio, float):

1512 + threshold_ratio = {device: threshold_ratio for device in devices}

1513 + elif isinstance(threshold_ratio, dict):

1514 + assert threshold_ratio.keys() == set(devices)

1515 + if current_platform.is_rocm() and threshold_ratio is not None:

1516 # ROCm can keep a small runtime/driver footprint resident even after

1517 # all model allocations are gone. On MI300 this has been observed

1518 # around 2.5 GiB, which is above a strict 1% idle threshold but nowhere

1519 # near the amount of free memory needed by the next vLLM runner.

1495 - min_threshold_bytes = 4 * 1024**3

1496 - threshold_bytes = max(threshold_bytes or 0, min_threshold_bytes)

1520 + min_threshold_b = 4 * 1024**3

1521 + if threshold_bytes is None:

1522 + threshold_bytes = {}

1523 + for device, ratio in threshold_ratio.items():

1524 + threshold_bytes[device] = max(

1525 + threshold_bytes.get(device, 0), min_threshold_b if ratio < 0.05 else 0

1526 + )

1527

1528 # Use nvml instead of pytorch to reduce measurement error from torch cuda

1529 # context.

1500 - devices = get_physical_device_indices(devices)

1530 start_time = time.time()

1531 while True:

1503 - output: dict[int, str] = {}

1504 - output_raw: dict[int, tuple[float, float]] = {}

1505 - for device in devices:

1506 - if current_platform.is_rocm():

1507 - dev_handle = amdsmi_get_processor_handles()[device]

1508 - mem_info = amdsmi_get_gpu_vram_usage(dev_handle)

1509 - gb_used = mem_info["vram_used"] / 2**10

1510 - gb_total = mem_info["vram_total"] / 2**10

1511 - else:

1512 - dev_handle = nvmlDeviceGetHandleByIndex(device)

1513 - mem_info = nvmlDeviceGetMemoryInfo(dev_handle)

1514 - gb_used = mem_info.used / 2**30

1515 - gb_total = mem_info.total / 2**30

1516 - output_raw[device] = (gb_used, gb_total)

1517 - output[device] = f"{gb_used:.02f}/{gb_total:.02f}"

1518 -

1532 + output_raw = record_gpu_memory_usage_stats(devices=devices)

1533 + output = {

1534 + device: f"{gb_used:.02f}/{gb_total:.02f}"

1535 + for device, (gb_used, gb_total) in output_raw.items()

1536 + }

1537 print("gpu memory used/total (GiB): ", end="")

1538 for k, v in output.items():

1539 print(f"{k}={v}; ", end="")

1540 print("")

1541

1542 if threshold_bytes is not None and threshold_ratio is not None:

1525 - threshold_gib = threshold_bytes / 2**30

1526 - threshold = f"max({threshold_gib:.2f} GiB, {threshold_ratio:.3f})"

1543 + threshold_gib = {

1544 + device: threshold_b / 2**30

1545 + for device, threshold_b in threshold_bytes.items()

1546 + }

1547 + threshold = "; ".join(

1548 + f"{device=}: max({threshold_gib[device]:.2f} GiB, "

1549 + f"{threshold_ratio[device]:.3f})"

1550 + for device in devices

1551 + )

1552 all_free = all(

1528 - used <= max(threshold_gib, total * threshold_ratio)

1529 - for used, total in output_raw.values()

1553 + used <= max(threshold_gib[device], total * threshold_ratio[device])

1554 + for device, (used, total) in output_raw.items()

1555 )

1556 elif threshold_bytes is not None:

1532 - threshold_gib = threshold_bytes / 2**30

1533 - threshold = f"{threshold_gib} GiB"

1534 - all_free = all(used <= threshold_gib for used, _ in output_raw.values())

1557 + threshold_gib = {

1558 + device: threshold_b / 2**30

1559 + for device, threshold_b in threshold_bytes.items()

1560 + }

1561 + threshold = "; ".join(

1562 + f"{device=}: {threshold_gib[device]:.2f} GiB" for device in devices

1563 + )

1564 + all_free = all(

1565 + used <= threshold_gib[device]

1566 + for device, (used, _) in output_raw.items()

1567 + )

1568 else:

1569 assert threshold_ratio is not None

1537 - threshold = f"{threshold_ratio:.3f}"

1570 + threshold = "; ".join(

1571 + f"{device=}: {threshold_ratio[device]:.3f}" for device in devices

1572 + )

1573 all_free = all(

1539 - used / total <= threshold_ratio for used, total in output_raw.values()

1574 + used / total <= threshold_ratio[device]

1575 + for device, (used, total) in output_raw.items()

1576 )

1577

1578 dur_s = time.time() - start_time

1579 if all_free:

1544 - print(

1545 - f"Done waiting for free GPU memory on devices {devices=} "

1546 - f"({threshold=}) {dur_s=:.02f}"

1547 - )

1580 + print(f"Done waiting for free GPU memory on ({threshold=}) {dur_s=:.02f}")

1581 break

1582

1583 if dur_s >= timeout_s:

1584 raise ValueError(

1552 - f"Memory of devices {devices=} not free after "

1553 - f"{dur_s=:.02f} ({threshold=})"

1585 + f"Memory of devices not free after {dur_s=:.02f} ({threshold=})"

1586 )

1587

1588 time.sleep(5)

@@ -1558,7 +1590,7 @@ def wait_for_gpu_memory_to_clear(

1590

1591 def wait_for_rocm_memory_to_settle(

1592 *,

1561 - threshold_ratio: float = 0.1,

1593 + threshold_ratio: float | dict[int, float] | None = 0.1,

1594 timeout_s: float = 240,

1595 ) -> None:

1596 """Block until ROCm device VRAM usage drops below ``threshold_ratio``.