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(