[GG] feat(dcp): support partial replicated sparse-indexer topology - #179
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds indexer-specific DCP/query-split groups, configurable DeepSeek indexer KV-cache sharding, CP-aware cache grouping and sizing, lockstep MLA COW handling, group-aware slot mapping, and expanded distributed and attention tests. ChangesDCP indexer cache and lockstep allocation
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant ModelRunner
participant DeepseekV32IndexerCache
participant SparseAttnIndexer
participant IndexerDCPGroup
ModelRunner->>DeepseekV32IndexerCache: derive KV shard layout
DeepseekV32IndexerCache->>SparseAttnIndexer: pass replication and shard metadata
SparseAttnIndexer->>IndexerDCPGroup: select indexer DCP group
SparseAttnIndexer->>IndexerDCPGroup: gather sparse top-k candidates
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 3
🧹 Nitpick comments (6)
vllm/distributed/parallel_state.py (1)
1451-1485: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd required Google-style docstrings to the new callables.
vllm/distributed/parallel_state.py#L1451-L1485: document accessor returns and_build_indexer_replica_group_ranksarguments, return groups, andValueError.vllm/model_executor/layers/sparse_attn_indexer.py#L2146-L2163: documentdcp_replicatedanddcp_kv_shard_countin the constructor API.tests/distributed/test_indexer_parallel_groups.py#L9-L43: add test docstrings.tests/model_executor/layers/test_sparse_attn_indexer_b12x.py#L39-L74: add a docstring describing replicated-constructor behavior.tests/model_executor/layers/test_sparse_attn_indexer_b12x.py#L452-L583: add a docstring describing replicated decode behavior.As per coding guidelines, Python code must use Google-style docstrings with
Args:,Returns:, andRaises:sections.🤖 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 1451 - 1485, The new callables and tests lack the required Google-style documentation. In vllm/distributed/parallel_state.py:1451-1485, document get_indexer_dcp_group and get_indexer_query_split_group return values, and document _build_indexer_replica_group_ranks arguments, returned group lists, and its ValueError condition using Args:, Returns:, and Raises:. In vllm/model_executor/layers/sparse_attn_indexer.py:2146-2163, document the constructor parameters dcp_replicated and dcp_kv_shard_count. Add descriptive docstrings to the tests in tests/distributed/test_indexer_parallel_groups.py:9-43, tests/model_executor/layers/test_sparse_attn_indexer_b12x.py:39-74 for replicated-constructor behavior, and tests/model_executor/layers/test_sparse_attn_indexer_b12x.py:452-583 for replicated decode behavior.Source: Coding guidelines
tests/models/test_dcp_shard_draft_defaults.py (1)
47-58: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_get_indexer_specduplicates the cache's own shard/replication derivation.
_get_indexer_specmanually recomputesdcp_replicated/dcp_kv_shard_count(lines 54-57) using the exact same formulas asDeepseekV32IndexerCache.__init__(deepseek_v2.py lines 702-708) instead of constructing a realDeepseekV32IndexerCacheand letting it self-populate these fields. If the production formula changes, this test helper could silently diverge and mask a regression.🤖 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/models/test_dcp_shard_draft_defaults.py` around lines 47 - 58, The test helper _get_indexer_spec should stop assigning dcp_shard_count, dcp_replicated, and dcp_kv_shard_count from duplicated formulas. Construct or obtain a real DeepseekV32IndexerCache through its normal initialization path, passing the layer and config inputs it requires, then use its self-populated fields when calling get_kv_cache_spec.vllm/model_executor/models/deepseek_v2.py (1)
675-685: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused DCP replication helper.
_replicate_indexer_cache_under_dcpis only referenced by the test import/calls, whileDeepseekV32IndexerCachecomputesdcp_replicateddirectly in__init__. Since there are no production call sites, this helper is redundant and can be removed.🤖 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/models/deepseek_v2.py` around lines 675 - 685, Remove the unused _replicate_indexer_cache_under_dcp helper, including its definition and any test imports or calls; preserve the existing direct dcp_replicated computation in DeepseekV32IndexerCache.__init__.vllm/v1/core/block_pool.py (1)
289-294: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value
assert get_group_id(...) != kv_cache_group_idturns an otherwise idempotent re-cache into a crash.
_insert_block_hashalready no-ops when the key matches, so a repeatedcache_full_blockscall for the same boundary from the same group is harmless today; this assert makes it fatal. Relaxing it to allow the same group id (while keeping the boundary/hash checks) keeps the lockstep invariant without adding a new failure mode.🤖 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/v1/core/block_pool.py` around lines 289 - 294, Update the existing-group assertion in _insert_block_hash to allow get_group_id(blk.block_hash) to equal kv_cache_group_id, while retaining the block token-count and hash equality assertions. Preserve idempotent repeated cache_full_blocks behavior without weakening those boundary/hash checks.vllm/v1/core/kv_cache_utils.py (2)
1339-1347: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer failing loudly over silently disabling lockstep on non-
intworld sizes.
isinstance(..., int)rejects e.g.numpy.int64, which would silently route allocation down the non-lockstep path instead of raising. Since the signature already types these asint, consider dropping the guard (or coercing withint(...)) so a bad caller surfaces immediately.♻️ Suggested simplification
- specs = list(kv_cache_specs) - if ( - len(specs) < 2 - or not isinstance(dcp_world_size, int) - or not isinstance(pcp_world_size, int) - ): - return False - cp_world_size = dcp_world_size * pcp_world_size + specs = list(kv_cache_specs) + if len(specs) < 2: + return False + cp_world_size = int(dcp_world_size) * int(pcp_world_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/v1/core/kv_cache_utils.py` around lines 1339 - 1347, Update the world-size validation in the lockstep decision logic around dcp_world_size and pcp_world_size to remove the isinstance(..., int) checks, or explicitly coerce the values with int(...). Preserve the existing len(specs) and cp_world_size <= 1 checks, while ensuring invalid caller values raise instead of silently returning False.
1860-1867: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftRe-implements CP-layout detection with raw
getattrinstead of the shared contract.This duplicates the intent of
has_nondefault_kv_cp_layout(already imported) and drifts from it:dcp_kv_shard_count == dcp * pcpis a default layout under the helper but "non-default" here. Consider plumbing the CP world sizes into this function and reusing the helper so page-size padding decisions stay consistent with grouping.🤖 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/v1/core/kv_cache_utils.py` around lines 1860 - 1867, Update the surrounding function to accept the CP world-size inputs required by has_nondefault_kv_cp_layout, then replace the local is_nondefault_cp_mla raw getattr-based detection with that shared helper for each relevant KV-cache spec. Preserve the helper’s default-layout semantics so page-size padding matches grouping behavior.
🤖 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/distributed/parallel_state.py`:
- Around line 2045-2065: Update the indexer shard validation before the branch
that initializes `_INDEXER_DCP` and `_INDEXER_QUERY_SPLIT`: for nonzero,
non-replicated shard counts, require a positive value that evenly divides
`decode_context_model_parallel_size`, and raise the existing topology error for
invalid values instead of falling back to the regular DCP group. Preserve the
current initialization path for valid shard counts.
In `@vllm/v1/core/kv_cache_coordinator.py`:
- Around line 303-322: The lockstep allocation path must mirror CoW replacements
rather than appending every returned block. Update the surrounding allocation
logic in the coordinator method using single_type_managers so each manager
applies CoW blocks at their replaced block_idx entries and only extends with
genuinely new blocks, preserving identical request block lengths and block IDs
across managers.
In `@vllm/v1/core/kv_cache_utils.py`:
- Around line 1727-1742: Update the non-MLA branch of the grouping key in the
custom-CP layout handling to include type(spec).__name__ alongside
spec.block_size, matching the type-aware MLA key. Keep the existing MLA key
fields and grouping behavior unchanged so different non-MLA spec types with the
same block size remain in separate buckets.
---
Nitpick comments:
In `@tests/models/test_dcp_shard_draft_defaults.py`:
- Around line 47-58: The test helper _get_indexer_spec should stop assigning
dcp_shard_count, dcp_replicated, and dcp_kv_shard_count from duplicated
formulas. Construct or obtain a real DeepseekV32IndexerCache through its normal
initialization path, passing the layer and config inputs it requires, then use
its self-populated fields when calling get_kv_cache_spec.
In `@vllm/distributed/parallel_state.py`:
- Around line 1451-1485: The new callables and tests lack the required
Google-style documentation. In vllm/distributed/parallel_state.py:1451-1485,
document get_indexer_dcp_group and get_indexer_query_split_group return values,
and document _build_indexer_replica_group_ranks arguments, returned group lists,
and its ValueError condition using Args:, Returns:, and Raises:. In
vllm/model_executor/layers/sparse_attn_indexer.py:2146-2163, document the
constructor parameters dcp_replicated and dcp_kv_shard_count. Add descriptive
docstrings to the tests in
tests/distributed/test_indexer_parallel_groups.py:9-43,
tests/model_executor/layers/test_sparse_attn_indexer_b12x.py:39-74 for
replicated-constructor behavior, and
tests/model_executor/layers/test_sparse_attn_indexer_b12x.py:452-583 for
replicated decode behavior.
In `@vllm/model_executor/models/deepseek_v2.py`:
- Around line 675-685: Remove the unused _replicate_indexer_cache_under_dcp
helper, including its definition and any test imports or calls; preserve the
existing direct dcp_replicated computation in DeepseekV32IndexerCache.__init__.
In `@vllm/v1/core/block_pool.py`:
- Around line 289-294: Update the existing-group assertion in _insert_block_hash
to allow get_group_id(blk.block_hash) to equal kv_cache_group_id, while
retaining the block token-count and hash equality assertions. Preserve
idempotent repeated cache_full_blocks behavior without weakening those
boundary/hash checks.
In `@vllm/v1/core/kv_cache_utils.py`:
- Around line 1339-1347: Update the world-size validation in the lockstep
decision logic around dcp_world_size and pcp_world_size to remove the
isinstance(..., int) checks, or explicitly coerce the values with int(...).
Preserve the existing len(specs) and cp_world_size <= 1 checks, while ensuring
invalid caller values raise instead of silently returning False.
- Around line 1860-1867: Update the surrounding function to accept the CP
world-size inputs required by has_nondefault_kv_cp_layout, then replace the
local is_nondefault_cp_mla raw getattr-based detection with that shared helper
for each relevant KV-cache spec. Preserve the helper’s default-layout semantics
so page-size padding matches grouping behavior.
🪄 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: 9c08b96b-9cc8-4b79-b0b7-a35a74191a61
📒 Files selected for processing (20)
tests/distributed/test_indexer_parallel_groups.pytests/model_executor/layers/test_sparse_attn_indexer_b12x.pytests/models/test_dcp_shard_draft_defaults.pytests/v1/attention/test_indexer_dcp_localize.pytests/v1/core/test_kv_cache_utils.pytests/v1/core/test_prefix_caching.pytests/v1/worker/test_attn_utils.pytests/v1/worker/test_gpu_block_table.pyvllm/distributed/parallel_state.pyvllm/envs.pyvllm/model_executor/layers/sparse_attn_indexer.pyvllm/model_executor/models/deepseek_v2.pyvllm/v1/attention/backends/mla/indexer.pyvllm/v1/core/block_pool.pyvllm/v1/core/kv_cache_coordinator.pyvllm/v1/core/kv_cache_utils.pyvllm/v1/core/single_type_kv_cache_manager.pyvllm/v1/kv_cache_interface.pyvllm/v1/worker/gpu/block_table.pyvllm/v1/worker/gpu/model_runner.py
|
Follow-up validation added in
No owner-merge or selected-record CKV code was added to this PR. |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
vllm/distributed/parallel_state.py (1)
1454-1460: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd Google-style API sections to both selector docstrings.
These public functions accept arguments, return coordinators, and can raise
RuntimeError, but their docstrings omitArgs:,Returns:, andRaises:.Proposed documentation update
def get_indexer_dcp_group( expected_world_size: int | None = None, ) -> GroupCoordinator: """Return the group matching an indexer's actual KV shard count. + Args: + expected_world_size: The indexer's resolved KV shard count. + + Returns: + The matching dedicated indexer or configured DCP group. + + Raises: + RuntimeError: If no group matches the requested shard count. + A partially replicated target indexer and a fully sharded speculative indexer can coexist in the same model. Callers that know their shard count must provide it so the speculative indexer does not inherit the target's smaller process group. """As per coding guidelines, Python docstrings must use Google-style
Args:,Returns:, andRaises:sections.Also applies to: 1485-1485
🤖 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 1454 - 1460, Update both public indexer-group selector docstrings to use Google-style sections: document every argument under Args:, describe the returned coordinator under Returns:, and state the RuntimeError conditions under Raises:. Keep the existing behavior and explanatory descriptions 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.
Nitpick comments:
In `@vllm/distributed/parallel_state.py`:
- Around line 1454-1460: Update both public indexer-group selector docstrings to
use Google-style sections: document every argument under Args:, describe the
returned coordinator under Returns:, and state the RuntimeError conditions under
Raises:. Keep the existing behavior and explanatory descriptions unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e4e74ba1-a8dc-44ab-9433-5b54e04c2ee9
📒 Files selected for processing (5)
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.pyvllm/v1/attention/backends/mla/indexer.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/model_executor/layers/test_sparse_attn_indexer_b12x.py
- vllm/v1/attention/backends/mla/indexer.py
- vllm/model_executor/layers/sparse_attn_indexer.py
|
Addressed the three actionable review findings in
Validation:
The broader |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
vllm/v1/core/single_type_kv_cache_manager.py (1)
390-414: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Args/Returns sections to the new CoW docstrings.
pending_cow_block_indexandapply_lockstep_cowtake parameters (request_id,block_idx,cow_block) but their docstrings omitArgs:/Returns:sections, unlike every other documented method in this class (add_local_computed_blocks,allocate_new_blocks,get_num_blocks_to_allocate, etc.).📝 Suggested docstring additions
def pending_cow_block_index(self, request_id: str) -> int | None: - """Return the pending partial-hit CoW index without consuming it.""" + """Return the pending partial-hit CoW index without consuming it. + + Args: + request_id: The request ID. + + Returns: + The pending block index, or ``None`` if there is no pending + partial-hit CoW for this request. + """"""Mirror a peer manager's CoW substitution using the same block ID. The primary manager owns the physical copy operation. A lockstep peer only transfers its request reference from the shared source block to the already allocated destination block. + + Args: + request_id: The request ID. + block_idx: The block index being substituted, must match this + manager's own pending CoW index for the request. + cow_block: The destination block allocated by the primary + manager for this lockstep group. """As per coding guidelines,
**/*.pycode should "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/v1/core/single_type_kv_cache_manager.py` around lines 390 - 414, Add Google-style Args and Returns sections to the docstrings of pending_cow_block_index and apply_lockstep_cow, documenting each parameter and the returned value where applicable. Keep the existing behavior and descriptions unchanged, and follow the class’s established docstring style.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/v1/core/single_type_kv_cache_manager.py`:
- Around line 390-414: Add Google-style Args and Returns sections to the
docstrings of pending_cow_block_index and apply_lockstep_cow, documenting each
parameter and the returned value where applicable. Keep the existing behavior
and descriptions unchanged, and follow the class’s established docstring style.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e90372d8-2ed9-4444-b861-3b2f7794d5cf
📒 Files selected for processing (8)
tests/distributed/test_indexer_parallel_groups.pytests/v1/attention/test_indexer_dcp_localize.pytests/v1/core/test_kv_cache_utils.pytests/v1/core/test_prefix_caching.pyvllm/distributed/parallel_state.pyvllm/v1/core/kv_cache_coordinator.pyvllm/v1/core/kv_cache_utils.pyvllm/v1/core/single_type_kv_cache_manager.py
🚧 Files skipped from review as they are similar to previous changes (4)
- vllm/distributed/parallel_state.py
- vllm/v1/core/kv_cache_coordinator.py
- tests/v1/attention/test_indexer_dcp_localize.py
- vllm/v1/core/kv_cache_utils.py
a9e5b90 to
673733e
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
vllm/v1/attention/backends/mla/indexer.py (1)
1191-1196: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDerive the CPU decode bound using the indexer shard group.
Line 1193 reuses
dcp_local_seq_lens_cpu, which belongs to the main CKV cache's configured DCP layout. With DCP4 and a partially replicated indexer using two shards, this can understate the indexer’s local maximum versusself.dcp_world_size/self.dcp_rank, causing B12X top-k to scan too small a KV window. Derive the CPU bound with the indexer group’s geometry (or safely useactive_width_tokens) and add a DCP4/indexer-shards=2 decode 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 `@vllm/v1/attention/backends/mla/indexer.py` around lines 1191 - 1196, Update the decode CPU bound in the surrounding indexer metadata construction to use the indexer shard group’s DCP geometry, based on its world size and rank, rather than slicing dcp_local_seq_lens_cpu from the main CKV layout; alternatively use active_width_tokens when it safely represents the bound. Preserve the None behavior when local DCP metadata is unavailable, and add coverage for DCP4 with two indexer shards verifying the B12X top-k scan uses the full local KV window.
🧹 Nitpick comments (3)
vllm/model_executor/layers/sparse_attn_indexer.py (2)
1170-1183: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
getattr-based probing forget_indexer_dcp_groupis dead defensiveness.The module already imports
get_indexer_dcp_groupunconditionally at line 14, so the selector can never beNoneand theget_dcp_group()fallback is unreachable.get_indexer_dcp_groupalso already fails closed on a shard-count mismatch, making the extra check redundant. Calling it directly keeps this helper consistent with the rest of the file. Same pattern at lines 1772-1777 for the query-split selector.♻️ Suggested simplification
- from vllm.distributed import parallel_state - - selector = getattr(parallel_state, "get_indexer_dcp_group", None) - group = ( - selector(expected_world_size) - if selector is not None - else parallel_state.get_dcp_group() - ) - if int(group.world_size) != int(expected_world_size): - raise RuntimeError( - "DCP owner top-k group does not match the indexer KV shard count: " - f"group={group.world_size}, shards={expected_world_size}" - ) - return group + return get_indexer_dcp_group(int(expected_world_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/model_executor/layers/sparse_attn_indexer.py` around lines 1170 - 1183, Update the DCP group helper to call the already imported get_indexer_dcp_group directly with expected_world_size, removing getattr probing, the get_dcp_group fallback, and the redundant world-size validation. Apply the same simplification to the query-split selector path around the corresponding helper.
446-478: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSilent
except Exception: passnow hides a real misconfiguration.With
expected_dcp_world_sizesupplied,get_indexer_dcp_group(group_world_size)raises when no group matches the requested shard count (fail-closed contract inparallel_state.get_indexer_dcp_group). Swallowing it here leavesdcp_world_size = expectedwithdcp_rank = 0on every rank, so warmup silently exercises the wrong rank geometry instead of surfacing the topology error. At minimum log at debug so it is diagnosable; ruff also flags both blind handlers (S110/BLE001).♻️ Suggested change
- except Exception: - pass + except Exception: + logger.debug("Failed to read DCP warmup config", exc_info=True) try: from vllm.distributed.parallel_state import get_indexer_dcp_group dcp_group = get_indexer_dcp_group(group_world_size) if int(dcp_group.world_size) > 1: dcp_world_size = int(dcp_group.world_size) dcp_rank = int(dcp_group.rank_in_group) - except Exception: - pass + except Exception: + logger.debug( + "No indexer DCP group for world size %s; using rank 0 for warmup", + group_world_size, + 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 446 - 478, Update _get_dcp_warmup_params to stop using blind except Exception: pass handlers. Preserve the fallback behavior only for genuinely unavailable optional configuration, but log caught exceptions at debug level so topology and configuration failures are diagnosable; ensure the get_indexer_dcp_group failure with an explicitly supplied expected_dcp_world_size is not silently hidden and satisfies Ruff’s S110/BLE001 checks.Source: Linters/SAST tools
tests/v1/core/test_prefix_caching.py (1)
180-217: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding a partial-replication (
dcp_kv_shard_count) lockstep coordinator test.All lockstep MLA tests in this file exercise full replication (
dcp_replicated=True) for the indexer group. The PR's headline feature — partial DCP replication viadcp_kv_shard_count(e.g. DCP8 with a 4-shard indexer) — is only exercised at the grouping-utility level intests/v1/core/test_kv_cache_utils.py(test_partial_replicated_mla_uses_lockstep_dcp8_layout), not through the actualKVCacheManager/coordinator allocation path (block-size equality assert, CoW mirroring, slot sharing). A variant ofmake_lockstep_mla_managerusingdcp_kv_shard_countinstead ofdcp_replicated=Truewould close this gap.🤖 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/v1/core/test_prefix_caching.py` around lines 180 - 217, Add a partial-replication lockstep coordinator test alongside make_lockstep_mla_manager, configuring the indexer MLAAttentionSpec with dcp_kv_shard_count=4 instead of dcp_replicated=True while retaining the DCP world size and lockstep block-size setup. Exercise the KVCacheManager allocation path, including block-size validation, copy-on-write mirroring, and shared slots, rather than relying only on grouping utilities.
🤖 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/v1/core/kv_cache_utils.py`:
- Around line 652-664: Update resolve_kv_cache_block_sizes to apply
get_kv_cache_cp_shard_count to wrapped attention specifications, including
UniformTypeKVCacheSpecs groups produced by lockstep_mla_layout, rather than
scaling only direct AttentionSpec instances. Use each spec’s dcp_replicated and
dcp_kv_shard_count properties and preserve raw block_size handling for unrelated
spec types.
- Around line 1687-1698: Update the `_is_lockstep_mla_spec_layout` input
filtering so it evaluates only MLA-family specifications, including
`MLAAttentionSpec`, `SlidingWindowMLASpec`, and other indexer/replicated-MLA
variants, while excluding unrelated KV-cache specs. Preserve the existing
lockstep layout logic and ensure `lockstep_mla_layout` is computed from this
filtered subset before `has_custom_cp` is evaluated.
---
Outside diff comments:
In `@vllm/v1/attention/backends/mla/indexer.py`:
- Around line 1191-1196: Update the decode CPU bound in the surrounding indexer
metadata construction to use the indexer shard group’s DCP geometry, based on
its world size and rank, rather than slicing dcp_local_seq_lens_cpu from the
main CKV layout; alternatively use active_width_tokens when it safely represents
the bound. Preserve the None behavior when local DCP metadata is unavailable,
and add coverage for DCP4 with two indexer shards verifying the B12X top-k scan
uses the full local KV window.
---
Nitpick comments:
In `@tests/v1/core/test_prefix_caching.py`:
- Around line 180-217: Add a partial-replication lockstep coordinator test
alongside make_lockstep_mla_manager, configuring the indexer MLAAttentionSpec
with dcp_kv_shard_count=4 instead of dcp_replicated=True while retaining the DCP
world size and lockstep block-size setup. Exercise the KVCacheManager allocation
path, including block-size validation, copy-on-write mirroring, and shared
slots, rather than relying only on grouping utilities.
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 1170-1183: Update the DCP group helper to call the already
imported get_indexer_dcp_group directly with expected_world_size, removing
getattr probing, the get_dcp_group fallback, and the redundant world-size
validation. Apply the same simplification to the query-split selector path
around the corresponding helper.
- Around line 446-478: Update _get_dcp_warmup_params to stop using blind except
Exception: pass handlers. Preserve the fallback behavior only for genuinely
unavailable optional configuration, but log caught exceptions at debug level so
topology and configuration failures are diagnosable; ensure the
get_indexer_dcp_group failure with an explicitly supplied
expected_dcp_world_size is not silently hidden and satisfies Ruff’s S110/BLE001
checks.
🪄 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: ff218887-3cf6-4d20-b4d4-e7e1913ae893
📒 Files selected for processing (18)
tests/distributed/test_indexer_parallel_groups.pytests/model_executor/layers/test_sparse_attn_indexer_b12x.pytests/models/test_dcp_shard_draft_defaults.pytests/v1/attention/test_indexer_dcp_localize.pytests/v1/core/test_kv_cache_utils.pytests/v1/core/test_prefix_caching.pytests/v1/worker/test_gpu_block_table.pyvllm/distributed/parallel_state.pyvllm/envs.pyvllm/model_executor/layers/sparse_attn_indexer.pyvllm/model_executor/models/deepseek_v2.pyvllm/v1/attention/backends/mla/indexer.pyvllm/v1/core/kv_cache_coordinator.pyvllm/v1/core/kv_cache_utils.pyvllm/v1/core/single_type_kv_cache_manager.pyvllm/v1/kv_cache_interface.pyvllm/v1/worker/gpu/block_table.pyvllm/v1/worker/gpu/model_runner.py
🚧 Files skipped from review as they are similar to previous changes (8)
- vllm/v1/worker/gpu/block_table.py
- tests/distributed/test_indexer_parallel_groups.py
- vllm/distributed/parallel_state.py
- vllm/v1/kv_cache_interface.py
- vllm/v1/core/single_type_kv_cache_manager.py
- tests/v1/attention/test_indexer_dcp_localize.py
- vllm/model_executor/models/deepseek_v2.py
- tests/models/test_dcp_shard_draft_defaults.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/v1/core/test_kv_cache_utils.py`:
- Around line 1801-1805: Update the test around
_get_kv_cache_groups_uniform_groups to assert that grouped_specs contains both
target and indexer groups before resolving or validating the result, ensuring
the multi-group LCM/GCD path is exercised rather than the single-group fast
path.
🪄 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: ed2d9eb1-92f8-46c3-8289-e2b7f48bbd85
📒 Files selected for processing (2)
tests/v1/core/test_kv_cache_utils.pyvllm/v1/core/kv_cache_utils.py
🚧 Files skipped from review as they are similar to previous changes (1)
- vllm/v1/core/kv_cache_utils.py
7847562
into
local-inference-lab:dev/gilded-gnosis
2 participants
Summary
Add an opt-in partial-replication topology for the GLM sparse-indexer cache.
This PR now builds directly on
dev/gilded-gnosisafter #159 was merged as4247d676; it no longer contains or duplicates any #159 commits.Design
VLLM_DCP_INDEXER_SHARDS=0keeps the indexer sharded by the configured DCP.VLLM_DCP_INDEXER_SHARDS=1fully replicates the indexer cache.4creates
{0,1,2,3}and{4,5,6,7}indexer groups with cross-replica querygroups
{0,4},{1,5},{2,6}, and{3,7}.physical page geometry, block tables, and slot mapping.
The legacy
VLLM_DCP_REPLICATE_INDEXER_CACHE=1switch remains backwardcompatible and maps to one indexer shard. All new behavior defaults off.
Measured impact
GLM-5.2 NVFP4 A16, TP8, MTP0, 8x RTX PRO 6000 Blackwell. Baseline and candidate
used the same image, helper, GPU group, and serial benchmark protocol. The
measured image also contained the independent prefill PRs #177 and #178; this
PR does not include their code.
DCP4 partial improves CC1 decode by 6.25-9.10% and prefill by 0.64-1.1% for a
5.0% KV-capacity cost. DCP8 partial improves prefill by 68.3-76.0%, keeps long
decode effectively in parity, and also costs about 5.0% KV capacity.
Validation
4247d676, with exactly five [GG] feat(dcp): support partial replicated sparse-indexer topology #179-specific commits.git range-diffconfirms semantic parity with the original unique commitrange; the only behavior adjustment keeps feat(kv-cache): support mixed MLA groups and native offload #159's supported external-load
contract instead of the obsolete rejection test.
CUDA platform detection unavailable in the CPU-only test container.
64k/400k prefill, and CC1 decode validation remains applicable because the
production diff is semantically unchanged.
The lockstep allocator and full-replication foundation are provided by merged
#159. This PR's review unit is the partial topology, dedicated process groups,
generalized cache geometry, query-split composition, and their tests.
Summary by CodeRabbit
New Features
VLLM_DCP_INDEXER_SHARDSto control sparse-indexer KV sharding/replication within DCP groups.Bug Fixes
Tests