diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index d1ae1c6de9..a227af909e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -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], ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py index f45074fff7..54b97d482e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hma.py @@ -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 = [ diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py new file mode 100644 index 0000000000..6811d60ee6 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py @@ -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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 1bc23cead5..34f99b7cfa 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -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( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index cba81cadd8..e506850bd7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -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: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index f13e216056..8c73461c5c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index 8b43709678..23c6e6cd35 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py @@ -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