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

[Bugfix] Fix Mooncake heterogeneous TP with replicated GQA heads (#52516)

Browse Source 
Signed-off-by: wangyicong <wangyicong@bytedance.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Yicong Wang
2026-09-09 14:18:32 +08:00
committed by GitHub
co-authored by OpenAI Codex
parent 95cf420ac1
commit d2906091bf
2 changed files with 156 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_mooncake_connector.py
+94 -1
View File
@@ -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 = [
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+62 -2
View File
@@ -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(