[Perf][PCP] Shard decode requests across PCP ranks

[pisceskkk]

Purpose

PCP currently replicates every decode request on every PCP rank. For PCP-only deployments (DCP == 1), every rank already has a complete KV-cache replica, so this repeats the same decode computation pcp_size times.

This PR:

• assigns decode-only rows round-robin across PCP ranks on every scheduler step;
• keeps the existing DualChunkSwap prefill partitioning unchanged;
• gathers the owner-produced MLA and sparse-indexer KV updates so every PCP rank retains a complete cache replica;
• preserves replicated decode when DCP > 1, because each DCP rank must consume its local KV shard.

Design trade-off. This approach distributes decode requests across the PCP group in a data-parallel fashion and removes redundant decode computation. However, keeping the replicated KV cache consistent across PCP ranks requires additional communication to propagate the owner-produced KV updates. This compute-communication trade-off makes the realized gains smaller and more workload-dependent than the ideal reduction in redundant computation would suggest.

Duplicate-work check

I searched open PCP/decode issues and PRs before submission. This does not duplicate:

• [MRv2] FlashAttention PCP support for GQA on MRv2 #49564, which adds FlashAttention/GQA PCP and PCP+DCP prefill support but does not shard PCP-only decode requests;
• [Kernel][Core][PCP][KV Update] Replace PCP update collectives with direct peer stores #49517, which replaces replicated KV-update collectives with opt-in, model-specific direct peer stores but does not reduce replicated decode computation;
• [Kernel][Core][Shared-PCP][Capacity] Owner-shard KV history for up to 4x capacity #49741, which owner-shards persistent KV history for capacity rather than sharding decode compute.

Related performance tracking: #31473.

AI assistance

OpenAI Codex was used to analyze, implement, test, benchmark, and prepare this draft. The commit includes an AI co-author trailer.

Test Result

DeepSeek-V2-Lite-Chat correctness and functionality

• Topologies: TP2+EP2 baseline versus TP1+PCP2+EP2 with decode sharding, on the same H20 GPUs.
• Dataset: full GSM8K test split, 1,319 questions, 5-shot.
• Fixed sampling: temperature 0, seed 42, max output 4,096, concurrency 256, and ignore_eos=false.

Topology	Correct	Accuracy	Invalid	Request errors
TP2+EP2	854 / 1,319	64.7460%	1	0
PCP2+EP2, decode sharded	853 / 1,319	64.6702%	1	0

Long-context PCP decode performance

Performance was measured on NVIDIA H20. Every measured case used 128k input tokens, 1k output tokens, concurrency 32, 128 requests, temperature 0, seed 0, and max_num_batched_tokens=32768, with identical prompts and request order for the compared states. Every case completed 128/128 requests with 0 failures.

PCP sharding headline

The direct comparison is PCP before sharding versus PCP after sharding.

• 4x H20 / PCP4+EP4: request and output throughput improve at every 90%-99% hit-rate point, by +1.98% to +6.17% req/s. At 95%, sharding raises throughput from 2.8343 to 3.0094 req/s and from 2,902.4 to 3,081.6 output tokens/s, while reducing mean TPOT from 8.84 to 8.39 ms/token.
• 8x H20 / PCP8+EP8: throughput improves at 95% and 97% by +3.73% and +1.33% req/s. At 95%, sharding raises throughput from 0.4833 to 0.5014 req/s and from 494.9 to 513.4 output tokens/s, while reducing mean TPOT from 58.87 to 57.08 ms/token.
• Boundary points are retained rather than hidden: the retested 4x-H20 80% point is -2.16% req/s, with nearly flat TPOT but higher TTFT/E2E than the pre-sharding point; the 8x-H20 99% single measurement is -0.86% req/s, with lower TPOT but higher TTFT.

4x H20 — random two-layer DeepSeek-V3.2 FP8 checkpoint (80/90/95/97/99% hit rates)

This random two-layer, 64-expert checkpoint is a controlled topology/performance probe, not a production-model estimate.

Throughput — PCP sharding before → after

Target hit	Actual hit	TP4 req/s	PCP4 req/s before → after	Req/s change	PCP4 output tok/s before → after	Output-throughput change
80%	79.9805%	1.0060	1.9696 → 1.9270	-2.16%	2,016.8 → 1,973.2	-2.16%
90%	89.9902%	1.7307	2.3989 → 2.4464	+1.98%	2,456.5 → 2,505.2	+1.98%
95%	94.9707%	2.4399	2.8343 → 3.0094	+6.17%	2,902.4 → 3,081.6	+6.17%
97%	96.9727%	2.8966	3.1382 → 3.2537	+3.68%	3,213.5 → 3,331.8	+3.68%
99%	98.9746%	3.3789	3.3710 → 3.4508	+2.37%	3,451.9 → 3,533.7	+2.37%

Latency — PCP sharding before → after

Target hit	Mean TPOT	TPOT change	Mean TTFT	Mean E2E
80%	12.90 → 12.88 ms/token	-0.18%	2,954 → 3,348 ms	16.15 → 16.52 s
90%	10.17 → 9.99 ms/token	-1.68%	2,851 → 2,761 ms	13.25 → 12.99 s
95%	8.84 → 8.39 ms/token	-5.18%	2,140 → 1,924 ms	11.19 → 10.50 s
97%	7.89 → 7.56 ms/token	-4.15%	1,996 → 1,987 ms	10.06 → 9.72 s
99%	7.25 → 6.95 ms/token	-4.09%	1,961 → 2,050 ms	9.38 → 9.16 s

The sharded PCP4 80% row is the retest on physical GPUs 4-7: 128/128 completed, 0 failed, actual hit rate 79.9805%, 1.927 req/s, 1,973.216 output tokens/s, TTFT mean/p99 3,347.783/10,309.820 ms, TPOT 12.875 ms/token, ITL 13.082 ms/token, and E2E mean/p99 16,518.981/22,152.288 ms.

8x H20 — 1/4-depth DeepSeek-V3.2 FP8 checkpoint (95/97/99% hit rates)

Throughput — PCP sharding before → after

Target hit	Actual hit	TP8 req/s	PCP8 req/s before → after	Req/s change	PCP8 output tok/s before → after	Output-throughput change
95%	94.9707%	0.4223	0.4833 → 0.5014	+3.73%	494.9 → 513.4	+3.73%
97%	96.9727%	0.4975	0.5395 → 0.5467	+1.33%	552.4 → 559.8	+1.33%
99%	98.9746%	0.5698	0.5661 → 0.5613	-0.86%	579.7 → 574.8	-0.86%

Latency — PCP sharding before → after

Target hit	Mean TPOT	TPOT change	Mean TTFT	Mean E2E
95%	58.87 → 57.08 ms/token	-3.04%	5,339 → 5,335 ms	65.56 → 63.73 s
97%	54.87 → 54.24 ms/token	-1.15%	3,086 → 2,962 ms	59.22 → 58.45 s
99%	52.76 → 52.48 ms/token	-0.53%	2,309 → 3,092 ms	56.28 → 56.77 s

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2330792aa2

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Gate decode sharding off for direct-write MLA models

When PCP > 1 and DCP == 1, this enables sharding for every MLA model, but the new cache replication only covers the generic MLAAttention/non-skipped indexer insertion paths. DeepSeek-V3.2 instead writes both MLA and indexer KV directly inside fused_norm_rope (vllm/models/deepseek_v32/attention.py:397-414) and then calls the indexer with skip_k_cache_insert=True; DeepSeek-V4 similarly inserts into its SWA and compressor caches directly. On non-owner ranks these paths therefore never receive the owner's decode KV, and because ownership can change after batch compaction, a later decode can read an incomplete cache and produce incorrect output. Either replicate writes in these direct-insertion paths or disable request sharding for those model implementations.

Useful? React with 👍 / 👎.

[github-actions]

✅ @pisceskkk, CI is now available for this PR.

• /ci run starts a CI build.
• /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.
• /ci cancel cancels scheduled or running CI builds for this PR branch.

[mergify]

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

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

[mergify]

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

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

[GirasoleY]

Thanks for upstream this. The change itself lgtm.

Just to understand the data. It looks like in your experiment this gives better TPOT than TP in a aggregated setup? I wonder if you could provide more details? Are these done in a larger batch size setup, and is it because the decode forward is actually faster, or is it because the mixed batch now runs faster?

I had a similar GLM5.2 experiment previously and the data shows in a pure decode setup, and it only provide minor throughput gain at high batch size (upto 6%), so just want to check if this aligns with what you see as well.

[pisceskkk]

Just to understand the data. It looks like in your experiment this gives better TPOT than TP in a aggregated setup? I wonder if you could provide more details? Are these done in a larger batch size setup, and is it because the decode forward is actually faster, or is it because the mixed batch now runs faster?
I had a similar GLM5.2 experiment previously and the data shows in a pure decode setup, and it only provide minor throughput gain at high batch size (upto 6%), so just want to check if this aligns with what you see as well.

Yes, these results were measured with large batches and a continuous stream of incoming requests. The TPOT improvement mainly comes from faster mixed-batch computation rather than pure decode. This is also the main reason why PCP can outperform TP in mixed deployment.

Even sharding decode tokens across PCP ranks may not necessarily outperform replicated-decode PCP in a pure-decode setup, since it introduces an additional all-gather and therefore comes with its own trade-off. :(