perf(dcp): split prefill queries and gather selected CKV - #110
Closed
voipmonitor wants to merge 7 commits into
Closed
perf(dcp): split prefill queries and gather selected CKV#110voipmonitor 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
📝 WalkthroughWalkthroughChangesThe PR adds environment-controlled DCP query splitting and transient full-CKV gathering for B12X sparse MLA. It introduces communication groups, metadata and workspace handling, query-split top-k processing, CKV remapping/prefetch, updated MLA dispatch, and CUDA unit coverage. DCP CKV gather flow
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant B12xMLASparseMetadataBuilder
participant B12xMLASparseImpl
participant DCPPrefetchGroup
participant MLAAttention
MLAAttention->>B12xMLASparseImpl: select eligible CKV gather path
B12xMLASparseImpl->>DCPPrefetchGroup: all-gather CKV cache
DCPPrefetchGroup-->>B12xMLASparseImpl: gathered CKV buffer
B12xMLASparseImpl-->>MLAAttention: bind gathered-cache attention plan
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 1905-1907: Update the query-split rank construction in the
surrounding parallel-state initialization to partition group_ranks by TP replica
before transposing. Build each query-split group only from ranks belonging to
the same TP row, preserving DP/PP/PCP replica and pipeline-stage boundaries
rather than iterating across the entire group_ranks list.
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 1516-1524: Remove the broad exception suppression around
get_query_split_group in the query-split initialization block. When
VLLM_DCP_QUERY_SPLIT is enabled with dcp_world_size greater than one, let
communicator acquisition or metadata errors propagate so initialization fails
instead of silently leaving qs_group disabled; preserve the existing world-size
and rank assignments for valid groups.
In `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py`:
- Around line 698-704: Update the CKV gather runtime condition in the relevant
MLA attention path to compare the minimum threshold against total batch tokens
via num_tokens, matching the builder behavior. Keep the other use_dcp,
decode-token, and prefill checks unchanged.
🪄 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: 3e614a91-3d94-4e35-8523-c3c6656ffbcc
📒 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
+1905
to
+1907
| 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]) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Keep query-split groups within each TP replica.
group_ranks spans all DP/PP/PCP replicas. Transposing the entire list combines unrelated replicas or pipeline stages, so their all-gathers can mix different requests or deadlock on differing shapes. Build query-split groups separately for each TP row.
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])
+ rank_layout = local_all_ranks if enable_elastic_ep else all_ranks
+ query_split_ranks: list[list[int]] = []
+ for tp_ranks in rank_layout.reshape(-1, tensor_model_parallel_size):
+ dcp_groups = tp_ranks.reshape(
+ -1, decode_context_model_parallel_size
+ )
+ query_split_ranks.extend(
+ dcp_groups[:, dcp_rank_idx].tolist()
+ for dcp_rank_idx in range(
+ decode_context_model_parallel_size
+ )
+ )📝 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
| 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]) | |
| rank_layout = local_all_ranks if enable_elastic_ep else all_ranks | |
| query_split_ranks: list[list[int]] = [] | |
| for tp_ranks in rank_layout.reshape(-1, tensor_model_parallel_size): | |
| dcp_groups = tp_ranks.reshape( | |
| -1, decode_context_model_parallel_size | |
| ) | |
| query_split_ranks.extend( | |
| dcp_groups[:, dcp_rank_idx].tolist() | |
| for dcp_rank_idx in range( | |
| decode_context_model_parallel_size | |
| ) | |
| ) |
🤖 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 1905 - 1907, Update the
query-split rank construction in the surrounding parallel-state initialization
to partition group_ranks by TP replica before transposing. Build each
query-split group only from ranks belonging to the same TP row, preserving
DP/PP/PCP replica and pipeline-stage boundaries rather than iterating across the
entire group_ranks list.
Comment on lines
+1516
to
+1524
| 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 query-split communicator is unavailable.
The flag and DCP topology guarantee this group exists. Swallowing every exception can silently disable the requested feature or produce rank-divergent collective participation and a 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 = 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)📝 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 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 | |
| if envs.VLLM_DCP_QUERY_SPLIT and dcp_world_size > 1: | |
| _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) |
🧰 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 1516 - 1524,
Remove the broad exception suppression around get_query_split_group in the
query-split initialization block. When VLLM_DCP_QUERY_SPLIT is enabled with
dcp_world_size greater than one, let communicator acquisition or metadata errors
propagate so initialization fails instead of silently leaving qs_group disabled;
preserve the existing world-size and rank assignments for valid groups.
Source: Linters/SAST tools
Comment on lines
+698
to
+704
| if ( | ||
| use_dcp | ||
| and envs_mod.VLLM_B12X_MLA_CKV_GATHER | ||
| and num_decode_tokens == 0 | ||
| and num_prefill_tokens == num_tokens | ||
| and cm.max_query_len > envs_mod.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS | ||
| ): |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Apply the minimum threshold to total batch tokens.
The runtime check uses num_tokens, but the builder uses cm.max_query_len. For example, two 16-token requests total 32 tokens yet fail the default minimum check here, unnecessarily disabling CKV gather.
Proposed fix
- and cm.max_query_len > envs_mod.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS
+ and num_tokens > envs_mod.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS📝 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 ( | |
| use_dcp | |
| and envs_mod.VLLM_B12X_MLA_CKV_GATHER | |
| and num_decode_tokens == 0 | |
| and num_prefill_tokens == num_tokens | |
| and cm.max_query_len > envs_mod.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS | |
| ): | |
| if ( | |
| use_dcp | |
| and envs_mod.VLLM_B12X_MLA_CKV_GATHER | |
| and num_decode_tokens == 0 | |
| and num_prefill_tokens == num_tokens | |
| and num_tokens > envs_mod.VLLM_B12X_MLA_CKV_GATHER_MIN_TOKENS | |
| ): |
🤖 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 698 - 704,
Update the CKV gather runtime condition in the relevant MLA attention path to
compare the minimum threshold against total batch tokens via num_tokens,
matching the builder behavior. Keep the other use_dcp, decode-token, and prefill
checks unchanged.
Author
|
Superseded by #111, which carries the same validated seven-commit stack on the consolidated dev/gilded-gnosis base with original authorship preserved. |
This was referenced Jul 19, 2026
This was referenced Aug 13, 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 adds two independent, opt-in B12X sparse-MLA DCP prefill optimizations:
VLLM_DCP_QUERY_SPLIT=1: split the replicated prefill query across DCP ranks before sparse MLA and restore the expected output layout afterward.VLLM_B12X_MLA_CKV_GATHER=1: gather only the selected compressed-KV rows needed by sparse MLA instead of running attention over each rank's full local CKV contribution.Both paths remain disabled by default and preserve the existing implementation as their fallback. The CKV path also includes the follow-up short-context correctness fix: because its attention buffer is global, selected-entry counts are capped by each token's global causal length rather than the local per-rank length.
Implementation notes
TP/DCP == 1, so query split is intentionally a functional no-op; the gain there is CKV gather.Measured performance
GLM-5.2 Luke NVFP4, A16, MTP off, TP8, identical v18 integration base on 8x RTX PRO 6000 Blackwell. Each cell is the median of three runs and passed first-token correctness checks.
Disabling both flags on the candidate showed no material regression versus baseline: -0.11% at 8k and -0.36% at 64k.
DCP8 KV capacity changed from 4,581,888 to 4,553,216 tokens, a 28,672-token / 0.63% reduction.
Correctness and verification
ruff check,ruff format --check,git diff --check, and Python bytecode compilation pass for the touched source and test.Result artifacts are stored locally under:
/root/bench-results/glm52-v18-ckv-ab-dcp4-r2-20260717/root/bench-results/glm52-v18-ckv-ab-dcp8-20260717Summary by CodeRabbit
New Features
Bug Fixes
Tests