[Attention] Add direct symmetric-memory DCP A2A by GirasoleY · Pull Request #48897 · vllm-project/vllm

GirasoleY · 2026-07-17T01:04:06Z

Context

Add an opt-in direct symmetric-memory A2A path for MLA DCP. With VLLM_USE_DIRECT_DCP_A2A=1 and
--dcp-comm-backend a2a, each rank dispatches its partial output and LSE straight into peer receive buffers, peer then waits and performs the LSE-weighted combine on the receive buffer.

This removes the stock pack -> NCCL all-to-all -> unpack/combine path from the hot operation.
The workspace is allocated once and shared across MLAlayers, with separate staging state per DBO ubatch.

Duplicate Check

	This PR	#48248
Exchange	One-sided writes directly into peer output/LSE buffers	FlashInfer LL128 all-to-all
Surrounding work	Eliminates stock pack/unpack and performs wait + LSE combine in CUDA	Retains pack/unpack and Triton combine
Runtime dependency	PyTorch symmetric memory in the local NVLink/NVSwitch domain	FlashInfer MNNVL workspace
Selection	Explicit opt-in; eager and graph-replay paths supported	Automatic Blackwell + full-CUDA-graph selection

This PR is an alternative of the full a2a lse reduce path, can be merged independently of #48248.

A2A + LSE-reduce latency reduction

Four GB200 GPUs, Kimi-K2.5 TP4/DCP4 geometry (64 total heads, 16 heads per
rank, head dimension 512, BF16). Each CUDA Graph contains 10 operations;
numbers are the median of seven repeats, with 100 graph replays per repeat,
reported as the maximum rank latency per operation. stock is the complete
current dcp_a2a_lse_reduce operation, including pack, NCCL A2A, and combine;
direct is the complete direct dispatch, signal, wait, and combine operation.

Batch tokens	Stock (us)	Direct (us)	Latency reduction	Speedup
1	18.776	12.504	33.4%	1.50x
8	20.944	12.711	39.3%	1.65x
16	23.944	12.947	45.9%	1.85x
32	26.203	14.369	45.2%	1.82x
64	31.369	18.059	42.4%	1.74x
128	42.327	29.098	31.3%	1.45x

InferenceX AgentX

Kimi-K2.5 DCP4, FlashInfer MLA AgentX

Concurrency	Per-GPU tok/s, stock -> direct	Total gain	Output tok/s gain	Request/s gain	ITL p50	ITL p95
8	5,734 -> 5,778	+0.77%	+0.20%	+0.56%	-4.05%	+0.57%
32	8,160 -> 8,341	+2.22%	+7.20%	+1.42%	-3.99%	+12.57%
64	10,421 -> 11,164	+7.14%	+4.45%	+4.71%	-5.86%	+18.80%

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

mergify · 2026-07-19T03:49:33Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @GirasoleY.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

LucasWilkinson

Overall looks pretty good to me; can this be on by default on CUDA devices? Ideally when we add new kernels we want them on by default in atleast one deployment scenario (can be HW/model/parallelism specific) just so the code is a bit more self documenting on why a kernel was added

LucasWilkinson · 2026-07-20T14:45:59Z

+        if self.dcp_direct_a2a:
+            self.dcp_direct_a2a_workspace = get_direct_dcp_a2a_workspace(
+                get_dcp_group().device_group,
+                kv_b_proj_weight.device,
+                self.dcp_direct_a2a_max_num_tokens,
+                self.num_heads,
+                self.kv_lora_rank,
+                act_dtype,
+                self.dcp_direct_a2a_num_ubatches,
+            )
+            logger.info_once("Using direct symmetric-memory DCP A2A for MLA.")
+


nit: why can't this be done at init time?

addressed in 1691322

LucasWilkinson · 2026-07-20T14:47:14Z

+                if self.dcp_direct_a2a_workspace is not None:
+                    attn_out = self.dcp_direct_a2a_workspace.lse_reduce(
+                        attn_out, lse, self.impl.lse_base_on_e
+                    )


nit: i think id be cleaner/more-robust if it was

if self.dcp_direct_a2a: assert self.dcp_direct_a2a_workspace is not None attn_out = self.dcp_direct_a2a_workspace.lse_reduce( attn_out, lse, self.impl.lse_base_on_e )

addressed in 1691322

LucasWilkinson · 2026-07-20T15:34:32Z

I think it might be worth refactoring this now so all the dispatching happens at init time (just a bit cleaner less booleans to track)

I think MLAAttention can take an explicit device: torch.device; MultiHeadLatentAttentionWrapper can pass self.kv_b_proj.weight.device. That lets the direct workspace be created in __init__, while the activation dtype is already available from the default dtype / model config:

self.dcp_combine: Callable[ [torch.Tensor, torch.Tensor], torch.Tensor ] | None = None if self.dcp_direct_a2a: self. dcp_direct_a2a_workspace = get_direct_dcp_a2a_workspace( group=get_dcp_group().device_group, device=device, max_num_tokens=min( vllm_config.scheduler_config.max_num_batched_tokens, max( vllm_config.scheduler_config.max_num_seqs * (1 + vllm_config.num_speculative_tokens), vllm_config.compilation_config.max_cudagraph_capture_size or 0, ), heads_per_rank=self.num_heads, head_dim=self.kv_lora_rank, dtype=dtype, num_ubatches=vllm_config.parallel_config.num_ubatches, ) self.dcp_combine = functools.partial( self. dcp_direct_a2a_workspace.lse_reduce, is_lse_base_on_e=self.impl.lse_base_on_e, ) elif self.impl.dcp_world_size > 1: dcp_combine = dcp_a2a_lse_reduce if self.dcp_a2a else cp_lse_ag_out_rs self.dcp_combine = functools.partial( dcp_combine, cp_group=get_dcp_group(), is_lse_base_on_e=self.impl.lse_base_on_e, )

The hot path then becomes:

assert lse is not None assert self.dcp_combine is not None attn_out = self.dcp_combine(attn_out, lse)

thoughts?

addressed in 1691322

Add an opt-in MLA path that dispatches partial output and LSE directly into peer symmetric-memory buffers, publishes system-scope completion signals, and combines received states with FP32 LSE weighting. Select the CUDA path by MLA capabilities rather than backend identity or sparsity. Support natural-log and base-2 LSE, FP16 and BF16 outputs, and isolated symmetric staging state per ubatch. Reuse one cached workspace across MLA layers and fail startup when explicitly requested initialization is unavailable. Cover DCP2/DCP4 correctness, both LSE bases and output dtypes, alternating ubatch slots, variable message sizes, workspace reuse, and CUDA Graph replay. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: girasoley <girasoley@inferact.ai>

pavanimajety

I am currently running into an issue that returns handle as None.
Do we need to set symmetric memory backend to NCCL explicitly?
Could you share your vLLM server launch command?

…efault on single-node CUDA Address review feedback on vllm-project#48897: - Move all DCP combine dispatching to MLAAttention.__init__: a single dcp_combine callable replaces the per-forward boolean chain, and the direct-A2A workspace is created at init (MLAAttention now accepts an explicit device, defaulting to kv_b_proj's device). - Resolve the MLA impl's dcp_world_size statically from the parallel config (the DCP group always has exactly decode_context_parallel_size ranks), removing the -1 sentinel and lazy resolution in forward. - VLLM_USE_DIRECT_DCP_A2A becomes tri-state: unset means auto (enabled on CUDA with fp16/bf16 when the DCP group is within one node), 1 forces it on (e.g. multi-node NVLink domains), 0 disables it. - Add CPU-only gating unit tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: girasoley <girasoley@inferact.ai>

Read CUTLASS-style padded token strides in the direct dispatch kernel while preserving packed head and feature vectorization. Cover contiguous and padded layouts across eager and CUDA graph execution. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: girasoley <girasoley@inferact.ai>

@pavanimajety

Without DBO, parallel_config.num_ubatches is 0, so the workspace allocated zero-byte symmetric buffers whose rendezvous returns None and the barrier crashes with a cryptic NoneType error. Clamp to one staging slot at the call site (dbo_current_ubatch_id() is 0 when DBO is off), raise a clear error in the workspace, and assert the rendezvous handle. Reported by @pavanimajety. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: girasoley <girasoley@inferact.ai>

GirasoleY · 2026-07-21T03:51:58Z

Overall looks pretty good to me; can this be on by default on CUDA devices? Ideally when we add new kernels we want them on by default in atleast one deployment scenario (can be HW/model/parallelism specific) just so the code is a bit more self documenting on why a kernel was added

The mode is default on with VLLM_USE_DIRECT_DCP_A2A unset. Set = 0 will explicit disable this.

pavanimajety

Validated the accuracy on my end, lgtm. Thanks for the PR!

Port of upstream vllm-project#48897 (branch proto/upstream-direct-dcp-mla, commits 6569754, 1691322, c15e0f1, 983f562). The CUDA kernel is byte-identical to the upstream head; the Python workspace differs only in hoisting function-local imports to module scope (repo style). - csrc/libtorch_stable/attention/dcp_direct_a2a.cu: single-kernel symmetric-memory output/LSE exchange + LSE-weighted combine with epoch-based signaling; persistent workspace, CUDA-graph safe, one staging slot per ubatch (min 1 -- num_ubatches is 0 without DBO). - vllm/v1/attention/ops/dcp_direct_a2a.py: DirectDCPA2AWorkspace and a cached per-(group, geometry) accessor. - VLLM_USE_DIRECT_DCP_A2A tri-state: unset = auto on single-node CUDA bf16/fp16; 1 = force (multi-node NVLink domains); 0 = off. - Generic MLAAttention: init-time dcp_combine dispatch (direct A2A when available, else packed dcp_a2a_lse_reduce for dcp_comm_backend="a2a", else cp_lse_ag_out_rs); dcp_world_size derived statically from parallel_config. Test: tests/distributed/test_dcp_direct_a2a.py -- fp16/bf16, with and without CUDA graphs, including Kimi-K3 decode geometry (6 heads/rank, D=512, world size 4 matching gathered H=24, non-contiguous 128-head backing storage). 8 passed on 2 GPUs; 12 passed on a 4-GPU GB200 node. No regression in the existing packed-A2A suite (1 pre-existing ordering failure reproduces on the unmodified base). docs/design/attention_backends.md is regenerated by the pre-commit hook; its one-line delta (FLASHINFER_MLA Non-Causal) is pre-existing drift from vllm-project#70, not a change of this commit. Amended: the combine kernel now masks empty-KV-shard rows itself -- the dispatch kernel takes the decode seq_lens (optional Tensor in the op schema), publishes LSE = -inf and skips the payload copy for empty shards, and the combine skips zero-weight sources; lse_reduce grows a seq_lens kwarg and direct_dcp_a2a_active() is exposed; the generic MLAAttention passes decode seq_lens to the direct combine. Folded from dcp-all-bucket 91fc3dba8 (empty-shard combine masking). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: girasoley <girasoley@inferact.ai>

ZJY0516 · 2026-08-07T07:14:15Z

/ci run

github-actions · 2026-08-07T07:14:32Z

✅ Triggered Buildkite CI #82809 for commit 894962bf4cce.

GirasoleY · 2026-08-11T02:40:33Z

Feature is merged in #50484, closing this.

GirasoleY requested review from Harry-Chen, LucasWilkinson, MatthewBonanni and tlrmchlsmth as code owners July 17, 2026 01:04

claude Bot reviewed Jul 17, 2026

View reviewed changes

mergify Bot added ci/build v1 labels Jul 17, 2026

mergify Bot added the needs-rebase label Jul 19, 2026

LucasWilkinson approved these changes Jul 20, 2026

View reviewed changes

pavanimajety reviewed Jul 21, 2026

View reviewed changes

Comment thread vllm/envs.py Resolved

Comment thread vllm/v1/attention/ops/dcp_direct_a2a.py Resolved

GirasoleY force-pushed the proto/upstream-direct-dcp-mla branch from 1145ea4 to 1691322 Compare July 21, 2026 02:09

pavanimajety added the verified Run pre-commit for new contributors without triggering other tests label Jul 21, 2026

mergify Bot removed the needs-rebase label Jul 21, 2026

pavanimajety reviewed Jul 21, 2026

View reviewed changes

Comment thread vllm/v1/attention/ops/dcp_direct_a2a.py Resolved

pavanimajety added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 21, 2026

pavanimajety approved these changes Jul 22, 2026

View reviewed changes

foraxe mentioned this pull request Jul 27, 2026

[DCP][Performance] Add Shared-DCP peer-addressable decode data paths #50009

Open

1 task done

Merge branch 'main' into proto/upstream-direct-dcp-mla

02fd63d

GirasoleY mentioned this pull request Jul 27, 2026

[Kimi-K3] DCP support #50055

Closed

chun-wan mentioned this pull request Jul 28, 2026

[Test][V1] Add sleep/wake correctness regression test for hybrid GDN/… #44972

Merged

GirasoleY mentioned this pull request Jul 30, 2026

[Kimi-K3] DCP support #50484

Merged

This was referenced Aug 4, 2026

[Attention][DCP] Add symm_a2a: single-node symmetric-memory A2A backend for MLA decode sgl-project/sglang#33364

Closed

[DCP][Performance] Speed up direct symmetric-memory Output/LSE GirasoleY/vllm#2

Draft

Merge branch 'main' into proto/upstream-direct-dcp-mla

894962b

ZJY0516 enabled auto-merge (squash) August 7, 2026 07:14

GirasoleY closed this Aug 11, 2026

auto-merge was automatically disabled August 11, 2026 02:40
Pull request was closed

malaiwah mentioned this pull request Aug 29, 2026

[Bug]: Direct DCP A2A crashes on GLM sparse-MLA strided output #54305

Open

Conversation

GirasoleY commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Context

Duplicate Check

A2A + LSE-reduce latency reduction

InferenceX AgentX

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Jul 19, 2026

Uh oh!

LucasWilkinson left a comment

Choose a reason for hiding this comment

Uh oh!

LucasWilkinson Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

GirasoleY Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

LucasWilkinson Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

GirasoleY Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

LucasWilkinson Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

GirasoleY Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

pavanimajety left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

GirasoleY commented Jul 21, 2026

Uh oh!

pavanimajety left a comment

Choose a reason for hiding this comment

Uh oh!

ZJY0516 commented Aug 7, 2026

Uh oh!

github-actions Bot commented Aug 7, 2026

Uh oh!

GirasoleY commented Aug 11, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

GirasoleY commented Jul 17, 2026 •

edited

Loading