[apex-mochen]

Purpose

Fixes #61762.

With Mamba align prefix caching (the default for hybrid models once prefix caching is on) and multi-module MTP (num_prefill_lookahead = num_speculative_tokens = L >= 2), a prompt that ends 1..L-1 tokens past a Mamba block boundary can stop being scheduled. Two clamps in Scheduler.schedule() disagree about the same position:

1. _mamba_block_aligned_split ends a chunk on a block boundary (a cacheable checkpoint, or a budget-limited chunk floored to a boundary).
2. _reserve_prefill_lookahead sees fewer than L tokens left and moves the end back by L - remaining, off the boundary.
3. On the next step the chunk starts mid-block, the split stops it at that same boundary, the reservation shortens it to 0, and the request is skipped. This repeats on every step: the request holds its KV blocks and never finishes, and no error is raised.

The report in #61762 (and my follow-up there) only showed this with disable_eagle_block_drop=True. While writing the fix I found it also happens with EAGLE block drop on (the default) whenever the chunk that reaches the last boundary is budget-limited. Block size 16, L=3, 49-token prompt, 16-token budget: the chunks end at 16, 32, then 46, and from then on 0 tokens per step. With a shared per-step budget, a long prompt's last full chunk can be budget-limited in this way.

Fix. The lookahead is now resolved inside _mamba_block_aligned_split, which knows where chunks may end. If a chunk would end inside the window, the split does the first of these that applies:

• Finish the prefill, if the budget allows it and it crosses no block boundary from a mid-block start. This gives up the checkpoint inside the window.
• Stop at the last block boundary before the window.
• Keep the aligned end. _reserve_prefill_lookahead now leaves Mamba-aligned ends alone, so the next chunk finishes the prefill. In that case the drafter reads sampled tokens at this one boundary, which is what the reservation tries to avoid. That is a draft-quality cost; the alternatives are a livelock, or a chunk ending mid-block followed by one crossing the boundary, which publishes a Mamba slot at the wrong offset ([Bug]: Accuracy drops ~20% when --enable-prefix-caching is used together with MTP speculative decoding (Qwen3.6 35B-A3B) #43559).

Unaligned (sub-block) chunk ends are still shortened by _reserve_prefill_lookahead as before. Nothing changes when num_prefill_lookahead <= 1 (EAGLE-family drafters, or no drafter). The window condition N - L < end < N cannot hold in that case.

Reachability

As far as I can tell, no public checkpoint reaches this today:

• Qwen3.5 and Qwen3-Next have one MTP layer.
• Ling/Ring-2.5 and Kimi-Linear have none.
• nemotron_h_mtp.py asserts a single MTP layer.

qwen3_next_mtp.py and bailing_moe_v3_mtp.py already iterate over num_nextn_predict_layers, and nothing rejects multi-module MTP with align. A hybrid checkpoint with two or more MTP layers would hit it. All evidence below is scheduler-level on CPU. I have not run a model, and I am not claiming end-to-end validation.

Test Plan

• tests/v1/core/test_mamba_align_chunk_split.py::test_mamba_split_respects_multi_module_mtp_prefill_lookahead (216 cases).
  • Each step runs the real _mamba_block_aligned_split and then _reserve_prefill_lookahead, allocating through a real hybrid KVCacheManager.
  • Parameters: Mamba block 64/1600, L 2/3/5, prompt tail inside/at the edge of/past the window, 1–2 full blocks, budget 4 blocks / 1 block / half a block (the sub-block mode), EAGLE block drop on/off.
  • It asserts that the prefill completes, that non-final chunk ends are block-aligned when the budget allows it, and that a chunk ends inside the window only on a boundary and only when it cannot finish. It also runs the file's existing _count_cached_boundary_states check, which verifies that every hash-cached Mamba slot holds the state its hash claims.
• test_mamba_split_lookahead_window_is_noop_without_multi_module_mtp: L=0/1 keep today's split.
• tests/v1/core/test_scheduler.py::test_mamba_align_multi_module_mtp_prefill_does_not_livelock: the real Scheduler, with the cases from [Bug]: Scheduler livelock with multi-module MTP + Mamba align prefix caching when prompt_len % block_size < num_speculative_tokens #61762.
• Stubs passed as self to _mamba_block_aligned_split in four test files gain num_prefill_lookahead=0, so those tests are unchanged.

Test Result

(CPU only, WSL Ubuntu 24.04, Python 3.12; my local environment lacks torchvision, which affects a few unrelated tests on both trees.)

• New tests on main's code (this PR's tests applied to 71536bd5f8): 124 failed, 143 passed. Of the failures, 72 are block-drop-off cases, 48 are block-drop-on (default) cases, all with a budget of one block or less, and 4 are the Scheduler cases. With this PR: 267 passed.
• tests/v1/core in full, main vs this PR: the failed/errored sets are identical (68 environment failures on both trees: missing torchvision and offline model configs), with 0 failures only on this PR; passes go from 940 to 1161 (the new tests).
• Repro from [Bug]: Scheduler livelock with multi-module MTP + Mamba align prefix caching when prompt_len % block_size < num_speculative_tokens #61762 (block 16, budget 512): main stalls at 46/49; this PR finishes, scheduling 0→49 in one chunk.
• Default EAGLE block drop, budget 16 (49/50-token prompts, L=3; 33-token prompt, L=2): main stalls at 46/47/31; this PR finishes. Budgets ≥ 32 show identical chunking before and after.
• I also ran a random-workload stall check: the real Scheduler on a hybrid attention+Mamba config, 40 seeds × 300 steps per config. It flags a running request that gets 0 tokens while the step has budget left for its whole remainder.
  • L=3 with block drop off: main has 16/40 seeds stalled, this PR has 0.
  • L=2 with block drop off: main has 14/40, this PR has 1. The one remaining seed is a harness artifact: a preempted 304-token request plus spec lookahead slots needs more blocks than the 24-block pool has, and its split/reservation results are the same as on main.
  • L=3 with block drop on, and L=1: 0 stalls on both trees, with identical per-run action counts, so scheduling is unchanged there.
• pre-commit run --files <changed files>: all hooks pass (ruff check/format, typos, mypy, SPDX, new-test CI wiring).

AI assistance

AI-assisted (Claude Code). I reviewed the change and ran the tests listed above myself.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[apex-mochen]

Hi @njhill — would you be able to take a look at this when you have a moment?

With Mamba "align" prefix caching and multi-module MTP (num_prefill_lookahead >= 2), _mamba_block_aligned_split ends a chunk on a block boundary and _reserve_prefill_lookahead then moves it back off that boundary; on the next step the chunk starts mid-block, stops at the same boundary, and is cut to 0 tokens, every step (#61762). With the default EAGLE block drop this happens whenever the chunk reaching the last boundary is budget-limited.

The fix resolves the lookahead inside the split (finish the prefill, else stop at the last boundary before the window, else keep the aligned end) and lets _reserve_prefill_lookahead leave Mamba-aligned ends alone. Verification on db58b6c4c1: the new tests fail on main (124 failed / 143 passed, including 48 with block drop on) and pass with the fix (267); tests/v1/core has the same failure set as main. All CPU, scheduler-level. No public checkpoint has a multi-layer MTP head on a hybrid model yet, so I could not run a real model.

One review focus for you: in the "keep the aligned end" case the drafter reads sampled tokens at that one boundary; I chose that over ending a chunk mid-block (which would publish a Mamba slot at the wrong offset, #43559). Happy to adjust if you'd prefer a different trade-off.

[apex-mochen]

Hi @ZJY0516 — this one hasn't had a review yet (first ask on Oct 16), and since #61181 put Mamba and KV cache under your CODEOWNERS and you've been merging Mamba changes lately (#61497), I wanted to ask you directly.

Quick context: with Mamba align prefix caching and multi-module MTP (num_prefill_lookahead >= 2), _mamba_block_aligned_split can end a chunk on a block boundary inside the MTP lookahead window, and _reserve_prefill_lookahead then shrinks it below that boundary, so the request livelocks (#61762). The fix keeps the block-aligned chunk end and skips the lookahead reserve only when the chunk ends block-aligned; non-Mamba and non-MTP paths are unchanged.

Head is db58b6c4c1 since Oct 16; CI is green and it still merges cleanly. Repro and before/after test output are in the PR description.