iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[KVConnector][P2P] Configurable unbound-store timeout and one-RTT rejection of a late fetch (#53453)

Browse Source 
Signed-off-by: Liran Schour <lirans@il.ibm.com>
This commit is contained in:
liranschour
2026-09-15 05:38:18 +03:00
committed by GitHub
parent a08dffe215
commit 17bc3e106a
3 changed files with 193 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/features/kv_offloading_usage.md
+22
View File
@@ -203,9 +203,31 @@ Block content hashes must match across instances for peers to exchange blocks (s
| `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. |
| `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. |
| `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. |
| `unbound_store_timeout_s` | no | `60` | Seconds a producer holds stored blocks for a consumer that has not fetched them yet. Raise it for deployments whose prefills outlast the default; the blocks keep primary-tier CPU slots pinned for the whole window. Once it expires a late fetch is rejected in a single round trip, so the consumer falls back to local prefill immediately. |
The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process.
A producer parks a request's blocks until the consumer's `FetchMsg` arrives; if none arrives
within `unbound_store_timeout_s` the blocks are released so they stop pinning primary-tier
slots. Raise it when prefills legitimately take longer than the default:
```bash
vllm serve <model> \
--kv-transfer-config '{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"spec_name": "TieringOffloadingSpec",
"secondary_tiers": [
{
"type": "p2p",
"unbound_store_timeout_s": 180
}
]
}
}'
```
#### Environment Variables
Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence.
tests/v1/kv_offload/tiering/p2p/test_manager.py
+75 -4
View File
@@ -117,6 +117,8 @@ def _make_manager() -> P2PSecondaryTierManager:
mgr._sessions = {}
mgr._kv_to_session = {}
mgr._unbound_stores = {}
mgr._reaped_stores = {}
mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S
mgr._failed_serve_ctxs = []
return mgr
@@ -759,9 +761,9 @@ class TestGetFinished:
assert "req-fresh" in mgr._unbound_stores
def test_unbound_store_reaped_after_timeout(self):
"""Unbound stores past _UNBOUND_STORE_TIMEOUT_S surface as failed
and their kv_request_id lands in _failed_req_ids so a late
FetchMsg/lookup doesn't try to satisfy them."""
"""Unbound stores past the reap deadline surface as failed and their
kv_request_id lands in _reaped_stores so a late FetchMsg is rejected
instead of parking demand nothing can satisfy."""
from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
mgr = self._make()
@@ -779,7 +781,7 @@ class TestGetFinished:
# 2 baseline + 2 buffered stores
assert JobResult(job_id=10, success=False) in results
assert JobResult(job_id=11, success=False) in results
assert "req-stale" in mgr._failed_req_ids
assert "req-stale" in mgr._reaped_stores
def test_submit_store_parks_unbound_batch(self):
"""submit_store on an unbound id appends a batch with a fresh
@@ -1204,6 +1206,75 @@ class TestBidirectionalManager:
assert 201 in b_ok, f"B loads succeeded: {b_ok}"
assert 200 in b_ok, f"B stores succeeded: {b_ok}"
def test_late_fetch_after_reap_fails_immediately(self):
"""A fetch for a reaped kv_request_id fails in one round trip.
The producer drops parked blocks once unbound_store_timeout_s
expires. A consumer whose fetch arrives after that must learn on the
next tick via TransferDoneMsg(success=False) — not stall for
_LOAD_TIMEOUT_S and then take the abort path.
"""
from vllm.v1.kv_offload.tiering.p2p.session.client import _LOAD_TIMEOUT_S
mgr_a, mgr_b = _build_paired_managers()
kv_id = "req-late"
# Record every message type leaving the consumer, so the assertion
# below can prove no abort was ever needed.
sent_from_a: list[str] = []
drain = mgr_a._control._drain_outbound_to
def recording_drain(peer_local_id: str):
out = drain(peer_local_id)
sent_from_a.extend(msg.get("type") for _, msg in out)
return out
mgr_a._control._drain_outbound_to = recording_drain # type: ignore[method-assign]
# Producer parks the blocks, then reaps them before any fetch lands.
mgr_b.submit_store(
_job_metadata(
job_id=200,
keys=[b"b-block"],
chunk_ids=[0],
kv_params={"remote_decoder": {"kv_request_id": kv_id}},
)
)
mgr_b._unbound_store_timeout_s = 0.0
reap_results = list(mgr_b.get_finished_jobs())
assert JobResult(job_id=200, success=False) in reap_results
assert kv_id in mgr_b._reaped_stores
# Consumer now asks for the blocks that no longer exist.
consumer_params = {
"remote_prefiller": {
"kv_request_id": kv_id,
"remote_host": "B",
"remote_port": 2,
},
}
mgr_a.on_new_request(_req_context(consumer_params))
mgr_a.submit_load(
_job_metadata(
job_id=101,
keys=[b"b-block"],
chunk_ids=[0],
kv_params=consumer_params,
)
)
started = time.monotonic()
all_a: list[JobResult] = []
for _ in range(8):
all_a.extend(list(mgr_a.get_finished_jobs()))
list(mgr_b.get_finished_jobs())
elapsed = time.monotonic() - started
assert JobResult(job_id=101, success=False) in all_a, all_a
assert kv_id in mgr_a._failed_req_ids
assert "abort_fetch" not in sent_from_a, sent_from_a
assert elapsed < _LOAD_TIMEOUT_S
# ---------------------------------------------------------------------------
# _accept_new_peers — duplicate connection rejection
vllm/v1/kv_offload/tiering/p2p/manager.py
+96 -10
View File
@@ -47,10 +47,18 @@ logger = init_logger(__name__)
# Reap unbound store batches that have been parked without a FetchMsg
# binding them to a session for longer than this. Protects against the
# prefiller buffering blocks for a decoder that never asks (decoder died,
# network partition, lost kv_request_id). Must be longer than the per-store
# deadline so the store-timeout path fires first for individual jobs.
# network partition, lost kv_request_id). Applies only while a batch is
# unbound: binding hands it to the session, which times the job out under
# _STORE_TIMEOUT_S instead.
# Overridable per tier via the ``unbound_store_timeout_s`` config key.
_UNBOUND_STORE_TIMEOUT_S = 60.0
# How long a reaped kv_request_id is remembered so a fetch that arrives
# after the reap can be rejected immediately. Must exceed the consumer's
# load timeout (session.client._LOAD_TIMEOUT_S): past that the consumer has
# already given up on its own, so there is nobody left to reject.
_REAPED_ID_RETENTION_S = 60.0
# Time we wait during shutdown for inflight transfers to drain via
# cancel(mode="wait") before falling back to mode="immediate". Bounded
# so a wedged peer can't hang shutdown.
@@ -208,6 +216,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
port: int | None = None,
backends: list[str] | None = None,
num_threads: int = 4,
unbound_store_timeout_s: float = _UNBOUND_STORE_TIMEOUT_S,
**kwargs: Any,
) -> None:
"""Initialize the P2P secondary tier manager.
@@ -245,9 +254,29 @@ class P2PSecondaryTierManager(SecondaryTierManager):
num_threads: NIXL agent worker threads for the UCX-only
branch. Ignored when ``backends`` contains a non-UCX
entry.
unbound_store_timeout_s: Seconds a producer holds stored blocks
for a consumer that has not fetched them yet, before
``_reap_unbound_stores`` drops them. Raise it for
deployments whose prefills outlast the default; the blocks
keep primary-tier slots pinned for the whole window.
**kwargs: Reserved for future tier-specific options.
Raises:
ValueError: If ``unbound_store_timeout_s`` is not a positive
number, or anything convertible to one.
"""
super().__init__(offloading_spec, primary_kv_view, tier_type)
try:
timeout_s = float(unbound_store_timeout_s)
except (TypeError, ValueError):
timeout_s = float("nan")
if not timeout_s > 0:
raise ValueError(
f"unbound_store_timeout_s must be a positive number, got "
f"{unbound_store_timeout_s!r}"
)
self._unbound_store_timeout_s = timeout_s
# Block hashes chain from NONE_HASH (see v1/core/kv_cache_utils.py).
# Peers whose seeds differ compute different hashes for identical
# content, so lookups silently miss and no KV crosses the wire. The
@@ -303,8 +332,16 @@ class P2PSecondaryTierManager(SecondaryTierManager):
# kv_request_id → list of batches submit_store'd before any peer
# asked for that id. Drained into a session by _on_session_fetch
# when the corresponding FetchMsg arrives, or surfaced as failures
# by _reap_unbound_stores after _UNBOUND_STORE_TIMEOUT_S.
# by _reap_unbound_stores after _unbound_store_timeout_s.
self._unbound_stores: dict[str, list[_UnboundStoreBatch]] = {}
# kv_request_id → time its parked batches were reaped. Nothing under
# one of these ids can be served any more, so a FetchMsg for it is
# rejected on the spot by _poll_once instead of parking demand until
# the consumer's own load timeout, and submit_store fails rather than
# re-parking under it — a prefill that outran the timeout would
# otherwise keep pinning fresh primary-tier slots, one timeout at a
# time. Entries are pruned after _REAPED_ID_RETENTION_S.
self._reaped_stores: dict[str, float] = {}
self._finished_jobs: list[JobResult] = []
# kv_request_ids that hit a transport/session failure; On load lookup()
@@ -390,7 +427,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
kv_transfer_params; if a session has bound the id, finish it. If
no session has bound the id yet, this is a no-op: parked batches
in `_unbound_stores` are left in place and cleaned up only by
`_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`.
`_reap_unbound_stores` after `_unbound_store_timeout_s`.
"""
source = req_context.get_state(P2PSourceInfo)
dest = req_context.get_state(P2PDestInfo)
@@ -447,6 +484,22 @@ class P2PSecondaryTierManager(SecondaryTierManager):
self._finished_jobs.append(JobResult(job_id=job_id, success=False))
return
# This id's earlier batches were already reaped, so the consumer has
# either given up or will be rejected by the fetch path. Parking more
# would pin primary-tier slots for another full timeout on blocks
# nobody can fetch, so fail the job now — a prefill that outran the
# timeout falls back to local recompute on the decoder either way.
if kv_request_id in self._reaped_stores:
logger.warning(
"P2P %s: submit_store for reaped kv_request_id=%s job_id=%d "
"— failing without parking",
self._local_id,
kv_request_id,
job_id,
)
self._finished_jobs.append(JobResult(job_id=job_id, success=False))
return
# Fast path: a session has already received FetchMsg for this id,
# so we can route the batch straight into its ServerRole.
session = self._kv_to_session.get(kv_request_id)
@@ -707,13 +760,27 @@ class P2PSecondaryTierManager(SecondaryTierManager):
"""Time out submit_store batches that no peer has ever fetched.
Walks `_unbound_stores` for entries whose oldest batch is older
than `_UNBOUND_STORE_TIMEOUT_S`. Drops the kv_request_id, surfaces
every batched job as failed, and adds the id to `_failed_req_ids`
so a late inbound FetchMsg short-circuits to a clean rejection.
than `_unbound_store_timeout_s`. Drops the kv_request_id, surfaces
every batched job as failed, and records the id in `_reaped_stores`
so a late inbound FetchMsg is rejected in one round trip and a late
`submit_store` fails instead of re-parking under the reaped id.
Also expires `_reaped_stores` entries whose consumer can no longer
be waiting, which is why this runs before the empty check below.
"""
now = time.monotonic()
if self._reaped_stores:
retention_deadline = now - _REAPED_ID_RETENTION_S
for kid in [
kid
for kid, reaped_at in self._reaped_stores.items()
if reaped_at <= retention_deadline
]:
del self._reaped_stores[kid]
if not self._unbound_stores:
return
deadline = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S
deadline = now - self._unbound_store_timeout_s
expired: list[str] | None = None
for kid, batches in self._unbound_stores.items():
# Batches are appended in arrival order, so the head is oldest.
@@ -725,7 +792,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
return
for kid in expired:
batches = self._unbound_stores.pop(kid)
self._failed_req_ids.add(kid)
self._reaped_stores[kid] = now
for batch in batches:
self._finished_jobs.append(
JobResult(job_id=batch.job_id, success=False)
@@ -735,7 +802,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
"without a fetch — failing %d job(s)",
self._local_id,
kid,
_UNBOUND_STORE_TIMEOUT_S,
self._unbound_store_timeout_s,
len(batches),
)
@@ -779,6 +846,24 @@ class P2PSecondaryTierManager(SecondaryTierManager):
# inline in dispatch, so the replayed add_stored_blocks calls
# match that demand and submit transfers immediately.
for kv_request_id in result.new_fetch_ids:
if kv_request_id in self._reaped_stores:
# This id's blocks were already reaped, so no
# submit_store will ever satisfy the demand on_fetch
# just recorded. Finalize the round now: the peer gets
# TransferDoneMsg(success=False) on this tick and falls
# back to local prefill, instead of parking until its
# own load timeout expires. The id is deliberately left
# unbound so any later submit_store keeps taking the
# reject path. Batches re-parked between the reap and the
# retention prune are failed here rather than left to a
# second reap, which would pin their slots for another
# full timeout.
session.finish_request(kv_request_id)
for batch in self._unbound_stores.pop(kv_request_id, ()):
self._finished_jobs.append(
JobResult(job_id=batch.job_id, success=False)
)
continue
self._kv_to_session[kv_request_id] = session
for batch in self._unbound_stores.pop(kv_request_id, ()):
session.add_stored_blocks(
@@ -810,6 +895,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
JobResult(job_id=batch.job_id, success=False)
)
self._unbound_stores.clear()
self._reaped_stores.clear()
self._control.close()
self._data.close()