design: investigate KV cache transfer latency as a predictor feature for P/D mode

[madhugoutham]

Summary

Investigate whether KV cache transfer metrics can improve latency predictions in prefill/decode disaggregated deployments. This is a design investigation, not an implementation proposal.

Problem

In P/D disagg mode, TTFT includes KV cache transfer time between prefill and decode pods. Transfer latency varies by interconnect type. The predictor has no features for transfer latency and predicts the same TTFT regardless of interconnect.

Design concern

Decode pods do not always pair with the same prefill pod. Transfer latency varies per prefill-decode pair depending on network topology. Training on per-pod average transfer latency when routing varies which prefill is chosen may introduce noise that makes predictions worse rather than better.

Investigation needed

1. Verify which NIXL metrics vLLM currently exposes and in what format
2. Determine whether decode pods consistently pair with specific prefill pods
3. Assess whether per-pod average transfer latency is meaningful despite pairing variance
4. Measure actual TTFT variance attributable to transfer in production P/D deployments

Note: Transfer latency ranges (sub-millisecond for NVLink to tens of milliseconds for TCP) are simulation estimates, not production measurements. Prefill-decode pairing varies per request, so per-pod average transfer latency may be a noisy signal.

[gkneighb]

I'd like to take this investigation on. Mapped it against the code to scope what's empirical vs. measurable, and landed on a hypothesis worth stating up front.

Lay of the land (from the code): the predictor already carries P/D-aware features -- pod_type_cat and engineered ones like prefill_tokens_in_flight -- but nothing for transfer. Transfer itself lives in the sidecar (pkg/sidecar/proxy/connector_nixlv2.go, NIXL V2 with a prefillPodHostPort), and the EPP's prefix_based_pd_decider selects the prefill per request, so a decode pod observes a mixture of prefill sources -- which is exactly the noise concern you flagged.

Splitting the four questions by type:

• Q1 (NIXL metrics) / Q2 (pairing stickiness) are empirical -- Q1 from vLLM + the nixlv2 connector, Q2 essentially answered by the per-request decider (confirmation welcome if there's any stickiness I'm missing).
• Q3 (is per-pod-avg meaningful?) / Q4 (TTFT variance from transfer) are a signal-vs-noise / variance-decomposition question, which is simulatable.

The hypothesis I'd test: per-pod average transfer latency is the wrong feature because it marginalizes over the prefill mixture -- its usefulness collapses as a function of (between-pod variance / within-pod pairing variance) and routing stickiness. But the EPP knows which prefill it picked before dispatch, so the right feature is the per-(chosen-prefill, decode) pair estimate -- a precise conditional available at predict time, not a noisy marginal. I'd expect the marginal feature to degrade predictions under heterogeneous interconnect + non-sticky routing, while the conditional one helps.

Method: a synthetic P/D fleet (N prefill x M decode, per-pair transfer drawn from an NVLink->TCP topology, a routing-stickiness knob), injecting transfer into TTFT, then comparing predictor quality (coverage / quantile-loss) with no transfer feature vs. per-pod-average vs. per-pair conditional. That gives data-backed answers to Q3/Q4 and a clear "when does this feature pay for itself" boundary.

Does this framing match the intent, and is there any production signal on the variance ranges (NVLink vs TCP) I should anchor the sim to rather than guessing? I'll follow up with the simulation results.

[gkneighb]

Followed up the plan with the simulation. Headline: a transfer feature helps, but the interesting question turned out to be which form, and the answer flips depending on whether the fleet is stable or autoscaling. It also walked back my own first instinct, so I'll show the steps.

Setup. Synthetic P/D fleet (8 prefill x 8 decode). Transfer is modeled physically -- overhead[p][d] + inv_bw[p][d]*(input_tokens/1000) + jitter -- where both endpoints' placement contributes (each pod has a fast/slow tier), so it's a genuine per-pair quantity, not a per-prefill constant. Routing is per-request (the llm-d case). I train an XGBoost p90 model (mirroring the predictor) and compare feature forms, 5 seeds:

• A -- compute features only
• A+IDs -- A + prefill_id/decode_id one-hots (lets the model learn the transfer matrix implicitly -- the honest competitor to an explicit feature)
• B -- A + per-decode-pod average transfer (the marginal; your noise concern)
• C_est -- A + per-(prefill,decode) transfer estimated from 5 samples/pair (realistic, not an oracle)
• C_truth -- oracle ceiling

Result 1 -- stable fleet (all pairs trained): pinball, lower=better

A	A+IDs	C_est	C_truth
1.005	0.707	0.709	0.667

Transfer is worth ~30% pinball -- but A+IDs ~= C_est. With stable pods you don't need an explicit transfer feature; pod-ID one-hots learn the matrix themselves. B lands at 0.865 -- it only captures the per-decode-pod component and misses the prefill-side variation that per-request routing scrambles, exactly the concern you raised.

Result 2 -- autoscaling churn (2/8 decode pods held out of training):

variant	warm-test	cold (churned) test
A	0.998	1.182
A+IDs	0.725	1.485 +/- 0.76
C_est	0.728	0.927 +/- 0.23
C_truth	0.691	0.677

On unseen pods, the one-hot approach collapses -- to worse than no feature at all, because a fresh decode_id column has no learned value and the model misfires. The estimated transfer feature generalizes (0.927), because it encodes transfer in a pod-agnostic way the model extrapolates from. And it's cheap -- 5 samples/pair ~= oracle.

So, answering Q3/Q4:

1. Transfer is a material TTFT signal under heterogeneous interconnect.
2. Per-pod-average is the wrong feature under per-request routing -- confirmed with numbers.
3. In a stable fleet, pod-ID one-hots match an explicit feature -- you don't strictly need one.
4. But llm-d fleets autoscale, and that's where an explicit per-(chosen-prefill, decode) transfer feature earns its keep -- one-hots can't value a new pod; an estimated transfer feature can. That's the case for building it.

Caveats, stated honestly:

• This conclusion is contingent on transfer having a per-(churned-pod) component. My first topology was prefill-dominated, which made decode churn a non-event and gave the opposite (misleading) read -- I fixed it to a both-endpoints model, which is more realistic but is an assumption worth validating against real topology.
• This establishes the decision structure; the production magnitudes (Q1 NIXL metrics, Q2 real pairing/variance ranges) still need empirical measurement -- the sim doesn't substitute for that. If anyone has NIXL transfer distributions from a real P/D deployment, I'd anchor the het/variance params to them.

Repro harness is small and self-contained (the three feature regimes + cold-start mode); happy to share.

[madhugoutham]

Thanks for taking this on. The simulation is well done and the self-correction on the topology model (prefill-dominated to both-endpoints) is the kind of rigor this investigation needs.

One thing the simulation doesn't account for

The latency predictor runs during runDataProducerPlugins (director.go:294), before scheduler.Schedule (director.go:305) picks any pods. At prediction time, neither prefill nor decode has been selected, so the per-(chosen-prefill, decode) pair estimate that C_est relies on is not available as a real-time input. The pair identity only becomes known after scheduling completes, making it useful for training data enrichment but not for live scoring.

Some observations worth checking (may affect scope)

• The sidecar (connector_nixlv2.go) records stage durations as OTel trace spans, not Prometheus metrics. We haven't confirmed whether a pipeline exists to feed this data into the predictor's training ingestion, but it's worth verifying.
• The predictor targets p90 quantile (QUANTILE_ALPHA=0.9), not mean. With 5 samples per pair, the p90 estimate has high variance, especially for non-Gaussian network latency.
• On InfiniBand clusters (common for P/D), transfer variance is microseconds against TTFT of 50 to 500ms. The signal becomes material mainly on multi-DC deployments with TCP fallback.
• llm-d/llm-d-router#1678 (topology-extractor) is building related infrastructure for locality-aware routing that may be worth coordinating with.

Some ideas for working within the current pipeline

• Enrich training data retrospectively with pair identity and transfer duration after each request completes, so the model learns transfer patterns over time without any pipeline changes.
• Use Kubernetes node topology labels (rack, leafgroup, data hall) to estimate transfer cost at prediction time, since those are always available and don't have a cold-pair problem.
• Narrow to top-K prefill candidates first, then score (prefill, decode) pairs with transfer estimates to recover near-global optimality.

Happy to have you drive this. @kaushikmitr would appreciate your thoughts on feasibility and priority.

[gkneighb]

Thanks -- and the timing correction is the important one. The predictor running in runDataProducerPlugins before scheduler.Schedule means the chosen (prefill, decode) pair isn't known at predict time, so C_est as I framed it can only be training-data enrichment, not a live feature. That revises the recommendation -- appreciate the catch.

The topology-labels idea is better than what I proposed, and my own cold-start result actually argues for it. My sim's main finding was that under pod churn an explicit transfer signal beats pod-ID one-hots because it generalizes to unseen pods. Topology labels (rack/leafgroup) are that generalizable encoding taken further: a new pod inherits its topology class's transfer characteristics with zero cold-pair history and no measurement pipeline -- strictly better than my per-pair estimate on both the cold-start and plumbing axes. Pooling samples across a topology class also directly answers your p90/5-samples variance point (far more samples per class than per pair).

So I'd revise the direction to: topology-derived transfer as the signal, used as (a) retrospective training enrichment now (pair + duration after completion, as you suggested -- no pipeline change), and (b) a candidate-time feature if the predictor is ever invoked with a known candidate during scoring. Which is the one architectural thing I'd want to confirm: is latency prediction only the pre-scheduling data-producer call, or is it also evaluated per-candidate during scoring (where the candidate's topology label would be available)? That determines whether (b) is even reachable.

Two of your other points scope the whole thing: if transfer is microseconds on InfiniBand and only material on multi-DC/TCP, the feature only earns its keep on those deployments -- worth gating any work on whether that's a target. And coordinating with #1678 (topology-extractor) makes sense if it's already surfacing the labels.

Happy to drive whichever slice is useful. @kaushikmitr -- I'd defer to you on feasibility and priority, especially whether the multi-DC/TCP case is common enough to warrant this over other predictor work.