[CI][ROCm] Expand weight loading test coverage on AMD and cap its KV cache (#54037)

Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com>

.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:

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

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)