[Bugfix][Core] Keep Mamba-aligned prefill chunks out of the MTP lookahead window

tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py

@@ -71,6 +71,7 @@ def _stub(manager, block_size, hash_block_size, *, block_drop=True):
         # The EAGLE adjustments key on the block-drop bit, not plain use_eagle:
         # they exist only to compensate for the drop.
         use_eagle_block_drop=block_drop,
+        num_prefill_lookahead=0,
         hash_block_size=hash_block_size,
         mamba_has_prefill_checkpoint_blocks=False,  # forced False under eagle
         mamba_partial_cache_hit=partial_hit,

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -189,6 +189,7 @@ def test_dcp_fine_hit_retention_uses_hash_alignment(
         max_num_scheduled_tokens=8192,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=use_eagle,
+        num_prefill_lookahead=0,
         mamba_partial_cache_hit=True,
         mamba_shared_prefix_checkpoint=False,
         mamba_has_prefill_checkpoint_blocks=bool(checkpoint_blocks),
@@ -244,6 +245,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
         max_num_scheduled_tokens=8192,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
+        num_prefill_lookahead=0,
         hash_block_size=hash_block_size,
         dcp_world_size=dcp_world_size,
         scheduler_block_size=scheduler_block_size,
@@ -294,6 +296,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():
         max_num_scheduled_tokens=token_budget,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
+        num_prefill_lookahead=0,
         hash_block_size=32,
         mamba_partial_cache_hit=False,
         mamba_shared_prefix_checkpoint=False,
@@ -335,6 +338,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
             long_prefill_token_threshold=long_prefill_threshold
         ),
         use_eagle_block_drop=False,
+        num_prefill_lookahead=0,
         hash_block_size=32,
         mamba_partial_cache_hit=False,
         mamba_shared_prefix_checkpoint=False,
@@ -1985,6 +1989,7 @@ def test_mamba_align_split_stops_below_eagle_proof_boundary():
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle=True,
         use_eagle_block_drop=True,
+        num_prefill_lookahead=0,
         hash_block_size=hash_block_size,
         mamba_partial_cache_hit=True,
         mamba_shared_prefix_checkpoint=False,

tests/v1/core/test_mamba_align_chunk_split.py

@@ -84,6 +84,7 @@ def test_mamba_prefill_checkpoint_valid(
 
 def _make_hybrid_kv_cache_manager(
     num_prefill_checkpoint_blocks: int = 0,
+    use_eagle: bool = True,
 ) -> KVCacheManager:
     config = KVCacheConfig(
         num_blocks=10000,
@@ -120,7 +121,7 @@ def _make_hybrid_kv_cache_manager(
         scheduler_block_size=MAMBA_BLOCK_SIZE,
         hash_block_size=ATTN_BLOCK_SIZE,
         enable_caching=True,
-        use_eagle=True,
+        use_eagle=use_eagle,
     )
 
 
@@ -132,6 +133,7 @@ def _split(
     partial_hit: bool = False,
     num_prefill_checkpoint_blocks: int = 0,
     max_num_scheduled_tokens: int = 16384,
+    num_prefill_lookahead: int = 0,
 ) -> int:
     """Call the real `Scheduler._mamba_block_aligned_split` on a stub self."""
     if use_eagle_block_drop is None:
@@ -140,6 +142,7 @@ def _split(
         block_size=MAMBA_BLOCK_SIZE,
         cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE),
         use_eagle_block_drop=use_eagle_block_drop,
+        num_prefill_lookahead=num_prefill_lookahead,
         max_num_scheduled_tokens=max_num_scheduled_tokens,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         # `prefix_match_unit` finer than the block size (#46384).
@@ -392,3 +395,113 @@ 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})"
         )
+
+
+def _prefill_with_prefill_lookahead(
+    prompt_len: int,
+    lookahead: int,
+    budget: int,
+    use_eagle_block_drop: bool,
+) -> tuple[list[int], int]:
+    """Chunk a prefill the way `Scheduler.schedule()` does for multi-module MTP.
+
+    Every step runs the real `_mamba_block_aligned_split` and then the real
+    `_reserve_prefill_lookahead` on the step's budget, allocating through a real
+    hybrid `KVCacheManager`. Returns the chunk ends and the number of cached
+    Mamba slots, each checked to hold the state its hash claims.
+    """
+    manager = _make_hybrid_kv_cache_manager(use_eagle=use_eagle_block_drop)
+    mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID]
+    (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
+    sched = SimpleNamespace(
+        num_prefill_lookahead=lookahead,
+        need_mamba_block_aligned_split=True,
+        cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE),
+    )
+    ends: list[int] = []
+    state_at: dict[int, int] = {}
+    # Bounded: with a constant budget a stalled prefill never recovers.
+    for _ in range(2 * cdiv(prompt_len, min(budget, MAMBA_BLOCK_SIZE)) + 2):
+        computed = request.num_computed_tokens
+        if computed >= prompt_len:
+            break
+        num_new = _split(
+            request,
+            min(prompt_len - computed, budget),
+            use_eagle_block_drop=use_eagle_block_drop,
+            max_num_scheduled_tokens=budget,
+            num_prefill_lookahead=lookahead,
+        )
+        num_new = Scheduler._reserve_prefill_lookahead(
+            sched, request, computed, num_new
+        )
+        if num_new == 0:
+            continue
+        assert (
+            manager.allocate_slots(request, num_new, num_lookahead_tokens=NUM_SPEC)
+            is not None
+        )
+        request.num_computed_tokens = computed + num_new
+        ends.append(request.num_computed_tokens)
+        blocks = mamba_manager.req_to_blocks[request.request_id]
+        running = cdiv(request.num_computed_tokens, MAMBA_BLOCK_SIZE) - 1
+        state_at[blocks[running].block_id] = request.num_computed_tokens
+    return ends, _count_cached_boundary_states(manager, request, state_at)
+
+
+@pytest.mark.parametrize("use_eagle_block_drop", [False, True])
+@pytest.mark.parametrize("budget_blocks", [4.0, 1.0, 0.5])
+@pytest.mark.parametrize("num_full_blocks", [1, 2])
+@pytest.mark.parametrize("tail", ["one", "window_edge", "past_window"])
+@pytest.mark.parametrize("lookahead", [2, 3, 5])
+@pytest.mark.parametrize("block_size", [64, 1600])
+def test_mamba_split_respects_multi_module_mtp_prefill_lookahead(
+    monkeypatch: pytest.MonkeyPatch,
+    block_size: int,
+    lookahead: int,
+    tail: str,
+    num_full_blocks: int,
+    budget_blocks: float,
+    use_eagle_block_drop: bool,
+) -> None:
+    """#61762: the Mamba split and the multi-module MTP lookahead must agree.
+
+    `_reserve_prefill_lookahead` shortened a chunk ending within
+    `num_prefill_lookahead` of the prompt end. Applied to a Mamba-aligned end,
+    that left the request just short of the boundary, and the next step (start
+    mid-block, stop at the boundary) was shortened to zero tokens, forever.
+    """
+    monkeypatch.setattr(sys.modules[__name__], "MAMBA_BLOCK_SIZE", block_size)
+    # Tokens past the last full block: inside the window, at its edge, past it.
+    tail_len = {"one": 1, "window_edge": lookahead - 1, "past_window": lookahead}[tail]
+    prompt_len = num_full_blocks * block_size + tail_len
+    budget = int(block_size * budget_blocks)
+
+    ends, _ = _prefill_with_prefill_lookahead(
+        prompt_len, lookahead, budget, use_eagle_block_drop
+    )
+
+    assert ends and ends[-1] == prompt_len, f"prefill stalled at {ends}"
+    if budget >= block_size:
+        assert all(end % block_size == 0 for end in ends[:-1]), ends
+    # A chunk ends inside the lookahead window only when no block boundary
+    # before the window can be reached and the prompt cannot be finished.
+    in_window = [end for end in ends[:-1] if prompt_len - lookahead < end]
+    assert all(end % block_size == 0 for end in in_window), ends
+    if budget >= prompt_len:
+        assert not in_window, ends
+
+
+def test_mamba_split_lookahead_window_is_noop_without_multi_module_mtp() -> None:
+    """EAGLE-family drafters (lookahead 1) and no drafter keep today's split."""
+    (request,) = create_requests(1, num_tokens=3201, block_size=ATTN_BLOCK_SIZE)
+    for lookahead in (0, 1):
+        assert (
+            _split(
+                request,
+                request.num_tokens,
+                use_eagle_block_drop=False,
+                num_prefill_lookahead=lookahead,
+            )
+            == 2 * MAMBA_BLOCK_SIZE
+        )

tests/v1/core/test_prefix_caching.py

@@ -2632,6 +2632,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
         max_num_scheduled_tokens=3 * block_size,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
+        num_prefill_lookahead=0,
         hash_block_size=block_size,
         mamba_partial_cache_hit=False,
         mamba_shared_prefix_checkpoint=False,

tests/v1/core/test_scheduler.py

@@ -4660,6 +4660,40 @@ def test_mamba_align_encoder_cache_cap_makes_progress():
     assert output.scheduled_encoder_inputs[request.request_id] == [1]
 
 
+@pytest.mark.parametrize(
+    ("prompt_len", "max_num_batched_tokens"),
+    [(49, 512), (50, 512), (33, 32), (65, 512)],
+)
+def test_mamba_align_multi_module_mtp_prefill_does_not_livelock(
+    prompt_len: int, max_num_batched_tokens: int
+):
+    """#61762: with multi-module MTP (prefill lookahead 3) and Mamba "align"
+    without EAGLE block drop, a prompt ending 1..lookahead-1 tokens past a
+    block boundary used to be scheduled 0 tokens on every step after its first
+    chunk, so it never finished.
+    """
+    scheduler = create_scheduler(
+        block_size=16,
+        enable_prefix_caching=True,
+        max_num_batched_tokens=max_num_batched_tokens,
+    )
+    scheduler.need_mamba_block_aligned_split = True
+    scheduler.num_prefill_lookahead = 3
+    (request,) = create_requests(num_requests=1, num_tokens=prompt_len, max_tokens=1)
+    scheduler.add_request(request)
+
+    for _ in range(8):
+        output = scheduler.schedule()
+        if request.request_id not in output.num_scheduled_tokens:
+            break
+        _model_output(
+            scheduler, output, [[] if request.is_prefill_chunk else [EOS_TOKEN_ID]]
+        )
+    assert request.is_finished(), (
+        f"stalled at {request.num_computed_tokens}/{prompt_len}"
+    )
+
+
 def test_mamba_align_eagle_schedules_encoder_at_boundary():
     """EAGLE lookahead at an aligned MM boundary requires encoder cache."""
     block_size = 512

vllm/v1/core/sched/scheduler.py

@@ -553,6 +553,26 @@ def _mamba_block_aligned_split(
         )
         # Stop at the earliest mandatory position strictly inside the chunk.
         end = min((s for s in stops if start < s < end), default=end)
+
+        # Multi-module MTP wants no chunk to end within `num_prefill_lookahead`
+        # of the end of the prompt (see `_reserve_prefill_lookahead`), which
+        # used to shorten the Mamba-aligned end chosen above: the request was
+        # left just short of the boundary, and the next step (start mid-block,
+        # stop at the boundary) was shortened to zero tokens, forever. Decide
+        # it here instead, giving up the checkpoint inside the window: finish
+        # the prefill if that crosses no boundary from mid-block, else stop at
+        # the last block boundary before the window. If neither exists, keep
+        # the aligned end: the next chunk finishes the prefill, and the drafter
+        # reads sampled tokens once rather than the slot state being poisoned.
+        window_start = request.num_tokens - self.num_prefill_lookahead
+        if window_start < end < request.num_tokens:
+            aligned_end = window_start // block_size * block_size
+            if start + num_new_tokens >= request.num_tokens and (
+                start % block_size == 0 or next_block_boundary >= request.num_tokens
+            ):
+                end = request.num_tokens
+            elif aligned_end > start:
+                end = aligned_end
         return max(end - start, 0)
 
     def _get_local_prefix_cache_hit(
@@ -581,8 +601,16 @@ def _reserve_prefill_lookahead(
         boundary closer to the end than that would make it fall back to
         sampled drafts, permanently polluting the trailing modules' KV caches.
         Either finish the prefill or leave at least num_prefill_lookahead for
-        the next chunk. No-op for eagle-family drafters (lookahead 1).
+        the next chunk. No-op for eagle-family drafters (lookahead 1). A
+        Mamba-aligned chunk end is kept: `_mamba_block_aligned_split` already
+        moved it out of the window wherever a block boundary allowed it.
         """
+        if (
+            self.need_mamba_block_aligned_split
+            and (num_computed_tokens + num_new_tokens) % self.cache_config.block_size
+            == 0
+        ):
+            return num_new_tokens
         remaining = request.num_tokens - num_computed_tokens - num_new_tokens
         if 0 < remaining < self.num_prefill_lookahead:
             num_new_tokens -= self.num_prefill_lookahead - remaining