feat: support decode context parallelism for PyTorch MLA

[CUHKSZzxy]

Summary

Add opt-in --dcp for PyTorch MLA: shard KV tokens across existing attention-TP ranks to increase cache capacity without increasing model world size.

• Decode: gather query heads, attend to local KV, and merge outputs using LSE weights; DSA merges local candidates into global top-k.
• Prefill: gather cached-prefix KV in bounded chunks.
• Support dense/sparse FlashMLA, MTP, CUDA graphs, prefix caching, BF16 KV, and sparse-MLA FP8 KV.

Scope: CUDA FlashMLA, BF16 activations, DP1/EP1, sparse top-k 512/2048. FP8 model weights are supported. dcp=1 is the default; TurboMind --cp is unchanged.

Accuracy

Full GLM-5.2-FP8, TP8+DCP4, MTP5, FP8 KV, CUDA graphs and prefix caching enabled; client concurrency 256:

• GSM8K: 1269/1319 (96.21%), zero request errors; no prefix-cache hits.
• Separate prefix-reuse probe: 8/8, with 28,672 cached tokens.

Performance

Full GLM-5.2-FP8 on 8×H200, FP8 KV, MTP5, prefix caching and parsers enabled. SWE-Smith: 16 conversations × 5 turns, concurrency 16, first-turn target 16,384 tokens, subsequent increments 2,048, output 1,024 tokens/request. Each run completed 80/80 requests without errors.

LMDeploy throughput: 15865bcf, existing symmetric-memory all-reduce enabled, experimental DCP communicator excluded. vLLM: 0.28.1rc1.dev649+g93911fcf6.

Configuration	Duration (s)	Output throughput (tok/s)	Usable KV token slots
LMDeploy TP8	93.81	873.27 (+4.18%)	796,160
LMDeploy TP8+DCP4	98.49	831.76 (+32.86%)	3,089,408
vLLM TP8¹	97.73	838.26	639,744
vLLM TP8+DCP4	130.85	626.05	1,920,512

Throughput percentages compare matching TP/DCP configurations. Capacity means total resident KV token slots, not per-request context length; do not multiply by TP. LMDeploy excludes its reserved block.

¹ vLLM TP8 used a 65,536-token context limit after 1,048,576 failed startup KV-capacity validation; other runs used 1,048,576.

Measurement notes and validation limits

• Accuracy results are single development runs, not rerun after the latest refactors. The prefix-reuse probe does not guarantee exact cold/warm response-text parity.
• Focused tests cover cache layouts, causal lengths, distributed top-k and CUDA-graph replay; the full dependency matrix remains unvalidated.
• Performance: one fresh-server trial per configuration, no client warmup or profiling; an unrelated connection probe precedes timing. Both vLLM trials included lazy MoE JIT during timing. These are not warmed steady-state or saturation measurements.
• Later conversation histories differ with model outputs. Runtime versions differ: LMDeploy Torch 2.10.0+cu128 / Triton 3.6.0; vLLM Torch 2.13.0+cu130 / Triton 3.7.1. Equal memory fractions have different reservation semantics.
• vLLM used explicit ag_rs because default a2a failed with this build's sparse FlashMLA backend.

Server and client commands

Run each configuration separately with a fresh server. MODEL points to the same full GLM-5.2-FP8 checkpoint.

LMDeploy at 15865bcf: run once with DCP=1, once with DCP=4. The distributed executor is auto-selected.

MODEL=/path/to/GLM-5.2-FP8
DCP=4  # 1 for TP8; 4 for TP8+DCP4

LMDEPLOY_ENABLE_SYMM_MEM_ALLREDUCE=1 \
LMDEPLOY_ENABLE_SYMM_MEM_LMHEAD=0 \
LMDEPLOY_ENABLE_FLASHINFER_ALLREDUCE=0 \
lmdeploy serve api_server "$MODEL" \
  --backend pytorch --tp 8 --dcp "$DCP" \
  --server-name 127.0.0.1 --server-port 23333 \
  --model-name GLM-5.2-FP8 \
  --max-batch-size 64 --session-len 1048576 \
  --max-prefill-token-num 8192 \
  --quant-policy fp8 --cache-max-entry-count 0.92 \
  --enable-prefix-caching \
  --tool-call-parser glm47 --reasoning-parser default \
  --speculative-algorithm deepseek_mtp \
  --speculative-num-draft-tokens 5 --trust-remote-code

vLLM: run with DCP=4, MAX_MODEL_LEN=1048576, then DCP=1, MAX_MODEL_LEN=65536. Block size, KV interleave, attention backend, prefix caching and chunked prefill are not explicitly overridden; this build selected block size 64, interleave 1, sparse FlashMLA, and enabled both caching and chunked prefill.

MODEL=/path/to/GLM-5.2-FP8
DCP=4
MAX_MODEL_LEN=1048576  # Use 65536 for DCP=1.

vllm serve "$MODEL" \
  --host 127.0.0.1 --port 23333 \
  --served-model-name GLM-5.2-FP8 \
  --tensor-parallel-size 8 --decode-context-parallel-size "$DCP" \
  --dcp-comm-backend ag_rs --kv-cache-dtype fp8 \
  --tool-call-parser glm47 --enable-auto-tool-choice \
  --reasoning-parser glm45 \
  --max-model-len "$MAX_MODEL_LEN" \
  --max-num-batched-tokens 8192 \
  --gpu-memory-utilization 0.92 --max-num-seqs 64 \
  --speculative-config.method mtp \
  --speculative-config.num_speculative_tokens 5

EvalScope 1.9.1 client arguments, identical for all four runs:

evalscope perf \
  --model GLM-5.2-FP8 --url http://127.0.0.1:23333/v1 --api openai \
  --dataset swe_smith --tokenizer-path "$MODEL" \
  --max-tokens 1024 --multi-turn --min-turns 5 --max-turns 5 \
  --multi-turn-args '{"first_turn_length":16384,"subsequent_turn_length":2048,"num_workers":1}' \
  --extra-args '{"ignore_eos":true}' \
  --number 16 --parallel 16 --rate 16 --seed 1024

ModelScope was unavailable, so the actual runs used a wrapper reading the same cached raw dataset with unchanged live construction. All 16 first-turn request hashes and dataset shard manifests matched across runs.

Assistance

Assisted with Codex + GPT-6-Astra High, reviewed manually

[Copilot]

🟡 Changes recommended

Five unresolved review findings, including four critical issues, must be addressed before approval.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds opt-in decode context parallelism for PyTorch MLA, including sharded KV handling, distributed attention merging, sparse top-k selection, scheduling, and CUDA-graph support.

Changes:

• Adds DCP configuration, process groups, logical-block scheduling, KV writes, and workspace management.
• Implements dense/sparse FlashMLA decode and prefill paths, MTP metadata, and CUDA-graph support.
• Adds CLI wiring, documentation, and focused correctness tests.

File summaries

File	Summary
tests/pytorch/test_dcp.py	Core DCP planning, allocation, causal-length, and metadata tests.
tests/pytorch/nn/test_nsa.py	NSA scoring and metadata coverage.
tests/pytorch/kernel/test_sparse_index_topk.py	Distributed sparse top-k tests.
tests/pytorch/kernel/test_mla_attention.py	MLA gathering, merging, and prefix-cache tests.
tests/pytorch/kernel/test_flash_attention.py	Partitioned attention and LSE merge tests.
tests/pytorch/kernel/test_fill_kv_cache.py	DCP KV-cache write tests.
tests/pytorch/kernel/test_dsa_indexer_preprocess.py	DCP DSA preprocessing tests.
tests/pytorch/engine/test_executor_base.py	DCP workspace setup tests.
lmdeploy/pytorch/paging/scheduler.py	DCP logical block sizing and scheduling.
lmdeploy/pytorch/paging/block_manager/default_block_manager.py	DCP-aware allocation; critical finding (1 vote) covers unsupported window-attention accounting.
lmdeploy/pytorch/paging/block_manager/__init__.py	Passes DCP settings to block managers.
lmdeploy/pytorch/models/utils/cudagraph.py	DCP CUDA-graph buffers and metadata.
lmdeploy/pytorch/models/deepseek_v32_mtp.py	MTP model state handling.
lmdeploy/pytorch/kernels/cuda/step_metadata/fill_graph_common_inputs.py	DCP graph input lengths.
lmdeploy/pytorch/kernels/cuda/step_metadata/fill_dsa_indexer_metadata.py	Per-query DCP DSA metadata.
lmdeploy/pytorch/kernels/cuda/sparse_index_topk.py	Sparse top-k constant updates.
lmdeploy/pytorch/kernels/cuda/sparse_index_dcp_topk.py	Distributed candidate packing and merging.
lmdeploy/pytorch/kernels/cuda/flashattention.py	Optional attention LSE output.
lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py	Interleaved DCP KV-cache writes.
lmdeploy/pytorch/kernels/cuda/dsa_indexer_preprocess.py	DCP-aware DSA cache preprocessing.
lmdeploy/pytorch/kernels/cuda/dcp.py	DCP compaction, merging, and reordering kernels.
lmdeploy/pytorch/engine/executor/base.py	DCP workspace reservation; moderate finding (1 vote) covers under-budgeted merge buffers.
lmdeploy/pytorch/engine/engine.py	DCP session-length and cache-capacity accounting.
lmdeploy/pytorch/engine/config_builder.py	Propagates DCP configuration.
lmdeploy/pytorch/distributed.py	DCP process groups and collectives.
lmdeploy/pytorch/config.py	DCP configuration validation.
lmdeploy/pytorch/backends/nsa.py	DCP NSA metadata.
lmdeploy/pytorch/backends/cuda/nsa.py	DCP DSA scoring and top-k handling.
lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py	Sparse MLA DCP paths; critical finding (2 votes) covers the TileLang override bypass.
lmdeploy/pytorch/backends/cuda/attention/mla.py	Dense MLA DCP execution.
lmdeploy/pytorch/backends/cuda/attention/default.py	DCP KV integration; critical finding (3 votes) covers non-FlashMLA decode paths.
lmdeploy/pytorch/backends/cp_utils.py	DCP planning, gathering, and merge helpers.
lmdeploy/messages.py	Public DCP engine configuration; critical finding (1 vote) covers non-CUDA validation.
lmdeploy/cli/utils.py	Defines the DCP CLI argument.
lmdeploy/cli/serve.py	Server CLI configuration wiring.
lmdeploy/cli/cli.py	Chat CLI configuration wiring.
docs/en/advance/context_parallel.md	Documents DCP usage and limitations.

Review details

Suppressed comments (2)

lmdeploy/messages.py:610

• DCP changes the scheduler/trie unit to block_size * dcp, so one allocated connector block represents a DCP-wide logical block. Mooncake still rounds and hashes by the physical block_size; after the first DCP-wide block, its hash count and request_blocks length diverge (and loads/saves raise the connector block-count errors). Since external connectors are documented as unsupported, reject dcp > 1 when kv_transfer_config is configured instead of accepting this combination.

        assert self.dcp >= 1, 'invalid dcp'
        assert self.tp % self.dcp == 0, 'tp must be divisible by dcp'

lmdeploy/pytorch/engine/executor/base.py:245

• The temporary merge buffers operate on the gathered head count (local_heads * dcp), but this reservation only budgets local_heads. With DCP>1 the cache-size calculation underestimates correction/LSE memory by the DCP factor, so it can allocate a KV cache that later OOMs during attention.

        attention_workspace += num_prefill_tokens * local_heads * (model.head_dim * 12 + 12)

• Files reviewed: 37/37 changed files
• Comments generated: 4
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[grimoire]

cuda kernel should not be called here.

[grimoire]

The build of this class would failed on dlinfer backend.

[grimoire]

Is prefix_output bf16? accumulation on bf16 might loss precision.

[CUHKSZzxy]

Yes, prefix_output is BF16. Merge arithmetic is FP32, but each result is stored back in BF16, matching vLLM’s MLA chunk merging. This reduces intermediate memory usage and traffic, but repeated rounding can accumulate. We can retain FP32 intermediates if preferred.

[grimoire]

The kernel requires that dcp*k should be multiple of 512， better add an assert in the kernel launcher.