A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 02:15 UTC
[Bugfix] Fix Ray compiled-DAG SHM channel stalls by detaching zero-copy `np.ndarray` logprobs buffers (#35736)
Signed-off-by: Nikita Shapovalov <nikita@poolside.ai>
Nikita Shapovalov committed
Apr 16, 2026 at 17:49 UTC
692db29cd45fcd2d2e0c3b7259b608ec1f5855ef
3 files changed
+109
-2
tests/v1/executor/test_ray_utils.py
new
+54
@@ -0,0 +1,54 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import numpy as np
5
+
6
+from vllm.v1.executor.ray_utils import detach_zero_copy_from_model_runner_output
7
+from vllm.v1.outputs import LogprobsLists, LogprobsTensors, ModelRunnerOutput
8
+
9
+
10
+def _make_readonly(arr: np.ndarray) -> np.ndarray:
11
+ arr.setflags(write=False)
12
+ return arr
13
+
14
+
15
+def test_detach_zero_copy_from_model_runner_output_copies_only_numpy_views():
16
+ cu_num_generated_tokens = [0, 2]
17
+ prompt_logprobs = LogprobsTensors.empty_cpu(1, 2)
18
+ output = ModelRunnerOutput(
19
+ req_ids=["req-0"],
20
+ req_id_to_index={"req-0": 0},
21
+ logprobs=LogprobsLists(
22
+ logprob_token_ids=_make_readonly(
23
+ np.array([[1, 2], [3, 4]], dtype=np.int32)
24
+ ),
25
+ logprobs=_make_readonly(
26
+ np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
27
+ ),
28
+ sampled_token_ranks=_make_readonly(np.array([1, 2], dtype=np.int32)),
29
+ cu_num_generated_tokens=cu_num_generated_tokens,
30
+ ),
31
+ prompt_logprobs_dict={"req-0": prompt_logprobs},
32
+ )
33
+
34
+ original_logprobs = output.logprobs
35
+ assert original_logprobs is not None
36
+
37
+ detach_zero_copy_from_model_runner_output(output)
38
+
39
+ detached_logprobs = output.logprobs
40
+ assert detached_logprobs is not None
41
+ assert detached_logprobs is not original_logprobs
42
+ assert (
43
+ detached_logprobs.logprob_token_ids is not original_logprobs.logprob_token_ids
44
+ )
45
+ assert detached_logprobs.logprobs is not original_logprobs.logprobs
46
+ assert (
47
+ detached_logprobs.sampled_token_ranks
48
+ is not original_logprobs.sampled_token_ranks
49
+ )
50
+ assert detached_logprobs.logprob_token_ids.flags.writeable
51
+ assert detached_logprobs.logprobs.flags.writeable
52
+ assert detached_logprobs.sampled_token_ranks.flags.writeable
53
+ assert detached_logprobs.cu_num_generated_tokens is cu_num_generated_tokens
54
+ assert output.prompt_logprobs_dict["req-0"] is prompt_logprobs
vllm/v1/executor/ray_executor.py
+8
-2
@@ -26,6 +26,7 @@ from vllm.v1.executor.ray_utils import (
26
WORKER_SPECIFIC_ENV_VARS,
27
FutureWrapper,
28
RayWorkerWrapper,
29
+ detach_zero_copy_from_model_runner_output,
30
initialize_ray_cluster,
31
ray,
32
)
@@ -463,7 +464,9 @@ class RayDistributedExecutor(Executor):
464
# Get output only from a single worker (output_rank)
465
# When PP is not used, we block here until the result is available.
466
if not non_block:
466
- return refs[0].get()
467
+ output = refs[0].get()
468
+ detach_zero_copy_from_model_runner_output(output)
469
+ return output
470
471
# When PP is used, we return a FutureWrapper immediately so that
472
# the scheduler can yield to the next batch.
@@ -473,7 +476,10 @@ class RayDistributedExecutor(Executor):
476
assert self.kv_output_aggregator is not None
477
if not non_block:
478
# Block and get results from all workers
476
- return self.kv_output_aggregator.aggregate(ray.get(refs))
479
+ outputs = ray.get(refs)
480
+ for output in outputs:
481
+ detach_zero_copy_from_model_runner_output(output)
482
+ return self.kv_output_aggregator.aggregate(outputs)
483
484
# Return a future that will aggregate outputs from all workers
485
return FutureWrapper(refs, self.kv_output_aggregator)
vllm/v1/executor/ray_utils.py
+47
@@ -7,6 +7,8 @@ from collections import defaultdict
7
from concurrent.futures import Future
8
from typing import TYPE_CHECKING, Union
9
10
+import numpy as np
11
+
12
import vllm.platforms
13
from vllm.config import ParallelConfig
14
from vllm.distributed import get_pp_group
@@ -190,6 +192,48 @@ except ImportError as e:
192
RayWorkerWrapper = None # type: ignore
193
194
195
+def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> None:
196
+ """Detach Ray SHM-channel zero-copy buffers from a ModelRunnerOutput in-place.
197
+
198
+ Ray compiled DAG SHM channels may return zero-copy objects (e.g. `np.ndarray`)
199
+ backed by Ray's shared-memory object store. Ray's channel docs explicitly
200
+ warn that subsequent reads may block if such an object is still in scope.
201
+
202
+ vLLM can return numpy-backed logprobs in `ModelRunnerOutput.logprobs`. If
203
+ those arrays are backed by Ray SHM (commonly read-only), retaining them in
204
+ scope across scheduler iterations can stall the channel and eventually hit
205
+ `RAY_CGRAPH_get_timeout`.
206
+
207
+ Copy read-only numpy arrays so the returned output no longer retains
208
+ references to Ray's shared-memory buffers.
209
+
210
+ We intentionally do not touch `prompt_logprobs_dict`: those entries are
211
+ `LogprobsTensors` backed by PyTorch-owned CPU tensors (`to_cpu_nonblocking`
212
+ or `empty_cpu`), not NumPy views decoded from Ray channels.
213
+ """
214
+ if output.logprobs is None:
215
+ return
216
+
217
+ token_ids, logprobs, ranks, cu_num_generated_tokens = output.logprobs
218
+
219
+ def _copy_if_readonly(arr):
220
+ if isinstance(arr, np.ndarray) and not arr.flags.writeable:
221
+ return arr.copy()
222
+ return arr
223
+
224
+ # `cu_num_generated_tokens` is already a plain Python list (or None), so it
225
+ # never aliases Ray SHM buffers and can be reused as-is.
226
+ token_ids_c = _copy_if_readonly(token_ids)
227
+ logprobs_c = _copy_if_readonly(logprobs)
228
+ ranks_c = _copy_if_readonly(ranks)
229
+ if token_ids_c is token_ids and logprobs_c is logprobs and ranks_c is ranks:
230
+ return
231
+
232
+ output.logprobs = type(output.logprobs)(
233
+ token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens
234
+ )
235
+
236
+
237
class FutureWrapper(Future):
238
"""A wrapper around Ray output reference to meet the interface
239
of .execute_model(): The top level (core busy loop) expects .result() api
@@ -207,8 +251,11 @@ class FutureWrapper(Future):
251
def result(self, timeout=None):
252
outputs = ray.get(self.ref_or_refs, timeout=timeout)
253
if self.aggregator is None:
254
+ detach_zero_copy_from_model_runner_output(outputs)
255
return outputs
256
257
+ for output in outputs:
258
+ detach_zero_copy_from_model_runner_output(output)
259
return self.aggregator.aggregate(outputs, output_rank=0)
260
261