[ROCm][CI] Fix nixl tests (#45219)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jun 23, 2026 at 13:11 UTC
ceae5bcbda06ac41ff76317c4475e8b1165e419e
8 files changed
+210
-15
.buildkite/test-amd.yaml
+21
-2
@@ -2279,11 +2279,28 @@ steps:
2279
commands:
2280
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
2281
2282
+- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
2283
+ timeout_in_minutes: 180
2284
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2285
+ agent_pool: mi300_2
2286
+ num_gpus: 2
2287
+ optional: true
2288
+ working_dir: "/vllm-workspace/tests"
2289
+ source_file_dependencies:
2290
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
2291
+ - vllm/v1/worker/kv_connector_model_runner_mixin.py
2292
+ - tests/v1/kv_connector/nixl_integration/
2293
+ - vllm/platforms/rocm.py
2294
+ commands:
2295
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
2296
+ - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
2297
+
2298
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
2299
timeout_in_minutes: 180
2300
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2301
agent_pool: mi300_4
2302
num_gpus: 4
2303
+ optional: true
2304
working_dir: "/vllm-workspace/tests"
2305
source_file_dependencies:
2306
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
@@ -2322,7 +2339,7 @@ steps:
2339
optional: true
2340
working_dir: "/vllm-workspace/tests"
2341
source_file_dependencies:
2325
- - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
2342
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
2343
- tests/v1/kv_connector/nixl_integration/
2344
- vllm/platforms/rocm.py
2345
commands:
@@ -2334,9 +2351,10 @@ steps:
2351
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2352
agent_pool: mi300_4
2353
num_gpus: 4
2354
+ optional: true
2355
working_dir: "/vllm-workspace/tests"
2356
source_file_dependencies:
2339
- - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
2357
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
2358
- tests/v1/kv_connector/nixl_integration/
2359
- vllm/platforms/rocm.py
2360
commands:
@@ -2348,6 +2366,7 @@ steps:
2366
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2367
agent_pool: mi300_4
2368
num_gpus: 4
2369
+ optional: true
2370
working_dir: "/vllm-workspace/tests"
2371
source_file_dependencies:
2372
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
.buildkite/test_areas/disaggregated.yaml
+67
@@ -13,6 +13,20 @@ steps:
13
commands:
14
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
15
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
16
+ mirror:
17
+ amd:
18
+ device: mi300_4
19
+ timeout_in_minutes: 110
20
+ depends_on:
21
+ - image-build-amd
22
+ source_file_dependencies:
23
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
24
+ - tests/v1/kv_connector/nixl_integration/
25
+ - vllm/platforms/rocm.py
26
+ commands:
27
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
28
+ - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
29
+
30
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
31
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
32
timeout_in_minutes: 30
@@ -36,6 +50,19 @@ steps:
50
commands:
51
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
52
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
53
+ mirror:
54
+ amd:
55
+ device: mi300_4
56
+ timeout_in_minutes: 50
57
+ depends_on:
58
+ - image-build-amd
59
+ source_file_dependencies:
60
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
61
+ - tests/v1/kv_connector/nixl_integration/
62
+ - vllm/platforms/rocm.py
63
+ commands:
64
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
65
+ - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
66
67
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
68
key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
@@ -48,6 +75,19 @@ steps:
75
commands:
76
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
77
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
78
+ mirror:
79
+ amd:
80
+ device: mi300_4
81
+ timeout_in_minutes: 110
82
+ depends_on:
83
+ - image-build-amd
84
+ source_file_dependencies:
85
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
86
+ - tests/v1/kv_connector/nixl_integration/
87
+ - vllm/platforms/rocm.py
88
+ commands:
89
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
90
+ - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
91
92
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
93
key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
@@ -60,6 +100,19 @@ steps:
100
commands:
101
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
102
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
103
+ mirror:
104
+ amd:
105
+ device: mi300_4
106
+ timeout_in_minutes: 60
107
+ depends_on:
108
+ - image-build-amd
109
+ source_file_dependencies:
110
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
111
+ - tests/v1/kv_connector/nixl_integration/
112
+ - vllm/platforms/rocm.py
113
+ commands:
114
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
115
+ - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
116
117
- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
118
key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
@@ -103,6 +156,20 @@ steps:
156
commands:
157
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
158
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
159
+ mirror:
160
+ amd:
161
+ device: mi300_2
162
+ timeout_in_minutes: 60
163
+ depends_on:
164
+ - image-build-amd
165
+ source_file_dependencies:
166
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
167
+ - vllm/v1/worker/kv_connector_model_runner_mixin.py
168
+ - tests/v1/kv_connector/nixl_integration/
169
+ - vllm/platforms/rocm.py
170
+ commands:
171
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
172
+ - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
173
174
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
175
key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
tests/config/test_model_arch_config.py
+17
@@ -6,6 +6,7 @@ import json
6
from pathlib import Path
7
8
import pytest
9
+from transformers import PretrainedConfig
10
11
from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
12
from vllm.transformers_utils.model_arch_config_convertor import (
@@ -114,6 +115,22 @@ def _assert_model_config_methods(
115
assert model_config.get_head_size() == expected["head_size"]
116
117
118
+def test_head_size_falls_back_when_head_dim_is_zero():
119
+ """Regression test for configs that materialize missing head_dim as 0."""
120
+ hf_config = PretrainedConfig(
121
+ model_type="deepseek_vl_v2",
122
+ hidden_size=1280,
123
+ num_attention_heads=10,
124
+ num_key_value_heads=10,
125
+ head_dim=0,
126
+ kv_lora_rank=None,
127
+ )
128
+
129
+ convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
130
+
131
+ assert convertor.get_head_size() == 128
132
+
133
+
134
@pytest.mark.parametrize("model", BASE_MODELS_TO_TEST)
135
def test_base_model_arch_config(model: str):
136
"""Test model architecture config for base models."""
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+3
-3
@@ -13,13 +13,13 @@ tp_configs=(
13
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2"
14
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2"
15
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1"
16
- "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA case
16
+ "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
17
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
18
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
19
)
20
dp_ep_configs=(
21
-"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
22
-"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
21
+"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # P-TP1, D-DPEP=2 (TP=1)
22
+"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # P-TP2, D-DPEP=2 (TP=1)
23
)
24
# We assume HMA enabled by default.
25
hybrid_ssm_configs=(
tests/v1/kv_connector/nixl_integration/test_nixl_imports.py
+34
-7
@@ -4,9 +4,9 @@
4
5
import importlib
6
import importlib.metadata as metadata
7
-import pathlib
7
import subprocess
8
import sys
9
+import types
10
11
import pytest
12
import torch
@@ -20,6 +20,34 @@ def _print_distribution_version(package_name: str) -> None:
20
print(f"{package_name}: {version}")
21
22
23
+def _import_nixl_ep_cpp(nixl_ep: types.ModuleType) -> types.ModuleType:
24
+ candidate_module_names = []
25
+
26
+ config_module_name = getattr(getattr(nixl_ep, "Config", None), "__module__", None)
27
+ if config_module_name and config_module_name.endswith("nixl_ep_cpp"):
28
+ candidate_module_names.append(config_module_name)
29
+
30
+ if torch.version.cuda is not None:
31
+ cuda_major = torch.version.cuda.split(".", maxsplit=1)[0]
32
+ candidate_module_names.append(f"nixl_ep_cu{cuda_major}.nixl_ep_cpp")
33
+
34
+ # Keep compatibility with the pre-dispatcher wheel layout.
35
+ candidate_module_names.append("nixl_ep.nixl_ep_cpp")
36
+
37
+ for module_name in dict.fromkeys(candidate_module_names):
38
+ try:
39
+ return importlib.import_module(module_name)
40
+ except ModuleNotFoundError as exc:
41
+ missing_module = exc.name
42
+ if missing_module not in (module_name, module_name.split(".", 1)[0]):
43
+ raise
44
+
45
+ raise AssertionError(
46
+ "No nixl_ep_cpp extension module found; tried "
47
+ f"{', '.join(dict.fromkeys(candidate_module_names))}"
48
+ )
49
+
50
+
51
@pytest.mark.skipif(torch.version.cuda is None, reason="CUDA NIXL EP canary")
52
def test_nixl_and_nixl_ep_imports() -> None:
53
"""Verify both core NIXL and the NIXL EP extension import successfully."""
@@ -38,14 +66,13 @@ def test_nixl_and_nixl_ep_imports() -> None:
66
nixl_ep = importlib.import_module("nixl_ep")
67
print(f"nixl_ep: {nixl_ep.__file__}")
68
41
- assert nixl_ep.__file__ is not None
42
- extension_dir = pathlib.Path(nixl_ep.__file__).parent
43
- extension_files = sorted(extension_dir.glob("nixl_ep_cpp*.so"))
44
- assert extension_files, f"No nixl_ep_cpp extension found in {extension_dir}"
69
+ nixl_ep_cpp = _import_nixl_ep_cpp(nixl_ep)
70
+ assert nixl_ep_cpp.__file__ is not None
71
+ extension_file = nixl_ep_cpp.__file__
72
+ print(f"nixl_ep_cpp: {extension_file}")
73
46
- extension_file = extension_files[0]
74
completed = subprocess.run(
48
- ["ldd", str(extension_file)],
75
+ ["ldd", extension_file],
76
capture_output=True,
77
check=False,
78
text=True,
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+44
@@ -94,6 +94,50 @@ def test_logical_to_kernel_block_ids_with_hma():
94
)
95
96
97
+@pytest.mark.cpu_test
98
+@pytest.mark.parametrize(
99
+ "is_rocm,has_mamba,use_host_buffer,done_recving,failed_recving,expected_syncs",
100
+ [
101
+ (True, True, False, {"req"}, set(), 1),
102
+ (False, True, False, {"req"}, set(), 0),
103
+ (True, False, False, {"req"}, set(), 0),
104
+ (True, True, True, {"req"}, set(), 0),
105
+ (True, True, False, set(), set(), 0),
106
+ (True, True, False, {"req"}, {"req"}, 0),
107
+ ],
108
+)
109
+def test_sync_device_after_mamba_recv_gates(
110
+ monkeypatch,
111
+ is_rocm,
112
+ has_mamba,
113
+ use_host_buffer,
114
+ done_recving,
115
+ failed_recving,
116
+ expected_syncs,
117
+):
118
+ """Only direct-GPU Mamba receives on ROCm need a device fence."""
119
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
120
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
121
+ NixlConnectorWorker,
122
+ )
123
+
124
+ worker = object.__new__(NixlConnectorWorker)
125
+ worker._has_mamba = has_mamba
126
+ worker.use_host_buffer = use_host_buffer
127
+
128
+ sync_calls = []
129
+ monkeypatch.setattr(base_worker.current_platform, "is_rocm", lambda: is_rocm)
130
+ monkeypatch.setattr(
131
+ base_worker.torch.accelerator,
132
+ "synchronize",
133
+ lambda: sync_calls.append(True),
134
+ )
135
+
136
+ worker._sync_device_after_mamba_recv(done_recving, failed_recving)
137
+
138
+ assert len(sync_calls) == expected_syncs
139
+
140
+
141
@pytest.mark.cpu_test
142
@pytest.mark.parametrize(
143
"group_spec_types,remote_physical_per_logical,"
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+18
@@ -1951,6 +1951,8 @@ class NixlBaseConnectorWorker:
1951
for block_ids in block_ids_for_heterogeneous_attn_post_process:
1952
self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
1953
1954
+ self._sync_device_after_mamba_recv(done_recving, failed_recv_reqs)
1955
+
1956
# Handle timeout to avoid stranding blocks on remote.
1957
now = time.perf_counter()
1958
while self._reqs_to_send:
@@ -1972,6 +1974,22 @@ class NixlBaseConnectorWorker:
1974
1975
return done_sending, done_recving
1976
1977
+ def _sync_device_after_mamba_recv(
1978
+ self,
1979
+ done_recving: set[str],
1980
+ failed_recv_reqs: set[str],
1981
+ ) -> None:
1982
+ """Synchronize ROCm direct-GPU Mamba receives before model execution."""
1983
+ if (
1984
+ not current_platform.is_rocm()
1985
+ or not self._has_mamba
1986
+ or self.use_host_buffer
1987
+ or not (done_recving - failed_recv_reqs)
1988
+ ):
1989
+ return
1990
+
1991
+ torch.accelerator.synchronize()
1992
+
1993
def _get_new_notifs(self) -> set[str]:
1994
"""Get req_ids which got a remote xfer notification.
1995
vllm/transformers_utils/model_arch_config_convertor.py
+6
-3
@@ -58,9 +58,12 @@ class ModelArchConfigConvertorBase:
58
if qk_rope_head_dim and qk_nope_head_dim:
59
return qk_rope_head_dim + qk_nope_head_dim
60
61
- # NOTE: Some configs may set head_dim=None in the config
62
- if getattr(self.hf_text_config, "head_dim", None) is not None:
63
- return self.hf_text_config.head_dim
61
+ # NOTE: Some config classes may set head_dim=None or materialize a missing
62
+ # head_dim as 0 (for example, DeepseekVLV2TextConfig).
63
+ if (
64
+ head_dim := getattr(self.hf_text_config, "head_dim", None)
65
+ ) is not None and head_dim > 0:
66
+ return head_dim
67
68
# NOTE: Some models (such as PLaMo2.1) use `hidden_size_per_head`
69
if getattr(self.hf_text_config, "hidden_size_per_head", None) is not None: