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

PD disagg with Mooncake Connector: GDN support (Qwen3.5) and MLA support (Deepseek-V4-Flash) (#46807)

Browse Source
This commit is contained in:
Dakai An
2026-06-29 23:29:04 -07:00
committed by GitHub
parent 97b5ce5c39
commit 0feca7ffa8
7 changed files with 735 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_mooncake_connector.py
+36 -4
View File
@@ -33,14 +33,37 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import
)
from vllm.utils.network_utils import get_open_port
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
)
from vllm.v1.request import RequestStatus
from .utils import create_request, create_scheduler, create_vllm_config
def _make_test_kv_cache_config() -> KVCacheConfig:
return KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[])
return KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
[
"model.layers.0.self_attn",
"model.layers.1.self_attn",
"model.layers.0.mla_attn",
"model.layers.1.eagle_attn",
],
FullAttentionSpec(
block_size=16,
num_kv_heads=4,
head_size=64,
dtype=torch.float16,
),
)
],
)
class FakeMooncakeWrapper:
@@ -126,6 +149,8 @@ async def test_build_transfer_params_separates_prefill_pp_layers():
worker.is_kv_producer = True
worker.tp_rank = 0
worker.tp_size = 1
worker.kv_cache_config = _make_test_kv_cache_config()
worker._physical_blocks_per_logical_kv_block = 1
worker.transfer_topo = SimpleNamespace(local_replicates_kv_cache=False)
block_len = 256
@@ -206,6 +231,7 @@ async def test_build_transfer_params_separates_prefill_pp_layers():
req_blocks={"d-req-pp": (transfer_id, [[20, 21]])},
kv_caches_base_addr=[region.base_addr for region in remote_regions],
block_lens=[region.block_len for region in remote_regions],
kv_block_lens=[region.kv_block_len for region in remote_regions],
registered_layer_names=[region.layer_name for region in remote_regions],
registered_layer_indices=[region.layer_index for region in remote_regions],
)
@@ -266,6 +292,7 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
kv_half = block_len // 2
prefill_worker.kv_caches_base_addr = [0x1000]
prefill_worker.block_len_per_layer = [block_len]
prefill_worker.kv_block_len_per_layer = [kv_half]
prefill_worker.registered_layer_names = ["model.layers.1.self_attn"]
prefill_worker.registered_layer_indices = [1]
@@ -294,6 +321,7 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
req_blocks={"d-req-layer-align": (transfer_id, [[20]])},
kv_caches_base_addr=[0xA000, 0xB000],
block_lens=[block_len, block_len],
kv_block_lens=[kv_half, kv_half],
registered_layer_names=[
"model.layers.0.self_attn",
"model.layers.1.self_attn",
@@ -804,7 +832,9 @@ async def test_kv_producer(monkeypatch):
prefill_worker = prefill_connector.connector_worker
prefill_worker.kv_caches_base_addr = [0x1000]
block_len = 4096
kv_half = block_len // 2
prefill_worker.block_len_per_layer = [block_len]
prefill_worker.kv_block_len_per_layer = [kv_half]
prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
prefill_worker.registered_layer_indices = [0]
@@ -832,6 +862,7 @@ async def test_kv_producer(monkeypatch):
req_blocks={"d-req-1": (transfer_id, [[20, 21]])},
kv_caches_base_addr=[0x2000],
block_lens=[block_len],
kv_block_lens=[kv_half],
registered_layer_names=["model.layers.0.self_attn"],
registered_layer_indices=[0],
)
@@ -845,8 +876,6 @@ async def test_kv_producer(monkeypatch):
) as mock_send_blocks:
# With blocks-first layout, each block is virtually split
# into K and V halves, producing non-coalesced transfers.
kv_half = block_len // 2
def expected_split_transfers(src_base, dst_base, src_blocks, dst_blocks):
"""Build expected (src_ptrs, dst_ptrs, lengths) for
virtual-split K/V transfers."""
@@ -981,6 +1010,7 @@ async def test_kv_consumuer(monkeypatch):
decode_worker = decode_connector.connector_worker
decode_worker.kv_caches_base_addr = [0x1000]
decode_worker.block_len_per_layer = [4096]
decode_worker.kv_block_len_per_layer = [4096]
decode_worker.registered_layer_names = ["model.layers.0.self_attn"]
decode_worker.registered_layer_indices = [0]
decode_worker.rpc_port = 54321
@@ -1236,6 +1266,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
prefill_worker.kv_caches_base_addr = [0x1000]
prefill_worker.block_len_per_layer = [local_block_len]
prefill_worker.kv_block_len_per_layer = [local_block_len // 2]
prefill_worker.registered_layer_names = ["model.layers.0.self_attn"]
prefill_worker.registered_layer_indices = [0]
@@ -1283,6 +1314,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
},
kv_caches_base_addr=[0x2000],
block_lens=[remote_block_len],
kv_block_lens=[remote_block_len // 2],
registered_layer_names=["model.layers.0.self_attn"],
registered_layer_indices=[0],
)
tests/v1/kv_connector/unit/test_mooncake_connector_hma.py
+2
View File
@@ -257,6 +257,7 @@ async def test_build_transfer_params_multi_group_trimming(monkeypatch):
},
kv_caches_base_addr=[0x2000],
block_lens=[block_len],
kv_block_lens=[block_len],
)
local_regions = [
@@ -348,6 +349,7 @@ async def test_build_transfer_params_group_count_mismatch(monkeypatch):
},
kv_caches_base_addr=[0x2000],
block_lens=[block_len],
kv_block_lens=[block_len],
)
local_regions = [
tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
+386
View File
@@ -0,0 +1,386 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for MooncakeConnector hybrid FA + GDN support.
GDN is represented as a MambaSpec in vLLM, so these tests exercise the
Mooncake MambaSpec path with mamba_type=GDN_ATTN. Mamba2 is intentionally not
validated by this test module.
"""
import asyncio
from types import SimpleNamespace
from unittest.mock import patch
import pytest
import torch
from vllm.config import set_current_vllm_config
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
KVConnectorRole,
MooncakeConnector,
MooncakeConnectorScheduler,
MooncakeConnectorWorker,
MooncakeXferMetadata,
SendBlockMeta,
TransferRegion,
)
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
MambaSpec,
)
from .test_mooncake_connector import patch_worker_dependencies
from .utils import create_request, create_vllm_config
def noop_shutdown():
pass
def make_hybrid_gdn_kv_cache_config(block_size: int) -> KVCacheConfig:
gdn_spec = MambaSpec(
block_size=block_size,
shapes=((6, 3), (1, 2, 2)),
dtypes=(torch.float16, torch.float16),
mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
)
assert gdn_spec.mamba_type == MambaAttentionBackendEnum.GDN_ATTN
return KVCacheConfig(
num_blocks=16,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(
["model.layers.0.self_attn"],
FullAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float16,
),
),
KVCacheGroupSpec(
["model.layers.1.linear_attn"],
gdn_spec,
),
],
)
def make_hybrid_gdn_scheduler(kv_role: str) -> MooncakeConnectorScheduler:
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector",
kv_role=kv_role,
)
vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
return MooncakeConnectorScheduler(
vllm_config=vllm_config,
engine_id="test-engine",
kv_cache_config=make_hybrid_gdn_kv_cache_config(
vllm_config.cache_config.block_size
),
)
@pytest.mark.cpu_test
def test_hybrid_gdn_remote_prefill_uses_mamba_n_minus_one():
scheduler = make_hybrid_gdn_scheduler(kv_role="kv_consumer")
request = create_request(num_tokens=10, do_remote_prefill=True)
num_new_tokens, is_async = scheduler.get_num_new_matched_tokens(
request, num_computed_tokens=0
)
assert num_new_tokens == request.num_prompt_tokens - 1
assert is_async is True
@pytest.mark.cpu_test
def test_hybrid_gdn_remote_decode_truncates_prefill_once():
scheduler = make_hybrid_gdn_scheduler(kv_role="kv_producer")
request = create_request(num_tokens=10, do_remote_decode=True)
original_tokens = list(request.prompt_token_ids)
num_new_tokens, is_async = scheduler.get_num_new_matched_tokens(
request, num_computed_tokens=0
)
assert num_new_tokens == 0
assert is_async is False
assert request.prompt_token_ids == original_tokens[:-1]
assert request._all_token_ids == original_tokens[:-1]
assert request.num_prompt_tokens == len(original_tokens) - 1
assert request.max_tokens == 1
assert request.kv_transfer_params["_p_side_truncated"] is True
scheduler.get_num_new_matched_tokens(request, num_computed_tokens=0)
assert request.prompt_token_ids == original_tokens[:-1]
def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector",
kv_role="kv_consumer",
)
kv_cache_config = make_hybrid_gdn_kv_cache_config(
vllm_config.cache_config.block_size
)
with set_current_vllm_config(vllm_config), patch_worker_dependencies():
connector = MooncakeConnector(
vllm_config,
KVConnectorRole.WORKER,
kv_cache_config,
)
worker = connector.connector_worker
fa_cache = torch.empty((2, 2, 11), dtype=torch.float16)
gdn_conv_state = torch.empty((2, 22), dtype=torch.float16)
gdn_ssm_state = torch.empty((2, 4), dtype=torch.float16)
worker.register_kv_caches(
{
"model.layers.0.self_attn": fa_cache,
"model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
}
)
assert worker.transfer_topo.is_mamba is True
assert worker.registered_layer_names == [
"model.layers.0.self_attn",
"model.layers.1.linear_attn",
]
assert worker.registered_group_indices == [0, 1]
assert worker.kv_caches_base_addr == [
fa_cache.data_ptr(),
gdn_conv_state.data_ptr(),
]
worker.shutdown()
worker.shutdown = noop_shutdown
connector.connector_worker = None
def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector",
kv_role="kv_consumer",
)
kv_cache_config = make_hybrid_gdn_kv_cache_config(
vllm_config.cache_config.block_size
)
with set_current_vllm_config(vllm_config), patch_worker_dependencies():
connector = MooncakeConnector(
vllm_config,
KVConnectorRole.WORKER,
kv_cache_config,
)
worker = connector.connector_worker
backing = torch.empty((4, 64), dtype=torch.float16)
fa_cache = backing[:2, :16]
gdn_conv_state = backing[:3]
gdn_ssm_state = torch.empty((3, 4), dtype=torch.float16)
with patch.object(
worker.engine, "batch_register_memory", return_value=0
) as batch_register_memory:
worker.register_kv_caches(
{
"model.layers.0.self_attn": fa_cache,
"model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
}
)
assert worker.kv_caches_base_addr == [
fa_cache.data_ptr(),
gdn_conv_state.data_ptr(),
]
batch_register_memory.assert_called_once()
registered_ptrs, registered_lens = batch_register_memory.call_args[0]
assert registered_ptrs == [backing.data_ptr()]
assert registered_lens == [backing.untyped_storage().nbytes()]
worker.shutdown()
worker.shutdown = noop_shutdown
connector.connector_worker = None
def test_hybrid_gdn_transfer_params_preserve_group_identity(monkeypatch):
monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector",
kv_role="kv_producer",
)
kv_cache_config = make_hybrid_gdn_kv_cache_config(
vllm_config.cache_config.block_size
)
with set_current_vllm_config(vllm_config), patch_worker_dependencies():
connector = MooncakeConnector(
vllm_config,
KVConnectorRole.WORKER,
kv_cache_config,
)
worker = connector.connector_worker
block_len = 0x100
transfer_id = "xfer-hybrid-gdn"
async def build_transfer_params():
send_meta = SendBlockMeta(
p_req_id="p-hybrid-gdn",
transfer_id=transfer_id,
local_block_ids=[
[10, 11],
[NULL_BLOCK_ID, 4],
],
ready=asyncio.Event(),
)
return await worker._build_transfer_params(
[("d-hybrid-gdn", send_meta)],
xfer_meta,
local_regions,
remote_regions,
)
xfer_meta = MooncakeXferMetadata(
remote_hostname="consumer-host",
remote_port=54321,
remote_tp_size=1,
remote_tp_rank=0,
req_blocks={
"d-hybrid-gdn": (
transfer_id,
[
[30, 31],
[NULL_BLOCK_ID, 7],
],
)
},
kv_caches_base_addr=[],
block_lens=[],
kv_block_lens=[],
)
local_regions = [
TransferRegion(
layer_name="model.layers.1.linear_attn",
layer_index=1,
base_addr=0x5000,
block_len=block_len,
kv_block_len=block_len,
group_index=1,
),
TransferRegion(
layer_name="model.layers.0.self_attn",
layer_index=0,
base_addr=0x1000,
block_len=block_len,
kv_block_len=block_len,
group_index=0,
),
]
remote_regions = [
TransferRegion(
layer_name="model.layers.1.linear_attn",
layer_index=1,
base_addr=0x6000,
block_len=block_len,
kv_block_len=block_len,
group_index=1,
),
TransferRegion(
layer_name="model.layers.0.self_attn",
layer_index=0,
base_addr=0x2000,
block_len=block_len,
kv_block_len=block_len,
group_index=0,
),
]
src_ptrs, dst_ptrs, lengths, err_reqs, err_msg = asyncio.run(
build_transfer_params()
)
assert err_reqs == []
assert err_msg is None
assert src_ptrs == [
0x5000 + 4 * block_len,
0x1000 + 10 * block_len,
]
assert dst_ptrs == [
0x6000 + 7 * block_len,
0x2000 + 30 * block_len,
]
assert lengths == [block_len, 2 * block_len]
worker.shutdown()
worker.shutdown = noop_shutdown
connector.connector_worker = None
def test_logical_to_kernel_block_ids_expands_fa_not_gdn():
worker = object.__new__(MooncakeConnectorWorker)
worker.shutdown = noop_shutdown
worker._physical_blocks_per_logical_kv_block = 17
worker.kv_cache_config = make_hybrid_gdn_kv_cache_config(block_size=544)
block_ids = [[2], [2]]
kernel_block_ids = worker._logical_to_kernel_block_ids(block_ids)
assert kernel_block_ids == [list(range(34, 51)), [2]]
def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
monkeypatch,
):
monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
vllm_config = create_vllm_config(
kv_connector="MooncakeConnector",
kv_role="kv_producer",
)
kv_cache_config = make_hybrid_gdn_kv_cache_config(
vllm_config.cache_config.block_size
)
with set_current_vllm_config(vllm_config), patch_worker_dependencies():
connector = MooncakeConnector(
vllm_config,
KVConnectorRole.WORKER,
kv_cache_config,
)
worker = connector.connector_worker
worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=True)
regions = worker._get_transfer_regions(
base_addrs=[0x1000, 0x2000],
block_lens=[0x100, 0x100],
kv_block_lens=[0x40, 0x100],
layer_names=[
"model.layers.0.self_attn",
"model.layers.1.linear_attn",
],
layer_indices=[0, 1],
group_indices=[0, 1],
)
assert [
(region.group_index, region.base_addr, region.kv_block_len)
for region in regions
] == [
(0, 0x1000, 0x40),
(0, 0x1040, 0x40),
(1, 0x2000, 0x100),
]
worker.shutdown()
worker.shutdown = noop_shutdown
connector.connector_worker = None
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+308 -96
View File
@@ -22,7 +22,6 @@ from vllm.config import VllmConfig
from vllm.distributed.kv_transfer.kv_connector.utils import (
EngineId,
TransferTopology,
get_current_attn_backend,
get_current_attn_backends,
)
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -51,10 +50,18 @@ from vllm.platforms import current_platform
from vllm.utils.math_utils import cdiv
from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
from vllm.v1.attention.backend import AttentionMetadata
from vllm.v1.attention.backends.utils import get_kv_cache_layout
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, get_kv_cache_layout
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.kv_cache_interface import FullAttentionSpec, SlidingWindowSpec
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheSpec,
MambaSpec,
MLAAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
)
from vllm.v1.request import RequestStatus
from vllm.v1.worker.block_table import BlockTable
from vllm.v1.worker.utils import select_common_block_size
logger = init_logger(__name__)
@@ -85,6 +92,7 @@ class TransferRegion:
base_addr: int
block_len: int
kv_block_len: int
group_index: int = 0
def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int:
@@ -111,24 +119,58 @@ def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int:
def _expand_transfer_regions(
base_addrs: list[int],
block_lens: list[int],
kv_block_lens: list[int],
layer_names: list[str],
layer_indices: list[int],
is_kv_layout_blocks_first: bool,
group_indices: list[int] | None = None,
split_kv_regions: list[bool] | None = None,
) -> list[TransferRegion]:
"""Expand registered KV tensors into the regions transferred by Mooncake."""
assert (
len(base_addrs) == len(block_lens) == len(layer_names) == len(layer_indices)
len(base_addrs)
== len(block_lens)
== len(kv_block_lens)
== len(layer_names)
== len(layer_indices)
), (
"Mooncake transfer regions require matching metadata lengths, got "
f"base_addrs={len(base_addrs)}, block_lens={len(block_lens)}, "
f"kv_block_lens={len(kv_block_lens)}, "
f"layer_names={len(layer_names)}, "
f"layer_indices={len(layer_indices)}."
)
if group_indices is None:
group_indices = [0] * len(layer_names)
assert len(group_indices) == len(layer_names), (
"Mooncake transfer regions require matching group metadata lengths, "
f"got group_indices={len(group_indices)}, layer_names={len(layer_names)}."
)
if split_kv_regions is None:
split_kv_regions = [is_kv_layout_blocks_first] * len(layer_names)
assert len(split_kv_regions) == len(layer_names), (
"Mooncake transfer regions require matching split metadata, "
f"got split_kv_regions={len(split_kv_regions)}, "
f"layer_names={len(layer_names)}."
)
regions: list[TransferRegion] = []
for base_addr, block_len, layer_name, layer_index in zip(
base_addrs, block_lens, layer_names, layer_indices
for (
base_addr,
block_len,
kv_block_len,
layer_name,
layer_index,
group_index,
split_kv_region,
) in zip(
base_addrs,
block_lens,
kv_block_lens,
layer_names,
layer_indices,
group_indices,
split_kv_regions,
):
kv_block_len = block_len // 2 if is_kv_layout_blocks_first else block_len
regions.append(
TransferRegion(
layer_name=layer_name,
@@ -136,9 +178,10 @@ def _expand_transfer_regions(
base_addr=base_addr,
block_len=block_len,
kv_block_len=kv_block_len,
group_index=group_index,
)
)
if is_kv_layout_blocks_first:
if split_kv_region:
regions.append(
TransferRegion(
layer_name=layer_name,
@@ -146,6 +189,7 @@ def _expand_transfer_regions(
base_addr=base_addr + kv_block_len,
block_len=block_len,
kv_block_len=kv_block_len,
group_index=group_index,
)
)
return regions
@@ -308,6 +352,17 @@ def _align_transfer_regions(
f"{remote_region.layer_index}."
),
)
if local_region.group_index != remote_region.group_index:
return (
[],
[],
(
"Mooncake registered group index mismatch for "
f"{local_region.layer_name}: producer="
f"{local_region.group_index}, consumer="
f"{remote_region.group_index}."
),
)
aligned_local.append(local_region)
aligned_remote.append(remote_region)
@@ -332,8 +387,10 @@ class MooncakeXferMetadata(
req_blocks: dict[ReqId, tuple[TransferId, list[list[int]]]]
kv_caches_base_addr: list[int]
block_lens: list[int]
kv_block_lens: list[int]
registered_layer_names: list[str] = msgspec.field(default_factory=list)
registered_layer_indices: list[int] = msgspec.field(default_factory=list)
registered_group_indices: list[int] = msgspec.field(default_factory=list)
class MooncakeXferResponseStatus(IntEnum):
@@ -581,6 +638,9 @@ class MooncakeConnectorScheduler:
for g in kv_cache_config.kv_cache_groups
)
)
# GDN is represented as a MambaSpec in vLLM. This Mooncake MambaSpec
# path is currently tested with GDN; Mamba2 is not validated yet.
self._has_mamba = kv_cache_config.has_mamba_layers
# Requests that need to start recv/send.
# New requests are added by update_state_after_alloc in
@@ -617,6 +677,38 @@ class MooncakeConnectorScheduler:
for i, blocks in enumerate(block_ids)
]
def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
"""D-side only. Returns N-1 for Mamba models since the decoder
always recomputes the last token and must start from h(N-1)."""
if self._has_mamba and num_prompt_tokens > 1:
return num_prompt_tokens - 1
return num_prompt_tokens
def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
"""P-side only: drop the last prompt token so the prefiller computes
h(N-1) instead of h(N). The decoder recomputes the last token to
derive h(N) correctly.
Guarded by ``_p_side_truncated`` to avoid repeated truncation if the
request is preempted and rescheduled."""
params = request.kv_transfer_params
if (
params is not None
and not params.get("_p_side_truncated")
and request.num_prompt_tokens > 1
):
if request.prompt_token_ids is not None:
request.prompt_token_ids.pop()
elif request.prompt_embeds is not None:
request.prompt_embeds = request.prompt_embeds[:-1]
else:
return
request._all_token_ids.pop()
request.num_prompt_tokens -= 1
request.max_tokens = 1
params["_p_side_truncated"] = True
def get_num_new_matched_tokens(
self, request: "Request", num_computed_tokens: int
) -> tuple[int, bool]:
@@ -650,10 +742,15 @@ class MooncakeConnectorScheduler:
# Remote prefill: get all prompt blocks from remote.
assert not self.is_kv_producer
token_ids = request.prompt_token_ids or []
count = len(token_ids) - num_computed_tokens
count = self._get_remote_prefill_token_count(len(token_ids)) - (
num_computed_tokens
)
if count > 0:
return count, True
if params.get("do_remote_decode") and self._has_mamba:
self._truncate_mamba_request_for_prefill(request)
# No remote prefill for this request.
return 0, False
@@ -802,7 +899,7 @@ class MooncakeConnectorWorker:
self,
vllm_config: VllmConfig,
engine_id: str,
kv_cache_config: "KVCacheConfig | None" = None,
kv_cache_config: "KVCacheConfig",
):
if TransferEngine is None:
logger.error("Mooncake is not available")
@@ -831,10 +928,15 @@ class MooncakeConnectorWorker:
protocol = kv_transfer_config.kv_connector_extra_config.get( # type: ignore[union-attr]
"mooncake_protocol", "rdma"
)
device_name = kv_transfer_config.kv_connector_extra_config.get( # type: ignore[union-attr]
"device_name", ""
)
logger.info(
"The Mooncake Transfer Engine is using %s as its protocol.", protocol
)
ret_value = self.engine.initialize(self.hostname, "P2PHANDSHAKE", protocol, "")
ret_value = self.engine.initialize(
self.hostname, "P2PHANDSHAKE", protocol, device_name
)
if ret_value != 0:
raise RuntimeError("Mooncake Transfer Engine initialization failed.")
@@ -852,10 +954,11 @@ class MooncakeConnectorWorker:
self.engine_id: EngineId = engine_id
self.tp_rank = get_tensor_model_parallel_rank()
self.tp_size = get_tensor_model_parallel_world_size()
self.num_blocks = 0
self.block_len_per_layer: list[int] = []
self.kv_block_len_per_layer: list[int] = []
self.registered_layer_names: list[str] = []
self.registered_layer_indices: list[int] = []
self.registered_group_indices: list[int] = []
self.seen_base_addresses: list[int] = []
assert (parallel_config := vllm_config.parallel_config)
@@ -916,26 +1019,40 @@ class MooncakeConnectorWorker:
self.cache_config = vllm_config.cache_config
self.kv_cache_config = kv_cache_config
self.use_mla = self.model_config.use_mla
self._physical_blocks_per_logical_kv_block = 1
self._sync_block_size_with_kernel()
# Get the attention backend from the first layer
# NOTE (NickLucche) models with multiple backends are not supported yet
backend = get_current_attn_backend(vllm_config)
self.backend_name = backend.get_name()
self.attn_backends = get_current_attn_backends(vllm_config)
self.kv_cache_layout = get_kv_cache_layout()
logger.debug("Detected attention backend %s", self.backend_name)
logger.debug(
"Detected attention backends %s",
[backend.get_name() for backend in self.attn_backends],
)
logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
self._layer_specs: dict[str, KVCacheSpec] = {}
for group in kv_cache_config.kv_cache_groups:
group_spec = group.kv_cache_spec
specs_by_layer = getattr(group_spec, "kv_cache_specs", {})
for layer_name in group.layer_names:
self._layer_specs[layer_name] = specs_by_layer.get(
layer_name, group_spec
)
self._layer_group_indices: dict[str, int] = {
layer: group_index
for group_index, group in enumerate(kv_cache_config.kv_cache_groups)
for layer in group.layer_names
}
self.transfer_topo = TransferTopology(
tp_rank=self.tp_rank,
tp_size=self.tp_size,
block_size=self.block_size,
engine_id=self.engine_id,
is_mla=self.use_mla,
is_mamba=False,
is_mamba=kv_cache_config.has_mamba_layers,
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
attn_backends=[backend],
attn_backends=self.attn_backends,
)
self.async_zmq_ctx = zmq.asyncio.Context()
@@ -958,6 +1075,9 @@ class MooncakeConnectorWorker:
kernel_block_size,
)
assert self.block_size > kernel_block_size
self._physical_blocks_per_logical_kv_block = (
self.block_size // kernel_block_size
)
self.block_size = kernel_block_size
def __del__(self):
@@ -1092,14 +1212,18 @@ class MooncakeConnectorWorker:
local_regions = self._get_transfer_regions(
self.kv_caches_base_addr,
self.block_len_per_layer,
self.kv_block_len_per_layer,
self.registered_layer_names,
self.registered_layer_indices,
self.registered_group_indices,
)
remote_regions = self._get_transfer_regions(
meta.kv_caches_base_addr,
meta.block_lens,
meta.kv_block_lens,
meta.registered_layer_names,
meta.registered_layer_indices,
meta.registered_group_indices,
)
local_regions, remote_regions, align_err = _align_transfer_regions(
local_regions, remote_regions
@@ -1271,6 +1395,32 @@ class MooncakeConnectorWorker:
remote_tp_ranks,
)
def _logical_to_kernel_block_ids(
self, block_ids: list[list[int]]
) -> list[list[int]]:
# For example, if a 544-token logical block is served by 32-token
# FA kernel blocks, FA block id k expands to [17k, ..., 17k + 16],
# while the matching Mamba/GDN state block remains k. Only attention
# groups need logical block ids expanded to kernel block ids; Mamba/GDN
# state block ids stay in the logical/page-id space.
if self._physical_blocks_per_logical_kv_block == 1:
return block_ids
block_arange = np.arange(self._physical_blocks_per_logical_kv_block).reshape(
1, -1
)
group_specs = self.kv_cache_config.kv_cache_groups
return [
BlockTable.map_to_kernel_blocks(
np.array(group),
self._physical_blocks_per_logical_kv_block,
block_arange,
).tolist()
if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
else group
for i, group in enumerate(block_ids)
]
async def _build_transfer_params(
self,
ready_reqs: list[tuple[ReqId, SendBlockMeta]],
@@ -1293,14 +1443,6 @@ class MooncakeConnectorWorker:
):
continue
# Per-group partial hit trimming, then flatten.
# With HMA, groups share the same KV tensor but use different
# block ranges. We trim and concatenate so the coalescer and
# address math see one flat block list — same as non-HMA, but
# now including blocks from every group.
local_block_ids: list[int] = []
remote_block_ids: list[int] = []
has_block_error = False
if len(send_meta.local_block_ids) != len(remote_block_ids_per_group):
logger.error(
"req %s: KV group count mismatch: local=%d, remote=%d",
@@ -1312,26 +1454,55 @@ class MooncakeConnectorWorker:
if err_msg is None:
err_msg = "KV group count mismatch"
continue
for local_group, remote_group in zip(
send_meta.local_block_ids, remote_block_ids_per_group
# Keep KV-cache group identity. Hybrid/HMA groups can carry
# different semantics (e.g. full-attention KV pages vs GDN/Mamba
# inner-state slots), so their block IDs must not be flattened and
# reused for every registered region.
local_block_ids_by_group: list[list[int]] = []
remote_block_ids_by_group: list[list[int]] = []
has_block_error = False
group_specs = self.kv_cache_config.kv_cache_groups
for group_index, (local_group, remote_group) in enumerate(
zip(send_meta.local_block_ids, remote_block_ids_per_group)
):
is_mamba_group = isinstance(
group_specs[group_index].kv_cache_spec,
MambaSpec,
)
if is_mamba_group:
# Mamba/GDN prefix caching can use null blocks only as
# align-mode placeholders. They do not carry transferable
# state, so skip them on both producer and consumer sides.
local_group = [
block_id
for block_id in local_group
if block_id != NULL_BLOCK_ID
]
remote_group = [
block_id
for block_id in remote_group
if block_id != NULL_BLOCK_ID
]
n_local = len(local_group)
n_remote = len(remote_group)
if n_local < n_remote:
logger.error(
"req %s: local blocks(%d) < remote blocks(%d) "
"in a KV cache group",
"in a KV cache group (is_mamba_group=%s)",
d_req_id,
n_local,
n_remote,
is_mamba_group,
)
has_block_error = True
break
if n_local > n_remote:
elif n_local > n_remote:
# Partial prefix cache hit: just read uncomputed blocks.
local_group = local_group[-n_remote:]
local_block_ids.extend(local_group)
remote_block_ids.extend(remote_group)
local_group = local_group[-n_remote:] if n_remote > 0 else []
local_block_ids_by_group.append(local_group)
remote_block_ids_by_group.append(remote_group)
if has_block_error:
err_reqs.append(d_req_id)
@@ -1339,22 +1510,44 @@ class MooncakeConnectorWorker:
err_msg = "P num blocks less than D"
continue
if not local_block_ids:
if not any(local_block_ids_by_group):
continue
# Group by indices
group_local_block_ids, group_remote_block_ids = group_concurrent_contiguous(
local_block_ids, remote_block_ids
local_block_ids_by_group = self._logical_to_kernel_block_ids(
local_block_ids_by_group
)
remote_block_ids_by_group = self._logical_to_kernel_block_ids(
remote_block_ids_by_group
)
for local_region, remote_region in zip(local_regions, remote_regions):
should_transfer, src_region_offset, dst_region_offset, transfer_len = (
self._get_sender_transfer_plan(
local_kv_block_len=local_region.kv_block_len,
remote_kv_block_len=remote_region.kv_block_len,
remote_tp_rank=agent_meta.remote_tp_rank,
remote_tp_size=agent_meta.remote_tp_size,
)
assert local_region.group_index == remote_region.group_index, (
"Aligned Mooncake transfer regions must belong to the same "
"KV group."
)
group_index = local_region.group_index
assert group_index < len(local_block_ids_by_group), (
"Transfer region references a missing KV group."
)
local_block_ids = local_block_ids_by_group[group_index]
remote_block_ids = remote_block_ids_by_group[group_index]
if not local_block_ids:
continue
# Group by indices within this region's KV-cache group only.
group_local_block_ids, group_remote_block_ids = (
group_concurrent_contiguous(local_block_ids, remote_block_ids)
)
(
should_transfer,
src_region_offset,
dst_region_offset,
transfer_len,
) = self._get_sender_transfer_plan(
local_kv_block_len=local_region.kv_block_len,
remote_kv_block_len=remote_region.kv_block_len,
remote_tp_rank=agent_meta.remote_tp_rank,
remote_tp_size=agent_meta.remote_tp_size,
)
if not should_transfer:
# Replicated KV cache: only one producer rank in the TP group
@@ -1368,7 +1561,7 @@ class MooncakeConnectorWorker:
"Computed source transfer region exceeds local KV block size."
)
assert dst_region_offset + transfer_len <= remote_region.kv_block_len, (
"Computed destination transfer region exceeds remote KV block size."
"Destination transfer region exceeds remote KV block size."
)
# Collapse one contiguous block group into a single larger
# transfer descriptor when the per-block copy is identical.
@@ -1411,28 +1604,10 @@ class MooncakeConnectorWorker:
)
lengths.append(transfer_len)
if local_region is local_regions[0]:
logger.debug(
"Mooncake transfer plan for request %s: local_tp=%d "
"remote_tp=%d remote_tp_rank=%d local_block_len=%d "
"remote_block_len=%d src_offset=%d dst_offset=%d "
"transfer_len=%d coalesce=%s",
d_req_id,
self.tp_size,
agent_meta.remote_tp_size,
agent_meta.remote_tp_rank,
local_region.block_len,
remote_region.block_len,
src_region_offset,
dst_region_offset,
transfer_len,
can_coalesce,
)
logger.debug(
"Sending kv_caches for request %s (%d blocks) to %s",
d_req_id,
len(local_block_ids),
sum(len(group) for group in local_block_ids_by_group),
remote_session,
)
@@ -1480,18 +1655,33 @@ class MooncakeConnectorWorker:
logger.info("Registering KV_Caches. use_mla: %s", self.use_mla)
kv_data_ptrs = []
kv_data_lens = []
seen_base_addresses = []
kv_data_ptrs: list[int] = []
kv_data_lens: list[int] = []
region_base_addresses: list[int] = []
seen_storage_ptrs: set[int] = set()
self.block_len_per_layer = []
self.kv_block_len_per_layer = []
self.registered_layer_names = []
self.registered_layer_indices = []
self.registered_group_indices = []
split_k_and_v = self.transfer_topo.split_k_and_v
tensor_size_bytes = None
for layer_name, cache_or_caches in kv_caches.items():
layer_index = extract_layer_index(layer_name)
cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
layer_spec = self._layer_specs.get(layer_name)
if layer_spec is None:
logger.debug(
"Skipping layer %s because no KV cache spec is present.",
layer_name,
)
continue
if isinstance(layer_spec, MambaSpec):
conv, _ = cache_or_caches
cache_list = [conv]
else:
cache_list = self.transfer_topo.get_transfer_cache_regions(
cache_or_caches, layer_spec
)
logger.debug(
"registering layer %s with %d cache tensor(s)",
layer_name,
@@ -1501,45 +1691,46 @@ class MooncakeConnectorWorker:
for cache in cache_list:
self._log_debug_cache_registration(layer_name, cache)
base_addr = cache.data_ptr()
if base_addr in seen_base_addresses:
continue
seen_base_addresses.append(base_addr)
if tensor_size_bytes is None:
tensor_size_bytes = cache.nbytes
self.num_blocks = cache.shape[0]
assert cache.shape[0] == self.num_blocks, (
"All kv cache tensors must have the same number of blocks"
)
# Use stride-based block length so RDMA reaches the last
# block's padding (e.g. DeepseekV4 MLA alignment). stride(0)
# reflects the actual byte distance between consecutive
# blocks in GPU memory, which matches or exceeds the
# shape-based size.
block_len = cache.stride(0) * cache.element_size()
region_base_addresses.append(base_addr)
if isinstance(layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)):
kv_block_len = layer_spec.page_size_bytes
elif self.transfer_topo.virtually_split_kv_in_blocks and not isinstance(
layer_spec, MambaSpec
):
kv_block_len = block_len // 2
else:
kv_block_len = block_len
self.block_len_per_layer.append(block_len)
self.kv_block_len_per_layer.append(kv_block_len)
self.registered_layer_names.append(layer_name)
self.registered_layer_indices.append(layer_index)
kv_data_ptrs.append(base_addr)
kv_data_lens.append(self.num_blocks * block_len)
self.registered_group_indices.append(
self._layer_group_indices[layer_name]
)
storage = cache.untyped_storage()
storage_addr = storage.data_ptr()
if storage_addr not in seen_storage_ptrs:
seen_storage_ptrs.add(storage_addr)
kv_data_ptrs.append(storage_addr)
kv_data_lens.append(storage.nbytes())
self.kv_caches_base_addr = seen_base_addresses
self.seen_base_addresses = seen_base_addresses
self.kv_caches_base_addr = region_base_addresses
self.seen_base_addresses = kv_data_ptrs
if not kv_data_ptrs:
raise RuntimeError("No KV cache tensors were registered with Mooncake.")
ret_value = self.engine.batch_register_memory(kv_data_ptrs, kv_data_lens)
if ret_value != 0:
raise RuntimeError("Mooncake batch memory registration failed.")
assert tensor_size_bytes is not None
assert self.num_blocks != 0
self.device_kv_caches = kv_caches
logger.debug(
"registered num_blocks=%d block_lens=%s",
self.num_blocks,
"registered block_lens=%s kv_block_lens=%s",
self.block_len_per_layer,
self.kv_block_len_per_layer,
)
# No need to launch server for D node.
@@ -1642,8 +1833,10 @@ class MooncakeConnectorWorker:
},
kv_caches_base_addr=self.kv_caches_base_addr,
block_lens=self.block_len_per_layer,
kv_block_lens=self.kv_block_len_per_layer,
registered_layer_names=self.registered_layer_names,
registered_layer_indices=self.registered_layer_indices,
registered_group_indices=self.registered_group_indices,
)
encoded_data = self._encoder.encode(metadata)
@@ -1852,15 +2045,34 @@ class MooncakeConnectorWorker:
self,
base_addrs: list[int],
block_lens: list[int],
kv_block_lens: list[int],
layer_names: list[str],
layer_indices: list[int],
group_indices: list[int] | None = None,
) -> list[TransferRegion]:
if not group_indices:
group_indices = [
self._layer_group_indices.get(layer_name, 0)
for layer_name in layer_names
]
split_kv_regions = None
if self.transfer_topo.virtually_split_kv_in_blocks:
split_kv_regions = [
not isinstance(
self._layer_specs[layer_name],
(MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec),
)
for layer_name in layer_names
]
return _expand_transfer_regions(
base_addrs=base_addrs,
block_lens=block_lens,
kv_block_lens=kv_block_lens,
layer_names=layer_names,
layer_indices=layer_indices,
is_kv_layout_blocks_first=self.transfer_topo.virtually_split_kv_in_blocks,
group_indices=group_indices,
split_kv_regions=split_kv_regions,
)
def _get_sender_transfer_plan(
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+1 -1
View File
@@ -320,7 +320,7 @@ class NixlBaseConnectorScheduler:
logger.warning("Connection listener got unexpected message %s", msg)
sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))
def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int:
def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
"""D-side only. Returns N-1 for Mamba models since the decoder
always recomputes the last token and must start from h(N-1)."""
if self._has_mamba and num_prompt_tokens > 1:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+1 -1
View File
@@ -60,7 +60,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
if params is not None and params.get("do_remote_prefill"):
# Remote prefill: get all prompt blocks from remote.
token_ids = request.prompt_token_ids or []
actual = self._mamba_prefill_token_count(len(token_ids))
actual = self._get_remote_prefill_token_count(len(token_ids))
count = actual - num_computed_tokens
if count > 0:
return count, True
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py
+1 -1
View File
@@ -116,7 +116,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
if params is not None and params.get("do_remote_prefill"):
token_ids = request.prompt_token_ids or []
actual = self._mamba_prefill_token_count(len(token_ids))
actual = self._get_remote_prefill_token_count(len(token_ids))
count = actual - num_computed_tokens
if count > 0:
return count, True