diff --git a/tests/models/kimi_k3/test_amd_kda_checkpoint.py b/tests/models/kimi_k3/test_amd_kda_checkpoint.py index 0a622bfa04..6f3129a224 100644 --- a/tests/models/kimi_k3/test_amd_kda_checkpoint.py +++ b/tests/models/kimi_k3/test_amd_kda_checkpoint.py @@ -75,7 +75,7 @@ def _builder( vllm_config.cache_config.prefix_match_unit = prefix_match_unit if num_spec: vllm_config.speculative_config = _StubSpeculativeConfig(num_spec) - return KimiK3ROCmKDAMetadataBuilder( + builder = KimiK3ROCmKDAMetadataBuilder( kv_cache_spec=MambaSpec( block_size=block_size, shapes=((16, 64),), @@ -91,6 +91,11 @@ def _builder( vllm_config=vllm_config, device=torch.device("cuda"), ) + vllm_config.cache_config.hash_block_size = prefix_match_unit or block_size + vllm_config.cache_config.cache_hit_alignment_tokens = ( + prefix_match_unit or block_size + ) + return builder def _build(batch: BatchSpec, block_size: int = BLOCK_SIZE, **kw): diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 7441427955..63024634d3 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -100,6 +100,8 @@ def _make_builder( prefix_match_unit: int | None = None, use_eagle: bool = False, disable_eagle_block_drop: bool = False, + cache_hit_alignment_tokens: int | None = None, + hash_block_size: int | None = None, ) -> AttentionMetadataBuilder: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -122,6 +124,11 @@ def _make_builder( vllm_config.cache_config.use_replayssm = use_recoverssm vllm_config.cache_config.use_kda_recoverssm = use_recoverssm vllm_config.cache_config.prefix_match_unit = prefix_match_unit + hash_block_size = hash_block_size or prefix_match_unit or mamba_block_size + vllm_config.cache_config.hash_block_size = hash_block_size + vllm_config.cache_config.cache_hit_alignment_tokens = ( + cache_hit_alignment_tokens or hash_block_size + ) builder = builder_cls( kv_cache_spec=MambaSpec( block_size=mamba_block_size, @@ -220,6 +227,8 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): vllm_config=builder_config, device=DEVICE, ) + builder_config.cache_config.hash_block_size = BLOCK_SIZE + builder_config.cache_config.cache_hit_alignment_tokens = BLOCK_SIZE # An all-prefill speculative batch used to leave an all-false spec mask # alive, then access active_non_spec_mask_cpu before it was initialized. @@ -287,18 +296,45 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_internal_checkpoint_metadata_targets_last_aligned_boundary(): +@pytest.mark.parametrize( + "batch,prefix_match_unit,hit_alignment,expected_offsets,expected_states", + [ + ( + BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]), + None, + 16, + [48, 0], + [2, NULL_BLOCK_ID], + ), + ( + BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]), + 16, + 16, + [48, 0], + [2, NULL_BLOCK_ID], + ), + (BatchSpec(seq_lens=[530], query_lens=[530]), None, 128, [512], [32]), + (BatchSpec(seq_lens=[530], query_lens=[50]), None, 128, [32], [32]), + (BatchSpec(seq_lens=[530], query_lens=[530]), 16, 16, [528], [32]), + (BatchSpec(seq_lens=[530], query_lens=[50]), 16, 16, [48], [32]), + ], +) +def test_internal_checkpoint_metadata_targets_last_aligned_boundary( + batch, prefix_match_unit, hit_alignment, expected_offsets, expected_states +): device = torch.device("cuda") - batch = BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]) common_attn_metadata = create_common_attn_metadata( batch, BLOCK_SIZE, device, arange_block_indices=True - ).replace(is_prefilling=torch.tensor([True, True])) + ).replace(is_prefilling=torch.tensor([True] * len(batch.seq_lens))) builder = _make_builder( KimiK3KDAMetadataBuilder, num_speculative_tokens=0, full_cuda_graph=False, mamba_cache_mode="align", num_prefill_checkpoint_blocks=1, + prefix_match_unit=prefix_match_unit, + cache_hit_alignment_tokens=hit_alignment, + hash_block_size=16, device=device, ) assert isinstance(builder, KimiK3KDAMetadataBuilder) @@ -313,26 +349,33 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary(): assert actual.checkpoint is not None torch.testing.assert_close( actual.checkpoint.state_indices, - torch.tensor([2, NULL_BLOCK_ID], dtype=torch.int32, device=device), + torch.tensor(expected_states, dtype=torch.int32, device=device), ) torch.testing.assert_close( actual.checkpoint.checkpoint_offsets, - torch.tensor([48, 0], dtype=torch.int32, device=device), + torch.tensor(expected_offsets, dtype=torch.int32, device=device), ) @pytest.mark.parametrize( - ("disable_eagle_block_drop", "prefix_match_unit", "expected_offset"), - [(False, 16, 80), (True, 16, 96), (False, 8, None)], + ("seq_len", "disable_eagle_block_drop", "prefix_match_unit", "expected_offset"), + [ + (100, False, 16, 80), + (100, True, 16, 96), + (100, False, 8, None), + (96, False, 16, 80), + (96, True, 16, 80), + ], ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_spec_internal_checkpoint_metadata_targets_replay_boundary( + seq_len: int, disable_eagle_block_drop: bool, prefix_match_unit: int, expected_offset: int | None, ) -> None: device = torch.device("cuda") - batch = BatchSpec(seq_lens=[100], query_lens=[100]) + batch = BatchSpec(seq_lens=[seq_len], query_lens=[seq_len]) common_attn_metadata = create_common_attn_metadata( batch, BLOCK_SIZE, device, arange_block_indices=True ) @@ -390,6 +433,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset(): full_cuda_graph=False, mamba_cache_mode="align", num_prefill_checkpoint_blocks=1, + prefix_match_unit=16, device=device, ) assert isinstance(builder, KimiK3KDAMetadataBuilder) diff --git a/tests/models/kimi_k3/test_prefix_cache.py b/tests/models/kimi_k3/test_prefix_cache.py index 3525550fe9..20e5d6670e 100644 --- a/tests/models/kimi_k3/test_prefix_cache.py +++ b/tests/models/kimi_k3/test_prefix_cache.py @@ -8,6 +8,7 @@ import random from collections.abc import Iterator from dataclasses import dataclass from math import lcm +from uuid import uuid4 import pytest import requests @@ -353,6 +354,45 @@ def _format(turns: list[Turn]) -> str: return "\n".join(rows) +@pytest.mark.parametrize( + "prompt_len", + [ + # Blocks=6144, PMU=128. P=7040 is PMU-aligned but not block-aligned: + # ensure lookup reads the cached attention proof at 7040 despite + # the final hit limit of 7039, then applies EAGLE's 128-token drop. + pytest.param(7040, id="eagle-proof-scan"), + # Blocks=6144, PMU=128. P=7296 is an exact PMU multiple: + # ensure the Mamba checkpoint is 7168, not 7040 from applying both + # the P - 1 cap and EAGLE's 128-token drop. + pytest.param(7296, id="eagle-double-cap"), + ], +) +def test_dspark_exact_resend_reuses_prompt_checkpoint(prompt_len: int) -> None: + """Resends reuse the Mamba checkpoint and its companion EAGLE attention proof.""" + mode = Mode(128, True) + salt = str(uuid4()) + rng = random.Random(0) + prompt = [rng.randint(1000, 150000) for _ in range(prompt_len)] + deployment = DEPLOYMENTS["plain"] + with _serve(deployment, mode) as servers: + url = servers[0].url_root + cold = _complete(url, prompt, salt, max_tokens=1) + resend = _complete(url, prompt, salt, max_tokens=1) + recompute = _complete(url, prompt, f"{salt}-recompute", max_tokens=1) + + # EAGLE's dropped hash unit already leaves tokens to recompute on resend. + expected = (prompt_len // 128 - 1) * 128 + actual = _cached(resend) + assert _cached(cold) == _cached(recompute) == 0 + assert actual == expected, f"{prompt_len=}: {actual=}, {expected=}" + check_logprobs_close( + outputs_0_lst=[_tokens_text_logprobs(recompute)], + outputs_1_lst=[_tokens_text_logprobs(resend)], + name_0="recompute", + name_1="resend", + ) + + @pytest.mark.parametrize( "name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES] ) diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index cd01fba49e..4b52eb9c09 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -385,7 +385,7 @@ def test_cudagraph_capture_batch_stays_decode_only(): def _create_checkpoint_builder_and_batch( - num_spec: int, prefix_match_unit: int = 16 + num_spec: int, prefix_match_unit: int | None = 16 ) -> tuple[GDNAttentionMetadataBuilder, CommonAttentionMetadata, dict]: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -413,6 +413,8 @@ def _create_checkpoint_builder_and_batch( vllm_config=vllm_config, device=DEVICE, ) + vllm_config.cache_config.hash_block_size = prefix_match_unit or 64 + vllm_config.cache_config.cache_hit_alignment_tokens = prefix_match_unit or 64 batch = BatchSpec( seq_lens=[65, 100, 100], query_lens=[3 if num_spec else 1, 100, 99] ) @@ -431,7 +433,9 @@ def _create_checkpoint_builder_and_batch( @pytest.mark.parametrize("num_spec", [0, 3]) -@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)]) +@pytest.mark.parametrize( + "prefix_match_unit, expected_offset", [(None, 64), (16, 96), (8, 96)] +) def test_checkpoint_metadata_preserves_non_spec_order( num_spec, prefix_match_unit, expected_offset ): diff --git a/tests/v1/attention/test_mamba2_prefill_checkpoints.py b/tests/v1/attention/test_mamba2_prefill_checkpoints.py index e5276786e2..a3968b5a0b 100644 --- a/tests/v1/attention/test_mamba2_prefill_checkpoints.py +++ b/tests/v1/attention/test_mamba2_prefill_checkpoints.py @@ -77,7 +77,9 @@ def _create_mamba2_builder( num_prefill_checkpoint_blocks: int = 1, prefill_checkpoint_alignment: int | None = 1, ) -> Mamba2AttentionMetadataBuilder: - vllm_config = create_vllm_config(block_size=MAMBA_BLOCK_SIZE) + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", block_size=MAMBA_BLOCK_SIZE + ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode # get_mamba_chunk_size() reads `mamba_chunk_size` then `chunk_size`. vllm_config.model_config.hf_text_config.mamba_chunk_size = CHUNK_SIZE @@ -89,12 +91,15 @@ def _create_mamba2_builder( num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, prefill_checkpoint_alignment=prefill_checkpoint_alignment, ) - return Mamba2AttentionMetadataBuilder( + builder = Mamba2AttentionMetadataBuilder( kv_cache_spec=spec, layer_names=["layer.0"], vllm_config=vllm_config, device=DEVICE, ) + vllm_config.cache_config.hash_block_size = MAMBA_BLOCK_SIZE + vllm_config.cache_config.cache_hit_alignment_tokens = MAMBA_BLOCK_SIZE + return builder def _build(builder, seq_lens, query_lens): diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 76beffce45..51317deb8d 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -22,6 +22,7 @@ from vllm.v1.core.kv_cache_utils import ( get_block_hash, get_group_id, init_none_hash, + resolve_cache_hit_alignment_tokens, ) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.kv_cache_interface import ( @@ -148,29 +149,84 @@ def make_full_mamba_manager( ) -def test_dcp_fine_hit_retention_uses_hash_alignment_without_eagle(): - """DCP must not discard a reusable Mamba state at hash alignment.""" - hash_block_size = 2 +@pytest.mark.parametrize( + "dcp_world_size,hash_block_size,block_size,prompt_tokens,checkpoint_blocks", + [ + (4, 2, 2, 7, 0), + (8, 16, 16, 512, 1), + (8, 16, 16, 513, 1), + *[(8, 128, 896, p, 1) for p in (127, 128, 129, 7040, 7168, 7296, 7449)], + ], +) +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("extension_tokens", [0, 32]) +def test_dcp_fine_hit_retention_uses_hash_alignment( + dcp_world_size, + hash_block_size, + block_size, + prompt_tokens, + checkpoint_blocks, + use_eagle, + extension_tokens, +): + """Resends and extensions reuse the retained hash-aligned checkpoint.""" manager = make_full_mamba_manager( - dcp_world_size=4, + dcp_world_size=dcp_world_size, hash_block_size=hash_block_size, - full_block_size=hash_block_size, - mamba_block_size=hash_block_size, - num_blocks=64, + full_block_size=block_size, + mamba_block_size=block_size, + num_blocks=128, + use_eagle=use_eagle, + num_speculative_blocks=4 if use_eagle else 0, + num_prefill_checkpoint_blocks=checkpoint_blocks, ) manager.coordinator.retention_interval = 0 - - token_ids = list(range(7)) - producer = make_request("producer", token_ids, hash_block_size, sha256) - computed_blocks, num_computed, _ = manager.get_computed_blocks(producer) - assert num_computed == 0 - assert manager.allocate_slots(producer, 6, 0, computed_blocks) is not None + assert manager.coordinator.enable_partial_hash_hits + scheduler = SimpleNamespace( + block_size=block_size * dcp_world_size, + cache_config=SimpleNamespace(block_size=block_size), + hash_block_size=hash_block_size, + max_num_scheduled_tokens=8192, + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), + use_eagle_block_drop=use_eagle, + mamba_partial_cache_hit=True, + mamba_shared_prefix_checkpoint=False, + mamba_has_prefill_checkpoint_blocks=bool(checkpoint_blocks), + mamba_prefill_checkpoint_alignment=16, + ) + producer = make_request( + "producer", list(range(prompt_tokens)), hash_block_size, sha256 + ) + while producer.num_computed_tokens < prompt_tokens: + scheduled = Scheduler._mamba_block_aligned_split( + scheduler, producer, prompt_tokens - producer.num_computed_tokens + ) + assert scheduled > 0 + manager.new_step_starts() + assert manager.allocate_slots(producer, scheduled) is not None + producer.num_computed_tokens += scheduled manager.free(producer) manager.new_step_starts() - consumer = make_request("consumer", token_ids, hash_block_size, sha256) - _, num_computed, _ = manager.get_computed_blocks(consumer) - assert num_computed == 6 + replay = make_request( + "replay", list(range(prompt_tokens + extension_tokens)), hash_block_size, sha256 + ) + _, hit, _ = manager.get_computed_blocks(replay) + expected = ( + max((prompt_tokens // hash_block_size - 1) * hash_block_size, 0) + if use_eagle + else ((prompt_tokens - 1) // hash_block_size * hash_block_size) + ) + assert hit == expected + cached_states = { + end + for end in range(hash_block_size, prompt_tokens + 1, hash_block_size) + if manager.block_pool.get_cached_block( + producer.block_hashes[end // hash_block_size - 1], [1] + ) + is not None + } + assert cached_states == ({expected} if expected else set()) @pytest.mark.parametrize("dcp_world_size", [1, 4]) @@ -339,7 +395,7 @@ def test_hybrid_mamba_align_partial_hash_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 blocks = manager.allocate_slots(req0, 6, num_computed, computed_blocks) @@ -488,7 +544,7 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -538,7 +594,7 @@ def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks(): num_prefill_checkpoint_blocks=1, ) - owner = make_request("owner", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + owner = make_request("owner", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(owner) assert manager.allocate_slots(owner, 6, num_computed, computed_blocks) is not None manager.free(owner) @@ -619,44 +675,70 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle(): assert num_computed == 112 -@pytest.mark.parametrize( - "retention_interval,transient_published", - [(None, True), (0, False)], -) +@pytest.mark.parametrize("retention_interval", [None, 0, 32, 64, 256]) +@pytest.mark.parametrize("hash_block_size", [16, 32]) +@pytest.mark.parametrize("use_eagle", [False, True]) def test_internal_checkpoint_publication_respects_retention( - retention_interval, transient_published + retention_interval, hash_block_size, use_eagle ): - """With retention_interval=0, a mid-prompt internal checkpoint is - request-local and must not enter the prefix cache; the prompt-end - checkpoint keeps the existing publication behavior. Dense retention - (None) is unchanged. - """ - hash_block_size = 16 + """Publish only retained block boundaries and the final prompt checkpoint.""" manager = make_full_mamba_manager( dcp_world_size=1, hash_block_size=hash_block_size, full_block_size=hash_block_size, mamba_block_size=32, num_prefill_checkpoint_blocks=1, + use_eagle=use_eagle, ) manager.coordinator.retention_interval = retention_interval request = make_request("producer", list(range(240)), hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(request) - # Mid-prompt chunk: exports a transient checkpoint at 112. + # The EAGLE adjustment can make a chunk checkpoint block-aligned or partial. assert manager.allocate_slots(request, 128, num_computed, computed_blocks) - transient_hash = request.block_hashes[112 // hash_block_size - 1] + checkpoint = (128 if use_eagle else 127) // hash_block_size * hash_block_size - ( + hash_block_size if use_eagle else 0 + ) + retained_checkpoint = checkpoint % 32 == 0 and ( + retention_interval is None + or (retention_interval > 0 and checkpoint % retention_interval == 0) + ) + transient_hash = request.block_hashes[checkpoint // hash_block_size - 1] transient_hit = manager.block_pool.get_cached_block(transient_hash, [1]) - assert (transient_hit is not None) == transient_published + assert (transient_hit is not None) == retained_checkpoint + offloads = drain_boundary_state_offloads(manager).get(request.request_id, []) + assert any(position == checkpoint for _, _, position in offloads) == ( + retained_checkpoint + ) + assert all( + position % 32 == 0 + and ( + retention_interval is None + or (retention_interval > 0 and position % retention_interval == 0) + ) + for _, _, position in offloads + ) + if retention_interval in (None, 32, 64): + retained_hash = request.block_hashes[128 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(retained_hash, [1]) is not None request.num_computed_tokens = 128 manager.new_step_starts() - # Final chunk: the prompt-end checkpoint at 224 stays published. + # The final checkpoint is reusable regardless of the retention interval. assert manager.allocate_slots(request, 112) is not None - end_hash = request.block_hashes[224 // hash_block_size - 1] + prompt_checkpoint = ( + 240 if use_eagle else 239 + ) // hash_block_size * hash_block_size - (hash_block_size if use_eagle else 0) + end_hash = request.block_hashes[prompt_checkpoint // hash_block_size - 1] assert manager.block_pool.get_cached_block(end_hash, [1]) is not None + assert any( + position == prompt_checkpoint + for _, _, position in drain_boundary_state_offloads(manager).get( + request.request_id, [] + ) + ) def test_transient_checkpoint_evicts_retained_boundary_hash(): @@ -802,9 +884,9 @@ def test_hash_aligned_query_end_uses_regular_partial_tail(): mamba_block_size=mamba_block_size, num_prefill_checkpoint_blocks=1, ) - request = make_request("producer", list(range(14)), hash_block_size, sha256) + request = make_request("producer", list(range(15)), hash_block_size, sha256) - new_blocks = manager.allocate_slots(request, request.num_tokens) + new_blocks = manager.allocate_slots(request, 14) assert new_blocks is not None mamba_manager = manager.coordinator.single_type_managers[1] @@ -928,7 +1010,7 @@ def test_boundary_state_offloads_returns_cow_target(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1004,7 +1086,7 @@ def test_finished_partial_tail_uses_table_source_once(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None source_block_id = manager.get_blocks("0").get_block_ids()[1][1] @@ -1024,7 +1106,7 @@ def test_finished_partial_tail_uses_table_source_once(): enable_caching=True, hash_block_size=hash_block_size, ) - req1 = make_request("1", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = advanced_manager.get_computed_blocks(req1) assert ( advanced_manager.allocate_slots(req1, 6, num_computed, computed_blocks) @@ -1212,7 +1294,7 @@ def test_truncate_computed_blocks_preserves_sparse_prefix_positions(): enable_caching=True, hash_block_size=hash_block_size, ) - producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + producer = make_request("producer", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) blocks, num_computed, _ = manager.get_computed_blocks(producer) assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None manager.free(producer) @@ -1267,7 +1349,7 @@ def test_truncate_computed_blocks_allows_short_mamba_group_only(): enable_caching=True, hash_block_size=hash_block_size, ) - producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + producer = make_request("producer", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) blocks, num_computed, _ = manager.get_computed_blocks(producer) assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None manager.free(producer) @@ -1331,7 +1413,7 @@ def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1411,7 +1493,7 @@ def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1470,7 +1552,7 @@ def test_hybrid_full_attention_partial_hash_hit_uses_cow(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1544,7 +1626,7 @@ def test_hybrid_partial_hit_cow_target_starts_uncached(): hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert num_computed == 0 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -1700,7 +1782,7 @@ def test_cow_retained_blocks_returned_for_release(): enable_caching=True, hash_block_size=hash_block_size, ) - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None @@ -2133,6 +2215,72 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed +@pytest.mark.parametrize("with_sliding_window", [False, True]) +def test_resolved_cache_hit_alignment_matches_coordinator(with_sliding_window): + """Workers place Mamba checkpoints on the engine-resolved alignment, so it + must be where the coordinator lands hits, including when a group without + fine-grained lookups disables partial hits.""" + hash_block_size = 2 + mamba_block_size = 4 * hash_block_size + kv_cache_groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + if with_sliding_window: + kv_cache_groups.append( + KVCacheGroupSpec( + ["swa"], + SlidingWindowSpec( + block_size=2 * hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=2 * hash_block_size, + ), + ) + ) + kv_cache_config = KVCacheConfig( + num_blocks=64, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups + ) + coordinator = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ).coordinator + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(enable_prefix_caching=True), + parallel_config=SimpleNamespace(decode_context_parallel_size=1), + ) + + assert coordinator.enable_partial_hash_hits != with_sliding_window + assert ( + resolve_cache_hit_alignment_tokens( + kv_cache_config, + vllm_config, + coordinator.scheduler_block_size, + hash_block_size, + ) + == coordinator._cache_hit_alignment_tokens + ) + + def test_opted_out_scratch_group_keeps_partial_hash_hits(): hash_block_size = 2 mamba_block_size = 2 * hash_block_size @@ -2179,7 +2327,7 @@ def test_opted_out_scratch_group_keeps_partial_hash_hits(): assert manager.coordinator.enable_partial_hash_hits - req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + req0 = make_request("0", [0, 0, 1, 1, 2, 2, 3], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None manager.free(req0) @@ -2240,7 +2388,7 @@ def test_kpool_tail_supports_128_token_partial_hash_hits(): assert manager.coordinator.enable_partial_hash_hits shared_prefix = [10] * (12 * hash_block_size) - req0 = make_request("0", shared_prefix, hash_block_size, sha256) + req0 = make_request("0", shared_prefix + [11], hash_block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert ( manager.allocate_slots(req0, len(shared_prefix), num_computed, computed_blocks) @@ -2266,7 +2414,7 @@ def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int): ) assert manager.coordinator.enable_partial_hash_hits - req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256) + req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2, 3], 2, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None manager.free(req0) @@ -2318,7 +2466,7 @@ def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot( assert manager.coordinator.single_type_managers[1].block_size == block_size prefix = list(range(12)) - req0 = make_request("snapshot-owner", prefix, block_size, sha256) + req0 = make_request("snapshot-owner", prefix + [12], block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None manager.free(req0) @@ -2355,7 +2503,7 @@ def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots(): mamba_block_size=block_size, ) prefix = list(range(12)) - req0 = make_request("joint-owner", prefix, block_size, sha256) + req0 = make_request("joint-owner", prefix + [12], block_size, sha256) computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None manager.new_step_starts() @@ -2702,6 +2850,7 @@ def test_checkpoint_reservation_matches_worker( mamba_block_size=mamba_block_size, checkpoint_alignment=16, drop_eagle_block=use_eagle, + cache_hit_alignment_tokens=mamba_manager.cache_hit_alignment_tokens, ) reserved = mamba_manager._checkpoints.get(request.request_id) if reserved is not None: diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 427201a7ce..a89fd085f1 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -5047,19 +5047,10 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(): def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): - """An identical resend and a longer sibling resume at DIFFERENT positions. + """EAGLE resend and extension reuse the same retained checkpoint. - How far a lookup matches depends on who is asking. A resend of the same - prompt caps its lookup at ``num_tokens - 1`` (the last token is recomputed - for logits), while a sibling whose prompt merely starts with this one caps - above the prompt. The two coincide unless the prompt length is an exact - multiple of the alignment -- there they differ by one alignment unit, and - under the EAGLE drop BOTH are reachable. - - Retaining only the higher one leaves the resend with every retained state - above every candidate its lookup can produce, and the reconciled hit - collapses to 0 -- the same zero-hit failure sparse retention already fixes - at unaligned prompt lengths. + Attention may inspect the prompt's final block before dropping it, so an + aligned resend must not lose another block to the P - 1 lookup limit. """ block_size = 32 num_spec = 3 @@ -5097,16 +5088,12 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): use_eagle=True, ) - # 128 tokens, an exact multiple of the 32-token alignment. A longer sibling - # matches 128 and drops to 96; this prompt's own resend caps at 127, matches - # 96 and drops to 64. Both states must survive retention. + # Both requests prove 128 attention tokens, drop one block, and reuse 96. token_ids = [i for i in range(4) for _ in range(block_size)] req0 = make_request("0", token_ids, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0) assert num_computed_tokens == 0 - # Prefill in block-aligned chunks the way the align-mode scheduler does: a - # state only materializes as a chunk's running-state block, so a - # single-shot prefill could not retain the lower one. + # Materialize the checkpoint through normal block-aligned prefill chunks. for chunk_end in (32, 64, 96, 128): blocks = manager.allocate_slots( req0, @@ -5118,8 +5105,7 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): assert blocks is not None req0.num_computed_tokens = chunk_end - # Block ``i`` ends at token ``(i + 1) * 32``, so positions 64 and 96 are - # mamba blocks 1 and 2. + # Both materialized replay checkpoints survive; resend and extension use 96. pool = manager.block_pool expected_mamba_cached = {1, 2} for i in range(4): @@ -5130,15 +5116,13 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): assert cached is None, f"mamba hash {i} should not be cached" manager.free(req0) - # The identical resend: full attention matches blocks 0-2 (96 tokens, capped - # by num_tokens - 1) and the EAGLE drop caps the candidate at 64. Without - # the lower state retained the reconciled hit would be 0. + # The identical resend must reuse the retained checkpoint, not miss to zero. req1 = make_request("1", token_ids, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1) - assert num_computed_tokens == 2 * block_size - assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 2] + assert num_computed_tokens == 3 * block_size + assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 3] - # The longer sibling resumes one alignment unit higher, off the same prompt. + # The longer sibling reuses that same checkpoint. longer = make_request("2", token_ids + [9] * block_size, block_size, sha256) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(longer) assert num_computed_tokens == 3 * block_size diff --git a/vllm/config/cache.py b/vllm/config/cache.py index e6436f98f9..5eb26726c7 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -230,6 +230,15 @@ class CacheConfig: for hybrid models where requests occupy multiple KV cache groups.""" kv_cache_max_concurrency: float | None = field(default=None, init=False) """Per-DP-engine maximum concurrency at max_model_len tokens.""" + # The two differ when hashes are finer than the scheduler block but partial + # hash hits are off (e.g. a sliding-window group cannot look up sub-block + # hits), so hits only land on the scheduler block. + # TODO(zjy0516): Merge into one value + hash_block_size: int | None = field(default=None, init=False) + """Tokens per prefix-cache block hash.""" + cache_hit_alignment_tokens: int | None = field(default=None, init=False) + """Token granularity at which prefix-cache hits land: `hash_block_size` + with partial hash hits, otherwise the scheduler block size.""" kv_sharing_fast_prefill: bool = False """In some KV sharing setups, e.g. YOCO (https://arxiv.org/abs/2405.05254), @@ -298,6 +307,8 @@ class CacheConfig: "effective_attention_block_size", "kv_cache_size_tokens", "kv_cache_max_concurrency", + "hash_block_size", + "cache_hit_alignment_tokens", # Feature toggles not impacting compiled graph shape "kv_sharing_fast_prefill", "swa_bounded_replay", diff --git a/vllm/model_executor/layers/mamba/checkpoint.py b/vllm/model_executor/layers/mamba/checkpoint.py index 45888113e1..2f3122f870 100644 --- a/vllm/model_executor/layers/mamba/checkpoint.py +++ b/vllm/model_executor/layers/mamba/checkpoint.py @@ -25,6 +25,7 @@ def compute_mamba_prefill_checkpoints( mamba_block_size: int, checkpoint_alignment: int | None, drop_eagle_block: bool, + cache_hit_alignment_tokens: int, ) -> tuple[list[int], list[int]]: """Per-row internal prefill checkpoint offsets and cache block columns. @@ -43,7 +44,7 @@ def compute_mamba_prefill_checkpoints( for seq_len, query_len in zip(seq_lens, query_lens): query_start = seq_len - query_len position = get_mamba_prefill_checkpoint_position( - seq_len, hash_block_size, drop_eagle_block=drop_eagle_block + seq_len, cache_hit_alignment_tokens, drop_eagle_block=drop_eagle_block ) valid = is_mamba_prefill_checkpoint_valid( query_start=query_start, @@ -106,12 +107,17 @@ class MambaPrefillCheckpointBuilder: return None if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0: return None + cache_config = self.vllm_config.cache_config + hash_block_size = cache_config.hash_block_size + cache_hit_alignment_tokens = cache_config.cache_hit_alignment_tokens + if hash_block_size is None or cache_hit_alignment_tokens is None: + # Unresolved only for the minimal KV cache built for profiling. + return None assert m.seq_lens_cpu_upper_bound is not None all_query_lens = m.query_start_loc_cpu.diff().tolist() query_lens = [all_query_lens[row] for row in request_rows] seq_lens = m.seq_lens_cpu_upper_bound.tolist() block_size = self.kv_cache_spec.block_size - hash_block_size = self.vllm_config.cache_config.prefix_match_unit or block_size speculative_config = self.vllm_config.speculative_config drop_eagle_block = ( speculative_config is not None and speculative_config.use_eagle_block_drop() @@ -123,6 +129,7 @@ class MambaPrefillCheckpointBuilder: mamba_block_size=block_size, checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment, drop_eagle_block=drop_eagle_block, + cache_hit_alignment_tokens=cache_hit_alignment_tokens, ) if not any(checkpoint_offsets): return None diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 75bdb40c24..3b0102fe84 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -479,8 +479,8 @@ class BlockPool: kv_cache_group_id: KV cache group that owns the partial entry. block_size: Cache block size for the owning group. The partial entry hash itself is always the prefix-chain hash at - ``num_tokens``; ``block_size`` is used to assert that the - entry is partial within the owning cache block. + ``num_tokens``; ``block_size`` must be a multiple of the hash + block size. replace_existing_hashes: Whether the block contents were replaced and all existing cache entries must be removed before the new entry is registered. @@ -494,9 +494,6 @@ class BlockPool: return None assert block_size % self.hash_block_size == 0 - assert replace_existing_hashes or ( - block_size > self.hash_block_size and num_tokens % block_size != 0 - ) block_hash = self._get_partial_block_hash(request, num_tokens) num_hash_blocks = num_tokens // self.hash_block_size block_hash_with_group_id = make_block_hash_with_group_id( diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index deb7f89967..c32173a635 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -4,13 +4,14 @@ from abc import ABC, abstractmethod from collections.abc import Sequence from typing import NamedTuple -from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + eagle_proof_margin, + partial_hash_hits_enabled, ) from vllm.v1.core.single_type_kv_cache_manager import ( CrossAttentionManager, @@ -28,8 +29,6 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.request import Request -logger = init_logger(__name__) - def _validate_prefix_cache_retention_interval( retention_interval: int | None, @@ -695,37 +694,15 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): # managers in every group. TP needs hashing finer than the Mamba block; # DCP accepts equality because it scales the effective full-attention # block instead. - has_partial_mamba_group = any( - isinstance(g.kv_cache_spec, MambaSpec) - and g.kv_cache_spec.mamba_cache_mode == "align" - and ( - (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size) - or ( - dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size - ) - ) - for g in kv_cache_config.kv_cache_groups - ) self.enable_partial_hash_hits = ( - allow_partial_hash_hits and has_partial_mamba_group + allow_partial_hash_hits + and partial_hash_hits_enabled( + kv_cache_config.kv_cache_groups, + hash_block_size, + dcp_world_size, + [type(manager) for manager in self.single_type_managers], + ) ) - if self.enable_partial_hash_hits: - unsupported_partial_hit_managers = { - type(manager).__name__ - for manager, group in zip( - self.single_type_managers, kv_cache_config.kv_cache_groups - ) - if group.kv_cache_spec.prefix_cacheable - and not manager.supports_fine_grained_hash_lookup - and manager.block_size != hash_block_size - } - if unsupported_partial_hit_managers: - self.enable_partial_hash_hits = False - logger.warning_once( - "Disabling fine-grained prefix-cache hits because these KV " - "cache managers require block-aligned lookups: %s.", - ", ".join(sorted(unsupported_partial_hit_managers)), - ) cache_hit_alignment_tokens = self._cache_hit_alignment_tokens for manager in self.single_type_managers: manager.cache_hit_alignment_tokens = cache_hit_alignment_tokens @@ -899,15 +876,15 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): # mamba: its finder never drops (draft models have no mamba # layers), so the hit would grow past the candidate. if drop_eagle_block and not isinstance(spec, MambaSpec): - eagle_margin = ( - self.hash_block_size - if self.enable_partial_hash_hits - and manager_cls.supports_fine_grained_hash_lookup - and group_block_size > self.hash_block_size - else group_block_size + eagle_margin = eagle_proof_margin( + group_block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup, ) _max_length = min( - curr_hit_length + eagle_margin, max_cache_hit_length + curr_hit_length + eagle_margin, + len(block_hashes) * self.hash_block_size, ) hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, @@ -983,9 +960,21 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): for spec, group_ids, manager_cls, use_eagle in self.attention_groups: manager = self.single_type_managers[group_ids[0]] + lookup_length = max_cache_hit_length + if use_eagle and not isinstance(spec, MambaSpec): + eagle_margin = eagle_proof_margin( + manager.block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup, + ) + lookup_length = min( + max_cache_hit_length + eagle_margin, + len(block_hashes) * self.hash_block_size, + ) blocks, group_hit = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, - max_length=max_cache_hit_length, + max_length=lookup_length, kv_cache_group_ids=group_ids, block_pool=manager.block_pool, kv_cache_spec=spec, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 79ec2fc5e4..53a01a5cab 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -53,6 +53,7 @@ from vllm.v1.utils import tensor_data if TYPE_CHECKING: from vllm.v1.core.block_pool import BlockPool + from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager # BlockHash represents the hash of a single KV-cache block used for @@ -747,6 +748,60 @@ def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCache return replace(spec, block_size=block_size) +def partial_hash_hits_enabled( + kv_cache_groups: Sequence[KVCacheGroupSpec], + hash_block_size: int, + dcp_world_size: int = 1, + manager_classes: Sequence[type["SingleTypeKVCacheManager"]] | None = None, +) -> bool: + """Whether aligned Mamba states support sub-block prefix-cache hits. + + If ``manager_classes`` (one per group) is given, also require every other + prefix-cacheable group to support fine-grained lookups. + """ + if not any( + isinstance(spec, MambaSpec) + and spec.mamba_cache_mode == "align" + and ( + (dcp_world_size == 1 and spec.block_size > hash_block_size) + or (dcp_world_size > 1 and spec.block_size >= hash_block_size) + ) + for group in kv_cache_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ): + return False + if manager_classes is None: + return True + unsupported_managers = { + manager_cls.__name__ + for group, manager_cls in zip(kv_cache_groups, manager_classes) + if group.kv_cache_spec.prefix_cacheable + and not manager_cls.supports_fine_grained_hash_lookup + and resolve_dcp_kv_block_size(group.kv_cache_spec, dcp_world_size) + != hash_block_size + } + if unsupported_managers: + logger.warning_once( + "Disabling fine-grained prefix-cache hits because these KV " + "cache managers require block-aligned lookups: %s.", + ", ".join(sorted(unsupported_managers)), + ) + return False + return True + + +def eagle_proof_margin( + block_size: int, hash_block_size: int, fine_grained_lookup: bool +) -> int: + """Tokens an EAGLE group matches past a cache hit before dropping them. + + Fine-grained lookups drop one hash unit; others drop one cache block. + """ + if fine_grained_lookup and block_size > hash_block_size: + return hash_block_size + return block_size + + def resolve_kv_cache_block_sizes( kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, @@ -775,8 +830,7 @@ def resolve_kv_cache_block_sizes( if groups and not groups[0].kv_cache_spec.dcp_sharded: dcp = 1 bs = cache_config.block_size * dcp - # The Mamba prefill checkpoint builder reads prefix_match_unit directly, - # so a value dropped here puts its checkpoints off the scheduler's grid. + # Reject a prefix_match_unit that would otherwise be silently ignored. if ( cache_config.prefix_match_unit not in (None, bs) and len(groups) == 1 @@ -797,8 +851,8 @@ def resolve_kv_cache_block_sizes( ] scheduler_block_size = math.lcm(*group_block_sizes) - logger.info("kv cache group sizes %s", group_block_sizes) - logger.info("kv lcm block sizes %s", scheduler_block_size) + logger.info_once("kv cache group sizes %s", tuple(group_block_sizes)) + logger.info_once("kv lcm block sizes %s", scheduler_block_size) # Block hashes are only consumed by prefix caching and KV connectors # (P/D, offloading); when neither is active, keep hash_block_size equal @@ -839,18 +893,10 @@ def resolve_kv_cache_block_sizes( and isinstance(spec.tokens_per_state, int) and spec.tokens_per_state > 1 } - has_partial_mamba_group = any( - isinstance(spec, MambaSpec) - and spec.mamba_cache_mode == "align" - and ( - (dcp == 1 and block_size > hash_block_size) - or (dcp > 1 and block_size >= hash_block_size) - ) - for group, block_size in zip(groups, group_block_sizes) - for spec in iter_layer_specs(group.kv_cache_spec) - ) cache_hit_alignment = ( - hash_block_size if has_partial_mamba_group else scheduler_block_size + hash_block_size + if partial_hash_hits_enabled(groups, hash_block_size, dcp) + else scheduler_block_size ) if any(cache_hit_alignment % alignment for alignment in prefix_alignments): raise ValueError( @@ -861,6 +907,38 @@ def resolve_kv_cache_block_sizes( return scheduler_block_size, hash_block_size +def resolve_cache_hit_alignment_tokens( + kv_cache_config: KVCacheConfig, + vllm_config: VllmConfig, + scheduler_block_size: int, + hash_block_size: int, +) -> int: + """Token granularity at which prefix-cache hits land. + + Mirrors the alignment the hybrid KV cache coordinator pushes to its + managers: fine-grained hits land on ``hash_block_size``, all others on + ``scheduler_block_size``. + """ + groups = kv_cache_config.kv_cache_groups + if not vllm_config.cache_config.enable_prefix_caching or len(groups) <= 1: + return scheduler_block_size + manager_classes = [] + for group in groups: + manager_cls = KVCacheSpecRegistry.get_manager_class( + group.kv_cache_spec, group.role + ) + assert manager_cls is not None + manager_classes.append(manager_cls) + if partial_hash_hits_enabled( + groups, + hash_block_size, + vllm_config.parallel_config.decode_context_parallel_size, + manager_classes, + ): + return hash_block_size + return scheduler_block_size + + def get_request_block_hasher( hash_block_size: int, caching_hash_fn: Callable[[Any], bytes], diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9739dbb3d5..0758c61225 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -381,9 +381,8 @@ class Scheduler(SchedulerInterface): or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0 for group in kv_cache_config.kv_cache_groups ) - # A finer prefix_match_unit is configured: a mamba partial tail entry - # can only be registered by a step ending exactly at the prompt's last - # hash boundary, so the split adds that stop. + # Fine-grained hits reuse the prompt's hash-aligned Mamba checkpoint. + # Split there only if the backend cannot save it within the final chunk. self.mamba_partial_cache_hit = ( self.need_mamba_block_aligned_split and self.hash_block_size < self.block_size @@ -460,14 +459,19 @@ class Scheduler(SchedulerInterface): # The last block-aligned position whose state can be cached. With # Eagle, FullAttn prunes the last matching block, so back off one # block to avoid a Mamba cache miss. - last_cache_position = request.num_tokens - request.num_tokens % block_size - if self.use_eagle_block_drop: - last_cache_position = max(last_cache_position - block_size, 0) + ckpt_block_size = ( + self.hash_block_size if self.mamba_partial_cache_hit else self.block_size + ) + last_cache_position = get_mamba_prefill_checkpoint_position( + request.num_tokens, + block_size, + self.use_eagle_block_drop, + ) end = start + num_new_tokens checkpoint_position = get_mamba_prefill_checkpoint_position( prefill_end, - self.hash_block_size, + ckpt_block_size, drop_eagle_block=self.use_eagle_block_drop, ) use_internal_checkpoint = ( @@ -500,18 +504,14 @@ class Scheduler(SchedulerInterface): next_block_boundary = (start // block_size + 1) * block_size tail_boundary = ( - request.num_prompt_tokens // self.hash_block_size * self.hash_block_size + get_mamba_prefill_checkpoint_position( + request.num_prompt_tokens, + self.hash_block_size, + self.use_eagle_block_drop, + ) if self.mamba_partial_cache_hit and not use_internal_checkpoint else 0 ) - if tail_boundary and self.use_eagle_block_drop: - # Eagle matches one hash unit past the candidate and drops it, so - # nothing proves the prompt's own last hash boundary. Materialize - # the state one unit lower, where the hit can actually land. Keyed on - # the block-drop bit, not plain use_eagle: this shift exists only to - # compensate for the drop, and the Mamba manager's matching gate - # reads the same bit (the coordinator is handed use_eagle_block_drop). - tail_boundary = max(tail_boundary - self.hash_block_size, 0) junction = request.shared_prefix_boundary # Block-floored: a sub-block junction's state is not separately cacheable. block_floored = start + (junction - start) // block_size * block_size @@ -532,8 +532,7 @@ class Scheduler(SchedulerInterface): else 0, # Never run past the last cacheable block boundary mid-chunk. last_cache_position, - # Fine-grained hits: the prompt's partial-tail entry can only be - # registered by a chunk ending exactly at its last hash boundary. + # Stop at the prompt checkpoint unless the backend saves it internally. tail_boundary if last_cache_position < tail_boundary < request.num_prompt_tokens else 0, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 5940a1cd0a..3e7ae1b706 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -876,7 +876,14 @@ class FullAttentionManager(SingleTypeKVCacheManager): block are intentionally skipped. """ hash_block_size = self.block_pool.hash_block_size - boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size + # A resend matches at most `prompt_len - 1` tokens, but EAGLE lookups + # read the whole prompt before dropping a hash unit. + token_limit = ( + request.num_prompt_tokens + if self.use_eagle + else request.num_prompt_tokens - 1 + ) + boundary_tokens = token_limit // hash_block_size * hash_block_size if boundary_tokens == 0 or boundary_tokens > num_tokens: return if boundary_tokens % self.block_size == 0: @@ -1788,7 +1795,7 @@ class MambaManager(SingleTypeKVCacheManager): # is kept, so `prefill_end` skips reserving blocks for the others. checkpoint_position = get_mamba_prefill_checkpoint_position( num_tokens, - self.block_pool.hash_block_size, + self.cache_hit_alignment_tokens, self.drop_eagle_checkpoint_block, ) if not self._needs_internal_checkpoint( @@ -2073,20 +2080,29 @@ class MambaManager(SingleTypeKVCacheManager): checkpoint_position == get_mamba_prefill_checkpoint_position( request.num_prompt_tokens, - hash_block_size, + self.cache_hit_alignment_tokens, self.drop_eagle_checkpoint_block, ) ) + is_retained_checkpoint = False + if checkpoint_position % self.block_size == 0: + block_idx = checkpoint_position // self.block_size - 1 + retained = self.reachable_block_mask( + block_idx, + block_idx + 1, + self.cache_hit_alignment_tokens, + self.kv_cache_spec, + self.drop_eagle_checkpoint_block, + retention_interval=retention_interval, + ) + is_retained_checkpoint = retained is None or retained[0] if ( - retention_interval == 0 - and num_tokens < request.num_prompt_tokens - and not is_prompt_checkpoint + not is_prompt_checkpoint + and not is_retained_checkpoint and checkpoint_position != request.shared_prefix_boundary ): - # retention_interval == 0 keeps this transient checkpoint - # request-local. The slot may carry a hash from this step's - # full-block pass; that must go too, since the checkpoint - # state is about to overwrite the block. + # The transient state overwrites this slot, so neither it nor + # the full-block hash registered earlier this step is reusable. self.block_pool._maybe_evict_cached_block(checkpoint_block) return None if checkpoint_block.block_hash_num_tokens == checkpoint_position: @@ -2099,21 +2115,13 @@ class MambaManager(SingleTypeKVCacheManager): block_size=self.block_size, replace_existing_hashes=True, ) - if self.block_size == hash_block_size: + if num_tokens <= 0 or num_tokens % hash_block_size != 0: return None - if num_tokens % self.block_size == 0: - return None - if num_tokens % hash_block_size != 0: - return None - latest_prompt_hash_boundary = ( - request.num_prompt_tokens // hash_block_size - ) * hash_block_size - if self.drop_eagle_checkpoint_block: - # Eagle groups match one hash unit past the candidate and drop it, - # so register the tail one unit lower. - latest_prompt_hash_boundary = max( - latest_prompt_hash_boundary - hash_block_size, 0 - ) + latest_prompt_hash_boundary = get_mamba_prefill_checkpoint_position( + request.num_prompt_tokens, + hash_block_size, + self.drop_eagle_checkpoint_block, + ) # The junction is the other position a sibling resumes at: where one was # observed to stop, and where the scheduler already ends a chunk. Bounded # to the prompt chunk being computed -- during decode the target is the @@ -2126,7 +2134,7 @@ class MambaManager(SingleTypeKVCacheManager): ): return None - block_idx = num_tokens // self.block_size + block_idx = cdiv(num_tokens, self.block_size) - 1 blocks = self.req_to_blocks[request.request_id] if block_idx >= len(blocks): return None @@ -2141,7 +2149,7 @@ class MambaManager(SingleTypeKVCacheManager): kv_cache_group_id=self.kv_cache_group_id, block_size=self.block_size, ) - if partial_hash is not None: + if partial_hash is not None and num_tokens % self.block_size: self._partial_hit_reqs[request.request_id] = (block_idx, source_block) self.num_cached_block[request.request_id] = block_idx # Producer of this partial tail: the boundary state currently lives diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index e0a003af62..679bbbfd24 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -53,6 +53,7 @@ from vllm.v1.core.kv_cache_utils import ( get_kv_cache_configs, get_request_block_hasher, init_none_hash, + resolve_cache_hit_alignment_tokens, resolve_kv_cache_block_sizes, update_kv_cache_capacity, ) @@ -352,6 +353,19 @@ class EngineCore: vllm_config.validate_block_size() + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( + scheduler_kv_cache_config, vllm_config + ) + cache_hit_alignment_tokens = resolve_cache_hit_alignment_tokens( + scheduler_kv_cache_config, + vllm_config, + scheduler_block_size, + hash_block_size, + ) + for kv_cache_config in kv_cache_configs: + kv_cache_config.hash_block_size = hash_block_size + kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens + self.model_executor.initialize_from_config(kv_cache_configs) if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: self.model_executor.compile_or_warm_up_model() diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index cf707611b4..ca3f485917 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1127,7 +1127,10 @@ def get_mamba_prefill_checkpoint_position( drop_eagle_block: bool, ) -> int: """Return the reusable Mamba checkpoint boundary for a prefill.""" - checkpoint_position = (num_tokens - 1) // hash_block_size * hash_block_size + # Without EAGLE, leave one token to recompute on resend. + # EAGLE's block drop already leaves tokens to recompute. + proof_limit = num_tokens if drop_eagle_block else num_tokens - 1 + checkpoint_position = proof_limit // hash_block_size * hash_block_size if drop_eagle_block: checkpoint_position -= hash_block_size return max(checkpoint_position, 0) @@ -1487,6 +1490,10 @@ class KVCacheConfig: """Resolved retention policy for local prefix-cache checkpoints.""" kv_cache_layout: str | None = None """The KV cache layout resolved by the engine core, adopted by all workers.""" + hash_block_size: int | None = None + """Tokens per prefix-cache block hash, resolved by the engine core.""" + cache_hit_alignment_tokens: int | None = None + """Token granularity of prefix-cache hits, resolved by the engine core.""" hisparse_host_num_blocks: int | None = None """Capacity of the dedicated HiSparse host-block manager, when enabled.""" diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 70427d85a1..f2dfc8a689 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -821,10 +821,15 @@ class Worker(WorkerBase): # so that it's available to the warmup stage. self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks - # Adopt the engine core's layout; workers spawned after resolution - # (e.g. elastic EP scale-up) only see it through the config. + # Adopt the engine core's layout and prefix-cache granularity; workers + # spawned after resolution (e.g. elastic EP scale-up) only see them + # through the config. if kv_cache_config.kv_cache_layout is not None: record_kv_cache_layout(self.cache_config, kv_cache_config.kv_cache_layout) + self.cache_config.hash_block_size = kv_cache_config.hash_block_size + self.cache_config.cache_hit_alignment_tokens = ( + kv_cache_config.cache_hit_alignment_tokens + ) # Init kv cache connector here, because it requires # `kv_cache_config`.