diff --git a/.github/vllm-main-verified.commit b/.github/vllm-main-verified.commit index 6bb1718ee3c..452da533896 100644 --- a/.github/vllm-main-verified.commit +++ b/.github/vllm-main-verified.commit @@ -1 +1 @@ -b9a7cd464c9ae9b1b450f8982b76d7be4de73724 \ No newline at end of file +ee0da84ab9e04ac7610e28580af62c365e898389 \ No newline at end of file diff --git a/examples/offline_data_parallel.py b/examples/offline_data_parallel.py index 0d66f2bc0fa..bcf8ee31006 100644 --- a/examples/offline_data_parallel.py +++ b/examples/offline_data_parallel.py @@ -123,8 +123,13 @@ def main( os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port) - # CUDA_VISIBLE_DEVICES for each DP rank is set automatically inside the - # engine processes. + from vllm_ascend.utils import vllm_version_is + + _dp_device_ids = None + if not vllm_version_is("0.23.0"): + import torch + + _dp_device_ids = [str(i) for i in range(torch.npu.device_count())] # Sample prompts. prompts = [ @@ -165,6 +170,7 @@ def start(rank): enable_expert_parallel=enable_expert_parallel, trust_remote_code=trust_remote_code, quantization=quantization, + **({} if _dp_device_ids is None else {"device_ids": _dp_device_ids}), ) outputs = llm.generate(prompts, sampling_params) # Print the outputs. diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 3105e95abeb..1460ebe022e 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -872,6 +872,24 @@ def _run_vllm_runner_dp_worker(conn, llm_kwargs: dict[str, Any], dp_rank: int, d os.environ["VLLM_DP_MASTER_PORT"] = str(master_port) os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + from vllm_ascend.utils import vllm_version_is + + if not vllm_version_is("0.23.0"): + import torch + + visible = os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "") + full_device_ids: list[str] = [d for d in visible.split(",") if d] + if not full_device_ids: + full_device_ids = [str(i) for i in range(torch.npu.device_count())] + + if llm_kwargs.get("distributed_executor_backend") == "ray": + devs = full_device_ids + chunk = max(len(devs) // dp_size, 1) + start = dp_rank * chunk + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = ",".join(devs[start : start + chunk]) + else: + llm_kwargs["device_ids"] = full_device_ids + llm = LLM(**llm_kwargs) conn.send({"status": "ready", "rank": dp_rank}) diff --git a/tests/ut/core/test_profiling_chunk.py b/tests/ut/core/test_profiling_chunk.py index 3c8a80c56f2..47649f6fedf 100644 --- a/tests/ut/core/test_profiling_chunk.py +++ b/tests/ut/core/test_profiling_chunk.py @@ -90,6 +90,7 @@ def test_invalid_min_chunk_raises(self): ProfilingChunkConfig({"min_chunk": 0}) @patch("vllm.config.VllmConfig.__post_init__", MagicMock()) + @patch("vllm.config.device.DeviceConfig.__post_init__", MagicMock()) @patch("vllm_ascend.platform.NPUPlatform._fix_incompatible_config") def test_enabled_without_pp_raises(self, _mock): clear_ascend_config() @@ -106,6 +107,7 @@ def test_enabled_without_pp_raises(self, _mock): clear_ascend_config() @patch("vllm.config.VllmConfig.__post_init__", MagicMock()) + @patch("vllm.config.device.DeviceConfig.__post_init__", MagicMock()) @patch("vllm_ascend.platform.NPUPlatform._fix_incompatible_config") def test_enabled_with_pp_ok(self, _mock): clear_ascend_config() @@ -121,6 +123,7 @@ def test_enabled_with_pp_ok(self, _mock): clear_ascend_config() @patch("vllm.config.VllmConfig.__post_init__", MagicMock()) + @patch("vllm.config.device.DeviceConfig.__post_init__", MagicMock()) @patch("vllm_ascend.platform.NPUPlatform._fix_incompatible_config") def test_disabled_without_pp_ok(self, _mock): clear_ascend_config() @@ -246,6 +249,7 @@ class TestProfilingChunkScheduler(TestBase): @patch("vllm_ascend.ascend_config.get_ascend_config") @patch("vllm.config.ModelConfig.__post_init__", MagicMock()) @patch("vllm.config.VllmConfig.__post_init__", MagicMock()) + @patch("vllm.config.device.DeviceConfig.__post_init__", MagicMock()) def create_scheduler(self, mock_get_ascend_config): profiling_cfg = MagicMock() profiling_cfg.enabled = True diff --git a/tests/ut/ops/test_gdn_attn_builder.py b/tests/ut/ops/test_gdn_attn_builder.py index 4376a9881de..20f8ae546a0 100644 --- a/tests/ut/ops/test_gdn_attn_builder.py +++ b/tests/ut/ops/test_gdn_attn_builder.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from types import SimpleNamespace from typing import cast +from unittest.mock import patch import pytest import torch @@ -35,6 +36,7 @@ from vllm_ascend.ops.triton.fla.utils import ( prepare_update_chunk_offsets as runtime_prepare_update_chunk_offsets, ) +from vllm_ascend.utils import vllm_version_is @pytest.fixture(autouse=True) @@ -48,6 +50,19 @@ def _patch_triton_cdiv(monkeypatch): ) +@pytest.fixture(autouse=True) +def _no_pin_memory(): + # compute_causal_conv1d_metadata uses np_to_pinned_tensor which reads + # PIN_MEMORY. Without physical NPU, t.pin_memory() raises + # "Please register PrivateUse1HooksInterface first". + with patch("vllm.utils.torch_utils.PIN_MEMORY", False): + if vllm_version_is("0.23.0"): + yield + else: + with patch("vllm.v1.attention.backends.utils.PIN_MEMORY", False): + yield + + @dataclass class BatchSpec: seq_lens: list[int] diff --git a/tests/ut/quantization/methods/test_w4a16_mxfp4.py b/tests/ut/quantization/methods/test_w4a16_mxfp4.py index 268ff8ba9a2..79624322969 100644 --- a/tests/ut/quantization/methods/test_w4a16_mxfp4.py +++ b/tests/ut/quantization/methods/test_w4a16_mxfp4.py @@ -1,4 +1,4 @@ -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch @@ -17,10 +17,12 @@ class TestAscendW4A16MXFP4MoEMethod(TestBase): @patch("vllm_ascend.quantization.methods.w4a16_mxfp4.ensure_mxfp4_moe_available") @patch("vllm_ascend.quantization.methods.w4a16_mxfp4.get_current_vllm_config") @patch("vllm_ascend.quantization.methods.w4a16_mxfp4.get_ascend_config") - def setUp(self, mock_ascend, mock_vllm, mock_ensure): + @patch("vllm_ascend.quantization.methods.w4a16_mxfp4.get_ep_group") + def setUp(self, mock_ep_group, mock_ascend, mock_vllm, mock_ensure): mock_vllm.return_value = create_mock_vllm_config() mock_ascend.return_value = create_mock_ascend_config() mock_ensure.return_value = None + mock_ep_group.return_value = Mock() self.scheme = AscendW4A16MXFP4FusedMoEMethod() @pytest.mark.skip("Execute after the issue is fixed") diff --git a/tests/ut/spec_decode/test_extract_hidden_states_proposer.py b/tests/ut/spec_decode/test_extract_hidden_states_proposer.py index 4a5426c2ee0..6763523984b 100644 --- a/tests/ut/spec_decode/test_extract_hidden_states_proposer.py +++ b/tests/ut/spec_decode/test_extract_hidden_states_proposer.py @@ -32,19 +32,23 @@ from vllm_ascend.spec_decode.extract_hidden_states_proposer import ( AscendExtractHiddenStatesProposer, ) +from vllm_ascend.utils import vllm_version_is @pytest.fixture(autouse=True) def _no_pin_memory(): - # On Ascend/NPU CI runners without physical hardware, torch.zeros(..., - # pin_memory=True) triggers aclInit and fails. Patch - # is_pin_memory_available so vllm's ExtractHiddenStatesProposer.__init__ - # creates CpuGpuBuffer with pin_memory=False. - with patch( - "vllm.v1.spec_decode.extract_hidden_states.is_pin_memory_available", - return_value=False, - ): - yield + if vllm_version_is("0.23.0"): + with patch( + "vllm.v1.spec_decode.extract_hidden_states.is_pin_memory_available", + return_value=False, + ): + yield + else: + with patch( + "vllm.v1.spec_decode.extract_hidden_states.PIN_MEMORY", + False, + ): + yield class MockCachedRequestState: diff --git a/tests/ut/worker/a2/test_worker_v1.py b/tests/ut/worker/a2/test_worker_v1.py index cee8c8cca23..cbc4c616d49 100644 --- a/tests/ut/worker/a2/test_worker_v1.py +++ b/tests/ut/worker/a2/test_worker_v1.py @@ -249,6 +249,7 @@ def test_wake_up_mode_enabled(self, mock_get_config, mock_allocator_class): mock_allocator.wake_up.assert_called_once_with(tags=["test_tag"]) worker.sleep_wakeup_manager.wakeup.assert_called_once_with(["test_tag"]) + @patch("vllm_ascend.worker.worker.current_platform") @patch("vllm_ascend.worker.worker.MemorySnapshot") @patch("vllm_ascend.worker.worker.NPUWorker._init_worker_distributed_environment") @patch("vllm_ascend.worker.worker.init_device_properties_triton") @@ -265,6 +266,7 @@ def test_init_device( mock_init_triton, mock_init_dist_env, mock_snapshot_cls, + mock_current_platform, ): """Test _init_device method""" from vllm_ascend.worker.worker import AscendDeviceType, NPUWorker @@ -279,15 +281,21 @@ def test_init_device( mock_snapshot.total_memory = 2000 mock_snapshot_cls.return_value = mock_snapshot + # Mock current_platform for v0.24.0 init_device path + mock_current_platform.logical_device_id_to_visible_device_id.return_value = 0 + mock_current_platform.device_type = "npu" + # Create worker mock with patch.object(NPUWorker, "__init__", lambda x, **kwargs: None): worker = NPUWorker() - worker.local_rank = 1 + worker.local_rank = 0 worker.model_config = MagicMock() worker.model_config.seed = 42 worker.parallel_config = MagicMock() worker.parallel_config.local_world_size = 0 worker.parallel_config.data_parallel_size = 1 + worker.parallel_config.assigned_physical_gpu_ids = None + worker.parallel_config.distributed_executor_backend = "ray" worker.vllm_config = MagicMock() worker.vllm_config.kv_transfer_config = None worker.cache_config = MagicMock() @@ -297,7 +305,7 @@ def test_init_device( result = worker._init_device() mock_init_dist_env.assert_called_once() - self.assertEqual(str(result), "npu:1") + self.assertEqual(str(result), "npu:0") self.assertEqual(worker.init_snapshot, mock_snapshot) self.assertEqual(worker.requested_memory, 2000 * 0.5) diff --git a/vllm_ascend/_310p/worker_310p.py b/vllm_ascend/_310p/worker_310p.py index 14251371d8d..49dc5b14941 100644 --- a/vllm_ascend/_310p/worker_310p.py +++ b/vllm_ascend/_310p/worker_310p.py @@ -25,7 +25,7 @@ from vllm.utils.torch_utils import set_random_seed # noqa: E402 from vllm_ascend._310p.model_runner_310p import NPUModelRunner310 -from vllm_ascend.utils import is_rc_device +from vllm_ascend.utils import is_rc_device, vllm_version_is from vllm_ascend.worker.worker import NPUWorker, init_workspace_manager @@ -137,7 +137,10 @@ def _init_device(self): torch.npu.empty_cache() # take current memory snapshot - self.init_snapshot = MemorySnapshot() + if vllm_version_is("0.23.0"): + self.init_snapshot = MemorySnapshot() + else: + self.init_snapshot = MemorySnapshot(device=device) self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization if is_rc_device(): self.init_snapshot.free_memory = psutil.virtual_memory().available diff --git a/vllm_ascend/__init__.py b/vllm_ascend/__init__.py index 4c70552fa05..3bd734af95b 100644 --- a/vllm_ascend/__init__.py +++ b/vllm_ascend/__init__.py @@ -15,8 +15,6 @@ # This file is a part of the vllm-ascend project. # -import vllm_ascend.logger # noqa: F401 - _GLOBAL_PATCH_APPLIED = False @@ -75,3 +73,6 @@ def register_model(): from .models import register_model register_model() + + +import vllm_ascend.logger # noqa: E402, F401 diff --git a/vllm_ascend/patch/platform/__init__.py b/vllm_ascend/patch/platform/__init__.py index cc352d6119b..c5822fb882f 100644 --- a/vllm_ascend/patch/platform/__init__.py +++ b/vllm_ascend/patch/platform/__init__.py @@ -51,3 +51,4 @@ if not vllm_version_is("0.23.0"): import vllm_ascend.patch.platform.patch_fused_moe # noqa + import vllm_ascend.patch.platform.patch_dp_device_ids # noqa diff --git a/vllm_ascend/patch/platform/patch_dp_device_ids.py b/vllm_ascend/patch/platform/patch_dp_device_ids.py new file mode 100644 index 00000000000..0820451e65f --- /dev/null +++ b/vllm_ascend/patch/platform/patch_dp_device_ids.py @@ -0,0 +1,72 @@ +# +# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# This file is a part of the vllm-ascend project. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Patch vLLM v0.24.0+ ``get_physical_gpu_ids_for_local_dp_rank`` so that it +# tolerates a pre-sharded ASCEND_RT_VISIBLE_DEVICES env var (one slice per +# DP rank), instead of unconditionally applying ``local_dp_rank * world_size`` +# as an offset into it. +# +# Background: +# PR #45026 removed the per-process device isolation that older vLLM +# versions performed internally. Application-level DP (e.g. +# ``offline_data_parallel.py``) now has to slice ASCEND_RT_VISIBLE_DEVICES +# per rank itself, but the upstream helper still expects the env var to +# contain ALL devices for ALL ranks and tries to read it with the +# ``local_dp_rank * world_size`` offset. With a sharded env var, that +# offset is out of range and the helper raises ``IndexError`` (wrapped in +# the user-facing "Error computing device indices for ..." message). + +from vllm_ascend.utils import vllm_version_is + +if not vllm_version_is("0.23.0"): + import os + + from vllm.platforms import current_platform + from vllm.v1.engine import utils as _engine_utils + + _original_get_physical_gpu_ids = _engine_utils.get_physical_gpu_ids_for_local_dp_rank + + def _patched_get_physical_gpu_ids_for_local_dp_rank( + device_control_env_var, + local_dp_rank, + world_size, + local_world_size=None, + user_assigned_gpu_ids=None, + ): + if local_world_size is None: + local_world_size = world_size + + # If the caller did not pass --device-ids and the env var has + # fewer devices than the full DP range expects, the env var has + # already been pre-sharded per rank by the caller. Use it + # directly from index 0 instead of applying the DP offset again. + if user_assigned_gpu_ids is None and device_control_env_var in os.environ: + visible = [d for d in os.environ[device_control_env_var].split(",") if d] + if local_dp_rank * world_size + local_world_size > len(visible): + return [ + current_platform.device_control_id_to_physical_device_id(visible[device_id]) + for device_id in range(local_world_size) + ] + + return _original_get_physical_gpu_ids( + device_control_env_var, + local_dp_rank, + world_size, + local_world_size, + user_assigned_gpu_ids, + ) + + _engine_utils.get_physical_gpu_ids_for_local_dp_rank = _patched_get_physical_gpu_ids_for_local_dp_rank diff --git a/vllm_ascend/patch/platform/patch_kv_cache_utils.py b/vllm_ascend/patch/platform/patch_kv_cache_utils.py index 510af1525dd..c08daf54238 100644 --- a/vllm_ascend/patch/platform/patch_kv_cache_utils.py +++ b/vllm_ascend/patch/platform/patch_kv_cache_utils.py @@ -17,6 +17,8 @@ UniformTypeKVCacheSpecs, ) +from vllm_ascend.utils import vllm_version_is + _orig_resolve_kv_cache_block_sizes = vllm.v1.core.kv_cache_utils.resolve_kv_cache_block_sizes @@ -249,8 +251,14 @@ def _get_kv_cache_config_deepseek_v4( vllm.v1.core.kv_cache_utils.resolve_kv_cache_block_sizes = _ascend_resolve_kv_cache_block_sizes vllm.v1.core.kv_cache_utils.group_and_unify_kv_cache_specs = group_and_unify_kv_cache_specs -vllm.v1.core.kv_cache_utils._get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_deepseek_v4 vllm.v1.core.kv_cache_utils._get_kv_cache_groups_uniform_groups = _get_kv_cache_groups_uniform_groups +# vllm v0.24.0 renamed _get_kv_cache_config_deepseek_v4 to _get_kv_cache_config_packed and +# get_kv_cache_config_from_groups now calls _get_kv_cache_config_packed directly, bypassing +# the alias patch above. Patch the canonical name so Ascend's non-packed layout is used. +if vllm_version_is("0.23.0"): + vllm.v1.core.kv_cache_utils._get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_deepseek_v4 +else: + vllm.v1.core.kv_cache_utils._get_kv_cache_config_packed = _get_kv_cache_config_deepseek_v4 # Also patch the reference used by engine/core.py which imports the function directly. import vllm.v1.engine.core # noqa: E402 diff --git a/vllm_ascend/patch/worker/patch_mamba_utils.py b/vllm_ascend/patch/worker/patch_mamba_utils.py index ff1066b3a29..54c66b6e330 100644 --- a/vllm_ascend/patch/worker/patch_mamba_utils.py +++ b/vllm_ascend/patch/worker/patch_mamba_utils.py @@ -192,6 +192,24 @@ def _batch_memcpy_unavailable(src_ptrs, dst_ptrs, sizes): mamba_utils.do_mamba_copy_block = _do_mamba_copy_block_torch mamba_utils.postprocess_mamba_align_gpu = _postprocess_mamba_align_gpu_cpu_fallback +# Ascend NPU does not support DT_UINT64 in aclnnInplaceZero. +# MambaCopyBuffers.create() uses torch.uint64 for src_ptrs/dst_ptrs, +# which triggers a runtime error. Remap to int64 at the source. +_original_create = MambaCopyBuffers.create + + +@classmethod +def _patched_create(cls, max_num_reqs, kv_cache_config, copy_funcs, make_buffer): + return _original_create( + max_num_reqs, + kv_cache_config, + copy_funcs, + lambda n, dtype: make_buffer(n, dtype=torch.int64 if dtype == torch.uint64 else dtype), + ) + + +MambaCopyBuffers.create = _patched_create + def preprocess_mamba( scheduler_output: SchedulerOutput, diff --git a/vllm_ascend/simple_kv_offload/copy_backend.py b/vllm_ascend/simple_kv_offload/copy_backend.py index 5572ee8f1cd..36051fbd65d 100644 --- a/vllm_ascend/simple_kv_offload/copy_backend.py +++ b/vllm_ascend/simple_kv_offload/copy_backend.py @@ -71,10 +71,11 @@ def launch_copy( is_store: bool, event_idx: int, events_list: list[tuple[int, torch.npu.Event]], + wait_event: torch.npu.Event | None = None, ) -> None: params = self._store_params if is_store else self._load_params assert params is not None and self._queue is not None - self._queue.put((src_blocks, dst_blocks, params, is_store, event_idx, events_list)) + self._queue.put((src_blocks, dst_blocks, params, is_store, event_idx, events_list, wait_event)) def shutdown(self) -> None: if self._shutdown: @@ -114,10 +115,13 @@ def _copy_loop(self) -> None: is_store, event_idx, events_list, + wait_event, ) = item stream = self._store_stream if is_store else self._load_stream with torch.npu.stream(stream): + if wait_event is not None: + stream.wait_event(wait_event) copy_blocks(src_blocks, dst_blocks, params) event = torch.npu.Event() event.record(stream) diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 1cbfb1c8062..dea225a4fc7 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -52,7 +52,7 @@ from vllm.utils.import_utils import LazyLoader from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import DeviceMemoryProfiler -from vllm.utils.torch_utils import get_dtype_size +from vllm.utils.torch_utils import PIN_MEMORY, get_dtype_size from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -280,6 +280,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) + if not vllm_version_is("0.23.0"): + self.pin_memory = PIN_MEMORY + # Replace the CUDA PrefetchOffloader set by parent __init__ with NPU version. offload_cfg = vllm_config.offload_config if (offload_cfg is not None diff --git a/vllm_ascend/worker/utils.py b/vllm_ascend/worker/utils.py index 570de76bec3..861d8734b9b 100644 --- a/vllm_ascend/worker/utils.py +++ b/vllm_ascend/worker/utils.py @@ -69,7 +69,7 @@ def __init__(self, device: torch.device, pin_memory: bool) -> None: def init_meta( self, attn_groups_iter: Iterable["AttentionGroup"], - kernel_block_sizes: list[int], + kernel_block_sizes: list[list[int]], cache_dtype: str, runner_only_attn_layers: set[str], static_forward_context: dict[str, Any], diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 65b553692a8..720d6a7a107 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -434,6 +434,7 @@ def _reshape_kv_cache_v2( cache_dtype: str, kernel_block_sizes: list[int], shared_kv_cache_layers: dict[str, str], + kv_cache_config: "KVCacheConfig | None" = None, ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: vllm_config = get_current_vllm_config() is_kv_consumer = ( diff --git a/vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py b/vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py index de082a3244f..37f0ab88621 100644 --- a/vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py +++ b/vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py @@ -337,6 +337,10 @@ def rejection_sample( # [num_speculative_steps] synthetic_conditional_rates: torch.Tensor | None = None, use_fp64: bool = False, + # TODO: refactor speculative decoding functionality in a future PR. + # `use_block_verification` is accepted but not yet implemented on NPU; + # wire it up when the block verification path is supported. + use_block_verification: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if use_fp64: raise NotImplementedError("FP64 rejection sampling is not supported on NPU.") diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py index 357e3517d0d..11f68bfcf6a 100644 --- a/vllm_ascend/worker/worker.py +++ b/vllm_ascend/worker/worker.py @@ -41,6 +41,7 @@ from vllm.distributed.parallel_state import Handle, get_pp_group, get_tp_group from vllm.logger import logger from vllm.lora.request import LoRARequest +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.utils.mem_constants import GiB_bytes @@ -395,7 +396,57 @@ def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: self.cache_config.num_cpu_blocks = num_cpu_blocks def _init_device(self): - device = torch.device(f"npu:{self.local_rank}") + if not vllm_version_is("0.23.0"): + # vLLM v0.24.0 (PR #45026) removed automatic per-process device + # isolation for DP workers. Mirror gpu_worker.py::init_device: + # shift self.local_rank by dp_local_rank * tp_pp_world_size so + # that each DP group binds to a distinct set of NPUs. + parallel_config = self.parallel_config + if ( + parallel_config.distributed_executor_backend not in ("ray", "external_launcher") + and parallel_config.data_parallel_backend != "ray" + and parallel_config.nnodes_within_dp == 1 + # vllm-ascend: when the user pre-shards devices via + # --device-ids (which becomes assigned_physical_gpu_ids), + # each child process already binds to its own NPU(s); the + # DP local_rank shift below would push local_rank past the + # length of the per-rank device list and trip the assert + # in this same method. Skip the shift in that case. + and parallel_config.assigned_physical_gpu_ids is None + ): + dp_local_rank = parallel_config.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = parallel_config.data_parallel_index + tp_pp_world_size = parallel_config.pipeline_parallel_size * parallel_config.tensor_parallel_size + self.local_rank += dp_local_rank * tp_pp_world_size + + # Publish the logical-to-physical mapping for topology queries. + assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + assert self.local_rank < len(assigned_physical_gpu_ids), ( + f"local_rank {self.local_rank} is out of bounds for " + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" + ) + if parallel_config.distributed_executor_backend not in ("ray", "external_launcher"): + assert parallel_config.local_world_size <= len(assigned_physical_gpu_ids), ( + f"local_world_size ({parallel_config.local_world_size}) " + f"exceeds assigned_physical_gpu_ids count " + f"({len(assigned_physical_gpu_ids)})" + ) + else: + visible_device_count = torch.npu.device_count() if torch.npu.is_available() else 0 + assert self.local_rank < visible_device_count, ( + f"DP adjusted local rank {self.local_rank} is out of bounds for {visible_device_count} devices." + ) + + visible_device_index = current_platform.logical_device_id_to_visible_device_id(self.local_rank) + device = torch.device(f"{current_platform.device_type}:{visible_device_index}") + else: + device = torch.device(f"npu:{self.local_rank}") + torch.npu.set_device(device) # Import _inductor for graph mode execution with triton @@ -414,7 +465,10 @@ def _init_device(self): setup_ascend_local_comm_res(self.local_rank, self.vllm_config.kv_transfer_config) # take current memory snapshot - self.init_snapshot = MemorySnapshot() + if vllm_version_is("0.23.0"): + self.init_snapshot = MemorySnapshot() + else: + self.init_snapshot = MemorySnapshot(device=device) self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization if self.init_snapshot.free_memory < self.requested_memory: GiB = lambda b: round(b / GiB_bytes, 2)