[GG] fix(dcp): build one-shard indexer query-split topology - #245
Conversation
VLLM_DCP_INDEXER_SHARDS=1 replicates the sparse-indexer K cache on every DCP rank, but the indexer topology was only built for 1 < shards < DCP, so one-shard runs fell back to _QUERY_SPLIT, which is singleton at TP == DCP. Every rank then scored every query row against the full indexer cache during prefill: DCP-wide redundant compute (~8x at DCP8, measurable as a long-context prefill throughput loss). With VLLM_DCP_QUERY_SPLIT=1 it was worse than redundant, since get_indexer_query_split_group(1) had no group to return and raised "No indexer query-split group matches the requested KV shard count". _build_indexer_replica_group_ranks already yields the correct layout for shards=1 -- single-rank owner groups and one TP-wide query-split group -- so let it run. Each rank scores totalQ/TP rows against the full replicated cache, and the existing indices-only all-gather restores the shared top-k buffer. No merge is needed or attempted: dcp_world_size stays 1 for a replicated cache, so every candidate-merge path still early-returns. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01J4BQPVrK6LdTUbXudyPkga
|
Warning Review limit reached
Next review available in: 53 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThe change centralizes detection of partial indexer replica-group layouts. It enables initialization when ChangesIndexer replica-group initialization
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 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: 1
🧹 Nitpick comments (1)
vllm/distributed/parallel_state.py (1)
1545-1547: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Google-style sections to the helper docstring.
The new helper has two arguments and a boolean return value. Add
Args:andReturns:sections to document them.As per coding guidelines, use Google-style docstrings with
Args:/Returns:/Raises:sections instead of reStructuredText/Sphinx fields.Proposed docstring update
def _needs_indexer_replica_groups(indexer_shards: int, dcp_size: int) -> bool: - """Return whether the indexer needs replica-specific process groups.""" + """Return whether the indexer needs replica-specific process groups. + + Args: + indexer_shards: Number of indexer shards. + dcp_size: Decode-context parallel size. + + Returns: + True when indexer_shards is positive and smaller than dcp_size. + """ return 1 <= indexer_shards < dcp_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 1545 - 1547, Update the _needs_indexer_replica_groups docstring with Google-style Args and Returns sections, documenting indexer_shards, dcp_size, and the boolean condition indicating whether replica-specific process groups are needed. Keep the helper logic unchanged.Source: Coding guidelines
🤖 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 `@tests/distributed/test_indexer_parallel_groups.py`:
- Line 62: Wrap the signature of
test_indexer_replica_group_initialization_gate_for_dcp8 across multiple lines so
no line exceeds the 88-character limit; keep the function name, parameters, and
behavior unchanged.
---
Nitpick comments:
In `@vllm/distributed/parallel_state.py`:
- Around line 1545-1547: Update the _needs_indexer_replica_groups docstring with
Google-style Args and Returns sections, documenting indexer_shards, dcp_size,
and the boolean condition indicating whether replica-specific process groups are
needed. Keep the helper logic unchanged.
🪄 Autofix
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: 3a25989e-2f30-4a84-9313-e89b3a48d91d
📥 Commits
Reviewing files that changed from the base of the PR and between 3003860 and 6d6d71a436e85957227f1c5c2ac2bea8c4abfd7f.
📒 Files selected for processing (2)
tests/distributed/test_indexer_parallel_groups.pyvllm/distributed/parallel_state.py
| ("indexer_shards", "expected"), | ||
| [(0, False), (1, True), (2, True), (4, True), (8, False)], | ||
| ) | ||
| def test_indexer_replica_group_initialization_gate_for_dcp8(indexer_shards, expected): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Wrap the test definition to meet the line-length limit.
Line 62 is 90 characters long. Split the function signature across lines.
As per coding guidelines, Python code must follow an 88-character line length limit.
Proposed line wrap
-def test_indexer_replica_group_initialization_gate_for_dcp8(indexer_shards, expected):
+def test_indexer_replica_group_initialization_gate_for_dcp8(
+ indexer_shards, expected
+):📝 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.
| def test_indexer_replica_group_initialization_gate_for_dcp8(indexer_shards, expected): | |
| def test_indexer_replica_group_initialization_gate_for_dcp8( | |
| indexer_shards, expected | |
| ): |
🤖 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 `@tests/distributed/test_indexer_parallel_groups.py` at line 62, Wrap the
signature of test_indexer_replica_group_initialization_gate_for_dcp8 across
multiple lines so no line exceeds the 88-character limit; keep the function
name, parameters, and behavior unchanged.
Source: Coding guidelines
6d6d71a to
9dbd6d0
Compare
a891f74
into
local-inference-lab:dev/gilded-gnosis
3 participants
Supersedes #244. This keeps Koush's functional commit unchanged and adds focused regression coverage. A replacement PR was necessary because GitHub denied direct maintainer writes to the fork head branch even though
maintainer_can_modifyis enabled.Purpose
VLLM_DCP_INDEXER_SHARDS=1fully replicates the sparse-indexer K cache on every DCP rank. The process-group initialization gate previously excluded this valid layout, so the dedicated singleton owner groups and DCP-wide query-split group were never created.That had two consequences:
VLLM_DCP_QUERY_SPLIT=1, lookup of the missing one-shard query-split group failed.The existing
_build_indexer_replica_group_ranksimplementation already produces the correct topology for one shard. This PR lets that path run for1 <= indexer_shards < DCP. It also extracts the gate into a pure helper so the boundary behavior is covered without initializing distributed backends.There is no behavior change for disabled replication (
0), existing partial layouts, orindexer_shards == DCP.Test plan
The parametrized regression test covers shard counts
0, 1, 2, 4, 8at DCP8. In particular, it fails if the one-shard boundary is changed back to1 < indexer_shards.The runtime path was also validated with a model-free TP4/DCP4 process-group test and an E2E GLM-5.2 prefill A-B-A comparison.
Test results
Unit and static checks
22 passedgit diff --checkpassedTP4/DCP4 topology
With
VLLM_DCP_INDEXER_SHARDS=1andVLLM_DCP_QUERY_SPLIT=1:indexer_query_split:0was created;E2E prefill A-B-A
Test profile: TP4/DCP4, MTP off,
MAX_NUM_SEQS=1, graph cap 6, CKV gather enabled, owner merge disabled, unique uncached prompts, three samples per point. The available GLM-5.2 EXL3 3.25bpw checkpoint was used because the validation host did not contain the release NVFP4 checkpoint. Absolute throughput is therefore not an NVFP4 release claim.The short-context difference is measurement noise. The 64K result was repeatably faster with the now-functional replicated query-split topology.
The root-port validation host required both variants to use the same Torch-NCCL fallback because its custom-AR/PyNCCL startup path has an unrelated environment-specific failure. This does not affect the process-group boundary fixed here and was held constant across the comparison.
Checklist
Summary by CodeRabbit