[Bugfix][DCP] Fix illegal memory access in DCP a2a decode under full CUDA graphs (#45487)

This commit is contained in:
Jeff (Junze) Ma
2026-06-14 00:14:31 -07:00
committed by GitHub
parent 4ef4492e9b
commit 9fd737badc
+10 -11
View File
@@ -26,10 +26,6 @@ import torch
import torch.distributed as dist
from vllm.triton_utils import tl, triton
from vllm.v1.worker.workspace import (
current_workspace_manager,
is_workspace_manager_initialized,
)
if TYPE_CHECKING:
from vllm.distributed.parallel_state import GroupCoordinator
@@ -117,13 +113,16 @@ def _dcp_a2a_send_recv_buffers(
device: torch.device,
dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
if is_workspace_manager_initialized():
send_buffer, recv_buffer = current_workspace_manager().get_simultaneous(
(shape, dtype),
(shape, dtype),
)
return send_buffer, recv_buffer
# Don't use the shared WorkspaceManager here. A FULL cudagraph bakes in the
# buffer address at capture, but the workspace is growable and sized only to
# the largest *captured* batch (the cudagraph capture cap). Any eager a2a
# with a bigger batch regrows it, freeing that address and poisoning every
# captured graph -> illegal memory access on replay. This bites the very
# first request: the post-capture warmup runs an eager decode at
# max_num_seqs (> the cap), so the graphs are already dangling before the
# server is ready. torch.empty buffers instead live in the graph's private
# pool and stay valid for its lifetime (as _dcp_a2a_unpack_combine and the
# AG+RS combine path already rely on).
return (
torch.empty(shape, device=device, dtype=dtype),
torch.empty(shape, device=device, dtype=dtype),