diff --git a/tests/distributed/aux_output_connector/test_store.py b/tests/distributed/aux_output_connector/test_store.py index 24d67c3d60..62d4472811 100644 --- a/tests/distributed/aux_output_connector/test_store.py +++ b/tests/distributed/aux_output_connector/test_store.py @@ -251,9 +251,6 @@ def test_next_step_does_not_consume_pending_output(monkeypatch): event = Mock() monkeypatch.setattr(torch.cuda, "Event", lambda **kwargs: event) monkeypatch.setattr(async_utils, "stream", lambda *args: nullcontext()) - monkeypatch.setattr( - async_utils, "async_copy_to_np", lambda tensor: tensor.numpy().copy() - ) worker = _make_worker(1) worker.begin_step( _metadata(0, [_request_metadata("request", 0, 1, 0, [])], {}).metadata diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index a10f54a97c..dec382725b 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -18,6 +18,7 @@ from vllm.v1.kv_cache_interface import ( MambaSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.worker.gpu.async_utils import async_copy_to_np from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner @@ -357,3 +358,12 @@ def test_get_drafter_hidden_states_tolerates_missing_target_buffer(target_buffer assert torch.equal(out, buffer[:4]) else: assert out is hidden_states + + +def test_async_copy_to_np_does_not_alias_reused_buffer(): + buffer = torch.zeros(4, dtype=torch.int64) + + snapshot = async_copy_to_np(buffer) + buffer.fill_(1) + + assert snapshot.tolist() == [0, 0, 0, 0] diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index 647b0c297c..0d6481585e 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -262,7 +262,7 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput): def async_copy_to_np(x: torch.Tensor) -> np.ndarray: - return x.to("cpu", non_blocking=True).numpy() + return x.to("cpu", non_blocking=True, copy=True).numpy() @contextlib.contextmanager