[GG] perf(indexer): shard DCP1 prefill queries across TP#176
[GG] perf(indexer): shard DCP1 prefill queries across TP#176voipmonitor wants to merge 2 commits into
Conversation
📝 WalkthroughWalkthroughQuery-split rank construction now uses a validated DCP layout helper. B12X sparse-attention prefill gathers top-k indices directly into shared storage through an alias-checked helper, without gathering scores. Tests cover rank layouts, validation errors, in-place gathering, and aliasing. ChangesQuery-split DCP flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant B12XPrefill
participant QuerySplitGroup
participant DCPGather
participant SharedIndexBuffer
B12XPrefill->>QuerySplitGroup: resolve query-split group
B12XPrefill->>DCPGather: gather top-k indices
DCPGather->>SharedIndexBuffer: write rank-local slices in place
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
|
Superseded by #175. During cleanup it became clear that DCP1 activation directly depends on the query-split rank-construction and exact index-gather contract introduced there. Keeping this as a second PR would either duplicate the same code or create a hidden stacked dependency. #175 now contains the cohesive implementation, current-GG rebase, review fixes, DCP1 qualification data, TP6 exclusions, and 27/27 targeted tests. |
![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/distributed/parallel_state.py (1)
1450-1472: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a
Raises:section to the new helper's docstring.
_get_query_split_group_ranksraisesValueErrorfor two distinct conditions (non-positivedcp_size, non-divisible TP group size) but the docstring is a bare one-liner with noArgs:/Raises:sections documenting them.📝 Proposed docstring
def _get_query_split_group_ranks( tp_group_ranks: list[list[int]], dcp_size: int ) -> list[list[int]]: - """Transpose DCP groups within each independent TP cohort.""" + """Transpose DCP groups within each independent TP cohort. + + Args: + tp_group_ranks: Ranks belonging to each independent TP group. + dcp_size: Number of DCP ranks within each TP group. + + Returns: + Query-split rank groups, one per DCP position across all TP cohorts. + + Raises: + ValueError: If `dcp_size` is not positive, or a TP group's size is + not divisible by `dcp_size`. + """As per coding guidelines, "Use Google-style docstrings in Python code, with
Args:/Returns:/Raises:sections instead of reStructuredText/Sphinx fields such as:param:,:return:, and:rtype:."🤖 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 - 1472, Update the docstring for _get_query_split_group_ranks with a Google-style Raises section documenting ValueError for non-positive dcp_size and for TP group sizes not divisible by dcp_size; preserve the existing summary and implementation.Source: Coding guidelines
vllm/model_executor/layers/sparse_attn_indexer.py (1)
375-404: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
Args:/Raises:sections to the docstring.
_query_split_all_gather_indicesraises three distinctRuntimeErrors (contiguity, row-count, aliasing) that aren't obvious from the one-line summary alone.📝 Proposed docstring
def _query_split_all_gather_indices( group, local_indices: torch.Tensor, gathered_indices: torch.Tensor, ) -> None: - """Restore query-split rows directly into their shared index buffer.""" + """Restore query-split rows directly into their shared index buffer. + + Args: + group: Query-split process group. + local_indices: This rank's contiguous slice of top-k indices; must + alias `gathered_indices` at this rank's offset. + gathered_indices: Shared buffer to gather all ranks' rows into. + + Raises: + RuntimeError: If the buffers are non-contiguous, have a mismatched + row count/trailing shape, or `local_indices` does not alias its + expected slot in `gathered_indices`. + """As per coding guidelines, "Use Google-style docstrings in Python code, with
Args:/Returns:/Raises:sections instead of reStructuredText/Sphinx fields such as:param:,:return:, and:rtype:."🤖 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 375 - 404, Expand the `_query_split_all_gather_indices` docstring with Google-style `Args:` entries for `group`, `local_indices`, and `gathered_indices`, and a `Raises:` section documenting the three `RuntimeError` conditions: non-contiguous buffers, invalid output shapes, and failure of `local_indices` to alias its rank-local output slice.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-1472: Update the docstring for _get_query_split_group_ranks
with a Google-style Raises section documenting ValueError for non-positive
dcp_size and for TP group sizes not divisible by dcp_size; preserve the existing
summary and implementation.
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 375-404: Expand the `_query_split_all_gather_indices` docstring
with Google-style `Args:` entries for `group`, `local_indices`, and
`gathered_indices`, and a `Raises:` section documenting the three `RuntimeError`
conditions: non-contiguous buffers, invalid output shapes, and failure of
`local_indices` to alias its rank-local output slice.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e3f49032-420e-493f-bf81-81fc7395a48f
📒 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
1 participant
Summary
Why this is exact
Each TP rank owns a disjoint query-row slice, computes the same sparse top-k operation it computed before, and all-gathers the resulting int32 indices. No score, index, CKV, or activation compression is introduced. A deterministic unique 64k prompt produced the same winning token, with a baseline-vs-query logprob difference of 7.15e-7 versus 9.54e-7 baseline repeat spread.
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 did not benefit: DCP1 regressed about 1.1%, DCP3 regressed 15.9% at 64k and 14.7% at 400k, and DCP6 was neutral at 400k. The image helper therefore keeps this opt-in path auto-disabled for all TP6 modes.
Tests
tests/distributed/test_query_split_groups.pytests/model_executor/layers/test_sparse_attn_indexer_b12x.pySummary by CodeRabbit
Bug Fixes
Tests