[CI][ROCm] Expand weight loading test coverage on AMD and cap its KV cache (#54037)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>
stefankoncarevic committed
Sep 1, 2026 at 09:57 UTC
b65af5e3396e161a64db6352241acf3afbdcd958
3 files changed
+21
-3
.buildkite/test_areas/weight_loading.yaml
+3
-3
@@ -17,10 +17,10 @@ steps:
17
- bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
18
mirror:
19
amd:
20
- label: ":amd: (MI300) Weight Loading Multi-GPU"
20
+ label: ":amd: (MI355) Weight Loading Multi-GPU"
21
dind: false
22
- device: mi300_2
23
- timeout_in_minutes: 35
22
+ device: mi355_2
23
+ timeout_in_minutes: 40
24
depends_on:
25
- image-build-amd
26
commands:
tests/weight_loading/models-amd.txt
+14
@@ -1,3 +1,17 @@
1
fp8, amd/Meta-Llama-3.1-8B-Instruct-FP8-KV, main
2
None, amd/Llama-3.2-1B-Instruct-FP8-KV, main
3
fp8, amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV, main
4
+gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main
5
+gptq, TheBloke/Llama-2-7B-GPTQ, main
6
+gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True
7
+compressed-tensors, nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change, main
8
+compressed-tensors, nm-testing/tinyllama-oneshot-w8-channel-a8-tensor, main
9
+compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2, main
10
+compressed-tensors, nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2, main
11
+compressed-tensors, nm-testing/tinyllama-oneshot-w4a16-group128-v2, main
12
+compressed-tensors, nm-testing/tinyllama-oneshot-w8a16-per-channel, main
13
+compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main
14
+compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main
15
+compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main
16
+awq, casperhansen/mixtral-instruct-awq, main
17
+fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main
tests/weight_loading/test_weight_loading.py
+4
@@ -7,6 +7,7 @@ import pytest
7
import torch
8
9
from vllm.platforms import current_platform
10
+from vllm.utils.mem_constants import GiB_bytes
11
12
MAX_MODEL_LEN = 1024
13
MODEL_NAME = os.environ.get(
@@ -41,6 +42,9 @@ def test_weight_loading(vllm_runner):
42
quantization=None if QUANTIZATION == "None" else QUANTIZATION,
43
max_model_len=MAX_MODEL_LEN,
44
tensor_parallel_size=2,
45
+ # Generating 20 tokens needs a tiny KV cache, while sizing and
46
+ # allocating a full-size one dominates this test on large devices.
47
+ kv_cache_memory_bytes=2 * GiB_bytes,
48
) as model:
49
output = model.generate_greedy("Hello world!", max_tokens=20)
50
print(output)