[KV Offload] Expose OffloadingManager.lock()/unlock() and service the tiering control plane off the scheduler thread

[liranschour]

Purpose

Two things, in priority order.

1. Expose exclusive access to an OffloadingManager as a public API. lock() / unlock(), plus a locked() helper, guard a manager's whole state subtree: the manager, the tiers it composes, and any front-end state those tiers own. Manager methods never acquire it themselves, so a caller holds it across as many operations as its own invariants need. That property is the point — several manager operations only make sense in sequence. prepare_load() documents that its keys are offloaded, a fact lookup() established one hook earlier, so a caller has to be able to span both. Internal per-method locking cannot offer that.

2. Use it for the first case that needs it (#55179). The p2p secondary tier answers a remote peer, which the engine does not drive, so its control plane only advanced from on_schedule_end. A peer's lookup or fetch therefore waited for a model-step boundary — seconds on a saturated rank, gating a transfer of well under one. Tiers now opt in with needs_control_plane_thread, and TieringOffloadingManager runs one thread that polls them and lets them serve between steps.

Not a duplicate of the open PRs on #55179

This overlaps open work and I want that on the record rather than buried.

PR	Author	Driver	Locking
this	me	thread owned by TieringOffloadingManager	public lock()/unlock() on OffloadingManager, held by callers
#58113	me	same thread, same sweep	private _enter_step/_exit_step/_observation_lock + *_unlocked variants
#57165	@nilig	EngineCore poller thread	new KVConnectorBase_V1.poll_pending_work
#55962	@Etelis	EngineCore + connector base API + executor changes	new base-API hook
#57139	me	—	the separate async peer-registration half, not this

This supersedes my own #58113, which I will close. The difference is not cosmetic:

• [KVConnector][Tiering] Drive the tiering control plane from a dedicated thread #58113 does not expose an API. It takes the lock inside the manager's public methods, which forces a *_unlocked variant of every method _SecondaryTierFacingParent can re-enter (lookup, on_new_request, create_store_job, on_request_finished, reset_cache) because the parent handle's caller already holds it. That duplication is most of its +393 lines in manager.py, and it is precisely what prevents an outside caller from saying "hold this across several operations." Moving the lock to the caller removes every one of those variants.
• [KVConnector][Tiering] Drive the tiering control plane from a dedicated thread #58113 releases at on_schedule_end, which is not the end of the step. build_connector_meta called on_schedule_end ~40 lines early, before _build_store_jobs → prepare_store (which evicts) and before on_request_finished. So its prepare_store re-acquired via _enter_step() and nothing released it: on every step that builds store jobs — the common case under load — the lock was held across future.result(), starving the control thread in exactly the window the change exists to exploit. That plausibly explains [KVConnector][Tiering] Drive the tiering control plane from a dedicated thread #58113's own "parity-within-noise" result, which was measured through the bug, and that conclusion should be re-tested.

Relative to #57165 / #55962: same goal, but no new KVConnectorBase_V1 hook and no EngineCore or executor changes. In #55179 @nilig asked for a read on "a thread touching the connector from inside EngineCore, even phase-exclusive, versus a control thread owned by the tier" — this is a working implementation of the second option, offered as input to that question. If maintainers prefer either existing PR, close this one.

Changes

Commit 1, on_schedule_end runs last — independent of all locking, reviewable on its own. Its docstring already claimed "called once at the end of each scheduler step"; it was called ~40 lines early. Moving it makes the contract true and gives callers a sound release point. Two ordering consequences, both handled: prepare_store's poll becomes unconditional (the per-step gate was previously reset mid-build_connector_meta, which is where its fresh completion data came from — without this, eviction decisions would run on data up to a step stale), and _flush_pending_cascades now re-parks keys this step parked, which is safe because _maybe_finalize_request holds off while pending_cascade_keys is non-empty.

Commit 2, the API and its consumer.

OffloadingConnectorScheduler holds one region per engine step, opened by the first manager-touching hook and closed at the end of build_connector_meta, so the manager is free while the engine runs the model. Hooks outside schedule() take a short region.

Two placements are load-bearing and easy to get wrong:

• has_pending_push_work must use the short region. EngineCore.step() returns early at core.py:597 when has_requests() is false, so a step-scoped acquire there would never reach the release and would hold the manager for good — the state an idle engine settles into. It does still need the lock: has_pending_work() iterates _req_state while the control thread may insert or delete via parent.on_new_request / on_request_finished.
• take_events drains inside the region. manager.take_events() is a generator whose backing list is cleared once exhausted, so iterating it lazily both pinned the manager for the consumer's duration and discarded events a concurrent complete_write appended, unpublished. It stays a generator so the caller's truthiness check at sched/scheduler.py is unchanged.

The thread yields unconditionally after unlock(), which is what bounds the scheduler's wait to one round — Python locks are not fair, so a release-then-reacquire loop could starve it. A failed round is logged, not fatal, and on_schedule_end still serves every tier, so the fallback is per-step servicing. _process_finished_jobs takes a tier filter so a round does not drag tiers that expect the scheduler thread. shutdown() stops and joins the thread before any tier teardown, since closing a transport under a live round can end the process rather than raise.

Config: control_plane_interval_s in kv_connector_extra_config, default 1 ms, 0 disables the thread and restores per-step servicing.

Docstrings that claimed single-threaded tier access are corrected, including the eviction-race argument in _pin_and_register_hits — which still holds, now because serve_external_requests runs inside one hold of the lock. async_lookup.py's "owned exclusively by the scheduler thread" claim was outright false once parent.lookup can fan out to fs.lookup from the control thread.

Test results

pytest tests/v1/kv_offload/tiering/ tests/v1/kv_connector/unit/offloading_connector/ -q
  901 passed, 12 skipped in 295s      (baseline on this tree: 889 passed, 12 skipped)

pytest tests/v1/kv_offload/tiering/p2p/ -q
  217 passed in 66s                   (subset of the above; re-run on its own)

pre-commit: all hooks pass, including ruff, ruff-format, mypy, and check-test-tethering.

Note for reviewers: a whole-directory run of tests/v1/kv_offload/ produces a large number of CUDA IMA errors that are pre-existing and unrelated (tests/v1/kv_offload/cpu/ poisons the rest of the session). Run the subdirectories separately.

12 new tests. The two that matter most were both verified to fail without the change, not just to pass with it:

• test_on_schedule_end_is_the_last_manager_call_of_a_step — without the reorder it observes ['on_schedule_end', 'prepare_store']. This is the check [KVConnector][Tiering] Drive the tiering control plane from a dedicated thread #58113's suite lacks, because its _step() helper drives on_schedule_end with no following prepare_store.
• test_control_thread_and_scheduler_never_overlap — with lock()/unlock() stubbed to no-ops it detects 86 overlaps. It also asserts both threads made progress, so neither a thread that never runs nor one that starves the scheduler passes it.

I wrote and then deleted a third test ("no eviction between a served lookup and its pin") after finding it still passed with the lock neutered: LRU keeps evicting the fresh keys rather than the constantly-looked-up hit keys, so it could never fail. The overlap test is the real guard for that property, and a test that cannot fail is worse than no test.

Model evaluation — the 4-case matrix from #55179

Run on two H200 pods, Llama-3.2-1B, --enforce-eager, --block-size 16, 96 GiB CPU
primary tier, secondary_tiers: [{"type": "p2p"}], max-model-len 32768.

Topology. Source = the pod holding the KV in its CPU offload tier. Consumer = the
other pod, which probes the source over P2P (remote_kv_source) and fetches. Every
measured fetch uses a freshly generated ~27k-token prompt that is populated on the source
first, so nothing hits a GPU prefix cache and the fetch is real work: each one moved
~1720 chunks (kv_offload_tiering_chunk_hits delta), i.e. the full prefix.

Busy = source saturated with 8 concurrent ~27k-token prefills. That took its
request latency from 0.018 s idle to 1.40 s, so the control plane really is gated
behind ~1.4 s of model work — the same regime as the issue's report.

Cold = no P2P session yet. A session lives as long as the engine, so only the first
fetch after a restart is genuinely cold; both pods were restarted between every phase and
the cold fetch measured first.

Signal = the consumer's vllm:kv_offload_lookup_async_delay_seconds delta for the
measured request (one observation per fetch). Warm figures are the mean of 3 reps.

case	A: baseline (7b077f36a8)	B: #58113	C: this PR
idle, cold	0.177 s	—	0.133 s
idle, warm	0.0455 s	—	0.0457 s
busy, cold	0.440 s	0.289 s	0.155 s
busy, warm	0.220 s	0.129 s	0.0534 s

The per-step gating penalty, which is what this PR removes:

arm	idle-warm	busy-warm	penalty
A: baseline	0.0455 s	0.2196 s	4.83x
C: this PR	0.0457 s	0.0534 s	1.17x

Saturating the source costs baseline 4.8x on the consumer's lookup delay. With the control
thread it costs 1.17x — the stall is essentially gone, and a busy source now behaves like
an idle one.

Supporting numbers on busy-warm: per-chunk tiering delay
(kv_offload_tiering_lookup_async_delay_seconds) 0.100 s → 0.0174 s (5.8x), and end-to-end
TTFT for the fetch 0.310 s → 0.144 s (2.2x). Idle is parity within noise in every metric,
as expected — there is no model step to hide behind.

Per-rep spread, busy-warm: baseline 0.1929 / 0.2532 / 0.2127, this PR
0.0388 / 0.0427 / 0.0786. The distributions do not overlap — the worst rep here is still
better than the best baseline rep.

No source-side slowdown. The concern that holding the manager lock slows the scheduler
does not show up: on the busy source, request latency was 1.4174 s (baseline) vs 1.4039 s
(this PR), and the populate-request mean 1.7056 s vs 1.7100 s. The background load
generator completed 158 vs 153 requests over comparable windows, a ~3% difference that is
within single-run noise and not corroborated by either latency measure. A dedicated
vllm bench serve throughput comparison is still worth running before merge; I did not
run one.

On #58113

Arm B is #58113 at b8ea135c71, measured on the same pods in the same session. It lands
between baseline and this PR: busy-warm 0.129 s vs baseline 0.220 s vs 0.0534 s here.

That is what the release-point bug predicts. #58113's thread is not dead — on steps that
build no store jobs the lock is released normally, so it recovers some of the win (1.7x
over baseline). But on any step that does build store jobs, prepare_store re-acquires
after on_schedule_end released and nothing releases it again, so the lock is held across
future.result() and the thread is shut out for exactly that step. Under load most steps
build store jobs, which is why it recovers only about a third of what is available (4.1x
over baseline here, 2.4x over #58113).

I am reporting this as consistent with the bug rather than as direct proof of the
mechanism; the mechanism itself is demonstrated by the unit test, which observes the call
order ['on_schedule_end', 'prepare_store'] without the reorder in commit 1.

Caveat on absolute magnitudes

These numbers are much smaller than the issue's 13.6 s, and deliberately so: this is
Llama-3.2-1B on a single H200 per side, not GLM-5.2-FP8 on DP8. What reproduces here is the
mechanism and its scaling — the consumer's lookup delay is proportional to the source's
step time under baseline, and roughly independent of it with the control thread. The issue's
13.6 s came from a far slower source; the ratio, not the absolute, is what transfers.

Raw per-rep JSON for all six runs is available on request.

NIXL thread-safety: checked, and it holds

Checked against the NIXL 1.4.1 sources (ai-dynamo/nixl at tag v1.4.1), which is the
deployed version. This corrects an earlier version of this description, which said
NIXL_THREAD_SYNC_RW does not exist and that the config offers no way to request a sync
mode. Both statements were true of 0.6.0, which is what I had installed; neither is true of
1.4.1.

Serialized use from a second thread is supported. makeMtType
(src/plugins/ucx/ucx_utils.cpp:383-392) returns only WORKER or SINGLE, never
CONTEXT -- and mt_mode_t::CONTEXT is the only value that maps to
UCS_THREAD_MODE_SINGLE (ucx_utils.cpp:476-492), the one UCX mode with creator-thread
affinity. Grepping confirms CONTEXT never appears as a produced value, only in switch
arms. So the weakest mode NIXL ever requests is UCS_THREAD_MODE_SERIALIZED, which UCX
documents as "Multiple threads can access, but only one at a time" -- precisely this
design. And because enable_prog_thread defaults on for this transport, it actually gets
UCS_THREAD_MODE_MULTI, with a hard library-support check that throws if unavailable.

The manager lock is load-bearing, not redundant. With syncMode = NONE -- which is what
this transport gets, since it sets neither enable_listen nor sync_mode -- NIXL's internal
agent lock compiles to no-ops (src/core/sync.h:36-39), leaving ~23 agent entry points
unguarded over shared, non-atomic containers. The external mutex is therefore the only
protection, and it must cover every entry point plus the lifetime of handles passed between
calls. It does.

Two caveats worth recording:

• NIXL binds a shared UCX worker per calling thread (tlsSharedWorkerMap,
  ucx_backend.cpp:821-825). That is harmless only while there is exactly one shared
  worker, which holds for this configuration (num_workers unset ⇒ numSharedWorkers_ == 1,
  so every thread maps to worker 0) but would not if a num_workers UCX backend param were
  ever set. The chosen worker also travels on the request handle rather than being
  re-derived per call, so a transfer prepped on one thread and polled on another still uses
  the same worker.
• 1.4.1 does accept an explicit sync_mode (including NIXL_THREAD_SYNC_RW, value 2),
  which would make NIXL's own locking real as belt-and-braces, and is what every upstream
  multi-threaded test uses. I have not added it, because the parameter does not exist on
  0.6.0 and passing it unconditionally would break older environments; it needs the same
  kind of version gate capture_telemetry already needs. Worth a follow-up.

Honest limit: ~/nixl is a 1.4.1 source checkout, not a build, and only 0.6.0 is installed
locally, so every 1.4.1 claim above is static source reading rather than runtime
introspection. A runtime confirmation (log the thread mode ucp_worker_query actually
grants, or run the two-thread pattern under TSan) would close the remaining gap. The
end-to-end matrix above ran on the deployed 1.4.1 across ~14k P2P chunk fetches with no
transfer failures, which is evidence but not a substitute.

pyzmq is fine under strict serialization: libzmq permits migrating a socket between threads
given a full memory barrier, which lock acquire/release provides, and nothing is bound
thread-locally at socket creation. The one hazard is zmq_ctx_destroy racing a socket call,
which the shutdown ordering handles.

AI assistance

AI assistance was used for this change (Claude Code): the code-path analysis, the implementation, and the tests. I am reviewing every changed line. The two items that were blocking this leaving draft are now done: the 4-case matrix has been run on H200 hardware (results above) and the NIXL question has been checked against 1.4.1. What remains before undraft is a dedicated throughput comparison and a human pass over the diff.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--58168.org.readthedocs.build/en/58168/

[orozery]

cc @mkhazraee wdyt?
This should help kvcr tier as well.
The purpose is to assure serve_external_requests is called more frequently (e.g. 1ms).
Today it is only called from the main thread, which may get to every O(s) if the model runner is busy.

[orozery]

Let's replace lock/unlock/locked with a single lock context manager property.

[liranschour]

Done

[orozery]

I think we should remove this.
The default should be no lock.
The CPU tier will implement its own threading lock.

[liranschour]

Done

[orozery]

spec.py imports this, and it's a user-facing default.
So it should be public.
We should check if this is the best place here.

[liranschour]

Made it public as DEFAULT_TIER_POLL_INTERVAL_S.
spec.py imports it from there as the fallback for the user-facing tier_poll_interval_s key.

[orozery]

How about we rename serve_control_plane / _control_plane_interval_s to poll_tiers() / tier_poll_interval_s?

[liranschour]

Done

[orozery]

how about serves_external_requests?

[liranschour]

Renamed

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @liranschour.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[nilig]

On my reproducer this PR fixes #55179: with a busy source rank the P2P lookup delay goes from 2.8-2.9 s to 0.6-0.7 s, the same value as with an idle source, and nothing else in the run got worse.

I ran this PR against the #55179 reproducer on a current nightly (e9757321, nixl 1.4.1), with the PR's files at 2fdfd08 applied on top of the same image for the second arm; the current head 271f52e only merges main and carries the same PR diff (an earlier run at a461602 gave the same busy-row numbers). Cell: GLM-5.2-FP8, two prefill pods and two decode pods, DP=8 per pod, 8 x H200 each, CPU offload tier 64 GiB per rank, P2P tier over UCX. Each case seeds a prefix on one prefill rank, then sends a follow-up to a different prefill rank that pulls it over P2P; "busy" means three 115,200-token requests were submitted to the source rank (the one that serves the lookup) 8 s before the follow-up; the source gauge showed two running when the follow-up was sent. The number is the destination rank's kv_offload_lookup_async_delay_seconds_sum delta for that one follow-up.

case	nightly	nightly + this PR
first P2P pull after boot, idle source	25.52 s (TTFT 33.5 s)	25.17 s (TTFT 33.2 s)
idle, warm pair x2	0.61 / 0.63 s	0.71 / 0.68 s
busy, warm pair x2	2.88 / 2.83 s	0.68 / 0.75 s
busy, cold pair	2.89 s	0.77 s
idle, cold pair	0.84 s	0.72 s

The busy rows land in the idle band with the PR, and it matches what I measured for the poller approach in #57165, so I think #57165 can be closed in favor of this one. Decode-side TPOT on the probed decode rank (request_time_per_output_token_seconds sum/count delta, 5 requests of 256 output tokens per arm) was 22.60 ms without and 22.17 ms with the PR (23.48 ms in the earlier run); at n=5 I would not read that as either equal or a regression, only as nothing gross. Every prefill engine logged KV offload control-plane thread started for tier(s) ['p2p'], interval 0.001s, so the thread was active during the cases.

The first pull after boot is a different problem (#57133: the destination rank's own first dummy batch JIT) and this PR does not change it: 25.5 s without, 25.2 s with, both on fresh nodes.

One row I cannot interpret yet: a gateway-routed follow-up with all eight source ranks loaded waited 9.7 s on its lookup with the PR (TTFT 11.8 s, 102K tokens pulled). n=1, no counterpart without the PR on this nightly, no JIT on the destination at that time, no lock-timeout warning from the control-plane thread. I will look at it separately; the direct-injection rows above do not show it.

One note on the setup, in case it matters for reproducing: prefill ran with --compilation-config '{"cudagraph_mode":"NONE"}' on both arms because this nightly captures CUDA graphs for deepep_high_throughput when breakable graphs are auto-enabled and crashes at startup (same early return as #53030).