Fix GLM-5.2 DCP MTP metadata, prefix cache, and global top-k - #30
Closed
voipmonitor wants to merge 99 commits into
Closed
Fix GLM-5.2 DCP MTP metadata, prefix cache, and global top-k#30voipmonitor wants to merge 99 commits into
voipmonitor wants to merge 99 commits into
Conversation
Tests: .venv/bin/python -m pytest tests/kernels/attention/test_cp_lse_ag_rs.py -v
Tests: .venv/bin/python -m pytest tests/config/test_virtual_tp.py tests/engine/test_arg_utils.py::test_virtual_tp_sharding_cli_alias -v
Tests: .venv/bin/python -m ruff check vllm/compilation/b12x_capture.py vllm/compilation/cuda_graph.py vllm/v1/worker/gpu/cudagraph_utils.py vllm/envs.py
Tests: .venv/bin/python -m ruff check vllm/config/kernel.py
Tests: .venv/bin/python -m ruff check vllm/model_executor/kernels/linear/__init__.py vllm/model_executor/kernels/linear/scaled_mm/__init__.py vllm/model_executor/kernels/linear/scaled_mm/b12x.py vllm/model_executor/layers/quantization/fp8.py vllm/model_executor/warmup/deep_gemm_warmup.py
Filter safetensors shards and tensors by checkpoint weight prefixes so draft model loading uses the index instead of scanning unrelated shards. Warm up DeepSeek V4 compressor Triton signatures before KV cache allocation and align sparse MLA decode metadata with speculative rows under DCP.
Route B12X sparse-indexer prefill through the paged compressed indexer with row-shared page tables, and size prefill chunks against the B12X scorer supertile so workspace reservations match runtime.
Pass the active request count into autoregressive draft sampling so padded FULL-cudagraph rows cannot overwrite draft logits for real requests.
GLM/Kimi DCP decode uses an uncompressed MLA indexer cache (compress_ratio == 1), but the sparse indexer still has to write and read through DCP rank-local cache pages. The previous path only remapped slots when compress_ratio > 1, so DCP4 wrote the indexer K cache with global slot ids while B12X attention consumed it through a DCP-local page table. Use DCP-local slot mapping whenever decode-context parallelism is active, and feed DCP-local seq_lens to the decode sparse-indexer metadata. For MTP/native expanded decode rows, compute the per-token lengths from global seq_lens first, then convert the expanded lengths to DCP-local lengths so query_start_loc arithmetic remains correct. Validated on GLM-5.1 NVFP4 DCP4 nomtp with B12X_MLA_SPARSE, V2 model runner, FULL+PIECEWISE CUDA graphs, 50k context smoke, and cc1 decode bench.
Prepare forced A8 and A16 FP4 MoE modes during vLLM weight loading so source tensors can be compacted after derived weights are built. Keep prepared W4A8 metadata through warmup and apply, and allow SwiGLU clamp parameters on A8. Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Chthonic's native B12X NVFP4 prepare path is correct for normal NVFP4, but forced W4A16 reuses the prepared ModelOpt weights through B12xExperts. That path needs the FI/B12X gated W13 reorder before b12x prepares the W4A16 weights; otherwise the gate/up halves are consumed in the wrong order and GLM emits corrupted output. Keep native B12X NVFP4 on prepare_nvfp4_moe_layer_for_b12x, switch to the FI/B12X prepare only when use_a16 or B12X_MOE_FORCE_A16=1 is active, and propagate ModelOpt's explicit use_a16 flag into the converter. Validated on the chthonic 5e83948 / b12x 465cb6e image with GLM v10 DCP1 MTP-off: B12X_MOE_FORCE_A16=1 on port 5329 and B12X_MOE_FORCE_A16=0 on port 5330 both start with B12X MLA sparse + B12X NvFp4 MoE and run /mnt/test.py -L coherently until timeout with CJK count 0. Also ran git diff --check and py_compile for the touched files.
Track vLLM-owned W4A16 MoE activation amax tensors and pass them into the B12X binding without using a workspace-owned buffer. Save activation-amax payloads from GPU runner post-forward hooks, add focused warmup tests, accept fp8_e4m3 for B12X MLA sparse, and add a GLM 5.2 serve helper. Co-authored-by: OpenAI Codex <codex@openai.com>
…ect#45895) Squashed from vllm-project#45895. Original commits: - c344fe2 add mtp_recycle_post_norm for MTP regression - 7d24aa6 skip indexer init when skip_topk - 4241fd6 make post norm recycle default for deepseek mtp - 8d4d038 fix ci Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
With --decode-context-parallel-size > 1, GLM-5.2 (DeepSeek-V4 sparse MLA + MTP spec decode) reported a permanent 0% prefix cache hit rate: HybridKVCacheCoordinator force-disabled prefix caching whenever a KV-cache group was dcp_replicated, and the MTP draft group is marked dcp_replicated (it replicates draft KV across DCP ranks). DCP-shard the MTP draft instead of replicating it, behind the opt-in VLLM_DCP_SHARD_DRAFT env flag (default off = original replication): - eagle/utils.py: build the draft with the parent's decode_context_parallel_size when sharding, so B12xMLASparseImpl.__init__ sizes its caller-owned-scratch plan for the DCP head all-gather and takes the cross-rank LSE-reduce path (the core fix). - mla_attention.py, deepseek_v2.py: gate the draft's dcp_replicated flag behind VLLM_DCP_SHARD_DRAFT so the draft KV is sharded like the target. - kv_cache_coordinator.py: only disable DCP prefix caching for the genuine DeepSeek-V4 MLA/SWA hybrid; assert against the DCP-scaled manager block size; relax the dcp==1 assert for non-hybrid layouts. - kv_cache_utils.py: use the GCD of effective block sizes for hash_block_size (scheduler size stays the LCM). - worker/utils.py: KVBlockZeroer supports non-uniform page sizes (one zeroing kernel per distinct page size). AI assistance (Claude Code) was used for this change. Signed-off-by: Florian Bernd <git@flobernd.de> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GGuk9dcSi3b1pk5JyN4Lxj
📝 WalkthroughWalkthroughThis PR adds ChangesDCP Draft Sharding for Speculative MLA
Sequence Diagram(s)sequenceDiagram
participant GPUModelRunner
participant AutoRegressiveSpeculator
participant PrefillCudaGraphManager
participant DraftModelSpeculator
participant sparse_attn_indexer
participant _dcp_global_topk_remap
rect rgba(70, 130, 180, 0.5)
Note over GPUModelRunner,PrefillCudaGraphManager: CUDA Graph Capture
GPUModelRunner->>AutoRegressiveSpeculator: capture(attn_states)
AutoRegressiveSpeculator->>PrefillCudaGraphManager: capture(attn_states=None, model_state, block_tables, ...) [if rebuild_prefill_attn_metadata]
PrefillCudaGraphManager->>PrefillCudaGraphManager: prepare_inputs_to_capture() -> attn_state
GPUModelRunner->>GPUModelRunner: _zero_cudagraph_capture_kv_blocks() -> zero block[0]
end
rect rgba(60, 160, 80, 0.5)
Note over DraftModelSpeculator,sparse_attn_indexer: Draft Propose (DCP active)
DraftModelSpeculator->>DraftModelSpeculator: set_attn() -> rebuild_prefill_attn_metadata=True
DraftModelSpeculator->>DraftModelSpeculator: _build_draft_attn_metadata(query_start_loc_cpu)
DraftModelSpeculator->>sparse_attn_indexer: _prefill(draft attn metadata, slot_mappings)
sparse_attn_indexer->>_dcp_global_topk_remap: topk_indices, logits, row_starts
_dcp_global_topk_remap->>_dcp_global_topk_remap: pack + all_gather across DCP group
_dcp_global_topk_remap-->>sparse_attn_indexer: remapped topk_indices (non-owned = -1)
end
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
vllm/v1/core/kv_cache_utils.py (1)
631-644:⚠️ Potential issue | 🟠 MajorPreserve
dcp_replicatedwhen reconstructing MLAAttentionSpec at line 1435.At line 1435 in
vllm/v1/core/kv_cache_utils.py, theMLAAttentionSpecreconstruction converts aSlidingWindowMLASpecwithout explicitly preservingdcp_replicated. WhileSlidingWindowMLASpecdoes not currently definedcp_replicated(it inherits fromSlidingWindowSpec, notFullAttentionSpec), this omission creates a data integrity risk: if the source spec ever gains this field or if future conversions involve specs withdcp_replicated=True, the field will silently default toFalse. This could cause the GCD hash layout computation at lines 638–644 to treat a replicated draft group as DCP-sharded, producing incorrect effective block sizes.Proposed fix
MLAAttentionSpec( block_size=uniform_block_size if uniform_block_size is not None else spec.block_size, num_kv_heads=spec.num_kv_heads, head_size=spec.head_size, dtype=spec.dtype, page_size_padded=spec.page_size_padded, cache_dtype_str=spec.cache_dtype_str, alignment=spec.alignment, compress_ratio=spec.compress_ratio, model_version=spec.model_version, + dcp_replicated=getattr(spec, "dcp_replicated", False), )🤖 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 631 - 644, At line 1435 where MLAAttentionSpec is being reconstructed from a SlidingWindowMLASpec, explicitly preserve the dcp_replicated attribute by passing it to the MLAAttentionSpec constructor or copying it after construction. This ensures that if dcp_replicated is present on the source spec, it is properly carried forward to the reconstructed spec and won't silently default to False, which could cause the GCD hash layout computation in the effective_block_sizes calculation to treat replicated draft groups as DCP-sharded and produce incorrect block sizes.vllm/v1/attention/backends/mla/indexer.py (2)
927-939:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPublish the effective DCP layout in returned metadata.
When
common_attn_metadata.dcp_local_seq_lensis absent, Lines 657-662 intentionally treat the KV as non-local (effective_dcp_world_size=1), but Lines 937-938 still advertiseself.dcp_world_size. Downstream global top-k remap uses this field to assume rank-local shards, which can turn valid replicated/full-KV indices into-1.Proposed fix
prefill=prefill_metadata, decode=decode_metadata, - dcp_world_size=self.dcp_world_size, - dcp_rank=self.dcp_rank, + dcp_world_size=effective_dcp_world_size, + dcp_rank=effective_dcp_rank, cp_interleave_size=self.cp_kv_cache_interleave_size,If callers still need the physical DCP size, add a separate explicit field instead of overloading the sharded-KV flag.
🤖 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 927 - 939, The DeepseekV32IndexerMetadata constructor at lines 927-939 is setting dcp_world_size to self.dcp_world_size, but this does not reflect the effective DCP layout that was intentionally set to 1 (non-local KV) in lines 657-662 when dcp_local_seq_lens is absent. This causes downstream issues in global top-k remap which expects the advertised dcp_world_size to match the actual KV sharding strategy. Replace the dcp_world_size assignment in the DeepseekV32IndexerMetadata constructor to use the effective_dcp_world_size value (which reflects whether KV is treated as replicated/non-local with size 1, or sharded with the physical size) instead of self.dcp_world_size. If the physical DCP size is still needed by callers, add a separate explicit field to the metadata rather than overloading the sharded-KV indicator field.
713-775:⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy liftKeep DCP prefill chunking rank-invariant before the global top-k collective.
Line 713 makes
compressed_seq_lens_cpu_nprank-local, and Line 745 uses it to buildchunk_specs. With DCP global top-k,sparse_attn_indexer.pyall-gathers once per chunk; if ranks split or skip chunks differently because local KV lengths differ, the collective can see different row shapes/order and hang. Build chunk boundaries from a rank-invariant global/max compressed length when global top-k is active, while still using rank-local lengths for the actual KV gather and emitting empty local chunks so every rank participates.🤖 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 713 - 775, When DCP global top-k is active, the prefill chunk boundaries must be rank-invariant to prevent collective operations from hanging due to mismatched shapes. Currently, chunk_specs is built from rank-local compressed_seq_lens_cpu_np on line 745 in the split_indexer_prefill_chunks calls, causing different ranks to create different chunk boundaries. Compute a rank-invariant global compressed length (e.g., through collective communication to get max or consensus values) and use it to build chunk_specs instead of the rank-local compressed_seq_lens_cpu_np. Keep the rank-local compressed_seq_lens_cpu_np values for the actual KV gather operations and metadata building in build_prefill_chunk_metadata. Ensure every rank emits chunks for all computed boundaries, even if some are empty locally, so all ranks participate in subsequent collectives consistently.
🧹 Nitpick comments (2)
vllm/v1/worker/utils.py (1)
167-167: Addstrict=Trueto thezip()call.Ruff flags line 167 with B905. Since
outerandouter_stridesare both derived fromouter_dims(line 20–25), they always have the same length, makingstrict=Trueboth safe and appropriate. However, the proposed diff exceeds the 88-character line limit. Reformat as:Proposed cleanup
- off_bytes = sum(i * s for i, s in zip(outer, outer_strides)) + off_bytes = sum( + i * s for i, s in zip(outer, outer_strides, strict=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/v1/worker/utils.py` at line 167, Add `strict=True` to the `zip()` call in the `off_bytes` calculation on line 167 where `outer` and `outer_strides` are zipped together. Since the line will exceed the 88-character limit after adding this parameter, reformat the code by breaking it across multiple lines (for example, by using an intermediate variable or restructuring the sum comprehension) to maintain compliance with the line length limit while adding the required `strict=True` parameter.Source: Linters/SAST tools
vllm/model_executor/layers/sparse_attn_indexer.py (1)
55-81: ⚡ Quick winUse Google-style docstrings for the new helpers.
These new docstrings describe behavior but omit
Args:/Returns:sections. Expand them to Google style, or remove the docstrings for trivial helpers. As per coding guidelines,**/*.py: “Use Google-style docstrings withArgs:/Returns:/Raises:sections instead of reStructuredText/Sphinx fields.”Also applies to: 198-207, 272-281, 316-330
🤖 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 55 - 81, The docstrings for the helper functions _use_persistent_topk_decode, _local_to_global_position, and _global_to_local_position do not follow Google-style format as required by the coding guidelines. Expand each of these docstrings to include Args: and Returns: sections that document all parameters and return values respectively. Also apply the same Google-style docstring format to the other helper functions mentioned in the comment (at lines 198-207, 272-281, and 316-330).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.
Inline comments:
In `@vllm/model_executor/layers/attention/mla_attention.py`:
- Around line 1059-1062: The extract_layer_index() helper function uses
assertions which raise AssertionError on parse failures, not ValueError, so the
current except ValueError clause does not catch the actual exception when
layer_name is invalid or ambiguous. Update the except block to catch
AssertionError instead of ValueError (or catch both exceptions) so that the
fallback assignment of layer_id = None properly handles all failure cases from
extract_layer_index().
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 12-16: The issue is that _dcp_global_topk_requested() defaults to
true, causing the B12X sparse indexer to be disabled even for DCP1 or non-DCP
runs where the total CP world size is 1. Locate the logic around line 440 where
the B12X indexer is conditionally disabled and modify the condition to check not
only _dcp_global_topk_requested() but also verify that the CP world size is
actually greater than 1. This ensures B12X is only disabled when truly
performing distributed sparse indexing across multiple processes, allowing
single-process runs to use B12X as intended.
In `@vllm/model_executor/models/deepseek_v2.py`:
- Around line 592-603: The extract_layer_index(self.prefix) call lacks error
handling for non-standard prefixes and can raise exceptions, unlike the fallback
protection added in MLAAttention. Wrap the extract_layer_index() call in a
try-except block to safely handle exceptions, setting layer_id to None when an
exception occurs. This ensures the indexer cache logic proceeds safely and the
MLA/indexer KV specs reach the merge/routing logic even when prefix parsing
fails.
In `@vllm/v1/core/kv_cache_coordinator.py`:
- Around line 552-576: The final assertion that checks for DCP prefix caching
compatibility with DeepseekV4 hybrid layouts requires
`self.disable_prefix_cache_for_dsv4_dcp` to be True, but this flag is set to
False when `enable_caching` is False. This incorrectly rejects valid
configurations where prefix caching is already disabled globally. Modify the
assertion condition to also permit the case where `enable_caching` is False,
since there is no prefix cache to disable in that scenario and the restriction
should not apply. The assertion should allow the configuration to pass when
either `self.disable_prefix_cache_for_dsv4_dcp` is True or when `enable_caching`
is False.
- Around line 558-568: The cache-hit lookup path uses unscaled spec.block_size
for hash conversions in _get_block_hashes() and hit-length calculations at lines
716, 726, 737, 758, but the assertion above correctly validates against
effective DCP-scaled block sizes. To fix this inconsistency, either pass
dcp_world_size and pcp_world_size to the manager_cls.find_longest_cache_hit()
calls (currently at line 728) to enable it to scale block sizes internally, or
convert spec.block_size to its effective DCP-scaled size before passing it to
_get_block_hashes() and before using it in the subsequent hit-length calculation
comparisons to ensure all hash and length computations use consistent DCP-scaled
dimensions.
In `@vllm/v1/worker/gpu/spec_decode/eagle/utils.py`:
- Line 24: The environment variable check on the line with _os.environ.get for
VLLM_DCP_SHARD_DRAFT exceeds the 88-character Python line limit. Refactor this
line by either breaking it across multiple lines using parentheses for implicit
line continuation, or by extracting the environment variable retrieval into a
separate variable on the line above the if statement to keep each line within
the 88-character limit.
---
Outside diff comments:
In `@vllm/v1/attention/backends/mla/indexer.py`:
- Around line 927-939: The DeepseekV32IndexerMetadata constructor at lines
927-939 is setting dcp_world_size to self.dcp_world_size, but this does not
reflect the effective DCP layout that was intentionally set to 1 (non-local KV)
in lines 657-662 when dcp_local_seq_lens is absent. This causes downstream
issues in global top-k remap which expects the advertised dcp_world_size to
match the actual KV sharding strategy. Replace the dcp_world_size assignment in
the DeepseekV32IndexerMetadata constructor to use the effective_dcp_world_size
value (which reflects whether KV is treated as replicated/non-local with size 1,
or sharded with the physical size) instead of self.dcp_world_size. If the
physical DCP size is still needed by callers, add a separate explicit field to
the metadata rather than overloading the sharded-KV indicator field.
- Around line 713-775: When DCP global top-k is active, the prefill chunk
boundaries must be rank-invariant to prevent collective operations from hanging
due to mismatched shapes. Currently, chunk_specs is built from rank-local
compressed_seq_lens_cpu_np on line 745 in the split_indexer_prefill_chunks
calls, causing different ranks to create different chunk boundaries. Compute a
rank-invariant global compressed length (e.g., through collective communication
to get max or consensus values) and use it to build chunk_specs instead of the
rank-local compressed_seq_lens_cpu_np. Keep the rank-local
compressed_seq_lens_cpu_np values for the actual KV gather operations and
metadata building in build_prefill_chunk_metadata. Ensure every rank emits
chunks for all computed boundaries, even if some are empty locally, so all ranks
participate in subsequent collectives consistently.
In `@vllm/v1/core/kv_cache_utils.py`:
- Around line 631-644: At line 1435 where MLAAttentionSpec is being
reconstructed from a SlidingWindowMLASpec, explicitly preserve the
dcp_replicated attribute by passing it to the MLAAttentionSpec constructor or
copying it after construction. This ensures that if dcp_replicated is present on
the source spec, it is properly carried forward to the reconstructed spec and
won't silently default to False, which could cause the GCD hash layout
computation in the effective_block_sizes calculation to treat replicated draft
groups as DCP-sharded and produce incorrect block sizes.
---
Nitpick comments:
In `@vllm/model_executor/layers/sparse_attn_indexer.py`:
- Around line 55-81: The docstrings for the helper functions
_use_persistent_topk_decode, _local_to_global_position, and
_global_to_local_position do not follow Google-style format as required by the
coding guidelines. Expand each of these docstrings to include Args: and Returns:
sections that document all parameters and return values respectively. Also apply
the same Google-style docstring format to the other helper functions mentioned
in the comment (at lines 198-207, 272-281, and 316-330).
In `@vllm/v1/worker/utils.py`:
- Line 167: Add `strict=True` to the `zip()` call in the `off_bytes` calculation
on line 167 where `outer` and `outer_strides` are zipped together. Since the
line will exceed the 88-character limit after adding this parameter, reformat
the code by breaking it across multiple lines (for example, by using an
intermediate variable or restructuring the sum comprehension) to maintain
compliance with the line length limit while adding the required `strict=True`
parameter.
🪄 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: 18983608-5cc4-4cb4-83b2-10e18522b544
📒 Files selected for processing (14)
vllm/model_executor/layers/attention/mla_attention.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/kv_cache_interface.pyvllm/v1/worker/gpu/attn_utils.pyvllm/v1/worker/gpu/model_runner.pyvllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.pyvllm/v1/worker/gpu/spec_decode/autoregressive/speculator.pyvllm/v1/worker/gpu/spec_decode/eagle/utils.pyvllm/v1/worker/gpu/spec_decode/speculator.pyvllm/v1/worker/utils.py
Comment on lines
+1059
to
+1062
| try: | ||
| layer_id = extract_layer_index(self.layer_name) | ||
| except ValueError: | ||
| layer_id = None |
There was a problem hiding this comment.
Catch the helper’s actual parse-failure path.
extract_layer_index() asserts on missing or ambiguous integer components, so the current except ValueError does not protect the fallback path; a default or non-standard layer_name can still abort KV-cache spec construction.
Proposed fix
try:
layer_id = extract_layer_index(self.layer_name)
- except ValueError:
+ except (AssertionError, IndexError, ValueError):
layer_id = None🤖 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/attention/mla_attention.py` around lines 1059 -
1062, The extract_layer_index() helper function uses assertions which raise
AssertionError on parse failures, not ValueError, so the current except
ValueError clause does not catch the actual exception when layer_name is invalid
or ambiguous. Update the except block to catch AssertionError instead of
ValueError (or catch both exceptions) so that the fallback assignment of
layer_id = None properly handles all failure cases from extract_layer_index().
Comment on lines
+12
to
+16
| from vllm.distributed.parallel_state import get_dcp_group | ||
| from vllm.forward_context import get_forward_context | ||
| from vllm.model_executor.custom_op import CustomOp | ||
| from vllm.platforms import current_platform | ||
| from vllm.triton_utils import HAS_TRITON, tl, triton |
There was a problem hiding this comment.
Don’t disable B12X for DCP1/non-DCP runs by default.
_dcp_global_topk_requested() defaults to true, so Line 440 disables the B12X sparse indexer even when total CP world size is 1. That contradicts the metadata-side gate and routes DCP1/non-DCP runs away from B12X unnecessarily.
Proposed fix
from vllm.model_executor.custom_op import CustomOp
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON, tl, triton
+from vllm.v1.worker.cp_utils import get_total_cp_world_size- if _dcp_global_topk_requested():
+ if _dcp_global_topk_requested() and get_total_cp_world_size() > 1:
return FalseAlso applies to: 436-441
🤖 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 12 - 16, The
issue is that _dcp_global_topk_requested() defaults to true, causing the B12X
sparse indexer to be disabled even for DCP1 or non-DCP runs where the total CP
world size is 1. Locate the logic around line 440 where the B12X indexer is
conditionally disabled and modify the condition to check not only
_dcp_global_topk_requested() but also verify that the CP world size is actually
greater than 1. This ensures B12X is only disabled when truly performing
distributed sparse indexing across multiple processes, allowing single-process
runs to use B12X as intended.
Comment on lines
+592
to
+603
| layer_id = extract_layer_index(self.prefix) | ||
| num_hidden_layers = getattr(vllm_config.model_config.hf_config, | ||
| "num_hidden_layers", None) | ||
| import os as _os | ||
| _shard_draft = _os.environ.get("VLLM_DCP_SHARD_DRAFT", "0").lower() in ( | ||
| "1", "true", "yes") | ||
| # R1 prototype: see MLAAttention.get_kv_cache_spec. Shard the indexer | ||
| # draft cache too when VLLM_DCP_SHARD_DRAFT is set. | ||
| dcp_replicated = (not _shard_draft | ||
| and layer_id is not None | ||
| and num_hidden_layers is not None | ||
| and int(layer_id) >= int(num_hidden_layers)) |
There was a problem hiding this comment.
Use the same safe layer-index parsing for the indexer cache.
This path lacks the fallback added in MLAAttention, but extract_layer_index() can raise on non-standard prefixes. If the indexer cache aborts here, the MLA/indexer KV specs never reach the merge/routing logic.
Proposed fix
- layer_id = extract_layer_index(self.prefix)
+ try:
+ layer_id = extract_layer_index(self.prefix)
+ except (AssertionError, IndexError, ValueError):
+ layer_id = None
num_hidden_layers = getattr(vllm_config.model_config.hf_config,
"num_hidden_layers", None)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| layer_id = extract_layer_index(self.prefix) | |
| num_hidden_layers = getattr(vllm_config.model_config.hf_config, | |
| "num_hidden_layers", None) | |
| import os as _os | |
| _shard_draft = _os.environ.get("VLLM_DCP_SHARD_DRAFT", "0").lower() in ( | |
| "1", "true", "yes") | |
| # R1 prototype: see MLAAttention.get_kv_cache_spec. Shard the indexer | |
| # draft cache too when VLLM_DCP_SHARD_DRAFT is set. | |
| dcp_replicated = (not _shard_draft | |
| and layer_id is not None | |
| and num_hidden_layers is not None | |
| and int(layer_id) >= int(num_hidden_layers)) | |
| try: | |
| layer_id = extract_layer_index(self.prefix) | |
| except (AssertionError, IndexError, ValueError): | |
| layer_id = None | |
| num_hidden_layers = getattr(vllm_config.model_config.hf_config, | |
| "num_hidden_layers", None) | |
| import os as _os | |
| _shard_draft = _os.environ.get("VLLM_DCP_SHARD_DRAFT", "0").lower() in ( | |
| "1", "true", "yes") | |
| # R1 prototype: see MLAAttention.get_kv_cache_spec. Shard the indexer | |
| # draft cache too when VLLM_DCP_SHARD_DRAFT is set. | |
| dcp_replicated = (not _shard_draft | |
| and layer_id is not None | |
| and num_hidden_layers is not None | |
| and int(layer_id) >= int(num_hidden_layers)) |
🤖 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 592 - 603, The
extract_layer_index(self.prefix) call lacks error handling for non-standard
prefixes and can raise exceptions, unlike the fallback protection added in
MLAAttention. Wrap the extract_layer_index() call in a try-except block to
safely handle exceptions, setting layer_id to None when an exception occurs.
This ensures the indexer cache logic proceeds safely and the MLA/indexer KV
specs reach the merge/routing logic even when prefix parsing fails.
Comment on lines
552
to
+576
| self.disable_prefix_cache_for_dsv4_dcp = ( | ||
| enable_caching | ||
| and dcp_world_size > 1 | ||
| and pcp_world_size == 1 | ||
| and ( | ||
| is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config) | ||
| or _has_dcp_replicated | ||
| ) | ||
| and is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config) | ||
| ) | ||
| if not self.disable_prefix_cache_for_dsv4_dcp: | ||
| # R3 fix: compare against the *effective* (DCP-scaled) block size the | ||
| # manager actually uses, not the unscaled spec block size. Under DCP | ||
| # a sharded group's manager.block_size == spec.block_size * dcp, which | ||
| # matches hash_block_size (= LCM of effective sizes). The original | ||
| # used g.kv_cache_spec.block_size (unscaled) and so was structurally | ||
| # unsatisfiable under DCP. Mirrors UnitaryKVCacheCoordinator. | ||
| assert all( | ||
| g.kv_cache_spec.block_size % hash_block_size == 0 | ||
| for g in kv_cache_config.kv_cache_groups | ||
| mgr.block_size % hash_block_size == 0 | ||
| for mgr in self.single_type_managers | ||
| ), "block_size must be divisible by hash_block_size" | ||
| assert dcp_world_size == 1 or self.disable_prefix_cache_for_dsv4_dcp, ( | ||
| "DCP not support hybrid attn now." | ||
| ) | ||
| # DCP>1 is allowed with prefix caching for non-hybrid (e.g. uniformly | ||
| # DCP-sharded) layouts; only the genuine DeepseekV4 MLA/SWA hybrid must | ||
| # keep caching disabled under DCP. | ||
| assert ( | ||
| dcp_world_size == 1 | ||
| or not is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config) | ||
| or self.disable_prefix_cache_for_dsv4_dcp | ||
| ), "DCP prefix caching unsupported for the DeepseekV4 MLA/SWA hybrid." |
There was a problem hiding this comment.
Don’t reject DSV4 DCP when prefix caching is already off.
self.disable_prefix_cache_for_dsv4_dcp is false when enable_caching is false, but the assertion still requires it for DCP + DeepseekV4 hybrid layouts. That makes a no-prefix-cache configuration fail initialization even though there is nothing to disable.
Proposed fix
assert (
dcp_world_size == 1
or not is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config)
+ or not enable_caching
or self.disable_prefix_cache_for_dsv4_dcp
), "DCP prefix caching unsupported for the DeepseekV4 MLA/SWA hybrid."📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| self.disable_prefix_cache_for_dsv4_dcp = ( | |
| enable_caching | |
| and dcp_world_size > 1 | |
| and pcp_world_size == 1 | |
| and ( | |
| is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config) | |
| or _has_dcp_replicated | |
| ) | |
| and is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config) | |
| ) | |
| if not self.disable_prefix_cache_for_dsv4_dcp: | |
| # R3 fix: compare against the *effective* (DCP-scaled) block size the | |
| # manager actually uses, not the unscaled spec block size. Under DCP | |
| # a sharded group's manager.block_size == spec.block_size * dcp, which | |
| # matches hash_block_size (= LCM of effective sizes). The original | |
| # used g.kv_cache_spec.block_size (unscaled) and so was structurally | |
| # unsatisfiable under DCP. Mirrors UnitaryKVCacheCoordinator. | |
| assert all( | |
| g.kv_cache_spec.block_size % hash_block_size == 0 | |
| for g in kv_cache_config.kv_cache_groups | |
| mgr.block_size % hash_block_size == 0 | |
| for mgr in self.single_type_managers | |
| ), "block_size must be divisible by hash_block_size" | |
| assert dcp_world_size == 1 or self.disable_prefix_cache_for_dsv4_dcp, ( | |
| "DCP not support hybrid attn now." | |
| ) | |
| # DCP>1 is allowed with prefix caching for non-hybrid (e.g. uniformly | |
| # DCP-sharded) layouts; only the genuine DeepseekV4 MLA/SWA hybrid must | |
| # keep caching disabled under DCP. | |
| assert ( | |
| dcp_world_size == 1 | |
| or not is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config) | |
| or self.disable_prefix_cache_for_dsv4_dcp | |
| ), "DCP prefix caching unsupported for the DeepseekV4 MLA/SWA hybrid." | |
| self.disable_prefix_cache_for_dsv4_dcp = ( | |
| enable_caching | |
| and dcp_world_size > 1 | |
| and pcp_world_size == 1 | |
| and is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config) | |
| ) | |
| if not self.disable_prefix_cache_for_dsv4_dcp: | |
| # R3 fix: compare against the *effective* (DCP-scaled) block size the | |
| # manager actually uses, not the unscaled spec block size. Under DCP | |
| # a sharded group's manager.block_size == spec.block_size * dcp, which | |
| # matches hash_block_size (= LCM of effective sizes). The original | |
| # used g.kv_cache_spec.block_size (unscaled) and so was structurally | |
| # unsatisfiable under DCP. Mirrors UnitaryKVCacheCoordinator. | |
| assert all( | |
| mgr.block_size % hash_block_size == 0 | |
| for mgr in self.single_type_managers | |
| ), "block_size must be divisible by hash_block_size" | |
| # DCP>1 is allowed with prefix caching for non-hybrid (e.g. uniformly | |
| # DCP-sharded) layouts; only the genuine DeepseekV4 MLA/SWA hybrid must | |
| # keep caching disabled under DCP. | |
| assert ( | |
| dcp_world_size == 1 | |
| or not is_deepseek_v4_hybrid_kv_cache_config(kv_cache_config) | |
| or not enable_caching | |
| or self.disable_prefix_cache_for_dsv4_dcp | |
| ), "DCP prefix caching unsupported for the DeepseekV4 MLA/SWA hybrid." |
🤖 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_coordinator.py` around lines 552 - 576, The final
assertion that checks for DCP prefix caching compatibility with DeepseekV4
hybrid layouts requires `self.disable_prefix_cache_for_dsv4_dcp` to be True, but
this flag is set to False when `enable_caching` is False. This incorrectly
rejects valid configurations where prefix caching is already disabled globally.
Modify the assertion condition to also permit the case where `enable_caching` is
False, since there is no prefix cache to disable in that scenario and the
restriction should not apply. The assertion should allow the configuration to
pass when either `self.disable_prefix_cache_for_dsv4_dcp` is True or when
`enable_caching` is False.
Comment on lines
558
to
568
| if not self.disable_prefix_cache_for_dsv4_dcp: | ||
| # R3 fix: compare against the *effective* (DCP-scaled) block size the | ||
| # manager actually uses, not the unscaled spec block size. Under DCP | ||
| # a sharded group's manager.block_size == spec.block_size * dcp, which | ||
| # matches hash_block_size (= LCM of effective sizes). The original | ||
| # used g.kv_cache_spec.block_size (unscaled) and so was structurally | ||
| # unsatisfiable under DCP. Mirrors UnitaryKVCacheCoordinator. | ||
| assert all( | ||
| g.kv_cache_spec.block_size % hash_block_size == 0 | ||
| for g in kv_cache_config.kv_cache_groups | ||
| mgr.block_size % hash_block_size == 0 | ||
| for mgr in self.single_type_managers | ||
| ), "block_size must be divisible by hash_block_size" |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect cache-hit paths for unscaled spec block-size usage.
# Expected: DCP-enabled hybrid lookup paths should use each manager's effective
# block_size, or explicitly pass an effective block size into class/static helpers.
rg -n -C5 \
'def find_longest_cache_hit|kv_cache_spec\.block_size|spec\.block_size|manager\.block_size|mgr\.block_size' \
vllm/v1/core/kv_cache_coordinator.py
fd -e py '.*manager.*' vllm/v1/core -x rg -n -C5 \
'def find_longest_cache_hit|kv_cache_spec\.block_size|spec\.block_size|block_size' {}Repository: local-inference-lab/vllm
Length of output: 37641
🏁 Script executed:
# Check HybridKVCacheCoordinator initialization and how specs are managed
rg -n -A10 -B5 'class HybridKVCacheCoordinator|def __init__' vllm/v1/core/kv_cache_coordinator.py | head -100Repository: local-inference-lab/vllm
Length of output: 2474
🏁 Script executed:
# Check how specs are stored and what information they contain
rg -n 'self.kv_cache_spec|self.kv_cache_config' vllm/v1/core/kv_cache_coordinator.py | head -50Repository: local-inference-lab/vllm
Length of output: 761
🏁 Script executed:
# Verify what happens in the manager classes when they use spec.block_size
ast-grep --pattern 'class FullAttentionManager {
$$$
find_longest_cache_hit($_) {
$$$
block_size = $_
$$$
}
$$$
}'Repository: local-inference-lab/vllm
Length of output: 50
🏁 Script executed:
# Check if spec blocks are DCP-aware by looking at any scaling operations on specs
rg -n 'dcp_world_size|pcp_world_size' vllm/v1/core/kv_cache_coordinator.py | grep -E '(spec|kv_cache_spec)' | head -20Repository: local-inference-lab/vllm
Length of output: 50
🏁 Script executed:
# View HybridKVCacheCoordinator.__init__ and how it creates managers
sed -n '512,600p' vllm/v1/core/kv_cache_coordinator.pyRepository: local-inference-lab/vllm
Length of output: 4397
🏁 Script executed:
# View the find_longest_cache_hit implementation in HybridKVCacheCoordinator and the spec usage
sed -n '653,750p' vllm/v1/core/kv_cache_coordinator.pyRepository: local-inference-lab/vllm
Length of output: 4431
🏁 Script executed:
# Check what spec is passed to manager_cls.find_longest_cache_hit and if it's the unscaled version
sed -n '728,735p' vllm/v1/core/kv_cache_coordinator.pyRepository: local-inference-lab/vllm
Length of output: 492
🏁 Script executed:
# Verify: are spec objects DCP-aware or unscaled? Check where they're created/stored
rg -n 'kv_cache_spec\s*=' vllm/v1/core/kv_cache_coordinator.py | head -20Repository: local-inference-lab/vllm
Length of output: 353
🏁 Script executed:
# Check the SpecGroup structure and what spec it holds
rg -n -B5 'class SpecGroup|`@dataclass`.*SpecGroup' vllm/v1/core/kv_cache_coordinator.pyRepository: local-inference-lab/vllm
Length of output: 245
🏁 Script executed:
# Check if specs are modified/scaled anywhere for DCP before being used in cache-hit
rg -n 'block_size.*=.*spec\.block_size|spec\.block_size.*=|block_size.*dcp' vllm/v1/core/kv_cache_coordinator.pyRepository: local-inference-lab/vllm
Length of output: 473
Fix cache-hit path to use effective (DCP-scaled) block sizes.
The assertion at lines 565-568 correctly compares mgr.block_size (the effective DCP-scaled size), but the cache-hit lookup below still uses unscaled spec.block_size for hash conversions (line 685) and hit-length calculations (lines 716, 726, 737, 758).
For DCP-sharded groups, the unscaled spec represents only one shard, so using it causes hashing and hit reporting at the wrong granularity. Either pass dcp_world_size and pcp_world_size to the manager_cls.find_longest_cache_hit() calls (line 728), or convert spec.block_size to effective size before using it in _get_block_hashes() and hit-length calculations.
🤖 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_coordinator.py` around lines 558 - 568, The cache-hit
lookup path uses unscaled spec.block_size for hash conversions in
_get_block_hashes() and hit-length calculations at lines 716, 726, 737, 758, but
the assertion above correctly validates against effective DCP-scaled block
sizes. To fix this inconsistency, either pass dcp_world_size and pcp_world_size
to the manager_cls.find_longest_cache_hit() calls (currently at line 728) to
enable it to scale block sizes internally, or convert spec.block_size to its
effective DCP-scaled size before passing it to _get_block_hashes() and before
using it in the subsequent hit-length calculation comparisons to ensure all hash
and length computations use consistent DCP-scaled dimensions.
| # DCP head all-gather and the cross-rank LSE-reduce path is taken. | ||
| import os as _os | ||
| _draft_parallel_config = speculative_config.draft_parallel_config | ||
| if _os.environ.get("VLLM_DCP_SHARD_DRAFT", "0").lower() in ("1", "true", "yes"): |
There was a problem hiding this comment.
Wrap the shard-draft environment check.
Line 24 exceeds the repository’s 88-character Python line limit.
Proposed cleanup
- if _os.environ.get("VLLM_DCP_SHARD_DRAFT", "0").lower() in ("1", "true", "yes"):
+ shard_draft = _os.environ.get("VLLM_DCP_SHARD_DRAFT", "0").lower()
+ if shard_draft in ("1", "true", "yes"):As per coding guidelines, “Maintain a line length limit of 88 characters for Python code.”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if _os.environ.get("VLLM_DCP_SHARD_DRAFT", "0").lower() in ("1", "true", "yes"): | |
| shard_draft = _os.environ.get("VLLM_DCP_SHARD_DRAFT", "0").lower() | |
| if shard_draft in ("1", "true", "yes"): |
🤖 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/worker/gpu/spec_decode/eagle/utils.py` at line 24, The environment
variable check on the line with _os.environ.get for VLLM_DCP_SHARD_DRAFT exceeds
the 88-character Python line limit. Refactor this line by either breaking it
across multiple lines using parentheses for implicit line continuation, or by
extracting the environment variable retrieval into a separate variable on the
line above the if statement to keep each line within the 88-character limit.
Source: Coding guidelines
lukealonso
force-pushed
the
dev/dark-devotion
branch
from
June 20, 2026 15:36
4b74109 to
2edcf9c
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Supersedes #26 and #28.
This consolidates the GLM-5.2 / DeepSeek-style DCP + MTP fixes into one branch on current
dev/dark-devotion, and adds exact DCP global top-k as the default DCP behavior.Included:
VLLM_DCP_GLOBAL_TOPK=0disables it).extract_layer_index()instead of an inline regex.Not included:
Why global top-k
The current local/per-rank sparse top-k union and exact global top-k choose materially different sparse routes under DCP. In the diagnostic comparison, DCP4 local-vs-global selection had Jaccard
0.237, precision0.274, recall0.659, so this is not just a tie-break difference. Exact global top-k is the more faithful DCP behavior.B12X follow-up
The current B12X sparse indexer path cannot directly implement exact global top-k because it only returns selected indices. The exact DCP global top-k path needs candidate scores/logits to all-gather and select globally, so this PR falls back to the logits-producing path for DCP global top-k. A native B12X follow-up should return
(global_kv_position, score)candidates or perform the global select/remap internally.Validation
Static validation:
Both passed.
Runtime context used for comparison:
voipmonitor/vllm:dark-devotion-pr28-globaltopk-ee6cd1e-20260619use_index_cache=true, full GLM-5.2FFFSSS...index pattern.Prior #28 MTP3 comparison on the same stack:
74.4, 0k C4211.4, 64k C147.6, 64k C4125.0capacity-limited; weighted acceptance51.4%; prefix cache0.0%.92.4, 0k C4288.6, 64k C164.7, 64k C4207.9; weighted acceptance59.0%; prefix cache reached61.7%.Summary by CodeRabbit
Release Notes
New Features
Bug Fixes