[Bugfix] MoRIIO toy P/D proxy: fix DP-rank index aliasing + harden for high-concurrency bursts (#46115)
Signed-off-by: Edwin Lim <edwin.lim@mangoboost.io> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: QinPR <1905873179@qq.com> Co-authored-by: Peiran Qin <66068739+QinPR@users.noreply.github.com>
limeward committed
Jul 17, 2026 at 12:35 UTC
c4cd2bd5440bca29039e5b6f33b695a95721bd3e
2 files changed
+268
-13
examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+63
-13
@@ -210,8 +210,31 @@ async def stream_decode_response(session, response, request_id):
210
await session.close()
211
212
213
-def example_round_robin_dp_loader(request_number, dp_size):
214
- return request_nums % dp_size
213
+def flat_interleaved_dp_route(request_number, instances):
214
+ """Flat round-robin over the full (instance, dp_rank) slot space.
215
+
216
+ ONE counter over (n_instances * dp_size) slots, so instance-selection and
217
+ DP-rank-selection are derived from the SAME index and can never alias. The
218
+ previous scheme computed instance = req % n and rank = req % dp from the
219
+ same counter with n | dp, which locked each instance to a stride-n subset
220
+ of its ranks (e.g. 2 prefill instances -> 4 of 8 ranks each -> half the
221
+ GPUs never receive a request, so the deployment falsely appears not to
222
+ scale).
223
+
224
+ Interleaved order — inst0_r0, inst1_r0, inst0_r1, inst1_r1, ... — so
225
+ consecutive requests alternate instances AND every rank gets walked.
226
+
227
+ Assumes homogeneous dp_size across a role's instances (true for the
228
+ DP<->DP and DP<->TP deployments this proxy targets). Returns
229
+ (instance_index, dp_rank); dp_rank is None when dp_size == 1 (e.g. a TP
230
+ decode), which avoids forwarding an out-of-range data-parallel rank.
231
+ """
232
+ n = len(instances)
233
+ dp = instances[0]["dp_size"]
234
+ slot = (request_number - 1) % (n * dp)
235
+ inst_idx = slot % n
236
+ dp_rank = (slot // n) if dp > 1 else None
237
+ return inst_idx, dp_rank
238
239
240
@app.route("/health", methods=["GET"])
@@ -252,18 +275,21 @@ async def handle_request(api: str, request: Request):
275
503,
276
)
277
)
255
- pid = request_nums % len(prefill_instances)
256
- did = request_nums % len(decode_instances)
278
+ # Flat interleaved round-robin (see flat_interleaved_dp_route): ONE
279
+ # counter over the full (instance, dp_rank) slot space per role, so
280
+ # instance-selection and DP-rank-selection derive from the same index
281
+ # and can never alias. The old scheme keyed both on request_nums with
282
+ # n_instances | dp_size, stranding half the ranks (e.g. in 2P_DP8EP).
283
+ pid, selected_prefill_dp_rank = flat_interleaved_dp_route(
284
+ request_nums, prefill_instances
285
+ )
286
+ # Decode instance selection uses the same interleaved walk; in READ
287
+ # mode the decode reads KV from selected_prefill_dp_rank, so the
288
+ # decode's own dp_rank is not forwarded here.
289
+ did, _ = flat_interleaved_dp_route(request_nums, decode_instances)
290
prefill_instance_endpoint = prefill_instances[pid]
291
decode_instance_endpoint = decode_instances[did]
292
260
- selected_prefill_dp_rank = None
261
- if prefill_instance_endpoint["dp_size"] > 1:
262
- selected_prefill_dp_rank = example_round_robin_dp_loader(
263
- request_nums // len(prefill_instance_endpoint),
264
- prefill_instance_endpoint["dp_size"],
265
- )
266
-
293
# Embed both zmq_addresses in the request_id so the connector can parse
294
# the peer's host/ports from it, similar to P2P-NCCL
295
uid = str(uuid.uuid4()).replace("-", "")
@@ -427,9 +453,33 @@ if __name__ == "__main__":
453
args = parser.parse_args()
454
455
t = start_service_discovery("0.0.0.0", 36367)
430
- app.debug = True
456
+ # High-concurrency hardening. Quart's app.run() uses a shallow listen
457
+ # backlog (100) and, with app.debug=True, adds per-request overhead that
458
+ # starves the single accept loop. Under a burst of ~512 simultaneous client
459
+ # connections the backlog overflows and the kernel RSTs the excess, so
460
+ # clients see "ClientOSError: [Errno 104] Connection reset by peer" before
461
+ # any response (~16% request loss at c=512). Serve via hypercorn with debug
462
+ # OFF and a deep backlog so the burst QUEUES (higher TTFT) instead of being
463
+ # reset -> 100% request success.
464
+ app.debug = False
465
app.config["BODY_TIMEOUT"] = 360000
466
app.config["RESPONSE_TIMEOUT"] = 360000
467
434
- app.run(host="0.0.0.0", port=args.port)
468
+ import asyncio
469
+ import os
470
+
471
+ from hypercorn.asyncio import serve as _hypercorn_serve
472
+ from hypercorn.config import Config as _HypercornConfig
473
+
474
+ _hcfg = _HypercornConfig()
475
+ _hcfg.bind = [f"0.0.0.0:{args.port}"]
476
+ # Deep listen backlog so a wide connection burst queues, not RSTs. NOTE:
477
+ # effective backlog is capped by the host's net.core.somaxconn (proxy runs
478
+ # --network host); kernel 6.x defaults to 4096. Override via
479
+ # PROXY_LISTEN_BACKLOG.
480
+ _hcfg.backlog = int(os.environ.get("PROXY_LISTEN_BACKLOG", "4096"))
481
+ # Long-lived SSE streams (8k1k decode ~5 min): never reap on keepalive.
482
+ _hcfg.keep_alive_timeout = 360000.0
483
+
484
+ asyncio.run(_hypercorn_serve(app, _hcfg))
485
t.join()
tests/v1/kv_connector/unit/test_moriio_proxy_routing.py
new
+205
@@ -0,0 +1,205 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Hardware-fair request routing in the MoRIIO toy P/D proxy.
4
+
5
+Exercises the REAL ``flat_interleaved_dp_route`` from the toy proxy (loaded from
6
+its ``examples/`` path with heavy deps stubbed) — no routing logic is copied
7
+here. This is the request-distribution half of end-to-end fairness: the proxy
8
+must hand every prefill/decode (instance, dp_rank) slot an equal share so no GPU
9
+is starved. The connector-side read routing is covered in
10
+``test_moriio_routing_fairness.py``.
11
+
12
+Regression target: the previous scheme derived ``instance = req % n_instances``
13
+and ``dp_rank = req % dp_size`` from the same counter, so when
14
+``n_instances | dp_size`` each instance was locked to a stride-``n`` subset of
15
+its ranks — e.g. 2 prefill instances x DP8 stranded 4 of every node's 8 GPUs.
16
+``flat_interleaved_dp_route`` walks ONE counter over the full
17
+``(instance, dp_rank)`` slot space, so the two selections can never alias.
18
+
19
+Role shapes below are exactly those in the RFC (#46107) deployments:
20
+ (1, 1) 1P/1D TP8 (2, 1) 2P/2D TP8 (4, 1) 4D TP8
21
+ (1, 8) 1D DP8EP (2, 8) 2P DP8EP (3, 8) 3D DP8EP
22
+"""
23
+
24
+import contextlib
25
+import importlib.util
26
+import sys
27
+import types
28
+from collections import Counter
29
+from pathlib import Path
30
+from typing import cast
31
+
32
+import pytest
33
+
34
+_MISSING = object()
35
+
36
+PROXY_REL = "examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py"
37
+
38
+
39
+def _module(name, **attrs):
40
+ module = types.ModuleType(name)
41
+ for attr, value in attrs.items():
42
+ setattr(module, attr, value)
43
+ return module
44
+
45
+
46
+def _package(name):
47
+ module = _module(name)
48
+ module.__path__ = []
49
+ return module
50
+
51
+
52
+class _QuartStub:
53
+ def __init__(self, *a, **k):
54
+ pass
55
+
56
+ def route(self, *a, **k):
57
+ return lambda fn: fn
58
+
59
+ def post(self, *a, **k):
60
+ return self.route()
61
+
62
+
63
+async def _make_response_stub(value):
64
+ return value
65
+
66
+
67
+@contextlib.contextmanager
68
+def _proxy_import_stubs():
69
+ """Stub the proxy's external deps so the module imports for a pure unit test.
70
+
71
+ Only third-party/vllm imports are stubbed; the routing function under test
72
+ is executed as-is from the real module.
73
+ """
74
+ common = "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common"
75
+
76
+ class _MoRIIOConstants:
77
+ TRANSFER_PREFIX = "moriio-transfer"
78
+
79
+ stubs = {
80
+ "aiohttp": _module("aiohttp"),
81
+ "msgpack": _module("msgpack"),
82
+ "zmq": _module("zmq"),
83
+ "quart": _module(
84
+ "quart",
85
+ Quart=_QuartStub,
86
+ Request=object,
87
+ make_response=_make_response_stub,
88
+ request=object(),
89
+ ),
90
+ "vllm": _package("vllm"),
91
+ "vllm.distributed": _package("vllm.distributed"),
92
+ "vllm.distributed.kv_transfer": _package("vllm.distributed.kv_transfer"),
93
+ "vllm.distributed.kv_transfer.kv_connector": _package(
94
+ "vllm.distributed.kv_transfer.kv_connector"
95
+ ),
96
+ "vllm.distributed.kv_transfer.kv_connector.v1": _package(
97
+ "vllm.distributed.kv_transfer.kv_connector.v1"
98
+ ),
99
+ "vllm.distributed.kv_transfer.kv_connector.v1.moriio": _package(
100
+ "vllm.distributed.kv_transfer.kv_connector.v1.moriio"
101
+ ),
102
+ common: _module(common, MoRIIOConstants=_MoRIIOConstants),
103
+ }
104
+ saved = {}
105
+ for name, module in stubs.items():
106
+ if name not in sys.modules:
107
+ saved[name] = _MISSING
108
+ sys.modules[name] = module
109
+ try:
110
+ yield
111
+ finally:
112
+ for name, previous in saved.items():
113
+ if previous is _MISSING:
114
+ sys.modules.pop(name, None)
115
+ else:
116
+ sys.modules[name] = cast(types.ModuleType, previous)
117
+
118
+
119
+def _load_proxy_module():
120
+ path = Path(__file__).parents[4] / PROXY_REL
121
+ spec = importlib.util.spec_from_file_location("moriio_proxy_under_test", path)
122
+ assert spec is not None and spec.loader is not None
123
+ module = importlib.util.module_from_spec(spec)
124
+ with _proxy_import_stubs():
125
+ spec.loader.exec_module(module)
126
+ return module
127
+
128
+
129
+@pytest.fixture(scope="module")
130
+def route():
131
+ return _load_proxy_module().flat_interleaved_dp_route
132
+
133
+
134
+def _instances(n: int, dp_size: int):
135
+ # Only dp_size is read by the router; tp_size carried for realism.
136
+ return [{"dp_size": dp_size, "tp_size": 8 // dp_size} for _ in range(n)]
137
+
138
+
139
+# (n_instances, dp_size) for every distinct P/D role shape in the RFC configs.
140
+ROLE_SHAPES = [(1, 1), (2, 1), (4, 1), (1, 8), (2, 8), (3, 8)]
141
+SHAPE_IDS = [f"n{n}_dp{dp}" for n, dp in ROLE_SHAPES]
142
+
143
+
144
+def _route_n(route, instances, count):
145
+ # Proxy uses 1-indexed request numbers (slot = (request_number - 1) % ...).
146
+ return [route(rn, instances) for rn in range(1, count + 1)]
147
+
148
+
149
+@pytest.mark.parametrize(("n", "dp"), ROLE_SHAPES, ids=SHAPE_IDS)
150
+def test_full_slot_space_is_covered_uniformly(route, n, dp):
151
+ instances = _instances(n, dp)
152
+ period = n * dp
153
+ # Three full cycles -> every (instance, dp_rank) slot must be hit the same
154
+ # number of times (exactly uniform, no starved slot, no aliasing).
155
+ hits = Counter(_route_n(route, instances, period * 3))
156
+
157
+ expected: set[tuple[int, int | None]]
158
+ if dp == 1:
159
+ expected = {(inst, None) for inst in range(n)}
160
+ else:
161
+ expected = {(inst, r) for inst in range(n) for r in range(dp)}
162
+ assert set(hits) == expected, f"missing slots: {expected - set(hits)}"
163
+ assert max(hits.values()) == min(hits.values())
164
+
165
+
166
+@pytest.mark.parametrize(("n", "dp"), ROLE_SHAPES, ids=SHAPE_IDS)
167
+def test_instance_and_dp_rank_marginals_are_balanced(route, n, dp):
168
+ instances = _instances(n, dp)
169
+ routed = _route_n(route, instances, n * dp * 5)
170
+
171
+ inst_counts = Counter(inst for inst, _ in routed)
172
+ assert set(inst_counts) == set(range(n))
173
+ assert max(inst_counts.values()) == min(inst_counts.values())
174
+
175
+ dp_counts = Counter(r for _, r in routed)
176
+ if dp == 1:
177
+ assert set(dp_counts) == {None}
178
+ else:
179
+ assert set(dp_counts) == set(range(dp))
180
+ assert max(dp_counts.values()) == min(dp_counts.values())
181
+
182
+
183
+def test_tp_instance_forwards_no_dp_rank(route):
184
+ # dp_size == 1 (a TP instance) must yield dp_rank None so the proxy never
185
+ # forwards an out-of-range data-parallel rank.
186
+ assert all(r is None for _, r in _route_n(route, _instances(2, 1), 8))
187
+
188
+
189
+def test_two_instance_dp8_gives_every_node_all_ranks(route):
190
+ # Direct regression for the stranded-GPU bug: with 2 prefill instances x DP8
191
+ # each instance must receive ALL 8 dp ranks (16 distinct slots), not a
192
+ # stride-2 subset of 4.
193
+ routed = _route_n(route, _instances(2, 8), 2 * 8)
194
+ per_instance: dict[int, set] = {0: set(), 1: set()}
195
+ for inst, dp_rank in routed:
196
+ per_instance[inst].add(dp_rank)
197
+ assert per_instance[0] == set(range(8))
198
+ assert per_instance[1] == set(range(8))
199
+
200
+
201
+def test_consecutive_requests_alternate_instances(route):
202
+ # Interleaved order spreads consecutive requests across instances rather
203
+ # than filling one instance's ranks before moving on.
204
+ insts = [inst for inst, _ in _route_n(route, _instances(3, 8), 3)]
205
+ assert insts == [0, 1, 2]