karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[ROCm][CI] Fix nixl tests (#45219)

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Andreas Karatzas
2026-06-23 13:11:40 -05:00
committed by GitHub
parent 6691f087a6
commit ceae5bcbda
8 changed files with 210 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test-amd.yaml
+21 -2
View File
@@ -2279,11 +2279,28 @@ steps:
commands:
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_2
num_gpus: 2
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_4
num_gpus: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
@@ -2322,7 +2339,7 @@ steps:
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
@@ -2334,9 +2351,10 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_4
num_gpus: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
@@ -2348,6 +2366,7 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_4
num_gpus: 4
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
.buildkite/test_areas/disaggregated.yaml
+67
View File
@@ -13,6 +13,20 @@ steps:
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
mirror:
amd:
device: mi300_4
timeout_in_minutes: 110
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
timeout_in_minutes: 30
@@ -36,6 +50,19 @@ steps:
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
mirror:
amd:
device: mi300_4
timeout_in_minutes: 50
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
@@ -48,6 +75,19 @@ steps:
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
mirror:
amd:
device: mi300_4
timeout_in_minutes: 110
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
@@ -60,6 +100,19 @@ steps:
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
mirror:
amd:
device: mi300_4
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
@@ -103,6 +156,20 @@ steps:
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
mirror:
amd:
device: mi300_2
timeout_in_minutes: 60
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
- vllm/platforms/rocm.py
commands:
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
tests/config/test_model_arch_config.py
+17
View File
@@ -6,6 +6,7 @@ import json
from pathlib import Path
import pytest
from transformers import PretrainedConfig
from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
from vllm.transformers_utils.model_arch_config_convertor import (
@@ -114,6 +115,22 @@ def _assert_model_config_methods(
assert model_config.get_head_size() == expected["head_size"]
def test_head_size_falls_back_when_head_dim_is_zero():
"""Regression test for configs that materialize missing head_dim as 0."""
hf_config = PretrainedConfig(
model_type="deepseek_vl_v2",
hidden_size=1280,
num_attention_heads=10,
num_key_value_heads=10,
head_dim=0,
kv_lora_rank=None,
)
convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
assert convertor.get_head_size() == 128
@pytest.mark.parametrize("model", BASE_MODELS_TO_TEST)
def test_base_model_arch_config(model: str):
"""Test model architecture config for base models."""
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+3 -3
View File
@@ -13,13 +13,13 @@ tp_configs=(
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2"
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2"
"GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1"
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA case
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
"GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
)
dp_ep_configs=(
"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)
"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)
"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)
"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)
)
# We assume HMA enabled by default.
hybrid_ssm_configs=(
tests/v1/kv_connector/nixl_integration/test_nixl_imports.py
+34 -7
View File
@@ -4,9 +4,9 @@
import importlib
import importlib.metadata as metadata
import pathlib
import subprocess
import sys
import types
import pytest
import torch
@@ -20,6 +20,34 @@ def _print_distribution_version(package_name: str) -> None:
print(f"{package_name}: {version}")
def _import_nixl_ep_cpp(nixl_ep: types.ModuleType) -> types.ModuleType:
candidate_module_names = []
config_module_name = getattr(getattr(nixl_ep, "Config", None), "__module__", None)
if config_module_name and config_module_name.endswith("nixl_ep_cpp"):
candidate_module_names.append(config_module_name)
if torch.version.cuda is not None:
cuda_major = torch.version.cuda.split(".", maxsplit=1)[0]
candidate_module_names.append(f"nixl_ep_cu{cuda_major}.nixl_ep_cpp")
# Keep compatibility with the pre-dispatcher wheel layout.
candidate_module_names.append("nixl_ep.nixl_ep_cpp")
for module_name in dict.fromkeys(candidate_module_names):
try:
return importlib.import_module(module_name)
except ModuleNotFoundError as exc:
missing_module = exc.name
if missing_module not in (module_name, module_name.split(".", 1)[0]):
raise
raise AssertionError(
"No nixl_ep_cpp extension module found; tried "
f"{', '.join(dict.fromkeys(candidate_module_names))}"
)
@pytest.mark.skipif(torch.version.cuda is None, reason="CUDA NIXL EP canary")
def test_nixl_and_nixl_ep_imports() -> None:
"""Verify both core NIXL and the NIXL EP extension import successfully."""
@@ -38,14 +66,13 @@ def test_nixl_and_nixl_ep_imports() -> None:
nixl_ep = importlib.import_module("nixl_ep")
print(f"nixl_ep: {nixl_ep.__file__}")
assert nixl_ep.__file__ is not None
extension_dir = pathlib.Path(nixl_ep.__file__).parent
extension_files = sorted(extension_dir.glob("nixl_ep_cpp*.so"))
assert extension_files, f"No nixl_ep_cpp extension found in {extension_dir}"
nixl_ep_cpp = _import_nixl_ep_cpp(nixl_ep)
assert nixl_ep_cpp.__file__ is not None
extension_file = nixl_ep_cpp.__file__
print(f"nixl_ep_cpp: {extension_file}")
extension_file = extension_files[0]
completed = subprocess.run(
["ldd", str(extension_file)],
["ldd", extension_file],
capture_output=True,
check=False,
text=True,
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+44
View File
@@ -94,6 +94,50 @@ def test_logical_to_kernel_block_ids_with_hma():
)
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"is_rocm,has_mamba,use_host_buffer,done_recving,failed_recving,expected_syncs",
[
(True, True, False, {"req"}, set(), 1),
(False, True, False, {"req"}, set(), 0),
(True, False, False, {"req"}, set(), 0),
(True, True, True, {"req"}, set(), 0),
(True, True, False, set(), set(), 0),
(True, True, False, {"req"}, {"req"}, 0),
],
)
def test_sync_device_after_mamba_recv_gates(
monkeypatch,
is_rocm,
has_mamba,
use_host_buffer,
done_recving,
failed_recving,
expected_syncs,
):
"""Only direct-GPU Mamba receives on ROCm need a device fence."""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
NixlConnectorWorker,
)
worker = object.__new__(NixlConnectorWorker)
worker._has_mamba = has_mamba
worker.use_host_buffer = use_host_buffer
sync_calls = []
monkeypatch.setattr(base_worker.current_platform, "is_rocm", lambda: is_rocm)
monkeypatch.setattr(
base_worker.torch.accelerator,
"synchronize",
lambda: sync_calls.append(True),
)
worker._sync_device_after_mamba_recv(done_recving, failed_recving)
assert len(sync_calls) == expected_syncs
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"group_spec_types,remote_physical_per_logical,"
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+18
View File
@@ -1951,6 +1951,8 @@ class NixlBaseConnectorWorker:
for block_ids in block_ids_for_heterogeneous_attn_post_process:
self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
self._sync_device_after_mamba_recv(done_recving, failed_recv_reqs)
# Handle timeout to avoid stranding blocks on remote.
now = time.perf_counter()
while self._reqs_to_send:
@@ -1972,6 +1974,22 @@ class NixlBaseConnectorWorker:
return done_sending, done_recving
def _sync_device_after_mamba_recv(
self,
done_recving: set[str],
failed_recv_reqs: set[str],
) -> None:
"""Synchronize ROCm direct-GPU Mamba receives before model execution."""
if (
not current_platform.is_rocm()
or not self._has_mamba
or self.use_host_buffer
or not (done_recving - failed_recv_reqs)
):
return
torch.accelerator.synchronize()
def _get_new_notifs(self) -> set[str]:
"""Get req_ids which got a remote xfer notification.
vllm/transformers_utils/model_arch_config_convertor.py
+6 -3
View File
@@ -58,9 +58,12 @@ class ModelArchConfigConvertorBase:
if qk_rope_head_dim and qk_nope_head_dim:
return qk_rope_head_dim + qk_nope_head_dim
# NOTE: Some configs may set head_dim=None in the config
if getattr(self.hf_text_config, "head_dim", None) is not None:
return self.hf_text_config.head_dim
# NOTE: Some config classes may set head_dim=None or materialize a missing
# head_dim as 0 (for example, DeepseekVLV2TextConfig).
if (
head_dim := getattr(self.hf_text_config, "head_dim", None)
) is not None and head_dim > 0:
return head_dim
# NOTE: Some models (such as PLaMo2.1) use `hidden_size_per_head`
if getattr(self.hf_text_config, "hidden_size_per_head", None) is not None: