[Bugfix] Fix Mooncake heterogeneous TP with replicated GQA heads (#52516)
Signed-off-by: wangyicong <wangyicong@bytedance.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Yicong Wang committed
Sep 9, 2026 at 14:18 UTC
d2906091bfc579cebefe3d8e8fb9077397ce9882
2 files changed
+156
-3
tests/v1/kv_connector/unit/test_mooncake_connector.py
+94
-1
@@ -26,6 +26,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector im
26
SendBlockMeta,
27
TransferRegion,
28
_align_transfer_regions,
29
+ _compute_sender_transfer_plan,
30
+ _validate_asymmetric_region_lengths,
31
get_mooncake_bootstrap_addr,
32
should_launch_bootstrap_server,
33
)
@@ -44,6 +46,93 @@ from vllm.v1.request import RequestStatus
46
from .utils import create_request, create_scheduler, create_vllm_config
47
48
49
+@pytest.mark.parametrize(
50
+ ("remote_tp_size", "expected_dst_offsets"),
51
+ [
52
+ (1, [0, None, 65536, None, 131072, None, 196608, None]),
53
+ (2, [0, None, 65536, None, 0, None, 65536, None]),
54
+ ],
55
+)
56
+@pytest.mark.parametrize("local_tp_rank", range(8))
57
+def test_sender_plan_replicated_gqa_tp8_to_smaller_tp(
58
+ remote_tp_size,
59
+ expected_dst_offsets,
60
+ local_tp_rank,
61
+):
62
+ expected_dst_offset = expected_dst_offsets[local_tp_rank]
63
+ plan = _compute_sender_transfer_plan(
64
+ local_tp_rank=local_tp_rank,
65
+ local_tp_size=8,
66
+ remote_tp_rank=local_tp_rank // (8 // remote_tp_size),
67
+ remote_tp_size=remote_tp_size,
68
+ local_kv_block_len=65536,
69
+ remote_kv_block_len=262144 // remote_tp_size,
70
+ producer_cache_replicated=True,
71
+ total_num_kv_heads=4,
72
+ )
73
+ if expected_dst_offset is None:
74
+ assert not plan[0]
75
+ else:
76
+ assert plan == (True, 0, expected_dst_offset, 65536)
77
+
78
+
79
+@pytest.mark.parametrize(
80
+ ("local_tp_size", "expected_src_offsets"),
81
+ [
82
+ (1, [0, 0, 65536, 65536, 131072, 131072, 196608, 196608]),
83
+ (2, [0, 0, 65536, 65536, 0, 0, 65536, 65536]),
84
+ (4, [0] * 8),
85
+ ],
86
+)
87
+@pytest.mark.parametrize("remote_tp_rank", range(8))
88
+def test_sender_plan_gqa_to_replicated_tp8(
89
+ local_tp_size,
90
+ expected_src_offsets,
91
+ remote_tp_rank,
92
+):
93
+ local_head_count = 4 // local_tp_size
94
+ plan = _compute_sender_transfer_plan(
95
+ local_tp_rank=remote_tp_rank // (8 // local_tp_size),
96
+ local_tp_size=local_tp_size,
97
+ remote_tp_rank=remote_tp_rank,
98
+ remote_tp_size=8,
99
+ local_kv_block_len=local_head_count * 65536,
100
+ remote_kv_block_len=65536,
101
+ producer_cache_replicated=False,
102
+ total_num_kv_heads=4,
103
+ )
104
+ assert plan == (True, expected_src_offsets[remote_tp_rank], 0, 65536)
105
+
106
+
107
+@pytest.mark.parametrize(
108
+ "local_tp,remote_tp,local_len,remote_len,valid",
109
+ [
110
+ (1, 8, 262144, 65536, True),
111
+ (8, 1, 65536, 262144, True),
112
+ (1, 8, 262144, 131072, False),
113
+ (8, 1, 131072, 262144, False),
114
+ ],
115
+)
116
+def test_region_length_validation_checks_replicated_gqa_heads(
117
+ local_tp, remote_tp, local_len, remote_len, valid
118
+):
119
+ """Replicated heads must have matching, whole per-head payloads."""
120
+ local_region = TransferRegion("layer", 0, 0, local_len, local_len)
121
+ remote_region = TransferRegion("layer", 0, 0, remote_len, remote_len)
122
+
123
+ assert (
124
+ _validate_asymmetric_region_lengths(
125
+ local_regions=[local_region],
126
+ remote_regions=[remote_region],
127
+ local_tp_size=local_tp,
128
+ remote_tp_size=remote_tp,
129
+ producer_cache_replicated=local_tp > 4,
130
+ total_num_kv_heads=4,
131
+ )
132
+ is None
133
+ ) == valid
134
+
135
+
136
def _make_test_kv_cache_config() -> KVCacheConfig:
137
return KVCacheConfig(
138
num_blocks=0,
@@ -150,9 +239,13 @@ async def test_build_transfer_params_separates_prefill_pp_layers():
239
worker.is_kv_producer = True
240
worker.tp_rank = 0
241
worker.tp_size = 1
242
+ worker.use_mla = False
243
worker.kv_cache_config = _make_test_kv_cache_config()
244
worker._physical_blocks_per_logical_kv_block = 1
155
- worker.transfer_topo = SimpleNamespace(local_replicates_kv_cache=False)
245
+ worker.transfer_topo = SimpleNamespace(
246
+ local_replicates_kv_cache=False,
247
+ total_num_kv_heads=4,
248
+ )
249
250
block_len = 256
251
remote_regions = [
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+62
-2
@@ -182,6 +182,7 @@ def _compute_sender_transfer_plan(
182
local_kv_block_len: int,
183
remote_kv_block_len: int,
184
producer_cache_replicated: bool,
185
+ total_num_kv_heads: int | None = None,
186
) -> tuple[bool, int, int, int]:
187
"""Plan one producer-rank to one consumer-rank copy for heterogeneous TP."""
188
tp_ratio = _get_tp_ratio(local_tp_size, remote_tp_size)
@@ -189,6 +190,23 @@ def _compute_sender_transfer_plan(
190
if tp_ratio == 1:
191
return True, 0, 0, local_kv_block_len
192
193
+ if total_num_kv_heads is not None:
194
+ consumer_cache_replicated = remote_tp_size > total_num_kv_heads
195
+ if producer_cache_replicated != consumer_cache_replicated:
196
+ local_head_count = max(total_num_kv_heads // local_tp_size, 1)
197
+ local_replica_count = max(local_tp_size // total_num_kv_heads, 1)
198
+
199
+ local_head = local_tp_rank * total_num_kv_heads // local_tp_size
200
+ remote_head = remote_tp_rank * total_num_kv_heads // remote_tp_size
201
+ bytes_per_head = local_kv_block_len // local_head_count
202
+
203
+ return (
204
+ local_tp_rank % local_replica_count == 0,
205
+ max(remote_head - local_head, 0) * bytes_per_head,
206
+ max(local_head - remote_head, 0) * bytes_per_head,
207
+ bytes_per_head,
208
+ )
209
+
210
if tp_ratio > 0:
211
if producer_cache_replicated:
212
return local_tp_rank % tp_ratio == 0, 0, 0, local_kv_block_len
@@ -233,6 +251,7 @@ def _validate_asymmetric_region_lengths(
251
local_tp_size: int,
252
remote_tp_size: int,
253
producer_cache_replicated: bool,
254
+ total_num_kv_heads: int | None = None,
255
) -> str | None:
256
"""Validate transfer-region metadata for a fixed producer/consumer pair.
257
@@ -246,7 +265,11 @@ def _validate_asymmetric_region_lengths(
265
"producer and consumer."
266
)
267
249
- if producer_cache_replicated:
268
+ if total_num_kv_heads is not None:
269
+ # TP ranks beyond the KV-head count replicate existing shards.
270
+ local_tp_size = min(local_tp_size, total_num_kv_heads)
271
+ remote_tp_size = min(remote_tp_size, total_num_kv_heads)
272
+ elif producer_cache_replicated:
273
return None
274
275
tp_ratio = _get_tp_ratio(local_tp_size, remote_tp_size)
@@ -1218,12 +1241,19 @@ class MooncakeConnectorWorker:
1241
)
1242
await sock.send_multipart((identity, self._encoder.encode(response)))
1243
return
1221
- validation_err = _validate_asymmetric_region_lengths(
1244
+ validation_err = self._validate_head_resharding_layout(
1245
+ meta.remote_tp_size, local_regions
1246
+ ) or _validate_asymmetric_region_lengths(
1247
local_regions=local_regions,
1248
remote_regions=remote_regions,
1249
local_tp_size=self.tp_size,
1250
remote_tp_size=meta.remote_tp_size,
1251
producer_cache_replicated=self._producer_cache_is_replicated(),
1252
+ total_num_kv_heads=(
1253
+ None
1254
+ if self.use_mla or self.kv_cache_config.has_mamba_layers
1255
+ else self.transfer_topo.total_num_kv_heads
1256
+ ),
1257
)
1258
if validation_err is not None:
1259
response = MooncakeXferResponse(
@@ -2035,6 +2065,31 @@ class MooncakeConnectorWorker:
2065
def _producer_cache_is_replicated(self) -> bool:
2066
return self.transfer_topo.local_replicates_kv_cache
2067
2068
+ def _validate_head_resharding_layout(
2069
+ self, remote_tp_size: int, local_regions: list[TransferRegion]
2070
+ ) -> str | None:
2071
+ """Reject unsupported layouts before splitting or gathering KV heads."""
2072
+ if self.use_mla or self.kv_cache_config.has_mamba_layers:
2073
+ return None
2074
+ num_kv_heads = self.transfer_topo.total_num_kv_heads
2075
+ if min(self.tp_size, num_kv_heads) == min(remote_tp_size, num_kv_heads):
2076
+ return None
2077
+
2078
+ for region in local_regions:
2079
+ spec = self._layer_specs[region.layer_name]
2080
+ if not isinstance(spec, AttentionSpec):
2081
+ continue
2082
+ if spec.page_size_bytes > spec.unpadded_page_size_bytes:
2083
+ return (
2084
+ "Mooncake KV-head re-sharding is not supported for padded "
2085
+ f"KV pages (layer {region.layer_name})."
2086
+ )
2087
+ if spec.kv_quant_mode.is_nvfp4:
2088
+ return (
2089
+ "Mooncake KV-head re-sharding is not supported for NVFP4 KV cache."
2090
+ )
2091
+ return None
2092
+
2093
def _get_transfer_regions(
2094
self,
2095
base_addrs: list[int],
@@ -2073,6 +2128,11 @@ class MooncakeConnectorWorker:
2128
local_kv_block_len=local_kv_block_len,
2129
remote_kv_block_len=remote_kv_block_len,
2130
producer_cache_replicated=self._producer_cache_is_replicated(),
2131
+ total_num_kv_heads=(
2132
+ None
2133
+ if self.use_mla or self.kv_cache_config.has_mamba_layers
2134
+ else self.transfer_topo.total_num_kv_heads
2135
+ ),
2136
)
2137
2138
def _log_debug_cache_registration(