iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][Mooncake] Save the EAGLE attention proof with hybrid Mamba partial tails (#60695)

Browse Source 
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Signed-off-by: Dao007forever <daole@inferact.ai>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
This commit is contained in:
Dao007forever
2026-10-10 00:08:08 +00:00
committed by GitHub
co-authored by Wei Zhao Claude Opus 5.5 zjy0516
parent 03a7d3776f
commit f8e6e365ae
11 changed files with 1142 additions and 390 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/models_basic.yaml
+2
View File
@@ -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/
tests/models/kimi_k3/test_prefix_cache.py
+114 -25
View File
@@ -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
)
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+9 -35
View File
@@ -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."""
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+66 -29
View File
@@ -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
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+281 -35
View File
@@ -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],),
)
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+364 -28
View File
@@ -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)]
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+69 -29
View File
@@ -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
)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+43 -8
View File
@@ -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,
)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+154 -48
View File
@@ -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,
)
)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+36 -152
View File
@@ -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)
)
vllm/v1/engine/core.py
+4 -1
View File
@@ -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