Repository navigation
Conversation
Add GQA decode shard merging and bounded cached-prefix prefill using shared DCP collectives. Extend paged attention to return LSE and FP8 KV fill to write rank-owned tokens. Share KV workspace sizing and update configuration, documentation, and regression tests.
CUHKSZzxy
force-pushed
the
feat/dcp-optimizations
branch
from
October 9, 2026 09:19
3ff3ca0 to
24ab1c7
Compare
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
GQA DCP currently permits configurations whose ALiBi or block-sparse semantics are silently discarded.
3 open findings
What changed in this PR
Extends PyTorch DCP from MLA to replicated-KV GQA models and introduces opt-in optimized CUDA communication.
Changes:
- Adds GQA DCP prefill/decode with FP8 KV-cache support.
- Centralizes DCP collectives and symmetric-memory workspaces.
- Adds
--communication-backend autowith native fallback.
| File | Description |
|---|---|
tests/pytorch/test_cuda_communicator.py |
Expands communicator and workspace tests. |
tests/pytorch/spec_decode/test_spec_agent.py |
Tests backend propagation to draft contexts. |
tests/pytorch/kernel/test_mla_attention.py |
Updates shared DCP manager tests. |
tests/pytorch/kernel/test_fill_kv_cache.py |
Tests sharded FP8 cache filling. |
tests/pytorch/engine/test_executor_base.py |
Updates DCP workspace fixture. |
lmdeploy/pytorch/third_party/flash_attn_interface.py |
Normalizes optional FA3 LSE returns. |
lmdeploy/pytorch/spec_decode/proposers/base.py |
Propagates draft-cache quantization policy. |
lmdeploy/pytorch/spec_decode/base.py |
Propagates communication backend. |
lmdeploy/pytorch/nn/norm.py |
Removes deferred fused all-reduce handling. |
lmdeploy/pytorch/nn/linear/base.py |
Routes reductions through the wrapper. |
lmdeploy/pytorch/nn/embedding.py |
Uses communicator-managed logits gathering. |
lmdeploy/pytorch/models/deepseek_v32.py |
Removes deferred reduction configuration. |
lmdeploy/pytorch/kernels/cuda/symm_mem_allgather.py |
Generalizes symmetric all-gather dimensions. |
lmdeploy/pytorch/kernels/cuda/pagedattention.py |
Adds natural-log LSE output. |
lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py |
Adds sharded FP8 cache writes. |
lmdeploy/pytorch/kernels/cuda/dcp.py |
Vectorizes prefix reordering. |
lmdeploy/pytorch/envs.py |
Removes legacy communicator environment flags. |
lmdeploy/pytorch/engine/executor/ray_executor.py |
Updates symmetric-memory device setup. |
lmdeploy/pytorch/engine/executor/base.py |
Accounts for GQA KV width. |
lmdeploy/pytorch/engine/executor/__init__.py |
Validates GQA DCP quantization. |
lmdeploy/pytorch/distributed.py |
Centralizes communicator and DCP ownership. |
lmdeploy/pytorch/config.py |
Adds backend configuration and GQA validation. |
lmdeploy/pytorch/backends/dlinfer/op_backend.py |
Removes obsolete CUDA-option handling. |
lmdeploy/pytorch/backends/default/linear.py |
Uses optimized reduction routing. |
lmdeploy/pytorch/backends/default/embedding.py |
Uses optimized reduction routing. |
lmdeploy/pytorch/backends/cuda/op_backend.py |
Propagates communicator group names. |
lmdeploy/pytorch/backends/cuda/nsa.py |
Uses shared candidate gathering. |
lmdeploy/pytorch/backends/cuda/comm/symm_mem_allreduce.py |
Revises provider preparation and fallback. |
lmdeploy/pytorch/backends/cuda/comm/symm_mem_allgather.py |
Generalizes reusable gather workspaces. |
lmdeploy/pytorch/backends/cuda/comm/flashinfer_allreduce.py |
Removes the FlashInfer provider. |
lmdeploy/pytorch/backends/cuda/comm/communicator.py |
Implements automatic collective selection. |
lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py |
Uses shared DCP collectives. |
lmdeploy/pytorch/backends/cuda/attention/mla.py |
Uses centralized DCP workspaces. |
lmdeploy/pytorch/backends/cuda/attention/cp.py |
Adds GQA DCP attention. |
lmdeploy/pytorch/backends/cuda/attention/__init__.py |
Selects GQA DCP implementation. |
lmdeploy/pytorch/backends/cp_utils.py |
Shares KV-width-based workspace planning. |
lmdeploy/pytorch/backends/communicator.py |
Adds generic all-gather APIs. |
lmdeploy/pytorch/backends/base.py |
Extends communicator construction. |
lmdeploy/messages.py |
Exposes communication backend configuration. |
lmdeploy/cli/utils.py |
Defines the backend CLI option. |
lmdeploy/cli/serve.py |
Wires the option into API serving. |
lmdeploy/cli/cli.py |
Wires the option into interactive chat. |
docs/en/advance/context_parallel.md |
Documents GQA DCP support and restrictions. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
# Conflicts: # lmdeploy/pytorch/nn/linear/base.py
# Conflicts: # lmdeploy/cli/cli.py
This branch has not been deployed
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
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.
Motivation
Depends on #4952. Merge that PR first. This extends its MLA DCP support to GQA models with replicated KV heads and adds opt-in optimized communication paths. The diff against
mainincludes #4952 until it is merged.Modification
--communication-backend auto, using symmetric-memory all-gather for query, LSE, top-k candidate, and cached-prefix payloads when eligible.Accuracy
Full GSM8K test split, zero-shot, thinking disabled, temperature 0, seed 20260921, output limit 2,048 tokens, concurrency 16. Both runs used FP8 KV and prefix caching on H200 GPUs.
Both mini checks passed 8/8. Qwen3.5 MTP acceptance was 92.71%. Unspecified engine settings used CLI defaults. These are single runs without a TP-only baseline. Prefix caching was enabled, but no cache hits were observed. The optimized communication paths were not enabled in these accuracy runs.
Checklist
Assistance
Assisted with Codex + GPT-6-Astra High; reviewed manually.