[RFC]: Fix for NIXL metadata exchange issue when the worker is multi-node

[GuanLuo]

Motivation.

This issue has been long identified and there are several PRs trying to address it. But those PRs don't get pushed to the finish line and would like to start yet another discussion to summarize the issue and propose a fix. Note that this fix doesn't intend to fix the other parallelism issues around NIXL (i.e. #22430), but only focus on the metadata exchange:

The issue is described in #19080 which is an attempt to fix the issue. In summary, the NIXL handshake is worker-to-worker and the decode work will try to reach the prefill worker based on the "remote host/port" in connector metadata. However, the "remote host" is the address of the head node, and therefore, for the prefill workers on the other nodes, the correct host address is not passed to decode worker and thus the handshake will fail. So we want to

• centralize the handshake metadata from all workers to the head node engine, so that any downstream worker can obtain desired metadata from the single "remote host" pointing to the head node.

#19080 addressed this but it encounters certain edge cases at the time and suggest that the centralization should be done as a separate step from the NIXL connector initialization. One of the edge case is that the rank 0 worker may not be initialized on the head node so the worker can't make assumption on whether it should nominate itself to be the metadata exchange listener.

To address this, #22274 shows a good approach to introduce a RPC to collect the metadata from all workers during engine initialization and later creates the listener for hosting the metadata. However #22274 creates the listener outside of vLLM engine which results in long propagation chain and implementation leak (the HTTP server needs to know about NIXL connector internals).

This proposal takes the idea from #22274 and tries to limit the change to be within EngineCore and NixlConnector (still use zmq for metadata exchange).

Proposed Change.

Add RPC get_kv_connector_handshake_metadata() in executor to collect metadata from all workers as shown in #22274

# Collect and store KV connector xfer metadata from workers
# (after KV cache registration)
self.xfer_handshake_metadata = (
    self.model_executor.get_kv_connector_handshake_metadata())

Instead of propagating it all the way up, feed it back to kv_cache_config for scheduler initialization. In this way, the metadata will be passed to NixlConnectorScheduler and the listener can be managed by the connector scheduler as proposed here.

With this modification, the user interaction remains the same and the NIXL handshake behavior is contained within the NIXLConnector (similar to #19080).

Feedback Period.

No response

CC List.

No response

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[wseaton]

Hi @GuanLuo, thanks for writing this up!

One of the rationales for #22274, and why I'm interested in continuing work on it is that the side-channel for kvconnector configuration ideally needs to be a protected interface, and there is concern about a potential issue with confused deputy attacks on that port. Pushing the KV connector handshake to North<->South through an API Server instead of East/West allows for the introduction of TLS (and potentially mTLS, provided by something like Istio) in a cluster type deployment. It's difficult to secure ZMQ in this way, maybe there are some potential mitigations but I'm not completely sure.

#22274 was also implemented in a way where the handshake strategy is backwards compatible, so I think there's a case to be made to simplify the sharing as you describe above in the P2P case, but also still support a North to South method as well.

[NickLucche]

Thanks for the RFC @GuanLuo !

However, the "remote host" is the address of the head node, and therefore, for the prefill workers on the other nodes, the correct host address is not passed to decode worker and thus the handshake will fail.

The system is relying on a inference gateway/load balancer to have knowledge about every P/D instance address, and those addresses are the ones being passed around for initiating the handshake.
Just to clarify, this is working as intended in multi-node.

Instead of propagating it all the way up, feed it back to kv_cache_config for scheduler initialization.

I don't think this would work as we need to have kv caches allocated for the metadata to be ready, which happens after scheduler init. Regardless, the same idea could be implemented differently.
However, I feel like at this point we should be looking for a more long-term solution which meets the requirements of a broader audience, taking into account security and keeping the cluster deployment as a pre-requisite as @wseaton explained.
i.e I believe we should assume vLLM to be a piece in a more complex deployment and avoid centralizing all parts of the flow to work independently of this system.

[GuanLuo]

The system is relying on an inference gateway/load balancer to have knowledge about every P/D instance address, and those addresses are the ones being passed around for initiating the handshake.
Just to clarify, this is working as intended in multi-node.

@NickLucche IIUC, here you are referring to a deployment of multiple P/D instances (an instance is a replica of the model). Then yes I agree that it is working as intended if the instance doesn't take more than one node.

Maybe my wording is not precise, let me write down my current understanding to make sure I didn't make some wrong assumptions. The "prefill workers" I referred to are the NIXL workers for each TP (/ DP) rank that collaboratively work as one instance. An example is if I have a model that set TP=2 and I only have two 1-GPU nodes, I have two NIXL workers, worker 0 for TP rank 0 on node 0, and worker 1 for rank 1 on node 1. Then the NIXL transfer issue emerges:

• When prefill finished, the KV transfer parameters returns from request_finished contains remote_host that is node 0 address.
• To start the KV transfer, decode needs to have NIXL agent to the prefill NIXL workers, which will contact remote_host to obtain the worker's metadata (NixlAgentMetadata).
• In the current implementation, worker 1's metadata is not obtainable through node 0 address (worker 1's listener is on node 1) and P/D disaggregation will fail. Both PRs mentioned in this RFC have addressed this by introducing a collective step in the initialization, and make all metadata available in remote_host. @wseaton 's solution is more robust as the other PR has worker init assumption that may not be true (detail)

However, I feel like at this point we should be looking for a more long-term solution which meets the requirements of a broader audience, taking into account security and keeping the cluster deployment as a pre-requisite as @wseaton explained.
i.e I believe we should assume vLLM to be a piece in a more complex deployment and avoid centralizing all parts of the flow to work independently of this system.

I agree. While I believe #22274 will solve the NIXL transfer issue, we have seen users hitting this issue in their deployments and blocked by it. We would love to have the issue resolved as soon as possible. I can make changes in our vLLM integration to accommodate #22274 if it is ready to be merged. But I am not sure about the current status of the PR, and thus proposed this RFC to fix the immediate issue with smaller code change while taking iterations on the long-term solution.

[GuanLuo]

I created a draft PR #26338 of this RFC for look and feel

[NickLucche]

Apologies for the delay and thanks for clarifying your point.
So, DP ranks are effectively spawned altogether as separate instances, so you should not run into this issue with dp=2 in your example setup. You can find multi-node wide EP examples (DP+EP, 16/32 ranks) in llm-d guides (see setting host here https://github.com/llm-d/llm-d/blob/main/guides/wide-ep-lws/manifests/modelserver/base/decode.yaml#L137C28-L137C40).
TP ranks however currently do assume to share the same host so you would indeed not be able to contact worker1.

I do not feel this use-case is particularly wide-spread for the kind of production deployment we're targeting, but this is absolutely hindering PP and other heterogenous setup, so I am also in favor of moving towards a single source of truth (eg rank0 or even igw) here.

[GuanLuo]

Thanks for pointing out the DP case, you are right that the issue is for TP. In fact the first PR was to solve DS R1 deployment (TP=16 on 2 H100x8) and we didn't push forward because we worked around it by DP+EP.

I will polish my PR to be ready for review since we agree that this is an issue for TP case.

[chaunceyjiang]

@GuanLuo Could you please take a look at this PR? I think it also addresses the issue with multiple nodes: #19591

[GuanLuo]

@chaunceyjiang I took a look at the PR and yes that it is addressing the same issue with multiple nodes. I believe I have tried similar implementation but your PR WAR the executor/scheduler initialization order that I encountered. One minor suggestion is that the metadata proxy may want to wait for all workers' metadata are received and then marked itself ready to complete scheduler initialization. The edge case is that the prefill instance can start handling requests before all worker metadata are received, and the decode instance start making handshakes, which is really unlikely to happen..

Is there a reason that there is no recent progress on that PR? It seems to be in draft for months.

[GuanLuo]

@NickLucche #26338 is ready for review. This PR also fixes incorrect device ID used for NIXL memory registration (009ec48)

[chaunceyjiang]

Is there a reason that there is no recent progress on that PR? It seems to be in draft for months.

When using Ray on multiple nodes, there's an issue that this PR (#24403) is attempting to address. I'm currently waiting for this PR to be merged.
@GuanLuo

[GuanLuo]

I see, I wonder if that is the issue that my PR also addressed. I would still love to have my PR reviewed and merged as it solely focuses on fixing bugs.

This PR also fixes incorrect device ID used for NIXL memory registration (009ec48)