From d2906091bfc579cebefe3d8e8fb9077397ce9882 Mon Sep 17 00:00:00 2001 From: Yicong Wang Date: Wed, 9 Sep 2026 14:18:32 +0800 Subject: [PATCH] [Bugfix] Fix Mooncake heterogeneous TP with replicated GQA heads (#52516) Signed-off-by: wangyicong Co-authored-by: OpenAI Codex --- .../unit/test_mooncake_connector.py | 95 ++++++++++++++++++- .../v1/mooncake/mooncake_connector.py | 64 ++++++++++++- 2 files changed, 156 insertions(+), 3 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index 7db886e808..f989fb4e32 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -26,6 +26,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector im SendBlockMeta, TransferRegion, _align_transfer_regions, + _compute_sender_transfer_plan, + _validate_asymmetric_region_lengths, get_mooncake_bootstrap_addr, should_launch_bootstrap_server, ) @@ -44,6 +46,93 @@ from vllm.v1.request import RequestStatus from .utils import create_request, create_scheduler, create_vllm_config +@pytest.mark.parametrize( + ("remote_tp_size", "expected_dst_offsets"), + [ + (1, [0, None, 65536, None, 131072, None, 196608, None]), + (2, [0, None, 65536, None, 0, None, 65536, None]), + ], +) +@pytest.mark.parametrize("local_tp_rank", range(8)) +def test_sender_plan_replicated_gqa_tp8_to_smaller_tp( + remote_tp_size, + expected_dst_offsets, + local_tp_rank, +): + expected_dst_offset = expected_dst_offsets[local_tp_rank] + plan = _compute_sender_transfer_plan( + local_tp_rank=local_tp_rank, + local_tp_size=8, + remote_tp_rank=local_tp_rank // (8 // remote_tp_size), + remote_tp_size=remote_tp_size, + local_kv_block_len=65536, + remote_kv_block_len=262144 // remote_tp_size, + producer_cache_replicated=True, + total_num_kv_heads=4, + ) + if expected_dst_offset is None: + assert not plan[0] + else: + assert plan == (True, 0, expected_dst_offset, 65536) + + +@pytest.mark.parametrize( + ("local_tp_size", "expected_src_offsets"), + [ + (1, [0, 0, 65536, 65536, 131072, 131072, 196608, 196608]), + (2, [0, 0, 65536, 65536, 0, 0, 65536, 65536]), + (4, [0] * 8), + ], +) +@pytest.mark.parametrize("remote_tp_rank", range(8)) +def test_sender_plan_gqa_to_replicated_tp8( + local_tp_size, + expected_src_offsets, + remote_tp_rank, +): + local_head_count = 4 // local_tp_size + plan = _compute_sender_transfer_plan( + local_tp_rank=remote_tp_rank // (8 // local_tp_size), + local_tp_size=local_tp_size, + remote_tp_rank=remote_tp_rank, + remote_tp_size=8, + local_kv_block_len=local_head_count * 65536, + remote_kv_block_len=65536, + producer_cache_replicated=False, + total_num_kv_heads=4, + ) + assert plan == (True, expected_src_offsets[remote_tp_rank], 0, 65536) + + +@pytest.mark.parametrize( + "local_tp,remote_tp,local_len,remote_len,valid", + [ + (1, 8, 262144, 65536, True), + (8, 1, 65536, 262144, True), + (1, 8, 262144, 131072, False), + (8, 1, 131072, 262144, False), + ], +) +def test_region_length_validation_checks_replicated_gqa_heads( + local_tp, remote_tp, local_len, remote_len, valid +): + """Replicated heads must have matching, whole per-head payloads.""" + local_region = TransferRegion("layer", 0, 0, local_len, local_len) + remote_region = TransferRegion("layer", 0, 0, remote_len, remote_len) + + assert ( + _validate_asymmetric_region_lengths( + local_regions=[local_region], + remote_regions=[remote_region], + local_tp_size=local_tp, + remote_tp_size=remote_tp, + producer_cache_replicated=local_tp > 4, + total_num_kv_heads=4, + ) + is None + ) == valid + + def _make_test_kv_cache_config() -> KVCacheConfig: return KVCacheConfig( num_blocks=0, @@ -150,9 +239,13 @@ async def test_build_transfer_params_separates_prefill_pp_layers(): worker.is_kv_producer = True worker.tp_rank = 0 worker.tp_size = 1 + worker.use_mla = False 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) + worker.transfer_topo = SimpleNamespace( + local_replicates_kv_cache=False, + total_num_kv_heads=4, + ) block_len = 256 remote_regions = [ 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 e5f2e4f48c..dc574a6219 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 @@ -182,6 +182,7 @@ def _compute_sender_transfer_plan( local_kv_block_len: int, remote_kv_block_len: int, producer_cache_replicated: bool, + total_num_kv_heads: int | None = None, ) -> tuple[bool, int, int, int]: """Plan one producer-rank to one consumer-rank copy for heterogeneous TP.""" tp_ratio = _get_tp_ratio(local_tp_size, remote_tp_size) @@ -189,6 +190,23 @@ def _compute_sender_transfer_plan( if tp_ratio == 1: return True, 0, 0, local_kv_block_len + if total_num_kv_heads is not None: + consumer_cache_replicated = remote_tp_size > total_num_kv_heads + if producer_cache_replicated != consumer_cache_replicated: + local_head_count = max(total_num_kv_heads // local_tp_size, 1) + local_replica_count = max(local_tp_size // total_num_kv_heads, 1) + + local_head = local_tp_rank * total_num_kv_heads // local_tp_size + remote_head = remote_tp_rank * total_num_kv_heads // remote_tp_size + bytes_per_head = local_kv_block_len // local_head_count + + return ( + local_tp_rank % local_replica_count == 0, + max(remote_head - local_head, 0) * bytes_per_head, + max(local_head - remote_head, 0) * bytes_per_head, + bytes_per_head, + ) + if tp_ratio > 0: if producer_cache_replicated: return local_tp_rank % tp_ratio == 0, 0, 0, local_kv_block_len @@ -233,6 +251,7 @@ def _validate_asymmetric_region_lengths( local_tp_size: int, remote_tp_size: int, producer_cache_replicated: bool, + total_num_kv_heads: int | None = None, ) -> str | None: """Validate transfer-region metadata for a fixed producer/consumer pair. @@ -246,7 +265,11 @@ def _validate_asymmetric_region_lengths( "producer and consumer." ) - if producer_cache_replicated: + if total_num_kv_heads is not None: + # TP ranks beyond the KV-head count replicate existing shards. + local_tp_size = min(local_tp_size, total_num_kv_heads) + remote_tp_size = min(remote_tp_size, total_num_kv_heads) + elif producer_cache_replicated: return None tp_ratio = _get_tp_ratio(local_tp_size, remote_tp_size) @@ -1218,12 +1241,19 @@ class MooncakeConnectorWorker: ) await sock.send_multipart((identity, self._encoder.encode(response))) return - validation_err = _validate_asymmetric_region_lengths( + validation_err = self._validate_head_resharding_layout( + meta.remote_tp_size, local_regions + ) or _validate_asymmetric_region_lengths( local_regions=local_regions, remote_regions=remote_regions, local_tp_size=self.tp_size, remote_tp_size=meta.remote_tp_size, producer_cache_replicated=self._producer_cache_is_replicated(), + total_num_kv_heads=( + None + if self.use_mla or self.kv_cache_config.has_mamba_layers + else self.transfer_topo.total_num_kv_heads + ), ) if validation_err is not None: response = MooncakeXferResponse( @@ -2035,6 +2065,31 @@ class MooncakeConnectorWorker: def _producer_cache_is_replicated(self) -> bool: return self.transfer_topo.local_replicates_kv_cache + def _validate_head_resharding_layout( + self, remote_tp_size: int, local_regions: list[TransferRegion] + ) -> str | None: + """Reject unsupported layouts before splitting or gathering KV heads.""" + if self.use_mla or self.kv_cache_config.has_mamba_layers: + return None + num_kv_heads = self.transfer_topo.total_num_kv_heads + if min(self.tp_size, num_kv_heads) == min(remote_tp_size, num_kv_heads): + return None + + for region in local_regions: + spec = self._layer_specs[region.layer_name] + if not isinstance(spec, AttentionSpec): + continue + if spec.page_size_bytes > spec.unpadded_page_size_bytes: + return ( + "Mooncake KV-head re-sharding is not supported for padded " + f"KV pages (layer {region.layer_name})." + ) + if spec.kv_quant_mode.is_nvfp4: + return ( + "Mooncake KV-head re-sharding is not supported for NVFP4 KV cache." + ) + return None + def _get_transfer_regions( self, base_addrs: list[int], @@ -2073,6 +2128,11 @@ class MooncakeConnectorWorker: local_kv_block_len=local_kv_block_len, remote_kv_block_len=remote_kv_block_len, producer_cache_replicated=self._producer_cache_is_replicated(), + total_num_kv_heads=( + None + if self.use_mla or self.kv_cache_config.has_mamba_layers + else self.transfer_topo.total_num_kv_heads + ), ) def _log_debug_cache_registration(