Skip to content

feat: extend DCP to GQA and optimize query communication - #4995

Draft
CUHKSZzxy wants to merge 38 commits into
InternLM:mainfrom
CUHKSZzxy:feat/dcp-optimizations
Draft

CUHKSZzxy wants to merge 38 commits into
InternLM:mainfrom
CUHKSZzxy:feat/dcp-optimizations

Conversation

@CUHKSZzxy

Copy link
Copy Markdown
Collaborator

Motivation

Depends on #4952; merge that PR first. This extends its MLA DCP support to GQA models with replicated KV heads and adds an optional optimized query gather. The diff against main includes #4952 until it is merged.

Modification

  • Add GQA DCP prefill/decode with bounded cached-prefix gathering and shared DCP components.
  • Support unquantized and per-tensor FP8 KV cache, including sharded FP8 cache fill. Propagate the allocated draft cache's quantization policy into MTP attention to fix FP8 warmup failure.
  • Add communicator-owned query gathering with opt-in symmetric memory (LMDEPLOY_ENABLE_SYMM_MEM_DCP=1) and a process-group fallback.
  • Share KV workspace sizing and document supported layouts and restrictions.

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.

Model Configuration Correct / total Accuracy Request / parser errors Truncated
Qwen3-235B-A22B TP8+DCP2 1253/1319 95.00% 0 / 0 1
Qwen3.5-397B-A17B-FP8 TP4+DCP2, MTP2, batch limit 128 1276/1319 96.74% 0 / 0 10

Both mini checks passed 8/8. Qwen3.5 MTP acceptance was 92.71%. Unspecified engine settings used CLI defaults. Qwen3 was tested before the final draft-policy fix; Qwen3.5 includes it. These are single runs without a TP-only baseline; prefix caching was enabled but no cache hits were observed. The optional symmetric-memory gather was not enabled in these accuracy runs.

Checklist

  • Pre-commit checks passed for the final fix.
  • Focused speculative-decoding tests: 24 passed, one CUDA-only test skipped.
  • Context-parallel documentation updated.
  • Full kernel/backend and dependency-version matrix remains unvalidated.

Assistance

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

Count speculative verification rows in DCP top-k workspace estimates and clarify the reserved buffer scope. Extend cached-prefill numerical coverage to DCP2/DCP4 and BF16/FP8 caches.

Validation: 53 executor/DCP CPU tests and 14 DCP attention tests passed on one H200 GPU. Ruff and git diff --check passed.
Preserve DCP cache ownership and metadata alongside upstream piecewise CUDA graph entry points. Retain both DCP and TileLang numerical coverage, reject the unsupported TileLang attention plus DCP combination, and remove the requested MTP loader test.

Validation: 233 CPU tests passed (16 skipped); all 40 MLA/indexer tests passed on one H200 GPU.
Fuse rank-local index filtering and ordered compaction, avoid packing contiguous queries before all-gather, and clarify DCP helper names.

Remove full-score NaN sanitation while preserving index masks and empty-shard attention handling. Widen candidate score addressing to int64 for large CUDA graph buffers and cover the affected paths with regression tests.
Keep LSE communication and softmax correction arithmetic in FP32, but store corrected contributions in the input dtype for reduce-scatter. Cover both BF16 and FP32 outputs in the correction regression test.

Validation: 24 DCP attention tests passed in an isolated dtype-only snapshot.
Cache composite keys across 11-bit radix passes, use block-wide prefix scans, and coalesce candidate loads across warp lanes. Ballot-based emission preserves exact candidate order and bit-preserved INT32 IDs.

Extend coverage to DCP8 and large-ID radix boundary ties. Validation: 46 MLA and selector tests passed in an isolated snapshot without communication optimizations; exact-reference, CUDA graph, and sanitizer checks also passed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant