Skip to content

[correctness] DCP A2A capture channel adoption: graphs sharing one channel bake host-side staging-slot parity; stale capture stream-key mappings can hand a graph's channel to later captures #95

Description

@malaiwah

Found during a performance/correctness review of the v20 r12 release image voipmonitor/vllm:gilded-gnosis-v20-vllmb46c3aa-si35aebc6-fi801d57a-cu132-20260730-r12. Code refs are against the composed r12 tree; all cited sparkinfer files are byte-identical to public commit de04125508c75e748016493f851e6bd4513b515b, so permalinks use that SHA. vLLM caller refs are from the composed r12 vLLM tree (base b46c3aa). Deployment context: GLM-5.2 753B MoE, 4x RTX 6000 Pro (sm120), TP4, DCP2/DCP4, FULL_AND_PIECEWISE CUDA graphs.

This is a proposal-tone hardening report: we could not construct a reachable corruption in the shipped GLM r12 configuration (reasons below), but the safety of the A2A channel under CUDA-graph replay currently rests on two invariants that nothing in the code enforces or asserts. Related memory-side report: #94 (section 1 covers the per-capture-pass channel slabs).

1. All graphs captured in one manager pass share one channel — with host-baked staging-slot parity

Where:

        if (
            not self.single_channel
            and _is_current_stream_capturing(self.device)
            and self._capture_channel_stack
        ):
            # Independent graph managers may reuse a torch-owned nested stream
            # key. The enclosing capture, not a stale key mapping, determines
            # which IPC channel is safe for this graph.
            channel = self._capture_channel_stack[-1]
            self._channels[key] = channel
            return channel
    // barrier); no host staging memcpys are issued.
    const int slot = slot_++ % 2;

Mechanism: vLLM enters pool.capture() once per cudagraph-manager pass (vllm/v1/worker/gpu_model_runner.py:6857, vllm/v1/worker/gpu/cudagraph_utils.py:464 wrap the entire PIECEWISE+FULL size loop in one graph_capture(...), which enters capture_b12x_dcp_a2a once). So all graphs of a pass (sizes 4..64, FULL and PIECEWISE) deliberately share one channel via the capture-stack adoption above — that is the by-design behavior already quantified in #94.

The subtle part is the staging-slot double buffer. The device barrier itself is replay-safe: self_counter[...] += 1 executes at kernel run time and the flag phase is value % 2 (pcie_dcp_a2a.cu:68-115). But the staging slot is selected on the host (slot_++ % 2) and baked into the captured kernel node's arguments. The protocol is single-barrier by design (.cu:3-8): after the barrier, peers read my staging slot with no tail barrier before my next op's copy phase. Slot alternation is therefore the only thing that prevents op N+1's staging copy on a fast rank from overwriting the slot a slower peer is still reading from op N.

Under graph replay, alternation across the executed op sequence holds only if every graph baked an even number of channel ops. In the shipped GLM config it does, by accident: each MLA layer under DCP decode issues an all_gather_heads + lse_reduce_scatter pair (vllm/model_executor/layers/attention/mla_attention.py:1429, 1545), 2 x 78 = 156 ops per graph, so every graph starts at slot parity 0 and any replay order alternates correctly. Nothing enforces this. A graph that bakes an odd op count (a config where only one of the pair is A2A-eligible over an odd layer count, a partial-layer/breakable graph, a future model) makes two consecutive replays issue back-to-back same-slot ops, re-creating exactly the overwrite window the double buffer exists to close. On-rank replays are stream-serialized, but the hazard window is cross-rank (peer's read phase vs my next copy phase), which stream order does not cover.

2. for_stream during an unwrapped capture can silently adopt another graph's channel

Where: pcie_dcp_a2a.py#L687-L708

        channel = self._channels.get(key)
        if channel is not None:
            return channel
        if _is_current_stream_capturing(self.device):
            ...
            if self._channels:
                channel = next(iter(self._channels.values()))
                self._channels[key] = channel
                return channel

Mechanism: two latent paths, both only reachable when a CUDA graph capture happens without entering pool.capture() (empty _capture_channel_stack):

  1. The self._channels.get(key) probe runs before the _is_current_stream_capturing check. Mappings written during earlier captures (lines 685, 696, 700) are never invalidated when the capture context exits, and torch recycles capture-stream handles (the code's own comments at L681-683 and L779-780 acknowledge this). A later unwrapped capture whose stream key collides therefore silently bakes a previous graph's channel.
  2. Failing that, next(iter(self._channels.values())) adopts an arbitrary channel. Since _channels is a per-rank local dict keyed by process-local stream pointers, two ranks can in principle adopt different channels for the same logical graph — channels are created collectively in order, so a rank-divergent adoption pairs mismatched IPC slabs, which is a deadlock/corruption, not just a race.

Reachability (honest): in r12 vLLM every A2A-bearing capture is wrapped by graph_capture() -> capture_b12x_dcp_a2a, and the memory-profiling probe uses checkpoint_channels/rollback_channels. We found no unwrapped capture path in the shipped image, so both paths are latent — but they are the API's default behavior for any embedder that captures graphs without the vLLM wrapper.

Suggested fixes

  • Derive the staging slot from device-side state instead of the host counter — e.g. from the parity of the existing self_counter (it already increments once per executed op), which makes replay sequences self-consistent regardless of baked arguments and removes the even-op-count invariant entirely. This looks like the smallest complete fix.
  • Alternatively: assert at capture exit that the channel's ops-baked count is even (cheap: snapshot slot_ in capture() enter/exit), and document the invariant.
  • Invalidate _channels entries created during a capture when the capture context pops, so recycled stream keys cannot resurrect a graph-owned channel (for_stream should reach the "no channel during CUDA graph capture" RuntimeError instead of the stale mapping / arbitrary adoption).

Verification hint

Unit-level, no vLLM needed: build a pool with 2 ranks, capture two graphs each containing an odd number of lse_reduce_scatter ops on the shared capture channel, then replay A, A while delaying one rank's read phase (e.g. stretch the reduction with a spin kernel between barrier and read). The second replay's copy lands in the same slot the delayed peer is still reading; comparing against the eager reference exposes the overwrite. With the device-side slot derivation the same test passes.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions