[ROCm][CI] Adding extract hs 2gpu (#47482)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 2, 2026 at 19:59 UTC
442ccc60983b4acc8a5a9c793c65509814dbc552
2 files changed
+4
-2
.buildkite/test-amd.yaml
+2
-1
@@ -1996,7 +1996,8 @@ steps:
1996
- label: Extract Hidden States Integration # TBD
1997
timeout_in_minutes: 180
1998
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
1999
- agent_pool: mi300_1
1999
+ agent_pool: mi300_2
2000
+ num_gpus: 2
2001
working_dir: "/vllm-workspace/tests"
2002
source_file_dependencies:
2003
- vllm/config/speculative.py
tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
+2
-1
@@ -12,6 +12,7 @@ from vllm import LLM, ModelRegistry, SamplingParams
12
from vllm.distributed.kv_transfer.kv_connector.v1 import (
13
example_hidden_states_connector,
14
)
15
+from vllm.platforms import current_platform
16
17
18
def get_and_check_output(output, expected_shape):
@@ -296,7 +297,7 @@ def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path):
297
)
298
299
299
-@pytest.mark.timeout(60)
300
+@pytest.mark.timeout(240 if current_platform.is_rocm() else 60)
301
@multi_gpu_test(num_gpus=2)
302
@create_new_process_for_each_test()
303
def test_extract_hidden_states_tp2():