[Bugfix][Prefix caching] Fix inconsistent hybrid Mamba prefix cache boundaries

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"),
     )
+    vllm_config.cache_config.hash_block_size = prefix_match_unit or block_size
+    vllm_config.cache_config.cache_hit_alignment_tokens = (
+        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,
+    cache_hit_alignment_tokens: 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,11 @@ 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
+    vllm_config.cache_config.hash_block_size = hash_block_size
+    vllm_config.cache_config.cache_hit_alignment_tokens = (
+        cache_hit_alignment_tokens or hash_block_size
+    )
     builder = builder_cls(
         kv_cache_spec=MambaSpec(
             block_size=mamba_block_size,
@@ -220,6 +227,8 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
         vllm_config=builder_config,
         device=DEVICE,
     )
+    builder_config.cache_config.hash_block_size = BLOCK_SIZE
+    builder_config.cache_config.cache_hit_alignment_tokens = 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 +296,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,hit_alignment,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, hit_alignment, 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,
+        cache_hit_alignment_tokens=hit_alignment,
+        hash_block_size=16,
         device=device,
     )
     assert isinstance(builder, KimiK3KDAMetadataBuilder)
@@ -313,26 +349,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 +433,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

@@ -8,6 +8,7 @@
 from collections.abc import Iterator
 from dataclasses import dataclass
 from math import lcm
+from uuid import uuid4
 
 import pytest
 import requests
@@ -353,6 +354,45 @@ def _format(turns: list[Turn]) -> str:
     return "\n".join(rows)
 
 
+@pytest.mark.parametrize(
+    "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(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(7296, id="eagle-double-cap"),
+    ],
+)
+def test_dspark_exact_resend_reuses_prompt_checkpoint(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"]
+    with _serve(deployment, mode) as servers:
+        url = servers[0].url_root
+        cold = _complete(url, prompt, salt, max_tokens=1)
+        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,8 @@ def _create_checkpoint_builder_and_batch(
         vllm_config=vllm_config,
         device=DEVICE,
     )
+    vllm_config.cache_config.hash_block_size = prefix_match_unit or 64
+    vllm_config.cache_config.cache_hit_alignment_tokens = prefix_match_unit or 64
     batch = BatchSpec(
         seq_lens=[65, 100, 100], query_lens=[3 if num_spec else 1, 100, 99]
     )
@@ -431,7 +433,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,15 @@ 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,
     )
+    vllm_config.cache_config.hash_block_size = MAMBA_BLOCK_SIZE
+    vllm_config.cache_config.cache_hit_alignment_tokens = MAMBA_BLOCK_SIZE
+    return builder
 
 
 def _build(builder, seq_lens, query_lens):