From b65af5e3396e161a64db6352241acf3afbdcd958 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Tue, 1 Sep 2026 09:57:22 +0200 Subject: [PATCH] [CI][ROCm] Expand weight loading test coverage on AMD and cap its KV cache (#54037) Signed-off-by: Stefan Koncarevic --- .buildkite/test_areas/weight_loading.yaml | 6 +++--- tests/weight_loading/models-amd.txt | 14 ++++++++++++++ tests/weight_loading/test_weight_loading.py | 4 ++++ 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml index 89861d9ac1..154b355778 100644 --- a/.buildkite/test_areas/weight_loading.yaml +++ b/.buildkite/test_areas/weight_loading.yaml @@ -17,10 +17,10 @@ steps: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt mirror: amd: - label: ":amd: (MI300) Weight Loading Multi-GPU" + label: ":amd: (MI355) Weight Loading Multi-GPU" dind: false - device: mi300_2 - timeout_in_minutes: 35 + device: mi355_2 + timeout_in_minutes: 40 depends_on: - image-build-amd commands: diff --git a/tests/weight_loading/models-amd.txt b/tests/weight_loading/models-amd.txt index e31e904c08..816d136839 100644 --- a/tests/weight_loading/models-amd.txt +++ b/tests/weight_loading/models-amd.txt @@ -1,3 +1,17 @@ fp8, amd/Meta-Llama-3.1-8B-Instruct-FP8-KV, main None, amd/Llama-3.2-1B-Instruct-FP8-KV, main fp8, amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV, main +gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main +gptq, TheBloke/Llama-2-7B-GPTQ, main +gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True +compressed-tensors, nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change, main +compressed-tensors, nm-testing/tinyllama-oneshot-w8-channel-a8-tensor, main +compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2, main +compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2, main +compressed-tensors, nm-testing/tinyllama-oneshot-w4a16-group128-v2, main +compressed-tensors, nm-testing/tinyllama-oneshot-w8a16-per-channel, main +compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main +compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main +compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main +awq, casperhansen/mixtral-instruct-awq, main +fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main diff --git a/tests/weight_loading/test_weight_loading.py b/tests/weight_loading/test_weight_loading.py index 6587730682..72c0a23367 100644 --- a/tests/weight_loading/test_weight_loading.py +++ b/tests/weight_loading/test_weight_loading.py @@ -7,6 +7,7 @@ import pytest import torch from vllm.platforms import current_platform +from vllm.utils.mem_constants import GiB_bytes MAX_MODEL_LEN = 1024 MODEL_NAME = os.environ.get( @@ -41,6 +42,9 @@ def test_weight_loading(vllm_runner): quantization=None if QUANTIZATION == "None" else QUANTIZATION, max_model_len=MAX_MODEL_LEN, tensor_parallel_size=2, + # Generating 20 tokens needs a tiny KV cache, while sizing and + # allocating a full-size one dominates this test on large devices. + kv_cache_memory_bytes=2 * GiB_bytes, ) as model: output = model.generate_greedy("Hello world!", max_tokens=20) print(output)