[john-evan08]

Overview

Fix the partial LSE must have packed heads error when using VLLM_USE_DIRECT_DCP_A2A=1 with transposed FlashAttention MLA LSE tensors, as encountered with DeepSeek-V2-Lite. The direct P2P kernel now reads both tensor strides, without a contiguous copy.

Unlike the NCCL fallback proposed in #54472, this enables the direct path for these LSE layouts.

Claims

• Enable direct P2P for strided LSE without changing reduction arithmetic or backend defaults.
• Improve DeepSeek-V2-Lite throughput at concurrency 32/128 on the measured long-context workload; c512 is inconclusive.

Validation

Serving performance

DeepSeek-V2-Lite, 8 x H100 80 GB, TP=DCP=8, expert parallelism enabled, Q/KV gather through NCCL. Both arms use the same image and Mooncake agentic workload; only VLLM_USE_DIRECT_DCP_A2A changes between 0 and 1.

Metric: generated output tokens/second/GPU; higher is better. Three nodes, two order-balanced repetitions per node: six paired measurements per concurrency.

Concurrency	A2A NCCL	A2A direct P2P	Paired throughput gain	95% CI
32	343.97	359.51	+4.52%	+2.93% to +6.13%
128	617.97	638.68	+3.35%	+2.29% to +4.41%
512	860.21	868.15	+0.92%	-1.47% to +3.36%

At c512, the interval includes small slowdowns and improvements: zero regression is not established. These points vary concurrency, not context length.

Performance was measured before rebasing, with the same kernel fix; the newer build was revalidated for correctness, not rebenchmarked.

Benchmark methodology and versions

Maximum context 163,840; 60 GiB KV cache/GPU; FlashInfer autotune off; AIperf 0.12.0; fresh server/cold prefix cache; no timed profiling. Request counts: 128/256/1,024; all 16,896 requests succeeded.

Trace: 64k_400_90kv_agent_new_noschedule_short_15perc.jsonl, filtered to input plus output lengths <=163,840. Model revision: 604d5664dddd88a0433dbae533b7fe9472482de0.

Throughput: arithmetic means. Gains: paired log ratios averaged per node, with a t(2) interval over three node means. Output-token totals differ by <=0.147% within pairs.

Benchmark upstream base: f3b77eff640b5ed5cc9221a4f5f402ab65adc021; patched source: 132c7f5fc41094d91aecb72cd03389326ff80431.

Correctness

• Four GPU tests passed, zero skips: packed/padded/transposed LSE, singleton tokens, empty shards, and graph replay; packed/strided kernel outputs are bitwise-identical for identical inputs.
• Latest DeepSeek rerun passed eager TP=2 and graph TP=2/8 numerical/path checks. Traces confirm direct kernels. Packed-copy/strided-P2P model controls produced identical saved outputs.
• Additional models are being tested; current model-level evidence covers DeepSeek-V2-Lite only.

Test command and numerical limitations

On at least four peer-access-capable H100 GPUs, with the patched build:

.venv/bin/python -m pytest \
  tests/distributed/test_dcp_direct_a2a_lse_reduce.py \
  --noconftest -k direct_a2a_matches_reference -v -s

--noconftest avoids unrelated dependencies. CI is configured for the two-rank cases on two H100s; four-rank cases passed locally. Upstream CI has not run for this branch. Lint, syntax, and git diff --check passed.

Six prompts, 32 tokens each; comparable logprob bounds: 0.5 maximum / 0.1 per-prompt mean. First divergences require mutual top-five membership and greedy gap <=0.25; later, different contexts are not compared. This is not broad accuracy evaluation or bitwise NCCL/P2P equivalence.

One default-scheduling comparison exceeded these bounds; unchanged reruns and fixed-scheduling controls passed. NCCL itself varied across runs, and its rejected output also exceeds the gap bound against packed-copy P2P. The precise cause remains unisolated; tolerances were not relaxed.

Details

FlashAttention MLA's [heads, tokens] LSE becomes a strided [tokens, heads] view after transpose. Reading both strides fixes addressing without adding copies, allocations, or synchronization. The receive buffer and combine kernel are unchanged.

Output-layout/capacity restrictions and oversized-prefill fallback remain. This does not fix the separate GLM attention-output layout issue #54305. If #54472 lands first, relax only its LSE guard for this newly supported layout.

Related test-fixture/CI work: #55492; consolidate overlapping support changes if it lands first.

AI assistance was used for implementation, tests, validation.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[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.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[john-evan08]

@zyongye @WoosukKwon @GirasoleY @ywang96. could you take a look, given your review of #54472 and #50484.

[lappazos]

@ZJY0516 @NickLucche please review and help us merge, thx!

[lappazos]

@zyongye @WoosukKwon @GirasoleY @ywang96 can u please review and merge?

[mergify]

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

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

[lappazos]

@hmellor @Harry-Chen @khluu can you please review and merge?

[wzhao18]

Looks good to me. Thanks!

[wzhao18]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95980 for commit b7d1b1bb69aa.

[github-actions]

✅ @john-evan08, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[wzhao18]

/ci retry

[github-actions]

✅ Queued 2 failed job(s) for retry in Buildkite CI #95980.