[BugFix][DiffusionGemma][CPU] Copy async output snapshots instead of aliasing reused buffers (#59829)

Signed-off-by: ojeda-e <eojeda@league.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

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

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]

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