[Kimi-K3] support DCP partial prefix cache hit (#50493)

Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

tests/distributed/test_kimi_linear_context_parallel.py

@@ -2,12 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import math
+from pathlib import Path
 
 import pytest
 
 from tests.conftest import VllmRunner
 from tests.utils import create_new_process_for_each_test
 from vllm import SamplingParams, TokensPrompt
+from vllm.config import CUDAGraphMode
+from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config
 
 MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
 
@@ -126,3 +129,111 @@ def test_kimi_linear_dcp_tiny(
             logprob_drifts.append(abs(baseline_logprob - dcp_logprob))
 
     assert max(logprob_drifts) <= 1e-2
+
+
+def _make_tiny_k3_config(model_dir: Path) -> str:
+    config = KimiK3Config(
+        text_config=_make_tiny_overrides(),
+        vision_config={
+            "vt_num_attention_heads": 2,
+            "vt_num_hidden_layers": 1,
+            "vt_hidden_size": 32,
+            "vt_intermediate_size": 64,
+            "qkv_hidden_size": 48,
+        },
+        architectures=["KimiK3ForConditionalGeneration"],
+    )
+    config.save_pretrained(model_dir)
+    return str(model_dir)
+
+
+def _run_k3_partial_prefix_reuse(
+    vllm_runner: type[VllmRunner], model_name: str, dcp_size: int
+) -> tuple[list[int], list[float], int]:
+    block_size = 256
+    common_prefix = [3 + (token_idx % 251) for token_idx in range(block_size)]
+    prime_prompt = TokensPrompt(prompt_token_ids=common_prefix + [17])
+    replay_prompt = TokensPrompt(prompt_token_ids=common_prefix + [18, 19])
+    sampling_params = SamplingParams(
+        temperature=0,
+        max_tokens=8,
+        seed=0,
+        ignore_eos=True,
+        logprobs=20,
+    )
+
+    with vllm_runner(
+        model_name=model_name,
+        skip_tokenizer_init=True,
+        load_format="dummy",
+        language_model_only=True,
+        tensor_parallel_size=2,
+        decode_context_parallel_size=dcp_size,
+        cp_kv_cache_interleave_size=1,
+        distributed_executor_backend="mp",
+        dtype="bfloat16",
+        seed=0,
+        max_model_len=512,
+        max_num_seqs=4,
+        max_num_batched_tokens=256,
+        gpu_memory_utilization=0.85,
+        block_size=block_size,
+        enable_chunked_prefill=True,
+        enable_prefix_caching=True,
+        mamba_cache_mode="align",
+    ) as runner:
+        assert (
+            runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode
+            == CUDAGraphMode.FULL_AND_PIECEWISE
+        )
+        runner.llm.generate(
+            [prime_prompt],
+            sampling_params,
+            use_tqdm=False,
+        )
+        replay_output = runner.llm.generate(
+            [replay_prompt],
+            sampling_params,
+            use_tqdm=False,
+        )[0]
+
+    completion = replay_output.outputs[0]
+    token_ids = list(completion.token_ids)
+    assert len(token_ids) == sampling_params.max_tokens
+    assert completion.logprobs is not None
+    assert len(completion.logprobs) == sampling_params.max_tokens
+    selected_logprobs = [
+        step_logprobs[token_id].logprob
+        for token_id, step_logprobs in zip(token_ids, completion.logprobs)
+    ]
+    return token_ids, selected_logprobs, replay_output.num_cached_tokens
+
+
+@create_new_process_for_each_test()
+@pytest.mark.distributed(num_gpus=2)
+def test_kimi_k3_dcp_partial_prefix_reuse(
+    vllm_runner: type[VllmRunner],
+    num_gpus_available: int,
+    tmp_path: Path,
+) -> None:
+    if num_gpus_available < 2:
+        pytest.skip("Need at least 2 GPUs")
+
+    model_name = _make_tiny_k3_config(tmp_path / "tiny-kimi-k3")
+    baseline_tokens, baseline_logprobs, baseline_cached = _run_k3_partial_prefix_reuse(
+        vllm_runner, model_name, dcp_size=1
+    )
+    dcp_tokens, dcp_logprobs, dcp_cached = _run_k3_partial_prefix_reuse(
+        vllm_runner, model_name, dcp_size=2
+    )
+
+    assert baseline_cached == 256
+    assert dcp_cached == 256
+    assert dcp_tokens == baseline_tokens
+    assert len(dcp_logprobs) == len(baseline_logprobs) == 8
+    logprob_drifts = []
+    for baseline_logprob, dcp_logprob in zip(baseline_logprobs, dcp_logprobs):
+        assert math.isfinite(baseline_logprob)
+        assert math.isfinite(dcp_logprob)
+        logprob_drifts.append(abs(baseline_logprob - dcp_logprob))
+    assert max(logprob_drifts) <= 1e-2

tests/v1/attention/test_mla_backends.py

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
     MLAAttention,
     QueryLenSupport,
     _DecodeConcatQuantFP8,
+    build_mla_chunked_context_metadata,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.platforms import current_platform
@@ -93,6 +94,48 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
         assert spec.page_size_bytes == 64 * 656
 
 
+@pytest.mark.cpu_test
+def test_dcp_chunked_context_accepts_non_virtual_block_aligned_prefix():
+    context_lens_cpu = torch.tensor([65, 96], dtype=torch.int32)
+    prefill_query_start_loc_cpu = torch.tensor([0, 1, 2], dtype=torch.int32)
+
+    metadata = build_mla_chunked_context_metadata(
+        context_lens_cpu=context_lens_cpu,
+        prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,
+        chunked_prefill_workspace=torch.empty(0),
+        chunked_prefill_workspace_size=128,
+        block_size=64,
+        align_chunk_to_block=True,
+        device=torch.device("cpu"),
+        dcp_world_size=4,
+        dcp_local_block_size=1,
+        dcp_virtual_block_size=4,
+    )
+
+    assert metadata is not None
+    assert metadata.context_lens.tolist() == [65, 96]
+    assert [chunk.seq_lens.tolist() for chunk in metadata.chunks] == [[65], [96]]
+    assert [chunk.starts.tolist() for chunk in metadata.chunks] == [[0], [0]]
+    assert [chunk.local_context_lens_allranks for chunk in metadata.chunks] == [
+        [[17, 16, 16, 16]],
+        [[24, 24, 24, 24]],
+    ]
+    assert [chunk.padded_local_seq_lens for chunk in metadata.chunks] == [
+        [17],
+        [24],
+    ]
+    assert [chunk.padded_local_cu_seq_lens.tolist() for chunk in metadata.chunks] == [
+        [0, 17],
+        [0, 24],
+    ]
+    assert [chunk.cu_seq_lens.tolist() for chunk in metadata.chunks] == [
+        [0, 65],
+        [0, 96],
+    ]
+    assert [chunk.num_context_tokens for chunk in metadata.chunks] == [65, 96]
+    assert [chunk.num_local_context_tokens for chunk in metadata.chunks] == [17, 24]
+
+
 # Remove sm100 backends from the list if not using sm100
 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
     BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -4,6 +4,7 @@
 "align") models: scheduler chunk splitting, partial tail registration, CoW
 on partial hits, and same-step deferral."""
 
+from math import lcm
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -68,12 +69,62 @@ def test_capable_connector_uses_divergent_partial_hit_lookup():
     manager.get_computed_blocks.assert_not_called()
 
 
-def test_mamba_align_split_partial_tail_schedule():
+def make_full_mamba_manager(
+    *,
+    dcp_world_size: int,
+    hash_block_size: int = 2,
+    full_block_size: int = 4,
+    mamba_block_size: int = 4,
+    num_blocks: int = 32,
+    use_eagle: bool = False,
+):
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=full_block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=mamba_block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    scheduler_block_size = lcm(
+        full_block_size * dcp_world_size,
+        mamba_block_size,
+    )
+    return make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        dcp_world_size=dcp_world_size,
+        scheduler_block_size=scheduler_block_size,
+        hash_block_size=hash_block_size,
+        use_eagle=use_eagle,
+    )
+
+
+@pytest.mark.parametrize("dcp_world_size", [1, 4])
+def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
     """Chunk ends with partial hits on: block-aligned chunks, one extra stop
     at the prompt's last hash boundary (registering the partial tail), then
     the remaining tokens. block=512, hash=32, prompt=10000, budget=8192:
     0 -> 8192 -> 9728 -> 9984 -> 10000."""
     block_size = 512
+    scheduler_block_size = block_size * dcp_world_size
     hash_block_size = 32
     mock = SimpleNamespace(
         cache_config=SimpleNamespace(block_size=block_size),
@@ -81,6 +132,8 @@ def test_mamba_align_split_partial_tail_schedule():
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle=False,
         hash_block_size=hash_block_size,
+        dcp_world_size=dcp_world_size,
+        scheduler_block_size=scheduler_block_size,
         mamba_partial_cache_hit=True,
     )
     split = Scheduler._mamba_block_aligned_split
@@ -109,6 +162,7 @@ def test_mamba_align_split_partial_tail_schedule():
     # stops at the next block boundary (10240), later chunk ends re-align.
     req2 = make_request("1", [0] * 12000, hash_block_size, sha256)
     req2.num_computed_tokens = 9984
+    assert req2.num_computed_tokens % scheduler_block_size != 0
     assert split(self=mock, request=req2, num_new_tokens=2016) == 256
     req2.num_computed_tokens = 10240
     assert split(self=mock, request=req2, num_new_tokens=1000) == 512
@@ -1499,3 +1553,151 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits():
 
     assert num_computed == mamba_block_size
     assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed
+
+
+@pytest.mark.parametrize("dcp_world_size", [1, 2, 4])
+def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int):
+    hash_block_size = 2
+    physical_block_size = 4
+    manager = make_full_mamba_manager(
+        dcp_world_size=dcp_world_size,
+        hash_block_size=hash_block_size,
+        full_block_size=physical_block_size,
+        mamba_block_size=physical_block_size,
+    )
+    assert manager.coordinator.enable_partial_hash_hits
+
+    req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
+    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+    manager.free(req0)
+    manager.new_step_starts()
+
+    partial_hash = req0.block_hashes[2]
+    partial_full_block = manager.block_pool.get_cached_block(partial_hash, [0])
+    partial_mamba_block = manager.block_pool.get_cached_block(partial_hash, [1])
+    assert partial_full_block is not None
+    assert partial_mamba_block is not None
+
+    req1 = make_request("dcp-replay", [0, 0, 1, 1, 2, 2, 3, 3], 2, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)
+    assert num_computed == 6
+    full_block_size = physical_block_size * dcp_world_size
+    assert [len(group) for group in computed_blocks.blocks] == [
+        (6 + full_block_size - 1) // full_block_size,
+        2,
+    ]
+
+    new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks)
+    assert new_blocks is not None
+    full_new_block_id = new_blocks.get_block_ids()[0][0]
+    mamba_new_block_id = new_blocks.get_block_ids()[1][0]
+    copies, retained = manager.take_kv_cache_block_copies()
+    assert KVCacheBlockCopy(partial_full_block[0].block_id, full_new_block_id) in copies
+    assert (
+        KVCacheBlockCopy(partial_mamba_block[0].block_id, mamba_new_block_id) in copies
+    )
+    manager.block_pool.free_blocks(retained)
+
+
+@pytest.mark.parametrize("dcp_world_size", [2, 4])
+def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot(
+    dcp_world_size: int,
+):
+    block_size = 4
+    manager = make_full_mamba_manager(
+        dcp_world_size=dcp_world_size,
+        hash_block_size=block_size,
+        full_block_size=block_size,
+        mamba_block_size=block_size,
+    )
+    assert manager.coordinator.enable_partial_hash_hits
+    assert manager.coordinator._cache_hit_alignment_tokens == block_size
+    assert manager.coordinator.single_type_managers[0].block_size == (
+        block_size * dcp_world_size
+    )
+    assert manager.coordinator.single_type_managers[1].block_size == block_size
+
+    prefix = list(range(12))
+    req0 = make_request("snapshot-owner", prefix, block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
+    assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None
+    manager.free(req0)
+    manager.new_step_starts()
+
+    req1 = make_request(
+        "snapshot-replay", prefix + list(range(12, 16)), block_size, sha256
+    )
+
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)
+    assert num_computed == 12
+    assert [len(group) for group in computed_blocks.blocks] == [
+        (12 + block_size * dcp_world_size - 1) // (block_size * dcp_world_size),
+        3,
+    ]
+    partial_full_block = computed_blocks.blocks[0][-1]
+    new_blocks = manager.allocate_slots(req1, 4, num_computed, computed_blocks)
+    assert new_blocks is not None
+    full_new_block_id = new_blocks.get_block_ids()[0][0]
+    copies, retained = manager.take_kv_cache_block_copies()
+    assert KVCacheBlockCopy(partial_full_block.block_id, full_new_block_id) in copies
+    assert all(
+        copy.src_block_id != computed_blocks.blocks[1][-1].block_id for copy in copies
+    )
+    manager.block_pool.free_blocks(retained)
+
+
+def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots():
+    block_size = 4
+    manager = make_full_mamba_manager(
+        dcp_world_size=2,
+        hash_block_size=block_size,
+        full_block_size=block_size,
+        mamba_block_size=block_size,
+    )
+    prefix = list(range(12))
+    req0 = make_request("joint-owner", prefix, block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
+    assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None
+    manager.new_step_starts()
+
+    replay_tokens = prefix + list(range(12, 16))
+    req1 = make_request("joint-before", replay_tokens, block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)
+    assert num_computed == 8
+    assert [len(group) for group in computed_blocks.blocks] == [1, 2]
+
+    req0.num_computed_tokens = 8
+    assert manager.allocate_slots(req0, 4) is not None
+    manager.new_step_starts()
+
+    req2 = make_request("joint-after", replay_tokens, block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req2)
+    assert num_computed == 12
+    assert [len(group) for group in computed_blocks.blocks] == [2, 3]
+
+
+def test_dcp_partial_hit_with_eagle_rewinds_one_hash_unit():
+    hash_block_size = 2
+    manager = make_full_mamba_manager(
+        dcp_world_size=2,
+        hash_block_size=hash_block_size,
+        full_block_size=4,
+        mamba_block_size=4,
+        use_eagle=True,
+    )
+
+    req0 = make_request("eagle-owner", [7] * 6, hash_block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
+    assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None
+    req0.num_computed_tokens = 4
+    manager.new_step_starts()
+    assert manager.allocate_slots(req0, 2) is not None
+    req0.num_computed_tokens = 6
+    manager.new_step_starts()
+
+    req1 = make_request("eagle-replay", [7] * 6 + [9] * 2, 2, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req1)
+    assert num_computed == 4
+    assert [len(group) for group in computed_blocks.blocks] == [1, 1]
+    assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None

tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py

@@ -55,9 +55,9 @@ def cache_full_block_and_partial_tail(
     token_ids: list[int],
     *,
     enable_kv_cache_events: bool = False,
+    block_size: int = 6,
 ) -> tuple[BlockPool, Request, list[KVCacheBlock], BlockHash]:
     hash_block_size = 2
-    block_size = 6
     kv_cache_group_id = 0
     req = make_request("0", token_ids, hash_block_size, sha256)
     pool = BlockPool(
@@ -389,27 +389,34 @@ def test_reset_prefix_cache_clears_partial_entry_metadata():
     assert pool.cached_block_hashes_by_block == {}
 
 
-def test_evict_cached_block_removes_full_hash_and_partial_entry():
-    pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail(
-        [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
+@pytest.mark.parametrize("dcp_world_size", [1, 2, 4])
+def test_evict_cached_block_removes_full_hash_and_partial_entry(
+    dcp_world_size: int,
+):
+    block_size = 6 * dcp_world_size
+    partial_num_tokens = 2 * block_size - 2
+    pool, req, blocks, partial_hash = cache_full_block_and_partial_tail(
+        list(range(partial_num_tokens)), block_size=block_size
     )
-    full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0]
+    full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, block_size)[0]
 
     assert pool.get_cached_block(full_hash, [0]) == [blocks[0]]
-    assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]]
+    assert pool.get_cached_block(partial_hash, [0]) == [blocks[1]]
 
     pool.evict_blocks({blocks[0].block_id, blocks[1].block_id})
 
     assert pool.get_cached_block(full_hash, [0]) is None
-    assert pool.get_cached_block(partial_hash_10, [0]) is None
+    assert pool.get_cached_block(partial_hash, [0]) is None
     assert pool.cached_block_hashes_by_block == {}
 
 
-def test_partial_block_promotes_to_direct_full_block_hash():
+@pytest.mark.parametrize("dcp_world_size", [1, 2, 4])
+def test_partial_block_promotes_to_direct_full_block_hash(dcp_world_size: int):
     hash_block_size = 2
-    block_size = 6
+    block_size = 6 * dcp_world_size
     kv_cache_group_id = 0
-    token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
+    partial_num_tokens = 2 * block_size - hash_block_size
+    token_ids = list(range(partial_num_tokens))
     req = make_request("0", token_ids, hash_block_size, sha256)
     pool = BlockPool(
         num_gpu_blocks=3,
@@ -426,27 +433,22 @@ def test_partial_block_promotes_to_direct_full_block_hash():
         block_size=block_size,
         kv_cache_group_id=kv_cache_group_id,
     )
-    partial_hash_10 = boundary_hash(req, hash_block_size, 10)
+    partial_hash = boundary_hash(req, hash_block_size, partial_num_tokens)
     assert pool.cache_partial_block(
         request=req,
         block=blocks[1],
-        num_tokens=10,
+        num_tokens=partial_num_tokens,
         kv_cache_group_id=kv_cache_group_id,
         block_size=block_size,
     )
-    assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]]
+    assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) == [blocks[1]]
 
-    req.append_output_token_ids([5, 5])
+    req.append_output_token_ids(list(range(partial_num_tokens, 2 * block_size)))
     full_hashes = BlockHashListWithBlockSize(
         req.block_hashes, hash_block_size, block_size
     )
     promoted_full_hash = full_hashes[1]
-    # The promoted full-block hash is the fine hash at the 12-token boundary,
-    # not a concatenation of the fine hashes inside the block.
-    assert promoted_full_hash == req.block_hashes[12 // hash_block_size - 1]
-    assert promoted_full_hash != BlockHash(
-        req.block_hashes[3] + req.block_hashes[4] + req.block_hashes[5]
-    )
+    assert promoted_full_hash == req.block_hashes[2 * block_size // hash_block_size - 1]
 
     pool.cache_full_blocks(
         request=req,
@@ -457,4 +459,4 @@ def test_partial_block_promotes_to_direct_full_block_hash():
         kv_cache_group_id=kv_cache_group_id,
     )
     assert pool.get_cached_block(promoted_full_hash, [kv_cache_group_id]) == [blocks[1]]
-    assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) is None
+    assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) is None

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm.v1.worker.gpu.model_runner as model_runner_module
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+)
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+
+@pytest.mark.parametrize(
+    ("mamba_cache_mode", "num_speculative_blocks", "expected"),
+    [
+        pytest.param("align", 0, 65_536, id="align-prefix-cache"),
+        pytest.param("none", 7, 8, id="no-prefix-cache-with-speculation"),
+    ],
+)
+def test_initialize_kv_cache_does_not_dcp_shard_mamba_block_table(
+    monkeypatch,
+    mamba_cache_mode: str,
+    num_speculative_blocks: int,
+    expected: int,
+):
+    """Mamba/GDN block-table rows index global positions, unlike DCP KV."""
+
+    max_model_len = 1_048_576
+    attention_block_size = 1_536
+    mamba_block_size = 16
+    dcp_size = 8
+    full_attention_spec = FullAttentionSpec(
+        block_size=attention_block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.bfloat16,
+    )
+    mamba_spec = MambaSpec(
+        shapes=((1,),),
+        dtypes=(torch.bfloat16,),
+        block_size=mamba_block_size,
+        mamba_cache_mode=mamba_cache_mode,
+        num_speculative_blocks=num_speculative_blocks,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=1,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["attention"], full_attention_spec),
+            KVCacheGroupSpec(["kda"], mamba_spec),
+        ],
+    )
+    vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size),
+        cache_config=SimpleNamespace(mamba_cache_mode=mamba_cache_mode),
+    )
+    runner = SimpleNamespace(
+        max_model_len=max_model_len,
+        is_encoder_decoder=False,
+        vllm_config=vllm_config,
+    )
+
+    class _CapturedWidths(Exception):
+        pass
+
+    captured: list[int] = []
+
+    def capture_width(max_num_blocks: int, *_args, **_kwargs) -> int:
+        captured.append(max_num_blocks)
+        if len(captured) == 2:
+            raise _CapturedWidths
+        return max_num_blocks
+
+    monkeypatch.setattr(model_runner_module, "get_block_table_width", capture_width)
+
+    with pytest.raises(_CapturedWidths):
+        GPUModelRunner.initialize_kv_cache(runner, kv_cache_config)
+
+    # Attention KV is local to one of eight DCP ranks; KDA state is replicated
+    # and therefore needs one table entry for every global 16-token page.
+    assert captured == [86, expected]
+
+
+def test_append_block_ids_rejects_write_past_row_capacity():
+    """Reject an oversized staged write before it can corrupt the next row."""
+
+    class _BlockTable:
+        gpu = torch.empty((2, 4), dtype=torch.int32)
+
+        def stage_write(self, *_args):
+            pytest.fail("an oversized write must not be staged")
+
+    block_tables = BlockTables.__new__(BlockTables)
+    block_tables.num_kv_cache_groups = 1
+    block_tables.blocks_per_kv_block = [1]
+    block_tables.block_tables = [_BlockTable()]
+    block_tables.num_blocks = SimpleNamespace(
+        np=torch.tensor([[0, 3]], dtype=torch.int32)
+    )
+
+    with pytest.raises(
+        RuntimeError,
+        match=r"request 1, group 0 exceeds row capacity \(5 > 4\)",
+    ):
+        block_tables.append_block_ids(
+            req_index=1,
+            new_block_ids=([4, 5],),
+            overwrite=False,
+        )
+
+    assert block_tables.num_blocks.np[0, 1] == 3

vllm/v1/core/kv_cache_coordinator.py

@@ -619,15 +619,22 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
                     "full-attention and Mamba groups, got: "
                     f"{type(g.kv_cache_spec).__name__}."
                 )
-        # Fine-grained hash hits require Mamba "align", no context
-        # parallelism, and compatible cache managers in every group.
+        # Fine-grained hash hits require Mamba "align" and compatible cache
+        # managers in every group. TP needs hashing finer than the Mamba block;
+        # DCP accepts equality because it scales the effective full-attention
+        # block instead.
         has_partial_mamba_group = any(
             isinstance(g.kv_cache_spec, MambaSpec)
             and g.kv_cache_spec.mamba_cache_mode == "align"
-            and g.kv_cache_spec.block_size > hash_block_size
+            and (
+                (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size)
+                or (
+                    dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size
+                )
+            )
             for g in kv_cache_config.kv_cache_groups
         )
-        self.enable_partial_hash_hits = dcp_world_size == 1 and has_partial_mamba_group
+        self.enable_partial_hash_hits = has_partial_mamba_group
         if self.enable_partial_hash_hits:
             unsupported_partial_hit_managers = {
                 type(manager).__name__

vllm/v1/worker/gpu/block_table.py

@@ -116,8 +116,15 @@ class BlockTables:
             bpk = self.blocks_per_kv_block[i]
             if bpk > 1:
                 block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]
+            end = start + len(block_ids)
+            row_capacity = self.block_tables[i].gpu.shape[1]
+            if end > row_capacity:
+                raise RuntimeError(
+                    f"Block table write for request {req_index}, group {i} exceeds "
+                    f"row capacity ({end} > {row_capacity})"
+                )
             self.block_tables[i].stage_write(req_index, start, block_ids)
-            self.num_blocks.np[i, req_index] = start + len(block_ids)
+            self.num_blocks.np[i, req_index] = end
 
     def apply_staged_writes(self) -> None:
         if self.num_kv_cache_groups == 0:

vllm/v1/worker/gpu/model_runner.py

@@ -60,7 +60,6 @@ from vllm.multimodal.encoder_budget import (
 )
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import SupportedTask
-from vllm.utils.math_utils import cdiv
 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -517,17 +516,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         for kv_cache_group in kv_cache_config.kv_cache_groups:
             spec = kv_cache_group.kv_cache_spec
             block_sizes.append(spec.block_size)
-            # When using DCP, each request's KV cache is sharded among different ranks.
-            # As a result, one block on the current rank covers `block_size * cp_size`
-            # tokens in the full, global (unsharded) sequence.
-            max_num_blocks = cdiv(
-                block_table_max_model_len, spec.block_size * self.dcp_size
+            # Let each cache type account for CP. Attention KV is DCP-sharded,
+            # while Mamba/GDN recurrent state is replicated across DCP ranks.
+            max_num_blocks = spec.max_num_blocks_per_req(
+                self.vllm_config, block_table_max_model_len
             )
-            # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens
+            # Preserve each cache type's alignment requirements after applying
+            # its topology-aware block-table width.
             if isinstance(spec, MambaSpec):
-                max_num_blocks = (
-                    max_num_blocks if self.cache_config.enable_prefix_caching else 1
-                ) + spec.num_speculative_blocks
                 max_num_blocks = get_block_table_width(
                     max_num_blocks, spec.block_size, token_alignment=None
                 )