[Draft, no merge] MVP for vLLM Disagg

[chunfangamd]

We prototype the PD Disagg on DeepSeek. So far, we've done

• Framework Integration: Successfully integrated vLLM’s Prefill/Decode disaggregated architecture into the InferenceX framework.
• MoRI/Deep EP Integration on DPSK V3: for MoE All-to-All and for KV Cache transfer.
• Confirmed MoRI-IO outperforms the standard NixlConnector in TTFT (Time to First Token).
• Stability & Bug Fixes: * Resolved hang issues at high concurrency (CONC512) by fixing KV cache leaks and optimizing the "Reaper" logic for memory block release.
• Fixed hardware compatibility issues, including PCI topology failures on Broadcom switches and error handling for mlx5 NICs.
• Cluster Validation: Verified the multi-node deployment recipe on the SA-9N (mia1) cluster.

co-authors:

• @ichbinblau
• @ChuanLi1101
• @billishyahao
  Thanks to @functionstackx

[github-actions]

Thanks for the contribution! For vLLM & SGLang, please ensure that your recipes is similar to the official vLLM recipes and/or the SGLang cookbook

• https://cookbook.sglang.io/
• https://docs.vllm.ai/projects/recipes/en/latest/

If it is not, please create a PR first before we can merge your PR into the master branch. Let's ensure that the documentation is first class such that the entire ML community can benefit from your hard work! Thank you

• https://github.com/vllm-project/recipes
• https://github.com/sgl-project/sgl-cookbook

[claude]

🔴 In handle_request(), req_data["max_tokens"] -= 1 is unconditional but the proxy registers both /v1/completions and /v1/chat/completions routes. Chat completion clients use max_completion_tokens instead of max_tokens, so any chat completions request without a max_tokens key raises a KeyError, returning a 500 error for every chat completions request routed through this proxy. The fix should mirror send_request_to_prefill(), which already handles both fields correctly.

Extended reasoning...

What the bug is and how it manifests

In moriio_proxy.py, the handle_request() function is registered for both /v1/completions and /v1/chat/completions routes. When forwarding a request to the decode worker, line 269 does:

req_data["max_tokens"] -= 1

This assumes max_tokens is always present in req_data. However, the OpenAI Chat Completions API uses max_completion_tokens as the field name, not max_tokens. A client calling /v1/chat/completions that omits max_tokens (the normal case, since it would instead set max_completion_tokens) will trigger a KeyError here.

The specific code path that triggers it

The route decorator:

@app.route("/v1/completions", methods=["POST"])
@app.route("/v1/chat/completions", methods=["POST"])
async def handle_request():

Directs both endpoint types to the same handler. Inside, before building kv_transfer_params for the decode worker, the code executes req_data["max_tokens"] -= 1 unconditionally. A chat completion payload like {"messages": [...], "max_completion_tokens": 512, "model": "..."} has no max_tokens key, causing a KeyError.

Why existing code does not prevent it

The try/except Exception block surrounding handle_request() catches the KeyError and returns a 500 Internal Server Error response — so the server stays alive, but every chat completions request fails. The bug is entirely silent from the server-operator perspective unless logs are checked.

What the impact would be

The /v1/chat/completions endpoint is advertised and registered but completely non-functional for any client using standard chat completions semantics (i.e., sending max_completion_tokens without also sending max_tokens). This is the standard OpenAI SDK behavior for chat completions.

How to fix it

The fix is to mirror what send_request_to_prefill() already does correctly:

# send_request_to_prefill — correct handling:
req_data_copy["max_tokens"] = 1
if "max_completion_tokens" in req_data_copy:
    req_data_copy["max_completion_tokens"] = 1

For the decode path, the equivalent fix would be:

if "max_tokens" in req_data:
    req_data["max_tokens"] -= 1
if "max_completion_tokens" in req_data:
    req_data["max_completion_tokens"] -= 1

Step-by-step proof

1. Client sends POST /v1/chat/completions with body {"model": "...", "messages": [{"role": "user", "content": "hello"}], "max_completion_tokens": 100}
2. handle_request() is invoked; req_data = {"model": ..., "messages": [...], "max_completion_tokens": 100}
3. Execution reaches line 269: req_data["max_tokens"] -= 1
4. "max_tokens" is not a key in req_data → KeyError: max_tokens
5. Outer except Exception as e catches it and returns (f"Internal Server Error: {e!s}", 500)
6. Client receives a 500 error; all chat completions requests through the proxy fail

The fact that send_request_to_prefill() explicitly handles both max_tokens and max_completion_tokens confirms the author was aware of the dual field naming — the decode path was simply not updated consistently.

[claude]

🟡 In setup_deps.sh, the if-branch that logs "All dependencies already present" prints the raw Unix epoch timestamp (~1748000000) as the wallclock time instead of the elapsed seconds. The else-branch correctly computes the elapsed time with $(( _SETUP_END - _SETUP_START ))s; the if-branch should do the same.

Extended reasoning...

What the bug is and how it manifests

At the very end of setup_deps.sh, the script records _SETUP_END=$(date +%s) and checks whether any components were newly installed. If nothing was installed (the common fast path on container re-use), it prints:

[SETUP] All dependencies already present (${_SETUP_END}s wallclock)

_SETUP_END is a Unix epoch timestamp — a large integer like 1748000000. The message therefore reads something like (1748000000s wallclock) instead of the intended elapsed time (e.g., (2s wallclock)).

The specific code path that triggers it

The final block of setup_deps.sh:

_SETUP_END=$(date +%s)
if [[ ${#_SETUP_INSTALLED[@]} -eq 0 ]]; then
    echo "[SETUP] All dependencies already present (${_SETUP_END}s wallclock)"  # BUG: should be elapsed
else
    echo "[SETUP] Installed: ${_SETUP_INSTALLED[*]} in $(( _SETUP_END - _SETUP_START ))s"
fi

The else branch correctly subtracts _SETUP_START from _SETUP_END; the if branch is missing that arithmetic.

Why existing code does not prevent it

There is no test for the log output and the script runs without error — the echo simply emits an incorrect number. The else branch was written correctly, suggesting the if branch was a copy-paste oversight.

Step-by-step proof

1. Script starts at Unix epoch T=1748000000, so _SETUP_START=1748000000.
2. All dependencies are already present; _SETUP_INSTALLED remains empty.
3. Script reaches the final block 2 seconds later: _SETUP_END=1748000002.
4. ${#_SETUP_INSTALLED[@]} -eq 0 is true, so the if branch executes.
5. Actual output: [SETUP] All dependencies already present (1748000002s wallclock) — clearly wrong.
6. Expected output: [SETUP] All dependencies already present (2s wallclock).

Impact

Purely cosmetic — no functional impact. However, the wallclock message is the primary timing signal operators use to verify the fast-path completed quickly. Displaying ~1.7 billion seconds makes the log useless for that purpose.

Fix

Replace ${_SETUP_END}s with $(( _SETUP_END - _SETUP_START ))s in the if-branch to match the else-branch.

[claude]

🟡 In _listen_for_register(), the guards data["request_address"] not in prefill_instances and data["request_address"] not in decode_instances are always True because they compare a string against a list of dicts — a string can never equal a dict in Python. The guards are dead code that provide no actual filtering. Fix by checking data["request_address"] not in [d["request_address"] for d in prefill_instances], or remove the redundant guards entirely since _append_whole_dict_unique() handles deduplication correctly.

Extended reasoning...

Bug: Type mismatch makes deduplication guards always True

data["request_address"] is a string (e.g., "http://192.168.1.1:2584"), while prefill_instances and decode_instances are declared as list[dict] (line ~35: prefill_instances: list[dict] = []). Python’s in operator for lists checks element equality via ==; since a string is never equal to a dict, data["request_address"] not in prefill_instances evaluates to True unconditionally, making the guard a no-op.

The same issue applies to both the prefill branch and the decode branch:

and data["request_address"] not in prefill_instances  # always True
...
and data["request_address"] not in decode_instances   # always True

The intent of the guard is clearly early-exit deduplication — skip calling _append_whole_dict_unique if an instance with the same address is already registered. However, because the type comparison is wrong, the guard never fires and _append_whole_dict_unique() is always called.

Functionally there is no impact because _append_whole_dict_unique() performs correct deduplication by comparing filtered dicts. As long as that function remains in place, no duplicates are registered. However, the guard is semantically broken dead code: a reader would reasonably assume it provides an O(n) fast path that skips the lock acquisition and full dict comparison, but it does not.

The risk materializes if _append_whole_dict_unique is ever removed or refactored away under the assumption the guard provides the deduplication — in that case, duplicate instances would be silently registered, causing requests to be routed to stale/duplicate endpoints.

Step-by-step proof:

1. A worker calls into the ZMQ router with data = {"type": "register", "role": "P", "request_address": "http://10.0.0.1:2584", ...}
2. prefill_instances at this point is [{"type": "register", "role": "P", "request_address": "http://10.0.0.1:2584", ...}] (previously registered)
3. Python evaluates "http://10.0.0.1:2584" not in [{"type": "register", ...}]
4. Python checks: "http://10.0.0.1:2584" == {"type": "register", ...} → False
5. Therefore not in → True, and the branch body executes — calling _append_whole_dict_unique again for the duplicate
6. _append_whole_dict_unique detects the duplicate and returns False without appending, saving the day

Fix: Replace the broken guard with a correct one:

and data["request_address"] not in [d["request_address"] for d in prefill_instances]

Or simply remove the guard entirely since _append_whole_dict_unique() already handles deduplication correctly under the lock.

[functionstackx]

should we use vllm router instead here? that seems like the more production ready version than this toy proxy?