A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][DCP] Fix illegal memory access in DCP a2a decode under full CUDA graphs (#45487)
Jeff (Junze) Ma committed
Jun 14, 2026 at 00:14 UTC
9fd737badcc5eaeb61fd1e7b894af02ba657f203
1 file changed
+10
-11
vllm/v1/attention/ops/dcp_alltoall.py
+10
-11
@@ -26,10 +26,6 @@ import torch
26
import torch.distributed as dist
27
28
from vllm.triton_utils import tl, triton
29
-from vllm.v1.worker.workspace import (
30
- current_workspace_manager,
31
- is_workspace_manager_initialized,
32
-)
29
30
if TYPE_CHECKING:
31
from vllm.distributed.parallel_state import GroupCoordinator
@@ -117,13 +113,16 @@ def _dcp_a2a_send_recv_buffers(
113
device: torch.device,
114
dtype: torch.dtype,
115
) -> tuple[torch.Tensor, torch.Tensor]:
120
- if is_workspace_manager_initialized():
121
- send_buffer, recv_buffer = current_workspace_manager().get_simultaneous(
122
- (shape, dtype),
123
- (shape, dtype),
124
- )
125
- return send_buffer, recv_buffer
126
-
116
+ # Don't use the shared WorkspaceManager here. A FULL cudagraph bakes in the
117
+ # buffer address at capture, but the workspace is growable and sized only to
118
+ # the largest *captured* batch (the cudagraph capture cap). Any eager a2a
119
+ # with a bigger batch regrows it, freeing that address and poisoning every
120
+ # captured graph -> illegal memory access on replay. This bites the very
121
+ # first request: the post-capture warmup runs an eager decode at
122
+ # max_num_seqs (> the cap), so the graphs are already dangling before the
123
+ # server is ready. torch.empty buffers instead live in the graph's private
124
+ # pool and stay valid for its lifetime (as _dcp_a2a_unpack_combine and the
125
+ # AG+RS combine path already rely on).
126
return (
127
torch.empty(shape, device=device, dtype=dtype),
128
torch.empty(shape, device=device, dtype=dtype),