[Kimi-K3] DCP support by GirasoleY · Pull Request #50484 · vllm-project/vllm

GirasoleY · 2026-07-30T21:59:28Z

Co-authored-by: @foraxe

Summary

Follow-up to #50000 adding decode context parallelism for Kimi-K3:

direct symmetric-memory DCP A2A output/LSE reduction, including empty-KV-shard masking
DCP support for the fused Kimi-K3 MLA layer
NVLS-multicast direct query gather and multimem chunked-context KV gather
direct publication of query shards into the consumer-final buffer, removing the staging-to-final materialization by @foraxe

This replaces #50055.

Query-buffer lifetime

The direct query gather returns a view of its persistent symmetric buffer. Reuse is ordered by the mandatory post-attention DCP combine:

query gather publishes the current epoch and waits for every rank
attention consumes the gathered query on the same CUDA stream
every combine implementation introduces a cross-rank dependency after attention
the next gather cannot execute until all ranks have passed the preceding attention read

DBO microbatches use independent workspace slots.

Accuracy

.buildkite/kimi-k3/rack5/gsm8k.yaml:

Arm	exact_match	Correct/N	Request errors
TP8	96.21%	1269/1319	0
DCP8	96.97%	1279/1319	0

Performance

120k regression workload: 114k shared prefix + 6k suffix, 400 output tokens, Mooncake disabled.

Arm	Conc	TPOT p50/p90/p99	HBM use mean/peak
TP8	c1	13.806/13.831/13.857ms	4.93%/7.17%
DCP8	c1	10.544/10.574/10.597ms	0.64%/1.10%
TP8	c2	16.192/23.591/34.101ms	9.18%/14.11%
DCP8	c2	11.828/19.437/27.562ms	1.30%/2.37%

HBM KV capacity per GPU:

TP8: 1.933M tokens / 25.12 GiB
DCP8: 19.752M tokens / 34.28 GiB

On 4x GB200 at matched Kimi-K3 geometry, the direct-to-final query path improved query-exchange latency by 10.4%-29.9%. Correctness and changing-input CUDA-graph replay passed.

Scope and overlap

This PR subsumes the direct symmetric-memory MLA A2A work in #48897 and extends it with Kimi-K3 fused MLA DCP, direct query gather, chunked-context KV gather, and the direct-to-final query optimization.

CUDA-graph padding follow-up

Commit 76b2e9d45e makes empty-shard LSE masking tolerate CUDA-graph padding rows.

Validation:

tests/distributed/test_dcp_a2a.py::TestLSEWeightedCombine::test_ag_rs_masks_empty_shard_and_padded_lse: 1 passed on GB300.
DCP2 partial-prefix reuse versus TP2 on 2x GB300, with direct A2A, query gather, and KV gather selected: 1 passed; 256 cached tokens, exact generated-token match, finite logprobs, and drift <= 1e-2.
Pre-commit hooks on both changed files: passed.
An exploratory full-file GB300 run had 25 passes and 7 return-LSE dtype assertion failures (float32 actual versus fp16/bf16 expected); those cases pass no masking metadata and do not execute this change.

AI assistance was used to diagnose, implement, and validate this follow-up fix.

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-31T13:54:40Z

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

zyongye · 2026-08-03T18:04:10Z

+    query_lens = query_start_loc[1:] - query_start_loc[:-1]
+    empty_rows = torch.repeat_interleave(
+        seq_lens == 0,
+        query_lens,
+        output_size=lse.shape[0],
+    )
+    lse.masked_fill_(empty_rows[:, None], float("-inf"))


I assume this will never be called if we use custom cuda kernel?

zyongye · 2026-08-03T18:07:49Z

    ),
+    # Direct DCP ops default on when applicable; set to 1 to enforce or 0 to disable.
+    "VLLM_USE_DIRECT_DCP_A2A": lambda: maybe_convert_bool(
+        os.getenv("VLLM_USE_DIRECT_DCP_A2A")


Can we set default value? Or I think we can make it an engine args.

we already have --dcp-comm-backend cuda-a2a (or a better name?)

Yeah I considered these, but the kernel version has some limitation. The intent here is

None = auto (default on when multicast is available)

1 = force use

0 = disable.

Which I feel it's actually cleaner this way... but happy to change.

prakhar-prakash-juspay · 2026-08-04T11:02:14Z

@GirasoleY do you plan to get this PR merged if it's working fine?

github-actions · 2026-08-04T21:19:24Z

✅ @GirasoleY, CI is now available for this PR. Comment /ci run to run full CI or /ci retry to retry failed jobs.

WoosukKwon · 2026-08-07T23:21:01Z

/ci run

github-actions · 2026-08-07T23:21:17Z

✅ Triggered Buildkite CI #82945 for commit 65328d0ee326.

WoosukKwon · 2026-08-10T16:32:17Z

/ci run

github-actions · 2026-08-10T16:32:34Z

✅ Triggered Buildkite CI #83171 for commit 4b4c1d0214d9.

Resolves the recurring vllm/envs.py structural conflict per docs/superpowers/specs/2026-05-14-envs-merge-conflict-resolution-design.md: main's legacy `if TYPE_CHECKING:` block and `environment_variables` dict are dropped wholesale (superseded by the pydantic BaseSettings tree on this branch), then main's semantic delta is ported field-by-field. 6 main-side commits touched vllm/envs.py since base e644c8c (+55 -0). All 10 new vars have already-merged callers, so every port is mandatory: - vllm-project#51447 VLLM_MAX_STOP_STRINGS (int=4), VLLM_MAX_NUM_BAD_WORDS (int=128), VLLM_MAX_BAD_WORDS_TOTAL_TOKENS (int=1024) -> ServerSettings - vllm-project#49948 VLLM_MAX_AUDIO_DECODE_BYTES (int=268_435_456) -> MediaSettings, carrying compile_factor=False to mirror main's ignore-set addition - vllm-project#50484 VLLM_USE_DIRECT_DCP_A2A / _Q_GATHER / _KV_GATHER (bool|None=None) -> QuantSettings, with one shared `_parse_direct_dcp` before-validator reproducing main's maybe_convert_bool exactly - vllm-project#52079 VLLM_KIMI_K3_GEMM_RS (bool=False) -> QuantSettings - vllm-project#49458 VLLM_USE_HW_AGNOSTIC (bool=False) -> UsageSettings - vllm-project#47808 VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN (int=8192) -> QuantSettings No deletions, modifications, or renames this window. Nothing was dropped silently: all 6 commits' envs.py deltas are covered above. Env var set parity after resolution: 292 branch fields vs 293 main runtime entries, sole difference VLLM_TRITON_ATTN_USE_TD -- the known deprecation shim divergence, re-confirmed untouched by this merge window. Verified: 54 tests pass across tests/test_envs.py, tests/test_envs_pydantic.py and tests/docs/test_env_vars_gen.py; `pre-commit run --files vllm/envs.py` clean; tests/test_request_input_bounds.py passes (22 tests). The audio, DCP and end-to-end hw-agnostic consumer suites need a CUDA box plus soundfile / multiprocess and were not run here. AI assistance was used to enumerate the port list and apply the resolution; see Appendix G of the playbook for the full audit trail. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Vinay Damodaran <vrdn@hey.com>

… refactor (vllm-project#52377) GLM-5.2 + DCP4 + MTP spec decode produced garbled output (repetition, acceptance collapsing to ~100%) after the vllm-project#50484 DCP Manager refactor, which only handled dense MLA. Per vllm-project#46514, per-token seq_lens must be localized AFTER the MTP expansion; the old path under-counted local KV on (W-1)/W ranks, the indexer picked a wrong sparse top-k, and decode degenerated into repetition. Squashed from the production bugfix chain: - 9a63d1e6: restore main MLADCPManager in mla_attention.decode (sparse-MLA DCP startup crash) - d5f5ca13: guard mask_dcp_empty_shards_ against empty local decode shard (num_decodes=0 -> index out of bounds) - 1e008a8a: expose FP8 sparse decode metadata via FlashMLASparseMetadata.decode - 3492ca7a: skip LSE empty-shard mask on sparse DCP decode (superseded by the official vllm-project#52377 fix below) - a05834a6 (vllm-project#52377 upstream): getattr(attn_metadata, "decode", None) guard, _direct_workspace_query_gather fallback for forced-MQA, use_forced_mqa flag (full seq_lens/query_start_loc for sparse) - c912d820: drop VLLM_DEBUG_DCP instrumentation Verified in production: 8+2 -> 10, 3+4 x3 identical, counting 1-174 consecutive, acceptance 0.90/0.78/0.63/0.56/0.47 (healthy reject curve). Co-authored-by: Sisyphus <sisyphus@openci.dev>

… memory Port the direct-final PCP KV idea from vllm-project#49517 onto the merged vllm-project#50484 symmetric-memory substrate. Producers write final Main-KV and Indexer-K rows into every replica through peer pointers and skip the vllm-project#52046 AllGather + insert path. Signed-off-by: 云挚 <ningyunxiao.nyx@antgroup.com>

…hrough PyTorch symmetric memory Port the direct-final PCP KV idea from vllm-project#49517 onto the merged vllm-project#50484 symmetric-memory substrate. Producers write final Main-KV and Indexer-K rows into every replica through peer pointers and skip the vllm-project#52046 AllGather + insert path. Signed-off-by: 云挚 <ningyunxiao.nyx@antgroup.com> (cherry picked from commit 748d7d4)

GirasoleY requested review from AndreasKaratzas, Harry-Chen, LucasWilkinson, MatthewBonanni, WoosukKwon, mgoin, pavanimajety, tlrmchlsmth, yewentao256 and zyongye as code owners July 30, 2026 21:59

claude Bot reviewed Jul 30, 2026

View reviewed changes

mergify Bot added ci/build kimi k3 nvidia v1 labels Jul 30, 2026

github-project-automation Bot added this to NVIDIA Jul 30, 2026

foraxe mentioned this pull request Jul 30, 2026

[Model Support] Kimi K3 Tracking Issue #50001

Open

GirasoleY mentioned this pull request Jul 30, 2026

[Kimi-K3] support DCP partial prefix cache hit #50493

Merged

foraxe mentioned this pull request Jul 31, 2026

[Kimi-K3] Add direct-final NVLS Query publication sgl-project/sglang#33069

Closed

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

yewentao256 mentioned this pull request Jul 31, 2026

[Feature]: Kimi K3 Performance Optimization #50587

Open

46 tasks

MatthewBonanni mentioned this pull request Jul 31, 2026

[Attention][MLA] Per-request scheduling for MLA chunked context #50613

Merged

zyongye reviewed Aug 3, 2026

View reviewed changes

GirasoleY force-pushed the k3-dcp branch 2 times, most recently from a526d4b to 5dfad64 Compare August 4, 2026 21:17

ywang96 added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 4, 2026

Merge branch 'main' into k3-dcp

65328d0

WoosukKwon added ready ONLY add when PR is ready to merge/full CI is needed and removed ready ONLY add when PR is ready to merge/full CI is needed labels Aug 7, 2026

Merge branch 'main' into k3-dcp

4b4c1d0

WoosukKwon merged commit 63ac04a into vllm-project:main Aug 10, 2026
242 of 248 checks passed

github-project-automation Bot moved this from Ready to Done in NVIDIA Aug 10, 2026

wzhao18 mentioned this pull request Aug 10, 2026

[Attention] Fix MLA prefill workspace allocation size #51733

Merged

4 tasks

pisceskkk mentioned this pull request Aug 11, 2026

bugfix(dcp,mla): Fix empty-shard MLA DCP decode merge #38416

Closed

5 tasks

GirasoleY mentioned this pull request Aug 11, 2026

[Attention] Add direct symmetric-memory DCP A2A #48897

Closed

foraxe mentioned this pull request Aug 12, 2026

[feature][Kimi-K3] Support Kimi-K3 DCP decode context parallelism on Ascend vllm-project/vllm-ascend#14126

Draft

GirasoleY mentioned this pull request Aug 14, 2026

[Kimi-K3][DCP] Publish prefill KV directly in MLA layout #52239

Open

This was referenced Aug 14, 2026

[Bugfix][DCP] Handle sparse MLA metadata after DCP Manager refactor #52377

Merged

[Attention][MLA] FlashMLA sparse: DCP on the fp8_ds_mla mixed-batch path + MTP #46514

Merged

foraxe mentioned this pull request Aug 19, 2026

[Kernel][PCP] Publish replicated KV updates through PyTorch symmetric memory #52863

Open

malaiwah mentioned this pull request Aug 29, 2026

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

Open

jin-amd mentioned this pull request Sep 16, 2026

[Bugfix][MLA][DCP] Round chunked-prefill workspace and max_context_chunk to satisfy DCP divisibility asserts #41229

Closed

4 tasks

Navigation Menu

Uh oh!

[Kimi-K3] DCP support - #50484

[Kimi-K3] DCP support#50484
WoosukKwon merged 9 commits into
vllm-project:mainfrom
GirasoleY:k3-dcp

GirasoleY commented Jul 30, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

mergify Bot commented Jul 31, 2026

Uh oh!

zyongye Aug 3, 2026

Uh oh!

zyongye Aug 3, 2026

Uh oh!

pavanimajety Aug 4, 2026

Uh oh!

GirasoleY Aug 4, 2026 •

edited

Loading

Uh oh!

prakhar-prakash-juspay commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

Uh oh!

WoosukKwon commented Aug 7, 2026

Uh oh!

github-actions Bot commented Aug 7, 2026

Uh oh!

WoosukKwon commented Aug 10, 2026

Uh oh!

github-actions Bot commented Aug 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Uh oh!

Conversation

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

Uh oh!

Summary

Query-buffer lifetime

Accuracy

Performance

Scope and overlap

CUDA-graph padding follow-up

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Jul 31, 2026

Uh oh!

zyongye Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

zyongye Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

pavanimajety Aug 4, 2026

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

prakhar-prakash-juspay commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

Uh oh!

WoosukKwon commented Aug 7, 2026

Uh oh!

github-actions Bot commented Aug 7, 2026

Uh oh!

WoosukKwon commented Aug 10, 2026

Uh oh!

github-actions Bot commented Aug 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

GirasoleY commented Jul 30, 2026 •

edited

Loading

GirasoleY Aug 4, 2026 •

edited

Loading