diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 4e78a3f626c..926475ecb13 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -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 diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index fb08feb2476..598558939c5 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -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 diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index e172983b54f..46790be6e4e 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -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.""" diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index bf9b15e7c78..57602289ce6 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -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=( diff --git a/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py b/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py index d88dfc31816..feb03d0d1a9 100644 --- a/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py +++ b/tests/v1/kv_connector/nixl_integration/test_nixl_imports.py @@ -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, diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index eed20e03668..d508f3cae2c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -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," diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 7ee072ceaf1..060fa5e3228 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -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. diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 37402dcaa0b..f8a30748b90 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -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: