@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 17bc3e106a

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

Signed-off-by: Liran Schour <lirans@il.ibm.com>

liranschour committed Sep 15, 2026 at 05:38 UTC 17bc3e106acbaeb552892893373659bb68d5020e

3 files changed +193 -14

docs/features/kv_offloading_usage.md

+22

@@ -203,9 +203,31 @@ Block content hashes must match across instances for peers to exchange blocks (s

203 | `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. |

204 | `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. |

205 | `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. |

206 +| `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. |

207

208 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.

209

210 +A producer parks a request's blocks until the consumer's `FetchMsg` arrives; if none arrives

211 +within `unbound_store_timeout_s` the blocks are released so they stop pinning primary-tier

212 +slots. Raise it when prefills legitimately take longer than the default:

213 +

214 +```bash

215 +vllm serve <model> \

216 + --kv-transfer-config '{

217 + "kv_connector": "OffloadingConnector",

218 + "kv_role": "kv_both",

219 + "kv_connector_extra_config": {

220 + "spec_name": "TieringOffloadingSpec",

221 + "secondary_tiers": [

222 + {

223 + "type": "p2p",

224 + "unbound_store_timeout_s": 180

225 + }

226 + ]

227 + }

228 + }'

229 +```

230 +

231 #### Environment Variables

232

233 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

@@ -117,6 +117,8 @@ def _make_manager() -> P2PSecondaryTierManager:

117 mgr._sessions = {}

118 mgr._kv_to_session = {}

119 mgr._unbound_stores = {}

120 + mgr._reaped_stores = {}

121 + mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S

122 mgr._failed_serve_ctxs = []

123 return mgr

124

@@ -759,9 +761,9 @@ class TestGetFinished:

761 assert "req-fresh" in mgr._unbound_stores

762

763 def test_unbound_store_reaped_after_timeout(self):

762 - """Unbound stores past _UNBOUND_STORE_TIMEOUT_S surface as failed

763 - and their kv_request_id lands in _failed_req_ids so a late

764 - FetchMsg/lookup doesn't try to satisfy them."""

764 + """Unbound stores past the reap deadline surface as failed and their

765 + kv_request_id lands in _reaped_stores so a late FetchMsg is rejected

766 + instead of parking demand nothing can satisfy."""

767 from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch

768

769 mgr = self._make()

@@ -779,7 +781,7 @@ class TestGetFinished:

781 # 2 baseline + 2 buffered stores

782 assert JobResult(job_id=10, success=False) in results

783 assert JobResult(job_id=11, success=False) in results

782 - assert "req-stale" in mgr._failed_req_ids

784 + assert "req-stale" in mgr._reaped_stores

785

786 def test_submit_store_parks_unbound_batch(self):

787 """submit_store on an unbound id appends a batch with a fresh

@@ -1204,6 +1206,75 @@ class TestBidirectionalManager:

1206 assert 201 in b_ok, f"B loads succeeded: {b_ok}"

1207 assert 200 in b_ok, f"B stores succeeded: {b_ok}"

1208

1209 + def test_late_fetch_after_reap_fails_immediately(self):

1210 + """A fetch for a reaped kv_request_id fails in one round trip.

1211 +

1212 + The producer drops parked blocks once unbound_store_timeout_s

1213 + expires. A consumer whose fetch arrives after that must learn on the

1214 + next tick via TransferDoneMsg(success=False) — not stall for

1215 + _LOAD_TIMEOUT_S and then take the abort path.

1216 + """

1217 + from vllm.v1.kv_offload.tiering.p2p.session.client import _LOAD_TIMEOUT_S

1218 +

1219 + mgr_a, mgr_b = _build_paired_managers()

1220 + kv_id = "req-late"

1221 +

1222 + # Record every message type leaving the consumer, so the assertion

1223 + # below can prove no abort was ever needed.

1224 + sent_from_a: list[str] = []

1225 + drain = mgr_a._control._drain_outbound_to

1226 +

1227 + def recording_drain(peer_local_id: str):

1228 + out = drain(peer_local_id)

1229 + sent_from_a.extend(msg.get("type") for _, msg in out)

1230 + return out

1231 +

1232 + mgr_a._control._drain_outbound_to = recording_drain # type: ignore[method-assign]

1233 +

1234 + # Producer parks the blocks, then reaps them before any fetch lands.

1235 + mgr_b.submit_store(

1236 + _job_metadata(

1237 + job_id=200,

1238 + keys=[b"b-block"],

1239 + chunk_ids=[0],

1240 + kv_params={"remote_decoder": {"kv_request_id": kv_id}},

1241 + )

1242 + )

1243 + mgr_b._unbound_store_timeout_s = 0.0

1244 + reap_results = list(mgr_b.get_finished_jobs())

1245 + assert JobResult(job_id=200, success=False) in reap_results

1246 + assert kv_id in mgr_b._reaped_stores

1247 +

1248 + # Consumer now asks for the blocks that no longer exist.

1249 + consumer_params = {

1250 + "remote_prefiller": {

1251 + "kv_request_id": kv_id,

1252 + "remote_host": "B",

1253 + "remote_port": 2,

1254 + },

1255 + }

1256 + mgr_a.on_new_request(_req_context(consumer_params))

1257 + mgr_a.submit_load(

1258 + _job_metadata(

1259 + job_id=101,

1260 + keys=[b"b-block"],

1261 + chunk_ids=[0],

1262 + kv_params=consumer_params,

1263 + )

1264 + )

1265 +

1266 + started = time.monotonic()

1267 + all_a: list[JobResult] = []

1268 + for _ in range(8):

1269 + all_a.extend(list(mgr_a.get_finished_jobs()))

1270 + list(mgr_b.get_finished_jobs())

1271 + elapsed = time.monotonic() - started

1272 +

1273 + assert JobResult(job_id=101, success=False) in all_a, all_a

1274 + assert kv_id in mgr_a._failed_req_ids

1275 + assert "abort_fetch" not in sent_from_a, sent_from_a

1276 + assert elapsed < _LOAD_TIMEOUT_S

1277 +

1278

1279 # ---------------------------------------------------------------------------

1280 # _accept_new_peers — duplicate connection rejection

vllm/v1/kv_offload/tiering/p2p/manager.py

+96 -10

@@ -47,10 +47,18 @@ logger = init_logger(__name__)

47 # Reap unbound store batches that have been parked without a FetchMsg

48 # binding them to a session for longer than this. Protects against the

49 # prefiller buffering blocks for a decoder that never asks (decoder died,

50 -# network partition, lost kv_request_id). Must be longer than the per-store

51 -# deadline so the store-timeout path fires first for individual jobs.

50 +# network partition, lost kv_request_id). Applies only while a batch is

51 +# unbound: binding hands it to the session, which times the job out under

52 +# _STORE_TIMEOUT_S instead.

53 +# Overridable per tier via the ``unbound_store_timeout_s`` config key.

54 _UNBOUND_STORE_TIMEOUT_S = 60.0

56 +# How long a reaped kv_request_id is remembered so a fetch that arrives

57 +# after the reap can be rejected immediately. Must exceed the consumer's

58 +# load timeout (session.client._LOAD_TIMEOUT_S): past that the consumer has

59 +# already given up on its own, so there is nobody left to reject.

60 +_REAPED_ID_RETENTION_S = 60.0

61 +

62 # Time we wait during shutdown for inflight transfers to drain via

63 # cancel(mode="wait") before falling back to mode="immediate". Bounded

64 # so a wedged peer can't hang shutdown.

@@ -208,6 +216,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

216 port: int | None = None,

217 backends: list[str] | None = None,

218 num_threads: int = 4,

219 + unbound_store_timeout_s: float = _UNBOUND_STORE_TIMEOUT_S,

220 **kwargs: Any,

221 ) -> None:

222 """Initialize the P2P secondary tier manager.

@@ -245,9 +254,29 @@ class P2PSecondaryTierManager(SecondaryTierManager):

254 num_threads: NIXL agent worker threads for the UCX-only

255 branch. Ignored when ``backends`` contains a non-UCX

256 entry.

257 + unbound_store_timeout_s: Seconds a producer holds stored blocks

258 + for a consumer that has not fetched them yet, before

259 + ``_reap_unbound_stores`` drops them. Raise it for

260 + deployments whose prefills outlast the default; the blocks

261 + keep primary-tier slots pinned for the whole window.

262 **kwargs: Reserved for future tier-specific options.

263 +

264 + Raises:

265 + ValueError: If ``unbound_store_timeout_s`` is not a positive

266 + number, or anything convertible to one.

267 """

268 super().__init__(offloading_spec, primary_kv_view, tier_type)

269 + try:

270 + timeout_s = float(unbound_store_timeout_s)

271 + except (TypeError, ValueError):

272 + timeout_s = float("nan")

273 + if not timeout_s > 0:

274 + raise ValueError(

275 + f"unbound_store_timeout_s must be a positive number, got "

276 + f"{unbound_store_timeout_s!r}"

277 + )

278 + self._unbound_store_timeout_s = timeout_s

279 +

280 # Block hashes chain from NONE_HASH (see v1/core/kv_cache_utils.py).

281 # Peers whose seeds differ compute different hashes for identical

282 # content, so lookups silently miss and no KV crosses the wire. The

@@ -303,8 +332,16 @@ class P2PSecondaryTierManager(SecondaryTierManager):

332 # kv_request_id → list of batches submit_store'd before any peer

333 # asked for that id. Drained into a session by _on_session_fetch

334 # when the corresponding FetchMsg arrives, or surfaced as failures

306 - # by _reap_unbound_stores after _UNBOUND_STORE_TIMEOUT_S.

335 + # by _reap_unbound_stores after _unbound_store_timeout_s.

336 self._unbound_stores: dict[str, list[_UnboundStoreBatch]] = {}

337 + # kv_request_id → time its parked batches were reaped. Nothing under

338 + # one of these ids can be served any more, so a FetchMsg for it is

339 + # rejected on the spot by _poll_once instead of parking demand until

340 + # the consumer's own load timeout, and submit_store fails rather than

341 + # re-parking under it — a prefill that outran the timeout would

342 + # otherwise keep pinning fresh primary-tier slots, one timeout at a

343 + # time. Entries are pruned after _REAPED_ID_RETENTION_S.

344 + self._reaped_stores: dict[str, float] = {}

345

346 self._finished_jobs: list[JobResult] = []

347 # kv_request_ids that hit a transport/session failure; On load lookup()

@@ -390,7 +427,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

427 kv_transfer_params; if a session has bound the id, finish it. If

428 no session has bound the id yet, this is a no-op: parked batches

429 in `_unbound_stores` are left in place and cleaned up only by

393 - `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`.

430 + `_reap_unbound_stores` after `_unbound_store_timeout_s`.

431 """

432 source = req_context.get_state(P2PSourceInfo)

433 dest = req_context.get_state(P2PDestInfo)

@@ -447,6 +484,22 @@ class P2PSecondaryTierManager(SecondaryTierManager):

484 self._finished_jobs.append(JobResult(job_id=job_id, success=False))

485 return

486

487 + # This id's earlier batches were already reaped, so the consumer has

488 + # either given up or will be rejected by the fetch path. Parking more

489 + # would pin primary-tier slots for another full timeout on blocks

490 + # nobody can fetch, so fail the job now — a prefill that outran the

491 + # timeout falls back to local recompute on the decoder either way.

492 + if kv_request_id in self._reaped_stores:

493 + logger.warning(

494 + "P2P %s: submit_store for reaped kv_request_id=%s job_id=%d "

495 + "— failing without parking",

496 + self._local_id,

497 + kv_request_id,

498 + job_id,

499 + )

500 + self._finished_jobs.append(JobResult(job_id=job_id, success=False))

501 + return

502 +

503 # Fast path: a session has already received FetchMsg for this id,

504 # so we can route the batch straight into its ServerRole.

505 session = self._kv_to_session.get(kv_request_id)

@@ -707,13 +760,27 @@ class P2PSecondaryTierManager(SecondaryTierManager):

760 """Time out submit_store batches that no peer has ever fetched.

761

762 Walks `_unbound_stores` for entries whose oldest batch is older

710 - than `_UNBOUND_STORE_TIMEOUT_S`. Drops the kv_request_id, surfaces

711 - every batched job as failed, and adds the id to `_failed_req_ids`

712 - so a late inbound FetchMsg short-circuits to a clean rejection.

763 + than `_unbound_store_timeout_s`. Drops the kv_request_id, surfaces

764 + every batched job as failed, and records the id in `_reaped_stores`

765 + so a late inbound FetchMsg is rejected in one round trip and a late

766 + `submit_store` fails instead of re-parking under the reaped id.

767 +

768 + Also expires `_reaped_stores` entries whose consumer can no longer

769 + be waiting, which is why this runs before the empty check below.

770 """

771 + now = time.monotonic()

772 + if self._reaped_stores:

773 + retention_deadline = now - _REAPED_ID_RETENTION_S

774 + for kid in [

775 + kid

776 + for kid, reaped_at in self._reaped_stores.items()

777 + if reaped_at <= retention_deadline

778 + ]:

779 + del self._reaped_stores[kid]

780 +

781 if not self._unbound_stores:

782 return

716 - deadline = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S

783 + deadline = now - self._unbound_store_timeout_s

784 expired: list[str] | None = None

785 for kid, batches in self._unbound_stores.items():

786 # Batches are appended in arrival order, so the head is oldest.

@@ -725,7 +792,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

792 return

793 for kid in expired:

794 batches = self._unbound_stores.pop(kid)

728 - self._failed_req_ids.add(kid)

795 + self._reaped_stores[kid] = now

796 for batch in batches:

797 self._finished_jobs.append(

798 JobResult(job_id=batch.job_id, success=False)

@@ -735,7 +802,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

802 "without a fetch — failing %d job(s)",

803 self._local_id,

804 kid,

738 - _UNBOUND_STORE_TIMEOUT_S,

805 + self._unbound_store_timeout_s,

806 len(batches),

807 )

808

@@ -779,6 +846,24 @@ class P2PSecondaryTierManager(SecondaryTierManager):

846 # inline in dispatch, so the replayed add_stored_blocks calls

847 # match that demand and submit transfers immediately.

848 for kv_request_id in result.new_fetch_ids:

849 + if kv_request_id in self._reaped_stores:

850 + # This id's blocks were already reaped, so no

851 + # submit_store will ever satisfy the demand on_fetch

852 + # just recorded. Finalize the round now: the peer gets

853 + # TransferDoneMsg(success=False) on this tick and falls

854 + # back to local prefill, instead of parking until its

855 + # own load timeout expires. The id is deliberately left

856 + # unbound so any later submit_store keeps taking the

857 + # reject path. Batches re-parked between the reap and the

858 + # retention prune are failed here rather than left to a

859 + # second reap, which would pin their slots for another

860 + # full timeout.

861 + session.finish_request(kv_request_id)

862 + for batch in self._unbound_stores.pop(kv_request_id, ()):

863 + self._finished_jobs.append(

864 + JobResult(job_id=batch.job_id, success=False)

865 + )

866 + continue

867 self._kv_to_session[kv_request_id] = session

868 for batch in self._unbound_stores.pop(kv_request_id, ()):

869 session.add_stored_blocks(

@@ -810,6 +895,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):

895 JobResult(job_id=batch.job_id, success=False)

896 )

897 self._unbound_stores.clear()

898 + self._reaped_stores.clear()

899 self._control.close()

900 self._data.close()

901