[Bug-fix] Measure KV transfer speed per actual transfer for disaggregation (Mooncake)

[hunhokim]

Motivation

The current implementation of KV-transfer metrics can substantially over-estimate transfer speed whenever a transfer overlaps computation. For example, in a setup with Infiniband NDR NICs, KV transfer speed is observed to be well above 100 GB/s while the line rate is only 50 GB/s. This makes diagnostics on KV transfer unreliable.

Cause

This happens only when KV transfer occurs multiple times on a single request: the recorded byte count spans the whole transfer while the recorded latency covers only the last transfer, resulting in the transfer speed being inflated. Two specific examples:

• Chunked prefill: latency is captured only for the last chunk, but the byte count covers all chunks — a large byte count divided by a small slice of time.
• Early-send of cached prefix: when available, the cached prefix KV is sent to the decode side at the very start of the prefill phase. Latency is measured only over the newly generated tokens transferred after prefill — pairing a full byte count with a partial latency window.

The fix measures time and bytes together, on every actual transfer, so the ratio stays correct no matter how a transfer overlaps computation or splits across chunks.

Implementation

Common interfaces are introduced in disaggregation/common/conn.py with implementations provided for Mooncake.

Inside transfer_worker, each KV transfer is recorded into a per-request _TransferRecord. On completion, get_transfer_metric() pops the record and sources both latency and bytes from it.

If no record exists, it falls back to the old index-derived byte count with None latency (unchanged behavior for backends that don't record, e.g. NIXL).

Some points to consider:

• Teardown: prefill.py snapshots the record onto the new Req.transfer_metric field before the scheduler's clear() drains it.
• Replica factor: the measured path sums bytes across destination ranks and does not re-apply the factor; the fallback path scales by get_kv_replica_factor(). Conflating them would reintroduce the MLA byte over-count.
• Concurrency: a single status_record_lock guards all sites that touch status and the record together (_record_transfer, update_status, clear()), preventing a late chunk from leaking a stale record into the next request reusing the same bootstrap_room.

Unit tests

python -m pytest test/registered/unit/disaggregation/test_kv_transfer_metrics.py

E2E tests

Environments are as follows:

• DeepSeek-V2-Lite used Prefill CP2 and Decode TP2
• Qwen3-8B used TP2 for both Prefill and Decode
• NVIDIA H100 and Infiniband NDR NICs were used

All the values are from prometheus metrics.

Before the patch, transfer speeds show unrealistic values.
After the patch, transfer speeds show reasonable values.

Chunked prefill

prompt_tokens = 8192+1655 (chunked)

Value = KV transfer speed (GiB/s)

Model	Before patch	After patch
DeepSeek-V2-Lite	106	18.5
Qwen3-8B	80	35

Early-send of cached prefix

prompt_tokens = 6554

Value = KV transfer speed (GiB/s)

Model	Before patch	After patch
DeepSeek-V2-Lite	231	19.2
Qwen3-8B	141	35

---

CI States

Latest PR Test (Base): ⏳ Run #29968160092
Latest PR Test (Extra): ⏳ Run #29968159916

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[zcfh]

May I ask why this PR was closed? Was it replaced by a new PR?

[hunhokim]

May I ask why this PR was closed? Was it replaced by a new PR?

I had a private conversation with one of the maintainers on Slack and was told to work on the transfer engine side (e.g. Mooncake) for accurate metrics collection.