[GG] perf(dcp): split prefill queries and gather selected CKV - #111
Closed
voipmonitor wants to merge 7 commits into
Closed
[GG] perf(dcp): split prefill queries and gather selected CKV#111voipmonitor wants to merge 7 commits into
voipmonitor wants to merge 7 commits into
Conversation
…h (Fix B)
Fix A — shard indexer queries across DCP pairs:
- New query-split process group in parallel_state.py (ranks sharing
the same dcp_rank form a group; e.g. {0,2,4,6}/{1,3,5,7} at TP=8/DCP=2)
- Row narrowing + cu_seqlen narrowing + finals all-gather in
sparse_attn_indexer.py
- Active page width fix for block_table (from v17 reference patch)
- Eliminates 8/dcp indexer compute redundancy at no memory cost
Fix B — gather full CKV for prefill, eliminating DCP collectives:
- Transient CKV gather workspace with dispose-recycle pooling
- _dcp_gather_ckv: cp_gather_cache local shard, zero-pad, NCCL
all-gather into rank-major gathered buffer
- _map_global_topk_to_gathered_ckv Triton kernel: maps global top-k
indices to physical slots in the gathered buffer
- Separate extend plan with local heads only (no q all-gather)
- Skip LSE return/merge, project_before_merge, workspace gather
- Generalized from v17 (nvfp4-only, TP4/DCP4) to support fp8_ds_mla
and nvfp4_ds_mla with KV_FP8_ROPE at any DCP topology
Layer prefetch (spec §5):
- Ping-pong workspace (2x), side CUDA stream, class-level events
- Layer L kicks off layer L+1's history gather during layer L compute
- _append_current_chunk_to_gathered writes current chunk BF16 KV
into gathered buffer via concat_and_cache_mla (history is in the
paged cache but current chunk is not yet written when prefetch runs)
- set_ckv_current_chunk_kv stores per-layer BF16 latent before
forward_mqa
- Prefetch gated on fp8_ds_mla or KV_FP8_ROPE (concat_and_cache_mla
quantization path)
Env vars: VLLM_DCP_QUERY_SPLIT, VLLM_B12X_MLA_CKV_GATHER,
VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS, VLLM_B12X_MLA_CKV_GATHER_MAX_TOKENS
Co-authored-by: opencode <opencode@anthropic.ai>
Co-authored-by: opencode <opencode@anthropic.ai>
The layer-prefetch path for the transient full-CKV DCP prefill gather was effectively dead and, once enabled, produced corrupted KV. Four fixes: - Enable the prefetch. It keyed on layer.layer_idx, which MLAAttention never sets (it exposes layer_name), so every layer gathered synchronously on the critical path. Resolve the index from layer_name via extract_layer_index (matching the prewarm fallback). - Fix the current-chunk append for multi-request chunks. global_pos was computed with the batch-global token count/index, correct only for a single request; with MAX_NUM_SEQS>1 it misplaced every request but the last, corrupting the appended KV. Compute position per request from query_start_loc. The synchronous gather is unaffected (it reads the current chunk from the paged cache), which is why only the prefetch path regressed. - Reset the cross-layer prefetch pipeline (_shared_gather_event, _shared_gather_buf_idx) once per step in the metadata builder. The state leaked across chunks, mis-scheduling layer 0 of later chunks onto a stale gathered buffer. - Order the side-stream prefetch after the default stream's KV-cache writes (stream.wait_stream) so it observes prior do_kv_cache_update history; and route the prefetch all-gather through a dedicated DCP communicator so it cannot collide with the indexer's DCP top-k merge on the default stream. Also fix a latent .view(-1) crash on the non-contiguous dcp_rank_req_starts slice (use reshape) and note the intentionally-unspecified compaction order in _map_global_topk_to_gathered_ckv. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Vdiz6KR1rDBfMkTWGFrHjL
The CKV-gather prefill path attends the full global sequence (cache_seqlens = global_cache_seq_lens_per_req), but clamped each token's selected-entry count (nsa_cache_seqlens) by the local per-rank causal length (~global/dcp) copied from the non-gather DCP path. The clamp is a no-op when local_len >= min(topk, global_len) (long context), but truncates the selected KV set to roughly global/dcp when local_len is the smaller bound (short context), dropping the other ranks' selected tokens and producing small-context-only garbage. Clamp by the global per-token causal length instead, computed from the global sequence lengths and query_start_loc. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ui7hs7Ygur2iSiaHkNW1Kv
📝 WalkthroughWalkthroughAdds optional B12X DCP query splitting and CKV gathering. The change introduces new process groups, metadata and Triton helpers, cross-layer KV prefetching, global top-k remapping, causal masking, updated MLA routing, and CUDA coverage for causal-length computation. ChangesB12X DCP attention flow
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant MLA_attention
participant B12xMLASparseImpl
participant DCP_group
participant CKV_extend_plan
MLA_attention->>B12xMLASparseImpl: evaluate CKV gather eligibility
B12xMLASparseImpl->>DCP_group: all-gather local paged KV
DCP_group-->>B12xMLASparseImpl: return gathered KV records
B12xMLASparseImpl->>CKV_extend_plan: bind gathered cache and global lengths
CKV_extend_plan-->>MLA_attention: return attention output
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@vllm/distributed/parallel_state.py`:
- Around line 1904-1908: Update the query-split construction in the
decode-context parallel initialization block so ranks are grouped only within
each TP cohort, not across all entries in group_ranks. Partition or iterate
group_ranks by the appropriate TP cohort boundaries, then collect each DCP
position within each cohort; preserve singleton groups as the intended no-op for
configurations such as DP2 with TP8/DCP8.
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 1513-1524: Update the query-split initialization around
get_query_split_group so exceptions are not swallowed when VLLM_DCP_QUERY_SPLIT
is enabled and dcp_world_size exceeds one. Fail initialization immediately by
propagating the exception, while preserving the existing assignment of qs_group,
qs_world_size, and qs_rank for valid multi-rank groups.
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py`:
- Around line 1604-1623: Update dcp_prefill_ckv_gather_eligible so it returns
false whenever the surrounding forward_mqa flow would select the global-head
_decode_plan under speculative extend-as-decode; require the CKV-specific extend
plan conditions instead, preventing CKV local-head output from entering a
DCP-gathered-head plan.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: c25801a2-90b1-4d7d-9e12-9945659f3179
📒 Files selected for processing (6)
tests/v1/attention/test_sparse_mla_backends.pyvllm/distributed/parallel_state.pyvllm/envs.pyvllm/model_executor/layers/attention/mla_attention.pyvllm/model_executor/layers/sparse_attn_indexer.pyvllm/v1/attention/backends/mla/b12x_mla_sparse.py
Comment on lines
+1904
to
+1908
| if decode_context_model_parallel_size > 1 and envs.VLLM_DCP_QUERY_SPLIT: | ||
| query_split_ranks: list[list[int]] = [] | ||
| for dcp_rank_idx in range(decode_context_model_parallel_size): | ||
| query_split_ranks.append([grp[dcp_rank_idx] for grp in group_ranks]) | ||
| _QUERY_SPLIT = init_model_parallel_group( |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Keep query-split groups inside each TP cohort.
group_ranks contains DCP groups from every DP/PP/PCP cohort. Collecting each DCP position across the entire list merges unrelated replicas. For DP2 with TP8/DCP8, this even creates two-rank query groups instead of the intended singleton no-op, risking incorrect outputs or collective hangs.
Proposed fix
- query_split_ranks: list[list[int]] = []
- for dcp_rank_idx in range(decode_context_model_parallel_size):
- query_split_ranks.append([grp[dcp_rank_idx] for grp in group_ranks])
+ query_split_source = local_all_ranks if enable_elastic_ep else all_ranks
+ dcp_groups_per_tp = (
+ tensor_model_parallel_size // decode_context_model_parallel_size
+ )
+ query_split_ranks = (
+ query_split_source.reshape(
+ -1,
+ dcp_groups_per_tp,
+ decode_context_model_parallel_size,
+ )
+ .transpose(1, 2)
+ .reshape(-1, dcp_groups_per_tp)
+ .tolist()
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if decode_context_model_parallel_size > 1 and envs.VLLM_DCP_QUERY_SPLIT: | |
| query_split_ranks: list[list[int]] = [] | |
| for dcp_rank_idx in range(decode_context_model_parallel_size): | |
| query_split_ranks.append([grp[dcp_rank_idx] for grp in group_ranks]) | |
| _QUERY_SPLIT = init_model_parallel_group( | |
| if decode_context_model_parallel_size > 1 and envs.VLLM_DCP_QUERY_SPLIT: | |
| query_split_source = local_all_ranks if enable_elastic_ep else all_ranks | |
| dcp_groups_per_tp = ( | |
| tensor_model_parallel_size // decode_context_model_parallel_size | |
| ) | |
| query_split_ranks = ( | |
| query_split_source.reshape( | |
| -1, | |
| dcp_groups_per_tp, | |
| decode_context_model_parallel_size, | |
| ) | |
| .transpose(1, 2) | |
| .reshape(-1, dcp_groups_per_tp) | |
| .tolist() | |
| ) | |
| _QUERY_SPLIT = init_model_parallel_group( |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/distributed/parallel_state.py` around lines 1904 - 1908, Update the
query-split construction in the decode-context parallel initialization block so
ranks are grouped only within each TP cohort, not across all entries in
group_ranks. Partition or iterate group_ranks by the appropriate TP cohort
boundaries, then collect each DCP position within each cohort; preserve
singleton groups as the intended no-op for configurations such as DP2 with
TP8/DCP8.
Comment on lines
+1513
to
+1524
| qs_group = None | ||
| qs_world_size = 1 | ||
| qs_rank = 0 | ||
| if envs.VLLM_DCP_QUERY_SPLIT and dcp_world_size > 1: | ||
| try: | ||
| _qs = get_query_split_group() | ||
| if int(_qs.world_size) > 1: | ||
| qs_group = _qs | ||
| qs_world_size = int(_qs.world_size) | ||
| qs_rank = int(_qs.rank_in_group) | ||
| except Exception: | ||
| pass |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Fail fast when the enabled query-split group is unavailable.
The initializer guarantees this group when the flag and DCP are active. Swallowing every exception can silently disable the feature—or make ranks diverge before later collectives and hang.
Proposed fix
if envs.VLLM_DCP_QUERY_SPLIT and dcp_world_size > 1:
- try:
- _qs = get_query_split_group()
- if int(_qs.world_size) > 1:
- qs_group = _qs
- qs_world_size = int(_qs.world_size)
- qs_rank = int(_qs.rank_in_group)
- except Exception:
- pass
+ qs_group = get_query_split_group()
+ if int(qs_group.world_size) > 1:
+ qs_world_size = int(qs_group.world_size)
+ qs_rank = int(qs_group.rank_in_group)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| qs_group = None | |
| qs_world_size = 1 | |
| qs_rank = 0 | |
| if envs.VLLM_DCP_QUERY_SPLIT and dcp_world_size > 1: | |
| try: | |
| _qs = get_query_split_group() | |
| if int(_qs.world_size) > 1: | |
| qs_group = _qs | |
| qs_world_size = int(_qs.world_size) | |
| qs_rank = int(_qs.rank_in_group) | |
| except Exception: | |
| pass | |
| qs_group = None | |
| qs_world_size = 1 | |
| qs_rank = 0 | |
| if envs.VLLM_DCP_QUERY_SPLIT and dcp_world_size > 1: | |
| qs_group = get_query_split_group() | |
| if int(qs_group.world_size) > 1: | |
| qs_world_size = int(qs_group.world_size) | |
| qs_rank = int(qs_group.rank_in_group) |
🧰 Tools
🪛 Ruff (0.15.21)
[error] 1523-1524: try-except-pass detected, consider logging the exception
(S110)
[warning] 1523-1523: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/model_executor/layers/sparse_attn_indexer.py` around lines 1513 - 1524,
Update the query-split initialization around get_query_split_group so exceptions
are not swallowed when VLLM_DCP_QUERY_SPLIT is enabled and dcp_world_size
exceeds one. Fail initialization immediately by propagating the exception, while
preserving the existing assignment of qs_group, qs_world_size, and qs_rank for
valid multi-rank groups.
Source: Linters/SAST tools
Comment on lines
+1604
to
+1623
| def dcp_prefill_ckv_gather_eligible( | ||
| self, | ||
| attn_metadata: B12xMLASparseMetadata, | ||
| num_tokens: int, | ||
| ) -> bool: | ||
| if not self._ckv_gather_enabled: | ||
| return False | ||
| if torch.cuda.is_current_stream_capturing(): | ||
| return False | ||
| if ( | ||
| not attn_metadata.dcp_ckv_gather_eligible | ||
| or attn_metadata.num_decode_tokens != 0 | ||
| or attn_metadata.num_prefill_tokens != attn_metadata.num_actual_tokens | ||
| or int(num_tokens) != attn_metadata.num_actual_tokens | ||
| or int(num_tokens) <= self._ckv_gather_min_tokens | ||
| or attn_metadata.dcp_padded_total_tokens > self._ckv_local_capacity | ||
| or attn_metadata.dcp_local_total_tokens | ||
| > attn_metadata.dcp_padded_total_tokens | ||
| ): | ||
| return False |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Prevent CKV gather from entering the global-head decode plan.
With speculative extend-as-decode enabled, a lowered CKV minimum or raised speculative threshold can make this return True while forward_mqa selects _decode_plan. CKV supplies local heads, but that plan expects DCP-gathered heads, causing a bind/geometry failure. Fall back or force the CKV extend plan.
Conservative eligibility fix
or int(num_tokens) <= self._ckv_gather_min_tokens
+ or (
+ self.spec_extend_as_decode
+ and attn_metadata.max_query_len <= self.spec_decode_max_q
+ )
or attn_metadata.dcp_padded_total_tokens > self._ckv_local_capacity📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def dcp_prefill_ckv_gather_eligible( | |
| self, | |
| attn_metadata: B12xMLASparseMetadata, | |
| num_tokens: int, | |
| ) -> bool: | |
| if not self._ckv_gather_enabled: | |
| return False | |
| if torch.cuda.is_current_stream_capturing(): | |
| return False | |
| if ( | |
| not attn_metadata.dcp_ckv_gather_eligible | |
| or attn_metadata.num_decode_tokens != 0 | |
| or attn_metadata.num_prefill_tokens != attn_metadata.num_actual_tokens | |
| or int(num_tokens) != attn_metadata.num_actual_tokens | |
| or int(num_tokens) <= self._ckv_gather_min_tokens | |
| or attn_metadata.dcp_padded_total_tokens > self._ckv_local_capacity | |
| or attn_metadata.dcp_local_total_tokens | |
| > attn_metadata.dcp_padded_total_tokens | |
| ): | |
| return False | |
| def dcp_prefill_ckv_gather_eligible( | |
| self, | |
| attn_metadata: B12xMLASparseMetadata, | |
| num_tokens: int, | |
| ) -> bool: | |
| if not self._ckv_gather_enabled: | |
| return False | |
| if torch.cuda.is_current_stream_capturing(): | |
| return False | |
| if ( | |
| not attn_metadata.dcp_ckv_gather_eligible | |
| or attn_metadata.num_decode_tokens != 0 | |
| or attn_metadata.num_prefill_tokens != attn_metadata.num_actual_tokens | |
| or int(num_tokens) != attn_metadata.num_actual_tokens | |
| or int(num_tokens) <= self._ckv_gather_min_tokens | |
| or ( | |
| self.spec_extend_as_decode | |
| and attn_metadata.max_query_len <= self.spec_decode_max_q | |
| ) | |
| or attn_metadata.dcp_padded_total_tokens > self._ckv_local_capacity | |
| or attn_metadata.dcp_local_total_tokens | |
| > attn_metadata.dcp_padded_total_tokens | |
| ): | |
| return False |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py` around lines 1604 - 1623,
Update dcp_prefill_ckv_gather_eligible so it returns false whenever the
surrounding forward_mqa flow would select the global-head _decode_plan under
speculative extend-as-decode; require the CKV-specific extend plan conditions
instead, preventing CKV local-head output from entering a DCP-gathered-head
plan.
Author
|
Additional GG verification: the complete focused sparse-MLA backend module passed in the v18 CUDA runtime image: 74 passed, 714 skipped ( |
Author
|
Final baked-image validation (vLLM
Machine-readable release results will be published with the GLM-5.2 v18 wiki page. |
voipmonitor
changed the title
Author
|
Superseded by the canonical dev/gilded-gnosis consolidation. The corresponding implementation is present in commit(s): a87f739,e408f0d18c,40f63da46e,9a017dab5a,0fd5178972,adf65d4269. Closing the old-base PR so future work targets the canonical GG branch. |
This was referenced Jul 20, 2026
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
2 participants
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.
Summary
This ports the already validated DCP prefill optimization stack to the consolidated
dev/gilded-gnosisbranch. It adds two independent, opt-in B12X sparse-MLA paths:VLLM_DCP_QUERY_SPLIT=1splits the replicated prefill query across DCP ranks before sparse MLA and restores the expected output layout afterward.VLLM_B12X_MLA_CKV_GATHER=1gathers only the selected compressed-KV rows needed by sparse MLA instead of running attention over each rank's full local CKV contribution.The seven commits preserve the original authorship and applied cleanly to GG. Both paths remain disabled by default and preserve the existing implementation as fallback.
Implementation
For TP8/DCP8,
TP/DCP == 1, so query split is intentionally a no-op; the measured gain there comes from selected-CKV gather.Performance
GLM-5.2 Luke NVFP4, A16, MTP off, TP8, identical v18 integration base on 8x RTX PRO 6000 Blackwell. Each value is the median of three runs and passed first-token checks.
With both flags disabled, the candidate differed from baseline by -0.11% at 8k and -0.36% at 64k. DCP8 KV capacity changed from 4,581,888 to 4,553,216 tokens (-0.63%).
Verification
ruff check,ruff format --check,git diff --check, and Python compilation pass on GG.Supersedes #110, which contains the same stack against the retired FF integration base.
Summary by CodeRabbit
New Features
Bug Fixes