[Attention] DCP sparse MLA output-merge optimizations - #48196
Draft
LucasWilkinson wants to merge 1 commit into
Draft
[Attention] DCP sparse MLA output-merge optimizations#48196LucasWilkinson wants to merge 1 commit into
LucasWilkinson wants to merge 1 commit into
Conversation
Optimize the DCP decode attention-output merge for sparse MLA: - Project the attention output through an all-gathered W_UV before the DCP merge, shrinking the merge payload from kv_lora_rank to v_head_dim per head. Exact: the LSE-weighted merge and the linear W_UV projection commute. Encapsulated in _dcp_merge_and_v_up_proj / _dcp_lse_merge. - Default the DCP A2A merge transport to FlashInfer's fused decode-CP all-to-all (helix) when available (measured 15.5us vs 45-74us for NCCL all_to_all_single at the decode merge payload), falling back to the Triton/NCCL path otherwise. - Propagate per-row valid counts into the merge so empty sparse rows are handled in the pack without extra zero-fill passes. - Auto-size the TRTLLM sparse-MLA workspace for the DCP head-gathered decode shape. AI assistance (Claude) was used to author this change; the submitter has reviewed every line and run the tests below. Tests: tests/distributed/test_dcp_a2a.py (4-rank pack/combine vs reference, incl. valid-count handling and CUDA graph capture); GSM8K-64 on GLM-5.1-NVFP4 TP4/DCP4 (accuracy parity); imports, ruff, manual-stage mypy. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
LucasWilkinson
requested review from
MatthewBonanni and
pavanimajety
as code owners
July 10, 2026 01:35
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Leoyzen
added a commit
to Leoyzen/vllm
that referenced
this pull request
Aug 14, 2026
…ap long lines - mla_attention.py: add get_current_vllm_config_or_none, cp_lse_ag_out_rs, dcp_a2a_lse_reduce imports for _dcp_merge_and_v_up_proj (vllm-project#48196); drop unused finalize_mla_pcp_decode import - test_mtp.py: add missing torch import for prefix-cache test (vllm-project#50897) - cli_args.py: wrap >88-col lines (vllm-project#50668) - flashmla_sparse.py: wrap long comment (vllm-project#49357) - flashinfer_mla_sparse.py: drop unused np_to_pinned_tensor import (vllm-project#48196) - kv_cache_coordinator.py: fix factory-function param indentation (vllm-project#50897) Co-authored-by: Sisyphus <sisyphus@openci.dev>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Optimizes the DCP decode attention-output merge for sparse MLA. This is one of two PRs split out of #47355 (which is now scoped to just the indexer side-stream overlap). The two are independent and target
maindirectly — they touch disjoint regions offorward_impl, so they can merge in either order.W_UVbefore the DCP merge, shrinking the exchanged payload fromkv_lora_ranktov_head_dimper head. Exact — the LSE-weighted merge and the linearW_UVprojection commute. Encapsulated in_dcp_merge_and_v_up_proj/_dcp_lse_merge.all_to_all_singleat the decode merge payload), with a Triton/NCCL fallback.Relationship to #47355
Split from #47355. This PR is the DCP-specific half; #47355 is the indexer side-stream overlap (which applies to all sparse-MLA decode, not just DCP). Independent; either can land first.
Not a duplicate
Extends the existing DCP a2a merge already in
main(dcp_a2a_lse_reduce,cp_lse_ag_out_rs) with the helix transport, project-before-merge, and valid-count handling. No open PR covers this.Testing
tests/distributed/test_dcp_a2a.py— 4-rank pack/combine vs reference incl. valid-count handling and CUDA graph capture.AI assistance
AI assistance (Claude) was used to author this change; the submitter has reviewed every line and run the tests above.