From 5da84df6b7c824caefcac2976843ae25283d7f8b Mon Sep 17 00:00:00 2001 From: MOVZX Date: Sun, 27 Sep 2026 10:43:44 +0700 Subject: [PATCH] [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 --- .../core/test_single_type_kv_cache_manager.py | 70 ++++++++++++++ .../v1/e2e/general/test_mamba_prefix_cache.py | 96 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 9 ++ 3 files changed, 175 insertions(+) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index f20a96f5cc..2705b87f80 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -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})" + ) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 5b6826a4b3..6433edac37 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -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) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 0dd9734847..4c56ba8d81 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -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(