Link Management

Function Description

The HCCL collective communication APIs used during NN model execution employ a bilateral communication model, requiring both sides to initiate link establishment simultaneously. In the PD disaggregation solution, this process is simplified to unilateral link establishment initiated by the client. Since dynamic scaling primarily occurs on the Decode side, the Prefill side is defined as the server and the Decode side as the client. Consequently, link establishment follows a procedure where the Decode side initiates link requests to the Prefill side.

Two core APIs are provided: link_clusters and unlink_clusters, both invoked by the Decode side. Link establishment operations are performed on a point-to-point basis.

link_clusters: Establishes links between nodes.
unlink_clusters: Disconnects links between nodes.

Application Scenarios

Link establishment is a prerequisite for KV cache transmission between Prefill and Decode nodes; therefore, links must be established before KV cache transmission is enabled.
In cluster reliability scenarios, if a node on the Prefill or Decode side experiences an anomaly, the faulty node can be taken offline by disconnecting its links, provided that overall cluster availability remains unaffected.
Dynamic adjustment of the PD cluster ratio is achieved through link establishment and disconnection: nodes are added or removed dynamically based on workload, where node addition requires link establishment and node removal requires link disconnection.

Example

The following pseudocode demonstrates the link establishment process between a Prefill node and a Decode node.

Launch the Prefill- and Decode-side scripts, where the LLM-DataDist initialization API is called. The device IP address and port for listening must be set on the Prefill side.

# Script on the Prefill side
from llm_datadist import LLMDataDist, LLMRole, LLMStatusCode, LLMClusterInfo

 # LLM-DataDist initialization
 llm_datadist = LLMDataDist(LLMRole.Prompt, cluster_id=0)
 llm_config = LLMConfig()
 llm_config.listen_ip_info ="192.168.1.1:26000"# local_ip + port
 llm_config.device_id =0  # device_id must match local_ip.
 llm_options = llm_config.generate_options()
 llm_datadist.init(llm_options)

# Script on the Decode side
from llm_datadist import LLMDataDist, LLMRole, LLMStatusCode, LLMClusterInfo

 # LLM-DataDist initialization
 llm_datadist = LLMDataDist(LLMRole.DECODER, cluster_id=0)
 llm_config = LLMConfig()
 llm_config.device_id =0 
 llm_options = llm_config.generate_options()
 llm_datadist.init(llm_options)

In the Decode-side script, call link_clusters to initiate link establishment. When the service exits, call unlink_clusters on the Decode side to disconnect the link.

 # Generate cluster information for link establishment.
 cluster = LLMClusterInfo()
 cluster.remote_cluster_id = 1  # remote_cluster_id must match the LLM-DataDist instance created on the Prefill side.
 cluster.append_local_ip_info("192.168.2.1", 26000) # local_ip_info indicates the IP address of the local device for link establishment.
 cluster.append_remote_ip_info("192.168.1.1", 26000) # remote_ip_info indicates the IP address of the peer device for link establishment.

 # Call link_clusters to establish a link.
 # ret is the return value of the API, and rets indicates the link establishment results of each cluster.
 ret, rets = llm_datadist.link_clusters([cluster], timeout=5000)
 # Check the link establishment result.
 if ret != LLMStatusCode.LLM_SUCCESS:
     raiseException("link failed.")
 for cluster_i in range(len(rets)):
     link_ret = rets[cluster_i]
     if link_ret != LLMStatusCode.LLM_SUCCESS:
         print(f"{cluster_i} link failed.")

On the Decode side, call check_link_status to quickly verify if data transmission over the link is normal.

try:
    llm_datadist.check_link_status(remote_cluster_id=1)
except LLMException as ex:
    print(f"check_link_status exception:{ex.status_code}")
    raise ex

After service termination, the Decode side disconnects the link. Both the Prefill and Decode sides call finalize of llm_datadist to release resources.

 # Script on the Prefill side
 # Call llm_datadist to allocate the KV cache.
 # Perform service inference.
 # ...

 # Call this API on service exit after the Decode side disconnects the link.
 llm_datadist.finalize()

 # Script on the Decode side
 # pull_cache and model inference
 # ...

 # Call unlink_clusters to disconnect the link on service exit.
 ret, rets = llm_datadist.unlink_clusters([cluster], timeout=5000)
 if ret != LLMStatusCode.LLM_SUCCESS:
     raiseRuntimeError(f'[unlink_cluster] failed, ret={ret}')
 llm_datadist.finalize()

To add a node or bring an offline node back online, the same procedure must be followed. To remove a node, under normal circumstances, the Decode side should actively call unlink_clusters. If the Decode side cannot call unlink_clusters, the Prefill side must call unlink_clusters instead.

By calling link_clusters for node addition and unlink_clusters for node removal, you can flexibly scale the distributed cluster dynamically.

Troubleshooting

If an exception occurs on the Decode side that prevents it from calling unlink_clusters, the Prefill side must call unlink_clusters to clean up resources; otherwise, link establishment cannot be performed again.

The unlink_clusters API provides a forcible disconnection capability, which is useful in case of link faults where normal disconnection may take a long time. When the forcible disconnection option is enabled (by setting force=True), the call must be initiated on both sides, and only the local link is cleaned up.

# Forcibly disconnect the link.
 ret, rets = llm_datadist.unlink_clusters([cluster], timeout=5000, force=True)
 if ret != LLMStatusCode.LLM_SUCCESS:
     raise RuntimeError(f'[unlink_clusters] failed, ret={ret}')

Parent topic: Functions