[Bugfix][V1] Use the Mamba cache group's block size for align-mode chunk splitting

tests/test_config.py

@@ -2565,6 +2565,26 @@ def test_eagle_block_drop_can_be_disabled_without_disabling_eagle(
     assert speculative_config.use_eagle_block_drop() is not disable_eagle_block_drop
 
 
+@pytest.mark.parametrize("method", ["dflash", "dspark"])
+def test_dflash_family_does_not_drop_trailing_prefix_cache_block(method: str):
+    # DFlash/DSpark draft from their own KV cache and never write target
+    # blocks, so the EAGLE volatile-trailing-block drop must not apply to
+    # them: applying it backs the last mamba-aligned cache position off one
+    # block, the final block-aligned mamba state never materializes, and
+    # prefix-cache / offload-tier lookups collapse ("stores but never serves
+    # a hit", #53505). Start from an ngram config to avoid loading model
+    # metadata: these predicates depend only on the speculative method.
+    speculative_config = SpeculativeConfig(
+        method="ngram",
+        num_speculative_tokens=3,
+    )
+    speculative_config.method = method
+
+    assert speculative_config.use_eagle()
+    assert not speculative_config.use_eagle_preserves_target_kv_cache()
+    assert speculative_config.use_eagle_block_drop() is False
+
+
 def test_draft_sample_method_gumbel_is_rejected():
     with pytest.raises(ValidationError):
         SpeculativeConfig(

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -167,16 +167,18 @@ def test_dcp_fine_hit_retention_uses_hash_alignment_without_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."""
+    """Chunk ends with partial hits on: a stop at every crossed block
+    boundary (each materializes a mamba state slot; no spanning 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 -> 512 -> ... -> 9728 -> 9984 -> 10000."""
     block_size = 512
     scheduler_block_size = block_size * dcp_world_size
     hash_block_size = 32
     mock = SimpleNamespace(
         block_size=scheduler_block_size,
         cache_config=SimpleNamespace(block_size=block_size),
+        mamba_state_block_size=block_size,
         max_num_scheduled_tokens=8192,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
@@ -191,10 +193,12 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
 
     req = make_request("0", [0] * 10000, hash_block_size, sha256)
     req.num_computed_tokens = 0
-    assert split(self=mock, request=req, num_new_tokens=8192) == 8192
-    req.num_computed_tokens = 8192
-    # Stop at the last block boundary (9728).
-    assert split(self=mock, request=req, num_new_tokens=1808) == 1536
+    # Stop at the next block boundary even from an aligned start (no
+    # spanning chunks: interior state slots must materialize).
+    assert split(self=mock, request=req, num_new_tokens=8192) == 512
+    req.num_computed_tokens = 9728 - 512
+    # Full-budget requests still advance exactly one block per step.
+    assert split(self=mock, request=req, num_new_tokens=10000) == 512
     req.num_computed_tokens = 9728
     # Extra stop at the prompt's last hash boundary (9984).
     assert split(self=mock, request=req, num_new_tokens=272) == 256
@@ -227,6 +231,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():
     mock = SimpleNamespace(
         block_size=block_size,
         cache_config=SimpleNamespace(block_size=block_size),
+        mamba_state_block_size=block_size,
         max_num_scheduled_tokens=token_budget,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
@@ -266,6 +271,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
     mock = SimpleNamespace(
         block_size=block_size,
         cache_config=SimpleNamespace(block_size=block_size),
+        mamba_state_block_size=block_size,
         max_num_scheduled_tokens=token_budget,
         scheduler_config=SimpleNamespace(
             long_prefill_token_threshold=long_prefill_threshold

tests/v1/core/test_mamba_align_chunk_split.py

@@ -139,6 +139,10 @@ def _split(
     stub = SimpleNamespace(
         block_size=MAMBA_BLOCK_SIZE,
         cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE),
+        # Set by `Scheduler.__init__` from the mamba group's spec; the equal
+        # geometry of these tests keeps it identical to `cache_config`.
+        mamba_state_block_size=MAMBA_BLOCK_SIZE,
+        use_eagle=use_eagle,
         use_eagle_block_drop=use_eagle_block_drop,
         max_num_scheduled_tokens=max_num_scheduled_tokens,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
@@ -213,14 +217,21 @@ def test_partial_checkpoint_resume_stops_at_mamba_block_boundary() -> None:
     )
 
 
-def test_disabling_eagle_block_drop_keeps_the_trailing_cache_boundary() -> None:
+def test_split_is_boundary_locked_independently_of_eagle_block_drop() -> None:
     (request,) = create_requests(1, num_tokens=3602, block_size=ATTN_BLOCK_SIZE)
 
     with_drop = _split(request, request.num_tokens, use_eagle_block_drop=True)
     without_drop = _split(request, request.num_tokens, use_eagle_block_drop=False)
 
+    # Unconditional boundary stops lock chunk ends to the state grid, so in
+    # non-checkpoint mode the split no longer depends on the Eagle
+    # block-drop back-off: an aligned start ends at the next state boundary
+    # either way instead of spanning to the trailing cache boundary (the
+    # spanning chunk is exactly the k-1-null-interior-slots case above).
+    # The back-off still lowers `last_cache_position` for the partial-tail
+    # and checkpoint interactions; it just can no longer move a chunk end.
     assert with_drop == MAMBA_BLOCK_SIZE
-    assert without_drop == 2 * MAMBA_BLOCK_SIZE
+    assert without_drop == MAMBA_BLOCK_SIZE
 
 
 def _run_chunked_prefill(
@@ -392,3 +403,217 @@ def test_unaligned_resume_never_runs_past_its_block(
             f"intermediate chunk end {end} is neither block-aligned nor the "
             f"partial-tail stop ({tail_stop})"
         )
+
+
+# Heterogeneous layouts: `cache_config.block_size` is the minimum over all
+# groups and can be finer than the mamba state grid (page-size matching, a
+# drafter/attention group with a smaller block, or an explicit --block-size;
+# see also #53142 for the worker-side sibling). Production geometry that
+# exposed this: draft attention block 816, mamba block 1648 — 816k equals
+# 1648m only at the LCM (84048), so chunk stops on the generic grid never
+# materialize a state and the mamba group publishes no prefix hashes.
+HETERO_ATTN_BLOCK_SIZE = 816
+HETERO_MAMBA_BLOCK_SIZE = 1648
+HETERO_SCHEDULER_BLOCK_SIZE = 84048  # lcm(816, 1648)
+HETERO_MAMBA_GROUP_ID = 1
+
+
+def _make_heterogeneous_kv_cache_manager() -> KVCacheManager:
+    config = KVCacheConfig(
+        num_blocks=10000,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["attention_layer"],
+                FullAttentionSpec(
+                    block_size=HETERO_ATTN_BLOCK_SIZE,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba_layer"],
+                MambaSpec(
+                    block_size=HETERO_MAMBA_BLOCK_SIZE,
+                    shapes=((1, 1),),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                    num_speculative_blocks=NUM_SPEC,
+                ),
+            ),
+        ],
+    )
+    return KVCacheManager(
+        config,
+        max_model_len=262144,
+        scheduler_block_size=HETERO_SCHEDULER_BLOCK_SIZE,
+        hash_block_size=ATTN_BLOCK_SIZE,
+        enable_caching=True,
+        use_eagle=True,
+    )
+
+
+def _hetero_split(request: Request, num_new_tokens: int) -> int:
+    """Real split on a heterogeneous-layout stub `self`.
+
+    `cache_config.block_size` is the group minimum (816); the mamba state
+    grid is 1648. `mamba_state_block_size` mirrors the attribute
+    `Scheduler.__init__` derives from the mamba group's spec.
+    """
+    stub = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=HETERO_ATTN_BLOCK_SIZE),
+        mamba_state_block_size=HETERO_MAMBA_BLOCK_SIZE,
+        use_eagle=True,
+        use_eagle_block_drop=True,
+        max_num_scheduled_tokens=16384,
+        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
+        mamba_partial_cache_hit=False,
+        hash_block_size=ATTN_BLOCK_SIZE,
+        mamba_has_prefill_checkpoint_blocks=False,
+        mamba_prefill_checkpoint_alignment=None,
+    )
+    return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens)
+
+
+def test_heterogeneous_block_sizes_stop_chunks_on_the_mamba_grid() -> None:
+    """Chunk ends must follow MambaSpec.block_size, not cache_config.block_size."""
+    prompt_len = 4 * HETERO_MAMBA_BLOCK_SIZE + 70
+    (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
+    pos, ends = 0, []
+    while pos < prompt_len:
+        request.num_computed_tokens = pos
+        num_new = _hetero_split(request, prompt_len - pos)
+        assert num_new > 0, f"no progress at {pos}"
+        pos += num_new
+        ends.append(pos)
+    for end in ends[:-1]:
+        assert end % HETERO_MAMBA_BLOCK_SIZE == 0, (
+            f"intermediate chunk end {end} is off the mamba state grid; the "
+            f"worker never commits a state there (816-grid stop)"
+        )
+    assert ends[0] % HETERO_MAMBA_BLOCK_SIZE == 0, (
+        f"first chunk ended at {ends[0]}, off the mamba state grid "
+        f"(816-grid stop); ends={ends}"
+    )
+
+
+def test_aligned_start_does_not_span_multiple_state_blocks() -> None:
+    """A full token budget must not skip interior state boundaries.
+
+    With the stop conditional on a mid-block start, a chunk beginning exactly
+    on a boundary could run to the budget-clamped end whenever the budget
+    exceeds one block, crossing k boundaries and leaving the k-1 interior
+    state slots permanently null (one state column is materialized per step).
+    """
+    prompt_len = 5 * HETERO_MAMBA_BLOCK_SIZE + 30
+    (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
+    # Request the FULL remaining prompt every step, as a solo prefill with a
+    # budget larger than 2 blocks would (no per-block rationing).
+    pos, ends = 0, []
+    while pos < prompt_len:
+        request.num_computed_tokens = pos
+        num_new = _hetero_split(request, prompt_len - pos)
+        assert num_new > 0, f"no progress at {pos}"
+        pos += num_new
+        ends.append(pos)
+    expected_grid_ends = [
+        (i + 1) * HETERO_MAMBA_BLOCK_SIZE
+        for i in range(prompt_len // HETERO_MAMBA_BLOCK_SIZE)
+    ]
+    materialized = [e for e in ends if e % HETERO_MAMBA_BLOCK_SIZE == 0]
+    assert materialized == expected_grid_ends, (
+        f"state-grid chunk ends {materialized} != consecutive boundaries "
+        f"{expected_grid_ends}; interior slots stayed null (spanning chunk)"
+    )
+
+
+def _hetero_prefill(prompt_len: int) -> tuple[KVCacheManager, Request, dict[int, int]]:
+    """Prefill one request through the real manager under the hetero layout.
+
+    Budgets one mamba block per step, the shape the fixed split produces for
+    a request sharing the token budget (and the equal-geometry tests above).
+    """
+    manager = _make_heterogeneous_kv_cache_manager()
+    (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
+    mamba_manager = manager.coordinator.single_type_managers[HETERO_MAMBA_GROUP_ID]
+    state_at: dict[int, int] = {}
+    while request.num_computed_tokens < request.num_tokens:
+        computed = request.num_computed_tokens
+        budget = min(HETERO_MAMBA_BLOCK_SIZE, request.num_tokens - computed)
+        num_new = _hetero_split(request, budget)
+        assert num_new > 0, f"no progress at {computed}"
+        assert (
+            manager.allocate_slots(request, num_new, num_lookahead_tokens=NUM_SPEC)
+            is not None
+        )
+        request.num_computed_tokens = computed + num_new
+        blocks = mamba_manager.req_to_blocks[request.request_id]
+        running = cdiv(request.num_computed_tokens, HETERO_MAMBA_BLOCK_SIZE) - 1
+        state_at[blocks[running].block_id] = request.num_computed_tokens
+    return manager, request, state_at
+
+
+def test_heterogeneous_block_sizes_publish_mamba_states_for_immediate_reuse() -> None:
+    """Mamba states must be published where the worker materializes them.
+
+    On the generic 816 grid no chunk ever ends on the 1648 state grid, so the
+    mamba group publishes zero prefix hashes and a request repeating the same
+    prompt immediately reuses nothing.
+    """
+    prompt_len = 4 * HETERO_MAMBA_BLOCK_SIZE + 70
+    manager, first, state_at = _hetero_prefill(prompt_len)
+    mamba_manager = manager.coordinator.single_type_managers[HETERO_MAMBA_GROUP_ID]
+
+    published = [
+        block.block_hash_num_tokens
+        for block in mamba_manager.req_to_blocks[first.request_id]
+        if not block.is_null and block.block_hash is not None
+    ]
+    full_block_entries = [p for p in published if p % HETERO_MAMBA_BLOCK_SIZE == 0]
+    assert full_block_entries, (
+        f"no mamba full-block state published (entries={published}); "
+        f"immediate reuse on this layout is impossible"
+    )
+    # Safety: every hashed slot holds the state its hash claims.
+    for block in mamba_manager.req_to_blocks[first.request_id]:
+        if block.is_null or block.block_hash is None:
+            continue
+        assert state_at.get(block.block_id) == block.block_hash_num_tokens, (
+            f"mamba slot hashed as state@{block.block_hash_num_tokens} but "
+            f"holds state@{state_at.get(block.block_id)}"
+        )
+
+    (second,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
+    # The scheduler hashes prompt token ids deterministically, so mirror it:
+    second.all_token_ids = first.all_token_ids
+    second.block_hashes = first.block_hashes
+    _, num_computed, _ = manager.get_computed_blocks(second)
+    assert num_computed >= HETERO_MAMBA_BLOCK_SIZE, (
+        f"immediate repeat of the same prompt reused {num_computed} tokens "
+        f"(expected at least one legal mamba boundary)"
+    )
+    assert num_computed % HETERO_MAMBA_BLOCK_SIZE == 0
+
+
+def test_equal_block_sizes_control_still_reuses() -> None:
+    """Control: with cache_config.block_size == MambaSpec.block_size the
+    behavior is the same before and after a heterogeneous-layout fix."""
+    prompt_len = 2 * MAMBA_BLOCK_SIZE + 70
+    manager = _make_hybrid_kv_cache_manager()
+    (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
+    while request.num_computed_tokens < request.num_tokens:
+        computed = request.num_computed_tokens
+        num_new = _split(request, request.num_tokens - computed)
+        assert num_new > 0
+        assert (
+            manager.allocate_slots(request, num_new, num_lookahead_tokens=NUM_SPEC)
+            is not None
+        )
+        request.num_computed_tokens = computed + num_new
+    (second,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
+    second.all_token_ids = request.all_token_ids
+    second.block_hashes = request.block_hashes
+    _, num_computed, _ = manager.get_computed_blocks(second)
+    assert num_computed >= MAMBA_BLOCK_SIZE
+    assert num_computed % MAMBA_BLOCK_SIZE == 0

tests/v1/core/test_prefix_caching.py

@@ -1164,6 +1164,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
     mock = SimpleNamespace(
         block_size=block_size,
         cache_config=SimpleNamespace(block_size=block_size),
+        mamba_state_block_size=block_size,
         max_num_scheduled_tokens=3 * block_size,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
@@ -1178,7 +1179,9 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
         request=req_2,
         num_new_tokens=3 * block_size,
     )
-    assert num_new_tokens_adjusted == 2 * block_size  # adjust to the common prefix
+    # The boundary stop is unconditional, so from num_computed 0 the chunk
+    # ends at the first block boundary covered by the junction (one block).
+    assert num_new_tokens_adjusted == block_size  # stop at the next boundary
 
     manager.allocate_slots(req_2, 3 * block_size, 0, computed_blocks)
     # Cleanup

vllm/config/speculative.py

@@ -1862,8 +1862,23 @@ def use_eagle(self) -> bool:
         return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark")
 
     def use_eagle_block_drop(self) -> bool:
-        """Whether volatile trailing cache blocks should be discarded."""
-        return self.use_eagle() and not self.disable_eagle_block_drop
+        """Whether volatile trailing cache blocks should be discarded.
+
+        Only eagle-family drafters share (and pollute via lookahead KV write)
+        the target's full-attention KV cache groups; DFlash/DSpark draft from
+        their own KV cache and never write target blocks (#53477), so the
+        drop applies to eagle/eagle3/mtp only, unless explicitly disabled.
+        """
+        return (
+            self.use_eagle_preserves_target_kv_cache()
+            and not self.disable_eagle_block_drop
+        )
+
+    def use_eagle_preserves_target_kv_cache(self) -> bool:
+        # Only eagle-family drafters share (and pollute via lookahead KV
+        # write) the target's full-attention KV cache groups; DFlash/DSpark
+        # draft from their own KV cache and never write target blocks.
+        return self.method in ("eagle", "eagle3", "mtp")
 
     def use_dflash(self) -> bool:
         return self.method == "dflash"