From f8e6e365aec519285d99961015ee270a3871d871 Mon Sep 17 00:00:00 2001 From: Dao007forever Date: Fri, 9 Oct 2026 17:08:08 -0700 Subject: [PATCH] [Bugfix][Mooncake] Save the EAGLE attention proof with hybrid Mamba partial tails (#60695) Signed-off-by: Dao Le Signed-off-by: Dao007forever Signed-off-by: zjy0516 Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: zjy0516 --- .buildkite/test_areas/models_basic.yaml | 2 + tests/models/kimi_k3/test_prefix_cache.py | 139 +++++-- .../unit/test_mooncake_store_coordinator.py | 44 +- .../unit/test_mooncake_store_hma_e2e.py | 95 +++-- .../unit/test_mooncake_store_scheduler.py | 316 ++++++++++++-- .../unit/test_mooncake_store_worker.py | 392 ++++++++++++++++-- .../v1/mooncake/store/coordinator.py | 98 +++-- .../kv_connector/v1/mooncake/store/data.py | 51 ++- .../v1/mooncake/store/scheduler.py | 202 ++++++--- .../kv_connector/v1/mooncake/store/worker.py | 188 ++------- vllm/v1/engine/core.py | 5 +- 11 files changed, 1142 insertions(+), 390 deletions(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index fd21324a7e..1d9a7fe2f3 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -165,6 +165,8 @@ steps: - vllm/v1/core/ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/ + - vllm/distributed/mooncake_store.py - vllm/v1/simple_kv_offload/ - vllm/v1/spec_decode/ - vllm/model_executor/layers/mamba/ diff --git a/tests/models/kimi_k3/test_prefix_cache.py b/tests/models/kimi_k3/test_prefix_cache.py index 20e5d6670e..d6819a507a 100644 --- a/tests/models/kimi_k3/test_prefix_cache.py +++ b/tests/models/kimi_k3/test_prefix_cache.py @@ -4,7 +4,12 @@ import contextlib import json +import os import random +import shutil +import socket +import subprocess +import time from collections.abc import Iterator from dataclasses import dataclass from math import lcm @@ -60,6 +65,16 @@ SPEC_CONFIG = { } NIXL = {"kv_connector": "NixlConnector", "kv_role": "kv_both"} +MOONCAKE = { + "kv_connector": "MooncakeStoreConnector", + "kv_role": "kv_both", + "kv_load_failure_policy": "recompute", + "kv_connector_extra_config": { + "load_async": True, + "lookup_async": True, + "enable_cross_layers_blocks": False, + }, +} # SimpleCPUOffloadConnector as in the Kimi-K3 recipe. OFFLOAD = { "kv_connector": "SimpleCPUOffloadConnector", @@ -243,10 +258,16 @@ def _tokens_text_logprobs(response: dict): def _reset_gpu_prefix_cache(url: str) -> None: - response = requests.post( - f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60 - ) - response.raise_for_status() + deadline = time.monotonic() + 60 + while True: + response = requests.post( + f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60 + ) + response.raise_for_status() + if response.json()["success"]: + return + assert time.monotonic() < deadline, "GPU cache reset timed out" + time.sleep(0.1) @dataclass(frozen=True) @@ -354,6 +375,35 @@ def _format(turns: list[Turn]) -> str: return "\n".join(rows) +def _check_resend_reuses_prompt_checkpoint( + deployment: Deployment, 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)] + with _serve(deployment, mode) as servers: + url = servers[0].url_root + cold = _complete(url, prompt, salt, max_tokens=1) + if deployment.offload: + _reset_gpu_prefix_cache(url) + 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( "prompt_len", [ @@ -368,28 +418,67 @@ def _format(turns: list[Turn]) -> str: ], ) 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) + _check_resend_reuses_prompt_checkpoint(DEPLOYMENTS["plain"], prompt_len) - # 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.fixture +def mooncake_store(tmp_path, monkeypatch) -> Iterator[None]: + """Run a Mooncake master and point the servers' store config at it.""" + if shutil.which("mooncake_master") is None: + pytest.skip("mooncake_master is not installed") + port = get_open_port() + log_path = tmp_path / "mooncake_master.log" + with open(log_path, "w") as log: + master = subprocess.Popen( + [ + "mooncake_master", + f"--port={port}", + f"--metrics_port={get_open_port()}", + "--enable_metric_reporting=false", + "--default_kv_lease_ttl=60000", + ], + stdout=log, + stderr=subprocess.STDOUT, + ) + try: + deadline = time.monotonic() + 60 + while True: + assert master.poll() is None, log_path.read_text() + with ( + contextlib.suppress(OSError), + socket.create_connection(("127.0.0.1", port)), + ): + break + assert time.monotonic() < deadline, "Mooncake master did not start" + time.sleep(0.1) + config_path = tmp_path / "mooncake.json" + config_path.write_text( + json.dumps( + { + "metadata_server": "P2PHANDSHAKE", + "master_server_address": f"127.0.0.1:{port}", + "protocol": "rdma", + "device_name": os.getenv("MOONCAKE_DEVICE_NAME", "mlx5_12"), + "global_segment_size": "4GB", + "local_buffer_size": "1GB", + } + ) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(config_path)) + monkeypatch.setenv("MC_MAX_MR_SIZE", str(4 << 30)) + yield + finally: + master.kill() + master.wait() + + +@pytest.mark.usefixtures("mooncake_store") +def test_mooncake_resend_reuses_prompt_checkpoint() -> None: + """Blocks=6144, PMU=128. P=7449 needs Mamba state at 7296 and attention proof + at 7424, both beyond the normal Mooncake save boundary of 6144: ensure they + remain reusable from the store after a GPU cache reset.""" + _check_resend_reuses_prompt_checkpoint( + Deployment(Instance(kv_config=MOONCAKE), offload=True), 7449 ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 8767b1c0d9..ac52ca07ab 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -30,7 +30,13 @@ def _mamba_align(block_size=32): ) -def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None): +def _make_coord( + groups, + hash_block_size, + use_eagle=False, + retention_interval=None, + enable_partial_hash_hits=False, +): """Construct a coordinator using the natural LCM of group block sizes as the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for the test fixtures.""" @@ -42,6 +48,7 @@ def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=Non hash_block_size=hash_block_size, use_eagle=use_eagle, retention_interval=retention_interval, + enable_partial_hash_hits=enable_partial_hash_hits, ) @@ -219,8 +226,7 @@ def test_coordinator_fine_grained_partial_tail_hit(): KVCacheGroupSpec(["L0"], _full(32)), KVCacheGroupSpec(["L1"], _mamba_align(32)), ] - coord = _make_coord(groups, hash_block_size=16) - assert coord.enable_partial_hash_hits + coord = _make_coord(groups, hash_block_size=16, enable_partial_hash_hits=True) hs = _hashes(4) # 4 hash units of 16 = 64 tokens; block 0 = [0,32), etc. # Both groups: full block 0 (key = last sub-hash hs[1]) + partial boundary # at token 48 (key = hs[2]). No hs[3] -> block 1 is not full. @@ -287,38 +293,6 @@ def test_store_mask_swa_only_window_around_each_lcm_boundary(): assert masks[1] == [False, False, False, True, False, False, False, True] -def test_store_mask_swa_does_not_double_scale_under_dcp(): - """The coordinator receives DCP-resolved specs (worker.py applies - resolve_dcp_kv_cache_spec) and indexes chunks in that already-scaled block - size. So dcp_world_size must NOT further scale reachable_block_mask; the - mask must be identical for any dcp_world_size given the same specs. - - Regression: previously the coordinator forwarded its dcp_world_size, which - double-scaled the SWA block size and produced a wrong mask when dcp > 1. - """ - full = _full(32) - swa = _swa(block_size=8, sliding_window=8) - groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] - scheduler_block_size = lcm(32, 8) - - def make(dcp_world_size): - return MooncakeStoreCoordinator( - groups, - scheduler_block_size=scheduler_block_size, - hash_block_size=8, - dcp_world_size=dcp_world_size, - ) - - mask_dcp1 = make(1).store_mask(64) - mask_dcp2 = make(2).store_mask(64) - - # dcp must not change the mask; both match the correct single-block-size - # result (chunks ending at 32 and 64 -> blocks 3 and 7). - expected = [False, False, False, True, False, False, False, True] - assert mask_dcp1[1] == expected - assert mask_dcp2[1] == expected - - def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm(): """Same hybrid layout but sliding_window=16 (= 2 SWA blocks). Each lcm boundary should now span two SWA tail blocks.""" diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index a8354390d7..0682f9f465 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -18,11 +18,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp MooncakeStoreCoordinator, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + BoundaryPut, ChunkedTokenDatabase, KeyMetadata, LoadSpec, ReqMeta, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( # noqa: E501 + _partial_tail_non_mamba_puts, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import ( # noqa: E501 KVCacheStoreRecvingThread, KVCacheStoreSendingThread, @@ -378,6 +382,15 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size(): assert out[1][2].hex() == fine_hashes[6].hex() +def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta: + """Add the tail's non-Mamba puts as the scheduler does before the worker.""" + req.boundary_puts = [BoundaryPut(*put) for put in req.boundary_puts or []] + req.boundary_puts[:0] = _partial_tail_non_mamba_puts( + thread.coord, req, [db.block_size for db in thread.token_databases] + ) + return req + + def test_sub_block_partial_tail_offload_reads_cow_block(): """Sub-block prompt (the 900/128/1536 shape, scaled to 12/4/16): the partial tail is offloaded for both groups under the boundary sub-hash. The @@ -394,7 +407,12 @@ def test_sub_block_partial_tail_offload_reads_cow_block(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) assert coord.enable_partial_hash_hits class _RecordingStore(_DictStore): @@ -441,11 +459,13 @@ def test_sub_block_partial_tail_offload_reads_cow_block(): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, mamba_cow_block, 12)], + num_prompt_tokens=13, + prefill_end_tokens=20, + completed_token_len=20, + boundary_puts=[(1, mamba_cow_block, 12)], ) - send._maybe_offload_boundary_states(req) + send._maybe_offload_boundary_states(_resolve_partial_tail(send, req)) # boundary = 12 // 4 * 4 = 12 -> keyed by hs[12 // 4 - 1] = hs[2]. partial_hash = hs[2] @@ -473,7 +493,12 @@ def test_offload_syncs_event_before_put(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) event = MagicMock() class _FencedStore(_DictStore): @@ -513,12 +538,13 @@ def test_offload_syncs_event_before_put(): block_hashes=hs, can_save=True, num_prompt_tokens=12, + completed_token_len=12, store_job_id=1, - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 8)], ) req.current_event = event - send.add_request(req) + send.add_request(_resolve_partial_tail(send, req)) send._handle_request(send.request_queue.get()) assert send.request_queue.qsize() == 0 assert store._data @@ -532,7 +558,8 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): offload must persist every FA block up to the boundary — the normal save floors to the lcm, so those blocks are otherwise never written and the consumer's per-group lookup would miss. The mamba boundary block still - reads the core-provided CoW target.""" + reads the core-provided CoW target. Without EAGLE, the 12-token prompt's + resend-safe checkpoint is 8, leaving two attention blocks to save.""" full = FullAttentionSpec(block_size=4, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( block_size=16, @@ -544,7 +571,12 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ] - coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4) + coord = MooncakeStoreCoordinator( + groups, + scheduler_block_size=16, + hash_block_size=4, + enable_partial_hash_hits=True, + ) assert coord.enable_partial_hash_hits class _RecordingStore(_DictStore): @@ -590,19 +622,19 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks(): block_hashes=hs, can_save=True, num_prompt_tokens=12, - boundary_state_offloads=[(1, mamba_cow_block, 12)], + completed_token_len=12, + boundary_puts=[(1, mamba_cow_block, 8)], ) - send._maybe_offload_boundary_states(req) + send._maybe_offload_boundary_states(_resolve_partial_tail(send, req)) - # FA (block 4): full blocks ending at 4, 8 and 12, keyed by their normal + # FA (block 4): full blocks ending at 4 and 8, keyed by their normal # block-end hashes; mamba (block 16): the partial boundary block under # the boundary sub-hash, read from the CoW target. expected = { token_dbs[0].key_for(hs[0]): [1 * 512], token_dbs[0].key_for(hs[1]): [2 * 512], - token_dbs[0].key_for(hs[2]): [3 * 512], - token_dbs[1].key_for(hs[2]): [10_000 + mamba_cow_block * 512], + token_dbs[1].key_for(hs[1]): [10_000 + mamba_cow_block * 512], } assert store.puts == expected @@ -611,11 +643,10 @@ def test_worker_lookup_hits_sub_block_partial_tail(): """worker.lookup must query sub-block keys when partial hash hits are on. Regression test: ``lookup`` hard-coded ``fine_grained = False``, so the - sub-block keys persisted by ``_sub_block_tail_puts`` were never probed and - partial prefix hits silently returned 0 while the store side kept writing - them. Here the mamba block (16) exceeds the hash unit (4), so - ``enable_partial_hash_hits`` is on and the lookup must find the stored - boundary at 12. + partial tail's sub-block keys were never probed and partial prefix hits + silently returned 0 while the store side kept writing them. Here the + mamba block (16) exceeds the hash unit (4), so ``enable_partial_hash_hits`` + is on and the lookup must find the stored boundary at 8. """ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) mamba = MambaSpec( @@ -644,6 +675,8 @@ def test_worker_lookup_hits_sub_block_partial_tail(): KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], mamba), ], + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=4, ) vllm_config = _minimal_vllm_config(cache_block_size=16) # Hash unit 4 < mamba block 16 -> partial hash hits are enabled. @@ -672,7 +705,7 @@ def test_worker_lookup_hits_sub_block_partial_tail(): replicate_config=MagicMock(), ) - # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//4-1]). + # A 12-token prompt leaves a resend-safe sub-block checkpoint at 8. hs = [BlockHash(bytes([i + 1]) * 4) for i in range(5)] req = ReqMeta( req_id="r0", @@ -680,15 +713,16 @@ def test_worker_lookup_hits_sub_block_partial_tail(): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, 7, 12)], + num_prompt_tokens=12, + completed_token_len=12, + boundary_puts=[(1, 7, 8)], ) - send_thread._maybe_offload_boundary_states(req) + send_thread._maybe_offload_boundary_states(_resolve_partial_tail(send_thread, req)) worker.store = store - # A 13-token prompt sharing the prefix must hit the stored boundary at 12. - assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 12 + # A 13-token prompt sharing the prefix must hit the stored boundary at 8. + assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8 def test_worker_setup_tolerates_finer_scratch_group(): @@ -744,6 +778,8 @@ def test_worker_setup_tolerates_finer_scratch_group(): KVCacheGroupSpec(["L1"], mamba), KVCacheGroupSpec(["L2"], scratch), ], + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=8, ) vllm_config = _minimal_vllm_config(cache_block_size=16) # Hash unit 8 divides the participating groups (16) but not the scratch @@ -777,7 +813,7 @@ def test_worker_setup_tolerates_finer_scratch_group(): group_participates=[True, True], ) - # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]). + # Persist the sub-block partial tail at boundary 8 (keyed by hs[8//8-1]). hs = [BlockHash(bytes([i + 1]) * 8) for i in range(3)] req = ReqMeta( req_id="r0", @@ -785,10 +821,11 @@ def test_worker_setup_tolerates_finer_scratch_group(): block_ids=([1], [2]), block_hashes=hs, can_save=True, - num_prompt_tokens=20, - boundary_state_offloads=[(1, 7, 12)], + num_prompt_tokens=12, + completed_token_len=12, + boundary_puts=[(1, 7, 8)], ) - send_thread._maybe_offload_boundary_states(req) + send_thread._maybe_offload_boundary_states(_resolve_partial_tail(send_thread, req)) worker.store = store diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 7ced38d65a..47796eb147 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -7,7 +7,11 @@ from unittest.mock import patch import pytest import torch +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( + MooncakeStoreCoordinator, +) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + BoundaryPut, LoadSpec, MooncakeLookupResult, MooncakeStoreConnectorMetadata, @@ -17,6 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( MooncakeStoreScheduler, + _partial_tail_non_mamba_puts, ) from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks @@ -62,6 +67,31 @@ def _make_bare_scheduler( scheduler._next_store_job_id = 0 scheduler._pinned_saves = {} scheduler._boundary_state_group_ids = frozenset({1}) + scheduler._store_coord = MooncakeStoreCoordinator( + [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=max(16, hash_block_size), + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=max(16, hash_block_size), + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + max(16, hash_block_size), + hash_block_size, + enable_partial_hash_hits=enable_partial_hash_hits, + ) return scheduler @@ -121,6 +151,7 @@ def _make_new_scheduler_output() -> SimpleNamespace: req_id="req-0", num_computed_tokens=0, prompt_token_ids=list(range(32)), + num_prompt_tokens=32, prefill_token_ids=None, block_ids=([0, 1],), block_hashes=[b"h0", b"h1"], @@ -199,6 +230,7 @@ def _add_unfinished_request( ) -> None: request = SimpleNamespace( all_token_ids=token_ids, + num_prompt_tokens=prefill_end_tokens, block_hashes=block_hashes, num_output_placeholders=0, ) @@ -268,6 +300,7 @@ def _make_qsa_hybrid_cache_config(): def test_scheduler_projects_nonprefix_groups_and_mamba_ids(): vllm_config = SimpleNamespace( + speculative_config=None, kv_transfer_config=SimpleNamespace( kv_role="kv_both", kv_connector_extra_config={} ), @@ -528,6 +561,7 @@ def _make_pending_load_unfinished_request( ) -> None: request = SimpleNamespace( num_tokens=num_tokens, + num_prompt_tokens=num_tokens, block_hashes=block_hashes, num_output_placeholders=0, ) @@ -594,6 +628,7 @@ def _make_resumed_unfinished_request( ) -> None: request = SimpleNamespace( all_token_ids=token_ids, + num_prompt_tokens=32, block_hashes=block_hashes, num_computed_tokens=num_computed_tokens, num_output_placeholders=0, @@ -796,13 +831,24 @@ def test_from_request_tracker_load_with_can_load_false_still_saves(): assert tracker.num_saved_tokens == 48 -def test_from_request_tracker_no_load_saves_normally(): +@pytest.mark.parametrize("token_len,save_partial_tail", [(48, False), (13, True)]) +def test_from_request_tracker_no_load_saves_normally(token_len, save_partial_tail): tracker = RequestTracker( req_id="req-0", - token_len=48, + token_len=token_len, allocated_block_ids=([0, 1, 2],), num_saved_tokens=0, + prefill_end_tokens=token_len, ) + if save_partial_tail: + tracker.token_len = 8 + assert ( + ReqMeta.from_request_tracker( + tracker, 16, save_partial_tail=True, num_prompt_tokens=token_len + ) + is None + ) + tracker.token_len = token_len req_meta = ReqMeta.from_request_tracker( tracker, @@ -810,12 +856,70 @@ def test_from_request_tracker_no_load_saves_normally(): load_spec=None, skip_save=False, block_hashes=[b"h0", b"h1", b"h2"], + save_partial_tail=save_partial_tail, + num_prompt_tokens=token_len, ) assert req_meta is not None assert req_meta.can_save is True assert req_meta.load_spec is None - assert tracker.num_saved_tokens == 48 + assert req_meta.token_len_chunk == token_len // 16 * 16 + assert req_meta.completed_token_len == token_len + assert req_meta.boundary_puts is None + assert tracker.num_saved_tokens == token_len // 16 * 16 + assert ( + ReqMeta.from_request_tracker( + tracker, + 16, + save_partial_tail=save_partial_tail, + num_prompt_tokens=token_len, + ) + is None + ) + + +def test_partial_tail_is_resolved_once_on_the_prompt_completing_save(): + scheduler = _make_bare_scheduler( + hash_block_size=4, enable_partial_hash_hits=True, save_decode_cache=True + ) + scheduler._store_group_ids = (0, 1) + _add_unfinished_request( + scheduler, + token_ids=list(range(48)), + block_hashes=[bytes([i]) for i in range(12)], + prefill_end_tokens=45, + ) + block_ids = ([0, 1, 2], [50]) + scheduler._request_trackers["req-0"].allocated_block_ids = block_ids + + def step(num_computed_tokens: int, num_scheduled_tokens: int): + return scheduler.build_connector_meta( + SimpleNamespace( + finished_req_ids=set(), + preempted_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace( + req_ids=["req-0"], + new_block_ids=[()], + num_computed_tokens=[num_computed_tokens], + resumed_req_ids=set(), + ), + num_scheduled_tokens={"req-0": num_scheduled_tokens}, + scheduled_spec_decode_tokens={}, + kv_connector_block_state=_make_connector_block_state(block_ids), + ) + ).requests + + # The save that completes the prompt carries the tail for the worker. + (req_meta,) = step(44, 1) + assert req_meta.publish_partial_tail + assert req_meta.boundary_puts == [(0, 2, 44)] + + # A later save does not recompute or resend it. + (req_meta,) = step(45, 3) + assert req_meta.token_len_chunk == 48 + assert not req_meta.publish_partial_tail + assert req_meta.boundary_puts is None class _StubLookupClient: @@ -1018,7 +1122,7 @@ def _add_pending_partial_tail_request( all_token_ids=list(range(num_tokens)), block_hashes=block_hashes, num_output_placeholders=0, - num_prompt_tokens=12, + num_prompt_tokens=13, ) scheduler._unfinished_requests["req-0"] = (request, block_ids) scheduler._request_trackers["req-0"] = RequestTracker( @@ -1057,7 +1161,7 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) out = _add_pending_partial_tail_request( scheduler, - num_tokens=12, + num_tokens=13, block_hashes=[b"h0", b"h1", b"h2"], block_ids=([0],), ) @@ -1069,8 +1173,8 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): assert req_meta.req_id == "req-0" assert req_meta.can_save is True assert req_meta.token_len_chunk == 0 - assert req_meta.boundary_state_offloads == [(1, 7, 12)] - assert req_meta.num_prompt_tokens == 12 + assert req_meta.boundary_puts == [(1, 7, 12)] + assert req_meta.num_prompt_tokens == 13 assert req_meta.block_ids == ([0],) store_job_id = req_meta.store_job_id assert scheduler._pinned_saves[store_job_id][0] == [7] @@ -1081,34 +1185,71 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): assert tracker.has_pending_offload is True -def test_finished_partial_tail_is_pre_pinned_as_store_job(): +@pytest.mark.parametrize("attention_block_size", [4, 8, 16]) +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("num_workers", [1, 2]) +def test_finished_partial_tail_is_pre_pinned_as_store_job( + attention_block_size, use_eagle, num_workers +): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) scheduler._store_group_ids = (0, 1) scheduler.client = SimpleNamespace(discard=lambda *_: None) + scheduler._num_workers = num_workers + groups = list(scheduler._store_coord.kv_cache_groups) + groups[0] = KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + is_eagle_group=use_eagle, + ) + scheduler._store_coord = MooncakeStoreCoordinator( + groups, 16, 4, use_eagle, enable_partial_hash_hits=True + ) + prompt_tokens = 49 if use_eagle else 45 + attention_ids = list( + range(1, (prompt_tokens + attention_block_size - 1) // attention_block_size + 1) + ) + block_ids = (attention_ids, [50]) request = SimpleNamespace( request_id="req-0", - block_hashes=[b"h0", b"h1", b"h2"], + block_hashes=[bytes([i]) for i in range(prompt_tokens // 4)], + num_computed_tokens=prompt_tokens, + num_prompt_tokens=prompt_tokens, ) scheduler._request_trackers["req-0"] = RequestTracker( req_id="req-0", - token_len=12, - allocated_block_ids=([3], [9]), - token_ids=list(range(12)), - prefill_end_tokens=12, + token_len=prompt_tokens, + allocated_block_ids=block_ids, + token_ids=list(range(prompt_tokens)), + prefill_end_tokens=prompt_tokens, ) - block_ids = ([3], [9]) + pool = scheduler._gpu_block_pool + owned_ids = attention_ids + [50] + pool.touch([pool.blocks[i] for i in owned_ids]) delay_free = scheduler.register_finished_partial_tail( request, block_ids, - [(1, 9, 12)], + [(1, 50, 44)], ) # The exact source is pinned immediately, so the request can be freed # before the next connector metadata build. assert delay_free is False - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1 - assert scheduler._gpu_block_pool.blocks[3].ref_cnt == 0 + proof_end = 48 if use_eagle else 44 + expected = attention_ids[ + 32 // attention_block_size : (proof_end + attention_block_size - 1) + // attention_block_size + ] + pool.free_blocks([pool.blocks[i] for i in owned_ids]) + assert all( + pool.blocks[i].ref_cnt == (1 if i in expected else 0) for i in attention_ids + ) + assert pool.blocks[50].ref_cnt == 1 out = SimpleNamespace( finished_req_ids={"req-0"}, @@ -1132,13 +1273,116 @@ def test_finished_partial_tail_is_pre_pinned_as_store_job(): assert req_meta.token_len_chunk == 0 assert req_meta.block_ids == block_ids assert req_meta.block_hashes == request.block_hashes - assert req_meta.boundary_state_offloads == [(1, 9, 12)] - assert scheduler._pinned_saves[req_meta.store_job_id][0] == [9] - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1 + assert req_meta.boundary_puts == [ + (0, attention_ids[idx], min((idx + 1) * attention_block_size, proof_end)) + for idx in range( + 32 // attention_block_size, + (proof_end + attention_block_size - 1) // attention_block_size, + ) + ] + [(1, 50, 44)] + assert set(scheduler._pinned_saves[req_meta.store_job_id][0]) == {50, *expected} + assert all(pool.blocks[i].ref_cnt == 1 for i in expected + [50]) assert scheduler._finished_partial_tail_metas == {} + for _ in range(num_workers - 1): + scheduler.update_connector_output( + _make_worker_output({req_meta.store_job_id: 1}) + ) + assert all(pool.blocks[i].ref_cnt == 1 for i in expected + [50]) scheduler.update_connector_output(_make_worker_output({req_meta.store_job_id: 1})) - assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 0 + assert all(pool.blocks[i].ref_cnt == 0 for i in owned_ids) + + +def test_partial_tail_boundary_follows_eagle_block_drop_not_group_flags(): + """Core drops the Mamba checkpoint whenever EAGLE block drop is on, even if + only the Mamba group carries the eagle flag.""" + groups = [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + is_eagle_group=True, + ), + ] + coord = MooncakeStoreCoordinator(groups, 16, 4, use_eagle=True) + # Core's checkpoint for a 49-token prompt: 48 minus one EAGLE hash unit. + req_meta = ReqMeta( + req_id="req-0", + token_len_chunk=0, + block_ids=([1, 2, 3, 4], [50]), + block_hashes=[bytes([i]) for i in range(12)], + num_prompt_tokens=49, + completed_token_len=49, + boundary_puts=[BoundaryPut(1, 50, 44)], + ) + + assert _partial_tail_non_mamba_puts(coord, req_meta, [16, 16]) == [ + BoundaryPut(0, 3, 44) + ] + + +@pytest.mark.parametrize( + "completed, mamba_tails, publish, expected", + [ + # The chunk ending at the junction (40) hands off its Mamba state. Its + # attention KV comes from the normal saves, so nothing is added here. + (40, [40], False, []), + # The prompt-completing save: blocks from the LCM floor 64 up to the + # checkpoint 68 plus the EAGLE proof unit, whether or not the junction + # hand-off is in the same save. + (75, [68], True, [BoundaryPut(0, 17, 68), BoundaryPut(0, 18, 72)]), + (75, [40, 68], True, [BoundaryPut(0, 17, 68), BoundaryPut(0, 18, 72)]), + ], +) +def test_shared_prefix_junction_handoff_adds_no_attention_tail( + completed, mamba_tails, publish, expected +): + """``--enable-mamba-shared-prefix-checkpoint`` also hands off a sub-block + Mamba state at the shared-prefix junction, below the prompt checkpoint.""" + groups = [ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + is_eagle_group=True, + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + coord = MooncakeStoreCoordinator( + groups, 16, 4, use_eagle=True, enable_partial_hash_hits=True + ) + # 75-token prompt: checkpoint 72 - 4 = 68. + req_meta = ReqMeta( + req_id="req-0", + token_len_chunk=0, + block_ids=(list(range(1, 20)), [50, 51, 52, 53, 54]), + block_hashes=[bytes([i]) for i in range(18)], + num_prompt_tokens=75, + completed_token_len=completed, + boundary_puts=[BoundaryPut(1, 60 + i, n) for i, n in enumerate(mamba_tails)], + publish_partial_tail=publish, + ) + + assert _partial_tail_non_mamba_puts(coord, req_meta, [4, 16]) == expected def test_decode_boundary_state_offload_dropped_unclaimed(): @@ -1235,7 +1479,7 @@ def test_boundary_state_group_ids_are_remapped_to_store_projection(): scheduler._handle_boundary_state_offloads({"req-0": [(2, 7, 800)]}, meta) - assert meta.requests[0].boundary_state_offloads == [(1, 7, 800)] + assert meta.requests[0].boundary_puts == [(1, 7, 800)] def test_resumed_prefill_claims_boundaries_past_prompt_length(): @@ -1245,13 +1489,13 @@ def test_resumed_prefill_claims_boundaries_past_prompt_length(): # missing for boundaries full attention does store and no joint hit could # complete there. scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) - _register_offload_request(scheduler, prefill_end_tokens=20, num_prompt_tokens=12) - out = _make_offload_only_output([(1, 7, 16), (1, 9, 20), (1, 11, 24)]) + _register_offload_request(scheduler, prefill_end_tokens=37, num_prompt_tokens=13) + out = _make_offload_only_output([(1, 7, 16), (1, 9, 32), (1, 11, 48)]) meta = scheduler.build_connector_meta(out) - # 16 and 20 are inside the resumed prefill; 24 is past it. - assert meta.requests[0].boundary_state_offloads == [(1, 7, 16), (1, 9, 20)] + # Aligned states at 16 and 32 are inside the resumed prefill; 48 is past it. + assert meta.requests[0].boundary_puts == [(1, 7, 16), (1, 9, 32)] store_job_id = meta.requests[0].store_job_id assert scheduler._pinned_saves[store_job_id][0] == [7, 9] @@ -1267,7 +1511,7 @@ def test_boundary_state_job_pins_exact_blocks_once(): meta = scheduler.build_connector_meta(out) req_meta = meta.requests[0] - assert req_meta.boundary_state_offloads == [ + assert req_meta.boundary_puts == [ (1, 7, 16), (1, 8, 32), (1, 9, 48), @@ -1294,6 +1538,7 @@ def test_store_job_pins_current_non_null_non_mamba_blocks(): request = SimpleNamespace( all_token_ids=list(range(48)), block_hashes=[bytes([i]) for i in range(12)], + num_prompt_tokens=48, num_output_placeholders=0, ) scheduler._unfinished_requests["req-0"] = (request, ([7, 2, 0], [21, 0, 22])) @@ -1379,7 +1624,7 @@ def test_preemption_and_request_id_reuse_do_not_release_inflight_job(): _register_offload_request(scheduler, prefill_end_tokens=64, num_prompt_tokens=64) second = scheduler.build_connector_meta(_make_offload_only_output([(1, 8, 32)])) second_job_id = second.requests[0].store_job_id - assert second.requests[0].boundary_state_offloads == [(1, 8, 32)] + assert second.requests[0].boundary_puts == [(1, 8, 32)] scheduler.update_connector_output(_make_worker_output({first_job_id: 1})) assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 0 @@ -1412,9 +1657,9 @@ def test_resumed_partial_tail_uses_exact_boundary(): meta = scheduler.build_connector_meta(out) assert len(meta.requests) == 1 - assert meta.requests[0].boundary_state_offloads == [(1, 7, 12)] - # Ordinary metadata retains the full resumed prefill range. - assert meta.requests[0].num_prompt_tokens == 20 + assert meta.requests[0].boundary_puts == [(1, 7, 12)] + assert meta.requests[0].num_prompt_tokens == 13 + assert meta.requests[0].prefill_end_tokens == 20 tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 0 assert tracker.has_pending_offload is True @@ -1427,7 +1672,7 @@ def test_resumed_partial_tail_attached_to_save_keeps_exact_boundary(): all_token_ids=list(range(48)), block_hashes=[b"h0", b"h1", b"h2"], num_output_placeholders=0, - num_prompt_tokens=36, + num_prompt_tokens=37, ) scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) scheduler._request_trackers["req-0"] = RequestTracker( @@ -1445,8 +1690,9 @@ def test_resumed_partial_tail_attached_to_save_keeps_exact_boundary(): assert len(meta.requests) == 1 assert meta.requests[0].can_save is True - assert meta.requests[0].boundary_state_offloads == [(0, 7, 36)] - assert meta.requests[0].num_prompt_tokens == 48 + assert meta.requests[0].boundary_puts == [(0, 7, 36)] + assert meta.requests[0].num_prompt_tokens == 37 + assert meta.requests[0].prefill_end_tokens == 48 # Ordinary saving still covers the full resumed prefill range. tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 48 @@ -1460,7 +1706,7 @@ def test_partial_tail_cow_block_is_referenced_for_the_job(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) out = _add_pending_partial_tail_request( scheduler, - num_tokens=12, + num_tokens=13, block_hashes=[b"h0", b"h1", b"h2"], block_ids=([0],), ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index cb6bd8d995..09d27eb947 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -30,6 +30,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import ( ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( BlobBlockHashes, + BoundaryPut, ChunkedTokenDatabase, KeyMetadata, LBHNCStoreLayout, @@ -44,6 +45,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import ( MooncakeStoreConnectorStats, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( # noqa: E501 + _partial_tail_non_mamba_puts, +) from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash from vllm.v1.kv_cache_interface import ( @@ -942,6 +946,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys +def _resolve_partial_tail(thread, req: ReqMeta) -> ReqMeta: + """Add the tail's non-Mamba puts as the scheduler does before the worker.""" + req.boundary_puts = [BoundaryPut(*put) for put in req.boundary_puts or []] + req.boundary_puts[:0] = _partial_tail_non_mamba_puts( + thread.coord, req, [db.block_size for db in thread.token_databases] + ) + return req + + def _make_partial_tail_send_thread( store, *, @@ -954,6 +967,8 @@ def _make_partial_tail_send_thread( hash_block_size=4, lcm_block_size=16, mamba_group_ids={1}, + use_eagle=False, + eagle_proof_margin_by_group={}, ) db = ChunkedTokenDatabase( KeyMetadata("test-model", 0, 0, 0, 0), @@ -987,17 +1002,74 @@ def _make_partial_tail_req(block_ids: list[int]) -> ReqMeta: block_ids=(block_ids, [1]), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 12)], + num_prompt_tokens=13, + completed_token_len=13, ) +@pytest.mark.parametrize("use_eagle", [False, True]) +def test_partial_tail_ignores_handoff_off_prompt_boundary(use_eagle): + """A Mamba hand-off away from the prompt checkpoint (a shared-prefix + junction) gets no attention tail; only its own Mamba put remains.""" + store = MagicMock() + thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = use_eagle + thread.coord.eagle_proof_margin_by_group = {0: 4} if use_eagle else {} + req = _make_partial_tail_req([1, 2, 3]) + req.num_prompt_tokens = 17 if use_eagle else 13 + req.boundary_puts = [(1, 7, 8)] + + assert _resolve_partial_tail(thread, req).boundary_puts == [(1, 7, 8)] + + +def test_eagle_attention_proof_published_after_checkpoint_handoff(): + store = MagicMock() + stored = set() + store.batch_is_exist.side_effect = lambda keys: [int(k in stored) for k in keys] + + def put(keys, *args): + stored.update(keys) + return [256] * len(keys) + + store.batch_put_from_multi_buffers.side_effect = put + thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = True + thread.coord.eagle_proof_margin_by_group = {0: 4} + metadata = _make_partial_tail_req([1, 2, 3]) + metadata.num_prompt_tokens = 13 + metadata.completed_token_len = 8 + metadata.boundary_puts = [(1, 7, 8)] + + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, metadata) + ) + mamba_key = thread.token_databases[1].key_for(b"a1") + attention_key = thread.token_databases[0].key_for(b"a2") + assert mamba_key in stored + assert attention_key not in stored + + metadata.completed_token_len = 13 + metadata.boundary_puts = None + metadata.publish_partial_tail = True + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, metadata) + ) + assert attention_key in stored + keys, addrs, *_ = store.batch_put_from_multi_buffers.call_args.args + assert addrs[keys.index(attention_key)] == [0x1000 + 3 * 256] + assert mamba_key not in keys + + def test_partial_tail_offload_skips_null_source_blocks(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) thread = _make_partial_tail_send_thread(store) - assert thread._maybe_offload_boundary_states(_make_partial_tail_req([0, 2, 3])) + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, _make_partial_tail_req([0, 2, 3])) + ) keys, addrs, _sizes, _replicate_config = ( store.batch_put_from_multi_buffers.call_args.args @@ -1131,9 +1203,11 @@ def test_partial_tail_offload_skips_cap_omitted_mamba_group(): block_hashes=[b"a0", b"a1", b"a2"], can_save=True, # The scheduler accepted group 1 and omitted group 2 at boundary 12. - boundary_state_offloads=[(1, 7, 12)], + boundary_puts=[(1, 7, 12)], + num_prompt_tokens=13, + completed_token_len=13, ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _replicate_config = ( store.batch_put_from_multi_buffers.call_args.args @@ -1156,7 +1230,9 @@ def test_partial_tail_offload_replaces_stale_group_ids_after_filtering(): supports_group_ids=True, ) - assert thread._maybe_offload_boundary_states(_make_partial_tail_req([1, 2, 3])) + assert thread._maybe_offload_boundary_states( + _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1190,13 +1266,17 @@ def test_partial_tail_put_failure_activates_pressure_gate(): store.batch_put_from_multi_buffers.return_value = [256, -200, 256] thread = _make_partial_tail_send_thread(store) - _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) + _run_store_req( + thread, _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) assert thread._store_pressure_active is True assert thread._skip_store_requests == {"req-a"} assert thread._saved_offset.get("req-a", 0) == 0 - _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) + _run_store_req( + thread, _resolve_partial_tail(thread, _make_partial_tail_req([1, 2, 3])) + ) assert store.batch_put_from_multi_buffers.call_count == 1 @@ -1277,9 +1357,9 @@ def test_block_aligned_snapshot_offload_uses_provided_block(): block_ids=([1, 2, 3], [5]), block_hashes=hs, can_save=True, - boundary_state_offloads=[(1, 7, 32)], + boundary_puts=[(1, 7, 32)], ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args # boundary 32 is block-aligned for the mamba group (block 16): one key, @@ -1289,7 +1369,9 @@ def test_block_aligned_snapshot_offload_uses_provided_block(): assert addrs == [[0x2000 + 7 * 256]] -def test_mixed_snapshot_and_sub_block_offloads(): +@pytest.mark.parametrize("saved_tokens", [0, 32]) +@pytest.mark.parametrize("use_eagle", [False, True]) +def test_mixed_snapshot_and_sub_block_offloads(saved_tokens, use_eagle): """A retention snapshot and the prompt-end sub-block CoW tail can arrive in one hand-off; the sub-block path covers FA gap blocks but reads the mamba boundary only from the CoW block, never positionally.""" @@ -1297,39 +1379,150 @@ def test_mixed_snapshot_and_sub_block_offloads(): store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) thread = _make_partial_tail_send_thread(store) + thread.coord.use_eagle = use_eagle + thread.coord.eagle_proof_margin_by_group = {0: 4} if use_eagle else {} + thread._saved_offset["req-a"] = saved_tokens - hs = [bytes([i + 1]) * 4 for i in range(11)] # 11 hash units = 44 tokens + hs = [bytes([i + 1]) * 4 for i in range(12)] req = ReqMeta( req_id="req-a", token_len_chunk=0, - block_ids=([1, 2, 3], [0, 0, 0]), + block_ids=(list(range(1, 14)), [0, 0, 0]), block_hashes=hs, can_save=True, - boundary_state_offloads=[(1, 9, 32), (1, 7, 44)], + boundary_puts=[(1, 9, 32), (1, 7, 44)], + num_prompt_tokens=49 if use_eagle else 45, + completed_token_len=49 if use_eagle else 45, ) - assert thread._maybe_offload_boundary_states(req) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args db_full, db_mamba = thread.token_databases + proof_end = 48 if use_eagle else 44 assert keys == [ + # The tail reads only the LCM gap, even when normal saves lag. + *[db_full.key_for(BlockHash(hs[i])) for i in range(8, proof_end // 4)], # Aligned snapshot: boundary 32 from the handed-off block 9. db_mamba.key_for(BlockHash(hs[7])), - # Sub-block boundary 44: FA gap blocks ending at 4, 8, 12. - db_full.key_for(BlockHash(hs[0])), - db_full.key_for(BlockHash(hs[1])), - db_full.key_for(BlockHash(hs[2])), # Mamba boundary block from the CoW hand-off (block 7). db_mamba.key_for(BlockHash(hs[10])), ] assert addrs == [ + *[[0x1000 + (i + 1) * 256] for i in range(8, proof_end // 4)], [0x2000 + 9 * 256], - [0x1000 + 1 * 256], - [0x1000 + 2 * 256], - [0x1000 + 3 * 256], [0x2000 + 7 * 256], ] +@pytest.mark.parametrize("tp_rank", [0, 1]) +def test_partial_tail_with_smaller_mamba_blocks_writes_one_mamba_key(tp_rank): + """Under DCP2 the attention block spans the LCM (16) and Mamba blocks are + half that. The prompt-completing save writes one attention block and one + Mamba key: the boundary state from the CoW block, never the interior + Mamba blocks.""" + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + coord = SimpleNamespace( + enable_partial_hash_hits=True, + hash_block_size=4, + lcm_block_size=16, + mamba_group_ids={1}, + use_eagle=False, + eagle_proof_margin_by_group={}, + ) + db_full = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0), block_size=16, hash_block_size=4 + ) + db_full.set_kv_caches_base_addr([0x1000]) + db_full.set_block_len([256]) + db_mamba = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=8, + hash_block_size=4, + ) + db_mamba.set_kv_caches_base_addr([0x2000]) + db_mamba.set_block_len([256]) + thread = _make_store_sending_thread( + store, coord=coord, token_databases=[db_full, db_mamba], tp_rank=tp_rank + ) + # Mamba is replicated across two ranks, which split its keys by block. + thread.group_put_steps = [1, 2] + + hs = [bytes([i + 1]) * 4 for i in range(11)] + req = ReqMeta( + req_id="req-a", + token_len_chunk=0, + block_ids=([1, 2, 3], list(range(20, 26))), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, 7, 44)], + num_prompt_tokens=45, + completed_token_len=45, + publish_partial_tail=True, + ) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) + assert req.boundary_puts == [(0, 3, 44), (1, 7, 44)] + + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args + # Mamba block cdiv(44, 8) - 1 = 5 belongs to put_step_rank 5 % 2 = 1, + # which is tp_rank 0 for group 1. + mamba_puts = [db_mamba.key_for(BlockHash(hs[10]))] if tp_rank == 0 else [] + assert keys == [db_full.key_for(BlockHash(hs[10])), *mamba_puts] + assert addrs == [[0x1000 + 3 * 256], *([[0x2000 + 7 * 256]] * len(mamba_puts))] + + +def test_shared_prefix_junction_handoff_writes_only_its_mamba_key(): + """A junction hand-off (``--enable-mamba-shared-prefix-checkpoint``) below + the prompt checkpoint stores its Mamba state under the junction hash and + no attention tail.""" + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + coord = SimpleNamespace( + enable_partial_hash_hits=True, + hash_block_size=4, + lcm_block_size=16, + mamba_group_ids={1}, + use_eagle=True, + eagle_proof_margin_by_group={0: 4}, + ) + db_full = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0), block_size=4, hash_block_size=4 + ) + db_full.set_kv_caches_base_addr([0x1000]) + db_full.set_block_len([256]) + db_mamba = ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), + block_size=16, + hash_block_size=4, + ) + db_mamba.set_kv_caches_base_addr([0x2000]) + db_mamba.set_block_len([256]) + thread = _make_store_sending_thread( + store, coord=coord, token_databases=[db_full, db_mamba] + ) + + hs = [bytes([i + 1]) * 4 for i in range(18)] + # 75-token prompt (checkpoint 68), chunk cut at the junction 40. + req = ReqMeta( + req_id="req-b", + token_len_chunk=0, + block_ids=(list(range(1, 11)), [20, 21, 22]), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, 7, 40)], + num_prompt_tokens=75, + completed_token_len=40, + ) + assert thread._maybe_offload_boundary_states(_resolve_partial_tail(thread, req)) + assert req.boundary_puts == [(1, 7, 40)] + + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args + assert keys == [db_mamba.key_for(BlockHash(hs[9]))] + assert addrs == [[0x2000 + 7 * 256]] + + def test_snapshot_offload_skips_null_handoff_block(): """A hand-off the core could not materialize (null block) carries no committed state; persisting it would poison the boundary key.""" @@ -1338,13 +1531,16 @@ def test_snapshot_offload_skips_null_handoff_block(): hs = [bytes([i + 1]) * 4 for i in range(8)] assert thread._maybe_offload_boundary_states( - ReqMeta( - req_id="req-a", - token_len_chunk=0, - block_ids=([1, 2, 3], [5]), - block_hashes=hs, - can_save=True, - boundary_state_offloads=[(1, NULL_BLOCK_ID, 32)], + _resolve_partial_tail( + thread, + ReqMeta( + req_id="req-a", + token_len_chunk=0, + block_ids=([1, 2, 3], [5]), + block_hashes=hs, + can_save=True, + boundary_puts=[(1, NULL_BLOCK_ID, 32)], + ), ) ) store.batch_is_exist.assert_not_called() @@ -1398,6 +1594,7 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): [True, False], ), ) + coord.store_mask = MagicMock(side_effect=coord.store_mask) db_full = ChunkedTokenDatabase( KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), @@ -1427,6 +1624,8 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): block_ids=([0, 1, 2, 3], [0, 1, 2, 3]), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, + num_prompt_tokens=33, + prefill_end_tokens=64, ), ) @@ -1436,6 +1635,7 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): assert full_hashes == [b"a2".hex(), b"a3".hex()] assert masked_hashes == [b"a2".hex()] + coord.store_mask.assert_called_once_with(64, 32, num_prompt_tokens=64) def test_store_sending_thread_prepares_missing_chunks_once_per_group(): @@ -3580,6 +3780,139 @@ def _make_bare_worker( return worker +@pytest.mark.parametrize("use_eagle", [False, True]) +@pytest.mark.parametrize("prefix_match_unit", [None, 8]) +@pytest.mark.parametrize("dcp_world_size,mamba_block_size", [(1, 32), (8, 16)]) +def test_mooncake_lookup_reuses_resolved_hash_checkpoint( + use_eagle, prefix_match_unit, dcp_world_size, mamba_block_size +): + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import ( + coordinator as mooncake_coordinator, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( + MooncakeStoreScheduler, + ) + + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + is_eagle_group=use_eagle, + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + config = SimpleNamespace( + speculative_config=None, + kv_transfer_config=SimpleNamespace( + kv_role="kv_both", kv_connector_extra_config={} + ), + kv_events_config=None, + cache_config=SimpleNamespace( + block_size=16, + enable_prefix_caching=True, + prefix_match_unit=prefix_match_unit, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=dcp_world_size, world_size=dcp_world_size + ), + ) + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "scheduler.LookupKeyClient" + ): + scheduler = MooncakeStoreScheduler( + config, + KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=groups, + # The hit alignment the engine core resolves for this config. + cache_hit_alignment_tokens=prefix_match_unit or 16, + ), + ) + assert scheduler.enable_partial_hash_hits + hash_block_size = prefix_match_unit or 16 + assert scheduler._hash_block_size == hash_block_size + coordinator = mooncake_coordinator.MooncakeStoreCoordinator( + groups, + scheduler_block_size=math.lcm(16 * dcp_world_size, mamba_block_size), + hash_block_size=hash_block_size, + enable_partial_hash_hits=True, + ) + hashes = [BlockHash(bytes([i]) * 16) for i in range(512 // hash_block_size)] + pool = mooncake_coordinator.ExternalCachedBlockPool( + hash_block_size=hash_block_size, + exists={(group, bytes(h)) for group in range(2) for h in hashes}, + ) + _, hit = coordinator.find_longest_cache_hit(hashes, 511, pool) + assert hit == 512 - hash_block_size + + +@pytest.mark.parametrize("alignment, expected", [(8, True), (32, False), (None, False)]) +def test_scheduler_follows_core_hit_alignment(alignment, expected): + """The store takes partial hits from core's resolved alignment, not from its + own groups: these Mamba specs alone would allow them, but core may still + keep hits on the block (e.g. for a sliding-window group it never sees).""" + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( + MooncakeStoreScheduler, + ) + + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=32, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=32, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ] + config = SimpleNamespace( + speculative_config=None, + kv_transfer_config=SimpleNamespace( + kv_role="kv_both", kv_connector_extra_config={} + ), + kv_events_config=None, + cache_config=SimpleNamespace( + block_size=32, enable_prefix_caching=True, prefix_match_unit=8 + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=1, world_size=1), + ) + with patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "scheduler.LookupKeyClient" + ): + scheduler = MooncakeStoreScheduler( + config, + KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=groups, + cache_hit_alignment_tokens=alignment, + ), + ) + + assert scheduler.enable_partial_hash_hits is expected + assert scheduler._store_coord.enable_partial_hash_hits is expected + + def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): worker = _make_bare_worker() worker.tp_size = 4 @@ -3873,6 +4206,7 @@ def test_lookup_partial_tail_uses_hash_alignment(): worker._kv_cache_groups, scheduler_block_size=16, hash_block_size=4, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) worker.store.batch_is_exist.return_value = [0, 0, 1, 0, 0, 1] @@ -3946,6 +4280,7 @@ def test_lookup_plan_resolves_group_tail_keys_from_existing_hashes(): scheduler_block_size=16, hash_block_size=4, use_eagle=True, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) hashes = [BlockHash(f"h{i}".encode()) for i in range(6)] @@ -4015,6 +4350,7 @@ def test_lookup_plan_recovers_tail_key_after_multi_chunk_convergence(): worker._kv_cache_groups, scheduler_block_size=16, hash_block_size=4, + enable_partial_hash_hits=True, ) _refresh_group_tp_replication_factors(worker) hashes = [BlockHash(f"h{i}".encode()) for i in range(12)] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index d05b553b64..48a36bc82f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -3,7 +3,8 @@ """External-store cache-hit coordinator for MooncakeStoreConnector.""" from collections.abc import Sequence -from typing import NamedTuple, cast +from dataclasses import replace +from typing import TYPE_CHECKING, NamedTuple, cast from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( chunk_hashes_for_block_size, @@ -13,12 +14,15 @@ from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, + eagle_proof_margin, + resolve_dcp_kv_cache_spec, ) from vllm.v1.core.single_type_kv_cache_manager import ( SingleTypeKVCacheManager, ) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, MambaSpec, @@ -26,6 +30,9 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +if TYPE_CHECKING: + from vllm.config import VllmConfig + class StoreSpecGroup(NamedTuple): spec: KVCacheSpec @@ -82,7 +89,7 @@ class MooncakeStoreCoordinator: hash_block_size: int, use_eagle: bool = False, retention_interval: int | None = None, - dcp_world_size: int = 1, + enable_partial_hash_hits: bool = False, ) -> None: # Mirrors core's resolve_kv_cache_block_sizes: the hash unit only has # to divide groups that participate in prefix caching. Non-shareable @@ -110,14 +117,51 @@ class MooncakeStoreCoordinator: } self.hash_block_size = hash_block_size self.lcm_block_size = scheduler_block_size - self.enable_partial_hash_hits = partial_hash_hits_enabled( - kv_cache_groups, hash_block_size, dcp_world_size - ) + # Core's decision: whether hits land on the hash unit, below the block. + self.enable_partial_hash_hits = enable_partial_hash_hits self.use_eagle = use_eagle # Mirror vLLM core's KVCacheCoordinator.retention_interval. self.retention_interval = retention_interval self._verify_and_split_kv_cache_groups() + @classmethod + def from_kv_cache_config( + cls, + kv_cache_config: KVCacheConfig, + vllm_config: "VllmConfig", + scheduler_block_size: int, + hash_block_size: int, + ) -> "MooncakeStoreCoordinator": + """Build the coordinator the connector's scheduler and workers share. + + Specs are DCP-resolved, so a block spans the tokens of one block-table + entry. Partial hits follow core's hit alignment, which checks every KV + cache group and its manager, not only the groups the store transfers; + unset (a config the engine core did not resolve) keeps hits on the + block. + """ + dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size + spec_config = vllm_config.speculative_config + alignment = kv_cache_config.cache_hit_alignment_tokens + return cls( + [ + replace( + group, + kv_cache_spec=resolve_dcp_kv_cache_spec( + group.kv_cache_spec, dcp_world_size + ), + ) + for group in kv_cache_config.prefix_cacheable_groups + ], + scheduler_block_size, + hash_block_size, + use_eagle=spec_config is not None and spec_config.use_eagle_block_drop(), + retention_interval=kv_cache_config.prefix_cache_retention_interval, + enable_partial_hash_hits=( + alignment is not None and alignment < scheduler_block_size + ), + ) + def align_lookup_length(self, length: int) -> int: alignment = ( self.hash_block_size @@ -167,6 +211,17 @@ class MooncakeStoreCoordinator: self.eagle_group_ids = { gid for g in attention_groups if g.use_eagle for gid in g.group_ids } + self.eagle_proof_margin_by_group = { + gid: eagle_proof_margin( + group.spec.block_size, + self.hash_block_size, + self.enable_partial_hash_hits + and group.manager_cls.supports_fine_grained_hash_lookup, + ) + for group in attention_groups + if group.use_eagle and not isinstance(group.spec, MambaSpec) + for gid in group.group_ids + } def find_longest_cache_hit( self, @@ -394,14 +449,16 @@ class MooncakeStoreCoordinator: # never drops a block, so a widened bound would match past the # attention-verified hit and resume from speculative state (#43559). 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 spec.block_size > self.hash_block_size - else spec.block_size + eagle_margin = eagle_proof_margin( + spec.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, + len(block_hashes) * self.hash_block_size, ) - _max_length = min(curr_hit_length + eagle_margin, max_length) hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, # type: ignore[arg-type] max_length=_max_length, @@ -449,20 +506,3 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec: if isinstance(spec, UniformTypeKVCacheSpecs): return next(iter(spec.kv_cache_specs.values())) return spec - - -def partial_hash_hits_enabled( - kv_cache_groups: Sequence[KVCacheGroupSpec], - hash_block_size: int, - dcp_world_size: int = 1, -) -> bool: - """Match core's DCP-aware Mamba partial-hit condition.""" - return any( - isinstance(spec := _unwrap_spec(g.kv_cache_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 g in kv_cache_groups - ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 2510a7ee8e..4cd0402b7b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -7,7 +7,7 @@ from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from typing import cast +from typing import NamedTuple, cast import numpy as np import torch @@ -712,6 +712,7 @@ class RequestTracker: # For a fresh request this is len(prompt). For a resumed-from-preemption # request it includes previously-generated tokens, which are re-prefilled. prefill_end_tokens: int = 0 + partial_tail_sent: bool = False def reset(self) -> None: self.token_len = 0 @@ -720,6 +721,7 @@ class RequestTracker: self.token_ids = None self.has_pending_offload = False self.prefill_end_tokens = 0 + self.partial_tail_sent = False def update( self, @@ -739,6 +741,14 @@ class RequestTracker: existing.extend(new) +class BoundaryPut(NamedTuple): + """A block stored under the hash of the prefix ending at ``num_tokens``.""" + + group_id: int + block_id: int + num_tokens: int + + @dataclass class ReqMeta: """Per-request metadata for store put/get operations.""" @@ -756,15 +766,24 @@ class ReqMeta: # Absolute request offset represented by token_ids[0]. token_ids_start: int = 0 num_prompt_tokens: int | None = None + # Total prefill tokens, including generated tokens replayed after preemption. + prefill_end_tokens: int | None = None # Identifies this store job for the engine's lifetime. A request id cannot # serve that purpose: it is reused once a preempted request resumes, so it # would release the wrong job's blocks. store_job_id: int | None = None - # Core-provided (group_id, block_id, boundary_tokens) mamba "align" - # boundary states. A block-aligned entry is a committed boundary snapshot; - # a non-aligned entry is the sub-block CoW tail. The store-job reference - # keeps each exact block alive until every worker rank finishes the job. - boundary_state_offloads: list[tuple[int, int, int]] | None = None + # Blocks stored under an explicit prefix hash rather than by position: + # core-provided mamba "align" states (a block-aligned entry is a committed + # snapshot, a non-aligned one the sub-block CoW prompt tail) and the + # scheduler-resolved non-Mamba blocks of the prompt's partial tail. The + # store-job reference keeps each block alive until every worker rank + # finishes the job. + boundary_puts: list[BoundaryPut] | None = None + # Total computed prefix length at the end of this step. + completed_token_len: int | None = None + # Set on the save that first covers the whole prompt, so the scheduler + # publishes the prompt's partial tail once rather than on every save. + publish_partial_tail: bool = False @staticmethod def from_request_tracker( @@ -773,8 +792,12 @@ class ReqMeta: load_spec: LoadSpec | None = None, skip_save: bool | None = False, block_hashes: list[BlockHash] | None = None, + save_partial_tail: bool = False, + num_prompt_tokens: int | None = None, ) -> "ReqMeta | None": """Create ReqMeta from a RequestTracker.""" + if save_partial_tail: + assert num_prompt_tokens is not None if block_hashes is None: block_hashes = [] input_token_len = tracker.token_len @@ -783,7 +806,15 @@ class ReqMeta: chunk_boundary = cdiv(token_ids_start + 1, block_size) * block_size num_tokens_to_save = input_token_len // block_size * block_size - skip_save = skip_save or num_tokens_to_save < chunk_boundary + publish_tail = ( + save_partial_tail + and not tracker.partial_tail_sent + and num_prompt_tokens is not None + and 0 < num_prompt_tokens <= input_token_len + ) + skip_save = skip_save or ( + num_tokens_to_save < chunk_boundary and not publish_tail + ) # A ReqMeta must never carry both a save AND a load. # The save would also be wasted work — the bytes are being looked up # in the store right now. Later cached_reqs steps save new tokens @@ -795,6 +826,7 @@ class ReqMeta: if not skip_save: tracker.num_saved_tokens = num_tokens_to_save + tracker.partial_tail_sent |= publish_tail token_ids = None if tracker.token_ids and not skip_save: @@ -826,7 +858,10 @@ class ReqMeta: block_hashes=block_hashes, token_ids=token_ids, token_ids_start=token_ids_start, - num_prompt_tokens=tracker.prefill_end_tokens, + num_prompt_tokens=num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + completed_token_len=input_token_len, + publish_partial_tail=publish_tail and not skip_save, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 8f277d7eea..239dad7364 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -5,14 +5,17 @@ # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/). """Scheduler-side logic for MooncakeStoreConnector.""" +from collections.abc import Sequence + from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, ) -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 - partial_hash_hits_enabled, +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( + MooncakeStoreCoordinator, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 + BoundaryPut, LoadSpec, MooncakeStoreConnectorMetadata, MooncakeStoreWorkerMetadata, @@ -23,12 +26,19 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import ( LookupKeyClient, ) from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks -from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes +from vllm.v1.core.kv_cache_utils import ( + resolve_kv_cache_block_sizes, +) from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + MambaSpec, + get_mamba_prefill_checkpoint_position, +) from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request @@ -49,6 +59,79 @@ def _new_req_prefill_tokens(request: NewRequestData) -> list[int]: return request.prompt_token_ids +def _partial_tail_non_mamba_puts( + coord: MooncakeStoreCoordinator, + req_meta: ReqMeta, + block_sizes: Sequence[int], +) -> list[BoundaryPut]: + """Locate the non-Mamba blocks a partial-tail save publishes. + + A later request resumes at the prompt's Mamba checkpoint (``boundary``) only + if both pieces of KV are stored: + + - Mamba groups: the state at ``boundary``, handed off by the core as a + boundary put of its own, so Mamba groups are not listed here. + - Other groups: the KV from the last LCM-aligned normal save up to + ``boundary`` plus the group's EAGLE proof margin (``proof_end``). Each + block is keyed by the hash at its end, clipped to ``proof_end``. + + Groups whose proof is not computed yet are omitted. For example, with LCM + 16, blocks of 4 tokens, and a checkpoint at 44, the full-attention blocks + are 8-10, keyed at 36, 40 and 44. + + A Mamba state handed off at any other position, i.e. a shared-prefix + junction (``--enable-mamba-shared-prefix-checkpoint``), adds no puts here. + + ``block_sizes`` are the per-group block sizes, indexed like + ``req_meta.block_ids``. + """ + mamba_tails = [ + put.num_tokens + for put in req_meta.boundary_puts or [] + if put.group_id in coord.mamba_group_ids + and put.num_tokens % block_sizes[put.group_id] + ] + # The tail is due on the save that first covers the prompt, or when Mamba + # hands off its checkpoint state. + if not mamba_tails and not req_meta.publish_partial_tail: + return [] + prompt_tokens = req_meta.num_prompt_tokens + completed = req_meta.completed_token_len + assert prompt_tokens is not None and completed is not None + hash_block_size = coord.hash_block_size + # Core drops the Mamba checkpoint whenever EAGLE block drop is on, whichever + # groups carry the eagle flag. + boundary = get_mamba_prefill_checkpoint_position( + prompt_tokens, hash_block_size, coord.use_eagle + ) + # Any other sub-block Mamba hand-off is a shared-prefix junction. Its + # attention KV, EAGLE proof included, ends on an attention block end no later + # than this checkpoint's proof, so the normal and prompt-tail saves store it. + if not req_meta.publish_partial_tail and boundary not in mamba_tails: + return [] + num_hashes = len(req_meta.block_hashes) + if boundary == 0 or boundary // hash_block_size > num_hashes: + return [] + start = boundary // coord.lcm_block_size * coord.lcm_block_size + puts: list[BoundaryPut] = [] + for group_id, block_size in enumerate(block_sizes): + if group_id in coord.mamba_group_ids: + continue + proof_end = boundary + coord.eagle_proof_margin_by_group.get(group_id, 0) + if proof_end > completed or proof_end // hash_block_size > num_hashes: + continue + group_blocks = req_meta.block_ids[group_id] + for block_idx in range(start // block_size, cdiv(proof_end, block_size)): + if block_idx >= len(group_blocks): + break + block_id = group_blocks[block_idx] + if block_id == NULL_BLOCK_ID: + continue + num_tokens = min((block_idx + 1) * block_size, proof_end) + puts.append(BoundaryPut(group_id, block_id, num_tokens)) + return puts + + class MooncakeStoreScheduler: """Scheduler-side component for MooncakeStoreConnector.""" @@ -84,11 +167,6 @@ class MooncakeStoreScheduler: self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes( kv_cache_config, vllm_config ) - self.enable_partial_hash_hits = partial_hash_hits_enabled( - store_groups, - self._hash_block_size, - vllm_config.parallel_config.decode_context_parallel_size, - ) mamba_groups = { group_id: group.kv_cache_spec for group_id, group in enumerate(store_groups) @@ -99,6 +177,11 @@ class MooncakeStoreScheduler: ), "MooncakeStoreScheduler requires mamba_cache_mode='align'" self._boundary_state_group_ids = frozenset(mamba_groups) + self._store_coord = MooncakeStoreCoordinator.from_kv_cache_config( + kv_cache_config, vllm_config, self._block_size, self._hash_block_size + ) + self.enable_partial_hash_hits = self._store_coord.enable_partial_hash_hits + self._gpu_block_pool: BlockPool | None = None self._num_workers = vllm_config.parallel_config.world_size self._next_store_job_id = 0 @@ -273,6 +356,8 @@ class MooncakeStoreScheduler: # producer. Loads are still carried by the same metadata. skip_save=is_consumer, block_hashes=request_real.block_hashes, + num_prompt_tokens=request_real.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) if req_meta is not None: meta.add_request(req_meta) @@ -323,6 +408,8 @@ class MooncakeStoreScheduler: load_spec=load_spec, skip_save=is_consumer, block_hashes=request_real.block_hashes, + num_prompt_tokens=request_real.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) else: # Decode/chunked request @@ -364,6 +451,8 @@ class MooncakeStoreScheduler: load_spec=None, skip_save=False, block_hashes=unfinished_req.block_hashes, + num_prompt_tokens=unfinished_req.num_prompt_tokens, + save_partial_tail=self.enable_partial_hash_hits, ) if req_meta is not None: @@ -398,6 +487,7 @@ class MooncakeStoreScheduler: load_spec=load_spec, skip_save=None, block_hashes=unfinished_req.block_hashes, + num_prompt_tokens=unfinished_req.num_prompt_tokens, ) if req_meta is not None: meta.add_request(req_meta) @@ -455,6 +545,8 @@ class MooncakeStoreScheduler: itself is freed, until every rank reports the job done. """ pool = self._gpu_block_pool + coord = self._store_coord + block_sizes = [g.kv_cache_spec.block_size for g in coord.kv_cache_groups] for req_meta in meta.requests: if not req_meta.can_save: continue @@ -464,27 +556,27 @@ class MooncakeStoreScheduler: if req_meta.store_job_id is not None: assert req_meta.store_job_id in self._pinned_saves continue + if tail := _partial_tail_non_mamba_puts(coord, req_meta, block_sizes): + req_meta.boundary_puts = [*tail, *(req_meta.boundary_puts or [])] req_meta.store_job_id = store_job_id = self._next_store_job_id self._next_store_job_id += 1 - block_ids: list[int] = [] - if req_meta.boundary_state_offloads: - block_ids.extend( - block_id for _, block_id, _ in req_meta.boundary_state_offloads - ) + block_ids = [put.block_id for put in req_meta.boundary_puts or []] assert NULL_BLOCK_ID not in block_ids, ( - "A null block cannot back a boundary-state offload" - ) - # Every allocated block is referenced, not just the ones covering - # this job's token range: a rank resumes from its own last - # successful offset, which lags the scheduler's whenever a save was - # skipped or failed, so it may read anywhere below the range. - block_ids.extend( - block_id - for group_id, group in enumerate(req_meta.block_ids) - if group_id not in self._boundary_state_group_ids - for block_id in group - if block_id != NULL_BLOCK_ID + "A null block cannot back a boundary put" ) + if req_meta.token_len_chunk: + # Normal prefix save: pin all non-Mamba sources for retries. + # Every allocated block is referenced, not just the ones covering + # this job's token range: a rank resumes from its own last + # successful offset, which lags the scheduler's whenever a save was + # skipped or failed, so it may read anywhere below the range. + block_ids.extend( + block_id + for group_id, group in enumerate(req_meta.block_ids) + if group_id not in self._boundary_state_group_ids + for block_id in group + if block_id != NULL_BLOCK_ID + ) # An aligned boundary block may also be present in the request's # block table. Take and release exactly one reference per block. block_ids = list(dict.fromkeys(block_ids)) @@ -515,17 +607,39 @@ class MooncakeStoreScheduler: if boundary_tokens > tracker.prefill_end_tokens: return False - pinned_block_ids: list[int] = [] - remapped_offloads: list[tuple[int, int, int]] = [] + mamba_puts: list[BoundaryPut] = [] for group_id, block_id, boundary in partial_tail_offloads: store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id) if store_group_id not in self._boundary_state_group_ids: return False if block_id == NULL_BLOCK_ID: return False - pinned_block_ids.append(block_id) - remapped_offloads.append((store_group_id, block_id, boundary)) - pinned_block_ids = list(dict.fromkeys(pinned_block_ids)) + mamba_puts.append(BoundaryPut(store_group_id, block_id, boundary)) + req_meta = ReqMeta( + req_id=request.request_id, + token_len_chunk=0, + block_ids=tuple( + block_ids[group_id].copy() for group_id in self._store_group_ids + ), + block_hashes=list(request.block_hashes), + can_save=True, + num_prompt_tokens=request.num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + boundary_puts=mamba_puts, + completed_token_len=request.num_computed_tokens, + ) + coord = self._store_coord + req_meta.boundary_puts = [ + *_partial_tail_non_mamba_puts( + coord, + req_meta, + [g.kv_cache_spec.block_size for g in coord.kv_cache_groups], + ), + *mamba_puts, + ] + pinned_block_ids = list( + dict.fromkeys(put.block_id for put in req_meta.boundary_puts) + ) pool = self._gpu_block_pool assert pool is not None, ( @@ -537,18 +651,8 @@ class MooncakeStoreScheduler: self._pinned_saves[store_job_id] = (pinned_block_ids, self._num_workers) pool.touch([pool.blocks[block_id] for block_id in pinned_block_ids]) - self._finished_partial_tail_metas[request.request_id] = ReqMeta( - req_id=request.request_id, - token_len_chunk=0, - block_ids=tuple( - block_ids[group_id].copy() for group_id in self._store_group_ids - ), - block_hashes=list(request.block_hashes), - can_save=True, - num_prompt_tokens=tracker.prefill_end_tokens, - store_job_id=store_job_id, - boundary_state_offloads=remapped_offloads, - ) + req_meta.store_job_id = store_job_id + self._finished_partial_tail_metas[request.request_id] = req_meta tracker.has_pending_offload = True # The store job owns exact block refs, so request cleanup need not wait. return False @@ -575,7 +679,7 @@ class MooncakeStoreScheduler: # going away, so the offload is conservatively dropped. logger.debug("Dropping boundary-state offload for request %s", req_id) continue - accepted: list[tuple[int, int, int]] = [] + accepted: list[BoundaryPut] = [] for group_id, block_id, boundary_tokens in entries: # Every other group stops saving at the end of this prefill, so # a mamba-only key past it can never complete a joint hybrid @@ -589,12 +693,12 @@ class MooncakeStoreScheduler: store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id) if store_group_id not in self._boundary_state_group_ids: continue - accepted.append((store_group_id, block_id, boundary_tokens)) + accepted.append(BoundaryPut(store_group_id, block_id, boundary_tokens)) if not accepted: continue tracker.has_pending_offload = True if (req_meta := save_metas.get(req_id)) is not None: - req_meta.boundary_state_offloads = accepted + req_meta.boundary_puts = accepted continue meta.add_request( ReqMeta( @@ -603,8 +707,10 @@ class MooncakeStoreScheduler: block_ids=tracker.allocated_block_ids, block_hashes=req_tuple[0].block_hashes, can_save=True, - num_prompt_tokens=tracker.prefill_end_tokens, - boundary_state_offloads=accepted, + num_prompt_tokens=req_tuple[0].num_prompt_tokens, + prefill_end_tokens=tracker.prefill_end_tokens, + boundary_puts=accepted, + completed_token_len=tracker.token_len, ) ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 173aa324e7..004b69286e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -74,7 +74,6 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import ( BlockHash, maybe_convert_block_hash, - resolve_dcp_kv_cache_spec, resolve_kv_cache_block_sizes, ) from vllm.v1.kv_cache_interface import ( @@ -569,40 +568,43 @@ class KVCacheStoreSendingThread(KVTransferThread): self._skip_store_requests.clear() return True - def _boundary_snapshot_puts( - self, req_meta: ReqMeta, entries: list[tuple[int, int, int]] + def _boundary_puts( + self, req_meta: ReqMeta ) -> list[tuple[str, list[int], list[int], KeyMetadata]]: - """Puts for committed mamba "align" boundary-state snapshots. + """Puts for ``req_meta.boundary_puts``: each block under the hash of + the prefix ending at its ``num_tokens``. - These are block-aligned boundaries, i.e. exactly what the normal save - would key — but ``store_mask`` masks mamba groups out of it entirely, so - this is their *only* writer. The exclusion is not an optimization: the - normal save resolves a chunk's address as - ``req_meta.block_ids[g][start // block_size]``, and ``block_ids`` is the - connector's append-only mirror of the core's per-group table. An - align-mode table is mutated in place (a superseded state block is freed - and nulled; speculative blocks relocate), and the connector is never - told, so a stale mirror entry is indistinguishable from a live one — a - retry of a failed or pressure-skipped chunk would read a block that now - belongs to another request. - - Each entry's handed-off block *is* the boundary state and is pinned by - the core, so it is uploaded under its boundary-end hash key and never - resolved positionally. + The scheduler resolves and pins every block, so none is looked up + positionally here. That matters for mamba "align" states: ``store_mask`` + masks mamba groups out of the normal save, which resolves a chunk's + address as ``req_meta.block_ids[g][start // block_size]``. An align-mode + table is mutated in place (a superseded state block is freed and + nulled; speculative blocks relocate) without telling the connector, so + a retried positional read could hit a block that now belongs to another + request. The handed-off block *is* the state. """ hash_block_size = self.coord.hash_block_size puts: list[tuple[str, list[int], list[int], KeyMetadata]] = [] - for group_id, block_id, boundary in entries: - if boundary == 0 or block_id == NULL_BLOCK_ID: + for group_id, block_id, num_tokens in req_meta.boundary_puts or []: + if ( + not self.group_participates[group_id] + or num_tokens == 0 + or block_id == NULL_BLOCK_ID + ): continue - hash_idx = boundary // hash_block_size - 1 + # A negative index would silently key the block by the last hash. + assert num_tokens % hash_block_size == 0, ( + f"Boundary put at {num_tokens} tokens is not a multiple of the " + f"hash block size {hash_block_size}" + ) + hash_idx = num_tokens // hash_block_size - 1 if hash_idx >= len(req_meta.block_hashes): continue db = self.token_databases[group_id] # Distribute across ranks by the same rule as normal chunks. put_step = self.group_put_steps[group_id] put_step_rank = (self.tp_rank + group_id) % put_step - if (boundary // db.block_size - 1) % put_step != put_step_rank: + if (cdiv(num_tokens, db.block_size) - 1) % put_step != put_step_rank: continue addr, size = db.prepare_value_for_block(block_id) puts.append( @@ -610,115 +612,20 @@ class KVCacheStoreSendingThread(KVTransferThread): ) return puts - def _sub_block_tail_puts( - self, req_meta: ReqMeta, entries: list[tuple[int, int, int]] - ) -> list[tuple[str, list[int], list[int], KeyMetadata]]: - """Puts for the request's sub-block partial tail (its last prompt hash - boundary), so a later request can hit the sub-block prefix. - - Covers every group's blocks from the normal save's lcm floor to the - boundary: the normal save floors to ``lcm_block_size``, so a - smaller-block group's full blocks in that gap are never persisted - elsewhere, and the consumer's lookup needs every group at every probed - boundary. Full blocks are keyed by their block-end hash and the partial - boundary block by the boundary sub-hash; a mamba "align" group - contributes only its boundary block, from the core-provided CoW block. - """ - boundaries = {boundary for _, _, boundary in entries} - if len(boundaries) != 1: - raise ValueError( - "Sub-block partial-tail offloads for one request must share a boundary" - ) - boundary = boundaries.pop() - hash_block_size = self.coord.hash_block_size - if boundary == 0 or boundary // hash_block_size - 1 >= len( - req_meta.block_hashes - ): - return [] - - mamba_offloads = {group_id: block_id for group_id, block_id, _ in entries} - saved = self._saved_offset.get(req_meta.req_id, 0) - puts: list[tuple[str, list[int], list[int], KeyMetadata]] = [] - for g_idx, db in enumerate(self.token_databases): - if not self.group_participates[g_idx]: - continue - group_blocks = req_meta.block_ids[g_idx] - # Distribute across ranks by the same rule as normal chunks. - put_step = self.group_put_steps[g_idx] - put_step_rank = (self.tp_rank + g_idx) % put_step - # Always include the boundary block: its sub-hash key is written - # only here, even if normal saves already advanced past it. - last_block = cdiv(boundary, db.block_size) - 1 - for block_idx in range( - min(saved // db.block_size, last_block), last_block + 1 - ): - if block_idx % put_step != put_step_rank: - continue - valid_end = min((block_idx + 1) * db.block_size, boundary) - key_hash = req_meta.block_hashes[valid_end // hash_block_size - 1] - if g_idx in mamba_offloads: - if valid_end != boundary: - # Interior align-mode state positions are null or - # stale (the block table is not append-only) and never - # valid gap content; only the boundary block is - # persisted, from the core-provided hand-off. - continue - block_id = mamba_offloads[g_idx] - elif g_idx in self.coord.mamba_group_ids: - continue - elif block_idx < len(group_blocks): - block_id = group_blocks[block_idx] - else: - continue - if block_id == NULL_BLOCK_ID: - logger.debug( - "Skipping unavailable partial-tail source block " - "(req=%s, group=%d, block=%d)", - req_meta.req_id, - g_idx, - block_idx, - ) - continue - addr, size = db.prepare_value_for_block(block_id) - puts.append((db.key_for(key_hash), addr, size, db.metadata)) - return puts - def _maybe_offload_boundary_states(self, req_meta: ReqMeta) -> bool: - """Persist connector-pinned mamba "align" boundary states handed off - for this request, deduped against the store. + """Persist this job's boundary puts, deduped against the store. - This is every mamba key the connector writes — ``store_mask`` excludes - mamba groups from the positional normal save, aligned boundaries - included (see :meth:`_boundary_snapshot_puts`). - - The two entry kinds are keyed and sourced differently, so they are - prepared separately and put in one batch: - - - block-aligned for its group: a committed boundary-state snapshot, - the handed-off block itself; - - not block-aligned: the sub-block CoW partial tail, which also has to - cover the other groups' blocks in the normal save's lcm gap. + These are every mamba key the connector writes (``store_mask`` excludes + mamba groups from the positional normal save) plus the other groups' + partial-tail blocks in the normal save's lcm gap. Returns: True when no put is needed or every put succeeds, False otherwise. """ - offloads = req_meta.boundary_state_offloads - if not offloads or not req_meta.block_hashes: + if not req_meta.boundary_puts or not req_meta.block_hashes: return True - - snapshots: list[tuple[int, int, int]] = [] - sub_block: list[tuple[int, int, int]] = [] - for group_id, block_id, boundary in offloads: - entry = (group_id, block_id, boundary) - if boundary % self.token_databases[group_id].block_size == 0: - snapshots.append(entry) - else: - sub_block.append(entry) - - puts = self._boundary_snapshot_puts(req_meta, snapshots) - if sub_block and self.coord.enable_partial_hash_hits: - puts.extend(self._sub_block_tail_puts(req_meta, sub_block)) + puts = self._boundary_puts(req_meta) if not puts: return True @@ -854,9 +761,7 @@ class KVCacheStoreSendingThread(KVTransferThread): # Offload the handed-off mamba boundary states (independent of the # normal positional save, which may be skipped this step). - if req_meta.boundary_state_offloads is not None and not ( - self._maybe_offload_boundary_states(req_meta) - ): + if not self._maybe_offload_boundary_states(req_meta): return if token_len == 0: @@ -870,7 +775,7 @@ class KVCacheStoreSendingThread(KVTransferThread): store_masks = self.coord.store_mask( token_len, save_start, - num_prompt_tokens=req_meta.num_prompt_tokens, + num_prompt_tokens=req_meta.prefill_end_tokens, ) starts: list[int] = [] @@ -1603,31 +1508,10 @@ class MooncakeStoreWorker: ) return - self._kv_cache_groups = [ - dataclasses.replace( - group, - kv_cache_spec=resolve_dcp_kv_cache_spec( - group.kv_cache_spec, - self.dcp_size, - ), - ) - for group in kv_cache_config.prefix_cacheable_groups - ] - spec_cfg = getattr(vllm_config, "speculative_config", None) - use_eagle_block_drop = bool( - spec_cfg.use_eagle_block_drop() - if spec_cfg is not None - and callable(getattr(spec_cfg, "use_eagle_block_drop", None)) - else False - ) - self.coord = MooncakeStoreCoordinator( - self._kv_cache_groups, - scheduler_block_size=self.block_size, - hash_block_size=self.hash_block_size, - use_eagle=use_eagle_block_drop, - retention_interval=kv_cache_config.prefix_cache_retention_interval, - dcp_world_size=self.dcp_size, + self.coord = MooncakeStoreCoordinator.from_kv_cache_config( + kv_cache_config, vllm_config, self.block_size, self.hash_block_size ) + self._kv_cache_groups = self.coord.kv_cache_groups self.store_tp_size, store_namespace, store_layout_cls = ( self._select_store_layout(extra_config) ) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index d0dadae27e..b13764abec 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -362,7 +362,10 @@ class EngineCore: scheduler_block_size, hash_block_size, ) - for kv_cache_config in kv_cache_configs: + # The scheduler's config is a deep copy taken before these are resolved, + # so set them there too: the KV connector's scheduler side reads the same + # fields as its workers. + for kv_cache_config in (*kv_cache_configs, scheduler_kv_cache_config): kv_cache_config.hash_block_size = hash_block_size kv_cache_config.cache_hit_alignment_tokens = cache_hit_alignment_tokens