[Bugfix][Prefix caching] Fix inconsistent hybrid Mamba prefix cache boundaries and incomplete tail saves

.buildkite/test_areas/models_basic.yaml

@@ -155,14 +155,52 @@ steps:
   - vllm/v1/core/
   - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
   - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+  - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/
+  - vllm/distributed/mooncake_store.py
   - vllm/v1/simple_kv_offload/
   - vllm/v1/spec_decode/
   - vllm/model_executor/layers/mamba/
   - tests/models/kimi_k3/test_prefix_cache.py
   - .buildkite/scripts/install-kv-connectors.sh
   commands:
   - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-  - pytest -v -s models/kimi_k3/test_prefix_cache.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+  - |
+    set -e
+    export MC_MAX_MR_SIZE=4294967296
+    export MOONCAKE_CONFIG_PATH=$$(mktemp /vllm-workspace/mooncake-prefix-cache.XXXXXX.json)
+    MOONCAKE_PORTS=$$(uv run --no-project python - <<'PY'
+    import json
+    import os
+    import socket
+    from pathlib import Path
+
+    with socket.socket() as rpc, socket.socket() as metrics:
+        rpc.bind(("127.0.0.1", 0))
+        metrics.bind(("127.0.0.1", 0))
+        port = rpc.getsockname()[1]
+        metrics_port = metrics.getsockname()[1]
+    config = {
+        "metadata_server": "P2PHANDSHAKE",
+        "master_server_address": f"127.0.0.1:{port}",
+        "protocol": "rdma",
+        "device_name": os.getenv("MOONCAKE_DEVICE_NAME", "mlx5_12"),
+        "global_segment_size": "32GB",
+        "local_buffer_size": "4GB",
+    }
+    Path(os.environ["MOONCAKE_CONFIG_PATH"]).write_text(json.dumps(config))
+    print(port, metrics_port)
+    PY
+    )
+    read -r MOONCAKE_MASTER_PORT MOONCAKE_METRICS_PORT <<EOF
+    $$MOONCAKE_PORTS
+    EOF
+    mooncake_master --port="$$MOONCAKE_MASTER_PORT" \
+      --metrics_port="$$MOONCAKE_METRICS_PORT" --enable_metric_reporting=false \
+      --default_kv_lease_ttl=60000 \
+      > "$${MOONCAKE_CONFIG_PATH}.master.log" 2>&1 &
+    MOONCAKE_MASTER_PID=$$!
+    trap 'kill "$$MOONCAKE_MASTER_PID" 2>/dev/null || true; wait "$$MOONCAKE_MASTER_PID" 2>/dev/null || true' EXIT
+    pytest -v -s models/kimi_k3/test_prefix_cache.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 2
 
 - label: ":nvidia: (H200) GLM5Next"

tests/models/kimi_k3/test_amd_kda_checkpoint.py

@@ -75,7 +75,7 @@ def _builder(
     vllm_config.cache_config.prefix_match_unit = prefix_match_unit
     if num_spec:
         vllm_config.speculative_config = _StubSpeculativeConfig(num_spec)
-    return KimiK3ROCmKDAMetadataBuilder(
+    builder = KimiK3ROCmKDAMetadataBuilder(
         kv_cache_spec=MambaSpec(
             block_size=block_size,
             shapes=((16, 64),),
@@ -91,6 +91,11 @@ def _builder(
         vllm_config=vllm_config,
         device=torch.device("cuda"),
     )
+    builder.checkpoint_builder.set_block_sizes(
+        hash_block_size=prefix_match_unit or block_size,
+        mamba_ckpt_block_size=prefix_match_unit or block_size,
+    )
+    return builder
 
 
 def _build(batch: BatchSpec, block_size: int = BLOCK_SIZE, **kw):

tests/models/kimi_k3/test_kda_metadata.py

@@ -100,6 +100,8 @@ def _make_builder(
     prefix_match_unit: int | None = None,
     use_eagle: bool = False,
     disable_eagle_block_drop: bool = False,
+    mamba_ckpt_block_size: int | None = None,
+    hash_block_size: int | None = None,
 ) -> AttentionMetadataBuilder:
     vllm_config = create_vllm_config(
         model_name="Qwen/Qwen3.5-0.8B",
@@ -122,6 +124,7 @@ def _make_builder(
     vllm_config.cache_config.use_replayssm = use_recoverssm
     vllm_config.cache_config.use_kda_recoverssm = use_recoverssm
     vllm_config.cache_config.prefix_match_unit = prefix_match_unit
+    hash_block_size = hash_block_size or prefix_match_unit or mamba_block_size
     builder = builder_cls(
         kv_cache_spec=MambaSpec(
             block_size=mamba_block_size,
@@ -138,6 +141,10 @@ def _make_builder(
         vllm_config=vllm_config,
         device=device,
     )
+    builder.checkpoint_builder.set_block_sizes(
+        hash_block_size=hash_block_size,
+        mamba_ckpt_block_size=mamba_ckpt_block_size or hash_block_size,
+    )
     if use_recoverssm:
         assert isinstance(builder, KimiK3KDAMetadataBuilder)
         builder.recoverssm_context = Mock()
@@ -220,6 +227,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
         vllm_config=builder_config,
         device=DEVICE,
     )
+    builder.checkpoint_builder.set_block_sizes(BLOCK_SIZE, BLOCK_SIZE)
 
     # An all-prefill speculative batch used to leave an all-false spec mask
     # alive, then access active_non_spec_mask_cpu before it was initialized.
@@ -287,18 +295,45 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
-def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
+@pytest.mark.parametrize(
+    "batch,prefix_match_unit,mamba_ckpt_block_size,expected_offsets,expected_states",
+    [
+        (
+            BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]),
+            None,
+            16,
+            [48, 0],
+            [2, NULL_BLOCK_ID],
+        ),
+        (
+            BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]),
+            16,
+            16,
+            [48, 0],
+            [2, NULL_BLOCK_ID],
+        ),
+        (BatchSpec(seq_lens=[530], query_lens=[530]), None, 128, [512], [32]),
+        (BatchSpec(seq_lens=[530], query_lens=[50]), None, 128, [32], [32]),
+        (BatchSpec(seq_lens=[530], query_lens=[530]), 16, 16, [528], [32]),
+        (BatchSpec(seq_lens=[530], query_lens=[50]), 16, 16, [48], [32]),
+    ],
+)
+def test_internal_checkpoint_metadata_targets_last_aligned_boundary(
+    batch, prefix_match_unit, mamba_ckpt_block_size, expected_offsets, expected_states
+):
     device = torch.device("cuda")
-    batch = BatchSpec(seq_lens=[50, 32], query_lens=[50, 16])
     common_attn_metadata = create_common_attn_metadata(
         batch, BLOCK_SIZE, device, arange_block_indices=True
-    ).replace(is_prefilling=torch.tensor([True, True]))
+    ).replace(is_prefilling=torch.tensor([True] * len(batch.seq_lens)))
     builder = _make_builder(
         KimiK3KDAMetadataBuilder,
         num_speculative_tokens=0,
         full_cuda_graph=False,
         mamba_cache_mode="align",
         num_prefill_checkpoint_blocks=1,
+        prefix_match_unit=prefix_match_unit,
+        mamba_ckpt_block_size=mamba_ckpt_block_size,
+        hash_block_size=16,
         device=device,
     )
     assert isinstance(builder, KimiK3KDAMetadataBuilder)
@@ -313,26 +348,33 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
     assert actual.checkpoint is not None
     torch.testing.assert_close(
         actual.checkpoint.state_indices,
-        torch.tensor([2, NULL_BLOCK_ID], dtype=torch.int32, device=device),
+        torch.tensor(expected_states, dtype=torch.int32, device=device),
     )
     torch.testing.assert_close(
         actual.checkpoint.checkpoint_offsets,
-        torch.tensor([48, 0], dtype=torch.int32, device=device),
+        torch.tensor(expected_offsets, dtype=torch.int32, device=device),
     )
 
 
 @pytest.mark.parametrize(
-    ("disable_eagle_block_drop", "prefix_match_unit", "expected_offset"),
-    [(False, 16, 80), (True, 16, 96), (False, 8, None)],
+    ("seq_len", "disable_eagle_block_drop", "prefix_match_unit", "expected_offset"),
+    [
+        (100, False, 16, 80),
+        (100, True, 16, 96),
+        (100, False, 8, None),
+        (96, False, 16, 80),
+        (96, True, 16, 80),
+    ],
 )
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_spec_internal_checkpoint_metadata_targets_replay_boundary(
+    seq_len: int,
     disable_eagle_block_drop: bool,
     prefix_match_unit: int,
     expected_offset: int | None,
 ) -> None:
     device = torch.device("cuda")
-    batch = BatchSpec(seq_lens=[100], query_lens=[100])
+    batch = BatchSpec(seq_lens=[seq_len], query_lens=[seq_len])
     common_attn_metadata = create_common_attn_metadata(
         batch, BLOCK_SIZE, device, arange_block_indices=True
     )
@@ -390,6 +432,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():
         full_cuda_graph=False,
         mamba_cache_mode="align",
         num_prefill_checkpoint_blocks=1,
+        prefix_match_unit=16,
         device=device,
     )
     assert isinstance(builder, KimiK3KDAMetadataBuilder)

tests/models/kimi_k3/test_prefix_cache.py

@@ -4,10 +4,13 @@
 
 import contextlib
 import json
+import os
 import random
+import time
 from collections.abc import Iterator
 from dataclasses import dataclass
 from math import lcm
+from uuid import uuid4
 
 import pytest
 import requests
@@ -59,6 +62,16 @@
 }
 
 NIXL = {"kv_connector": "NixlConnector", "kv_role": "kv_both"}
+MOONCAKE = {
+    "kv_connector": "MooncakeStoreConnector",
+    "kv_role": "kv_both",
+    "kv_load_failure_policy": "recompute",
+    "kv_connector_extra_config": {
+        "load_async": True,
+        "lookup_async": True,
+        "enable_cross_layers_blocks": False,
+    },
+}
 # SimpleCPUOffloadConnector as in the Kimi-K3 recipe.
 OFFLOAD = {
     "kv_connector": "SimpleCPUOffloadConnector",
@@ -242,10 +255,16 @@ def _tokens_text_logprobs(response: dict):
 
 
 def _reset_gpu_prefix_cache(url: str) -> None:
-    response = requests.post(
-        f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60
-    )
-    response.raise_for_status()
+    deadline = time.monotonic() + 60
+    while True:
+        response = requests.post(
+            f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60
+        )
+        response.raise_for_status()
+        if response.json()["success"]:
+            return
+        assert time.monotonic() < deadline, "GPU cache reset timed out"
+        time.sleep(0.1)
 
 
 @dataclass(frozen=True)
@@ -353,6 +372,57 @@ def _format(turns: list[Turn]) -> str:
     return "\n".join(rows)
 
 
+@pytest.mark.parametrize(
+    "deployment_name,prompt_len",
+    [
+        # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned:
+        # ensure lookup reads the cached attention proof at 7040 despite
+        # the final hit limit of 7039, then applies EAGLE's 128-token drop.
+        pytest.param("plain", 7040, id="eagle-proof-scan"),
+        # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple:
+        # ensure the Mamba checkpoint is 7168, not 7040 from applying both
+        # the P - 1 cap and EAGLE's 128-token drop.
+        pytest.param("plain", 7296, id="eagle-double-cap"),
+        # Blocks=6144, PMU=128. P=7449 needs Mamba state at 7296 and attention
+        # proof at 7424, both beyond the normal Mooncake save boundary of 6144:
+        # ensure they remain reusable after GPU cache reset.
+        pytest.param("mooncake", 7449, id="mooncake-tail-proof"),
+    ],
+)
+def test_dspark_exact_resend_reuses_prompt_checkpoint(
+    deployment_name: str, prompt_len: int
+) -> None:
+    """Resends reuse the Mamba checkpoint and its companion EAGLE attention proof."""
+    mode = Mode(128, True)
+    salt = str(uuid4())
+    rng = random.Random(0)
+    prompt = [rng.randint(1000, 150000) for _ in range(prompt_len)]
+    deployment = DEPLOYMENTS["plain"]
+    if deployment_name == "mooncake":
+        if not os.getenv("MOONCAKE_CONFIG_PATH"):
+            pytest.skip("Mooncake Store requires MOONCAKE_CONFIG_PATH and a master")
+        deployment = Deployment(Instance(kv_config=MOONCAKE), offload=True)
+    with _serve(deployment, mode) as servers:
+        url = servers[0].url_root
+        cold = _complete(url, prompt, salt, max_tokens=1)
+        if deployment.offload:
+            _reset_gpu_prefix_cache(url)
+        resend = _complete(url, prompt, salt, max_tokens=1)
+        recompute = _complete(url, prompt, f"{salt}-recompute", max_tokens=1)
+
+    # EAGLE's dropped hash unit already leaves tokens to recompute on resend.
+    expected = (prompt_len // 128 - 1) * 128
+    actual = _cached(resend)
+    assert _cached(cold) == _cached(recompute) == 0
+    assert actual == expected, f"{prompt_len=}: {actual=}, {expected=}"
+    check_logprobs_close(
+        outputs_0_lst=[_tokens_text_logprobs(recompute)],
+        outputs_1_lst=[_tokens_text_logprobs(resend)],
+        name_0="recompute",
+        name_1="resend",
+    )
+
+
 @pytest.mark.parametrize(
     "name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES]
 )

tests/v1/attention/test_gdn_metadata_builder.py

@@ -385,7 +385,7 @@ def test_cudagraph_capture_batch_stays_decode_only():
 
 
 def _create_checkpoint_builder_and_batch(
-    num_spec: int, prefix_match_unit: int = 16
+    num_spec: int, prefix_match_unit: int | None = 16
 ) -> tuple[GDNAttentionMetadataBuilder, CommonAttentionMetadata, dict]:
     vllm_config = create_vllm_config(
         model_name="Qwen/Qwen3.5-0.8B",
@@ -413,6 +413,10 @@ def _create_checkpoint_builder_and_batch(
         vllm_config=vllm_config,
         device=DEVICE,
     )
+    builder.checkpoint_builder.set_block_sizes(
+        hash_block_size=prefix_match_unit or 64,
+        mamba_ckpt_block_size=prefix_match_unit or 64,
+    )
     batch = BatchSpec(
         seq_lens=[65, 100, 100], query_lens=[3 if num_spec else 1, 100, 99]
     )
@@ -431,7 +435,9 @@ def _create_checkpoint_builder_and_batch(
 
 
 @pytest.mark.parametrize("num_spec", [0, 3])
-@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)])
+@pytest.mark.parametrize(
+    "prefix_match_unit, expected_offset", [(None, 64), (16, 96), (8, 96)]
+)
 def test_checkpoint_metadata_preserves_non_spec_order(
     num_spec, prefix_match_unit, expected_offset
 ):

tests/v1/attention/test_mamba2_prefill_checkpoints.py

@@ -77,7 +77,9 @@ def _create_mamba2_builder(
     num_prefill_checkpoint_blocks: int = 1,
     prefill_checkpoint_alignment: int | None = 1,
 ) -> Mamba2AttentionMetadataBuilder:
-    vllm_config = create_vllm_config(block_size=MAMBA_BLOCK_SIZE)
+    vllm_config = create_vllm_config(
+        model_name="Qwen/Qwen3.5-0.8B", block_size=MAMBA_BLOCK_SIZE
+    )
     vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
     # get_mamba_chunk_size() reads `mamba_chunk_size` then `chunk_size`.
     vllm_config.model_config.hf_text_config.mamba_chunk_size = CHUNK_SIZE
@@ -89,12 +91,17 @@ def _create_mamba2_builder(
         num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks,
         prefill_checkpoint_alignment=prefill_checkpoint_alignment,
     )
-    return Mamba2AttentionMetadataBuilder(
+    builder = Mamba2AttentionMetadataBuilder(
         kv_cache_spec=spec,
         layer_names=["layer.0"],
         vllm_config=vllm_config,
         device=DEVICE,
     )
+    builder.checkpoint_builder.set_block_sizes(
+        hash_block_size=MAMBA_BLOCK_SIZE,
+        mamba_ckpt_block_size=MAMBA_BLOCK_SIZE,
+    )
+    return builder
 
 
 def _build(builder, seq_lens, query_lens):