[GG] perf(indexer): shard sparse prefill queries and halve result traffic#175
Conversation
|
Warning Review limit reached
Next review available in: 3 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 (4)
📝 WalkthroughWalkthroughChangesThe PR centralizes query-split rank-group construction with validation, adds in-place index all-gather support for B12X sparse prefill, removes score gathering in that path, and adds tests for rank layouts, buffer aliasing, and in-place writes. Query-split integration
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant B12XPrefill
participant QuerySplitGroup
participant GatheredIndices
B12XPrefill->>QuerySplitGroup: resolve query-split group
B12XPrefill->>GatheredIndices: gather top-k indices in place
QuerySplitGroup-->>GatheredIndices: provide rank-local slices
GatheredIndices-->>B12XPrefill: restored gathered indices
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.
🧹 Nitpick comments (2)
vllm/model_executor/layers/sparse_attn_indexer.py (1)
399-402: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCover the non-PyNCCL aliasing path.
The added test exercises only
pynccl_comm.all_gather; the fallback passes an aliased rank-local view totorch.distributed.all_gather_into_tensor. Add backend-level coverage or verify the deployed PyTorch/NCCL combination supports this exact in-place layout before relying on it. NCCL documents the rank-local condition assendbuff == recvbuff + rank * sendcount, while PyTorch’s API documentation specifies sizing but not this aliasing guarantee. (github.com)🤖 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 399 - 402, Update _dcp_all_gather_first_dim_into to cover or validate the aliased rank-local send-buffer path used by torch.distributed.all_gather_into_tensor, not only pynccl_comm.all_gather. Add backend-level coverage for the exact in-place layout, or establish an explicit PyTorch/NCCL capability check and provide a safe non-aliased fallback when unsupported.vllm/distributed/parallel_state.py (1)
1450-1453: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument both new helpers with Google-style sections.
Both helpers accept arguments and raise validation errors, but their docstrings contain only a summary.
vllm/distributed/parallel_state.py#L1450-L1453: addArgs:,Returns:, andRaises:sections describing the TP cohorts, DCP size, output groups, andValueError.vllm/model_executor/layers/sparse_attn_indexer.py#L375-L380: addArgs:,Returns:, andRaises:sections describing the group, buffers, and validation failures.🤖 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 1450 - 1453, Document both helpers with Google-style Args, Returns, and Raises sections: in vllm/distributed/parallel_state.py lines 1450-1453, update _get_query_split_group_ranks to describe TP cohorts, dcp_size, transposed output groups, and ValueError validation; in vllm/model_executor/layers/sparse_attn_indexer.py lines 375-380, update the corresponding helper’s docstring to describe its group and buffers, return value, and validation failures.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.
Nitpick comments:
In `@vllm/distributed/parallel_state.py`:
- Around line 1450-1453: Document both helpers with Google-style Args, Returns,
and Raises sections: in vllm/distributed/parallel_state.py lines 1450-1453,
update _get_query_split_group_ranks to describe TP cohorts, dcp_size, transposed
output groups, and ValueError validation; in
vllm/model_executor/layers/sparse_attn_indexer.py lines 375-380, update the
corresponding helper’s docstring to describe its group and buffers, return
value, and validation failures.
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 399-402: Update _dcp_all_gather_first_dim_into to cover or
validate the aliased rank-local send-buffer path used by
torch.distributed.all_gather_into_tensor, not only pynccl_comm.all_gather. Add
backend-level coverage for the exact in-place layout, or establish an explicit
PyTorch/NCCL capability check and provide a safe non-aliased fallback when
unsupported.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 0d9320bf-f356-453b-9328-050d15b6c3fc
📒 Files selected for processing (4)
tests/distributed/test_query_split_groups.pytests/model_executor/layers/test_sparse_attn_indexer_b12x.pyvllm/distributed/parallel_state.pyvllm/model_executor/layers/sparse_attn_indexer.py
6cf014d to
f8c0cb4
Compare
f8c0cb4 to
599d3f6
Compare
voipmonitor
changed the title
|
Addressed the review findings while rebasing onto current
Validation: 27/27 targeted tests, Ruff, and |
1 participant
Summary
The feature remains gated by
VLLM_DCP_QUERY_SPLIT. This PR does not globally enable it; the measured launcher policy enables only qualified TP4/TP8 topologies and retains an explicit kill switch.Why this is exact
Each query-split rank computes a disjoint row slice with the existing sparse top-k kernel, then all-gathers the resulting int32 indices. It does not compress scores, indices, CKV, weights, or activations. A deterministic unique 64k prompt produced the same winning token; baseline-vs-query logprob delta was 7.15e-7 versus 9.54e-7 baseline repeat spread.
Communication impact
For TP8/DCP4, 8192 rows, top-k 2048, and query-split size 2, eliminating the fp32 score collective reduces result traffic per GPU/layer/direction from 64 MiB to 32 MiB. The isolated collective median fell from 2.726 ms to 1.412 ms per active layer and 8k chunk.
E2E results
GLM-5.2, MTP off, exact raw prefill, no concurrent model loading or benchmark traffic:
Virtual TP6 has 11 padded heads per rank and is intentionally excluded from auto-enable:
Validation
git diff --checkpass