[BugFix] Honor drop_eagle_block in MambaManager

On a hybrid model (full-attention + Mamba/GDN layers, e.g. Qwen3.6-35B-A3B),
enabling MTP/EAGLE speculative decoding together with prefix caching corrupts
the recurrent state, and the server serves wrong answers. #43559 reports
tool-eval accuracy dropping from ~90% to ~50%. Nothing raises an exception
and nothing shows up in the logs, so an operator has no signal that anything
is wrong, and prefix caching then spreads it, because the corrupted state is
hash-reachable and every later request sharing that prefix restores from it.

The root cause is the find_longest_cache_hit method on MambaManager,
at vllm/v1/core/single_type_kv_cache_manager.py:1049 on main, which takes
drop_eagle_block in its signature and never acts on it. Its body searches
right-to-left for the last cached block, builds [null]*i + [cached], and
breaks, referencing neither drop_eagle_block nor .pop(.

Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
2026-09-27 10:45:10 +07:00
parent 5e6656e7a0
commit 5da84df6b7
3 changed files with 175 additions and 0 deletions
@@ -914,3 +914,73 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
f"but allocator pulled {len(new_blocks)}"
)
total_computed = num_tokens
def test_mamba_honors_drop_eagle_block():
"""MambaManager must drop the EAGLE lookahead block, like FullAttentionManager.
EAGLE/MTP requires the final matched block of a cache hit be dropped and
recomputed: it can hold state written over draft positions that verification
later rejects. FullAttentionManager pops it. Before this fix MambaManager
accepted `drop_eagle_block` and ignored it, so its hit came back one block
longer, and the retained block's recurrent-state snapshot -- taken over the
unverified positions -- stayed hash-reachable for every later request sharing
the prefix (#43559).
block_size == alignment_tokens here, so the alignment loop is a no-op in both
managers and the eagle drop is the only behavior under test.
"""
block_size = 16
num_blocks = 5
full_gid, mamba_gid = 0, 1
block_pool = BlockPool(
num_gpu_blocks=64, enable_caching=True, hash_block_size=block_size
)
block_hashes = [BlockHash(f"blk{i:04d}".encode()) for i in range(num_blocks)]
# Cache the same prefix chain under both groups, each on real blocks.
for gid in (full_gid, mamba_gid):
for i, (block_hash, block) in enumerate(
zip(block_hashes, block_pool.get_new_blocks(num_blocks))
):
block_pool._insert_block_hash(
make_block_hash_with_group_id(block_hash, gid),
block,
num_tokens=(i + 1) * block_size,
)
full_spec = FullAttentionSpec(
block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float16
)
mamba_spec = MambaSpec(
block_size=block_size,
shapes=((1,),),
dtypes=(torch.float16,),
mamba_cache_mode="align",
)
def hit_blocks(manager, gid, spec, drop_eagle_block):
return len(
manager.find_longest_cache_hit(
block_hashes=block_hashes,
max_length=num_blocks * block_size,
kv_cache_group_ids=[gid],
block_pool=block_pool,
kv_cache_spec=spec,
drop_eagle_block=drop_eagle_block,
alignment_tokens=block_size,
)[0]
)
# Control: with no eagle drop, both groups see the whole cached prefix.
assert hit_blocks(FullAttentionManager, full_gid, full_spec, False) == num_blocks
assert hit_blocks(MambaManager, mamba_gid, mamba_spec, False) == num_blocks
# With the eagle drop requested, both must give up the final block.
full_hit = hit_blocks(FullAttentionManager, full_gid, full_spec, True)
mamba_hit = hit_blocks(MambaManager, mamba_gid, mamba_spec, True)
assert full_hit == num_blocks - 1
assert mamba_hit == full_hit, (
f"MambaManager ignored drop_eagle_block: full-attention dropped the "
f"lookahead block (hit={full_hit}), mamba kept it (hit={mamba_hit})"
)
@@ -937,6 +937,102 @@ def _run_mamba_prefix_cache_mrv1(
cleanup_dist_env_and_memory()
def _run_mamba_eagle_block_reuse(
monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
):
"""Recreate the #43559 failure conditions: a request finishing right after a
block boundary it crossed with every draft rejected, then a second request
whose cache hit lands on that final (eagle) block.
Request A ends a few tokens past the 2nd aligned boundary; with
num_accepted_tokens=1 every boundary crossing happens with rejected drafts
in flight, and A finishes before any later step could repair its cached
snapshot. Request B shares the prefix: pre-fix MambaManager returns the
boundary block as part of the hit (B resumes on the possibly-poisoned
snapshot), post-fix the hit stops one block earlier and B recomputes it.
Assertions, all deterministic:
1. B's num_cached_tokens == BLOCK_SIZE (the eagle block was dropped).
2. A's cached snapshot at the boundary matches the no-spec reference.
3. B's state at its next boundary matches the reference -- corruption
resumed from the reused block would surface here (the user-visible
failure, at state level).
"""
global async_scheduling_mode
async_scheduling_mode = async_scheduling
run_ref_mamba_state_in_subprocess()
apply_patch(monkeypatch)
full_prompt = datasets.load_dataset("heheda/a_long_article")["train"][0]["text"]
engine = LLM(
model=MODEL,
load_format="dummy",
enable_prefix_caching=True,
block_size=BLOCK_SIZE,
mamba_cache_mode="align",
speculative_config={
"method": "qwen3_next_mtp",
"num_speculative_tokens": num_speculative_tokens,
},
max_num_batched_tokens=3072,
hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
async_scheduling=async_scheduling,
seed=42,
)
global prompt_token_ids
prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
global num_accepted_tokens
num_accepted_tokens = 1
global cur_step_action_idx, step_actions
cur_step_action_idx = 0
step_actions = []
boundary = BLOCK_SIZE * 2
# A decodes across BOTH boundaries (align mode only snapshots state at
# decode/chunk crossings, and B needs a fallback state block once the
# eagle block is dropped), rejecting 3 drafts every step, and finishes
# a few tokens past the second boundary -- no later step can repair the
# cached snapshot.
engine.generate(
[TokensPrompt(prompt_token_ids=prompt_token_ids[: BLOCK_SIZE - 2])],
SamplingParams(temperature=0.0, max_tokens=BLOCK_SIZE + 10),
)
a_boundary_state = mamba_kv_cache_dict.get(boundary)
assert a_boundary_state is not None, "A never cached the boundary block"
mamba_kv_cache_dict.clear()
# B: prefix through the boundary, prompt ending shy of the next one so a
# few decode steps cross it and its state gets captured for comparison.
next_boundary = BLOCK_SIZE * 3
(out_b,) = engine.generate(
[TokensPrompt(prompt_token_ids=prompt_token_ids[: next_boundary - 10])],
SamplingParams(temperature=0.0, max_tokens=15),
)
assert out_b.num_cached_tokens == BLOCK_SIZE, (
f"cache hit must drop the eagle block: expected {BLOCK_SIZE} cached "
f"tokens, got {out_b.num_cached_tokens}"
)
mamba_state_ref = torch.load("mamba_kv_cache_dict_ref.pth")
check_mamba_state_equal(mamba_state_ref, {boundary: a_boundary_state}, [boundary])
check_mamba_state_equal(mamba_state_ref, mamba_kv_cache_dict, [next_boundary])
mamba_kv_cache_dict.clear()
del engine
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
@create_new_process_for_each_test("spawn")
def test_mamba_eagle_block_reuse_mrv1(monkeypatch: pytest.MonkeyPatch):
_run_mamba_eagle_block_reuse(monkeypatch, async_scheduling=False)
@create_new_process_for_each_test("spawn")
def test_mamba_eagle_block_reuse_mrv1_async(monkeypatch: pytest.MonkeyPatch):
_run_mamba_eagle_block_reuse(monkeypatch, async_scheduling=True)
@create_new_process_for_each_test("spawn")
def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
_run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=False)
@@ -1538,6 +1538,15 @@ class MambaManager(SingleTypeKVCacheManager):
return computed_blocks, hit_length
max_num_blocks = max_length // block_size
if drop_eagle_block and max_num_blocks > 0:
# EAGLE/MTP: drop the final matched page. Its recurrent-state snapshot
# may have been taken over draft tokens that verification later rejects,
# so the state must be restored from the previous page boundary and
# recomputed forward. FullAttentionManager (see find_longest_cache_hit
# above) pops the last block; Mamba keeps only the rightmost real block
# with the prefix null-padded, so a literal pop would delete the state
# block itself -- lower the search ceiling by one page instead.
max_num_blocks -= 1
# Search from right to left and early stop when a match is found.
for i in range(max_num_blocks - 1, -1, -1):
if cached_block := block_pool.get_cached_block(