forked from Karylab-cklius/vllm
[Bugfix][DCP] Fix illegal memory access in DCP a2a decode under full CUDA graphs (#45487)
This commit is contained in:
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user