[GG] perf(dcp): merge sparse prefill top-k by row owner#178
[GG] perf(dcp): merge sparse prefill top-k by row owner#178voipmonitor wants to merge 3 commits into
Conversation
|
Warning Review limit reached
Next review available in: 2 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. 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)
📝 WalkthroughWalkthroughChangesB12X DCP owner merge
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Prefill as B12X prefill
participant Merge as _merge_b12x_prefill_dcp_topk
participant Owner as _merge_b12x_dcp_topk_by_owner
participant DCP as DCP collectives
participant TP as TP all-gather
Prefill->>Merge: merge DCP top-k candidates
Merge->>Owner: invoke when flag is enabled
Owner->>DCP: route candidates to owners
Owner->>TP: gather owner results
Merge-->>Prefill: return indices and owner-merge status
Possibly related issues
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 |
|
Follow-up fixes the integration contract with partial indexer replication (#179). Owner merge now selects its collective by the indexer KV shard count instead of always using the configured full DCP group, and validates the selected world size. This fixes TP8/DCP8 with a 4-shard partial indexer, which previously failed CUDA-graph warmup with |
![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 (3)
tests/model_executor/layers/test_sparse_attn_indexer_b12x.py (2)
899-917: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider covering the fail-closed branches too.
The
RuntimeErrorpaths (TP/DCP rank-layout disagreement, invalid trailing shape, non-aliased local rows) are the ones most likely to fire on a misconfigured deployment and currently have no test.🤖 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/model_executor/layers/test_sparse_attn_indexer_b12x.py` around lines 899 - 917, Extend test_b12x_dcp_owner_merge_falls_back_for_tail to cover the fail-closed RuntimeError branches of _merge_b12x_dcp_topk_by_owner: TP/DCP rank-layout disagreement, invalid trailing shape, and non-aliased local rows. Assert each misconfigured input raises RuntimeError while preserving the existing tail-fallback assertion.
944-948: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
monkeypatch.setattron the lazyenvsmodule leaves a concrete attribute behind.
vllm.envsresolves flags through module__getattr__; patching then restoring plants a real module attribute that permanently shadows the lazy getter for the rest of the session, so later tests readingVLLM_DCP_TOPK_OWNER_MERGEno longer see env changes. Patching the attribute on the consuming module (indexer_mod) is not possible here since it readsenvs.X, so prefermonkeypatch.setenv("VLLM_DCP_TOPK_OWNER_MERGE", ...)plusenvs.disable_envs_cache(), or patchindexer_mod.envswith a lightweight stub.🤖 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/model_executor/layers/test_sparse_attn_indexer_b12x.py` around lines 944 - 948, Replace the monkeypatch of indexer_mod.envs.VLLM_DCP_TOPK_OWNER_MERGE with environment-variable patching via monkeypatch.setenv, using the expected enabled/disabled value, and call envs.disable_envs_cache() so lazy flag resolution reflects each test’s environment. Ensure cleanup restores the environment and cache state without leaving a concrete attribute on the lazy envs module.vllm/model_executor/layers/sparse_attn_indexer.py (1)
388-393: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer an explicit injection point over duck-typed
getattr(group, "all_to_all_single").If a real
GroupCoordinatorever gains anall_to_all_singleattribute (with different semantics/signature), this silently routes production traffic through it instead of the NCCL path. An explicit keyword arg (e.g.all_to_all_fn=None) or a dedicated fake type check makes the test seam unambiguous.🤖 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 388 - 393, Replace the duck-typed getattr(group, "all_to_all_single") test seam in the all-to-all helper with an explicit injected hook, such as an all_to_all_fn keyword argument defaulting to None. Invoke the injected function only when provided, otherwise preserve the existing GroupCoordinator NCCL path; do not infer test behavior from arbitrary group attributes.
🤖 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/model_executor/layers/sparse_attn_indexer.py`:
- Around line 1279-1296: Update the DCP query-split prefill path so merged
scores from _merge_b12x_dcp_topk_by_owner() are copied from owner_values into
the corresponding owner rows of topk_scores_buffer. Include the used_owner_merge
path in the query-split all-gather and write-back logic, keeping gathered
indices and scores synchronized for every query partition.
---
Nitpick comments:
In `@tests/model_executor/layers/test_sparse_attn_indexer_b12x.py`:
- Around line 899-917: Extend test_b12x_dcp_owner_merge_falls_back_for_tail to
cover the fail-closed RuntimeError branches of _merge_b12x_dcp_topk_by_owner:
TP/DCP rank-layout disagreement, invalid trailing shape, and non-aliased local
rows. Assert each misconfigured input raises RuntimeError while preserving the
existing tail-fallback assertion.
- Around line 944-948: Replace the monkeypatch of
indexer_mod.envs.VLLM_DCP_TOPK_OWNER_MERGE with environment-variable patching
via monkeypatch.setenv, using the expected enabled/disabled value, and call
envs.disable_envs_cache() so lazy flag resolution reflects each test’s
environment. Ensure cleanup restores the environment and cache state without
leaving a concrete attribute on the lazy envs module.
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 388-393: Replace the duck-typed getattr(group,
"all_to_all_single") test seam in the all-to-all helper with an explicit
injected hook, such as an all_to_all_fn keyword argument defaulting to None.
Invoke the injected function only when provided, otherwise preserve the existing
GroupCoordinator NCCL path; do not infer test behavior from arbitrary group
attributes.
🪄 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: 5d8a5093-ad69-4d96-a0e0-f5366328e42a
📒 Files selected for processing (3)
tests/model_executor/layers/test_sparse_attn_indexer_b12x.pyvllm/envs.pyvllm/model_executor/layers/sparse_attn_indexer.py
1 participant
Summary
Add an opt-in exact 2D owner merge for sharded sparse-indexer top-k during DCP prefill.
Instead of replicating every rank's candidate rows to every other rank, the path:
VLLM_DCP_TOPK_OWNER_MERGE=0remains the default. Unsupported shapes and tail rows return to the existing replicated NCCL oracle.Why
Query split reduces indexer compute, but the existing global top-k merge still replicates all candidate scores and indices. The owner layout preserves sharded scoring and exact FP32 selection while reducing candidate workspace and communication.
Results
GLM-5.2 NVFP4, 8x RTX PRO 6000 Blackwell, TP8/DCP4/MTP0/A16, query split plus full-CKV gather:
At the 131k planning profile, KV capacity increased from 2,140,416 to 2,186,496 tokens (+2.15%) because the owner path reserves a smaller candidate-merge workspace.
Correctness and fallback
Validation
22 passedintests/model_executor/layers/test_sparse_attn_indexer_b12x.pyagainst the exact installed overlay.ruff checkandruff format --checkpass for all modified files.The separate SparkInfer candidate transport is intentionally not part of this PR. This change uses standard vLLM collectives and is independently useful.
Summary by CodeRabbit
New Features
Tests