Conversation
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.
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.
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 an optional optimized query gather. The diff against
mainincludes #4952 until it is merged.Modification
LMDEPLOY_ENABLE_SYMM_MEM_DCP=1) and a process-group fallback.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. 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
Assistance
Assisted with Codex + GPT-6-Astra High, reviewed manually