Topology Aware Routing to lower KV Cache transfer latency

[RishabhSaini]

Problem

Prefill and decode pods are currently selected independently without network locality awareness. KV cache transfer between pods on different nodes incurs inter-node network latency (infiniband upto 400 GB/s)
vs. intra-node transfer (NVLink upto 900GB/s), directly impacting TTFT SLO compliance.

Proposal

Add topology-aware prefill pod selection:

• After selecting decode pod, prefer prefill pods on same node (giving them a higher score)
• Fallback to global selection if non-colocated pods still have a higher score

[elevran]

xref #8

[elevran]

You can use existing k8s labels for hostname, zone, and region.
We already collect labels for pods.

[RishabhSaini]

/assign

[elevran]

/assign @vishbhat

[RishabhSaini]

You can use existing k8s labels for hostname, zone, and region. We already collect labels for pods.

@elevran
Could you point to me in code where are these node, zone, regions labels are set in modelService pods? I can only seem to find nodeName in the spec of the modelService?

Currently the solutions include:

1. Query nodeName from the modelService pods and add it to the EndpointMetdata. After the Decode Profiler runs, the Prefill Profiller can look at the cycleState to find the nodeName of the decode pod. This avoids making any K8S API request to find the node, region, zone. However, only nodeName is used to make scoring decision and not region and zone.

2. Scorer can identify the nodeName from the modelservice pod, and make K8S API requests to find the zone and region using the known nodeName. Adds latency since querying the K8S API

3. Populate the labels of the ModelService pods with the node, zone, region. The scorer then just uses them to make scoring decisions. However, existing helm based CRDs cannot achieve this. It requires a K8S controller to populate the pods acording to the nodes they are assigned to.

My suggestion is 1, where we just use Node name to give scores based on colocation

[elevran]

A scorer should not make runtime calls to the k8s API. It is expected to operate on existing data only to keep scheduling overheads minimal.

The node, zone and region are Pod labels. We already store labels for the Pod so retrieving those (if defined) should be trivial.
These are (presumably) standard k8s labels, independent of any ModelServer abstraction we may have.
The scorer shuold set higher P/D scores based on location affinity. If co-location of P and D nodes can not be determined the scorer should given a zero score to all Pods equally.

Therefore, my suggestion is to implement the scorer based on existing labels when available and assign zero score when not available.
If multiple labels exist (node, zone and region) give higher marks to closer entities (e.g., 1.0 for same node, 0.5 for same zone and 0.25 for same region).

[RishabhSaini]

llm-d does not currently store the hostname, zone and region labels at the Pod level