Enable prefix caching under DCP for GLM-5.2 / DeepSeek-V4 by DCP-sharding the MTP draft

[flobernd]

Problem

With --decode-context-parallel-size > 1, GLM-5.2 (DeepSeek-V4 sparse MLA + MTP
spec decode) reports Prefix cache hit rate: 0.0% permanently:
HybridKVCacheCoordinator force-disables prefix caching whenever a KV-cache
group is dcp_replicated, and the base branch marks the MTP draft group
dcp_replicated=True (it replicates draft KV across DCP ranks because the draft
broke when naively sharded). So MTP + DCP ⇒ no prefix caching — very costly for
long shared prefixes (agentic / multi-turn).

Fix

DCP-shard the MTP draft (instead of replicating it), behind the opt-in
VLLM_DCP_SHARD_DRAFT env flag (default off = original replication). The
non-obvious part is why the draft broke when sharded, and the fix for it.

Root cause of the sharded-draft breakage. The draft model is built via
eagle/utils.py:_create_draft_vllm_config using
speculative_config.draft_parallel_config, whose decode_context_parallel_size
is 1 (correct for a replicated draft). So B12xMLASparseImpl.__init__ —
which reads parallel_config.decode_context_parallel_size and sizes its
caller-owned-scratch plan via _workspace_num_heads = num_heads * max(1, dcp) —
built the draft attention impl non-DCP (dcp_world_size=1, need_lse=False,
scratch for the un-gathered head count). Meanwhile the draft KV was sharded and
the metadata builder used dcp=4. Result: each rank read its 1/dcp KV shard as
the full context and never did the cross-rank LSE reduce → garbage proposals
(spec acceptance 4.0 → 1.6; the target still verifies, so output stayed correct
— only speed was lost).

The fix: when the draft is DCP-sharded, build it with the parent's DCP world
size so the draft impl is constructed with dcp_world_size=4 — its plan/scratch
is sized for the head all-gather and the cross-rank LSE-reduce path is taken. (A
mid-forward override does NOT work: the b12x decode/extend plan is sized once per
mode in __init__, so flipping dcp_world_size later throws
q heads 32 do not match scratch heads 16.)

Note on B12X: this does not change the B12X scratch ownership model. The
B12xMLASparseImpl plan is still built once per mode in __init__ and the
scratch is still caller-owned, fetched per call from
current_workspace_manager() with plan.bind(scratch=...) (per
docs/contributing/b12x-vllm-bindings.md). This change only routes the correct
decode_context_parallel_size into the draft's config so its plan is sized
like the target's.

Changes (6 files)

1. v1/worker/gpu/spec_decode/eagle/utils.py — _create_draft_vllm_config:
   inherit the target's decode_context_parallel_size for the draft when
   sharding. The core fix.
2. model_executor/layers/attention/mla_attention.py,
   model_executor/models/deepseek_v2.py — gate the draft's
   dcp_replicated (layer_id >= num_hidden_layers) behind
   VLLM_DCP_SHARD_DRAFT; when set, the draft KV is DCP-sharded like the target.
3. v1/core/kv_cache_coordinator.py — (a) disable prefix caching under DCP
   only for the genuine DeepSeek-V4 MLA/SWA hybrid, not for a dcp_replicated
   group; (b) block-size divisibility assert fix — compare the DCP-scaled
   manager.block_size vs hash_block_size, not the unscaled spec.block_size
   (the original was structurally unsatisfiable under DCP; Unitary already does
   this); (c) relax assert dcp==1 or disable to allow DCP>1 + caching for
   non-hybrid layouts.
4. v1/core/kv_cache_utils.py — resolve_kv_cache_block_sizes uses the GCD
   of effective block sizes for hash_block_size (scheduler size stays the LCM,
   so hits still align to whole cross-rank blocks).
5. v1/worker/utils.py — KVBlockZeroer supports non-uniform page sizes
   (DeepSeek-V4 has two: MLA latent vs DSA indexer); one zeroing kernel per
   distinct page size. Independently useful robustness fix.

For a follow-up, the env flag should be promoted to a speculative_config field.

Results (author hardware: 8×, TP=8, DCP=4, MTP=5, GLM-5.2-NVFP4; VLLM_DCP_SHARD_DRAFT=1)

metric	replicated baseline	sharded + fix
prefix caching	off (0%)	works (hits climb; cross-request reuse)
spec acceptance length	3.3–4.0	3.25–4.21 (preserved)
GPU KV cache	1,642,855 tok / 6.42×	1,717,248 / 6.71× (+5%, draft VRAM reclaimed)
decode tok/s (single)	22.7	30.1 (+33%)
decode tok/s (conc=4)	54.3	73.0 (+34%)

The sharded draft is faster than replicated (¼ the KV per rank → cheaper b12x
decode per step at equal acceptance), on top of enabling prefix caching and
reclaiming VRAM.

Note on determinism

At temp=0 the engine is not bit-deterministic in this config (DCP reductions

• async scheduling + spec decode → FP/ordering variance): repeated identical
  requests — cache hit or miss — produce minor synonym-level differences. This is
  pre-existing (warm-vs-warm runs also differ) and independent of this change;
  prefix-cache reuse is correct (coherent outputs, preserved acceptance, hashing
  verified by climbing hits). A strict cold==warm byte check is therefore not a
  valid gate here.

Why this is not a duplicate

This PR is additive on top of the base branch
codex/glm52-dcp-mtp-replicated-kv-20260618 (open PR #26, "Fix GLM 5.2 DCP MTP
metadata and graph capture state"), which fixes DCP MTP metadata/graph-capture
for a replicated draft. This PR takes the opposite layout decision — it
DCP-shards the draft to enable prefix caching — and is gated behind a new env
flag so the base branch's replicated path remains the default. No other open PR
on this fork addresses DCP prefix caching / draft sharding (#27 spec proposer
metadata cleanup, #22 MiMo-V2 streaming, #8 Step3.5 MTP argmax are
unrelated).

Tests run

• In this environment (no CUDA/GPU, no built vLLM venv): static validation
  only — python -m py_compile passes on all 6 modified files; all added lines
  are ≤ 88 chars; the diff is exactly the 6-file change described above.
• On author hardware (8×GPU, config above, VLLM_DCP_SHARD_DRAFT=1):
  functional + performance validation per the results table — prefix-cache hit
  rate climbs across requests, spec acceptance length preserved, decode
  throughput +33–34%. Correctness validated via output coherence + acceptance +
  hit-rate (see determinism note for why a byte-exact cold==warm check is not a
  valid gate).
• Not run here: pre-commit (ruff/mypy) and the pytest suite — they require
  the uv/.venv toolchain and, for the DCP/MTP paths, multi-GPU hardware not
  available in this environment. Recommend running pre-commit run --all-files
  and the relevant GPU tests before merge.

Flags / repro

VLLM_DCP_SHARD_DRAFT=1 enables the sharded draft. VLLM_DCP_DEBUG=1 logs
per-layer DCP attention setup.

---

AI assistance disclosure: This change was prepared with AI assistance
(Claude Code). The patch was authored/validated by a human submitter who
reviews every changed line and is accountable for the change end-to-end.

🤖 Generated with Claude Code

https://claude.ai/code/session_01GGuk9dcSi3b1pk5JyN4Lxj

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• dev/*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 91fed6a8-520c-4431-9e88-eb68d3503ba0

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀