why does P need to know the dp and tp size of D for transfer?

This is due to the RDMA-based push(write) mode.
If using pull (read) mode, this is not required.

can you also provide a proxy example for the pull mode?

@kouroshHakha do you have any other comments other than this example request?

Let's add this example in a follow up PR. In this PR @inkcherry has added an example script examples/online_serving/disaggregated_serving/moriio_toy_proxy_server.py.

Both pull mode and push mode use a unified proxy. For pull mode, you can refer to 1P1D TP8 (PULL MODE) in the test plan of this pr.

why is there a p_endpoint? Isn't that supposed to be the d_endpoint?

yes, thanks, renamed

Undefined variable when UnicodeDecodeError is raised

High Severity

In _handle_message, when msg.decode("UTF-8") raises a UnicodeDecodeError, the variable msg_str is never assigned because the exception occurs during its assignment. However, the exception handler on line 532 and line 534 both try to use msg_str, which will cause a NameError crash. The logging statement should use msg directly or handle the case where msg_str is not defined.

Undefined variables when handshake future already exists

High Severity

In _background_moriio_handshake, the variables host, port, tp_size, and remote_dp_size are only assigned inside the if fut is None block. However, these variables are used unconditionally on line 1036 in range(remote_dp_size) and line 1039. If a handshake future already exists for this remote_engine_id, fut will not be None, the variables won't be assigned, and the function will crash with a NameError. This can occur when a second request for the same remote engine arrives while handshaking is in progress.

Wrong variable used for KV cache size calculation

Medium Severity

In register_kv_caches, line 1120 uses cache from a previous loop instead of kv_cache from the current iteration. The variable cache was last assigned in the loop on lines 1103-1109, and always points to the last cache of the last layer processed. This causes incorrect size calculations where all layers are recorded with the same size from the final cache, rather than their actual sizes. The line should use kv_cache.nelement() * kv_cache.element_size().

Function ignores parameter, uses global instead

Medium Severity

The function example_round_robin_dp_loader accepts a parameter request_number but ignores it entirely, instead using the global variable request_nums. This makes the function parameter meaningless and the call on line 231 passes a calculated value that is never used. The function body should use request_number instead of request_nums.

# Test to verify the bug
request_nums = 10  # Global

def example_round_robin_dp_loader(request_number, dp_size):
    return request_nums % dp_size  # Bug: uses global, not parameter

# The parameter is ignored - result is based on global `request_nums`, not `request_number`
result = example_round_robin_dp_loader(5, 3)
print(f"Called with request_number=5, dp_size=3")
print(f"Expected: 5 % 3 = 2")
print(f"Actual: {result} (because it uses global request_nums={request_nums})")
print(f"Bug confirmed: {result != 2}")

python3 -c "
request_nums = 10
def example_round_robin_dp_loader(request_number, dp_size):
    return request_nums % dp_size
result = example_round_robin_dp_loader(5, 3)
print(f'Called with request_number=5, dp_size=3')
print(f'Expected: 5 % 3 = 2')
print(f'Actual: {result} (uses global request_nums={request_nums})')
print(f'Bug confirmed: {result != 2}')
"

Called with request_number=5, dp_size=3
Expected: 5 % 3 = 2
Actual: 1 (uses global request_nums=10)
Bug confirmed: True

Race condition in write worker thread initialization

Medium Severity

In ensure_worker_started, there's a TOCTOU (time-of-check-time-of-use) race condition. The check if self._write_worker_started and the flag assignment self._write_worker_started = True both happen outside the lock. Multiple threads could pass the initial check simultaneously, then each would set the flag and proceed to acquire the lock and start separate worker threads. The flag should be checked and set inside the lock with a double-check pattern.

Infinite loop when no requests need saving

High Severity

In save_kv_layer, when metadata.reqs_to_save is empty, remote_engine_id remains None from its initialization. The while loop at line 1249 then checks remote_engine_id not in self.write_ready_flags, which is always True when remote_engine_id is None. Combined with an empty _ready_requests queue, the first condition is satisfied indefinitely, causing an infinite spin loop that will hang the entire system. Unlike start_load_kv which has a wait_handshake_readd_req guard, this function lacks such protection.

len() called on dictionary instead of list

Medium Severity

The expression len(prefill_instance_endpoint) is called on a dictionary (the endpoint configuration dict containing keys like "dp_size", "tp_size", etc.), not on the list of prefill instances. This returns the number of dictionary keys (around 7-8) rather than the intended number of prefill instances. The load balancing calculation is therefore incorrect. This should likely be len(prefill_instances) to get the actual instance count.

# Simulating the proxy server data structures
prefill_instance_endpoint = {
    "dp_size": 2,
    "tp_size": 4,
    "request_address": "http://10.0.0.1:20005/v1/completions",
    "handshake_port": 6301,
    "notify_port": 61005,
}
prefill_instances = [prefill_instance_endpoint]  # List with 1 instance

request_nums = 100
# Bug: uses dict length instead of list length
result = request_nums // len(prefill_instance_endpoint)
correct_result = request_nums // len(prefill_instances)

print(f"len(prefill_instance_endpoint) = {len(prefill_instance_endpoint)} (dict keys)")
print(f"len(prefill_instances) = {len(prefill_instances)} (actual instances)")
print(f"Bug result: {result}, Correct result: {correct_result}")

python3 -c "
prefill_instance_endpoint = {'dp_size': 2, 'tp_size': 4, 'request_address': 'http://10.0.0.1:20005', 'handshake_port': 6301, 'notify_port': 61005}
prefill_instances = [prefill_instance_endpoint]
request_nums = 100
result = request_nums // len(prefill_instance_endpoint)
correct_result = request_nums // len(prefill_instances)
print(f'len(prefill_instance_endpoint) = {len(prefill_instance_endpoint)} (dict keys)')
print(f'len(prefill_instances) = {len(prefill_instances)} (actual instances)')
print(f'Bug result: {result}, Correct result: {correct_result}')
"

len(prefill_instance_endpoint) = 5 (dict keys)
len(prefill_instances) = 1 (actual instances)
Bug result: 20, Correct result: 100

Missing key check causes KeyError for non-transfer requests

High Severity

In build_connector_meta, the code iterates over all requests in scheduled_cached_reqs.req_ids (which includes ALL running and resumed requests), then directly accesses self._reqs_need_pending_save[req_id] without checking if the key exists. However, _reqs_need_pending_save only contains requests that have do_remote_decode=True and are in chunked prefill mode. For any regular request without KV transfer enabled, this will crash with a KeyError. The code needs to check if req_id in self._reqs_need_pending_save before accessing it.

Function ignores parameter and uses global variable instead

Medium Severity

The function example_round_robin_dp_loader accepts a parameter request_number but completely ignores it, instead using the global variable request_nums. This makes the function's parameter meaningless and could cause incorrect round-robin behavior when the caller expects the passed value to be used.

Using len() on dict instead of list for sizing

High Severity

The code calls len(prefill_instance_endpoint) where prefill_instance_endpoint is a single dict (one prefill instance), not the list of all instances. This returns the number of keys in the dict rather than the number of prefill instances. The intent appears to be len(prefill_instances) to distribute requests across multiple prefill instances.

Variable undefined when UnicodeDecodeError is raised

High Severity

If msg.decode("UTF-8") on line 527 raises a UnicodeDecodeError, the variable msg_str is never assigned. The except block on line 532 attempts to log msg_str, and line 534 also references it. This causes a NameError to be raised instead of the intended MoRIIOError, masking the actual problem and potentially crashing the message handling thread.

Variables undefined when handshake future already exists

High Severity

The variables host, port, tp_size, and remote_dp_size are only assigned inside the if fut is None block (lines 1020-1024), but they are used unconditionally on lines 1036-1039. If a handshake future already exists in _handshake_futures (i.e., fut is not None), these variables will be undefined, causing a NameError when the code attempts to iterate range(remote_dp_size) or submit the handshake.

Race condition in worker thread startup check

Medium Severity

The ensure_worker_started method checks and sets _write_worker_started (lines 95-97) before acquiring _write_worker_lock. If two threads call this method concurrently, both could pass the initial check before either sets the flag to True, resulting in multiple background worker threads being started. The flag check and set should occur inside the lock.

Wrong variable used for cache size calculation

High Severity

The line uses cache (a leftover variable from the previous nested loop at lines 1103-1109) instead of kv_cache (the current loop variable from line 1111). This causes local_kv_cache_size to be populated with the size of the last cache from the earlier loop for every layer, rather than each layer's actual cache size. This could lead to incorrect size calculations and potential data corruption during KV cache transfers.

Are these libraries needed in runtime or just in the build phase?

runtime