@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 332 lines 16.4 KB
Rendered Raw
1 # NIXL push-mode KV transfer
2
3 The default NIXL connector is **pull-based**: the decode (D) instance
4 reads KV blocks from the prefill (P) instance via `NIXL READ` after
5 prefill completes. `NixlPushConnector` adds a **push-based** alternative
6 in which P writes the KV blocks directly into D's pre-allocated memory
7 via `NIXL WRITE`.
8
9 This document describes the threading, queues, and scheduling
10 interactions specific to the push design. The pull-mode design is
11 unchanged; the push connector reuses the same handshake, NIXL agent
12 setup, and metadata path wherever possible.
13
14 ## High-level flow
15
16 ```mermaid
17 sequenceDiagram
18 autonumber
19 participant Client
20 participant Proxy
21 participant DSched as D Scheduler
22 participant DWorker as D Worker (main)
23 participant DWriter as D Writer
24 participant PWriter as P Writer
25 participant PWorker as P Worker (main)
26 participant PSched as P Scheduler
27
28 Client->>Proxy: POST /v1/completions
29 Proxy->>PSched: prefill leg (do_remote_decode=True, max_tokens=1)
30 Proxy->>DSched: decode leg (do_remote_prefill=True, P coordinates)
31
32 note over DSched,DWriter: D side - register blocks with P
33 DSched->>DSched: update_state_after_alloc, stash registration, arm watchdog
34 DSched->>DWorker: build_connector_meta -> meta.push_registrations
35 DWorker->>DWriter: enqueue (req_id, reg_data) on _reg_send_inbox
36 DWriter->>PWriter: NIXL send_notif PUSH_REG msgpack
37
38 note over PSched,PWriter: P side - prefill, stage finished blocks
39 PSched->>PSched: request_finished, stash blocks
40 PSched->>PWorker: build_connector_meta -> meta.push_finished_blocks
41 PWorker->>PWriter: enqueue (req_id, blocks) on _finished_blocks_inbox
42
43 note over PWriter: P writer matches and WRITEs
44 PWriter->>PWriter: get_new_notifs returns PUSH_REG, route via _handle_push_reg_notif
45 alt PUSH_REG and finished blocks both present
46 PWriter->>PWriter: pop matching pair, fire WRITE
47 else only one side present
48 PWriter->>PWriter: stash and wait, self-poll only when blocks unmatched
49 end
50 PWriter->>PWriter: _ensure_handshake to D (async; defer WRITE)
51 PWriter->>PWriter: handshake callback re-queues on _deferred_push_inbox, wake
52 PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif
53
54 note over DWorker,DWriter: D side - completion accounting
55 DWriter-->>DWorker: forward HB and completion notifs via _pending_completion_notifs
56 DWorker->>DWorker: _get_new_notifs drains, HB extends lease, completion marks recv done
57 DWorker->>DSched: update_connector_output(finished_recving)
58 DSched->>DSched: clear watchdog deadline
59
60 note over PWorker,PWriter: P side - reclaim
61 PWorker->>PWorker: get_finished, drain _sending_transfers, queue eviction
62 PWriter->>PWriter: drain _evict_finished_inbox, drop stale state
63 PWorker->>PSched: update_connector_output(finished_sending)
64 PSched->>PSched: free lease
65
66 DWorker-->>Proxy: stream decode tokens
67 Proxy-->>Client: response
68 ```
69
70 ## Threads
71
72 ``NixlPushConnectorWorker`` introduces a single dedicated background
73 thread per worker (i.e. per TP rank), named ``nixl-push-writer``.
74 Each owns the new push-specific NIXL operations on its rank:
75
76 * ``nixl_wrapper.get_new_notifs()`` — receive notifications.
77 * ``nixl_wrapper.send_notif(...)`` for the ``PUSH_REG:<msgpack>`` (D
78 side) and for the per-WRITE completion notif (P side).
79 * ``nixl_wrapper.make_prepped_xfer(...) / transfer(...)`` — submit the
80 WRITE itself.
81
82 Heartbeats continue to go out from the engine main thread via the
83 existing base-worker ``_send_heartbeats`` plumbing inside
84 ``start_load_kv``.
85
86 ### Wake model
87
88 The writer thread blocks on ``_push_writer_wake`` (a
89 ``threading.Event``) when it has no work. Three callers set the
90 event:
91
92 1. **``start_load_kv``** (worker main thread, called once per engine
93 step with the scheduler's metadata) — sets the wake only when the
94 step actually hands the writer new work, i.e. when
95 ``meta.push_registrations`` or ``meta.push_finished_blocks`` is
96 non-empty. This is the wake for new transfers.
97 2. **``get_finished``** (worker main thread, called once per engine
98 step to report completions) — always sets the wake. The writer is
99 the sole consumer of ``nixl_wrapper.get_new_notifs()`` for push,
100 so this gives it a chance to drain inbound notifs (heartbeats from
101 D, completion notifs after a WRITE, late-arriving ``PUSH_REG``)
102 even when there is no new metadata to act on.
103 3. **Handshake-completion callback** (background handshake executor
104 thread) — both handshakes run on the executor and never block the
105 writer; their done-callbacks re-enqueue the deferred op and set the
106 wake, since neither ``send_notif`` nor the NIXL WRITE may run off the
107 writer thread:
108 * the **D→P** handshake (before sending ``PUSH_REG``) re-enqueues the
109 registration onto ``_reg_send_inbox``;
110 * the **P→D** handshake (before a WRITE) re-enqueues the matched
111 ``(req_id, blocks, reg_data)`` onto ``_deferred_push_inbox``.
112
113 On this second pass ``_ensure_handshake`` returns ``None`` (the agent
114 is now connected), so the writer sends the ``PUSH_REG`` / issues the
115 WRITE directly. If a handshake *failed*, the callback fails or drops
116 the request instead of re-enqueuing, so there is no retry loop (see
117 Failure handling).
118
119 In addition to event-driven wakes, the writer self-polls at
120 ``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side
121 finished blocks waiting for an unmatched ``PUSH_REG``.
122
123 When a request completes on P (lease expires or the WRITE finishes),
124 ``get_finished`` enqueues the request id onto ``_evict_finished_inbox``,
125 which the writer drains to drop stale ``_push_finished_blocks`` /
126 ``_pending_d_registrations`` and stop self-polling.
127
128 ## Writer-local matching tables
129
130 | Table | Owner | Holds |
131 |--------------------------------|------------------|------------------------------------------------------------------------|
132 | `_pending_d_registrations` | writer | D registrations received from a remote D, waiting for P's blocks |
133 | `_push_finished_blocks` | writer | P blocks staged by the scheduler, waiting for a remote D registration |
134
135 Either side can arrive first. The writer matches in both directions:
136 when a ``PUSH_REG`` arrives we look up ``_push_finished_blocks``, and
137 when finished blocks arrive we look up ``_pending_d_registrations``.
138 Both lookups try an exact ``request_id`` match first, then fall back
139 to comparing the ids after stripping the trailing per-engine random
140 suffix (via ``get_base_request_id``). The fallback exists because the
141 proxy hands the same ``X-Request-Id`` to both legs, so P and D wrap it
142 into the same ``cmpl-<uuid>-<index>`` form and differ only by the
143 8-hex randomization suffix that ``input_processor.assign_request_id``
144 appends per engine. Stripping just that suffix normalizes both sides
145 to the same id while preserving the completion index (so multi-prompt
146 sub-requests stay distinct). It also works whether or not
147 ``VLLM_DISABLE_REQUEST_ID_RANDOMIZATION`` is set, which matters since
148 that env var is slated for removal upstream.
149
150 ## Wire format
151
152 A push registration is sent as a NIXL notification:
153
154 ```text
155 PUSH_REG:<msgpack-encoded dict>
156 ```
157
158 Fields in the dict:
159
160 | Field | Set by | Meaning |
161 |----------------------|--------|------------------------------------------------------------------------|
162 | ``request_id`` | D | D's own vLLM request id; P's match key, echoed in the completion notif |
163 | ``decode_engine_id`` | D | D's engine id (P uses this for the reverse handshake) |
164 | ``decode_host`` | D | D's NIXL side-channel host |
165 | ``decode_port`` | D | D's NIXL side-channel port |
166 | ``decode_tp_size`` | D | D's tensor-parallel size |
167 | ``local_block_ids`` | D | per-group lists of D's *logical* block ids (preallocated) |
168 | ``remote_engine_id`` | D | P's engine id (for the existing P-side handshake) |
169 | ``remote_host`` | D | P's NIXL side-channel host |
170 | ``remote_port`` | D | P's NIXL side-channel port |
171 | ``remote_tp_size`` | D | P's tensor-parallel size |
172
173 D ships **logical** block ids; P expands them to physical block ids at
174 WRITE-submission time using the ratio learned during the NIXL
175 handshake (`remote_physical_blocks_per_logical`). This matches the
176 pull-mode contract — schedulers ship logical ids, workers expand to
177 physical at submission.
178
179 The completion notif sent from P to D after a WRITE is the existing
180 `<request_id>:<tp_size>` format used in pull mode (here ``request_id``
181 is D's own request id, taken from the registration), so the D-side
182 accounting code is unchanged.
183
184 ## Pipeline parallelism and hybrid KV caches
185
186 The pull connector requires the local and remote workers to expose a
187 congruent list of KV regions — region *i* here corresponds to region
188 *i* there. That assumption breaks under pipeline parallelism (PP)
189 combined with a hybrid (HMA) KV layout:
190
191 * a PP-sharded prefiller (**P**) holds only a slice of the model's
192 layers while the `PP=1` decoder (**D**) holds them all, so region
193 counts differ;
194 * with HMA, several layer names are pooled into one region and the layer
195 that represents a pooled region can differ between P and D.
196
197 `NixlPushConnector` handles this for PP-sharded producers by routing
198 **by layer name** instead of by region index. Each
199 worker advertises which layer names back each of its NIXL regions in the
200 handshake metadata (`NixlAgentMetadata.region_members`). A producer that
201 needs layer routing derives its transfer layer order once, when it
202 registers its KV caches, and every transfer it issues uses that order.
203 `add_remote_agent` then selects exactly the remote regions this stage
204 owns and reorders them to match, so both sides stay paired regardless of
205 how each remote rank happens to order its metadata.
206
207 Addresses, block lengths, strides, and per-region capacities follow the same
208 layer order. Descriptor offsets use each layer's region capacity, so P and
209 D need not allocate the same number of blocks. Physical allocations are still
210 registered once, even when multiple layers share them.
211
212 Packed MLA caches (block-outermost layouts such as DeepSeek-V4's) store the
213 pages of every layer for one block in a single row. A PP stage holds only its
214 own layers, so its rows are shorter and each layer sits at a different offset:
215
216 ```text
217 D (PP=1), block b: | L0 | L1 | L2 | L3 | L4 | L5 |
218 P (stage 2), block b: | L3 | L4 | L5 |
219 ```
220
221 `PP=1` peers transfer whole rows. They also advertise where each layer's page
222 sits in the row (`NixlAgentMetadata.packed_member_layouts`), so a PP producer
223 writes each of its pages to the matching slot in D's row. Page sizes must
224 match; offsets, row strides, and block counts may differ. Because PP stages can
225 list the same attention backends in a different order, packed push compares the
226 sorted set of backend names during the handshake. Pull mode is unchanged.
227
228 Invariants enforced when the remote regions are aligned:
229
230 * every locally owned layer must be advertised exactly once by the
231 remote; a missing layer fails the handshake rather than silently
232 leaving that layer's KV stale, and remote-only layers (owned by other
233 PP stages) are ignored;
234 * a remote that omits layer metadata while the local layout requires
235 layer routing fails the handshake instead of falling back to
236 region-index routing;
237 * layer order is a property of the local layout alone, so the same local
238 source descriptors serve every remote engine and TP rank; only the
239 remote descriptor list is rebuilt per rank.
240
241 Decode-side PP is unsupported because completions are counted per
242 consumer rank. Mamba/SSM hybrids are unsupported under PP. Packed PP push
243 requires MLA caches and equal P/D block sizes. MLA pages are replicated
244 across TP ranks, so producer TP1 to consumer TP2 is supported.
245
246 ## Scheduler-side responsibilities
247
248 `NixlPushConnectorScheduler` extends the base scheduler with:
249
250 * **D side** — `update_state_after_alloc` stashes registration data in
251 `_push_pending_registrations` and arms a soft watchdog
252 (`_push_registration_deadlines`). `build_connector_meta` drains the
253 stash into `meta.push_registrations` and any expired entries are
254 dropped with a warning.
255 * **P side** — `request_finished` stashes block IDs in
256 `_finished_request_blocks` (for the lease and for
257 `has_pending_push_work`) and `_newly_finished_push_blocks` (for the
258 next worker step via `meta.push_finished_blocks`).
259 * **Both sides** — `has_pending_push_work` keeps the engine main loop
260 stepping while there is in-flight push state, so the writer always
261 gets at least one wake per step.
262
263 `update_connector_output`:
264
265 * `finished_sending` (P side) clears the lease entry.
266 * `finished_recving` (D side) clears the watchdog deadline.
267
268 ## Timeouts and watchdogs
269
270 Two per-request timers are armed on the scheduler:
271
272 * **D-side registration watchdog** — ``_push_registration_deadlines``.
273 If a registered request does not see a push completion within
274 ``push_registration_timeout`` seconds (defaults to
275 ``decoder_kv_blocks_ttl``), ``build_connector_meta`` drops the stale
276 registration and the pending entry, logs a warning, and stops trying
277 to resend the registration. The corresponding request remains tracked
278 in ``_reqs_need_recv``; it is the engine's request-level abort path
279 (or the user / proxy timing out the HTTP call) that ultimately fails
280 the request.
281 * **P-side block lease** — same ``_kv_lease_duration`` used by pull
282 mode. ``request_finished`` sets the expiration in ``_reqs_need_send``
283 and ``update_connector_output(finished_sending=...)`` clears it on
284 successful WRITE. Stale leases are reaped by ``get_finished`` in the
285 base worker, which then enqueues the eviction onto
286 ``_evict_finished_inbox`` so the writer also stops self-polling.
287
288 ## Failure handling
289
290 * **D-side handshake failure (P→D handshake before sending PUSH_REG)** —
291 the future's done-callback calls ``_handle_failed_transfer(rid, None)``,
292 which marks D's pre-allocated blocks invalid and enqueues onto
293 ``_failed_recv_reqs`` so the next ``get_finished`` reports the
294 request as a failed recv. Same recv-side accounting as pull mode.
295 * **D-side ``send_notif`` failure when shipping the PUSH_REG to P** —
296 identical handling: ``_handle_failed_transfer`` marks the recv as
297 failed.
298 * **P-side handshake failure (P→D handshake before a WRITE)** — the
299 future's done-callback logs ``push_handshake_failed`` and drops the
300 request without re-queuing. It deliberately does *not* call
301 ``_handle_failed_transfer`` (there is no ``_recving_metadata`` entry to
302 invalidate on the producer side, same reasoning as the WRITE-submission
303 failure below). P's blocks are reclaimed by the ``_kv_lease_duration``
304 lease and D's stale registration by its watchdog.
305 * **P-side WRITE submission failure** — the WRITE handle (if any) is
306 released and ``xfer_stats.record_failed_transfer()`` bumps the
307 failure counter. We deliberately do not call
308 ``_handle_failed_transfer`` here: ``req_id`` on the P side has no
309 entry in ``_recving_metadata`` (P is not the receiver), so the
310 helper would put a P-local request id into ``_failed_recv_reqs``
311 and trip the assertion in the base worker's ``get_finished``. The
312 outbound WRITE is dropped on the floor; D's lease watchdog handles
313 the missing completion.
314
315 ## Summary
316
317 The push design is a small, well-contained extension on top of the
318 existing NIXL connector:
319
320 * one new connector class, one new scheduler class, one new worker
321 class — all subclasses of the existing base classes;
322 * one dedicated background thread per worker;
323 * a few cross-thread queues, each with a single consumer (the writer);
324 most have one producer, except the two replay queues fed by both the
325 engine main thread and a handshake-completion callback:
326 ``_reg_send_inbox`` (registrations replayed after their D→P handshake)
327 and ``_deferred_push_inbox`` (matched pushes replayed after their P→D
328 handshake);
329 * one new notification type (`PUSH_REG:<msgpack>`).
330
331 Behavior on the engine main thread is otherwise unchanged. The writer
332 thread is event-driven and idle when there is no push work.