[GG] perf(indexer): shard sparse prefill queries and halve result traffic - #175
Conversation
📝 WalkthroughWalkthroughChangesThe PR adds indexer-specific DCP and query-split group topology, introduces alias-safe in-place index gathering, updates B12X prefill integration, and adds coverage for group layouts and collective buffer contracts. Indexer DCP and query-split integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant initialize_model_parallel
participant _build_indexer_replica_group_ranks
participant _query_split_all_gather_indices
participant _dcp_all_gather_first_dim_into
participant B12X_prefill
initialize_model_parallel->>_build_indexer_replica_group_ranks: build DCP and query-split rank groups
_build_indexer_replica_group_ranks-->>initialize_model_parallel: return rank layouts
B12X_prefill->>_query_split_all_gather_indices: submit local and gathered index buffers
_query_split_all_gather_indices->>_dcp_all_gather_first_dim_into: gather aliased indices
_dcp_all_gather_first_dim_into-->>B12X_prefill: restore gathered indices
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 |
![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
f8c0cb4 to
599d3f6
Compare
voipmonitor
changed the title
|
Addressed the review findings while rebasing onto current
Validation: 27/27 targeted tests, Ruff, and |
2c14a58 to
a563e57
Compare
|
Refreshed this PR onto current |
443e1d6
into
local-inference-lab:dev/gilded-gnosis
There was a problem hiding this comment.
🧹 Nitpick comments (1)
vllm/model_executor/layers/sparse_attn_indexer.py (1)
504-536: 🩺 Stability & Availability | 🔵 TrivialSilent exception swallowing flagged by static analysis.
Both
try/except Exception: passblocks here (including the one spanning the changed lines 511-524) silently discard any failure resolvingvllm_config/parallel_configor the indexer DCP group, which could mask real misconfiguration during warmup instead of surfacing it. Ruff flags this pattern (S110/BLE001).♻️ Suggested fix: log instead of silently swallowing
- except Exception: - pass + except Exception: + logger.debug("Failed to resolve DCP warmup config", exc_info=True)🤖 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 504 - 536, Replace both silent `except Exception: pass` handlers in `_get_dcp_warmup_params` with appropriate warning or debug logging that records the failure and preserves the existing fallback values. Keep the warmup parameter resolution flow unchanged while ensuring failures from `get_current_vllm_config_or_none`/`parallel_config` and `get_indexer_dcp_group` are no longer discarded and satisfy the static-analysis rule.Source: Linters/SAST tools
🤖 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/model_executor/layers/sparse_attn_indexer.py`:
- Around line 504-536: Replace both silent `except Exception: pass` handlers in
`_get_dcp_warmup_params` with appropriate warning or debug logging that records
the failure and preserves the existing fallback values. Keep the warmup
parameter resolution flow unchanged while ensuring failures from
`get_current_vllm_config_or_none`/`parallel_config` and `get_indexer_dcp_group`
are no longer discarded and satisfy the static-analysis rule.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 805274a8-6a23-4976-8166-029dc79583d8
📒 Files selected for processing (4)
tests/distributed/test_indexer_parallel_groups.pytests/model_executor/layers/test_sparse_attn_indexer_b12x.pyvllm/distributed/parallel_state.pyvllm/model_executor/layers/sparse_attn_indexer.py
2 participants
Summary
int32top-k indices directly into the caller-owned buffer;torch.distributedbackends a safe copied send slice.The feature remains gated by
VLLM_DCP_QUERY_SPLIT. This PR does not globally enable it; deployment policy retains an explicit kill switch and excludes unqualified virtual-TP6 layouts.Current GG integration
The PR was reconstructed on 2026-07-30 from current
dev/gilded-gnosis@2d9c9801094842960b94c053539d62892e35a0bc.The earlier branch carried its own query-split topology helper. Current GG now has canonical full and partially replicated indexer topology through
_build_indexer_replica_group_ranks,get_indexer_dcp_group, andget_indexer_query_split_group. This refresh removes the duplicate implementation and uses that canonical builder for DCP1, partial DCP, and full DCP.The remaining diff is therefore limited to the still-missing communication optimization, DCP1 group initialization, and focused regression coverage. It composes with current GG's partial-indexer groups and context-crossover policy instead of replacing either one.
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
int32indices. 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 was7.15e-7versus9.54e-7baseline 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.
Qualified E2E results
GLM-5.2, MTP off, exact raw prefill, no concurrent model loading or benchmark traffic:
The immutable v20 r14 image contains the earlier, E2E-qualified implementation at
2c14a58a21c724f08d01d988b9972d10c77ed4c0. This refresh preserves that transport behavior while replacing obsolete topology code with current GG's implementation.Virtual TP6 has 11 padded heads per rank and remains excluded from auto-enable:
Validation
Current-GG refresh:
47/47focused tests pass;git diff --checkpasses;Earlier E2E qualification retained by the unchanged transport core:
Summary by CodeRabbit
New Features
Bug Fixes
Tests