@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 88ed636218
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[KV Connector]: Support KV push from Prefill to Decode node using Nixl KV Connector (#35264)

Signed-off-by: Sunita Nadampalli <nadampal@amazon.com> Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

snadampal committed Jun 12, 2026 at 05:38 UTC 88ed63621866d1e4bdaacc560c911f7b8859c53d
26 files changed +6335 -3370
docs/design/nixl_kv_push_connector.md new
+256
@@ -0,0 +1,256 @@
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 D handshake (one-time)
51 + PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif
52 +
53 + note over DWorker,DWriter: D side - completion accounting
54 + DWriter-->>DWorker: forward HB and completion notifs via _pending_completion_notifs
55 + DWorker->>DWorker: _get_new_notifs drains, HB extends lease, completion marks recv done
56 + DWorker->>DSched: update_connector_output(finished_recving)
57 + DSched->>DSched: clear watchdog deadline
58 +
59 + note over PWorker,PWriter: P side - reclaim
60 + PWorker->>PWorker: get_finished, drain _sending_transfers, queue eviction
61 + PWriter->>PWriter: drain _evict_finished_inbox, drop stale state
62 + PWorker->>PSched: update_connector_output(finished_sending)
63 + PSched->>PSched: free lease
64 +
65 + DWorker-->>Proxy: stream decode tokens
66 + Proxy-->>Client: response
67 +```
68 +
69 +## Threads
70 +
71 +``NixlPushConnectorWorker`` introduces a single dedicated background
72 +thread per worker (i.e. per TP rank), named ``nixl-push-writer``.
73 +Each owns the new push-specific NIXL operations on its rank:
74 +
75 +* ``nixl_wrapper.get_new_notifs()`` — receive notifications.
76 +* ``nixl_wrapper.send_notif(...)`` for the ``PUSH_REG:<msgpack>`` (D
77 + side) and for the per-WRITE completion notif (P side).
78 +* ``nixl_wrapper.make_prepped_xfer(...) / transfer(...)`` — submit the
79 + WRITE itself.
80 +
81 +Heartbeats continue to go out from the engine main thread via the
82 +existing base-worker ``_send_heartbeats`` plumbing inside
83 +``start_load_kv``.
84 +
85 +### Wake model
86 +
87 +The writer thread blocks on ``_push_writer_wake`` (a
88 +``threading.Event``) when it has no work. Three callers set the
89 +event:
90 +
91 +1. **``start_load_kv``** (worker main thread, called once per engine
92 + step with the scheduler's metadata) — sets the wake only when the
93 + step actually hands the writer new work, i.e. when
94 + ``meta.push_registrations`` or ``meta.push_finished_blocks`` is
95 + non-empty. This is the wake for new transfers.
96 +2. **``get_finished``** (worker main thread, called once per engine
97 + step to report completions) — always sets the wake. The writer is
98 + the sole consumer of ``nixl_wrapper.get_new_notifs()`` for push,
99 + so this gives it a chance to drain inbound notifs (heartbeats from
100 + D, completion notifs after a WRITE, late-arriving ``PUSH_REG``)
101 + even when there is no new metadata to act on.
102 +3. **Handshake-completion callback** (background handshake executor
103 + thread) — when a deferred D→P handshake finishes successfully, the
104 + future's done-callback re-enqueues the registration onto
105 + ``_reg_send_inbox`` and sets the wake so the corresponding
106 + ``send_notif`` runs on the writer (we never call ``send_notif`` from
107 + the executor thread). On this second pass ``_ensure_handshake``
108 + returns ``None`` (the agent is now connected), so the writer sends
109 + the ``PUSH_REG`` directly. If the handshake *failed*, the callback
110 + fails the request instead of re-enqueuing, so there is no retry
111 + loop.
112 +
113 +In addition to event-driven wakes, the writer self-polls at
114 +``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side
115 +finished blocks waiting for an unmatched ``PUSH_REG``.
116 +
117 +When a request completes on P (lease expires or the WRITE finishes),
118 +``get_finished`` enqueues the request id onto ``_evict_finished_inbox``,
119 +which the writer drains to drop stale ``_push_finished_blocks`` /
120 +``_pending_d_registrations`` and stop self-polling.
121 +
122 +## Writer-local matching tables
123 +
124 +| Table | Owner | Holds |
125 +|--------------------------------|------------------|------------------------------------------------------------------------|
126 +| `_pending_d_registrations` | writer | D registrations received from a remote D, waiting for P's blocks |
127 +| `_push_finished_blocks` | writer | P blocks staged by the scheduler, waiting for a remote D registration |
128 +
129 +Either side can arrive first. The writer matches in both directions:
130 +when a ``PUSH_REG`` arrives we look up ``_push_finished_blocks``, and
131 +when finished blocks arrive we look up ``_pending_d_registrations``.
132 +Both lookups try an exact ``request_id`` match first, then fall back
133 +to comparing the ids after stripping the trailing per-engine random
134 +suffix (via ``get_base_request_id``). The fallback exists because the
135 +proxy hands the same ``X-Request-Id`` to both legs, so P and D wrap it
136 +into the same ``cmpl-<uuid>-<index>`` form and differ only by the
137 +8-hex randomization suffix that ``input_processor.assign_request_id``
138 +appends per engine. Stripping just that suffix normalizes both sides
139 +to the same id while preserving the completion index (so multi-prompt
140 +sub-requests stay distinct). It also works whether or not
141 +``VLLM_DISABLE_REQUEST_ID_RANDOMIZATION`` is set, which matters since
142 +that env var is slated for removal upstream.
143 +
144 +## Wire format
145 +
146 +A push registration is sent as a NIXL notification:
147 +
148 +```text
149 +PUSH_REG:<msgpack-encoded dict>
150 +```
151 +
152 +Fields in the dict:
153 +
154 +| Field | Set by | Meaning |
155 +|----------------------|--------|------------------------------------------------------------------------|
156 +| ``request_id`` | D | D's own vLLM request id; P's match key, echoed in the completion notif |
157 +| ``decode_engine_id`` | D | D's engine id (P uses this for the reverse handshake) |
158 +| ``decode_host`` | D | D's NIXL side-channel host |
159 +| ``decode_port`` | D | D's NIXL side-channel port |
160 +| ``decode_tp_size`` | D | D's tensor-parallel size |
161 +| ``local_block_ids`` | D | per-group lists of D's *logical* block ids (preallocated) |
162 +| ``remote_engine_id`` | D | P's engine id (for the existing P-side handshake) |
163 +| ``remote_host`` | D | P's NIXL side-channel host |
164 +| ``remote_port`` | D | P's NIXL side-channel port |
165 +| ``remote_tp_size`` | D | P's tensor-parallel size |
166 +
167 +D ships **logical** block ids; P expands them to physical block ids at
168 +WRITE-submission time using the ratio learned during the NIXL
169 +handshake (`remote_physical_blocks_per_logical`). This matches the
170 +pull-mode contract — schedulers ship logical ids, workers expand to
171 +physical at submission.
172 +
173 +The completion notif sent from P to D after a WRITE is the existing
174 +`<request_id>:<tp_size>` format used in pull mode (here ``request_id``
175 +is D's own request id, taken from the registration), so the D-side
176 +accounting code is unchanged.
177 +
178 +## Scheduler-side responsibilities
179 +
180 +`NixlPushConnectorScheduler` extends the base scheduler with:
181 +
182 +* **D side** — `update_state_after_alloc` stashes registration data in
183 + `_push_pending_registrations` and arms a soft watchdog
184 + (`_push_registration_deadlines`). `build_connector_meta` drains the
185 + stash into `meta.push_registrations` and any expired entries are
186 + dropped with a warning.
187 +* **P side** — `request_finished` stashes block IDs in
188 + `_finished_request_blocks` (for the lease and for
189 + `has_pending_push_work`) and `_newly_finished_push_blocks` (for the
190 + next worker step via `meta.push_finished_blocks`).
191 +* **Both sides** — `has_pending_push_work` keeps the engine main loop
192 + stepping while there is in-flight push state, so the writer always
193 + gets at least one wake per step.
194 +
195 +`update_connector_output`:
196 +
197 +* `finished_sending` (P side) clears the lease entry.
198 +* `finished_recving` (D side) clears the watchdog deadline.
199 +
200 +## Timeouts and watchdogs
201 +
202 +Two per-request timers are armed on the scheduler:
203 +
204 +* **D-side registration watchdog** — ``_push_registration_deadlines``.
205 + If a registered request does not see a push completion within
206 + ``push_registration_timeout`` seconds (defaults to
207 + ``decoder_kv_blocks_ttl``), ``build_connector_meta`` drops the stale
208 + registration and the pending entry, logs a warning, and stops trying
209 + to resend the registration. The corresponding request remains tracked
210 + in ``_reqs_need_recv``; it is the engine's request-level abort path
211 + (or the user / proxy timing out the HTTP call) that ultimately fails
212 + the request.
213 +* **P-side block lease** — same ``_kv_lease_duration`` used by pull
214 + mode. ``request_finished`` sets the expiration in ``_reqs_need_send``
215 + and ``update_connector_output(finished_sending=...)`` clears it on
216 + successful WRITE. Stale leases are reaped by ``get_finished`` in the
217 + base worker, which then enqueues the eviction onto
218 + ``_evict_finished_inbox`` so the writer also stops self-polling.
219 +
220 +## Failure handling
221 +
222 +* **D-side handshake failure (P→D handshake before sending PUSH_REG)** —
223 + the future's done-callback calls ``_handle_failed_transfer(rid, None)``,
224 + which marks D's pre-allocated blocks invalid and enqueues onto
225 + ``_failed_recv_reqs`` so the next ``get_finished`` reports the
226 + request as a failed recv. Same recv-side accounting as pull mode.
227 +* **D-side ``send_notif`` failure when shipping the PUSH_REG to P** —
228 + identical handling: ``_handle_failed_transfer`` marks the recv as
229 + failed.
230 +* **P-side WRITE submission failure** — the WRITE handle (if any) is
231 + released and ``xfer_stats.record_failed_transfer()`` bumps the
232 + failure counter. We deliberately do not call
233 + ``_handle_failed_transfer`` here: ``req_id`` on the P side has no
234 + entry in ``_recving_metadata`` (P is not the receiver), so the
235 + helper would put a P-local request id into ``_failed_recv_reqs``
236 + and trip the assertion in the base worker's ``get_finished``. The
237 + outbound WRITE is dropped on the floor; D's lease watchdog handles
238 + the missing completion.
239 +
240 +## Summary
241 +
242 +The push design is a small, well-contained extension on top of the
243 +existing NIXL connector:
244 +
245 +* one new connector class, one new scheduler class, one new worker
246 + class — all subclasses of the existing base classes;
247 +* one dedicated background thread per worker;
248 +* a few cross-thread queues, each with a single consumer (the writer);
249 + most have one producer, except ``_reg_send_inbox``, which is fed both
250 + by the engine main thread (new registrations) and by the
251 + handshake-completion callback (registrations replayed after their
252 + D→P handshake finishes);
253 +* one new notification type (`PUSH_REG:<msgpack>`).
254 +
255 +Behavior on the engine main thread is otherwise unchanged. The writer
256 +thread is event-driven and idle when there is no push work.
examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py new
+429
@@ -0,0 +1,429 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Push-mode disaggregated prefilling proxy demo.
5 +
6 +Companion to ``disagg_proxy_demo.py`` (pull mode). The client-facing API is
7 +the same; the difference is in how P and D coordinate the KV transfer:
8 +
9 +* Pull mode: proxy forwards P's ``kv_transfer_params`` (including
10 + ``remote_block_ids``) to D, and D pulls KV from P via NIXL READ.
11 +* Push mode: proxy hands D **only** P's coordinates
12 + (``remote_engine_id``, ``remote_host``, ``remote_port``, ``tp_size``)
13 + and the shared ``remote_request_id``. D registers its locally allocated
14 + blocks with P over a NIXL notification; P then pushes the KV to D via
15 + NIXL WRITE.
16 +
17 +Launch multiple vLLM instances configured with ``NixlPushConnector`` and
18 +matching ``engine_id`` / ``side_channel_port``, then start this proxy:
19 +
20 + python3 examples/disaggregated/disaggregated_serving/\
21 +disagg_proxy_pushconnector_demo.py \
22 + --model $model_name \
23 + --prefill localhost:8100 \
24 + --decode localhost:8200 \
25 + --prefill-engine-id prefill-engine-001 \
26 + --prefill-kv-host 10.0.0.1 \
27 + --prefill-side-channel-port 5600 \
28 + --prefill-tp-size 1 \
29 + --port 8000
30 +"""
31 +
32 +import argparse
33 +import contextlib
34 +import ipaddress
35 +import itertools
36 +import json
37 +import logging
38 +import os
39 +import sys
40 +import uuid
41 +from abc import ABC, abstractmethod
42 +from collections.abc import Callable
43 +
44 +import aiohttp
45 +import uvicorn
46 +from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
47 +from fastapi.responses import JSONResponse, StreamingResponse
48 +
49 +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
50 +logger = logging.getLogger()
51 +logging.basicConfig(level=logging.INFO)
52 +
53 +
54 +class SchedulingPolicy(ABC):
55 + @abstractmethod
56 + def schedule(self, cycler: itertools.cycle):
57 + raise NotImplementedError("Scheduling Proxy is not set.")
58 +
59 +
60 +class RoundRobinSchedulingPolicy(SchedulingPolicy):
61 + def schedule(self, cycler: itertools.cycle) -> str:
62 + return next(cycler)
63 +
64 +
65 +class PushProxy:
66 + """Push-mode proxy.
67 +
68 + The structure mirrors the pull-mode ``Proxy`` in
69 + ``disagg_proxy_demo.py``: an APIRouter with ``/v1/completions``,
70 + ``/v1/chat/completions``, ``/status`` and ``/instances/add``, plus
71 + round-robin scheduling across multiple P / D instances.
72 +
73 + Push-specific differences are confined to the request-handling
74 + methods (``create_completion`` / ``create_chat_completion``):
75 +
76 + * D's ``kv_transfer_params`` is built from CLI-provided P
77 + coordinates instead of being derived from P's response.
78 + * P and D requests are issued concurrently — D registers blocks and
79 + waits while P prefills and pushes.
80 + """
81 +
82 + def __init__(
83 + self,
84 + prefill_instances: list[str],
85 + decode_instances: list[str],
86 + model: str,
87 + scheduling_policy: SchedulingPolicy,
88 + prefill_engine_id: str,
89 + prefill_kv_host: str,
90 + prefill_side_channel_port: int,
91 + prefill_tp_size: int,
92 + custom_create_completion: Callable[[Request], StreamingResponse] | None = None,
93 + custom_create_chat_completion: Callable[[Request], StreamingResponse]
94 + | None = None,
95 + ):
96 + self.prefill_instances = prefill_instances
97 + self.decode_instances = decode_instances
98 + self.prefill_cycler = itertools.cycle(prefill_instances)
99 + self.decode_cycler = itertools.cycle(decode_instances)
100 + self.model = model
101 + self.scheduling_policy = scheduling_policy
102 +
103 + # Push-mode metadata: D needs P's coordinates up-front. Pull mode
104 + # learns these from P's response; push mode uses CLI args because
105 + # D issues its registration before P responds.
106 + self.push_metadata = {
107 + "do_remote_decode": False,
108 + "do_remote_prefill": True,
109 + "remote_engine_id": prefill_engine_id,
110 + "remote_host": prefill_kv_host,
111 + "remote_port": prefill_side_channel_port,
112 + "tp_size": prefill_tp_size,
113 + }
114 +
115 + self.custom_create_completion = custom_create_completion
116 + self.custom_create_chat_completion = custom_create_chat_completion
117 + self.router = APIRouter()
118 + self.setup_routes()
119 +
120 + # ── routes ──────────────────────────────────────────────────────── #
121 +
122 + def setup_routes(self):
123 + self.router.post(
124 + "/v1/completions", dependencies=[Depends(self.validate_json_request)]
125 + )(
126 + self.custom_create_completion
127 + if self.custom_create_completion
128 + else self.create_completion
129 + )
130 + self.router.post(
131 + "/v1/chat/completions", dependencies=[Depends(self.validate_json_request)]
132 + )(
133 + self.custom_create_chat_completion
134 + if self.custom_create_chat_completion
135 + else self.create_chat_completion
136 + )
137 + self.router.get("/status", response_class=JSONResponse)(self.get_status)
138 +
139 + async def validate_json_request(self, raw_request: Request):
140 + content_type = raw_request.headers.get("content-type", "").lower()
141 + if content_type != "application/json":
142 + raise HTTPException(
143 + status_code=415,
144 + detail="Unsupported Media Type: Only 'application/json' is allowed",
145 + )
146 +
147 + # ── HTTP forwarding ─────────────────────────────────────────────── #
148 +
149 + async def forward_request(self, url, data, headers, use_chunked=True):
150 + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
151 + try:
152 + async with session.post(
153 + url=url, json=data, headers=headers
154 + ) as response:
155 + if 200 <= response.status < 300 or 400 <= response.status < 500:
156 + if use_chunked:
157 + async for chunk_bytes in response.content.iter_chunked(
158 + 1024
159 + ):
160 + yield chunk_bytes
161 + else:
162 + yield await response.read()
163 + else:
164 + error_content = await response.text()
165 + with contextlib.suppress(json.JSONDecodeError):
166 + error_content = json.loads(error_content)
167 + logger.error(
168 + "Request failed with status %s: %s",
169 + response.status,
170 + error_content,
171 + )
172 + raise HTTPException(
173 + status_code=response.status,
174 + detail=f"Request failed with status {response.status}: "
175 + f"{error_content}",
176 + )
177 + except aiohttp.ClientError as e:
178 + logger.error("ClientError occurred: %s", str(e))
179 + raise HTTPException(
180 + status_code=502,
181 + detail="Bad Gateway: Error communicating with upstream server.",
182 + ) from e
183 + except Exception as e:
184 + logger.error("Unexpected error: %s", str(e))
185 + raise HTTPException(status_code=500, detail=str(e)) from e
186 +
187 + def schedule(self, cycler: itertools.cycle) -> str:
188 + return self.scheduling_policy.schedule(cycler)
189 +
190 + async def get_status(self):
191 + return {
192 + "mode": "push",
193 + "prefill_node_count": len(self.prefill_instances),
194 + "decode_node_count": len(self.decode_instances),
195 + "prefill_nodes": self.prefill_instances,
196 + "decode_nodes": self.decode_instances,
197 + "prefill_engine_id": self.push_metadata["remote_engine_id"],
198 + "prefill_kv_host": self.push_metadata["remote_host"],
199 + "prefill_side_channel_port": self.push_metadata["remote_port"],
200 + "prefill_tp_size": self.push_metadata["tp_size"],
201 + }
202 +
203 + # ── push-mode request handling ──────────────────────────────────── #
204 +
205 + def _build_decode_kv_params(self, request_id: str) -> dict:
206 + """Push-mode kv_transfer_params for D.
207 +
208 + ``remote_block_ids`` is intentionally omitted: D allocates its
209 + own blocks and registers them with P; P determines the
210 + prefill-side block IDs and ships them via the WRITE.
211 + """
212 + params = self.push_metadata.copy()
213 + params["remote_request_id"] = request_id
214 + return params
215 +
216 + def _common_headers(self, request_id: str) -> dict:
217 + h = {"X-Request-Id": request_id}
218 + api_key = os.environ.get("OPENAI_API_KEY")
219 + if api_key:
220 + h["Authorization"] = f"Bearer {api_key}"
221 + return h
222 +
223 + async def _push_completion(self, raw_request: Request, path: str):
224 + """Shared body for /v1/completions and /v1/chat/completions.
225 +
226 + Push mode fires P and D concurrently:
227 + * P runs a normal prefill (max_tokens=1, do_remote_decode=True).
228 + * D runs the decode (do_remote_prefill=True, no remote_block_ids).
229 +
230 + D blocks waiting for P's WRITE; the response streamed back to the
231 + client is the decode output from D.
232 + """
233 + request = await raw_request.json()
234 + request_id = str(uuid.uuid4())
235 +
236 + # Prefill leg (max_tokens=1, signals P to keep KV around for D).
237 + prefill_request = request.copy()
238 + prefill_request["max_tokens"] = 1
239 + if "max_completion_tokens" in prefill_request:
240 + prefill_request["max_completion_tokens"] = 1
241 + prefill_request["kv_transfer_params"] = {
242 + "do_remote_decode": True,
243 + "do_remote_prefill": False,
244 + "remote_engine_id": None,
245 + "remote_block_ids": None,
246 + "remote_host": None,
247 + "remote_port": None,
248 + }
249 +
250 + # Decode leg (push mode: no remote_block_ids).
251 + decode_request = request.copy()
252 + decode_request["kv_transfer_params"] = self._build_decode_kv_params(request_id)
253 +
254 + prefill_instance = self.schedule(self.prefill_cycler)
255 + decode_instance = self.schedule(self.decode_cycler)
256 + headers = self._common_headers(request_id)
257 +
258 + # Fire prefill; we don't read its body but must drain the
259 + # connection so the upstream server can free its slot.
260 + async for _ in self.forward_request(
261 + f"http://{prefill_instance}{path}", prefill_request, headers
262 + ):
263 + continue
264 +
265 + generator = self.forward_request(
266 + f"http://{decode_instance}{path}", decode_request, headers
267 + )
268 + return StreamingResponse(generator)
269 +
270 + async def create_completion(self, raw_request: Request):
271 + try:
272 + return await self._push_completion(raw_request, "/v1/completions")
273 + except HTTPException:
274 + raise
275 + except Exception:
276 + exc_info = sys.exc_info()
277 + print("Error occurred in disagg push proxy server")
278 + print(exc_info)
279 + raise
280 +
281 + async def create_chat_completion(self, raw_request: Request):
282 + try:
283 + return await self._push_completion(raw_request, "/v1/chat/completions")
284 + except HTTPException:
285 + raise
286 + except Exception:
287 + exc_info = sys.exc_info()
288 + error_messages = [str(e) for e in exc_info if e]
289 + print("Error occurred in disagg push proxy server")
290 + print(error_messages)
291 + return StreamingResponse(
292 + content=iter(error_messages), media_type="text/event-stream"
293 + )
294 +
295 +
296 +class PushProxyServer:
297 + def __init__(
298 + self,
299 + args: argparse.Namespace,
300 + scheduling_policy: SchedulingPolicy | None = None,
301 + create_completion: Callable[[Request], StreamingResponse] | None = None,
302 + create_chat_completion: Callable[[Request], StreamingResponse] | None = None,
303 + ):
304 + self.validate_parsed_serve_args(args)
305 + self.port = args.port
306 + self.proxy_instance = PushProxy(
307 + prefill_instances=[] if args.prefill is None else args.prefill,
308 + decode_instances=[] if args.decode is None else args.decode,
309 + model=args.model,
310 + scheduling_policy=(
311 + scheduling_policy
312 + if scheduling_policy is not None
313 + else RoundRobinSchedulingPolicy()
314 + ),
315 + prefill_engine_id=args.prefill_engine_id,
316 + prefill_kv_host=args.prefill_kv_host,
317 + prefill_side_channel_port=args.prefill_side_channel_port,
318 + prefill_tp_size=args.prefill_tp_size,
319 + custom_create_completion=create_completion,
320 + custom_create_chat_completion=create_chat_completion,
321 + )
322 +
323 + def validate_parsed_serve_args(self, args: argparse.Namespace):
324 + if not args.prefill:
325 + raise ValueError("Please specify at least one prefill node.")
326 + if not args.decode:
327 + raise ValueError("Please specify at least one decode node.")
328 + if not args.prefill_engine_id:
329 + raise ValueError(
330 + "--prefill-engine-id is required in push mode (it must match "
331 + "the engine_id passed to the prefill vLLM instance via "
332 + "--kv-transfer-config)."
333 + )
334 + if not args.prefill_kv_host:
335 + raise ValueError(
336 + "--prefill-kv-host is required in push mode (the IP / host "
337 + "that the prefill vLLM advertises on its NIXL side channel)."
338 + )
339 + self.validate_instances(args.prefill)
340 + self.validate_instances(args.decode)
341 +
342 + def validate_instances(self, instances: list):
343 + for instance in instances:
344 + if len(instance.split(":")) != 2:
345 + raise ValueError(f"Invalid instance format: {instance}")
346 + host, port = instance.split(":")
347 + try:
348 + if host != "localhost":
349 + ipaddress.ip_address(host)
350 + port = int(port)
351 + if not (0 < port < 65536):
352 + raise ValueError(f"Invalid port number in instance: {instance}")
353 + except Exception as e:
354 + raise ValueError(f"Invalid instance {instance}: {str(e)}") from e
355 +
356 + def run_server(self):
357 + app = FastAPI()
358 + app.include_router(self.proxy_instance.router)
359 + config = uvicorn.Config(app, port=self.port, loop="uvloop")
360 + server = uvicorn.Server(config)
361 + server.run()
362 +
363 +
364 +def parse_args():
365 + parser = argparse.ArgumentParser("vLLM disaggregated push-mode proxy server.")
366 + parser.add_argument("--model", "-m", type=str, required=True, help="Model name")
367 +
368 + parser.add_argument(
369 + "--prefill",
370 + "-p",
371 + type=str,
372 + nargs="+",
373 + help="List of prefill node URLs (host:port)",
374 + )
375 +
376 + parser.add_argument(
377 + "--decode",
378 + "-d",
379 + type=str,
380 + nargs="+",
381 + help="List of decode node URLs (host:port)",
382 + )
383 +
384 + parser.add_argument(
385 + "--port",
386 + type=int,
387 + default=8000,
388 + help="Server port number",
389 + )
390 +
391 + # Push-mode specific: P's coordinates that D needs in advance.
392 + parser.add_argument(
393 + "--prefill-engine-id",
394 + type=str,
395 + required=True,
396 + help=(
397 + "engine_id of the prefill vLLM instance (must match "
398 + "--kv-transfer-config engine_id on the prefill server)"
399 + ),
400 + )
401 + parser.add_argument(
402 + "--prefill-kv-host",
403 + type=str,
404 + required=True,
405 + help=(
406 + "IP / host the prefill vLLM advertises on its NIXL side "
407 + "channel (VLLM_NIXL_SIDE_CHANNEL_HOST)"
408 + ),
409 + )
410 + parser.add_argument(
411 + "--prefill-side-channel-port",
412 + type=int,
413 + default=5600,
414 + help="NIXL side channel port on the prefill node "
415 + "(VLLM_NIXL_SIDE_CHANNEL_PORT, default 5600)",
416 + )
417 + parser.add_argument(
418 + "--prefill-tp-size",
419 + type=int,
420 + default=1,
421 + help="Tensor parallel size of the prefill vLLM instance",
422 + )
423 + return parser.parse_args()
424 +
425 +
426 +if __name__ == "__main__":
427 + args = parse_args()
428 + proxy_server = PushProxyServer(args=args)
429 + proxy_server.run_server()
tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py
+4 -4
@@ -32,7 +32,7 @@ from unittest.mock import patch
32 import pytest
33
34 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
35 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import (
35 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
36 NixlConnector,
37 NixlConnectorMetadata,
38 )
@@ -436,7 +436,7 @@ def test_build_connector_meta_multiple_requests():
436
437
438 @patch(
439 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
439 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
440 FakeNixlWrapper,
441 )
442 def test_p_node_pull_kv_from_d(dist_init):
@@ -450,7 +450,7 @@ def test_p_node_pull_kv_from_d(dist_init):
450
451
452 @patch(
453 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
453 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
454 FakeNixlWrapper,
455 )
456 def test_p_node_pull_then_send_kv(dist_init):
@@ -472,7 +472,7 @@ def test_p_node_pull_then_send_kv(dist_init):
472
473
474 @patch(
475 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
475 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
476 FakeNixlWrapper,
477 )
478 def test_p_node_deferred_pull_on_no_handshake(dist_init):
tests/v1/kv_connector/unit/test_multi_connector.py
+5 -2
@@ -366,7 +366,10 @@ def test_multi_example_connector_consistency():
366
367
368 def _ignore_event_collection(events: list[str]) -> list[str]:
369 - return [event for event in events if event != "take_events"]
369 + # Filter out per-step polling hooks that the scheduler calls repeatedly
370 + # and which are not meaningful state transitions for these assertions.
371 + ignored = {"take_events", "has_pending_push_work"}
372 + return [event for event in events if event not in ignored]
373
374
375 def get_connector_events() -> dict[str, list[str]]:
@@ -1072,7 +1075,7 @@ def test_multi_connector_mixed_hma_disables_hybrid_kv_cache(monkeypatch):
1075 )
1076
1077 with patch(
1075 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1078 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1079 FakeNixlWrapper,
1080 ):
1081 llm = LLM(
tests/v1/kv_connector/unit/test_nixl_connector.py
+43 -36
@@ -344,7 +344,7 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv():
344
345
346 @patch(
347 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
347 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
348 FakeNixlWrapper,
349 )
350 def test_kv_transfer_handshake(dist_init):
@@ -560,7 +560,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
560
561 class TestNixlHandshake:
562 @patch(
563 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
563 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
564 FakeNixlWrapper,
565 )
566 def test_multi_xfer_one_engine(
@@ -643,7 +643,7 @@ class TestNixlHandshake:
643 connector.clear_connector_metadata()
644
645 @patch(
646 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
646 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
647 FakeNixlWrapper,
648 )
649 @pytest.mark.parametrize(
@@ -713,7 +713,7 @@ class TestNixlHandshake:
713 raise TimeoutError("Took too long to complete async handshake.")
714
715 @patch(
716 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
716 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
717 FakeNixlWrapper,
718 )
719 @pytest.mark.parametrize("local_tp_size", [1, 2])
@@ -725,7 +725,7 @@ class TestNixlHandshake:
725 remote configurations.
726 """
727 monkeypatch.setattr(
728 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size",
728 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size",
729 lambda: local_tp_size,
730 )
731
@@ -784,7 +784,7 @@ class TestNixlHandshake:
784 check_handshake(6)
785
786 @patch(
787 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
787 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
788 FakeNixlWrapper,
789 )
790 def test_prefill_tp_size_greater_than_decode_tp_size_mla(
@@ -887,7 +887,7 @@ class TestNixlHandshake:
887 assert req_id not in conn_p1.connector_worker._reqs_to_process
888
889 @patch(
890 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
890 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
891 FakeNixlWrapper,
892 )
893 def test_concurrent_load_kv(
@@ -952,7 +952,7 @@ class TestNixlHandshake:
952 raise TimeoutError("Took too long to complete async handshake.")
953
954 @patch(
955 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
955 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
956 FakeNixlWrapper,
957 )
958 def test_handshake_fails_on_kv_cache_layout_mismatch(
@@ -967,7 +967,7 @@ class TestNixlHandshake:
967 # Mock TP world size to 2 to force heterogeneous TP when
968 # remote_tp_size=1
969 with patch(
970 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501
970 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
971 return_value=2,
972 ):
973 # Initialize connector and worker (with fake NIXL wrapper)
@@ -1007,7 +1007,7 @@ class TestNixlHandshake:
1007 worker.add_remote_agent(meta, remote_tp_size=1)
1008
1009 @patch(
1010 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1010 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1011 FakeNixlWrapper,
1012 )
1013 def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental(
@@ -1022,7 +1022,7 @@ class TestNixlHandshake:
1022 # Mock TP world size to 2 to force heterogeneous TP when
1023 # remote_tp_size=1
1024 with patch(
1025 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501
1025 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1026 return_value=2,
1027 ):
1028 # Initialize connector and worker (with fake NIXL wrapper)
@@ -1064,7 +1064,7 @@ class TestNixlHandshake:
1064 worker.add_remote_agent(meta, remote_tp_size=1)
1065
1066 @patch(
1067 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1067 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1068 FakeNixlWrapper,
1069 )
1070 def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init):
@@ -1074,7 +1074,7 @@ class TestNixlHandshake:
1074 """
1075 vllm_config = create_vllm_config()
1076 with patch(
1077 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501
1077 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.get_tensor_model_parallel_world_size", # noqa: E501
1078 return_value=2,
1079 ):
1080 connector = NixlConnector(
@@ -1149,7 +1149,7 @@ class TestNixlHandshake:
1149 # we put here is important. First run ray, it will clean up the resources, then
1150 # the rest of the tests.
1151 @patch(
1152 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1152 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1153 FakeNixlWrapper,
1154 )
1155 def test_kv_connector_stats(default_vllm_config, dist_init):
@@ -1363,7 +1363,7 @@ def test_multi_kv_connector_stats_aggregation():
1363
1364
1365 @patch(
1366 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1366 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1367 FakeNixlWrapper,
1368 )
1369 def test_scheduler_kv_connector_stats_aggregation():
@@ -1428,7 +1428,7 @@ def test_scheduler_kv_connector_stats_aggregation():
1428
1429 @pytest.mark.parametrize("distributed_executor_backend", ["ray", None])
1430 @patch(
1431 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1431 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1432 FakeNixlWrapper,
1433 )
1434 def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
@@ -1615,7 +1615,7 @@ def test_register_kv_caches(
1615
1616 backend_cls = TritonAttentionBackend
1617
1618 - nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker"
1618 + nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
1619 nixl_connector = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector"
1620 with (
1621 patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
@@ -1865,15 +1865,17 @@ def test_kv_buffer_to_nixl_memory_types(
1865 _NIXL_SUPPORTED_DEVICE.update(FakePlatform.get_nixl_supported_devices())
1866
1867 with (
1868 - patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper"),
1868 patch(
1870 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.threading.Event"
1869 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper"
1870 ),
1871 patch(
1873 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.threading.Thread"
1872 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.threading.Event"
1873 ),
1874 patch(
1876 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.current_platform",
1875 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.threading.Thread"
1876 + ),
1877 + patch(
1878 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.current_platform",
1879 FakePlatform,
1880 ),
1881 patch(
@@ -1892,7 +1894,7 @@ def test_kv_buffer_to_nixl_memory_types(
1894
1895
1896 @patch(
1895 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1897 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1898 FakeNixlWrapper,
1899 )
1900 def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
@@ -1991,7 +1993,7 @@ def _setup_worker_with_remote_engine(
1993
1994
1995 @patch(
1994 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1996 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
1997 FakeNixlWrapper,
1998 )
1999 def test_engine_ttl_eviction(default_vllm_config, dist_init):
@@ -2026,7 +2028,7 @@ def test_engine_ttl_eviction(default_vllm_config, dist_init):
2028
2029
2030 @patch(
2029 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2031 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2032 FakeNixlWrapper,
2033 )
2034 def test_engine_ttl_disabled(default_vllm_config, dist_init):
@@ -2074,7 +2076,7 @@ def test_transfer_topology_unregister():
2076
2077
2078 @patch(
2077 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2079 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2080 FakeNixlWrapper,
2081 )
2082 def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_init):
@@ -2194,7 +2196,7 @@ class FailingNixlWrapper(FakeNixlWrapper):
2196
2197
2198 @patch(
2197 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2199 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2200 FailingNixlWrapper,
2201 )
2202 @pytest.mark.parametrize(
@@ -2284,10 +2286,13 @@ def test_transfer_failure_logging(
2286 slot_mapping={},
2287 )
2288
2287 - # Capture logs from the nixl.worker logger specifically
2289 + # Capture logs from the nixl connector loggers
2290 # vLLM loggers have propagate=False, so we need to capture directly
2291 nixl_logger = logging.getLogger(
2290 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker"
2292 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
2293 + )
2294 + pull_logger = logging.getLogger(
2295 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker"
2296 )
2297 captured_logs: list[logging.LogRecord] = []
2298
@@ -2298,6 +2303,7 @@ def test_transfer_failure_logging(
2303 handler = LogCapture()
2304 handler.setLevel(logging.ERROR)
2305 nixl_logger.addHandler(handler)
2306 + pull_logger.addHandler(handler)
2307
2308 try:
2309 connector.start_load_kv(dummy_ctx)
@@ -2313,6 +2319,7 @@ def test_transfer_failure_logging(
2319 connector.get_finished(finished_req_ids=set())
2320 finally:
2321 nixl_logger.removeHandler(handler)
2322 + pull_logger.removeHandler(handler)
2323
2324 # Print logs for manual comparison between commits
2325 error_logs = [r for r in captured_logs if r.levelno >= logging.ERROR]
@@ -2349,7 +2356,7 @@ def test_transfer_failure_logging(
2356
2357
2358 @patch(
2352 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2359 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2360 FailingNixlWrapper,
2361 )
2362 def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
@@ -2400,7 +2407,7 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
2407
2408
2409 @patch(
2403 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2410 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2411 FailingNixlWrapper,
2412 )
2413 def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init):
@@ -2454,7 +2461,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
2461
2462
2463 @patch(
2457 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2464 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2465 FailingNixlWrapper,
2466 )
2467 @pytest.mark.parametrize(
@@ -2597,7 +2604,7 @@ def test_failed_request_skips_kv_postprocessing(
2604 ],
2605 )
2606 @patch(
2600 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2607 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2608 FakeNixlWrapper,
2609 )
2610 def test_compatibility_hash_validation(
@@ -2706,7 +2713,7 @@ def test_compatibility_hash_validation(
2713 # Patch zmq_ctx to return our mock socket
2714 with (
2715 patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),
2709 - patch.object(nixl.worker, "zmq_ctx") as mock_zmq_ctx,
2716 + patch.object(nixl.base_worker, "zmq_ctx") as mock_zmq_ctx,
2717 ):
2718 mock_zmq_ctx.return_value.__enter__.return_value = mock_socket
2719
@@ -2740,7 +2747,7 @@ def test_compatibility_hash_validation(
2747 ],
2748 )
2749 @patch(
2743 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2750 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2751 FakeNixlWrapper,
2752 )
2753 def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario):
@@ -2806,7 +2813,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
2813 mock_socket.recv.return_value = msg_bytes
2814 with (
2815 patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),
2809 - patch.object(nixl.worker, "zmq_ctx") as mock_zmq_ctx,
2816 + patch.object(nixl.base_worker, "zmq_ctx") as mock_zmq_ctx,
2817 ):
2818 mock_zmq_ctx.return_value.__enter__.return_value = mock_socket
2819
@@ -2819,7 +2826,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
2826 )
2827
2828 @patch(
2822 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
2829 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2830 FakeNixlWrapper,
2831 )
2832 def test_mla_broadcast_notif_uses_remote_request_id(
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+6 -2
@@ -34,7 +34,9 @@ from .utils import (
34 (False, [0]),
35 ],
36 )
37 -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform")
37 +@patch(
38 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler.current_platform"
39 +)
40 def test_sw_sizes(mock_platform, swa_enabled, expected_sw_sizes):
41 """Test sw_sizes is correctly computed based on SWA enabled/disabled."""
42 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
@@ -782,7 +784,9 @@ def test_mamba_n1_p_side_truncation():
784 ],
785 ids=["fa_swa_mamba", "fa_swa_only", "fa_only"],
786 )
785 -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform")
787 +@patch(
788 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler.current_platform"
789 +)
790 def test_has_mamba_init(
791 mock_platform,
792 swa_enabled,
tests/v1/kv_connector/unit/test_nixl_push_connector.py new
+815
@@ -0,0 +1,815 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for NixlPushConnector (scheduler + worker).
4 +
5 +These tests cover the end-to-end mechanics of the push design without
6 +requiring a real NIXL agent or network:
7 +
8 +* Scheduler stages D registrations on ``update_state_after_alloc`` and
9 + P finished blocks on ``request_finished``.
10 +* ``build_connector_meta`` drains them onto
11 + ``meta.push_registrations`` / ``meta.push_finished_blocks``.
12 +* ``has_pending_push_work`` reports True/False over the lifecycle.
13 +* ``update_connector_output`` clears state on ``finished_sending`` and
14 + ``finished_recving``.
15 +* The worker matches D registrations against P finished blocks (both
16 + scenario directions) and forwards non-PUSH_REG NIXL notifs to the main
17 + thread's ``_get_new_notifs``.
18 +* ``get_finished`` enqueues evictions for the writer.
19 +"""
20 +
21 +from __future__ import annotations
22 +
23 +import logging
24 +import queue
25 +import threading
26 +import time
27 +from collections import defaultdict
28 +from typing import Any
29 +from unittest.mock import MagicMock, patch
30 +
31 +import msgspec
32 +
33 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
34 + PUSH_REG_NOTIF_PREFIX,
35 + NixlConnectorMetadata,
36 +)
37 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
38 + NixlPushConnectorWorker,
39 +)
40 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
41 + get_base_request_id,
42 +)
43 +from vllm.v1.outputs import KVConnectorOutput
44 +
45 +from .utils import make_nixl_push_scheduler
46 +
47 +# ----------------------------------------------------------------- #
48 +# Helpers / fakes #
49 +# ----------------------------------------------------------------- #
50 +
51 +
52 +def _make_request(
53 + *,
54 + request_id: str,
55 + is_d_side: bool = True,
56 + remote_engine_id: str = "prefill-engine",
57 + remote_request_id: str | None = None,
58 + remote_host: str = "10.0.0.1",
59 + remote_port: int = 5601,
60 + tp_size: int = 1,
61 + finished: bool = True,
62 +) -> MagicMock:
63 + """Build a minimal Request mock used by request_finished."""
64 + from vllm.v1.request import RequestStatus
65 +
66 + req = MagicMock()
67 + req.request_id = request_id
68 + req.num_computed_tokens = 64
69 +
70 + if is_d_side:
71 + # D-side request: do_remote_prefill=True -> prefill on a remote P.
72 + params: dict[str, Any] = {
73 + "do_remote_prefill": True,
74 + "do_remote_decode": False,
75 + "remote_engine_id": remote_engine_id,
76 + "remote_request_id": remote_request_id or f"prefill-{request_id}",
77 + "remote_host": remote_host,
78 + "remote_port": remote_port,
79 + "tp_size": tp_size,
80 + }
81 + else:
82 + # P-side request: do_remote_decode=True (we are the prefiller).
83 + params = {
84 + "do_remote_prefill": False,
85 + "do_remote_decode": True,
86 + }
87 + req.kv_transfer_params = params
88 + req.status = (
89 + RequestStatus.FINISHED_LENGTH_CAPPED if finished else RequestStatus.RUNNING
90 + )
91 + return req
92 +
93 +
94 +class _BlocksMock:
95 + """Minimal stand-in for ``KVCacheBlocks`` used in update_state_after_alloc."""
96 +
97 + def __init__(self, block_ids: tuple[list[int], ...]):
98 + self._block_ids = block_ids
99 +
100 + def get_unhashed_block_ids_all_groups(self) -> tuple[list[int], ...]:
101 + return self._block_ids
102 +
103 +
104 +def _stub_sw_clipping(scheduler) -> None:
105 + """Make ``get_sw_clipped_blocks`` a passthrough so tests don't need
106 + the full sliding-window machinery."""
107 + scheduler.get_sw_clipped_blocks = lambda block_ids: block_ids
108 +
109 +
110 +# ----------------------------------------------------------------- #
111 +# Scheduler-side tests #
112 +# ----------------------------------------------------------------- #
113 +
114 +
115 +class TestPushScheduler:
116 + def test_d_side_update_state_after_alloc_stages_registration(self):
117 + """D scheduler stashes registration data + arms watchdog deadline."""
118 + sched = make_nixl_push_scheduler()
119 + _stub_sw_clipping(sched)
120 +
121 + request = _make_request(request_id="req-d-1")
122 + blocks = _BlocksMock(block_ids=([10, 11, 12],))
123 +
124 + sched.update_state_after_alloc(request, blocks, num_external_tokens=48)
125 +
126 + assert request.request_id in sched._push_pending_registrations
127 + reg = sched._push_pending_registrations[request.request_id]
128 + # ``request_id`` is D's own vLLM request id; plus our own (D) coords.
129 + assert reg["request_id"] == request.request_id
130 + assert reg["decode_engine_id"] == sched.engine_id
131 + assert reg["decode_host"] == sched.side_channel_host
132 + assert reg["decode_port"] == sched.side_channel_port
133 + assert reg["local_block_ids"] == ([10, 11, 12],)
134 + assert reg["remote_engine_id"] == "prefill-engine"
135 +
136 + # Watchdog deadline set in the future.
137 + deadline = sched._push_registration_deadlines[request.request_id]
138 + assert deadline > time.perf_counter()
139 + # do_remote_prefill flipped off so the request isn't reprocessed.
140 + assert request.kv_transfer_params["do_remote_prefill"] is False
141 + # Tracked as awaiting a recv.
142 + assert request.request_id in sched._reqs_need_recv
143 +
144 + def test_p_side_request_finished_stages_blocks(self):
145 + """P scheduler pushes blocks into both _finished_request_blocks (lease)
146 + and _newly_finished_push_blocks (metadata for next step)."""
147 + sched = make_nixl_push_scheduler()
148 + _stub_sw_clipping(sched)
149 +
150 + request = _make_request(request_id="req-p-1", is_d_side=False)
151 + block_ids = ([20, 21, 22, 23],)
152 +
153 + delay, ret_params = sched.request_finished(request, block_ids)
154 +
155 + assert delay is True
156 + assert ret_params is not None
157 + assert ret_params["do_remote_prefill"] is True
158 + assert ret_params["do_remote_decode"] is False
159 + assert request.request_id in sched._finished_request_blocks
160 + assert request.request_id in sched._newly_finished_push_blocks
161 + assert request.request_id in sched._reqs_need_send # lease armed
162 +
163 + def test_build_connector_meta_drains_both_sides(self):
164 + """meta.push_registrations and meta.push_finished_blocks are filled
165 + from the staging dicts and the staging dicts are cleared."""
166 + sched = make_nixl_push_scheduler()
167 + _stub_sw_clipping(sched)
168 +
169 + # Stage one D registration and one P finished entry.
170 + d_req = _make_request(request_id="req-d-9")
171 + sched.update_state_after_alloc(
172 + d_req, _BlocksMock(([1, 2, 3],)), num_external_tokens=48
173 + )
174 + p_req = _make_request(request_id="req-p-9", is_d_side=False)
175 + sched.request_finished(p_req, ([4, 5, 6],))
176 +
177 + scheduler_output = MagicMock()
178 + scheduler_output.scheduled_new_reqs = []
179 + scheduler_output.scheduled_cached_reqs = MagicMock(
180 + req_ids=[], resumed_req_ids=set()
181 + )
182 +
183 + # Patch parent build_connector_meta so we don't have to set up
184 + # all the base scheduler plumbing.
185 + with patch.object(
186 + sched.__class__.__mro__[1],
187 + "build_connector_meta",
188 + return_value=NixlConnectorMetadata(),
189 + ):
190 + meta = sched.build_connector_meta(scheduler_output)
191 +
192 + assert isinstance(meta, NixlConnectorMetadata)
193 + assert "req-d-9" in meta.push_registrations
194 + assert "req-p-9" in meta.push_finished_blocks
195 + # Staging dicts cleared.
196 + assert sched._push_pending_registrations == {}
197 + assert sched._newly_finished_push_blocks == {}
198 + # Lease bookkeeping kept until the WRITE completes.
199 + assert "req-p-9" in sched._finished_request_blocks
200 +
201 + def test_has_pending_push_work_lifecycle(self):
202 + sched = make_nixl_push_scheduler()
203 + _stub_sw_clipping(sched)
204 +
205 + assert sched.has_pending_push_work() is False
206 +
207 + # P finished blocks waiting for WRITE completion.
208 + p_req = _make_request(request_id="req-p-7", is_d_side=False)
209 + sched.request_finished(p_req, ([0, 1],))
210 + assert sched.has_pending_push_work() is True
211 +
212 + # Drain via build_connector_meta - lease still pending until WRITE.
213 + scheduler_output = MagicMock()
214 + scheduler_output.scheduled_new_reqs = []
215 + scheduler_output.scheduled_cached_reqs = MagicMock(
216 + req_ids=[], resumed_req_ids=set()
217 + )
218 + with patch.object(
219 + sched.__class__.__mro__[1],
220 + "build_connector_meta",
221 + return_value=NixlConnectorMetadata(),
222 + ):
223 + sched.build_connector_meta(scheduler_output)
224 + # Lease is pending until WRITE completes -> still True.
225 + assert sched.has_pending_push_work() is True
226 +
227 + # Simulate WRITE completion via update_connector_output.
228 + sched.update_connector_output(
229 + KVConnectorOutput(
230 + finished_sending={"req-p-7"},
231 + finished_recving=set(),
232 + invalid_block_ids=set(),
233 + )
234 + )
235 + assert sched.has_pending_push_work() is False
236 +
237 + def test_update_connector_output_clears_lease_and_watchdog(self):
238 + sched = make_nixl_push_scheduler()
239 + _stub_sw_clipping(sched)
240 +
241 + d_req = _make_request(request_id="req-d-x")
242 + sched.update_state_after_alloc(
243 + d_req, _BlocksMock(([1, 2],)), num_external_tokens=32
244 + )
245 + p_req = _make_request(request_id="req-p-x", is_d_side=False)
246 + sched.request_finished(p_req, ([3, 4],))
247 +
248 + sched.update_connector_output(
249 + KVConnectorOutput(
250 + finished_sending={"req-p-x"},
251 + finished_recving={"req-d-x"},
252 + invalid_block_ids=set(),
253 + )
254 + )
255 + assert "req-p-x" not in sched._finished_request_blocks
256 + assert "req-d-x" not in sched._push_registration_deadlines
257 +
258 + def test_registration_watchdog_expires(self, caplog):
259 + """Stale D registrations whose deadline has passed are dropped at
260 + ``build_connector_meta`` time."""
261 + # Watchdog logs a WARNING when it drops the stale entry; that's
262 + # what this test is verifying, so silence it in the test report.
263 + caplog.set_level(
264 + logging.CRITICAL,
265 + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler"),
266 + )
267 + sched = make_nixl_push_scheduler()
268 + _stub_sw_clipping(sched)
269 +
270 + d_req = _make_request(request_id="req-d-stale")
271 + sched.update_state_after_alloc(
272 + d_req, _BlocksMock(([7, 8],)), num_external_tokens=32
273 + )
274 + # Force the deadline into the past.
275 + sched._push_registration_deadlines[d_req.request_id] = time.perf_counter() - 1.0
276 +
277 + scheduler_output = MagicMock()
278 + scheduler_output.scheduled_new_reqs = []
279 + scheduler_output.scheduled_cached_reqs = MagicMock(
280 + req_ids=[], resumed_req_ids=set()
281 + )
282 + with patch.object(
283 + sched.__class__.__mro__[1],
284 + "build_connector_meta",
285 + return_value=NixlConnectorMetadata(),
286 + ):
287 + meta = sched.build_connector_meta(scheduler_output)
288 +
289 + assert d_req.request_id not in sched._push_registration_deadlines
290 + assert d_req.request_id not in sched._push_pending_registrations
291 + assert d_req.request_id not in meta.push_registrations
292 +
293 +
294 +# ----------------------------------------------------------------- #
295 +# Worker-side tests #
296 +# ----------------------------------------------------------------- #
297 +
298 +
299 +class _StubWriterWorker(NixlPushConnectorWorker):
300 + """Construct a worker without invoking ``__init__`` so we can drive
301 + the matching/notif logic without bringing up NIXL or torch."""
302 +
303 + @classmethod
304 + def fresh(cls) -> _StubWriterWorker:
305 + w = object.__new__(cls)
306 +
307 + # Push-specific state managed by NixlPushConnectorWorker.
308 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
309 + ReqId,
310 + TransferHandle,
311 + )
312 +
313 + w._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list)
314 + w._sending_transfers_lock = threading.Lock()
315 + w._push_finished_blocks = {}
316 + w._pending_d_registrations = {}
317 + w._reg_send_inbox = queue.Queue()
318 + w._finished_blocks_inbox = queue.Queue()
319 + w._pending_completion_notifs = queue.Queue()
320 + w._evict_finished_inbox = queue.Queue()
321 + w._push_writer_wake = threading.Event()
322 + w._push_writer_stop = threading.Event()
323 + w._push_writer_thread = None
324 +
325 + # Base worker fields touched by start_load_kv / _get_new_notifs.
326 + w._recving_metadata = {}
327 + w._recving_transfers = defaultdict(list)
328 + w._reqs_to_process = set()
329 + w._reqs_to_send = {}
330 + w.consumer_notification_counts_by_req = defaultdict(int)
331 + w.tp_rank = 0
332 + w.world_size = 1
333 + w.engine_id = "test-decode-engine"
334 + w._remote_agents = {}
335 +
336 + # Track _do_start_push_kv invocations.
337 + calls: list[tuple[str, Any, dict[str, Any]]] = []
338 + w.start_push_calls = calls
339 + return w
340 +
341 + def _do_start_push_kv(
342 + self,
343 + request_id: str,
344 + local_block_ids,
345 + registration_data: dict[str, Any],
346 + ) -> None: # pragma: no cover - exercised through tests
347 + # Track the call instead of issuing real WRITEs.
348 + self.start_push_calls.append((request_id, local_block_ids, registration_data))
349 +
350 +
351 +def _registration_data(
352 + request_id: str,
353 + *,
354 + decode_engine_id: str = "decode-engine",
355 + decode_host: str = "10.0.0.2",
356 + decode_port: int = 5602,
357 + decode_tp_size: int = 1,
358 + local_block_ids=((100, 101, 102),),
359 + remote_engine_id: str = "prefill-engine",
360 + remote_host: str = "10.0.0.1",
361 + remote_port: int = 5601,
362 + remote_tp_size: int = 1,
363 +) -> dict[str, Any]:
364 + return {
365 + "request_id": request_id,
366 + "decode_engine_id": decode_engine_id,
367 + "decode_host": decode_host,
368 + "decode_port": decode_port,
369 + "decode_tp_size": decode_tp_size,
370 + "local_block_ids": local_block_ids,
371 + "remote_engine_id": remote_engine_id,
372 + "remote_host": remote_host,
373 + "remote_port": remote_port,
374 + "remote_tp_size": remote_tp_size,
375 + }
376 +
377 +
378 +class TestPushWriterMatching:
379 + def test_handle_push_reg_matches_existing_finished_blocks(self):
380 + """PUSH_REG arrives second (P finished first): match + fire."""
381 + w = _StubWriterWorker.fresh()
382 + # P had already finished; its blocks were stashed via metadata.
383 + w._push_finished_blocks["req-A"] = ([200, 201, 202],)
384 +
385 + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(
386 + _registration_data("req-A")
387 + )
388 + w._handle_push_reg_notif(notif)
389 +
390 + assert len(w.start_push_calls) == 1
391 + rid, blocks, reg = w.start_push_calls[0]
392 + assert rid == "req-A"
393 + assert blocks == ([200, 201, 202],)
394 + assert reg["decode_engine_id"] == "decode-engine"
395 + # Finished blocks consumed.
396 + assert "req-A" not in w._push_finished_blocks
397 + assert w._pending_d_registrations == {}
398 +
399 + def test_handle_push_reg_stashes_when_no_finished_blocks_yet(self):
400 + """PUSH_REG arrives first (D registered first): stash, no fire."""
401 + w = _StubWriterWorker.fresh()
402 +
403 + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(
404 + _registration_data("req-B")
405 + )
406 + w._handle_push_reg_notif(notif)
407 +
408 + assert len(w.start_push_calls) == 0
409 + assert "req-B" in w._pending_d_registrations
410 +
411 + def test_handle_push_reg_matches_after_stripping_random_suffix(self):
412 + """P and D assign the same logical request the same
413 + ``cmpl-<uuid>-<index>`` but different per-engine random suffixes;
414 + the writer should still match P's finished blocks via the
415 + suffix-stripping fallback in ``_pop_matching_finished_blocks``.
416 + """
417 + w = _StubWriterWorker.fresh()
418 + # Same base id + completion index; differ only in the trailing
419 + # ``-<8 hex>`` randomization suffix.
420 + p_id = "cmpl-12345678-aaaa-bbbb-cccc-1234567890ab-0-aaaaaaaa"
421 + d_id = "cmpl-12345678-aaaa-bbbb-cccc-1234567890ab-0-bbbbbbbb"
422 + # Sanity: same base id under the helper used by the connector.
423 + assert get_base_request_id(p_id) == get_base_request_id(d_id)
424 +
425 + w._push_finished_blocks[p_id] = ([1, 2, 3],)
426 + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(_registration_data(d_id))
427 + w._handle_push_reg_notif(notif)
428 +
429 + # Suffix-stripped fallback matched and fired.
430 + assert len(w.start_push_calls) == 1
431 + assert w.start_push_calls[0][0] == p_id
432 + assert p_id not in w._push_finished_blocks
433 +
434 + def test_handle_push_reg_drops_malformed(self, caplog):
435 + # The writer logs WARNING/ERROR when it sees these bad payloads;
436 + # that's the desired behavior, so suppress the noise from test
437 + # output rather than letting it look like a failure.
438 + caplog.set_level(
439 + logging.CRITICAL,
440 + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"),
441 + )
442 + w = _StubWriterWorker.fresh()
443 + # Missing request_id -> should drop without raising.
444 + bad = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode({"decode_engine_id": "x"})
445 + w._handle_push_reg_notif(bad)
446 + assert w._pending_d_registrations == {}
447 + assert w.start_push_calls == []
448 +
449 + # Undecodable payload also dropped.
450 + w._handle_push_reg_notif(PUSH_REG_NOTIF_PREFIX + b"\xff\xff\xff")
451 + assert w.start_push_calls == []
452 +
453 +
454 +class TestPushWriterStartLoadKv:
455 + def test_finished_blocks_inbox_matches_stashed_registration(self):
456 + """Run the writer-loop's finished-blocks drain against a
457 + pre-populated _pending_d_registrations entry."""
458 + w = _StubWriterWorker.fresh()
459 + w._pending_d_registrations["req-C"] = _registration_data("req-C")
460 +
461 + # Simulate start_load_kv enqueuing finished blocks.
462 + w._finished_blocks_inbox.put(("req-C", ([10, 11, 12],)))
463 +
464 + # Drain like the writer loop does.
465 + while True:
466 + try:
467 + rid, blocks = w._finished_blocks_inbox.get_nowait()
468 + except queue.Empty:
469 + break
470 + matched = w._pop_matching_registration(rid)
471 + if matched is not None:
472 + w._do_start_push_kv(rid, blocks, matched)
473 + else:
474 + w._push_finished_blocks[rid] = blocks
475 +
476 + assert len(w.start_push_calls) == 1
477 + assert w.start_push_calls[0][0] == "req-C"
478 + assert "req-C" not in w._pending_d_registrations
479 +
480 + def test_start_load_kv_enqueues_to_writer(self):
481 + """``start_load_kv`` should hand registrations + finished blocks
482 + to the writer queues without doing matching itself."""
483 + w = _StubWriterWorker.fresh()
484 + # Stub heartbeats to a no-op; tests don't exercise the heartbeat
485 + # path here.
486 + w._send_heartbeats = lambda metadata: None
487 + # Stub logical-to-kernel mapping used by reqs_to_recv.
488 + w._logical_to_kernel_block_ids = lambda x: x
489 +
490 + meta = NixlConnectorMetadata()
491 + meta.push_registrations = {
492 + "req-D": _registration_data("req-D"),
493 + }
494 + meta.push_finished_blocks = {
495 + "req-E": ([5, 6, 7],),
496 + }
497 +
498 + w.start_load_kv(meta)
499 +
500 + # Things are queued for the writer; nothing fires yet.
501 + assert w._reg_send_inbox.qsize() == 1
502 + assert w._finished_blocks_inbox.qsize() == 1
503 + assert w._push_writer_wake.is_set()
504 + assert w.start_push_calls == []
505 +
506 +
507 +class TestPushWriterNotifs:
508 + def test_get_new_notifs_processes_forwarded_completion_notif(self):
509 + """Non-PUSH_REG notifs forwarded by the writer thread are drained
510 + on the engine main thread inside ``_get_new_notifs``."""
511 + w = _StubWriterWorker.fresh()
512 + # Pretend the writer thread already forwarded a completion notif
513 + # for a request whose KV is being received.
514 + request_id = "req-recv-1"
515 + w._recving_metadata[request_id] = MagicMock()
516 + # Compose the standard completion notif: req_id:tp_size.
517 + notif_msg = f"{request_id}:1".encode()
518 + w._pending_completion_notifs.put(notif_msg)
519 +
520 + # transfer_topo is consulted only for the producer-side path; we
521 + # make it a MagicMock because the D-side branch returns early.
522 + w.transfer_topo = MagicMock()
523 +
524 + notified = w._get_new_notifs()
525 +
526 + # Notif consumed; D-side just touches _recving_transfers.
527 + assert notified == set()
528 + assert request_id in w._recving_transfers
529 +
530 + def test_get_finished_evicts_completed_state(self):
531 + """``get_finished`` should enqueue evictions and wake the writer."""
532 + w = _StubWriterWorker.fresh()
533 +
534 + # Stub the base ``get_finished`` to return one done_sending entry.
535 + # Patch via the MRO's parent class.
536 + with patch.object(
537 + NixlPushConnectorWorker.__mro__[1],
538 + "get_finished",
539 + return_value=({"req-done"}, set()),
540 + ):
541 + done_sending, done_recving = w.get_finished()
542 +
543 + assert "req-done" in done_sending
544 + assert done_recving == set()
545 + # Eviction enqueued for the writer.
546 + evicted = []
547 + while True:
548 + try:
549 + evicted.append(w._evict_finished_inbox.get_nowait())
550 + except queue.Empty:
551 + break
552 + assert evicted == ["req-done"]
553 + assert w._push_writer_wake.is_set()
554 +
555 +
556 +# ----------------------------------------------------------------- #
557 +# Negative / error-path tests #
558 +# ----------------------------------------------------------------- #
559 +
560 +
561 +class TestPushSchedulerNegative:
562 + """Failure / no-op paths on the scheduler side."""
563 +
564 + def test_update_state_after_alloc_no_kv_transfer_params_is_noop(self):
565 + """Requests without kv_transfer_params must not register anything."""
566 + sched = make_nixl_push_scheduler()
567 + _stub_sw_clipping(sched)
568 +
569 + request = MagicMock()
570 + request.request_id = "req-no-params"
571 + request.kv_transfer_params = None
572 +
573 + sched.update_state_after_alloc(
574 + request, _BlocksMock(([1, 2, 3],)), num_external_tokens=64
575 + )
576 +
577 + assert sched._push_pending_registrations == {}
578 + assert sched._push_registration_deadlines == {}
579 + assert sched._reqs_need_recv == {}
580 +
581 + def test_update_state_after_alloc_zero_external_tokens_does_not_register(self):
582 + """num_external_tokens=0 should not stage a D registration."""
583 + sched = make_nixl_push_scheduler()
584 + _stub_sw_clipping(sched)
585 +
586 + request = _make_request(request_id="req-zero-ext")
587 + sched.update_state_after_alloc(
588 + request, _BlocksMock(([1, 2, 3],)), num_external_tokens=0
589 + )
590 +
591 + assert sched._push_pending_registrations == {}
592 + assert sched._push_registration_deadlines == {}
593 +
594 + def test_request_finished_unfinished_status_does_not_stage(self):
595 + """If a request is still RUNNING, request_finished must not stash
596 + blocks for the worker (no push needed)."""
597 + sched = make_nixl_push_scheduler()
598 + _stub_sw_clipping(sched)
599 +
600 + request = _make_request(
601 + request_id="req-running", is_d_side=False, finished=False
602 + )
603 +
604 + delay, ret = sched.request_finished(request, ([1, 2, 3],))
605 +
606 + assert delay is False
607 + assert ret is None
608 + assert sched._finished_request_blocks == {}
609 + assert sched._newly_finished_push_blocks == {}
610 +
611 + def test_request_finished_empty_blocks_does_not_arm_lease(self):
612 + """Empty block-id groups should still complete cleanly without
613 + arming the lease/finished maps."""
614 + sched = make_nixl_push_scheduler()
615 + _stub_sw_clipping(sched)
616 +
617 + request = _make_request(request_id="req-empty", is_d_side=False)
618 + delay, ret = sched.request_finished(request, ((),))
619 +
620 + assert delay is False
621 + assert ret is not None
622 + assert "req-empty" not in sched._finished_request_blocks
623 + assert "req-empty" not in sched._newly_finished_push_blocks
624 + assert "req-empty" not in sched._reqs_need_send
625 +
626 + def test_update_connector_output_unknown_request_is_noop(self):
627 + """Idempotent cleanup: clearing a request that was never staged
628 + must not raise or mutate other state."""
629 + sched = make_nixl_push_scheduler()
630 + _stub_sw_clipping(sched)
631 +
632 + # Stage one real request to ensure it's NOT touched.
633 + live = _make_request(request_id="req-live", is_d_side=False)
634 + sched.request_finished(live, ([1],))
635 +
636 + sched.update_connector_output(
637 + KVConnectorOutput(
638 + finished_sending={"unknown-1"},
639 + finished_recving={"unknown-2"},
640 + invalid_block_ids=set(),
641 + )
642 + )
643 +
644 + # Live entry untouched.
645 + assert "req-live" in sched._finished_request_blocks
646 +
647 +
648 +class TestPushWriterNegative:
649 + """Failure / drop / idempotence paths in the writer thread."""
650 +
651 + def test_pop_matching_registration_returns_none_when_empty(self):
652 + w = _StubWriterWorker.fresh()
653 + assert w._pop_matching_registration("nope") is None
654 +
655 + def test_pop_matching_finished_blocks_returns_none_when_empty(self):
656 + w = _StubWriterWorker.fresh()
657 + assert w._pop_matching_finished_blocks("nope") is None
658 +
659 + def test_pop_matching_registration_no_match_when_base_ids_differ(self):
660 + """A registration whose base id (after stripping the random suffix)
661 + does NOT match the lookup request_id must not be popped."""
662 + w = _StubWriterWorker.fresh()
663 + # Two unrelated requests: different base UUIDs, so stripping the
664 + # trailing ``-<8 hex>`` suffix still yields different base ids.
665 + unrelated_d = "cmpl-aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa-0-11111111"
666 + lookup = "cmpl-bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb-0-22222222"
667 + assert get_base_request_id(unrelated_d) != get_base_request_id(lookup)
668 +
669 + w._pending_d_registrations[unrelated_d] = _registration_data(unrelated_d)
670 + result = w._pop_matching_registration(lookup)
671 + assert result is None
672 + # Original entry untouched.
673 + assert unrelated_d in w._pending_d_registrations
674 +
675 + def test_handle_push_reg_with_non_dict_payload_is_dropped(self, caplog):
676 + """msgpack-encoded non-dict payload (e.g. a list) should be
677 + dropped without raising."""
678 + caplog.set_level(
679 + logging.CRITICAL,
680 + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"),
681 + )
682 + w = _StubWriterWorker.fresh()
683 + bad = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode([1, 2, 3])
684 + w._handle_push_reg_notif(bad)
685 + assert w._pending_d_registrations == {}
686 + assert w.start_push_calls == []
687 +
688 + def test_handle_push_reg_with_non_string_request_id_is_dropped(self, caplog):
689 + """request_id must be a str; integers, None, etc. must drop."""
690 + caplog.set_level(
691 + logging.CRITICAL,
692 + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"),
693 + )
694 + w = _StubWriterWorker.fresh()
695 + for bogus_rid in (123, None, 4.5, b"bytes-not-str"):
696 + payload = _registration_data("placeholder")
697 + payload["request_id"] = bogus_rid # type: ignore[assignment]
698 + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(payload)
699 + w._handle_push_reg_notif(notif)
700 + assert w._pending_d_registrations == {}
701 + assert w.start_push_calls == []
702 +
703 + def test_handle_push_reg_idempotent_for_same_request_id(self):
704 + """Receiving the same PUSH_REG twice (e.g. P retries after a
705 + flake) keeps the entry staged exactly once and never fires."""
706 + w = _StubWriterWorker.fresh()
707 + notif = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(
708 + _registration_data("req-dup")
709 + )
710 + w._handle_push_reg_notif(notif)
711 + w._handle_push_reg_notif(notif)
712 + assert "req-dup" in w._pending_d_registrations
713 + assert len(w._pending_d_registrations) == 1
714 + assert w.start_push_calls == []
715 +
716 + def test_get_finished_enqueues_eviction_for_each_done_request(self):
717 + """``get_finished`` must enqueue an eviction for every request
718 + in ``done_sending`` so the writer can drop stale matching state.
719 + Unlike the happy-path test, this verifies the *cardinality*: N
720 + completed requests -> N evictions, in order."""
721 + w = _StubWriterWorker.fresh()
722 + with patch.object(
723 + NixlPushConnectorWorker.__mro__[1],
724 + "get_finished",
725 + return_value=({"req-1", "req-2", "req-3"}, set()),
726 + ):
727 + done_sending, _ = w.get_finished()
728 + assert done_sending == {"req-1", "req-2", "req-3"}
729 +
730 + evicted: list[str] = []
731 + while True:
732 + try:
733 + evicted.append(w._evict_finished_inbox.get_nowait())
734 + except queue.Empty:
735 + break
736 + assert sorted(evicted) == ["req-1", "req-2", "req-3"]
737 +
738 + def test_get_finished_with_no_completions_does_not_enqueue_eviction(self):
739 + """If there's nothing newly done, no eviction should be enqueued.
740 + The wake event IS still set because ``get_finished`` always wakes
741 + the writer to drain notifs."""
742 + w = _StubWriterWorker.fresh()
743 + with patch.object(
744 + NixlPushConnectorWorker.__mro__[1],
745 + "get_finished",
746 + return_value=(set(), set()),
747 + ):
748 + done_sending, done_recving = w.get_finished()
749 + assert done_sending == set()
750 + assert done_recving == set()
751 + assert w._evict_finished_inbox.qsize() == 0
752 + # Wake set so the writer drains NIXL notifs even when idle.
753 + assert w._push_writer_wake.is_set()
754 +
755 + def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog):
756 + """A completion notif for a request the worker doesn't know
757 + about should be logged but not crash."""
758 + caplog.set_level(
759 + logging.CRITICAL,
760 + logger=("vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker"),
761 + )
762 + w = _StubWriterWorker.fresh()
763 + w.transfer_topo = MagicMock()
764 + # Forward a completion notif for an unknown request_id.
765 + w._pending_completion_notifs.put(b"never-heard-of-you:1")
766 +
767 + notified = w._get_new_notifs()
768 + assert notified == set()
769 + # Did not register anywhere.
770 + assert "never-heard-of-you" not in w._recving_transfers
771 +
772 + def test_start_load_kv_with_empty_metadata_is_noop(self):
773 + """Empty metadata must not wake the writer or enqueue anything."""
774 + w = _StubWriterWorker.fresh()
775 + w._send_heartbeats = lambda metadata: None
776 + w._logical_to_kernel_block_ids = lambda x: x
777 +
778 + meta = NixlConnectorMetadata()
779 + w.start_load_kv(meta)
780 +
781 + assert w._reg_send_inbox.qsize() == 0
782 + assert w._finished_blocks_inbox.qsize() == 0
783 + # Wake should NOT be set if there was nothing to push.
784 + assert not w._push_writer_wake.is_set()
785 +
786 + def test_get_new_notifs_extends_lease_on_heartbeat(self):
787 + """``HB:`` notifs forwarded by the writer thread must extend the
788 + leases of tracked P-side requests on the engine main thread, and
789 + ignore request IDs that aren't being tracked."""
790 + w = _StubWriterWorker.fresh()
791 + w.transfer_topo = MagicMock()
792 + # _handle_heartbeat reads ``self._lease_extension`` (set in the
793 + # real ``__init__``).
794 + w._lease_extension = 10
795 +
796 + # Tracked P-side requests with a lease about to expire.
797 + old_expiry = time.perf_counter() - 5.0
798 + w._reqs_to_send["req-a"] = old_expiry
799 + w._reqs_to_send["req-b"] = old_expiry
800 +
801 + # Forwarded heartbeat covers a tracked request, an unknown one,
802 + # and another tracked one.
803 + w._pending_completion_notifs.put(b"HB:req-a,req-unknown,req-b")
804 +
805 + notified = w._get_new_notifs()
806 + assert notified == set()
807 +
808 + # Tracked leases were renewed strictly forward in time.
809 + now = time.perf_counter()
810 + for rid in ("req-a", "req-b"):
811 + assert w._reqs_to_send[rid] > old_expiry
812 + # New expiry must be roughly now + _lease_extension.
813 + assert w._reqs_to_send[rid] >= now
814 + # Unknown request must not be inserted by the heartbeat path.
815 + assert "req-unknown" not in w._reqs_to_send
tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py
+1 -1
@@ -44,7 +44,7 @@ from vllm.v1.simple_kv_offload.metadata import (
44 )
45
46 NIXL_WRAPPER_PATCH = (
47 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper"
47 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper"
48 )
49
50
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+3 -1
@@ -587,7 +587,9 @@ def test_cannot_recv():
587 assert_scheduler_empty(scheduler)
588
589
590 -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform")
590 +@patch(
591 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler.current_platform"
592 +)
593 def test_p_side_chunked_prefill_mamba(mock_platform):
594 """P-side integration: Mamba N-1 truncation + chunked prefill completes.
595
tests/v1/kv_connector/unit/utils.py
+63
@@ -524,3 +524,66 @@ def make_nixl_scheduler(
524 sched.blocks_per_sw = []
525 sched.is_bidirectional_kv_xfer_enabled = False
526 return sched
527 +
528 +
529 +def make_nixl_push_scheduler(
530 + *,
531 + decoder_kv_blocks_ttl: float = 30.0,
532 + push_registration_timeout: float | None = None,
533 + is_bidirectional_kv_xfer_enabled: bool = False,
534 + has_mamba: bool = False,
535 +):
536 + """Create a NixlPushConnectorScheduler via __new__ (skipping __init__).
537 +
538 + The push scheduler can't reuse :func:`make_nixl_scheduler` because it
539 + is a different class (``NixlPushConnectorScheduler`` vs
540 + ``NixlConnectorScheduler``) and carries push-specific state. Only the
541 + fields touched by the unit tests are populated.
542 + """
543 + from unittest.mock import MagicMock
544 +
545 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import (
546 + NixlPushConnectorScheduler,
547 + )
548 +
549 + sched = object.__new__(NixlPushConnectorScheduler)
550 +
551 + # Base scheduler fields (shared with pull / heartbeat path).
552 + sched._reqs_need_recv = {}
553 + sched._reqs_need_send = {}
554 + sched._reqs_in_batch = set()
555 + sched._reqs_not_processed = set()
556 + sched._reqs_need_save = {}
557 + sched._kv_lease_duration = 30
558 + sched.decoder_kv_blocks_ttl = decoder_kv_blocks_ttl
559 + sched.use_host_buffer = False
560 + sched.engine_id = "decode-engine"
561 + sched.side_channel_host = "127.0.0.1"
562 + sched.side_channel_port = 5600
563 + sched.is_bidirectional_kv_xfer_enabled = is_bidirectional_kv_xfer_enabled
564 + sched._has_mamba = has_mamba
565 +
566 + # vllm_config is consulted for parallel_config.tensor_parallel_size.
567 + vllm_config = MagicMock()
568 + vllm_config.parallel_config.tensor_parallel_size = 1
569 + sched.vllm_config = vllm_config
570 +
571 + # Push-specific state.
572 + sched._push_pending_registrations = {}
573 + sched._push_registration_deadlines = {}
574 + sched._finished_request_blocks = {}
575 + sched._newly_finished_push_blocks = {}
576 + sched._push_registration_timeout = (
577 + push_registration_timeout
578 + if push_registration_timeout is not None
579 + else decoder_kv_blocks_ttl
580 + )
581 +
582 + # Heartbeat fields touched by base request_finished /
583 + # update_connector_output.
584 + sched._heartbeat_by_engine = {}
585 + sched._heartbeat_req_engine = {}
586 + sched._last_heartbeat_time = 0.0
587 + sched.blocks_per_sw = []
588 +
589 + return sched
vllm/distributed/kv_transfer/kv_connector/factory.py
+12
@@ -179,6 +179,18 @@ KVConnectorFactory.register_connector(
179 "NixlConnector",
180 )
181
182 +KVConnectorFactory.register_connector(
183 + "NixlPullConnector",
184 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl",
185 + "NixlPullConnector",
186 +)
187 +
188 +KVConnectorFactory.register_connector(
189 + "NixlPushConnector",
190 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl",
191 + "NixlPushConnector",
192 +)
193 +
194 KVConnectorFactory.register_connector(
195 "MultiConnector",
196 "vllm.distributed.kv_transfer.kv_connector.v1.multi_connector",
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+12
@@ -569,6 +569,18 @@ class KVConnectorBase_V1(ABC):
569 """
570 return ()
571
572 + def has_pending_push_work(self) -> bool:
573 + """Return True if the connector has push-mode work that requires
574 + the engine main loop to keep stepping (e.g. a P-side request whose
575 + KV blocks are waiting to be WRITTEN to a D node).
576 +
577 + Connectors that don't implement push-based KV transfer should
578 + leave this as False.
579 + """
580 + # TODO: replace with a more general connector hook for keeping the
581 + # scheduler alive (e.g. extend has_unfinished_requests).
582 + return False
583 +
584 @classmethod
585 def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
586 """
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+3
@@ -538,6 +538,9 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
538 for c in self._connectors:
539 yield from c.take_events()
540
541 + def has_pending_push_work(self) -> bool:
542 + return any(c.has_pending_push_work() for c in self._connectors)
543 +
544 @classmethod
545 def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
546 """
vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py
+30
@@ -2,14 +2,35 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """NIXL KV-cache transfer connector (disaggregated prefill / decode)."""
4
5 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (
6 + NixlBaseConnectorScheduler,
7 +)
8 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
9 + NixlBaseConnectorWorker,
10 +)
11 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import (
12 + NixlBaseConnector,
13 NixlConnector,
14 + NixlPullConnector,
15 + NixlPushConnector,
16 )
17 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
18 NixlAgentMetadata,
19 NixlConnectorMetadata,
20 NixlHandshakePayload,
21 )
22 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
23 + NixlPullConnectorScheduler,
24 +)
25 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import (
26 + NixlPullConnectorWorker,
27 +)
28 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import (
29 + NixlPushConnectorScheduler,
30 +)
31 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
32 + NixlPushConnectorWorker,
33 +)
34 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
35 NixlConnectorScheduler,
36 )
@@ -22,10 +43,19 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
43
44 __all__ = [
45 "NixlAgentMetadata",
46 + "NixlBaseConnector",
47 + "NixlBaseConnectorScheduler",
48 + "NixlBaseConnectorWorker",
49 "NixlConnector",
50 "NixlConnectorMetadata",
51 "NixlConnectorScheduler",
52 "NixlConnectorWorker",
53 "NixlHandshakePayload",
54 "NixlKVConnectorStats",
55 + "NixlPullConnector",
56 + "NixlPullConnectorScheduler",
57 + "NixlPullConnectorWorker",
58 + "NixlPushConnector",
59 + "NixlPushConnectorScheduler",
60 + "NixlPushConnectorWorker",
61 ]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py new
+455
@@ -0,0 +1,455 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Base scheduler-side logic for the NIXL connector."""
4 +
5 +import threading
6 +import time
7 +from typing import TYPE_CHECKING, Any
8 +
9 +import msgspec
10 +import zmq
11 +
12 +from vllm import envs
13 +from vllm.distributed.kv_transfer.kv_connector.utils import (
14 + BlockIds,
15 + EngineId,
16 + yield_req_data,
17 +)
18 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19 + KVConnectorHandshakeMetadata,
20 + KVConnectorMetadata,
21 +)
22 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
23 + GET_META_MSG,
24 + HeartbeatInfo,
25 + NixlConnectorMetadata,
26 + NixlHandshakePayload,
27 + ReqId,
28 +)
29 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx
30 +from vllm.logger import init_logger
31 +from vllm.platforms import current_platform
32 +from vllm.utils.math_utils import cdiv
33 +from vllm.utils.network_utils import make_zmq_path
34 +from vllm.v1.core.sched.output import SchedulerOutput
35 +from vllm.v1.kv_cache_interface import (
36 + FullAttentionSpec,
37 + MambaSpec,
38 + SlidingWindowSpec,
39 +)
40 +
41 +if TYPE_CHECKING:
42 + from vllm.config import VllmConfig
43 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
44 + from vllm.v1.kv_cache_interface import KVCacheConfig
45 + from vllm.v1.outputs import KVConnectorOutput
46 + from vllm.v1.request import Request
47 +
48 +logger = init_logger(__name__)
49 +
50 +
51 +class NixlBaseConnectorScheduler:
52 + """Base implementation of Scheduler side methods shared by pull and push."""
53 +
54 + def __init__(
55 + self,
56 + vllm_config: "VllmConfig",
57 + engine_id: str,
58 + kv_cache_config: "KVCacheConfig",
59 + ):
60 + self.vllm_config = vllm_config
61 + self.block_size = vllm_config.cache_config.block_size
62 + self.engine_id: EngineId = engine_id
63 + self.kv_cache_config = kv_cache_config
64 + self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
65 + self.side_channel_port = (
66 + envs.VLLM_NIXL_SIDE_CHANNEL_PORT
67 + + vllm_config.parallel_config.data_parallel_index
68 + )
69 + assert vllm_config.kv_transfer_config is not None
70 + self._kv_lease_duration: int = (
71 + vllm_config.kv_transfer_config.get_from_extra_config(
72 + "kv_lease_duration", 30
73 + )
74 + )
75 + # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.
76 + self._heartbeat_interval = self._kv_lease_duration // 6
77 + if current_platform.device_type == "cpu":
78 + self.use_host_buffer = False
79 + else:
80 + self.use_host_buffer = (
81 + vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
82 + )
83 + self._is_hma_required = (
84 + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
85 + # Also handle unlikely SW-only model case instead of checking num_groups>1.
86 + and any(
87 + not isinstance(g.kv_cache_spec, FullAttentionSpec)
88 + for g in kv_cache_config.kv_cache_groups
89 + )
90 + )
91 + self._has_mamba = any(
92 + isinstance(g.kv_cache_spec, MambaSpec)
93 + for g in kv_cache_config.kv_cache_groups
94 + )
95 +
96 + logger.info("Initializing NIXL Scheduler %s", engine_id)
97 + if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
98 + logger.info("Hybrid Memory Allocator is enabled with NIXL")
99 +
100 + # Background thread for handling new handshake requests.
101 + self._nixl_handshake_listener_t: threading.Thread | None = None
102 + self._stop_event = threading.Event()
103 +
104 + # Requests that need to start recv/send.
105 + # New requests are added by update_state_after_alloc in
106 + # the scheduler. Used to make metadata passed to Worker.
107 + self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
108 + self._reqs_need_save: dict[ReqId, Request] = {}
109 + # Reqs to send and their expiration time
110 + self._reqs_need_send: dict[ReqId, float] = {}
111 + self._reqs_in_batch: set[ReqId] = set()
112 + # Reqs to remove from processed set because they're not to send after
113 + # remote prefill or aborted.
114 + self._reqs_not_processed: set[ReqId] = set()
115 +
116 + # Heartbeat tracking: requests needing periodic lease-renewal heartbeats to
117 + # remote P-side, stored as ready-to-send HeartbeatInfo grouped by remote engine
118 + self._heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {}
119 + # Reverse lookup: local req_id -> (engine_id, remote_req_id) for O(1) removal
120 + self._heartbeat_req_engine: dict[ReqId, tuple[EngineId, ReqId]] = {}
121 + self._last_heartbeat_time: float = 0.0
122 +
123 + # Gather Sliding Window sizes for each kv cache group (if any) in number of
124 + # blocks per KV cache group. This is used to clip the local attention window.
125 + sw_sizes_tokens: list[tuple[int, int]] = [
126 + (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
127 + if isinstance(g.kv_cache_spec, SlidingWindowSpec)
128 + else (0, self.block_size)
129 + for g in kv_cache_config.kv_cache_groups
130 + ]
131 + # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively
132 + # account for boundary overlap eg window isn't fully aligned with blocks.
133 + self.blocks_per_sw = [
134 + cdiv(n_tokens, block_size) + 1 if n_tokens else 0
135 + for n_tokens, block_size in sw_sizes_tokens
136 + ]
137 +
138 + # Threshold to decide whether to compute kv cache locally
139 + # or pull from a remote node: minimum number of remote
140 + # tokens to amortize the xfer latencies
141 + self.kv_recompute_threshold: int = int(
142 + vllm_config.kv_transfer_config.get_from_extra_config(
143 + "kv_recompute_threshold", 64
144 + )
145 + )
146 +
147 + # Bi-directional KV transfer feature supports KV block
148 + # transfers from D node to P node
149 + self.is_bidirectional_kv_xfer_enabled = (
150 + vllm_config.kv_transfer_config.get_from_extra_config(
151 + "bidirectional_kv_xfer", False
152 + )
153 + )
154 + self.decoder_kv_blocks_ttl = (
155 + vllm_config.kv_transfer_config.get_from_extra_config(
156 + "decoder_kv_blocks_ttl", 480
157 + )
158 + )
159 +
160 + if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0:
161 + logger.info(
162 + "Bidirectional KV transfer is enabled and the kv "
163 + "recompute threshold is set to %d tokens."
164 + "KV blocks on D are released after a TTL of %d seconds.",
165 + self.kv_recompute_threshold,
166 + self.decoder_kv_blocks_ttl,
167 + )
168 +
169 + def shutdown(self):
170 + self._stop_event.set()
171 + if self._nixl_handshake_listener_t is not None:
172 + self._nixl_handshake_listener_t.join()
173 + self._nixl_handshake_listener_t = None
174 +
175 + def on_new_request(self, request: "Request") -> None:
176 + """Track a request that may need heartbeats."""
177 + params = request.kv_transfer_params
178 + # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are
179 + # effectively disabled for Bidirectional KV transfer.
180 + if params is None or not params.get("do_remote_prefill"):
181 + return
182 + # Only track if all required remote fields are present.
183 + remote_engine_id = params.get("remote_engine_id")
184 + remote_request_id = params.get("remote_request_id")
185 + host = params.get("remote_host")
186 + port = params.get("remote_port")
187 + tp_size = params.get("tp_size")
188 + if (
189 + remote_engine_id is None
190 + or remote_request_id is None
191 + or host is None
192 + or port is None
193 + or tp_size is None
194 + ):
195 + return
196 + if remote_engine_id not in self._heartbeat_by_engine:
197 + self._heartbeat_by_engine[remote_engine_id] = HeartbeatInfo(
198 + req_ids=set(),
199 + host=host,
200 + port=port,
201 + tp_size=tp_size,
202 + )
203 + self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id)
204 + self._heartbeat_req_engine[request.request_id] = (
205 + remote_engine_id,
206 + remote_request_id,
207 + )
208 +
209 + def _stop_heartbeat(self, req_id: ReqId) -> None:
210 + """Remove *req_id* from heartbeat tracking (if tracked)."""
211 + if key := self._heartbeat_req_engine.pop(req_id, None):
212 + engine_id, remote_id = key
213 + if info := self._heartbeat_by_engine.get(engine_id):
214 + info.req_ids.discard(remote_id)
215 + if not info.req_ids:
216 + # Clean up empty engines so we don't leak a key when remote dies.
217 + del self._heartbeat_by_engine[engine_id]
218 +
219 + def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
220 + """
221 + Clip the number of blocks to the sliding window size for each kv cache group
222 + that employs SWA.
223 + This is necessary because the KV Cache manager initially allocates blocks for
224 + the entire sequence length, and successively cleans up blocks that are outside
225 + the window prior to the `request_finished_all_groups` hook.
226 + """
227 + if len(block_ids) == 0 or not self._is_hma_required:
228 + # No blocks to clip eg Full prefix cache hit or not a hybrid model.
229 + return block_ids
230 + # NOTE (NickLucche) This logic is currently handled at the connector level
231 + # because offloading connectors might want to receive the whole sequence even
232 + # for SWA groups. We will abstract this logic once the interface is more stable
233 + assert len(block_ids) == len(self.blocks_per_sw), (
234 + "Number of KV cache groups must match"
235 + )
236 + # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged
237 + return tuple(
238 + [
239 + blocks[-self.blocks_per_sw[i] :]
240 + if self.blocks_per_sw[i] > 0
241 + else blocks
242 + for i, blocks in enumerate(block_ids)
243 + ]
244 + )
245 +
246 + def set_xfer_handshake_metadata(
247 + self, metadata: dict[int, KVConnectorHandshakeMetadata]
248 + ) -> None:
249 + """
250 + Set the KV connector handshake metadata for this connector.
251 +
252 + Args:
253 + metadata (dict): the handshake metadata to set.
254 + """
255 + encoded_data: dict[int, bytes] = {}
256 + encoder = msgspec.msgpack.Encoder()
257 + for tp_rank, rank_metadata in metadata.items():
258 + if not isinstance(rank_metadata, NixlHandshakePayload):
259 + raise ValueError(
260 + "NixlConnectorScheduler expects NixlHandshakePayload for "
261 + "handshake metadata."
262 + )
263 + encoded_data[tp_rank] = encoder.encode(rank_metadata)
264 + logger.debug(
265 + "Tp rank %d: encoded NixlHandshakePayload size: %s bytes",
266 + tp_rank,
267 + str(len(encoded_data[tp_rank])),
268 + )
269 +
270 + # Only start the listener when we have metadata to serve.
271 + if self._nixl_handshake_listener_t is None:
272 + ready_event = threading.Event()
273 + self._nixl_handshake_listener_t = threading.Thread(
274 + target=self._nixl_handshake_listener,
275 + args=(
276 + encoded_data,
277 + ready_event,
278 + self._stop_event,
279 + self.side_channel_host,
280 + self.side_channel_port,
281 + ),
282 + daemon=True,
283 + name="nixl_handshake_listener",
284 + )
285 + self._nixl_handshake_listener_t.start()
286 + ready_event.wait() # Wait for listener ZMQ socket to be ready.
287 +
288 + @staticmethod
289 + def _nixl_handshake_listener(
290 + encoded_data: dict[int, Any],
291 + ready_event: threading.Event,
292 + stop_event: threading.Event,
293 + host: str,
294 + port: int,
295 + ):
296 + """Background thread for getting new NIXL handshakes."""
297 + # NOTE(rob): this is a simple implementation. We will move
298 + # to a better approach via HTTP endpoint soon.
299 +
300 + # Listen for new requests for metadata.
301 + path = make_zmq_path("tcp", host, port)
302 + logger.debug("Starting listening on path: %s", path)
303 + with zmq_ctx(zmq.ROUTER, path) as sock:
304 + sock.setsockopt(zmq.RCVTIMEO, 1000)
305 + ready_event.set()
306 + while True:
307 + try:
308 + identity, _, msg = sock.recv_multipart()
309 + except zmq.Again:
310 + if stop_event.is_set():
311 + break
312 + continue
313 + # Decode the message which contains (GET_META_MSG, rank)
314 + msg, target_tp_rank = msgspec.msgpack.decode(msg)
315 + logger.debug(
316 + "Received message for tp rank %s",
317 + target_tp_rank,
318 + )
319 + if msg != GET_META_MSG:
320 + logger.warning("Connection listener got unexpected message %s", msg)
321 + sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))
322 +
323 + def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int:
324 + """D-side only. Returns N-1 for Mamba models since the decoder
325 + always recomputes the last token and must start from h(N-1)."""
326 + if self._has_mamba and num_prompt_tokens > 1:
327 + return num_prompt_tokens - 1
328 + return num_prompt_tokens
329 +
330 + def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
331 + """P-side only: drop the last prompt token so the prefiller computes
332 + h(N-1) instead of h(N). The decoder recomputes the last token to
333 + derive h(N) correctly.
334 +
335 + Guarded by ``_p_side_truncated`` to avoid repeated truncation if the
336 + request is preempted and rescheduled."""
337 + params = request.kv_transfer_params
338 + if (
339 + params is not None
340 + # Guard against repeated truncation after preemption/reschedule.
341 + and not params.get("_p_side_truncated")
342 + and request.num_prompt_tokens > 1
343 + ):
344 + if request.prompt_token_ids is not None:
345 + request.prompt_token_ids.pop()
346 + elif request.prompt_embeds is not None:
347 + request.prompt_embeds = request.prompt_embeds[:-1]
348 + else:
349 + return
350 +
351 + request._all_token_ids.pop()
352 + request.num_prompt_tokens -= 1
353 + request.max_tokens = 1
354 + params["_p_side_truncated"] = True
355 +
356 + def _build_save_meta(
357 + self,
358 + meta: NixlConnectorMetadata,
359 + scheduler_output: SchedulerOutput,
360 + ) -> None:
361 + # only called when use_host_buffer is True to build the save metadata
362 +
363 + # NOTE: For the prefill side, there might be a chance that an early added
364 + # request is a chunked prefill, so we need to check if new blocks are added
365 + for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output):
366 + req_to_save = self._reqs_need_save.get(req_id)
367 + if req_to_save is None or new_block_id_groups is None:
368 + continue
369 + req = req_to_save
370 +
371 + assert req.kv_transfer_params is not None
372 + clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups)
373 + meta.add_new_req_to_save(
374 + request_id=req_id,
375 + local_block_ids=clipped_block_id_groups,
376 + kv_transfer_params=req.kv_transfer_params,
377 + )
378 + assert scheduler_output.num_scheduled_tokens is not None
379 + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
380 + is_partial = (
381 + req.num_computed_tokens + num_scheduled_tokens
382 + ) < req.num_prompt_tokens
383 + if not is_partial:
384 + # For non-partial prefills, once new req_meta is scheduled, it
385 + # can be removed from _reqs_need_save.
386 + # For partial prefill case, we will retain the request in
387 + # _reqs_need_save until all blocks are scheduled with req_meta.
388 + # Therefore, only pop if `not is_partial`.
389 + self._reqs_need_save.pop(req_id)
390 +
391 + def build_connector_meta(
392 + self,
393 + scheduler_output: SchedulerOutput,
394 + ) -> KVConnectorMetadata:
395 + meta = NixlConnectorMetadata()
396 +
397 + # Loop through scheduled reqs and convert to ReqMeta.
398 + for req_id, (req, block_ids) in self._reqs_need_recv.items():
399 + assert req.kv_transfer_params is not None
400 + meta.add_new_req_to_recv(
401 + request_id=req_id,
402 + local_block_ids=block_ids,
403 + kv_transfer_params=req.kv_transfer_params,
404 + )
405 +
406 + if self.use_host_buffer:
407 + self._build_save_meta(meta, scheduler_output)
408 +
409 + meta.reqs_to_send = self._reqs_need_send
410 + meta.reqs_in_batch = self._reqs_in_batch
411 + meta.reqs_not_processed = self._reqs_not_processed
412 +
413 + # Package heartbeats, throttled by heartbeat_interval.
414 + if self._heartbeat_by_engine:
415 + now = time.perf_counter()
416 + if now - self._last_heartbeat_time >= self._heartbeat_interval:
417 + self._last_heartbeat_time = now
418 + meta.heartbeat_by_engine = self._heartbeat_by_engine
419 +
420 + # Clear the list once workers start the transfers
421 + self._reqs_need_recv.clear()
422 + self._reqs_in_batch = set()
423 + self._reqs_not_processed = set()
424 + self._reqs_need_send = {}
425 +
426 + return meta
427 +
428 + def update_connector_output(self, connector_output: "KVConnectorOutput") -> None:
429 + """Stop heartbeating for requests whose KV transfer completed."""
430 + for req_id in connector_output.finished_recving or ():
431 + self._stop_heartbeat(req_id)
432 +
433 + def has_pending_push_work(self) -> bool:
434 + return False
435 +
436 + ############################################################
437 + # Abstract methods that subclasses must implement
438 + ############################################################
439 +
440 + def get_num_new_matched_tokens(
441 + self, request: "Request", num_computed_tokens: int
442 + ) -> tuple[int, bool]:
443 + raise NotImplementedError
444 +
445 + def update_state_after_alloc(
446 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
447 + ):
448 + raise NotImplementedError
449 +
450 + def request_finished(
451 + self,
452 + request: "Request",
453 + block_ids: BlockIds,
454 + ) -> tuple[bool, dict[str, Any] | None]:
455 + raise NotImplementedError
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py new
+2286
@@ -0,0 +1,2286 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Base worker-side logic for the NIXL connector."""
4 +
5 +import logging
6 +import os
7 +import queue
8 +import threading
9 +import time
10 +import uuid
11 +from collections import defaultdict
12 +from collections.abc import Iterator
13 +from concurrent.futures import Future, ThreadPoolExecutor
14 +from typing import TYPE_CHECKING, Any, cast
15 +
16 +import msgspec
17 +import numpy as np
18 +import torch
19 +import zmq
20 +
21 +from vllm.distributed.kv_transfer.kv_connector.utils import (
22 + BlockIds,
23 + EngineId,
24 + EngineTransferInfo,
25 + TransferTopology,
26 + get_current_attn_backends,
27 + kv_postprocess_blksize_and_layout_on_receive,
28 + kv_postprocess_blksize_on_receive,
29 + kv_postprocess_layout_on_receive,
30 +)
31 +from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp
32 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
33 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
34 + GET_META_MSG,
35 + NixlAgentMetadata,
36 + NixlConnectorMetadata,
37 + NixlHandshakePayload,
38 + ReqId,
39 + ReqMeta,
40 + TransferHandle,
41 + compute_nixl_compatibility_hash,
42 +)
43 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
44 + NixlKVConnectorStats,
45 +)
46 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
47 + TPMapping,
48 + _is_attention_spec,
49 + _is_ssm_spec,
50 + compute_tp_mapping,
51 +)
52 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
53 + _NIXL_SUPPORTED_DEVICE,
54 + get_representative_spec_type,
55 + zmq_ctx,
56 +)
57 +from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
58 + MambaConvSplitInfo,
59 + derive_mamba_conv_split,
60 +)
61 +from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
62 +from vllm.distributed.parallel_state import (
63 + get_tensor_model_parallel_rank,
64 + get_tensor_model_parallel_world_size,
65 +)
66 +from vllm.logger import init_logger
67 +from vllm.platforms import current_platform
68 +from vllm.utils.network_utils import make_zmq_path
69 +from vllm.v1.attention.backends.utils import get_kv_cache_layout
70 +from vllm.v1.kv_cache_interface import (
71 + FullAttentionSpec,
72 + MambaSpec,
73 + MLAAttentionSpec,
74 + UniformTypeKVCacheSpecs,
75 +)
76 +from vllm.v1.worker.block_table import BlockTable
77 +from vllm.v1.worker.utils import select_common_block_size
78 +
79 +if TYPE_CHECKING:
80 + from vllm.config import VllmConfig
81 + from vllm.v1.kv_cache_interface import KVCacheConfig
82 +
83 +logger = init_logger(__name__)
84 +
85 +
86 +class NixlBaseConnectorWorker:
87 + """Base implementation of Worker side methods shared by pull and push."""
88 +
89 + def _compute_desc_ids(
90 + self,
91 + block_ids: BlockIds,
92 + dst_num_blocks: int,
93 + block_size_ratio: float | None,
94 + physical_blocks_per_logical: int,
95 + ) -> np.ndarray:
96 + """Compute NIXL descriptor IDs for given block IDs."""
97 + num_fa_regions = self.num_regions
98 + num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0
99 +
100 + num_blocks = dst_num_blocks
101 + if block_size_ratio is not None:
102 + num_blocks = int(num_blocks * block_size_ratio)
103 + num_fa_descs = num_fa_regions * num_blocks
104 +
105 + # All-attention fast path: single vectorized broadcast.
106 + if num_ssm_regions == 0:
107 + # NOTE (NickLucche) With HMA, every kv group has the same number of layers
108 + # and layers from different groups share the same kv tensor.
109 + # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be
110 + # read across all regions, same for [3], but group0-group1 blocks will
111 + # always differ (different areas). Therefore we can just flatten the
112 + # block_ids and compute the descs ids for all groups at once.
113 + block_arr = np.concatenate(block_ids)[None, :]
114 + region_ids = np.arange(num_fa_regions)[:, None]
115 + return (region_ids * num_blocks + block_arr).flatten()
116 +
117 + # Compute desc ids per group using the right stride: FA descs have
118 + # num_blocks entries per region (kernel granularity), SSM descs have
119 + # logical_blocks entries per region (no kernel splitting).
120 + logical_blocks = num_blocks // physical_blocks_per_logical
121 + all_descs: list[np.ndarray] = []
122 + for i, group in enumerate(block_ids):
123 + group_arr = np.asarray(group)
124 + if _is_attention_spec(self._group_spec_types[i]):
125 + fa_region_ids = np.arange(num_fa_regions)[:, None]
126 + all_descs.append(
127 + (fa_region_ids * num_blocks + group_arr[None, :]).flatten()
128 + )
129 + elif _is_ssm_spec(self._group_spec_types[i]):
130 + # NOTE (NickLucche) SSM and Attention block regions can
131 + # be exchanged arbitrarily by manager. Therefore, descs
132 + # are laid out as:
133 + # [descs_fa (all regions) | descs_ssm (all regions)].
134 + # num_fa_descs offset must be computed per-engine since
135 + # P and D can have different num_blocks (and thus
136 + # different FA desc counts).
137 + ssm_region_ids = np.arange(num_ssm_regions)[:, None]
138 + all_descs.append(
139 + (
140 + ssm_region_ids * logical_blocks
141 + + group_arr[None, :]
142 + + num_fa_descs
143 + ).flatten()
144 + )
145 + else:
146 + raise ValueError(
147 + f"Unknown spec type {self._group_spec_types[i]} at index {i}"
148 + )
149 +
150 + return np.concatenate(all_descs)
151 +
152 + def _build_local_splits_from_plan(
153 + self,
154 + plan: TPMapping,
155 + src_blocks_data: list[tuple[int, int, int]],
156 + num_fa_descs: int,
157 + ) -> Iterator[list[tuple[int, int, int]]]:
158 + """Build split handle data for P_TP > D_TP scenario.
159 +
160 + num_fa_descs is the boundary between FA and SSM descriptors.
161 + Split counts are derived from source_ranks_per_group lengths.
162 + FA uses rank_to_attention_slot for the slot offset;
163 + SSM uses the rank's positional index.
164 + """
165 + fa_idx = next(
166 + i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
167 + )
168 + fa_num_splits = len(plan.source_ranks_per_group[fa_idx])
169 +
170 + has_ssm_descs = num_fa_descs < len(src_blocks_data)
171 + ssm_idx = next(
172 + (i for i, t in enumerate(self._group_spec_types) if _is_ssm_spec(t)),
173 + None,
174 + )
175 + ssm_num_splits = (
176 + len(plan.source_ranks_per_group[ssm_idx])
177 + if has_ssm_descs and ssm_idx is not None
178 + else 0
179 + )
180 +
181 + # Per-FA-descriptor replicate flag, in _build_fa_local emission order.
182 + fa_desc_replicated = self._fa_desc_replicated(num_fa_descs)
183 +
184 + for p_idx, p_rank in enumerate(plan.all_source_ranks):
185 + fa_slot = plan.rank_to_attention_slot.get(p_rank, 0)
186 +
187 + handle: list[tuple[int, int, int]] = []
188 + for j, (addr, local_len, dev) in enumerate(src_blocks_data):
189 + if j < num_fa_descs:
190 + if fa_desc_replicated[j]:
191 + # REPLICATE (MLA): whole block written on every rank.
192 + handle.append((addr, local_len, dev))
193 + else:
194 + # SPLIT (full-attn): this rank's head slice.
195 + chunk = local_len // fa_num_splits
196 + handle.append((addr + fa_slot * chunk, chunk, dev))
197 + else:
198 + chunk = local_len // ssm_num_splits
199 + handle.append((addr + p_idx * chunk, chunk, dev))
200 + yield handle
201 +
202 + def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]:
203 + """Per-FA-descriptor replicate flag, in _build_fa_local emission order
204 + (region-major; K then optional V per region). Length ``num_fa_descs``.
205 + """
206 + assert self.transfer_topo is not None
207 + n_regions = len(self.block_len_per_layer)
208 + if n_regions == 0 or self.num_regions == 0:
209 + return [False] * num_fa_descs
210 + nblk = num_fa_descs // self.num_regions
211 + virtually_split = self.transfer_topo.virtually_split_kv_in_blocks
212 + flags: list[bool] = []
213 + for i in range(n_regions):
214 + replicated = self._is_region_replicated(i)
215 + num_streams = 1 if replicated or not virtually_split else 2
216 + flags.extend([replicated] * (num_streams * nblk))
217 + assert len(flags) == num_fa_descs, (
218 + f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}"
219 + )
220 + return flags
221 +
222 + def _is_region_replicated(self, region_idx: int) -> bool:
223 + """Whether region ``region_idx`` is transferred REPLICATE vs SPLIT.
224 +
225 + REPLICATE (MLA): identical on every rank, whole block read from one
226 + rank at offset 0, key-only. SPLIT (full-attn): head-sharded across TP.
227 + Defaults to SPLIT when the per-region map is unset (e.g. tests that set
228 + block_len_per_layer without register_kv_caches).
229 + """
230 + return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx]
231 +
232 + def __init__(
233 + self,
234 + vllm_config: "VllmConfig",
235 + engine_id: str,
236 + kv_cache_config: "KVCacheConfig",
237 + ):
238 + nixl_wrapper_cls = NixlWrapper
239 + if nixl_wrapper_cls is None:
240 + logger.error("NIXL is not available")
241 + raise RuntimeError("NIXL is not available")
242 + logger.info("Initializing NIXL wrapper")
243 + logger.info("Initializing NIXL worker %s", engine_id)
244 +
245 + # Config.
246 + self.vllm_config = vllm_config
247 + # mypy will complain on re-assignment otherwise.
248 + self.block_size: int = cast(int, vllm_config.cache_config.block_size)
249 +
250 + if vllm_config.kv_transfer_config is None:
251 + raise ValueError("kv_transfer_config must be set for NixlConnector")
252 + self.kv_transfer_config = vllm_config.kv_transfer_config
253 +
254 + self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config(
255 + "backends", ["UCX"]
256 + )
257 + kv_lease_duration: int = vllm_config.kv_transfer_config.get_from_extra_config(
258 + "kv_lease_duration", 30
259 + )
260 + # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.
261 + self._lease_extension = kv_lease_duration * 2 // 3
262 +
263 + self._is_hma_required = (
264 + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
265 + and any(
266 + not isinstance(g.kv_cache_spec, FullAttentionSpec)
267 + for g in kv_cache_config.kv_cache_groups
268 + )
269 + )
270 + self.kv_cache_config = kv_cache_config
271 + self._layer_specs = {
272 + layer: group.kv_cache_spec
273 + for group in kv_cache_config.kv_cache_groups
274 + for layer in group.layer_names
275 + }
276 + self.hma_group_size = len(kv_cache_config.kv_cache_tensors)
277 +
278 + # ---- Model state (derived from model config) ----
279 + mamba_ssm_size = (0, 0)
280 + # Conv state sub-projection decomposition (None when no Mamba).
281 + # The 3-read transfer requires DS (dim, state_len) conv layout so
282 + # that x/B/C sub-projections are contiguous in memory.
283 + self._conv_decomp: MambaConvSplitInfo | None = None
284 + self._has_mamba = any(
285 + isinstance(g.kv_cache_spec, MambaSpec)
286 + for g in kv_cache_config.kv_cache_groups
287 + )
288 + if self._has_mamba:
289 + assert self._is_hma_required
290 + from vllm.model_executor.layers.mamba.mamba_utils import (
291 + is_conv_state_dim_first,
292 + )
293 +
294 + assert is_conv_state_dim_first(), (
295 + "3-read Mamba conv transfer requires DS conv state layout. "
296 + "Set VLLM_SSM_CONV_STATE_LAYOUT=DS"
297 + )
298 + mamba_spec = next(
299 + spec
300 + for spec in self._layer_specs.values()
301 + if isinstance(spec, MambaSpec)
302 + )
303 + self._conv_decomp = derive_mamba_conv_split(
304 + mamba_spec,
305 + vllm_config.parallel_config.tensor_parallel_size,
306 + )
307 + mamba_ssm_size = self._conv_decomp.ssm_sizes
308 + self._mamba_ssm_size = mamba_ssm_size
309 +
310 + # Agent.
311 + non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
312 + # Configure NIXL num_threads to avoid UAR exhaustion on Mellanox NICs.
313 + # Each UCX thread allocates UARs (doorbell pages) via DevX, and
314 + # excessive NIXL UAR usage can exhaust NIC UAR space. This can cause
315 + # components like NVSHMEM (used by DeepEP kernels) to fail during RDMA
316 + # initialization with "mlx5dv_devx_alloc_uar" errors.
317 + # Ref: https://network.nvidia.com/files/doc-2020/ethernet-adapters-programming-manual.pdf#page=63
318 + num_threads = vllm_config.kv_transfer_config.get_from_extra_config(
319 + "num_threads", 4
320 + )
321 + if nixl_agent_config is None:
322 + config = None
323 + else:
324 + # Enable telemetry by default for NIXL 0.7.1 and above.
325 + config = (
326 + nixl_agent_config(backends=self.nixl_backends, capture_telemetry=True)
327 + if len(non_ucx_backends) > 0
328 + else nixl_agent_config(num_threads=num_threads, capture_telemetry=True)
329 + )
330 +
331 + self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config)
332 + # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.
333 + self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict)
334 +
335 + # Metadata.
336 + self.engine_id: EngineId = engine_id
337 + self.tp_rank = get_tensor_model_parallel_rank()
338 + self.world_size = get_tensor_model_parallel_world_size()
339 +
340 + self.num_blocks = kv_cache_config.num_blocks
341 + self.enable_permute_local_kv = False
342 + self.enable_heterogeneous_attn_post_process = False
343 +
344 + # KV Caches and nixl tracking data.
345 + self.device_type = current_platform.device_type
346 + self.kv_buffer_device: str = vllm_config.kv_transfer_config.kv_buffer_device
347 + if self.device_type not in _NIXL_SUPPORTED_DEVICE:
348 + raise RuntimeError(f"{self.device_type} is not supported.")
349 + elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[self.device_type]:
350 + raise RuntimeError(
351 + f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
352 + "is not supported."
353 + )
354 + self.device_kv_caches: dict[str, torch.Tensor] = {}
355 +
356 + # cpu kv buffer for xfer
357 + # used when device memory can not be registered under nixl
358 + self.host_xfer_buffers: dict[str, torch.Tensor] = {}
359 + if self.device_type == "cpu":
360 + self.use_host_buffer = False
361 + else:
362 + self.use_host_buffer = self.kv_buffer_device == "cpu"
363 +
364 + # reserve different cores for start_load_kv() from model_forward()
365 + if self.device_type == "cpu":
366 + numa_core_list = current_platform.discover_numa_topology()
367 + # setup one last core in each numa for kv transfer.
368 + rsv_cores_for_kv = [
369 + max(each_numa_core_list) for each_numa_core_list in numa_core_list
370 + ]
371 +
372 + if rsv_cores_for_kv:
373 + if not hasattr(os, "sched_setaffinity"):
374 + raise NotImplementedError(
375 + "os.sched_setaffinity is not available on this platform"
376 + )
377 + os.sched_setaffinity(0, rsv_cores_for_kv)
378 +
379 + # support for oot platform which can't register nixl memory
380 + # type based on kv_buffer_device
381 + nixl_memory_type = current_platform.get_nixl_memory_type()
382 + if nixl_memory_type is None:
383 + if self.kv_buffer_device in ["cuda", "xpu"]:
384 + nixl_memory_type = "VRAM"
385 + elif self.kv_buffer_device == "cpu":
386 + nixl_memory_type = "DRAM"
387 + if nixl_memory_type is None:
388 + raise RuntimeError(
389 + f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
390 + "is not supported."
391 + )
392 + self.nixl_memory_type = nixl_memory_type
393 +
394 + # Note: host xfer buffer ops when use_host_buffer is True
395 + self.copy_blocks: CopyBlocksOp | None = None
396 +
397 + # Map of engine_id -> kv_caches_base_addr. For TP case, each local
398 + self.device_id: int = 0
399 + # Current rank may pull from multiple remote TP workers.
400 + # EngineId, dict[int, list[int]] -> engine_id, tp_rank, base_addr_for_layer
401 + self.kv_caches_base_addr = defaultdict[EngineId, dict[int, list[int]]](dict)
402 +
403 + # Number of NIXL regions. Currently one region per cache
404 + # (so 1 per layer for MLA, otherwise 2 per layer)
405 + self.num_regions = 0
406 +
407 + # nixl_prepped_dlist_handle.
408 + self.src_xfer_handles_by_block_size: dict[int, int] = {}
409 + # Populated dynamically during handshake based on remote configuration.
410 + # Keep track of regions at different tp_ratio values. tp_ratio->handles
411 + self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {}
412 + # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}.
413 + self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict)
414 +
415 + # Map of engine_id -> num_blocks. All ranks in the same deployment will
416 + # have the same number of blocks.
417 + self.dst_num_blocks: dict[EngineId, int] = {}
418 + self._registered_descs: list[Any] = []
419 +
420 + # In progress transfers.
421 + # [req_id -> list[handle]]
422 + self._recving_metadata: dict[ReqId, ReqMeta] = {}
423 + self._recving_transfers = defaultdict[ReqId, list[TransferHandle]](list)
424 + # Track the expiration time of requests that are waiting to be sent.
425 + self._reqs_to_send: dict[ReqId, float] = {}
426 + # Set of requests that have been part of a batch, regardless of status.
427 + self._reqs_to_process: set[ReqId] = set()
428 +
429 + # Invalid blocks from failed NIXL operations (thread-safe queue of block ids)
430 + self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue()
431 + # requests that skipped transfer (handshake or transfer failures)
432 + # Uses Queue for thread-safe cross-thread coordination with the
433 + # background handshake thread, matching the _ready_requests pattern.
434 + self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
435 +
436 + # Handshake metadata of this worker for NIXL transfers.
437 + self.xfer_handshake_metadata: NixlHandshakePayload | None = None
438 + # Background thread for initializing new NIXL handshakes.
439 + self._handshake_initiation_executor = ThreadPoolExecutor(
440 + # NIXL is not guaranteed to be thread-safe, limit 1 worker.
441 + max_workers=1,
442 + thread_name_prefix="vllm-nixl-handshake-initiator",
443 + )
444 + self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]()
445 + self._handshake_futures: dict[EngineId, Future[dict[int, str]]] = {}
446 + # Protects _handshake_futures and _remote_agents.
447 + self._handshake_lock = threading.RLock()
448 +
449 + # TTL-based eviction of stale remote engine state.
450 + self._engine_last_active: dict[EngineId, float] = {}
451 + self._engine_ttl: float = vllm_config.kv_transfer_config.get_from_extra_config(
452 + "engine_ttl", 3600.0
453 + )
454 +
455 + self.block_size = vllm_config.cache_config.block_size
456 + self.model_config = vllm_config.model_config
457 +
458 + self.use_mla = self.model_config.use_mla
459 +
460 + # Get the attention backend from the first layer
461 + # NOTE (NickLucche) models with multiple backends are not supported yet
462 + self.attn_backends = get_current_attn_backends(vllm_config)
463 + self.backend_name = self.attn_backends[0].get_name()
464 +
465 + self.kv_cache_layout = get_kv_cache_layout()
466 + self.host_buffer_kv_cache_layout = self.kv_cache_layout
467 + logger.info(
468 + "Detected attention backend(s) %s",
469 + [backend.get_name() for backend in self.attn_backends],
470 + )
471 + logger.info("Detected kv cache layout %s", self.kv_cache_layout)
472 +
473 + # lazy initialized in register_kv_caches
474 + self.compat_hash: str | None = None
475 + self.transfer_topo: TransferTopology | None = None
476 +
477 + # With heterogeneous TP, P must wait for all assigned D TP workers to
478 + # finish reading before safely freeing the blocks.
479 + self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
480 + self.xfer_stats = NixlKVConnectorStats()
481 +
482 + self._physical_blocks_per_logical_kv_block = 1
483 + self._sync_block_size_with_kernel()
484 +
485 + # Unwrap UniformTypeKVCacheSpecs to get the representative spec type
486 + self._group_spec_types = tuple(
487 + get_representative_spec_type(g.kv_cache_spec)
488 + for g in self.kv_cache_config.kv_cache_groups
489 + )
490 +
491 + # Per-region MLA flag, 1:1 with block_len_per_layer. True -> REPLICATE
492 + # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models
493 + # combining both (e.g. GQA main + MLA Eagle-3 draft).
494 + self._region_is_mla = list[bool]()
495 +
496 + # Enable different block lengths for different layers *only* when MLA is used.
497 + # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.
498 + self.block_len_per_layer = list[int]()
499 +
500 + # Per-engine TP mappings. Generated during handshake.
501 + self.tp_mappings: dict[EngineId, TPMapping] = {}
502 +
503 + self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
504 + "enforce_handshake_compat", True
505 + )
506 +
507 + def _sync_block_size_with_kernel(self) -> None:
508 + backends = get_current_attn_backends(self.vllm_config)
509 + kernel_block_size = select_common_block_size(self.block_size, backends)
510 + # Number of blocks not accounting for kernel block mismatches
511 + self._logical_num_blocks = self.num_blocks
512 + if self.block_size != kernel_block_size:
513 + logger.info_once(
514 + "User-specified logical block size (%s) does not match"
515 + " physical kernel block size (%s). Using the latter.",
516 + self.block_size,
517 + kernel_block_size,
518 + )
519 + assert self.block_size > kernel_block_size
520 + self._physical_blocks_per_logical_kv_block = (
521 + self.block_size // kernel_block_size
522 + )
523 + self.block_size = kernel_block_size
524 + self.num_blocks *= self._physical_blocks_per_logical_kv_block
525 +
526 + def _nixl_handshake(
527 + self,
528 + host: str,
529 + port: int,
530 + remote_tp_size: int,
531 + expected_engine_id: str,
532 + ) -> dict[int, str]:
533 + """Do a NIXL handshake with a remote instance."""
534 +
535 + # the first time we connect to a remote agent.
536 + # be careful, the handshake happens in a background thread.
537 + # it does not have an active cuda context until any cuda runtime
538 + # call is made. when UCX fails to find a valid cuda context, it will
539 + # disable any cuda ipc communication, essentially disabling any NVLink
540 + # communication.
541 + # when we are using device buffers, we need to set the device
542 + # explicitly to make sure the handshake background thread has a valid
543 + # cuda context.
544 + if not self.use_host_buffer:
545 + current_platform.set_device(self.device_id)
546 +
547 + # When target instance TP > local TP, we need to perform multiple
548 + # handshakes. Do it in a single background job for simplicity.
549 + # Regardless, only handshake with the remote TP rank(s) that current
550 + # local rank will read from. Note that With homogeneous TP,
551 + # this happens to be the same single rank_i.
552 + assert self.transfer_topo is not None
553 + p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size)
554 + remote_rank_to_agent_name = {}
555 + path = make_zmq_path("tcp", host, port)
556 +
557 + with zmq_ctx(zmq.REQ, path) as sock:
558 + for remote_rank in p_remote_ranks:
559 + logger.debug(
560 + "Querying metadata on path: %s at remote tp rank %s",
561 + path,
562 + remote_rank,
563 + )
564 +
565 + start_time = time.perf_counter()
566 + # Send query for the request.
567 + msg = msgspec.msgpack.encode((GET_META_MSG, remote_rank))
568 + # Set receive timeout to 5 seconds to avoid hanging on dead server
569 + sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds
570 + sock.send(msg)
571 + handshake_bytes = sock.recv()
572 +
573 + # Decode handshake payload to get compatibility hash
574 + handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload)
575 + try:
576 + handshake_payload = handshake_decoder.decode(handshake_bytes)
577 + except (msgspec.DecodeError, msgspec.ValidationError) as e:
578 + raise RuntimeError(
579 + f"Failed to decode NixlHandshakePayload. This likely indicates "
580 + f"an incompatibility between connector version. Error: {e}"
581 + ) from e
582 +
583 + got_metadata_time = time.perf_counter()
584 + logger.debug(
585 + "NIXL handshake: get metadata took: %s",
586 + got_metadata_time - start_time,
587 + )
588 +
589 + # Check compatibility hash BEFORE decoding agent metadata
590 + assert self.compat_hash is not None
591 + if (
592 + self.enforce_compat_hash
593 + and handshake_payload.compatibility_hash != self.compat_hash
594 + ):
595 + raise RuntimeError(
596 + f"NIXL compatibility hash mismatch. "
597 + f"Local: {self.compat_hash}, "
598 + f"Remote: {handshake_payload.compatibility_hash}. "
599 + f"Prefill and decode instances have incompatible "
600 + f"configurations. This may be due to: different vLLM versions,"
601 + f" models, dtypes, KV cache layouts, attention backends, etc. "
602 + f"Both instances must use identical configurations."
603 + f"Disable this check using "
604 + f'--kv-transfer-config \'{{"kv_connector_extra_config": '
605 + f'{{"enforce_handshake_compat": false}}}}\''
606 + )
607 +
608 + logger.info(
609 + "NIXL compatibility check passed (hash: %s)",
610 + handshake_payload.compatibility_hash,
611 + )
612 +
613 + # Decode agent metadata
614 + metadata_decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
615 + try:
616 + metadata = metadata_decoder.decode(
617 + handshake_payload.agent_metadata_bytes
618 + )
619 + except (msgspec.DecodeError, msgspec.ValidationError) as e:
620 + # This should not happen if hash matched
621 + raise RuntimeError(
622 + f"Failed to decode NixlAgentMetadata. Error: {e}"
623 + ) from e
624 +
625 + # Ensure engine id matches.
626 + if metadata.engine_id != expected_engine_id:
627 + raise RuntimeError(
628 + f"Remote NIXL agent engine ID mismatch. "
629 + f"Expected {expected_engine_id},"
630 + f"received {metadata.engine_id}."
631 + )
632 +
633 + # Register Remote agent.
634 + remote_agent_name = self.add_remote_agent(
635 + metadata, remote_rank, remote_tp_size
636 + )
637 + setup_agent_time = time.perf_counter()
638 + logger.debug(
639 + "NIXL handshake: add agent took: %s",
640 + setup_agent_time - got_metadata_time,
641 + )
642 + remote_rank_to_agent_name[remote_rank] = remote_agent_name
643 + return remote_rank_to_agent_name
644 +
645 + def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> None:
646 + """
647 + Initialize transfer buffer in CPU mem for accelerators
648 + NOT directly supported by NIXL (e.g., tpu)
649 + """
650 + xfer_buffers: dict[str, torch.Tensor] = {}
651 + inv_order = [0, 1, 3, 2, 4]
652 + try:
653 + for layer_name, kv_cache in kv_caches.items():
654 + kv_shape = kv_cache.shape
655 + kv_dtype = kv_cache.dtype
656 + permute_shape = False
657 + if (
658 + self.kv_cache_layout == "NHD"
659 + and self.vllm_config.kv_transfer_config is not None
660 + and self.vllm_config.kv_transfer_config.enable_permute_local_kv
661 + ):
662 + logger.info_once(
663 + "'enable_permute_local_kv' flag is enabled while "
664 + "device KV Layout is NHD. Init host buffer with"
665 + " HND to better support Decode/Prefill TP_ratio > 1."
666 + )
667 + # Since NHD will not support Decode/Prefill TP_ratio > 1,
668 + # we can leverage host_buffer for permute
669 + self.host_buffer_kv_cache_layout = "HND"
670 + kv_shape = (
671 + tuple(kv_shape[i] for i in inv_order)
672 + if not self.use_mla
673 + else kv_shape
674 + )
675 + permute_shape = not self.use_mla
676 +
677 + xfer_buffers[layer_name] = torch.empty(
678 + kv_shape, dtype=kv_dtype, device="cpu"
679 + )
680 + if permute_shape:
681 + xfer_buffers[layer_name] = xfer_buffers[layer_name].permute(
682 + inv_order
683 + )
684 + except MemoryError as e:
685 + logger.error("NIXLConnectorWorker gets %s.", e)
686 + raise
687 +
688 + self.host_xfer_buffers = xfer_buffers
689 +
690 + def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
691 + """Assign copy (d2h, h2d) operations when host buffer is used."""
692 + # Set a no-op if the host buffer is not cpu.
693 + if self.kv_buffer_device != "cpu":
694 + return
695 + # Set a no-op if self.device_type is 'cpu'.
696 + if self.device_type == "cpu":
697 + return
698 + assert self.use_host_buffer
699 + self.copy_blocks = copy_operation
700 +
701 + def _log_failure(
702 + self,
703 + failure_type: str,
704 + req_id: str | None,
705 + msg: str = "",
706 + error: Exception | None = None,
707 + meta: ReqMeta | None = None,
708 + **extra_context,
709 + ):
710 + """Log transfer failure with structured context for easier debugging."""
711 + context: dict[str, Any] = {
712 + "failure_type": failure_type,
713 + "request_id": req_id,
714 + "engine_id": self.engine_id,
715 + }
716 + if meta is None and req_id is not None:
717 + # Try to get metadata from in progress transfers when not provided
718 + meta = self._recving_metadata.get(req_id)
719 +
720 + if meta and meta.remote:
721 + context.update(
722 + {
723 + "remote_engine_id": meta.remote.engine_id,
724 + "remote_request_id": meta.remote.request_id,
725 + "remote_host": meta.remote.host,
726 + "remote_port": meta.remote.port,
727 + "num_local_blocks": sum(
728 + len(group) for group in meta.local_block_ids
729 + ),
730 + "num_remote_blocks": sum(
731 + len(group) for group in meta.remote.block_ids
732 + ),
733 + "local_block_ids_sample": meta.local_block_ids[0][:10]
734 + if meta.local_block_ids
735 + else [],
736 + }
737 + )
738 +
739 + context.update(extra_context)
740 + if msg:
741 + failure_type = f"{failure_type}. {msg}"
742 +
743 + logger.error(
744 + "NIXL transfer failure: %s | Context: %s",
745 + failure_type,
746 + context,
747 + exc_info=error is not None,
748 + stacklevel=2,
749 + )
750 +
751 + def _ensure_handshake(
752 + self,
753 + engine_id: EngineId,
754 + host: str,
755 + port: int,
756 + tp_size: int,
757 + ) -> Future[dict[int, str]] | None:
758 + """
759 + Ensure a handshake is in-flight (or already done) for *engine_id*.
760 +
761 + Returns the ``Future`` if a handshake is pending (or was just
762 + started), or ``None`` if the handshake already completed
763 + successfully. Callers can attach per-request callbacks to the
764 + returned future.
765 + Failures to handshake are logged and the request is marked as failed.
766 + """
767 + self._evict_stale_engines()
768 + with self._handshake_lock:
769 + if engine_id in self._remote_agents:
770 + return None
771 + fut = self._handshake_futures.get(engine_id)
772 + if fut is not None:
773 + return fut
774 + fut = self._handshake_initiation_executor.submit(
775 + self._nixl_handshake,
776 + host,
777 + port,
778 + tp_size,
779 + engine_id,
780 + )
781 + self._handshake_futures[engine_id] = fut
782 +
783 + def done_callback(f: Future[dict[int, str]], eid=engine_id):
784 + with self._handshake_lock:
785 + del self._handshake_futures[eid]
786 + try:
787 + self._remote_agents[eid] = f.result()
788 + self._engine_last_active[eid] = time.perf_counter()
789 + except Exception as e:
790 + self._log_failure(
791 + failure_type="handshake_setup_failed",
792 + req_id=None,
793 + error=e,
794 + remote_engine_id=eid,
795 + )
796 +
797 + fut.add_done_callback(done_callback)
798 + return fut
799 +
800 + def _background_nixl_handshake(
801 + self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta
802 + ):
803 + # Do NIXL handshake in background and add to _ready_requests when done.
804 + assert meta.remote is not None
805 + fut = self._ensure_handshake(
806 + remote_engine_id,
807 + meta.remote.host,
808 + meta.remote.port,
809 + meta.tp_size,
810 + )
811 + if fut is None:
812 + # Already handshaked — only happens if caller does not pre-check.
813 + self._ready_requests.put((req_id, meta))
814 + return
815 +
816 + # Check handshake success before proceeding with request.
817 + def request_ready(f: Future[Any], entry=(req_id, meta)):
818 + try:
819 + f.result()
820 + self._ready_requests.put(entry)
821 + except Exception as e:
822 + self._log_failure(
823 + failure_type="handshake_failed",
824 + req_id=req_id,
825 + error=e,
826 + meta=meta,
827 + )
828 + self._handle_failed_transfer(req_id, None)
829 +
830 + fut.add_done_callback(request_ready)
831 +
832 + def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
833 + """Register a cross-layers KV cache tensor with NIXL.
834 +
835 + `use_uniform_kv_cache()` guarantees a single KV cache group whose
836 + layers all share the same `AttentionSpec`, so any layer name from
837 + `_layer_specs` yields the correct per-layer spec for `page_size_bytes`.
838 + """
839 + first_layer = next(iter(self._layer_specs))
840 + # Forwarding a real layer name rather than a synthetic key
841 + self.register_kv_caches({first_layer: kv_cache})
842 +
843 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
844 + """Register the KV Cache data in nixl."""
845 + self.transfer_topo = TransferTopology(
846 + tp_rank=self.tp_rank,
847 + tp_size=self.world_size,
848 + block_size=self.block_size,
849 + engine_id=self.engine_id,
850 + is_mla=self.use_mla,
851 + total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
852 + attn_backends=self.attn_backends,
853 + # SSM States come in tuples (ssm, conv)
854 + tensor_shape=next(iter(kv_caches.values())).shape
855 + if not self._has_mamba
856 + else None,
857 + is_mamba=self._has_mamba,
858 + )
859 + self.compat_hash = compute_nixl_compatibility_hash(
860 + self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks
861 + )
862 +
863 + if self.use_host_buffer:
864 + self.initialize_host_xfer_buffer(kv_caches=kv_caches)
865 + assert len(self.host_xfer_buffers) == len(kv_caches), (
866 + f"host_buffer: {len(self.host_xfer_buffers)}, "
867 + f"kv_caches: {len(kv_caches)}"
868 + )
869 + xfer_buffers = self.host_xfer_buffers
870 + else:
871 + xfer_buffers = kv_caches
872 + assert not self.host_xfer_buffers, (
873 + "host_xfer_buffer should not be initialized when "
874 + f"kv_buffer_device is {self.kv_buffer_device}"
875 + )
876 +
877 + logger.info(
878 + "Registering KV_Caches. use_mla: %s, kv_buffer_device: %s, "
879 + "use_host_buffer: %s",
880 + self.use_mla,
881 + self.kv_buffer_device,
882 + self.use_host_buffer,
883 + )
884 +
885 + caches_data = []
886 + # With hybrid allocator, layers can share a kv cache tensor
887 + seen_base_addresses = []
888 +
889 + # Note(tms): I modified this from the original region setup code.
890 + # K and V are now in different regions. Advantage is that we can
891 + # elegantly support MLA and any cases where the K and V tensors
892 + # are non-contiguous (it's not locally guaranteed that they will be)
893 + # Disadvantage is that the encoded NixlAgentMetadata is now larger
894 + # (roughly 8KB vs 5KB).
895 + # Conversely for FlashInfer, K and V are registered in the same region
896 + # to better exploit the memory layout (ie num_blocks is the first dim).
897 + tensor_size_bytes = None
898 +
899 + for layer_name, cache_or_caches in xfer_buffers.items():
900 + # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
901 + # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
902 + # However, physical page_size may differ when kernel requires a specific
903 + # block size. This leads to SSM and FA layers having different num_blocks.
904 + # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.
905 + layer_spec = self._layer_specs.get(layer_name)
906 + if layer_spec is None:
907 + logger.debug(
908 + "Skipping layer %s as no KVCache spec is present. "
909 + "This is likely because the layer is sharing its KV cache",
910 + layer_name,
911 + )
912 + continue
913 + if isinstance(layer_spec, UniformTypeKVCacheSpecs):
914 + # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
915 + layer_spec = layer_spec.kv_cache_specs[layer_name]
916 + cache_list = self.transfer_topo.get_transfer_cache_regions(
917 + cache_or_caches, layer_spec
918 + )
919 + # `layer_spec.page_size_bytes` only accounts for logical page_size, that is
920 + # the page_size assuming constant `self._logical_num_blocks`.
921 + physical_page_size = (
922 + layer_spec.page_size_bytes
923 + if isinstance(layer_spec, MambaSpec)
924 + else layer_spec.page_size_bytes
925 + // self._physical_blocks_per_logical_kv_block
926 + )
927 + # For when registering multiple tensors eg K/V in separate regions.
928 + physical_page_size = physical_page_size // len(cache_list)
929 + if self.transfer_topo._cross_layers_blocks:
930 + # When cross-layers blocks are used, multiply by number of layers
931 + physical_page_size = physical_page_size * len(
932 + self.kv_cache_config.kv_cache_tensors
933 + )
934 + num_blocks = (
935 + self._logical_num_blocks
936 + if isinstance(layer_spec, MambaSpec)
937 + else self.num_blocks
938 + )
939 + # `page_size` accounts for physical blocks, st KVCache is always
940 + # [`num_blocks` * `page_size`]
941 + curr_tensor_size_bytes = num_blocks * physical_page_size
942 +
943 + # TODO (NickLucche) we could eventually unify how we handle FA/FI regions,
944 + # registering a single tensor for both K/V and splitting logically like FI.
945 + for cache in cache_list:
946 + base_addr = cache.data_ptr()
947 + if base_addr in seen_base_addresses:
948 + # NOTE (NickLucche) HMA employs memory pooling to share tensors
949 + # across groups. This results in skipping all tensors but the ones
950 + # pointed to by group0. Also, generally we will have more blocks
951 + # per tensor but fewer regions.
952 + logger.debug("Skipping %s because it's already seen", layer_name)
953 + continue
954 + logger.debug(
955 + "Registering layer %s with cache shape: %s", layer_name, cache.shape
956 + )
957 + seen_base_addresses.append(base_addr)
958 + # Only record non-Mamba page sizes.
959 + if isinstance(layer_spec, MambaSpec):
960 + self.block_len_per_layer.append(
961 + physical_page_size // self._physical_blocks_per_logical_kv_block
962 + )
963 + else:
964 + self.block_len_per_layer.append(physical_page_size)
965 + is_mla_region = isinstance(layer_spec, MLAAttentionSpec)
966 + self._region_is_mla.append(is_mla_region)
967 +
968 + if not is_mla_region:
969 + if tensor_size_bytes is None:
970 + tensor_size_bytes = curr_tensor_size_bytes
971 + assert tensor_size_bytes == curr_tensor_size_bytes, (
972 + "All non-MLA kv cache tensors must have the same size"
973 + )
974 +
975 + if cache.shape[0] != num_blocks:
976 + raise AssertionError(
977 + "All kv cache tensors must have the same number of "
978 + f"blocks; layer={layer_name}, "
979 + f"expected_num_blocks={num_blocks}, "
980 + f"cache_shape={tuple(cache.shape)}, "
981 + f"cache_stride={tuple(cache.stride())}, "
982 + f"layer_spec={type(layer_spec).__name__}, "
983 + f"backend={self.backend_name}, "
984 + "all_backends="
985 + f"{[backend.get_name() for backend in self.attn_backends]}, "
986 + f"kv_cache_layout={self.kv_cache_layout}, "
987 + "blocks_first="
988 + f"{self.transfer_topo.is_kv_layout_blocks_first}"
989 + )
990 +
991 + # Need to make sure the device ID is non-negative for NIXL,
992 + # Torch uses -1 to indicate CPU tensors.
993 + self.device_id = max(cache.get_device(), 0)
994 + caches_data.append(
995 + (base_addr, curr_tensor_size_bytes, self.device_id, "")
996 + )
997 +
998 + logger.debug(
999 + "Different block lengths collected: %s", set(self.block_len_per_layer)
1000 + )
1001 + assert (
1002 + len(self.block_len_per_layer)
1003 + == len(seen_base_addresses)
1004 + == len(self._region_is_mla)
1005 + )
1006 +
1007 + self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
1008 + self.num_regions = len(caches_data)
1009 +
1010 + if self.transfer_topo.virtually_split_kv_in_blocks:
1011 + # NOTE (NickLucche) When FlashInfer is used, memory is registered
1012 + # with joint KV for each block. This minimizes the overhead in
1013 + # registerMem allowing faster descs queries. In order to be able to
1014 + # split on kv_heads dim as required by heterogeneous TP, one must
1015 + # be able to index K/V separately. Hence we double the number
1016 + # of 'virtual' regions here and halve `block_len` below.
1017 + # Similarly for Mamba layers, we register SSM+Conv as a single region and
1018 + # then duplicate it logically to be able to index SSM/Conv separately.
1019 + # Exception: key-only REPLICATE regions (MLA) have no V half, so
1020 + # they contribute a single desc stream and are not doubled.
1021 + self.num_regions = sum(
1022 + 1 if self._is_region_replicated(i) else 2
1023 + for i in range(len(self._region_is_mla))
1024 + )
1025 +
1026 + # Total local FA descriptors (boundary between FA and mamba descs).
1027 + self.num_descs = self.num_regions * self.num_blocks
1028 +
1029 + descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
1030 + logger.debug("Registering descs: %s", caches_data)
1031 + self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
1032 + logger.debug("Done registering descs")
1033 + self._registered_descs.append(descs)
1034 +
1035 + self.device_kv_caches = kv_caches
1036 + self.dst_num_blocks[self.engine_id] = self.num_blocks
1037 +
1038 + if self._has_mamba:
1039 + logger.info(
1040 + "Hybrid SSM registration: num_blocks=%s, "
1041 + "logical_num_blocks=%s, ratio=%s, num_regions=%s, "
1042 + "num_descs=%s, mamba_ssm_size=%s, block_len_per_layer=%s",
1043 + self.num_blocks,
1044 + self._logical_num_blocks,
1045 + self._physical_blocks_per_logical_kv_block,
1046 + self.num_regions,
1047 + self.num_descs,
1048 + self._mamba_ssm_size,
1049 + set(self.block_len_per_layer),
1050 + )
1051 +
1052 + # Register local/src descr for NIXL xfer.
1053 + self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = (
1054 + self.register_local_xfer_handler(self.block_size)
1055 + )
1056 +
1057 + # After KV Caches registered, listen for new connections.
1058 + agent_metadata = NixlAgentMetadata(
1059 + engine_id=self.engine_id,
1060 + agent_metadata=self.nixl_wrapper.get_agent_metadata(),
1061 + device_id=self.device_id,
1062 + kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank],
1063 + num_blocks=self.num_blocks,
1064 + block_lens=self.block_len_per_layer,
1065 + kv_cache_layout=self.kv_cache_layout
1066 + if not self.use_host_buffer
1067 + else self.host_buffer_kv_cache_layout,
1068 + block_size=self.block_size,
1069 + ssm_sizes=self._mamba_ssm_size,
1070 + attn_backend_name=self.backend_name,
1071 + physical_blocks_per_logical_kv_block=(
1072 + self._physical_blocks_per_logical_kv_block
1073 + ),
1074 + )
1075 + # Wrap metadata in payload with hash for defensive decoding
1076 + assert self.compat_hash is not None
1077 + encoder = msgspec.msgpack.Encoder()
1078 + self.xfer_handshake_metadata = NixlHandshakePayload(
1079 + compatibility_hash=self.compat_hash,
1080 + agent_metadata_bytes=encoder.encode(agent_metadata),
1081 + )
1082 +
1083 + def _build_mamba_local(
1084 + self,
1085 + base_addresses: list[int],
1086 + block_size_ratio: int,
1087 + ) -> list[tuple[int, int, int]]:
1088 + """Build 4 desc regions (x, B, C, ssm) per layer for local mamba
1089 + blocks, enabling the 3-read transfer with DS conv layout."""
1090 + assert block_size_ratio == 1, (
1091 + "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
1092 + f"Got block_size_ratio={block_size_ratio}."
1093 + )
1094 + assert self._conv_decomp is not None
1095 + conv_offsets = self._conv_decomp.local_conv_offsets
1096 + conv_size, ssm_size = self._mamba_ssm_size
1097 + num_blocks = self._logical_num_blocks * block_size_ratio
1098 + physical_per_logical = self._physical_blocks_per_logical_kv_block
1099 +
1100 + result: list[tuple[int, int, int]] = []
1101 + for i, base_addr in enumerate(base_addresses):
1102 + # Jump one page_size, but ssm page_size may be bigger when kernel
1103 + # locks block size to a specific value (physical_per_logical scale).
1104 + page_stride = (
1105 + self.block_len_per_layer[i] // block_size_ratio * physical_per_logical
1106 + )
1107 + for off, sz in conv_offsets:
1108 + for blk in range(num_blocks):
1109 + result.append(
1110 + (base_addr + blk * page_stride + off, sz, self.device_id)
1111 + )
1112 + # SSM temporal state follows the conv state.
1113 + for blk in range(num_blocks):
1114 + result.append(
1115 + (
1116 + base_addr + blk * page_stride + conv_size,
1117 + ssm_size,
1118 + self.device_id,
1119 + )
1120 + )
1121 + return result
1122 +
1123 + def _build_mamba_remote(
1124 + self,
1125 + nixl_agent_meta: NixlAgentMetadata,
1126 + tp_ratio: int,
1127 + transfer_info: EngineTransferInfo,
1128 + ) -> list[tuple[int, int, int]]:
1129 + """Build 4 remote desc regions (proj0, proj1, proj2, ssm) per layer
1130 + for the 3-read transfer. For hetero-TP, each D rank reads only its
1131 + sub-projection slice from the P rank."""
1132 + assert self._conv_decomp is not None
1133 + effective_ratio = max(tp_ratio, 1)
1134 + # Mamba conv state is always TP-sharded, even when attention KV
1135 + # is replicated (num_kv_heads < tp_size).
1136 + local_offset = self.tp_rank % effective_ratio
1137 + conv_size_remote = nixl_agent_meta.ssm_sizes[0]
1138 +
1139 + conv_offsets = self._conv_decomp.remote_conv_offsets(local_offset, tp_ratio)
1140 + if tp_ratio >= 1:
1141 + ssm_read_size = self._mamba_ssm_size[1]
1142 + else:
1143 + ssm_read_size = nixl_agent_meta.ssm_sizes[1]
1144 +
1145 + remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical
1146 + num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical
1147 + device_id = nixl_agent_meta.device_id
1148 +
1149 + result: list[tuple[int, int, int]] = []
1150 + # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
1151 + # block lengths vary across layers (e.g. MLA).
1152 + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1153 + page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical
1154 + for off, sz in conv_offsets:
1155 + for blk in range(num_blocks):
1156 + result.append((base_addr + blk * page_stride + off, sz, device_id))
1157 + # SSM temporal state is also TP-sharded on the heads dimension.
1158 + for blk in range(num_blocks):
1159 + ssm_addr = (
1160 + base_addr
1161 + + blk * page_stride
1162 + + conv_size_remote
1163 + + local_offset * ssm_read_size
1164 + )
1165 + result.append((ssm_addr, ssm_read_size, device_id))
1166 + return result
1167 +
1168 + def _build_fa_local(
1169 + self,
1170 + base_addresses: list[int],
1171 + block_size_ratio: int,
1172 + ) -> list[tuple[int, int, int]]:
1173 + """Build local FA descriptors for all layers."""
1174 + assert self.transfer_topo is not None
1175 + num_blocks = self.num_blocks * block_size_ratio
1176 + result: list[tuple[int, int, int]] = []
1177 + for i, base_addr in enumerate(base_addresses):
1178 + kv_block_len = (
1179 + self.get_backend_aware_kv_block_len(
1180 + layer_idx=i, first_split=True, mamba_view=False
1181 + )
1182 + // block_size_ratio
1183 + )
1184 + page_stride = self.block_len_per_layer[i] // block_size_ratio
1185 + for block_id in range(num_blocks):
1186 + block_offset = block_id * page_stride
1187 + addr = base_addr + block_offset
1188 + result.append((addr, kv_block_len, self.device_id))
1189 +
1190 + if (
1191 + self.transfer_topo.virtually_split_kv_in_blocks
1192 + and not self._is_region_replicated(i)
1193 + ):
1194 + # Separate and interleave K/V regions to maintain the same
1195 + # descs ordering. This is needed for selecting contiguous heads
1196 + # when split across TP ranks. (Skipped for key-only REPLICATE.)
1197 + second_split = self.get_backend_aware_kv_block_len(
1198 + layer_idx=i, first_split=False, mamba_view=False
1199 + )
1200 + for block_id in range(num_blocks):
1201 + block_offset = block_id * page_stride
1202 + addr = base_addr + block_offset
1203 + v_addr = addr + kv_block_len
1204 + result.append((v_addr, second_split, self.device_id))
1205 + return result
1206 +
1207 + def _build_fa_remote(
1208 + self,
1209 + plan: TPMapping,
1210 + nixl_agent_meta: NixlAgentMetadata,
1211 + block_size_ratio: int,
1212 + ) -> list[tuple[int, int, int]]:
1213 + """Build remote FA descriptors for all layers."""
1214 + assert self.transfer_topo is not None
1215 + fa_group_idx = next(
1216 + i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
1217 + )
1218 + # SPLIT regions read their head slice from this many remote ranks at a
1219 + # per-rank offset; REPLICATE regions read the whole block once.
1220 + split_reads = len(plan.source_ranks_per_group[fa_group_idx])
1221 + num_blocks = nixl_agent_meta.num_blocks
1222 + result: list[tuple[int, int, int]] = []
1223 + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1224 + replicated = self._is_region_replicated(i)
1225 + # Read our whole local region size from remote..
1226 + local_block_len = self.get_backend_aware_kv_block_len(
1227 + layer_idx=i, first_split=True, mamba_view=False
1228 + )
1229 + remote_kv_block_len = local_block_len // block_size_ratio
1230 + if block_size_ratio > 1:
1231 + # ..using remote kv_block_len as transfer unit
1232 + local_block_len = remote_kv_block_len
1233 +
1234 + # REPLICATE reads the whole block once at offset 0; SPLIT gathers
1235 + # its head slice from `split_reads` remote ranks at a per-rank offset.
1236 + num_reads = 1 if replicated else split_reads
1237 + rank_offset = (
1238 + 0 if replicated else plan.rank_offset_factor * remote_kv_block_len
1239 + )
1240 + local_block_len = local_block_len // num_reads
1241 +
1242 + page_size = nixl_agent_meta.block_lens[i]
1243 + for block_id in range(num_blocks):
1244 + block_offset = block_id * page_size
1245 + # For each block, grab the kv heads chunk belonging to current local
1246 + # tp rank of size local_block_len.
1247 + addr = base_addr + block_offset + rank_offset
1248 + result.append((addr, local_block_len, nixl_agent_meta.device_id))
1249 +
1250 + emits_v = self.transfer_topo.virtually_split_kv_in_blocks and not replicated
1251 + if emits_v:
1252 + # With FlashInfer index V separately to allow head splitting.
1253 + second_split = self.get_backend_aware_kv_block_len(
1254 + layer_idx=i, first_split=False, mamba_view=False
1255 + )
1256 + second_split = second_split // num_reads
1257 + for block_id in range(num_blocks):
1258 + block_offset = block_id * page_size
1259 + addr = base_addr + block_offset + rank_offset
1260 + # Hop over the first split of remote page, K, to read V.
1261 + v_addr = addr + nixl_agent_meta.block_lens[i] // 2
1262 + result.append((v_addr, second_split, nixl_agent_meta.device_id))
1263 + return result
1264 +
1265 + def register_local_xfer_handler(
1266 + self,
1267 + block_size: int,
1268 + ) -> tuple[int, list[tuple[int, int, int]]]:
1269 + """
1270 + Function used for register local xfer handler with local block_size or
1271 + Remote block_size.
1272 +
1273 + When local block_size is same as remote block_size, we use local block_size
1274 + to register local_xfer_handler during init.
1275 +
1276 + When remote block size is less than local block size, we need to use
1277 + register another local_xfer_handler using remote block len to ensure
1278 + data copy correctness.
1279 + """
1280 + assert self.transfer_topo is not None
1281 + block_size_ratio = self.block_size // block_size
1282 + local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
1283 +
1284 + blocks_data = self._build_fa_local(local_base_addresses, block_size_ratio)
1285 + logger.debug(
1286 + "Created %s blocks for src engine %s and rank %s on device id %s",
1287 + len(blocks_data),
1288 + self.engine_id,
1289 + self.tp_rank,
1290 + self.device_id,
1291 + )
1292 + if self._has_mamba:
1293 + assert self.num_descs == len(blocks_data)
1294 + # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split
1295 + # is unnecessary — a single conv desc per block suffices. Consider
1296 + # adding a fast path that falls back to the standard 2-region
1297 + # registration (_build_fa_local mamba=True) when no hetero-TP
1298 + # remote has been seen. Currently we always register 4 regions
1299 + # because local descs are created before knowing the remote TP.
1300 + logger.debug("Registering local Mamba descriptors (4 regions/layer)")
1301 + blocks_data.extend(
1302 + self._build_mamba_local(local_base_addresses, block_size_ratio)
1303 + )
1304 +
1305 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
1306 + # NIXL_INIT_AGENT to be used for preparations of local descs.
1307 + return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data
1308 +
1309 + def add_remote_agent(
1310 + self,
1311 + nixl_agent_meta: NixlAgentMetadata,
1312 + remote_tp_rank: int = 0,
1313 + remote_tp_size: int = 1,
1314 + ) -> str:
1315 + """
1316 + Add the remote NIXL agent and prepare the descriptors for reading cache
1317 + blocks from remote.
1318 +
1319 + In particular, handle both homogeneous and heterogeneous TP. The former
1320 + requires local rank_i to read from remote rank_i.
1321 + The latter, in the case of D.world_size < P.world_size, requires that a
1322 + local (D) TP worker reads from multiple remote (P) TP workers.
1323 + Conversely, assuming D.world_size > P.world_size, two or more local TP
1324 + workers will read from a single remote TP worker.
1325 +
1326 + Here's an example for the last case described above (non-MLA):
1327 +
1328 + rank_offset p_remote_tp_rank
1329 + (kv split no)
1330 + --------------------------------
1331 + 0 0 Worker0 ---- 1st half of KV ----> Worker0 [ KV Cache ]
1332 + /
1333 + 1 0 Worker1 ---- 2nd half of KV -----/
1334 +
1335 + 0 1 Worker2 ---- 1st half of KV ----> Worker1 [ KV Cache ]
1336 + /
1337 + 1 1 Worker3 ---- 2nd half of KV -----/
1338 +
1339 +
1340 + Decoder TP workers Prefix TP workers
1341 + (world_size=4) (world_size=2)
1342 + tp_ratio = 4 // 2 = 2
1343 +
1344 + Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim]
1345 + then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "HND" layout format.
1346 + Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio
1347 + first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split
1348 + along the kv_heads dimension, and so forth until "tp_ratio" D TP workers have pulled from P-Worker0.
1349 +
1350 + Note that the above will also hold true for the homogeneous TP case, where tp_ratio evaluates to 1.
1351 +
1352 + Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0
1353 + so that the whole cache is shared by "tp_ratio" D TP workers.
1354 +
1355 + For Mamba hetero-TP, both tp_ratio > 0 (D_TP > P_TP) and
1356 + tp_ratio < 0 (P_TP > D_TP) are supported by the 3-read transfer.
1357 + """ # noqa: E501
1358 + engine_id = nixl_agent_meta.engine_id
1359 + # TODO re-evaluate refreshing for scaling/recovery
1360 + if remote_tp_rank in self._remote_agents.get(engine_id, {}):
1361 + logger.debug(
1362 + "Remote agent with engine_id %s and rank"
1363 + "%s already exchanged metadata, skip handshake.",
1364 + engine_id,
1365 + remote_tp_rank,
1366 + )
1367 + return self._remote_agents[engine_id][remote_tp_rank]
1368 +
1369 + ### Register remote engine in TransferTopology (idempotent).
1370 + assert self.transfer_topo is not None
1371 + transfer_topo = self.transfer_topo
1372 + physical_blocks_per_logical = (
1373 + nixl_agent_meta.physical_blocks_per_logical_kv_block
1374 + )
1375 + transfer_info = EngineTransferInfo(
1376 + remote_tp_size=remote_tp_size,
1377 + remote_block_size=nixl_agent_meta.block_size,
1378 + remote_block_len=nixl_agent_meta.block_lens[0],
1379 + remote_physical_blocks_per_logical=physical_blocks_per_logical,
1380 + )
1381 + transfer_topo.register_remote_engine(engine_id, transfer_info)
1382 + logger.info("Transfer plan: %s", transfer_topo.describe(engine_id))
1383 +
1384 + self.tp_mappings[engine_id] = compute_tp_mapping(
1385 + transfer_topology=transfer_topo,
1386 + remote_tp_size=remote_tp_size,
1387 + group_spec_types=self._group_spec_types,
1388 + )
1389 +
1390 + remote_agent_name = self.nixl_wrapper.add_remote_agent(
1391 + nixl_agent_meta.agent_metadata
1392 + )
1393 +
1394 + # Create dst descs and xfer side handles. TP workers have same #blocks
1395 + # so we only register once per engine_id.
1396 + # Example:
1397 + # block_size_ratio > 1:
1398 + # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
1399 + # local origin:| 0| 1| 8| 12|
1400 + # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
1401 + block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size)
1402 +
1403 + if engine_id not in self.dst_num_blocks:
1404 + self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
1405 +
1406 + # Keep track of remote agent kv caches base addresses.
1407 + self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
1408 + nixl_agent_meta.kv_caches_base_addr
1409 + )
1410 + self._validate_remote_agent_handshake(nixl_agent_meta, remote_tp_size)
1411 +
1412 + # This is 1 when P and D `--tensor-parallel-size` match. Otherwise,
1413 + # this is the ratio between the two sizes.
1414 + tp_ratio = transfer_topo.tp_ratio(remote_tp_size)
1415 +
1416 + logger.debug(
1417 + "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s",
1418 + engine_id,
1419 + remote_tp_rank,
1420 + tp_ratio,
1421 + )
1422 +
1423 + plan = self.tp_mappings[engine_id]
1424 +
1425 + ### (Optional) Register local agent memory regions. MLA is not split.
1426 + if (
1427 + tp_ratio < 0
1428 + and not self.use_mla
1429 + and tp_ratio not in self.src_xfer_handles_by_tp_ratio
1430 + ):
1431 + # Remote tp_size > local tp_size: read from multiple remote ranks.
1432 + # Logically "split" own regions into |tp_ratio| chunks. Mind that
1433 + # we only do this once per remote tp_size (replica-friendly).
1434 + self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
1435 +
1436 + for handle_data in self._build_local_splits_from_plan(
1437 + plan,
1438 + self.src_blocks_data,
1439 + self.num_descs,
1440 + ):
1441 + descs = self.nixl_wrapper.get_xfer_descs(
1442 + handle_data, self.nixl_memory_type
1443 + )
1444 + handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
1445 + self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
1446 +
1447 + ### Register remote agent memory regions
1448 + # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With
1449 + # heterogeneous TP, prepare the descriptors by splitting the P KV cache along
1450 + # kv_head dim, of D worker's kv_head size (D>P).
1451 + # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].
1452 +
1453 + # Register all remote blocks, but only the corresponding kv heads.
1454 + blocks_data = self._build_fa_remote(
1455 + plan,
1456 + nixl_agent_meta,
1457 + block_size_ratio,
1458 + )
1459 + logger.debug(
1460 + "Created %s blocks for dst engine %s with remote rank %s and local rank %s",
1461 + len(blocks_data),
1462 + engine_id,
1463 + remote_tp_rank,
1464 + self.tp_rank,
1465 + )
1466 + if self._has_mamba:
1467 + logger.debug(
1468 + "Registering remote Mamba blocks for engine %s rank %s",
1469 + engine_id,
1470 + remote_tp_rank,
1471 + )
1472 + blocks_data.extend(
1473 + self._build_mamba_remote(
1474 + nixl_agent_meta,
1475 + tp_ratio,
1476 + transfer_info,
1477 + )
1478 + )
1479 +
1480 + # Register with NIXL.
1481 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
1482 + self.dst_xfer_side_handles[engine_id][remote_tp_rank] = (
1483 + self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs)
1484 + )
1485 +
1486 + if block_size_ratio > 1:
1487 + # when prefill with smaller block_size, we need to init a
1488 + # new handler with same block_len to match
1489 + self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = (
1490 + self.register_local_xfer_handler(nixl_agent_meta.block_size)[0]
1491 + )
1492 +
1493 + return remote_agent_name
1494 +
1495 + def _validate_remote_agent_handshake(
1496 + self, nixl_agent_meta: NixlAgentMetadata, remote_tp_size: int
1497 + ):
1498 + """
1499 + Validate the remote agent handshake metadata ensuring the
1500 + invariants hold true.
1501 + """
1502 + remote_engine_id = nixl_agent_meta.engine_id
1503 +
1504 + assert self.transfer_topo is not None
1505 + remote_info = self.transfer_topo.get_engine_info(remote_engine_id)
1506 + assert remote_info.remote_tp_size == remote_tp_size
1507 +
1508 + tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
1509 + block_size_ratio = self.transfer_topo.block_size_ratio(
1510 + nixl_agent_meta.block_size
1511 + )
1512 + # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba.
1513 + # Mamba models can have replicated FA KV with tp_ratio < 0.
1514 + # MLA models do not need to handle kv replication.
1515 + if not self.use_mla and not self._has_mamba:
1516 + assert not (
1517 + tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id)
1518 + )
1519 +
1520 + remote_physical_per_logical = (
1521 + nixl_agent_meta.physical_blocks_per_logical_kv_block
1522 + )
1523 + if (
1524 + self._has_mamba
1525 + and remote_physical_per_logical
1526 + != self._physical_blocks_per_logical_kv_block
1527 + and self.vllm_config.cache_config.enable_prefix_caching
1528 + ):
1529 + raise RuntimeError(
1530 + "Prefix caching with heterogeneous physical_blocks_per_logical "
1531 + "is not supported for Mamba hybrid models. "
1532 + f"Local: {self._physical_blocks_per_logical_kv_block}, "
1533 + f"Remote: {remote_physical_per_logical}. "
1534 + "Disable prefix caching with --no-enable-prefix-caching."
1535 + )
1536 +
1537 + if self._is_hma_required:
1538 + assert block_size_ratio == 1, (
1539 + "HMA does not support different remote block size yet"
1540 + )
1541 + kv_cache_layout = (
1542 + self.kv_cache_layout
1543 + if not self.use_host_buffer
1544 + else self.host_buffer_kv_cache_layout
1545 + )
1546 + if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout:
1547 + if (
1548 + self.kv_transfer_config.enable_permute_local_kv
1549 + and nixl_agent_meta.kv_cache_layout == "HND"
1550 + ):
1551 + logger.info(
1552 + "Remote is HND and local is NHD, enabled additional permute "
1553 + "on local device KV."
1554 + )
1555 + assert not self._is_hma_required, (
1556 + "HMA does not support block size post processing"
1557 + )
1558 + self.enable_permute_local_kv = True
1559 + else:
1560 + raise RuntimeError(
1561 + "Heterogeneous TP expects same kv_cache_layout. "
1562 + "Or enable experimental feature to use HND to NHD support by "
1563 + "setting 'enable_permute_local_kv'=True in --kv-transfer-config."
1564 + )
1565 + # if remote_agent used attn is not same as local,
1566 + # hint heterogenuous attn post process
1567 + if (
1568 + nixl_agent_meta.attn_backend_name != self.backend_name
1569 + and self.backend_name in ["CPU_ATTN"]
1570 + ):
1571 + if self._is_hma_required:
1572 + raise RuntimeError(
1573 + "heterogeneous attn post process is not supported with HMA"
1574 + )
1575 + logger.info(
1576 + "[Experimental] CPU_ATTN backend is used, "
1577 + "hint heterogeneous attn post process"
1578 + )
1579 + self.enable_heterogeneous_attn_post_process = True
1580 +
1581 + # Heterogeneous TP requires head-splitting, which only works with
1582 + # HND layout. MLA and replicated-KV cases don't split on heads.
1583 + # Mamba doesn't support heterogeneous TP.
1584 + if (
1585 + abs(tp_ratio) != 1
1586 + and not self.use_mla
1587 + and not self.transfer_topo.is_kv_replicated(remote_engine_id)
1588 + and kv_cache_layout != "HND"
1589 + and not self.enable_permute_local_kv
1590 + ):
1591 + raise RuntimeError(
1592 + "Heterogeneous TP head-dimension splitting requires contiguous heads. "
1593 + "Use HND layout on the prefill side."
1594 + )
1595 +
1596 + # Per-region block_len validation enforcing the P/D invariant.
1597 + # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer)
1598 + # only allow the number of blocks to differ; SPLIT regions scale with
1599 + # tp_ratio. Mamba uses the ssm_sizes counterpart, so skip block_len here.
1600 + if not self._has_mamba:
1601 + assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (
1602 + "Number of KV layers must match between prefill and decode"
1603 + )
1604 + model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated(
1605 + remote_engine_id
1606 + )
1607 + for i, local_len in enumerate(self.block_len_per_layer):
1608 + replicated = model_replicated or self._is_region_replicated(i)
1609 + remote_len = nixl_agent_meta.block_lens[i]
1610 + if replicated:
1611 + assert local_len // block_size_ratio == remote_len, (
1612 + "KV cache sizes must match between P and D when "
1613 + f"replicated (region {i}: local={local_len}, "
1614 + f"remote={remote_len}, bsr={block_size_ratio})."
1615 + )
1616 + elif tp_ratio > 0:
1617 + assert remote_len == (local_len * tp_ratio) // block_size_ratio, (
1618 + f"SPLIT region {i}: remote P KV block_len {remote_len} "
1619 + f"must equal local {local_len} * tp_ratio {tp_ratio} "
1620 + f"// block_size_ratio {block_size_ratio}."
1621 + )
1622 + else:
1623 + assert block_size_ratio == 1, (
1624 + "Different local/remote block sizes are not supported "
1625 + "when P TP > D TP."
1626 + )
1627 + assert remote_len == local_len // (-tp_ratio), (
1628 + f"SPLIT region {i}: remote P KV block_len "
1629 + f"{remote_len} must equal local {local_len} "
1630 + f"// |tp_ratio| {-tp_ratio}."
1631 + )
1632 +
1633 + # TP workers that handhshake with same remote have same #blocks.
1634 + assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
1635 + # Same number of regions/~layers.
1636 + assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer)
1637 +
1638 + def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
1639 + """copy recved kv from host buffer to device."""
1640 + assert self.use_host_buffer
1641 + assert self.copy_blocks is not None
1642 +
1643 + local_block_ids = meta.local_physical_block_ids
1644 + # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups
1645 + for group_block_ids in local_block_ids:
1646 + self.copy_blocks(
1647 + self.host_xfer_buffers,
1648 + self.device_kv_caches,
1649 + group_block_ids,
1650 + group_block_ids,
1651 + "h2d",
1652 + )
1653 + if logger.isEnabledFor(logging.DEBUG):
1654 + logger.debug(
1655 + "synced recved kv of request[%s] to device kv buffer,"
1656 + "local_block_ids: %s. ",
1657 + req_id,
1658 + ",".join(map(str, local_block_ids)),
1659 + )
1660 +
1661 + def save_kv_to_host(self, metadata: NixlConnectorMetadata):
1662 + """copy kv from device to host buffer."""
1663 + assert self.use_host_buffer
1664 + assert self.copy_blocks is not None
1665 +
1666 + for req_id, meta in metadata.reqs_to_save.items():
1667 + meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
1668 + meta.local_block_ids
1669 + )
1670 + if logger.isEnabledFor(logging.DEBUG):
1671 + logger.debug(
1672 + "save_load_kv for request[%s] to host xfer buffer."
1673 + "local_block_ids: %s. ",
1674 + req_id,
1675 + ",".join(map(str, meta.local_physical_block_ids)),
1676 + )
1677 + # blocking
1678 + for group_block_ids in meta.local_physical_block_ids:
1679 + self.copy_blocks(
1680 + self.device_kv_caches,
1681 + self.host_xfer_buffers,
1682 + group_block_ids,
1683 + group_block_ids,
1684 + "d2h",
1685 + )
1686 +
1687 + def post_process_device_kv_on_receive(
1688 + self,
1689 + block_size_ratio: int,
1690 + block_ids_list: list[list[int]],
1691 + ):
1692 + """
1693 + Post process device kv cache after receiving from remote.
1694 +
1695 + 3 types of post processing supported:
1696 + * kv_cache_postprocess_layout => convert from HND to NHD
1697 + * kv_cache_postprocess_blksize => convert from small block size
1698 + to large block size
1699 + * kv_cache_postprocess_blksize_and_layout => convert from small
1700 + block size to large block size and convert from HND to NHD
1701 +
1702 + """
1703 + if len(self.device_kv_caches) == 0:
1704 + return
1705 + assert block_size_ratio >= 1, "Only nP < nD supported currently."
1706 + assert self.transfer_topo is not None
1707 + if self.enable_permute_local_kv and block_size_ratio > 1:
1708 + logger.debug(
1709 + "Post-processing device kv cache on receive by converting "
1710 + "block_size with %sx bigger and permuting layout from HND"
1711 + " to NHD.",
1712 + block_size_ratio,
1713 + )
1714 + elif self.enable_permute_local_kv:
1715 + logger.debug(
1716 + "Post-processing device kv cache on receive by permuting layout"
1717 + "from HND to NHD."
1718 + )
1719 + else:
1720 + logger.debug(
1721 + "Post-processing device kv cache on receive by converting "
1722 + "block_size with %sx bigger.",
1723 + block_size_ratio,
1724 + )
1725 +
1726 + split_k_and_v = self.transfer_topo.split_k_and_v
1727 +
1728 + for block_ids in block_ids_list:
1729 + indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
1730 +
1731 + for _, cache_or_caches in self.device_kv_caches.items():
1732 + cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1733 + for cache in cache_list:
1734 + if self.enable_permute_local_kv and block_size_ratio > 1:
1735 + kv_postprocess_blksize_and_layout_on_receive(
1736 + cache, indices, block_size_ratio
1737 + )
1738 + elif self.enable_permute_local_kv:
1739 + kv_postprocess_layout_on_receive(cache, indices)
1740 + else:
1741 + kv_postprocess_blksize_on_receive(
1742 + cache, indices, block_size_ratio
1743 + )
1744 +
1745 + def post_process_device_kv_on_receive_heterogeneous_attn(
1746 + self, block_ids: list[int]
1747 + ):
1748 + """
1749 + Post process device kv cache after receiving from remote
1750 + for heterogeneous attention.
1751 + """
1752 + assert self.enable_heterogeneous_attn_post_process
1753 +
1754 + indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
1755 +
1756 + for _, cache_or_caches in self.device_kv_caches.items():
1757 + blocks_to_update = cache_or_caches.index_select(1, indices)
1758 + current_platform.pack_kv_cache(
1759 + key=blocks_to_update[0],
1760 + value=blocks_to_update[1],
1761 + key_cache=cache_or_caches[0],
1762 + value_cache=cache_or_caches[1],
1763 + block_ids=block_ids,
1764 + indices=indices,
1765 + )
1766 +
1767 + def get_finished(self) -> tuple[set[str], set[str]]:
1768 + """
1769 + Get requests that are done sending or recving on this specific worker.
1770 + The scheduler process (via the MultiprocExecutor) will use this output
1771 + to track which workers are done.
1772 + """
1773 + assert self.transfer_topo is not None
1774 + done_sending = self._get_new_notifs()
1775 + done_recving = self._pop_done_transfers(self._recving_transfers)
1776 +
1777 + # Drain queue of requests where handshake or transfer setup failed.
1778 + failed_recv_reqs = set[ReqId]()
1779 + while not self._failed_recv_reqs.empty():
1780 + try:
1781 + failed_recv_reqs.add(self._failed_recv_reqs.get_nowait())
1782 + except queue.Empty:
1783 + break
1784 +
1785 + # Add failed requests to done_recving for scheduler tracking
1786 + # (blocks are already marked invalid, scheduler will handle recompute)
1787 + done_recving.update(failed_recv_reqs)
1788 +
1789 + if len(done_sending) > 0 or len(done_recving) > 0:
1790 + logger.debug(
1791 + "Rank %s, get_finished: %s requests done sending "
1792 + "and %s requests done recving (%s failed)",
1793 + self.tp_rank,
1794 + len(done_sending),
1795 + len(done_recving),
1796 + len(failed_recv_reqs),
1797 + )
1798 +
1799 + block_ids_for_blocksize_post_process = defaultdict(list)
1800 + block_ids_for_heterogeneous_attn_post_process = list[list[int]]()
1801 + for req_id in done_recving:
1802 + # clean up metadata for completed requests
1803 + meta = self._recving_metadata.pop(req_id, None)
1804 + assert meta is not None, f"{req_id} not found in recving_metadata list"
1805 +
1806 + # Skip KV sync and post-processing for failed requests
1807 + if req_id in failed_recv_reqs:
1808 + logger.warning(
1809 + "Skipping KV post-processing for failed request %s",
1810 + req_id,
1811 + )
1812 + continue
1813 +
1814 + assert meta.remote is not None
1815 + if self.use_host_buffer:
1816 + self.sync_recved_kv_to_device(req_id, meta)
1817 +
1818 + # post processing for heteroblocksize
1819 + remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id)
1820 + block_size_ratio = self.transfer_topo.block_size_ratio(
1821 + remote_info.remote_block_size
1822 + )
1823 + if not self.use_mla and (
1824 + block_size_ratio > 1 or self.enable_permute_local_kv
1825 + ):
1826 + assert not self._is_hma_required
1827 + block_ids_for_blocksize_post_process[block_size_ratio].append(
1828 + meta.local_physical_block_ids[0]
1829 + )
1830 + # post processing for heterogeneous attention
1831 + if self.enable_heterogeneous_attn_post_process:
1832 + block_ids_for_heterogeneous_attn_post_process.append(
1833 + meta.local_physical_block_ids[0]
1834 + )
1835 + for (
1836 + block_size_ratio,
1837 + block_ids_list,
1838 + ) in block_ids_for_blocksize_post_process.items():
1839 + self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list)
1840 +
1841 + for block_ids in block_ids_for_heterogeneous_attn_post_process:
1842 + self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
1843 +
1844 + # Handle timeout to avoid stranding blocks on remote.
1845 + now = time.perf_counter()
1846 + while self._reqs_to_send:
1847 + req_id, expires = next(iter(self._reqs_to_send.items()))
1848 + # Sorted dict, oldest requests are put first so we can exit early.
1849 + if now < expires:
1850 + break
1851 + count = self.consumer_notification_counts_by_req.pop(req_id, 0)
1852 + self.xfer_stats.record_kv_expired_req()
1853 + logger.warning(
1854 + "Releasing expired KV blocks for request %s which were "
1855 + "retrieved by %d remote worker(s) before lease expired.",
1856 + req_id,
1857 + count,
1858 + )
1859 + self._reqs_to_process.remove(req_id)
1860 + del self._reqs_to_send[req_id]
1861 + done_sending.add(req_id)
1862 +
1863 + return done_sending, done_recving
1864 +
1865 + def _get_new_notifs(self) -> set[str]:
1866 + """Get req_ids which got a remote xfer notification.
1867 +
1868 + Subclasses must implement this to handle mode-specific notifications.
1869 + """
1870 + raise NotImplementedError
1871 +
1872 + def _handle_heartbeat(self, payload: str) -> None:
1873 + """Extend leases for requests referenced in a heartbeat.
1874 +
1875 + Args:
1876 + payload: comma-separated P-side request IDs, e.g.
1877 + "req_abc,req_def".
1878 + """
1879 + new_expiry = time.perf_counter() + self._lease_extension
1880 + for req_id in payload.split(","):
1881 + if req_id in self._reqs_to_send:
1882 + old = self._reqs_to_send[req_id]
1883 + self._reqs_to_send[req_id] = max(old, new_expiry)
1884 + logger.debug(
1885 + "Heartbeat extended lease for request %s "
1886 + "by %ds (old_expiry=%.1f, new_expiry=%.1f)",
1887 + req_id,
1888 + self._lease_extension,
1889 + old,
1890 + new_expiry,
1891 + )
1892 +
1893 + def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:
1894 + """
1895 + Pop completed xfers by checking for DONE state.
1896 + Args:
1897 + transfers: dict of req_id -> list[running_xfer]
1898 + Returns:
1899 + set of req_ids that have all done xfers
1900 + """
1901 + done_req_ids: set[str] = set()
1902 + for req_id, handles in list(transfers.items()):
1903 + in_progress = []
1904 + for handle in handles:
1905 + try:
1906 + xfer_state = self.nixl_wrapper.check_xfer_state(handle)
1907 + if xfer_state == "DONE":
1908 + # Get telemetry from NIXL
1909 + res = self.nixl_wrapper.get_xfer_telemetry(handle)
1910 + self.xfer_stats.record_transfer(res)
1911 + self.nixl_wrapper.release_xfer_handle(handle)
1912 + elif xfer_state == "PROC":
1913 + in_progress.append(handle)
1914 + continue
1915 + else:
1916 + self._log_failure(
1917 + failure_type="transfer_failed",
1918 + msg="Marking blocks as invalid",
1919 + req_id=req_id,
1920 + xfer_state=xfer_state,
1921 + )
1922 + self._handle_failed_transfer(req_id, handle)
1923 + except Exception as e:
1924 + self._log_failure(
1925 + failure_type="transfer_exception",
1926 + msg="Marking blocks as invalid",
1927 + req_id=req_id,
1928 + error=e,
1929 + )
1930 + self._handle_failed_transfer(req_id, handle)
1931 +
1932 + if not in_progress:
1933 + # Only report request as completed when all transfers are done.
1934 + done_req_ids.add(req_id)
1935 + del transfers[req_id]
1936 + else:
1937 + transfers[req_id] = in_progress
1938 + return done_req_ids
1939 +
1940 + def _handle_failed_transfer(self, req_id: str, handle: int | None):
1941 + """
1942 + Handle a failed transfer by marking all (logical) blocks as invalid and
1943 + recording the failure.
1944 +
1945 + Args:
1946 + req_id: The request ID.
1947 + handle: The transfer handle.
1948 + """
1949 + # Use .get() here as the metadata cleanup is handled by get_finished()
1950 + # TODO (NickLucche) handle failed transfer for HMA.
1951 + if (meta := self._recving_metadata.get(req_id)) and not self._is_hma_required:
1952 + self._invalid_block_ids.put(set(meta.local_block_ids[0]))
1953 + self._failed_recv_reqs.put(req_id)
1954 + if handle is not None:
1955 + self.nixl_wrapper.release_xfer_handle(handle)
1956 + self.xfer_stats.record_failed_transfer()
1957 +
1958 + def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
1959 + """
1960 + Send heartbeat notifications to remote engines, extending lease on KV blocks.
1961 + """
1962 + for engine_id, hb_info in metadata.heartbeat_by_engine.items():
1963 + # Proactive handshake (this request may still be in waiting queue) so
1964 + # the **next** heartbeat for this remote can go through.
1965 + if (
1966 + self._ensure_handshake(
1967 + engine_id, hb_info.host, hb_info.port, hb_info.tp_size
1968 + )
1969 + is not None
1970 + ):
1971 + continue # handshake is still pending
1972 +
1973 + # Build the heartbeat message: "HB:req1,req2,..."
1974 + hb_msg = ("HB:" + ",".join(hb_info.req_ids)).encode()
1975 + for agent_name in self._remote_agents[engine_id].values():
1976 + try:
1977 + self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg)
1978 + except Exception:
1979 + logger.debug(
1980 + "Failed to send heartbeat to engine %s",
1981 + engine_id,
1982 + exc_info=True,
1983 + )
1984 +
1985 + def get_mapped_blocks(
1986 + self, block_ids: np.ndarray, block_size_ratio: int
1987 + ) -> np.ndarray:
1988 + """
1989 + Calculates the new set of block IDs by mapping every element
1990 + in the (potentially sparse) input array.
1991 + Example: block_ids=[0, 2], block_size_ratio=2
1992 + get_mapped_blocks 0 1 [2 3] 4 5
1993 + # remote is |h0-b0|h1-b0||h0-b1|h1-b1||h0-b1|h1-b1||
1994 + # local is |h0-b0......||h1-b0......||h2-b0........
1995 + local_block_ids 0 [1] 2
1996 + """
1997 + if block_ids.size == 0:
1998 + return np.array([], dtype=np.int64)
1999 +
2000 + start_ids = block_ids * block_size_ratio
2001 + offsets = np.arange(block_size_ratio)
2002 + mapped_2d = start_ids[:, None] + offsets[None, :]
2003 +
2004 + return mapped_2d.flatten().astype(np.int64)
2005 +
2006 + def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:
2007 + """
2008 + Convert logical block ids to kernel physical block ids.
2009 + This is required when the logical block size (the one set by the user)
2010 + does not match the one required by the attn backend.
2011 + """
2012 + if self._physical_blocks_per_logical_kv_block == 1:
2013 + # Noop when physical and logical block sizes are the same
2014 + return block_ids
2015 + block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape(
2016 + 1, -1
2017 + )
2018 + # Mamba blocks have no logical<>physical discrepancy
2019 + group_specs = self.kv_cache_config.kv_cache_groups
2020 + return [
2021 + BlockTable.map_to_kernel_blocks(
2022 + np.array(group),
2023 + self._physical_blocks_per_logical_kv_block,
2024 + block_arange,
2025 + ).tolist()
2026 + if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
2027 + else group
2028 + for i, group in enumerate(block_ids)
2029 + ]
2030 +
2031 + def _apply_prefix_caching(
2032 + self,
2033 + local_block_ids: BlockIds,
2034 + remote_block_ids: BlockIds,
2035 + remote_physical_per_logical: int,
2036 + ) -> tuple[BlockIds, list]:
2037 + """Apply prefix caching by trimming local/remote block ID lists.
2038 +
2039 + For non-Mamba models: end-trim remote to match local count, so that
2040 + already-cached prefix blocks are skipped in the transfer.
2041 +
2042 + For Mamba hybrid (prefix caching not yet supported): front-trim both
2043 + to the minimum count to handle kernel block count discrepancies from
2044 + logical block rounding in heterogeneous TP.
2045 + """
2046 + # Partial prefix cache hit: just read uncomputed blocks.
2047 + # Skip mamba groups — their blocks represent full state (conv+ssm),
2048 + # not per-token data, so trimming would corrupt the transfer.
2049 + remote_block_ids = list(remote_block_ids)
2050 + if not self._has_mamba:
2051 + for i, remote_group in enumerate(remote_block_ids):
2052 + num_local_blocks = len(local_block_ids[i])
2053 + assert num_local_blocks <= len(remote_group)
2054 + if num_local_blocks < len(remote_group):
2055 + remote_block_ids[i] = remote_group[-num_local_blocks:]
2056 + else:
2057 + # (NOTE: ZhanqiuHu) Mamba hybrid: no prefix caching support so far.HeteroTP
2058 + # can cause different kernel block counts due to logical block rounding.
2059 + # Example: 640 prompt tokens, kernel_block_size=64
2060 + # remote physical_per_logical=10, local physical_per_logical=6
2061 + # remote logical ids from kv_transfer_params = [0]
2062 + # local logical ids allocated = [0, 1]
2063 + # remote kernel blocks: [0..9] (1*10=10)
2064 + # local kernel blocks: [0..11] (2*6=12)
2065 + # actual data blocks = ceil(640/64) = 10, trim both to 10
2066 + # Vice versa (remote physical_per_logical=6, local=10):
2067 + # remote logical ids = [0, 1], local logical ids = [0]
2068 + # remote kernel blocks: [0..11] (2*6=12)
2069 + # local kernel blocks: [0..9] (1*10=10)
2070 + # actual data blocks = ceil(640/64) = 10, trim both to 10
2071 + local_block_ids = list(local_block_ids)
2072 + for i, remote_group in enumerate(remote_block_ids):
2073 + num_local_blocks = len(local_block_ids[i])
2074 + num_remote_blocks = len(remote_group)
2075 + if (
2076 + _is_ssm_spec(self._group_spec_types[i])
2077 + and num_local_blocks < num_remote_blocks
2078 + ):
2079 + # NOTE (NickLucche): With prefix caching on SSM, (remote) blocks
2080 + # prior to the last one are placeholders (null blocks). Mind that
2081 + # this doesn't really impact transfer, as we only still care about
2082 + # the last "block", the full in-place state.
2083 + assert num_local_blocks == 1, "SSM can only have one local block"
2084 + remote_block_ids[i] = remote_group[-num_local_blocks:]
2085 + elif (
2086 + self._physical_blocks_per_logical_kv_block
2087 + == remote_physical_per_logical
2088 + and num_local_blocks < num_remote_blocks
2089 + ):
2090 + # Partial prefix cache hit for FA group.
2091 + remote_block_ids[i] = remote_group[-num_local_blocks:]
2092 + else:
2093 + # TODO Handle prefix caching with different block_sizes
2094 + max_padding = max(
2095 + self._physical_blocks_per_logical_kv_block,
2096 + remote_physical_per_logical,
2097 + )
2098 + assert abs(num_local_blocks - num_remote_blocks) < max_padding, (
2099 + f"Group {i}: |{num_local_blocks} - "
2100 + f"{num_remote_blocks}| >= {max_padding}"
2101 + )
2102 + num_blocks = min(num_local_blocks, num_remote_blocks)
2103 + local_block_ids[i] = local_block_ids[i][:num_blocks]
2104 + remote_block_ids[i] = remote_group[:num_blocks]
2105 + return local_block_ids, remote_block_ids
2106 +
2107 + def _logical_to_remote_kernel_block_ids(
2108 + self, block_ids: BlockIds, remote_physical_per_logical: int
2109 + ) -> BlockIds:
2110 + """Map logical block IDs to physical kernel block IDs on the remote.
2111 +
2112 + Args:
2113 + block_ids: per-group lists of logical block IDs.
2114 + remote_physical_per_logical: remote engine's physical blocks
2115 + per logical block.
2116 +
2117 + Returns:
2118 + Same structure with FA groups expanded (each logical block L
2119 + becomes kernel blocks [L*remote_physical_per_logical, ..
2120 + L*remote_physical_per_logical +
2121 + remote_physical_per_logical - 1]).
2122 + Mamba groups are passed through unchanged.
2123 + """
2124 + if remote_physical_per_logical == 1:
2125 + return block_ids
2126 + remote_arange = np.arange(remote_physical_per_logical).reshape(1, -1)
2127 + group_specs = self.kv_cache_config.kv_cache_groups
2128 + result = [
2129 + BlockTable.map_to_kernel_blocks(
2130 + np.array(group),
2131 + remote_physical_per_logical,
2132 + remote_arange,
2133 + ).tolist()
2134 + if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
2135 + else group
2136 + for i, group in enumerate(block_ids)
2137 + ]
2138 + return result
2139 +
2140 + def get_backend_aware_kv_block_len(
2141 + self, layer_idx: int, first_split: bool = True, mamba_view: bool = False
2142 + ) -> int:
2143 + """
2144 + Get the block length for one K/V element (K and V have the same size).
2145 +
2146 + For FA and other backends, this is equal to the length of the whole
2147 + block, as K and V are in separate regions.
2148 + For FlashInfer, this is half the length of the whole block, as K and V
2149 + share the same region.
2150 + Similarly, for SSM-based models, state and conv are interleaved, but crucially
2151 + the their size differs.
2152 + Reference diagram:
2153 + KVCacheTensor (Shared)
2154 + / \\
2155 + / \\
2156 + / \\
2157 + Attention (FlashInfer) View Mamba View
2158 + | |
2159 + | |
2160 + +-------------------+ +-------------------+
2161 + | KVCacheTensor | | KVCacheTensor |
2162 + | | | |
2163 + |<----- page ------>| |<----- page ------->|
2164 + | size | | size |
2165 + | Key 0 | Val 0 | |Conv 0 | SSM 0 |
2166 + | Key 1 | Val 1 | |Conv 1 | SSM 1 |
2167 + | ... | ... | | ... | ... |
2168 + | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 |
2169 + | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 |
2170 + +-------------------+ +--------------------+
2171 + |1st_split-2nd_split| |1st_split-2nd_split |
2172 + """
2173 + assert self.transfer_topo is not None
2174 + virtually_split = self.transfer_topo.virtually_split_kv_in_blocks
2175 + if virtually_split and mamba_view:
2176 + block_len = self._mamba_ssm_size[not first_split]
2177 + else:
2178 + half_block = virtually_split and not self._is_region_replicated(layer_idx)
2179 + block_len = self.block_len_per_layer[layer_idx] // (2 if half_block else 1)
2180 + return block_len
2181 +
2182 + def get_kv_connector_stats(self) -> KVConnectorStats | None:
2183 + """
2184 + Get the KV transfer stats for the connector.
2185 + """
2186 + # Clear stats for next iteration
2187 + if not self.xfer_stats.is_empty():
2188 + return self.xfer_stats.clone_and_reset()
2189 + return None
2190 +
2191 + def get_block_ids_with_load_errors(self) -> set[int]:
2192 + """
2193 + Return and clear the set of block IDs that failed to load.
2194 +
2195 + This is called by the scheduler to identify blocks that need
2196 + to be retried after a NIXL transfer failure.
2197 + """
2198 + # Drain the queue (thread-safe, no lock needed).
2199 + result: set[int] = set()
2200 + while not self._invalid_block_ids.empty():
2201 + try:
2202 + result.update(self._invalid_block_ids.get_nowait())
2203 + except queue.Empty:
2204 + break
2205 + return result
2206 +
2207 + def _evict_stale_engines(self) -> None:
2208 + """Scan for and evict remote engines that have exceeded their TTL.
2209 +
2210 + Called from the main thread in when a new remote engine appears.
2211 + We can only go OOM as we discover and register a new remote, therefore we make
2212 + sure we clean up stale engine data structures before then. This invariant
2213 + prevents us from using background threads, though memory usage is not guaranteed
2214 + to be "optimal" until a new handshake is performed.
2215 +
2216 + Engines with active transfers or pending handshakes cannot be stale:
2217 + - Active transfers touch _engine_last_active in start_load_kv.
2218 + - Pending handshakes don't have an _engine_last_active entry yet
2219 + """
2220 + # NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number
2221 + # of remote engines is registered all at once (adding a background cleanup
2222 + # thread wouldnt help either).
2223 + # If that scenario is plausible, we can follow up with an LRU eviction policy.
2224 + if self._engine_ttl <= 0:
2225 + return
2226 +
2227 + now = time.perf_counter()
2228 + for eid, last_active in list(self._engine_last_active.items()):
2229 + if now - last_active > self._engine_ttl:
2230 + self._cleanup_remote_engine(eid)
2231 +
2232 + def _cleanup_remote_engine(
2233 + self, engine_id: EngineId, *, log_eviction: bool = True
2234 + ) -> None:
2235 + """Remove all state for a single remote engine.
2236 +
2237 + Releases NIXL resources (dlist handles, remote agents) and clears
2238 + all per-engine data structures. Used by both TTL eviction and
2239 + shutdown.
2240 + """
2241 + assert engine_id in self._remote_agents
2242 +
2243 + for handle in self.dst_xfer_side_handles.pop(engine_id).values():
2244 + self.nixl_wrapper.release_dlist_handle(handle)
2245 + for agent_name in self._remote_agents.pop(engine_id).values():
2246 + self.nixl_wrapper.remove_remote_agent(agent_name)
2247 +
2248 + del self.kv_caches_base_addr[engine_id]
2249 + del self.dst_num_blocks[engine_id]
2250 + del self.tp_mappings[engine_id]
2251 + if self.transfer_topo is not None:
2252 + self.transfer_topo.unregister_remote_engine(engine_id)
2253 +
2254 + last_active = self._engine_last_active.pop(engine_id)
2255 + if log_eviction:
2256 + logger.info(
2257 + "Evicted stale remote engine %s (inactive for %.1fs).",
2258 + engine_id,
2259 + time.perf_counter() - last_active,
2260 + )
2261 +
2262 + def __del__(self):
2263 + self.shutdown()
2264 +
2265 + def shutdown(self):
2266 + """Shutdown the connector worker."""
2267 + if not hasattr(self, "_handshake_initiation_executor"):
2268 + # error happens during init, no need to shutdown
2269 + return
2270 + self._handshake_initiation_executor.shutdown(wait=False)
2271 + for handles in self._recving_transfers.values():
2272 + for handle in handles:
2273 + self.nixl_wrapper.release_xfer_handle(handle)
2274 + self._recving_transfers.clear()
2275 + for handle in self.src_xfer_handles_by_block_size.values():
2276 + self.nixl_wrapper.release_dlist_handle(handle)
2277 + self.src_xfer_handles_by_block_size.clear()
2278 + for handles in self.src_xfer_handles_by_tp_ratio.values():
2279 + for handle in handles:
2280 + self.nixl_wrapper.release_dlist_handle(handle)
2281 + self.src_xfer_handles_by_tp_ratio.clear()
2282 + for engine_id in list(self._remote_agents):
2283 + self._cleanup_remote_engine(engine_id, log_eviction=False)
2284 + for desc in self._registered_descs:
2285 + self.nixl_wrapper.deregister_memory(desc)
2286 + self._registered_descs.clear()
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+115 -22
@@ -1,6 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""NixlConnector – thin facade that delegates to scheduler / worker."""
3 +"""NIXL connector facades.
4 +
5 +This module hosts the thin facade classes that vLLM's KV-connector layer
6 +instantiates. Almost all the real work lives in the per-mode scheduler
7 +and worker classes; the connector classes here only forward calls.
8 +
9 +* :class:`NixlBaseConnector` – common logic shared by pull and push.
10 +* :class:`NixlPullConnector` – pull-based (READ) KV transfer.
11 +* :class:`NixlPushConnector` – push-based (WRITE) KV transfer.
12 +* ``NixlConnector`` – backward-compatible alias for :class:`NixlPullConnector`.
13 +"""
14
15 from typing import TYPE_CHECKING, Any
16
@@ -28,16 +38,22 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
38 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
39 NixlConnectorMetadata,
40 )
31 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
32 - NixlConnectorScheduler,
41 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
42 + NixlPullConnectorScheduler,
43 +)
44 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import (
45 + NixlPullConnectorWorker,
46 +)
47 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import (
48 + NixlPushConnectorScheduler,
49 +)
50 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
51 + NixlPushConnectorWorker,
52 )
53 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
54 NixlKVConnectorStats,
55 NixlPromMetrics,
56 )
38 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
39 - NixlConnectorWorker,
40 -)
57 from vllm.forward_context import ForwardContext
58 from vllm.logger import init_logger
59 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
@@ -47,6 +63,12 @@ from vllm.v1.kv_cache_interface import MambaSpec
63 from vllm.v1.outputs import KVConnectorOutput
64
65 if TYPE_CHECKING:
66 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (
67 + NixlBaseConnectorScheduler,
68 + )
69 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
70 + NixlBaseConnectorWorker,
71 + )
72 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
73 from vllm.v1.kv_cache_interface import KVCacheConfig
74 from vllm.v1.request import Request
@@ -54,7 +76,9 @@ if TYPE_CHECKING:
76 logger = init_logger(__name__)
77
78
57 -class NixlConnector(KVConnectorBase_V1, SupportsHMA):
79 +class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
80 + """Base connector with common logic shared by pull and push modes."""
81 +
82 @property
83 def prefer_cross_layer_blocks(self) -> bool:
84 if any(
@@ -106,16 +130,9 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
130 self.kv_cache_config = kv_cache_config
131 self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
132 self.kv_transfer_config = vllm_config.kv_transfer_config
109 - if role == KVConnectorRole.SCHEDULER:
110 - self.connector_scheduler: NixlConnectorScheduler | None = (
111 - NixlConnectorScheduler(vllm_config, self.engine_id, kv_cache_config)
112 - )
113 - self.connector_worker: NixlConnectorWorker | None = None
114 - elif role == KVConnectorRole.WORKER:
115 - self.connector_scheduler = None
116 - self.connector_worker = NixlConnectorWorker(
117 - vllm_config, self.engine_id, kv_cache_config
118 - )
133 + # Subclasses must set self.connector_scheduler and self.connector_worker
134 + self.connector_scheduler: NixlBaseConnectorScheduler | None = None
135 + self.connector_worker: NixlBaseConnectorWorker | None = None
136
137 ############################################################
138 # Class Methods
@@ -256,11 +273,6 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
273 vllm_config, metric_types, labelnames, per_engine_labelvalues
274 )
275
259 - def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
260 - assert self.connector_worker is not None
261 - assert isinstance(self._connector_metadata, NixlConnectorMetadata)
262 - self.connector_worker.start_load_kv(self._connector_metadata)
263 -
276 def wait_for_layer_load(self, layer_name: str) -> None:
277 """NixlConnector does not do layerwise saving."""
278 pass
@@ -281,6 +293,11 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
293 if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks:
294 self.connector_worker.save_kv_to_host(self._connector_metadata)
295
296 + def has_pending_push_work(self) -> bool:
297 + if self.connector_scheduler is not None:
298 + return self.connector_scheduler.has_pending_push_work()
299 + return False
300 +
301 def shutdown(self):
302 if self.connector_worker is not None:
303 self.connector_worker.shutdown()
@@ -299,3 +316,79 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
316 """
317 assert self.connector_worker is not None
318 return self.connector_worker.xfer_handshake_metadata
319 +
320 +
321 +class NixlPullConnector(NixlBaseConnector):
322 + """Pull-based (READ) NIXL KV transfer connector."""
323 +
324 + def __init__(
325 + self,
326 + vllm_config: VllmConfig,
327 + role: KVConnectorRole,
328 + kv_cache_config: "KVCacheConfig",
329 + ):
330 + super().__init__(vllm_config, role, kv_cache_config)
331 + if role == KVConnectorRole.SCHEDULER:
332 + self.connector_scheduler = NixlPullConnectorScheduler(
333 + vllm_config, self.engine_id, kv_cache_config
334 + )
335 + self.connector_worker = None
336 + elif role == KVConnectorRole.WORKER:
337 + self.connector_scheduler = None
338 + self.connector_worker = NixlPullConnectorWorker(
339 + vllm_config, self.engine_id, kv_cache_config
340 + )
341 +
342 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
343 + assert self.connector_worker is not None
344 + assert isinstance(self.connector_worker, NixlPullConnectorWorker)
345 + assert isinstance(self._connector_metadata, NixlConnectorMetadata)
346 + self.connector_worker.start_load_kv(self._connector_metadata)
347 +
348 +
349 +class NixlPushConnector(NixlBaseConnector):
350 + """Push-based (WRITE) NIXL KV transfer connector."""
351 +
352 + def __init__(
353 + self,
354 + vllm_config: VllmConfig,
355 + role: KVConnectorRole,
356 + kv_cache_config: "KVCacheConfig",
357 + ):
358 + super().__init__(vllm_config, role, kv_cache_config)
359 + self.connector_scheduler: NixlPushConnectorScheduler | None = None
360 + self.connector_worker: NixlPushConnectorWorker | None = None
361 + if role == KVConnectorRole.SCHEDULER:
362 + self.connector_scheduler = NixlPushConnectorScheduler(
363 + vllm_config, self.engine_id, kv_cache_config
364 + )
365 + elif role == KVConnectorRole.WORKER:
366 + self.connector_worker = NixlPushConnectorWorker(
367 + vllm_config, self.engine_id, kv_cache_config
368 + )
369 + else:
370 + raise ValueError(f"Unsupported KVConnectorRole: {role}")
371 +
372 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
373 + """Drive push processing on the worker.
374 +
375 + The worker enqueues registrations / finished blocks for the
376 + background ``nixl-push-writer`` thread; the writer issues the
377 + WRITE transfers and polls NIXL notifs without further
378 + engine-thread involvement.
379 + """
380 + assert self.connector_worker is not None
381 + assert isinstance(self._connector_metadata, NixlConnectorMetadata)
382 + self.connector_worker.start_load_kv(self._connector_metadata)
383 +
384 +
385 +# Backward compatibility: NixlConnector is the pull-based connector.
386 +NixlConnector = NixlPullConnector
387 +
388 +
389 +__all__ = [
390 + "NixlBaseConnector",
391 + "NixlConnector",
392 + "NixlPullConnector",
393 + "NixlPushConnector",
394 +]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+14
@@ -19,6 +19,11 @@ TransferHandle = int
19 ReqId = str
20
21 GET_META_MSG = b"get_meta_msg"
22 +
23 +# Push-mode (WRITE-based) registration notification.
24 +# Sent worker-to-worker over NIXL: D worker -> P worker, encoded as
25 +# PUSH_REG_NOTIF_PREFIX + msgpack(registration_data).
26 +PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
27 #
28 # NIXL Connector Version
29 #
@@ -160,6 +165,8 @@ class ReqMeta:
165 local_physical_block_ids: BlockIds
166 tp_size: int
167 remote: RemoteMeta | None = None
168 + # Remote block size, discovered during NIXL handshake (push mode).
169 + remote_block_size: int | None = None
170
171
172 class NixlConnectorMetadata(KVConnectorMetadata):
@@ -171,6 +178,12 @@ class NixlConnectorMetadata(KVConnectorMetadata):
178 self.reqs_not_processed: set[ReqId] = set()
179 # Heartbeat data grouped by remote engine, sent by D worker to P.
180 self.heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {}
181 + # Push mode (D side): registration data the D worker should send to
182 + # P workers via NIXL notification on this step.
183 + self.push_registrations: dict[ReqId, dict[str, Any]] = {}
184 + # Push mode (P side): newly finished request blocks to be matched
185 + # against pending D registrations on the P worker.
186 + self.push_finished_blocks: dict[ReqId, BlockIds] = {}
187
188 def _add_new_req(
189 self,
@@ -182,6 +195,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):
195 local_physical_block_ids=local_block_ids,
196 # P workers don't need to receive tp_size from proxy here.
197 tp_size=kv_transfer_params.get("tp_size", 1),
198 + remote_block_size=kv_transfer_params.get("remote_block_size"),
199 )
200
201 def add_new_req_to_save(
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py new
+275
@@ -0,0 +1,275 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Pull-specific scheduler-side logic for the NIXL connector."""
4 +
5 +import time
6 +from typing import TYPE_CHECKING, Any
7 +
8 +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
9 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (
10 + NixlBaseConnectorScheduler,
11 +)
12 +from vllm.logger import init_logger
13 +
14 +if TYPE_CHECKING:
15 + from vllm.config import VllmConfig
16 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
17 + from vllm.v1.kv_cache_interface import KVCacheConfig
18 + from vllm.v1.request import Request
19 +
20 +logger = init_logger(__name__)
21 +
22 +
23 +class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
24 + """Pull-specific scheduler logic (READ-based KV transfer)."""
25 +
26 + def __init__(
27 + self,
28 + vllm_config: "VllmConfig",
29 + engine_id: str,
30 + kv_cache_config: "KVCacheConfig",
31 + ):
32 + super().__init__(vllm_config, engine_id, kv_cache_config)
33 +
34 + def get_num_new_matched_tokens(
35 + self, request: "Request", num_computed_tokens: int
36 + ) -> tuple[int, bool]:
37 + """
38 + For remote prefill, pull all prompt blocks from remote
39 + asynchronously relative to engine execution.
40 +
41 + Args:
42 + request (Request): the request object.
43 + num_computed_tokens (int): the number of locally
44 + computed tokens for this request
45 + Returns:
46 + * the number of tokens that can be loaded from the
47 + external KV cache beyond what is already computed.
48 + * true if the external KV cache tokens will be loaded
49 + asynchronously (between scheduler steps).
50 + """
51 +
52 + params = request.kv_transfer_params
53 + logger.debug(
54 + "NIXLConnector get_num_new_matched_tokens: "
55 + "num_computed_tokens=%s, kv_transfer_params=%s",
56 + num_computed_tokens,
57 + params,
58 + )
59 +
60 + if params is not None and params.get("do_remote_prefill"):
61 + # Remote prefill: get all prompt blocks from remote.
62 + token_ids = request.prompt_token_ids or []
63 + actual = self._mamba_prefill_token_count(len(token_ids))
64 + count = actual - num_computed_tokens
65 + if count > 0:
66 + return count, True
67 +
68 + if params is not None and params.get("do_remote_decode") and self._has_mamba:
69 + self._truncate_mamba_request_for_prefill(request)
70 +
71 + if (
72 + params is not None
73 + and params.get("do_remote_decode")
74 + and params.get("remote_block_ids")
75 + and all(
76 + p in params
77 + for p in (
78 + "remote_engine_id",
79 + "remote_request_id",
80 + "remote_host",
81 + "remote_port",
82 + )
83 + )
84 + ):
85 + # Decode node has kv blocks for part of prefill request, so, provide them
86 + # as an external token count to scheduler.
87 + # The tokens will be loaded if not already present
88 + # in the prefill node local cache
89 + remote_num_tokens = params.get("remote_num_tokens") or 0
90 + count = (
91 + min(remote_num_tokens, request.num_prompt_tokens) - num_computed_tokens
92 + )
93 + if count > 0:
94 + # Check kv_recompute_threshold: skip pull if
95 + # remote tokens are below the threshold.
96 + if (
97 + self.kv_recompute_threshold > 0
98 + and count < self.kv_recompute_threshold
99 + ):
100 + logger.debug(
101 + "Skipping remote pull for %s: %d remote tokens < threshold %d",
102 + request.request_id,
103 + count,
104 + self.kv_recompute_threshold,
105 + )
106 + return 0, False
107 + return count, True
108 +
109 + # No remote prefill for this request.
110 + return 0, False
111 +
112 + def update_state_after_alloc(
113 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
114 + ):
115 + params = request.kv_transfer_params
116 + logger.debug(
117 + "NIXLConnector update_state_after_alloc: "
118 + "num_external_tokens=%s, kv_transfer_params=%s",
119 + num_external_tokens,
120 + params,
121 + )
122 +
123 + if not params:
124 + return
125 +
126 + if params.get("do_remote_decode") or (
127 + params.get("do_remote_prefill") and self.is_bidirectional_kv_xfer_enabled
128 + ):
129 + self._reqs_in_batch.add(request.request_id)
130 + if self.use_host_buffer and params.get("do_remote_decode"):
131 + # NOTE: when accelerator is not directly supported by Nixl,
132 + # prefilled blocks need to be saved to host memory before transfer.
133 + self._reqs_need_save[request.request_id] = request
134 + elif params.get("do_remote_prefill") or (
135 + params.get("do_remote_decode")
136 + and self.is_bidirectional_kv_xfer_enabled
137 + and not params.get("_remote_blocks_processed")
138 + ):
139 + if params.get("remote_block_ids"):
140 + if all(
141 + p in params
142 + for p in (
143 + "remote_engine_id",
144 + "remote_request_id",
145 + "remote_host",
146 + "remote_port",
147 + )
148 + ):
149 + # If remote_blocks and num_external_tokens = 0, we have
150 + # a full prefix cache hit on the local node. We need to call
151 + # send_notif in _read_blocks to free the memory on the remote node.
152 +
153 + unhashed_local_block_ids: BlockIds = (
154 + blocks.get_unhashed_block_ids_all_groups()
155 + if num_external_tokens > 0
156 + else ()
157 + )
158 + local_block_ids = self.get_sw_clipped_blocks(
159 + unhashed_local_block_ids
160 + )
161 +
162 + # Get unhashed blocks to pull from remote. Mind that a full prefix
163 + # cache hit is indicated with an empty list.
164 + self._reqs_need_recv[request.request_id] = (
165 + request,
166 + local_block_ids,
167 + )
168 +
169 + else:
170 + logger.warning(
171 + "Got invalid KVTransferParams: %s. This "
172 + "request will not utilize KVTransfer",
173 + params,
174 + )
175 + else:
176 + assert num_external_tokens == 0
177 + # Only trigger 1 KV transfer per request.
178 + params["do_remote_prefill"] = False
179 + params["_remote_blocks_processed"] = True
180 +
181 + def request_finished(
182 + self,
183 + request: "Request",
184 + block_ids: "BlockIds",
185 + ) -> tuple[bool, dict[str, Any] | None]:
186 + """
187 + Once a request is finished, determine whether request blocks
188 + should be freed now or will be sent asynchronously and freed later.
189 + """
190 + from vllm.v1.request import RequestStatus
191 +
192 + params = request.kv_transfer_params
193 + logger.debug(
194 + "NIXLConnector request_finished(%s), request_status=%s, "
195 + "kv_transfer_params=%s",
196 + request.request_id,
197 + request.status,
198 + params,
199 + )
200 + if not params:
201 + return False, None
202 +
203 + is_p_node = bool(params.get("do_remote_decode"))
204 + is_d_node = not is_p_node
205 +
206 + # Stop heartbeating for aborted requests that never reached finished_recving:
207 + # normal path cleans up in update_connector_output.
208 + self._stop_heartbeat(request.request_id)
209 +
210 + if params.get("do_remote_prefill"):
211 + # If do_remote_prefill is still True when the request is finished,
212 + # update_state_after_alloc must not have been called (the request
213 + # must have been aborted before it was scheduled, e.g. via the
214 + # abort_immediately path used to clean up KV-transfer requests
215 + # rejected at the D-side serving layer).
216 + # To avoid stranding the prefill blocks in the prefill instance,
217 + # we must add empty block_ids to _reqs_need_recv so that our
218 + # worker side will notify and free blocks in the prefill instance.
219 + self._reqs_need_recv[request.request_id] = (request, [])
220 + params["do_remote_prefill"] = False
221 + return False, None
222 +
223 + if is_d_node and not self.is_bidirectional_kv_xfer_enabled:
224 + return False, None
225 +
226 + if request.status not in (
227 + RequestStatus.FINISHED_LENGTH_CAPPED,
228 + RequestStatus.FINISHED_STOPPED,
229 + ):
230 + # Also include the case of a P/D Prefill request with immediate
231 + # block free (eg abort). Stop tracking this request.
232 + self._reqs_not_processed.add(request.request_id)
233 + # Clear _reqs_need_save if a request is aborted as partial prefill.
234 + self._reqs_need_save.pop(request.request_id, None)
235 + return False, None
236 +
237 + # TODO: check whether block_ids actually ever be 0. If not we could
238 + # remove the conditional below
239 + delay_free_blocks = any(len(group) > 0 for group in block_ids)
240 + remote_num_tokens = 0
241 + if delay_free_blocks:
242 + # Prefill request on remote. It will be read from D upon completion
243 + request_kv_blocks_ttl = self._kv_lease_duration
244 + if is_d_node:
245 + # For blocks pinned on D, use a simpler timeout for now instead of a
246 + # lease mechanism as turn2 request is client-driven.
247 + request_kv_blocks_ttl = self.decoder_kv_blocks_ttl
248 + logger.debug(
249 + "NIXLConnector request_finished(%s) waiting for %d seconds "
250 + "before releasing blocks",
251 + request.request_id,
252 + request_kv_blocks_ttl,
253 + )
254 + self._reqs_need_send[request.request_id] = (
255 + time.perf_counter() + request_kv_blocks_ttl
256 + )
257 + # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
258 + # trimming down after allocating for the whole sequence length. Empty
259 + # blocks are always at the start of the list.
260 + # Here we "unpad" blocks to send the actual remote blocks to be read.
261 + block_ids = self.get_sw_clipped_blocks(block_ids)
262 +
263 + remote_num_tokens = request.num_computed_tokens
264 +
265 + return delay_free_blocks, dict(
266 + do_remote_prefill=is_p_node,
267 + do_remote_decode=is_d_node,
268 + remote_block_ids=block_ids,
269 + remote_engine_id=self.engine_id,
270 + remote_request_id=request.request_id,
271 + remote_host=self.side_channel_host,
272 + remote_port=self.side_channel_port,
273 + tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
274 + remote_num_tokens=remote_num_tokens,
275 + )
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py new
+382
@@ -0,0 +1,382 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Pull-specific (READ) worker-side logic for the NIXL connector."""
4 +
5 +import time
6 +from typing import TYPE_CHECKING
7 +
8 +import numpy as np
9 +
10 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
11 + NixlBaseConnectorWorker,
12 +)
13 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
14 + NixlConnectorMetadata,
15 + ReqMeta,
16 +)
17 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
18 + ReadSpec,
19 +)
20 +from vllm.logger import init_logger
21 +
22 +if TYPE_CHECKING:
23 + from vllm.config import VllmConfig
24 + from vllm.v1.kv_cache_interface import KVCacheConfig
25 +
26 +logger = init_logger(__name__)
27 +
28 +
29 +class NixlPullConnectorWorker(NixlBaseConnectorWorker):
30 + """Pull-specific (READ) worker logic."""
31 +
32 + def __init__(
33 + self,
34 + vllm_config: "VllmConfig",
35 + engine_id: str,
36 + kv_cache_config: "KVCacheConfig",
37 + ):
38 + super().__init__(vllm_config, engine_id, kv_cache_config)
39 +
40 + def start_load_kv(self, metadata: NixlConnectorMetadata):
41 + """
42 + Start loading by triggering non-blocking nixl_xfer.
43 + We check for these trnxs to complete in each step().
44 + """
45 + for req_id, meta in metadata.reqs_to_recv.items():
46 + meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
47 + meta.local_block_ids
48 + )
49 + assert meta.remote is not None
50 + # Remote block IDs are kept logical here; expanded in
51 + # _read_blocks_for_req using the remote engine's phys ratio.
52 + remote_engine_id = meta.remote.engine_id
53 + logger.debug(
54 + "start_load_kv for request %s from remote engine %s. "
55 + "Num local_block_ids: %s. Num remote_block_ids: %s. ",
56 + req_id,
57 + remote_engine_id,
58 + len(meta.local_physical_block_ids),
59 + len(meta.remote.block_ids),
60 + )
61 + # always store metadata for failure recovery
62 + self._recving_metadata[req_id] = meta
63 + if remote_engine_id not in self._remote_agents:
64 + # Initiate handshake with remote engine to exchange metadata.
65 + with self._handshake_lock:
66 + if remote_engine_id not in self._remote_agents:
67 + self._background_nixl_handshake(req_id, remote_engine_id, meta)
68 + continue
69 +
70 + # Handshake already completed, start async read xfer.
71 + self._read_blocks_for_req(req_id, meta)
72 +
73 + # Start transfers for requests whose handshakes have now finished.
74 + while not self._ready_requests.empty():
75 + self._read_blocks_for_req(*self._ready_requests.get_nowait())
76 +
77 + # Keep around the requests that have been part of a batch. This is
78 + # needed because async scheduling pushes the misalignment between the
79 + # moment in which requests expiration is set (P side) and the moment in
80 + # which blocks are read from D. As P can now more easily lag behind D
81 + # while processing the next batch, we make sure to only set an
82 + # expiration for requests that have not been read from D yet.
83 + for req_id in metadata.reqs_in_batch:
84 + self._reqs_to_process.add(req_id)
85 +
86 + # Remove all requests that are not to be processed (eg aborted).
87 + for req_id in metadata.reqs_not_processed:
88 + self._reqs_to_process.discard(req_id)
89 + # We should never get an abort after setting an expiry timer
90 + assert req_id not in self._reqs_to_send
91 +
92 + # Add to requests that are waiting to be read and track expiration.
93 + for req_id, expiration_time in metadata.reqs_to_send.items():
94 + if req_id in self._reqs_to_process:
95 + self._reqs_to_send[req_id] = expiration_time
96 +
97 + # Send heartbeats to P-side engines to keep KV blocks alive while
98 + # requests sit in the D scheduler WAITING queue.
99 + self._send_heartbeats(metadata)
100 +
101 + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
102 + assert meta.remote is not None and self.transfer_topo is not None
103 + engine_id = meta.remote.engine_id
104 + # Update last activity from this remote. Mind that cleanup is done on main
105 + # thread (this one), so we don't race on this structure.
106 + self._engine_last_active[engine_id] = time.perf_counter()
107 + plan = self.tp_mappings[engine_id]
108 + remote_info = self.transfer_topo.get_engine_info(engine_id)
109 + tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
110 +
111 + meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
112 + meta.remote.block_ids,
113 + remote_info.remote_physical_blocks_per_logical,
114 + )
115 + remote_block_ids = meta.remote.block_ids
116 + local_block_ids = meta.local_physical_block_ids
117 + num_groups = len(local_block_ids)
118 + read_specs = [
119 + ReadSpec(
120 + remote_rank=rank,
121 + local_block_ids=[
122 + list(local_block_ids[g])
123 + if rank in plan.source_ranks_per_group[g]
124 + else []
125 + for g in range(num_groups)
126 + ],
127 + remote_block_ids=[
128 + list(remote_block_ids[g])
129 + if rank in plan.source_ranks_per_group[g]
130 + else []
131 + for g in range(num_groups)
132 + ],
133 + )
134 + for rank in plan.all_source_ranks
135 + ]
136 +
137 + # D may have to perform multiple reads from different remote ranks.
138 + # MLA opt: when P TP > D TP, only a single read is executed for
139 + # the first remote rank (cache is duplicated)..
140 + if self.use_mla and tp_ratio < 0:
141 + assert len(read_specs) == 1
142 +
143 + for i, spec in enumerate(read_specs):
144 + remote_block_size = remote_info.remote_block_size
145 + logger.debug(
146 + "Remote agent %s available, calling _read_blocks"
147 + " on remote rank %s with remote block size %s for req %s",
148 + meta.remote.engine_id,
149 + spec.remote_rank,
150 + remote_block_size,
151 + req_id,
152 + )
153 + # Get side handles.
154 + if tp_ratio < 0 and not self.use_mla:
155 + assert remote_block_size == self.block_size
156 + # Remote tp_size > local tp_size: we must perform multiple
157 + # reads. Get the memory chunk onto which we will write to.
158 + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
159 + else:
160 + # Single read from remote, we write to the whole memory region.
161 + # Also handle remote block size different from local block size.
162 + local_xfer_side_handle = self.src_xfer_handles_by_block_size[
163 + remote_block_size
164 + ]
165 +
166 + # Destination handle: remote_engine_id -> remote_rank -> handle.
167 + remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
168 + spec.remote_rank
169 + ]
170 +
171 + self._read_blocks(
172 + read_spec=spec,
173 + request_id=req_id,
174 + dst_engine_id=meta.remote.engine_id,
175 + remote_request_id=meta.remote.request_id,
176 + local_xfer_side_handle=local_xfer_side_handle,
177 + remote_xfer_side_handle=remote_xfer_side_handle,
178 + )
179 +
180 + if self.use_mla and tp_ratio < 0 and read_specs:
181 + # ..but we still need to notify the other remote ranks that we
182 + # have the blocks we need so they can update the request state.
183 + notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()
184 + remote_agents = self._remote_agents[meta.remote.engine_id]
185 + for rank_to_notify, agent in remote_agents.items():
186 + if rank_to_notify != read_specs[0].remote_rank:
187 + self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
188 +
189 + def _read_blocks(
190 + self,
191 + read_spec: ReadSpec,
192 + dst_engine_id: str,
193 + request_id: str,
194 + remote_request_id: str,
195 + local_xfer_side_handle: int,
196 + remote_xfer_side_handle: int,
197 + ):
198 + """
199 + Post a READ point-to-point xfer request from a single local worker to
200 + a single remote worker.
201 + """
202 + assert self.transfer_topo is not None
203 + remote_rank = read_spec.remote_rank
204 + local_block_ids = read_spec.local_block_ids
205 + remote_block_ids = read_spec.remote_block_ids
206 +
207 + remote_info = self.transfer_topo.get_engine_info(dst_engine_id)
208 + block_size_ratio = self.transfer_topo.block_size_ratio(
209 + remote_info.remote_block_size
210 + )
211 + if block_size_ratio > 1:
212 + # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.
213 + assert not self._is_hma_required
214 + local_block_ids0 = local_block_ids[0] if local_block_ids else []
215 + remote_block_ids0 = remote_block_ids[0]
216 + local_block_ids_mapped = self.get_mapped_blocks(
217 + np.asarray(local_block_ids0), block_size_ratio
218 + ).tolist()
219 + if len(local_block_ids_mapped) > len(remote_block_ids0):
220 + # NOTE:
221 + # get_mapped_blocks will always expand block_ids for n times.
222 + # ex:
223 + # prefill block_ids with block_size as 4:
224 + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
225 + # Local decode block_ids with block_size as 16: [1, 2, 3]
226 + # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to
227 + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
228 + # Then we clip local to align with prefill
229 + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to
230 + # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
231 + local_block_ids_mapped = local_block_ids_mapped[
232 + : len(remote_block_ids0)
233 + ]
234 + local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
235 + remote_block_ids = [remote_block_ids0]
236 + # NOTE(rob): having the staging blocks be on the READER side is
237 + # not going to work well (since we will have to call rearrange tensors).
238 + # after we detect the txn is complete (which means we cannot make the
239 + # read trxn async easily). If we want to make "READ" happen cleanly,
240 + # then we will need to have the staging blocks on the remote side.
241 +
242 + # NOTE(rob): according to nvidia the staging blocks are used to
243 + # saturate IB with heterogeneous TP sizes.
244 +
245 + # Number of D TP workers that will read from dst P. Propagate info
246 + # on notification so that dst worker can wait before freeing blocks.
247 + notif_id = f"{remote_request_id}:{self.world_size}".encode()
248 +
249 + # Full prefix cache hit: do not need to read remote blocks,
250 + # just notify P worker that we have the blocks we need.
251 + if len(local_block_ids) == 0:
252 + # A full prefix cache hit is indicated with an empty list.
253 + agent_name = self._remote_agents[dst_engine_id][remote_rank]
254 + try:
255 + self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
256 + except Exception as e:
257 + self._log_failure(
258 + failure_type="notification_failed",
259 + msg="P worker blocks will be freed after timeout. "
260 + "This may indicate network issues.",
261 + req_id=request_id,
262 + error=e,
263 + dst_engine_id=dst_engine_id,
264 + remote_rank=remote_rank,
265 + remote_agent_name=agent_name,
266 + )
267 + self.xfer_stats.record_failed_notification()
268 + return
269 +
270 + assert (
271 + len(remote_block_ids)
272 + == len(local_block_ids)
273 + == len(self.kv_cache_config.kv_cache_groups)
274 + )
275 + remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical
276 + local_block_ids, remote_block_ids = self._apply_prefix_caching(
277 + local_block_ids, remote_block_ids, remote_physical_per_logical
278 + )
279 +
280 + # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
281 + # corresponding rank. With heterogeneous TP, fixing D>P, the D tp
282 + # workers will issue xfers to parts of the P worker remote kv caches.
283 +
284 + # Get descs ids.
285 + remote_block_descs_ids = self._compute_desc_ids(
286 + block_ids=remote_block_ids,
287 + dst_num_blocks=self.dst_num_blocks[dst_engine_id],
288 + block_size_ratio=None,
289 + physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
290 + )
291 + local_block_descs_ids = self._compute_desc_ids(
292 + block_ids=local_block_ids,
293 + dst_num_blocks=self.dst_num_blocks[self.engine_id],
294 + block_size_ratio=block_size_ratio,
295 + physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
296 + )
297 +
298 + assert len(local_block_descs_ids) == len(remote_block_descs_ids)
299 +
300 + # Prepare transfer with Nixl.
301 + handle = None
302 + try:
303 + handle = self.nixl_wrapper.make_prepped_xfer(
304 + "READ",
305 + local_xfer_side_handle,
306 + local_block_descs_ids,
307 + remote_xfer_side_handle,
308 + remote_block_descs_ids,
309 + notif_msg=notif_id,
310 + )
311 +
312 + # Begin async xfer.
313 + self.nixl_wrapper.transfer(handle)
314 +
315 + # Use handle to check completion in future step().
316 + self._recving_transfers[request_id].append(handle)
317 + except Exception as e:
318 + # mark all (logical) blocks for this request as invalid
319 + self._log_failure(
320 + failure_type="transfer_setup_failed",
321 + req_id=request_id,
322 + msg="Marking blocks as invalid",
323 + error=e,
324 + dst_engine_id=dst_engine_id,
325 + remote_rank=remote_rank,
326 + )
327 + self._handle_failed_transfer(request_id, handle)
328 +
329 + def _get_new_notifs(self) -> set[str]:
330 + """
331 + Get req_ids which got a remote xfer message. When multiple consumers
332 + are reading from the same producer (heterogeneous TP scenario), wait
333 + for all consumers to be done pulling.
334 +
335 + Also handles heartbeat notifications ("HB:req1,req2,...") by
336 + extending the lease on the referenced requests.
337 + """
338 + assert self.transfer_topo is not None
339 + notified_req_ids: set[str] = set()
340 + for notifs in self.nixl_wrapper.get_new_notifs().values():
341 + for notif in notifs:
342 + msg = notif.decode("utf-8")
343 +
344 + # Handle heartbeat messages from D-side.
345 + if msg.startswith("HB:"):
346 + self._handle_heartbeat(msg[3:])
347 + continue
348 +
349 + req_id, tp_size = msg.rsplit(":", 1)
350 + if (
351 + req_id not in self._reqs_to_send
352 + and req_id not in self._reqs_to_process
353 + ):
354 + logger.error(
355 + "Potentially invalid KV blocks for "
356 + "unrecognized request %s were retrieved by "
357 + "a decode worker. They may have expired.",
358 + req_id,
359 + )
360 + continue
361 +
362 + # NOTE: `tp_ratio` is the opposite when swapping local<>remote
363 + n_consumers = int(tp_size)
364 + tp_ratio = self.transfer_topo.tp_ratio(n_consumers)
365 +
366 + # Number of reads *per producer* to wait for.
367 + # When remote D TP > local P TP we expect `tp_ratio` reads.
368 + consumers_per_producer = (
369 + -tp_ratio if n_consumers > self.world_size else 1
370 + )
371 +
372 + self.consumer_notification_counts_by_req[req_id] += 1
373 + # Wait all consumers (D) to be done reading before freeing.
374 + if (
375 + self.consumer_notification_counts_by_req[req_id]
376 + == consumers_per_producer
377 + ):
378 + notified_req_ids.add(req_id)
379 + del self.consumer_notification_counts_by_req[req_id]
380 + self._reqs_to_process.remove(req_id)
381 + self._reqs_to_send.pop(req_id, None)
382 + return notified_req_ids
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py new
+348
@@ -0,0 +1,348 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Push-specific scheduler-side logic for the NIXL connector.
4 +
5 +In push mode, scheduler-side responsibilities are:
6 +
7 +* D side (decode): on ``update_state_after_alloc``, stash registration data
8 + (D's identity + locally allocated block IDs) into
9 + ``_push_pending_registrations``. The D worker drains it from
10 + ``meta.push_registrations`` next step and sends a NIXL notification to the
11 + P worker (no scheduler-level networking).
12 +* P side (prefill): on ``request_finished``, stash the finished block IDs
13 + into ``_finished_request_blocks`` for the lease, and into
14 + ``_newly_finished_push_blocks`` so the P worker picks them up via
15 + ``meta.push_finished_blocks`` and matches against any D registrations
16 + it already received via NIXL notifications.
17 +* Both sides: ``has_pending_push_work`` keeps the engine main loop stepping
18 + while pushes are in flight. ``update_connector_output`` cleans up
19 + ``_finished_request_blocks`` once the WRITE completes.
20 +
21 +A soft per-registration watchdog on the D scheduler fails requests that have
22 +been registered but not fulfilled within a configurable timeout.
23 +"""
24 +
25 +from __future__ import annotations
26 +
27 +import time
28 +from typing import TYPE_CHECKING, Any
29 +
30 +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
31 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
32 + KVConnectorMetadata,
33 +)
34 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (
35 + NixlBaseConnectorScheduler,
36 +)
37 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
38 + NixlConnectorMetadata,
39 + ReqId,
40 +)
41 +from vllm.logger import init_logger
42 +
43 +if TYPE_CHECKING:
44 + from vllm.config import VllmConfig
45 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
46 + from vllm.v1.core.sched.output import SchedulerOutput
47 + from vllm.v1.kv_cache_interface import KVCacheConfig
48 + from vllm.v1.outputs import KVConnectorOutput
49 + from vllm.v1.request import Request
50 +
51 +logger = init_logger(__name__)
52 +
53 +
54 +class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
55 + """Push-specific scheduler logic (WRITE-based KV transfer).
56 +
57 + All P2P communication is deferred to the worker level via NIXL
58 + notifications. The scheduler communicates with workers only through
59 + the standard ``build_connector_meta`` / ``update_connector_output``
60 + hooks.
61 + """
62 +
63 + def __init__(
64 + self,
65 + vllm_config: VllmConfig,
66 + engine_id: str,
67 + kv_cache_config: KVCacheConfig,
68 + ):
69 + super().__init__(vllm_config, engine_id, kv_cache_config)
70 +
71 + # D-side: registration data to pass to D workers via metadata on
72 + # the next ``build_connector_meta`` call.
73 + self._push_pending_registrations: dict[ReqId, dict[str, Any]] = {}
74 +
75 + # D-side: track the wall-clock deadline for each registered request
76 + # to detect "registered but never fulfilled" failures (e.g. the P
77 + # node disappeared after registration). Keyed by D request_id.
78 + self._push_registration_deadlines: dict[ReqId, float] = {}
79 +
80 + # P-side: block IDs for finished requests, kept for the lease and
81 + # used to drive ``has_pending_push_work``.
82 + self._finished_request_blocks: dict[ReqId, BlockIds] = {}
83 + # P-side: newly finished blocks to ship to P workers on next step.
84 + self._newly_finished_push_blocks: dict[ReqId, BlockIds] = {}
85 +
86 + # Soft watchdog timeout (seconds) for D-side registrations that
87 + # never receive a push completion. Defaults to the existing
88 + # decoder KV blocks TTL so behaviour matches the lease.
89 + assert vllm_config.kv_transfer_config is not None
90 + self._push_registration_timeout: float = float(
91 + vllm_config.kv_transfer_config.get_from_extra_config(
92 + "push_registration_timeout",
93 + self.decoder_kv_blocks_ttl,
94 + )
95 + )
96 +
97 + def get_num_new_matched_tokens(
98 + self, request: Request, num_computed_tokens: int
99 + ) -> tuple[int, bool]:
100 + """In push mode, D doesn't pull — it registers blocks and waits.
101 +
102 + However, we still need to handle the do_remote_prefill case where D
103 + needs to know how many tokens will be pushed.
104 + """
105 + params = request.kv_transfer_params
106 + logger.debug(
107 + "NixlPushConnector get_num_new_matched_tokens: "
108 + "num_computed_tokens=%s, kv_transfer_params=%s",
109 + num_computed_tokens,
110 + params,
111 + )
112 +
113 + if params is not None and params.get("do_remote_prefill"):
114 + token_ids = request.prompt_token_ids or []
115 + actual = self._mamba_prefill_token_count(len(token_ids))
116 + count = actual - num_computed_tokens
117 + if count > 0:
118 + return count, True
119 +
120 + if params is not None and params.get("do_remote_decode") and self._has_mamba:
121 + self._truncate_mamba_request_for_prefill(request)
122 +
123 + return 0, False
124 +
125 + def update_state_after_alloc(
126 + self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
127 + ):
128 + """In push mode, D stores registration data for the worker to send
129 + to P via NIXL notification (deferred to ``build_connector_meta``).
130 + """
131 + params = request.kv_transfer_params
132 + logger.debug(
133 + "NixlPushConnector update_state_after_alloc: "
134 + "num_external_tokens=%s, kv_transfer_params=%s",
135 + num_external_tokens,
136 + params,
137 + )
138 +
139 + if not params:
140 + return
141 +
142 + # P side: track the request as in-batch so the lease accounting
143 + # matches what the worker expects on the next step.
144 + if params.get("do_remote_decode"):
145 + self._reqs_in_batch.add(request.request_id)
146 +
147 + # P side with host-buffer offload: defer save to the worker.
148 + if self.use_host_buffer and params.get("do_remote_decode"):
149 + self._reqs_need_save[request.request_id] = request
150 + return
151 +
152 + # D side: only act on the first call (``do_remote_prefill`` is
153 + # unset on re-entry by the marker below).
154 + if not params.get("do_remote_prefill"):
155 + return
156 +
157 + if num_external_tokens <= 0:
158 + # Nothing to receive: full prefix-cache hit on D, no
159 + # registration to stage.
160 + return
161 +
162 + # First-pass D path: stash registration data the worker will
163 + # ship to P on the next ``build_connector_meta`` cycle.
164 + logger.debug(
165 + "KV PUSH mode: D node storing registration for request %s",
166 + request.request_id,
167 + )
168 + local_block_ids: BlockIds = blocks.get_unhashed_block_ids_all_groups()
169 + local_block_ids = self.get_sw_clipped_blocks(local_block_ids)
170 +
171 + # ``remote_*`` fields are P's coordinates (from D's perspective).
172 + # ``decode_*`` fields are D's own info that P needs for the
173 + # reverse handshake before WRITE-ing.
174 + self._push_pending_registrations[request.request_id] = {
175 + "request_id": request.request_id,
176 + "decode_engine_id": self.engine_id,
177 + "decode_host": self.side_channel_host,
178 + "decode_port": self.side_channel_port,
179 + "decode_tp_size": (self.vllm_config.parallel_config.tensor_parallel_size),
180 + "local_block_ids": local_block_ids,
181 + "remote_engine_id": params["remote_engine_id"],
182 + "remote_host": params["remote_host"],
183 + "remote_port": params["remote_port"],
184 + "remote_tp_size": params["tp_size"],
185 + }
186 + self._push_registration_deadlines[request.request_id] = (
187 + time.perf_counter() + self._push_registration_timeout
188 + )
189 + # In push mode D doesn't know P's blocks; P determines them
190 + # from the registration. We still track the request as
191 + # needing recv so the engine waits for P's WRITE completion.
192 + # ``remote_block_ids`` is also seeded to an empty tuple so the
193 + # base scheduler's ``add_new_req_to_recv`` can build the
194 + # ReqMeta without a KeyError — the actual remote block IDs are
195 + # learned by P over the NIXL handshake at WRITE time.
196 + params["remote_block_ids"] = ()
197 + self._reqs_need_recv[request.request_id] = (request, local_block_ids)
198 +
199 + # Mark as processed so a re-entry (e.g. preemption + reschedule)
200 + # doesn't re-stage the registration.
201 + params["do_remote_prefill"] = False
202 +
203 + def request_finished(
204 + self,
205 + request: Request,
206 + block_ids: BlockIds,
207 + ) -> tuple[bool, dict[str, Any] | None]:
208 + """Push-mode request_finished: stores blocks for workers."""
209 + from vllm.v1.request import RequestStatus
210 +
211 + params = request.kv_transfer_params
212 + logger.debug(
213 + "NixlPushConnector request_finished(%s), request_status=%s, "
214 + "kv_transfer_params=%s",
215 + request.request_id,
216 + request.status,
217 + params,
218 + )
219 + if not params:
220 + return False, None
221 +
222 + is_p_node = bool(params.get("do_remote_decode"))
223 +
224 + self._stop_heartbeat(request.request_id)
225 + # Drop any pending registration deadline; the request either
226 + # completed or was cancelled.
227 + self._push_registration_deadlines.pop(request.request_id, None)
228 +
229 + if params.get("do_remote_prefill"):
230 + # ``do_remote_prefill`` is still set, which means
231 + # ``update_state_after_alloc`` never ran (it would have
232 + # flipped this flag to False). The request was aborted
233 + # before it could be scheduled — e.g. rejected at the D
234 + # serving layer via abort_immediately. To keep P from
235 + # stranding the prefill blocks, we still register an empty
236 + # recv so the worker emits a notif that lets P free them.
237 + self._reqs_need_recv[request.request_id] = (request, [])
238 + params["do_remote_prefill"] = False
239 + return False, None
240 +
241 + # Push connector only acts on the P-side terminal path; D-side
242 + # finishing without a remote prefill is a no-op.
243 + if not is_p_node:
244 + return False, None
245 +
246 + if request.status not in (
247 + RequestStatus.FINISHED_LENGTH_CAPPED,
248 + RequestStatus.FINISHED_STOPPED,
249 + ):
250 + self._reqs_not_processed.add(request.request_id)
251 + self._reqs_need_save.pop(request.request_id, None)
252 + return False, None
253 +
254 + delay_free_blocks = any(len(group) > 0 for group in block_ids)
255 + remote_num_tokens = 0
256 + if delay_free_blocks:
257 + logger.debug(
258 + "NixlPushConnector request_finished(%s) waiting for %d seconds "
259 + "before releasing blocks",
260 + request.request_id,
261 + self._kv_lease_duration,
262 + )
263 + self._reqs_need_send[request.request_id] = (
264 + time.perf_counter() + self._kv_lease_duration
265 + )
266 +
267 + block_ids = self.get_sw_clipped_blocks(block_ids)
268 + remote_num_tokens = request.num_computed_tokens
269 +
270 + # Store finished blocks for worker-level matching with D
271 + # registrations (via NIXL notifications).
272 + self._finished_request_blocks[request.request_id] = block_ids
273 + self._newly_finished_push_blocks[request.request_id] = block_ids
274 +
275 + return delay_free_blocks, dict(
276 + do_remote_prefill=True,
277 + do_remote_decode=False,
278 + remote_block_ids=block_ids,
279 + remote_engine_id=self.engine_id,
280 + remote_request_id=request.request_id,
281 + remote_host=self.side_channel_host,
282 + remote_port=self.side_channel_port,
283 + tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
284 + remote_num_tokens=remote_num_tokens,
285 + )
286 +
287 + def build_connector_meta(
288 + self,
289 + scheduler_output: SchedulerOutput,
290 + ) -> KVConnectorMetadata:
291 + meta = super().build_connector_meta(scheduler_output)
292 + assert isinstance(meta, NixlConnectorMetadata)
293 +
294 + # Watchdog: any D-side registration whose deadline has passed without
295 + # a corresponding push completion is treated as failed and cleaned up.
296 + # The corresponding request is already tracked via _reqs_need_recv;
297 + # the engine layer will eventually time it out via the lease, but we
298 + # at least drop the stale registration so we don't keep retrying.
299 + now = time.perf_counter()
300 + # Deadlines are inserted in non-decreasing order (monotonic clock +
301 + # constant timeout, armed once per request), and dict insertion order
302 + # is preserved across key deletions, so we can stop at the first
303 + # not-yet-expired entry instead of scanning the whole dict.
304 + expired = []
305 + for rid, deadline in self._push_registration_deadlines.items():
306 + if deadline > now:
307 + break
308 + expired.append(rid)
309 + for rid in expired:
310 + self._push_registration_deadlines.pop(rid, None)
311 + # Avoid resending a registration that already timed out.
312 + self._push_pending_registrations.pop(rid, None)
313 + logger.warning(
314 + "NixlPushConnector: registration for request %s timed out "
315 + "after %.1fs without a push completion",
316 + rid,
317 + self._push_registration_timeout,
318 + )
319 +
320 + # D side: package pending registrations for D workers to send out.
321 + if self._push_pending_registrations:
322 + meta.push_registrations = dict(self._push_pending_registrations)
323 + self._push_pending_registrations.clear()
324 +
325 + # P side: package newly finished blocks for P workers to match against
326 + # any D registrations they have received via NIXL notifications.
327 + if self._newly_finished_push_blocks:
328 + meta.push_finished_blocks = dict(self._newly_finished_push_blocks)
329 + self._newly_finished_push_blocks.clear()
330 +
331 + return meta
332 +
333 + def has_pending_push_work(self) -> bool:
334 + # Keep the engine main loop alive while we have:
335 + # - finished P blocks awaiting WRITE completion, or
336 + # - pending D registrations the worker has not yet shipped, or
337 + # - newly finished blocks not yet shipped to P workers.
338 + return bool(self._finished_request_blocks or self._push_pending_registrations)
339 +
340 + def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
341 + """Clean up finished request blocks after push completes."""
342 + super().update_connector_output(connector_output)
343 + for req_id in connector_output.finished_sending or ():
344 + self._finished_request_blocks.pop(req_id, None)
345 + # On D side, finished_recving means the push completed; clear the
346 + # watchdog so we don't trip an expiration on a fulfilled request.
347 + for req_id in connector_output.finished_recving or ():
348 + self._push_registration_deadlines.pop(req_id, None)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py new
+742
@@ -0,0 +1,742 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Push-specific (WRITE) worker-side logic for the NIXL connector.
4 +
5 +A dedicated ``nixl-push-writer`` thread owns all push-related NIXL ops:
6 +calls ``get_new_notifs`` (routing PUSH_REG internally; HB / completion
7 +notifs are forwarded to the engine main thread), sends PUSH_REG via
8 +``send_notif``, matches D registrations with P finished blocks, and
9 +issues WRITE transfers via ``make_prepped_xfer`` / ``transfer``.
10 +
11 +The engine main thread feeds the writer through three queues:
12 +``_reg_send_inbox`` (D-side regs to send), ``_finished_blocks_inbox``
13 +(P-side blocks from metadata) and ``_pending_completion_notifs``
14 +(non-PUSH_REG notifs forwarded back for HB / completion accounting).
15 +
16 +Wake model: the writer self-polls every
17 +``_PUSH_WRITER_POLL_INTERVAL_MS`` only while it has unmatched
18 +``_push_finished_blocks`` (i.e. P-side blocks waiting for a D PUSH_REG
19 +notif that has no other wake source). All other progress is
20 +event-driven: the engine main thread sets ``_push_writer_wake`` from
21 +``start_load_kv`` (when handing it new work) and from ``get_finished``
22 +(so each engine step gives the writer a chance to drain NIXL notifs);
23 +the handshake-completion callback sets the same event after a deferred
24 +PUSH_REG send has been queued. When a request's lease expires (the base
25 +worker reports it via ``done_sending``) or the WRITE completes,
26 +``get_finished`` enqueues an eviction onto ``_evict_finished_inbox`` so
27 +the writer drops any leftover ``_push_finished_blocks`` /
28 +``_pending_d_registrations`` and stops self-polling.
29 +"""
30 +
31 +import queue
32 +import threading
33 +import time
34 +from collections import defaultdict
35 +from concurrent.futures import Future
36 +from typing import TYPE_CHECKING, Any
37 +
38 +import msgspec
39 +import numpy as np
40 +
41 +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
42 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
43 + NixlBaseConnectorWorker,
44 +)
45 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
46 + PUSH_REG_NOTIF_PREFIX,
47 + NixlConnectorMetadata,
48 + RemoteMeta,
49 + ReqId,
50 + ReqMeta,
51 + TransferHandle,
52 +)
53 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ReadSpec
54 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import get_base_request_id
55 +from vllm.logger import init_logger
56 +
57 +if TYPE_CHECKING:
58 + import torch
59 +
60 + from vllm.config import VllmConfig
61 + from vllm.v1.kv_cache_interface import KVCacheConfig
62 +
63 +logger = init_logger(__name__)
64 +
65 +# Writer-thread poll cadence while there is in-flight push state. When
66 +# fully idle, the writer blocks on a wake event signalled by the engine
67 +# main thread (start_load_kv / get_finished). Smaller -> lower latency
68 +# while active, slightly more CPU.
69 +_PUSH_WRITER_POLL_INTERVAL_MS = 1.0
70 +
71 +
72 +class NixlPushConnectorWorker(NixlBaseConnectorWorker):
73 + """Push-specific (WRITE) worker logic. See module docstring."""
74 +
75 + def __init__(
76 + self,
77 + vllm_config: "VllmConfig",
78 + engine_id: str,
79 + kv_cache_config: "KVCacheConfig",
80 + ):
81 + super().__init__(vllm_config, engine_id, kv_cache_config)
82 +
83 + # Push-specific state.
84 + # P-side: outgoing WRITE handles awaiting completion, keyed by
85 + # request_id. Mutated by writer (submit) and main thread
86 + # (``_pop_done_transfers``); guarded by
87 + # ``_sending_transfers_lock``.
88 + self._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list)
89 + self._sending_transfers_lock = threading.Lock()
90 +
91 + # Writer-thread owned matching state.
92 + # P-side: finished request blocks received from scheduler metadata
93 + # that have not yet been matched with an incoming D registration.
94 + self._push_finished_blocks: dict[ReqId, BlockIds] = {}
95 + # P-side: D registrations received via NIXL notification that have
96 + # not yet been matched with a finished P request.
97 + self._pending_d_registrations: dict[ReqId, dict[str, Any]] = {}
98 +
99 + # Cross-thread channels.
100 + self._reg_send_inbox: queue.Queue[tuple[str, dict[str, Any]]] = queue.Queue()
101 + self._finished_blocks_inbox: queue.Queue[tuple[str, BlockIds]] = queue.Queue()
102 + self._pending_completion_notifs: queue.Queue[bytes] = queue.Queue()
103 + # Main thread → writer: req_ids whose lease has expired or whose
104 + # WRITE has completed. Writer drops them from
105 + # ``_push_finished_blocks`` so an unmatched entry doesn't keep the
106 + # writer busy-polling forever.
107 + self._evict_finished_inbox: queue.Queue[str] = queue.Queue()
108 +
109 + # Wake signal from engine main thread (start_load_kv / get_finished).
110 + # Writer self-polls at _PUSH_WRITER_POLL_INTERVAL_MS while it has
111 + # active in-flight state; otherwise it blocks until signalled.
112 + self._push_writer_wake = threading.Event()
113 +
114 + self._push_writer_stop = threading.Event()
115 + self._push_writer_thread: threading.Thread | None = None
116 +
117 + # --- Lifecycle ----------------------------------------------------- #
118 +
119 + def register_kv_caches(self, kv_caches: dict[str, "torch.Tensor"]):
120 + super().register_kv_caches(kv_caches)
121 + if self._push_writer_thread is None:
122 + self._push_writer_thread = threading.Thread(
123 + target=self._push_writer_loop,
124 + daemon=True,
125 + name="nixl-push-writer",
126 + )
127 + self._push_writer_thread.start()
128 + logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank)
129 +
130 + def shutdown(self):
131 + self._push_writer_stop.set()
132 + # Unblock the writer if it's waiting in the no-active-state branch.
133 + self._push_writer_wake.set()
134 + if self._push_writer_thread is not None:
135 + self._push_writer_thread.join(timeout=2)
136 + self._push_writer_thread = None
137 + with self._sending_transfers_lock:
138 + for handles in self._sending_transfers.values():
139 + for handle in handles:
140 + self.nixl_wrapper.release_xfer_handle(handle)
141 + self._sending_transfers.clear()
142 + super().shutdown()
143 +
144 + # --- Engine-main-thread entry point -------------------------------- #
145 +
146 + def start_load_kv(self, metadata: NixlConnectorMetadata):
147 + """Pre-process metadata; defer NIXL ops to the writer thread."""
148 + # D-side: track reqs waiting for P to push.
149 + for req_id, meta in metadata.reqs_to_recv.items():
150 + meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
151 + meta.local_block_ids
152 + )
153 + assert meta.remote is not None
154 + remote_engine_id = meta.remote.engine_id
155 + logger.debug(
156 + "start_load_kv (push) for request %s from remote engine %s. "
157 + "Num local_block_ids: %s. Num remote_block_ids: %s. ",
158 + req_id,
159 + remote_engine_id,
160 + len(meta.local_physical_block_ids),
161 + len(meta.remote.block_ids),
162 + )
163 + self._recving_metadata[req_id] = meta
164 +
165 + # --- D-side: registrations to send to P via NIXL ---
166 + if metadata.push_registrations:
167 + for req_id, reg_data in metadata.push_registrations.items():
168 + self._reg_send_inbox.put((req_id, reg_data))
169 + self._push_writer_wake.set()
170 +
171 + # --- P-side: newly finished blocks awaiting a D registration match ---
172 + if metadata.push_finished_blocks:
173 + for req_id, block_ids in metadata.push_finished_blocks.items():
174 + self._finished_blocks_inbox.put((req_id, block_ids))
175 + self._push_writer_wake.set()
176 +
177 + # Batch + lease tracking (same as pull).
178 + for req_id in metadata.reqs_in_batch:
179 + self._reqs_to_process.add(req_id)
180 + for req_id in metadata.reqs_not_processed:
181 + self._reqs_to_process.discard(req_id)
182 + assert req_id not in self._reqs_to_send
183 + for req_id, expiration_time in metadata.reqs_to_send.items():
184 + if req_id in self._reqs_to_process:
185 + self._reqs_to_send[req_id] = expiration_time
186 +
187 + # Heartbeats still leave from the main thread (base worker behaviour).
188 + self._send_heartbeats(metadata)
189 +
190 + # --- Writer thread ------------------------------------------------- #
191 +
192 + def _push_writer_loop(self) -> None:
193 + sleep_s = _PUSH_WRITER_POLL_INTERVAL_MS / 1000.0
194 +
195 + while not self._push_writer_stop.is_set():
196 + try:
197 + # 1. D registrations to send.
198 + while True:
199 + try:
200 + rid, rd = self._reg_send_inbox.get_nowait()
201 + except queue.Empty:
202 + break
203 + self._send_registration_to_p(rid, rd)
204 +
205 + # 2. P-side finished blocks; match against pending regs.
206 + while True:
207 + try:
208 + rid, blocks = self._finished_blocks_inbox.get_nowait()
209 + except queue.Empty:
210 + break
211 + matched = self._pop_matching_registration(rid)
212 + if matched is not None:
213 + self._do_start_push_kv(rid, blocks, matched)
214 + else:
215 + self._push_finished_blocks[rid] = blocks
216 +
217 + # 2b. Evict finished blocks for requests that have either
218 + # completed (WRITE acknowledged) or whose lease expired
219 + # without a D registration. Drop pending registrations
220 + # for the same reason so we don't leak state.
221 + while True:
222 + try:
223 + rid = self._evict_finished_inbox.get_nowait()
224 + except queue.Empty:
225 + break
226 + self._push_finished_blocks.pop(rid, None)
227 + self._pending_d_registrations.pop(rid, None)
228 +
229 + # 3. NIXL notifs: route PUSH_REG; forward the rest.
230 + for notifs in self.nixl_wrapper.get_new_notifs().values():
231 + for notif in notifs:
232 + if notif.startswith(PUSH_REG_NOTIF_PREFIX):
233 + self._handle_push_reg_notif(notif)
234 + else:
235 + self._pending_completion_notifs.put(notif)
236 + except Exception:
237 + logger.exception("nixl-push-writer error; continuing")
238 +
239 + # Self-poll only while there is no other wake source: P-side
240 + # finished blocks waiting for a D PUSH_REG match. All other
241 + # progress is event-driven (see module docstring).
242 + if self._push_finished_blocks:
243 + self._push_writer_stop.wait(timeout=sleep_s)
244 + else:
245 + self._push_writer_wake.wait()
246 + self._push_writer_wake.clear()
247 +
248 + def _handle_push_reg_notif(self, notif: bytes) -> None:
249 + try:
250 + reg_data = msgspec.msgpack.decode(notif[len(PUSH_REG_NOTIF_PREFIX) :])
251 + except Exception:
252 + logger.exception("Failed to decode PUSH_REG notification payload")
253 + return
254 + rid = reg_data.get("request_id") if isinstance(reg_data, dict) else None
255 + if not isinstance(rid, str):
256 + logger.warning("PUSH_REG notif missing request_id; dropping")
257 + return
258 +
259 + match = self._pop_matching_finished_blocks(rid)
260 + if match is not None:
261 + fin_id, blocks = match
262 + self._do_start_push_kv(fin_id, blocks, reg_data)
263 + else:
264 + self._pending_d_registrations[rid] = reg_data
265 +
266 + # --- D-side registration send (writer thread) ---------------------- #
267 +
268 + def _send_registration_to_p(
269 + self,
270 + req_id: str,
271 + reg_data: dict[str, Any],
272 + ) -> None:
273 + """Handshake (if needed) then send PUSH_REG. ``send_notif`` always
274 + executes on the writer; the handshake runs on the background executor
275 + and the request is re-queued onto ``_reg_send_inbox`` once it
276 + completes (at which point ``_ensure_handshake`` returns ``None`` and we
277 + send directly)."""
278 + fut = self._ensure_handshake(
279 + reg_data["remote_engine_id"],
280 + reg_data["remote_host"],
281 + reg_data["remote_port"],
282 + reg_data["remote_tp_size"],
283 + )
284 + if fut is None:
285 + self._do_send_reg_notif(req_id, reg_data)
286 + return
287 +
288 + def _on_handshake(
289 + f: Future[dict[int, str]],
290 + rid: str = req_id,
291 + rd: dict[str, Any] = reg_data,
292 + ) -> None:
293 + try:
294 + f.result()
295 + except Exception as e:
296 + self._log_failure(
297 + failure_type="push_reg_handshake_failed", req_id=rid, error=e
298 + )
299 + self._handle_failed_transfer(rid, None)
300 + return
301 + # Re-queue for the writer to send now that the handshake is done.
302 + self._reg_send_inbox.put((rid, rd))
303 + # Wake the writer so it sends the PUSH_REG promptly even if
304 + # otherwise parked.
305 + self._push_writer_wake.set()
306 +
307 + fut.add_done_callback(_on_handshake)
308 +
309 + def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None:
310 + engine_id = reg_data["remote_engine_id"]
311 + notif_msg = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(reg_data)
312 + agents = self._remote_agents.get(engine_id)
313 + if not agents:
314 + logger.error(
315 + "No remote agents for engine %s; cannot send registration for %s",
316 + engine_id,
317 + req_id,
318 + )
319 + self._handle_failed_transfer(req_id, None)
320 + return
321 + for rank, agent_name in agents.items():
322 + try:
323 + self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_msg)
324 + except Exception as e:
325 + self._log_failure(
326 + failure_type="push_reg_notif_failed",
327 + req_id=req_id,
328 + error=e,
329 + remote_rank=rank,
330 + )
331 + logger.debug(
332 + "Sent PUSH_REG for %s to engine %s (%dB)", req_id, engine_id, len(notif_msg)
333 + )
334 +
335 + # --- Matching helpers --------------------------------------------- #
336 +
337 + def _pop_matching_registration(self, request_id: str) -> dict[str, Any] | None:
338 + """Pop the D-side registration matching *request_id*.
339 +
340 + Exact key first, then a match after stripping the random suffix from
341 + both sides. No match leaves the request unmatched (push not started).
342 + """
343 + data = self._pending_d_registrations.pop(request_id, None)
344 + if data is not None:
345 + return data
346 + base_id = get_base_request_id(request_id)
347 + for reg_id in list(self._pending_d_registrations):
348 + if get_base_request_id(reg_id) == base_id:
349 + return self._pending_d_registrations.pop(reg_id)
350 + return None
351 +
352 + def _pop_matching_finished_blocks(
353 + self, request_id: str
354 + ) -> tuple[str, BlockIds] | None:
355 + """Pop the P-side finished blocks matching *request_id*.
356 +
357 + Same lookup as ``_pop_matching_registration``: exact key, then a
358 + match after stripping the random suffix from both sides.
359 + """
360 + blocks = self._push_finished_blocks.pop(request_id, None)
361 + if blocks is not None:
362 + return request_id, blocks
363 + base_id = get_base_request_id(request_id)
364 + for fin_id in list(self._push_finished_blocks):
365 + if get_base_request_id(fin_id) == base_id:
366 + return fin_id, self._push_finished_blocks.pop(fin_id)
367 + return None
368 +
369 + # --- WRITE transfer logic (writer thread) ------------------------- #
370 +
371 + def _do_start_push_kv(
372 + self,
373 + request_id: str,
374 + local_block_ids: BlockIds,
375 + registration_data: dict[str, Any],
376 + ) -> None:
377 + """Start push-based KV transfer from P worker to D node.
378 +
379 + ``local_block_ids`` are P's *logical* block IDs (from the P
380 + scheduler's metadata). ``registration_data["local_block_ids"]``
381 + are D's *logical* block IDs (from D's scheduler, sent over the
382 + PUSH_REG notif). All conversion to physical block IDs is
383 + deferred to ``_xfer_blocks_for_req`` so each side uses its own
384 + physical-blocks-per-logical ratio (P uses
385 + ``self._physical_blocks_per_logical_kv_block``; D's ratio is
386 + learned during the NIXL handshake)."""
387 + decode_engine_id = registration_data["decode_engine_id"]
388 + remote_block_ids = registration_data["local_block_ids"]
389 + decode_host = registration_data["decode_host"]
390 + decode_port = registration_data["decode_port"]
391 + decode_request_id = registration_data["request_id"]
392 + if not local_block_ids:
393 + logger.warning("No local blocks to push for request %s", request_id)
394 + return
395 +
396 + if not self._ensure_d_handshake(
397 + decode_engine_id,
398 + decode_host,
399 + decode_port,
400 + registration_data["decode_tp_size"],
401 + request_id,
402 + ):
403 + return
404 +
405 + # Both sides are kept in logical form here; ``_xfer_blocks_for_req``
406 + # expands each side using the appropriate ratio.
407 + logical_local = self._as_grouped_block_ids(local_block_ids)
408 + logical_remote = self._as_grouped_block_ids(remote_block_ids)
409 + physical_local = self._logical_to_kernel_block_ids(logical_local)
410 +
411 + push_meta = ReqMeta(
412 + local_block_ids=logical_local,
413 + local_physical_block_ids=physical_local,
414 + tp_size=self.world_size,
415 + remote=RemoteMeta(
416 + block_ids=logical_remote,
417 + host="",
418 + port=0,
419 + engine_id=decode_engine_id,
420 + request_id=decode_request_id,
421 + ),
422 + )
423 +
424 + t0 = time.perf_counter()
425 + self._xfer_blocks_for_req(req_id=request_id, meta=push_meta)
426 + elapsed_ms = (time.perf_counter() - t0) * 1000.0
427 + if elapsed_ms > 200.0:
428 + logger.warning(
429 + "_do_start_push_kv for %s took %.1fms (slow NIXL submission)",
430 + request_id,
431 + elapsed_ms,
432 + )
433 +
434 + def _ensure_d_handshake(
435 + self,
436 + decode_engine_id: str,
437 + decode_host: str,
438 + decode_port: int,
439 + decode_tp_size: int,
440 + request_id: str,
441 + ) -> bool:
442 + """First-time P→D handshake. Blocking call on the writer thread.
443 +
444 + Returns True iff the handshake succeeded (or had already been
445 + completed). Returns False if the handshake raised; the request is
446 + skipped in that case (the engine layer will reschedule or fail it
447 + via the standard lease/timeout path)."""
448 + if decode_engine_id in self._remote_agents:
449 + return True
450 + try:
451 + remote_agents = self._nixl_handshake(
452 + decode_host,
453 + decode_port,
454 + decode_tp_size,
455 + decode_engine_id,
456 + )
457 + except Exception:
458 + logger.exception(
459 + "Failed handshake to D %s for push %s",
460 + decode_engine_id,
461 + request_id,
462 + )
463 + return False
464 + with self._handshake_lock:
465 + self._remote_agents[decode_engine_id] = remote_agents
466 + logger.info(
467 + "Push handshake to D %s done (%d agents)",
468 + decode_engine_id,
469 + len(remote_agents),
470 + )
471 + return True
472 +
473 + @staticmethod
474 + def _as_grouped_block_ids(block_ids: BlockIds) -> BlockIds:
475 + """Normalise a sequence of block IDs to a tuple-of-groups shape.
476 +
477 + ``BlockIds`` is canonically a tuple of per-group lists, but some
478 + registration payloads collapse a single-group case to a flat
479 + list. Re-wrap that case so downstream group-aware helpers see a
480 + consistent shape."""
481 + if block_ids and not isinstance(block_ids[0], (list, tuple)):
482 + return (list(block_ids),)
483 + return block_ids
484 +
485 + def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta):
486 + """Issue WRITE transfers to one or more remote TP ranks."""
487 + assert meta.remote is not None and self.transfer_topo is not None
488 + engine_id = meta.remote.engine_id
489 + plan = self.tp_mappings[engine_id]
490 + remote_info = self.transfer_topo.get_engine_info(engine_id)
491 + tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
492 +
493 + # Expand D's logical IDs using the ratio learned during the
494 + # NIXL handshake. ``meta`` is freshly built by
495 + # ``_do_start_push_kv`` so mutating it here is safe.
496 + meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
497 + meta.remote.block_ids,
498 + remote_info.remote_physical_blocks_per_logical,
499 + )
500 + remote_block_ids = meta.remote.block_ids
501 + local_block_ids = meta.local_physical_block_ids
502 + num_groups = len(local_block_ids)
503 + read_specs = [
504 + ReadSpec(
505 + remote_rank=rank,
506 + local_block_ids=[
507 + list(local_block_ids[g])
508 + if rank in plan.source_ranks_per_group[g]
509 + else []
510 + for g in range(num_groups)
511 + ],
512 + remote_block_ids=[
513 + list(remote_block_ids[g])
514 + if rank in plan.source_ranks_per_group[g]
515 + else []
516 + for g in range(num_groups)
517 + ],
518 + )
519 + for rank in plan.all_source_ranks
520 + ]
521 +
522 + if self.use_mla and tp_ratio < 0:
523 + assert len(read_specs) == 1
524 +
525 + for i, spec in enumerate(read_specs):
526 + remote_block_size = remote_info.remote_block_size
527 + logger.debug(
528 + "Remote agent %s available, calling _xfer_blocks"
529 + " on remote rank %s with remote block size %s for req %s",
530 + meta.remote.engine_id,
531 + spec.remote_rank,
532 + remote_block_size,
533 + req_id,
534 + )
535 + if tp_ratio < 0 and not self.use_mla:
536 + assert remote_block_size == self.block_size
537 + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
538 + else:
539 + local_xfer_side_handle = self.src_xfer_handles_by_block_size[
540 + remote_block_size
541 + ]
542 +
543 + remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
544 + spec.remote_rank
545 + ]
546 +
547 + self._xfer_blocks(
548 + read_spec=spec,
549 + request_id=req_id,
550 + dst_engine_id=meta.remote.engine_id,
551 + remote_request_id=meta.remote.request_id,
552 + local_xfer_side_handle=local_xfer_side_handle,
553 + remote_xfer_side_handle=remote_xfer_side_handle,
554 + )
555 +
556 + if self.use_mla and tp_ratio < 0 and read_specs:
557 + notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()
558 + remote_agents = self._remote_agents[meta.remote.engine_id]
559 + for rank_to_notify, agent in remote_agents.items():
560 + if rank_to_notify != read_specs[0].remote_rank:
561 + self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
562 +
563 + def _xfer_blocks(
564 + self,
565 + read_spec: ReadSpec,
566 + dst_engine_id: str,
567 + request_id: str,
568 + remote_request_id: str,
569 + local_xfer_side_handle: int,
570 + remote_xfer_side_handle: int,
571 + ):
572 + """Post a WRITE point-to-point xfer request."""
573 + assert self.transfer_topo is not None
574 + remote_rank = read_spec.remote_rank
575 + local_block_ids = read_spec.local_block_ids
576 + remote_block_ids = read_spec.remote_block_ids
577 +
578 + remote_info = self.transfer_topo.get_engine_info(dst_engine_id)
579 + block_size_ratio = self.transfer_topo.block_size_ratio(
580 + remote_info.remote_block_size
581 + )
582 + if block_size_ratio > 1:
583 + assert not self._is_hma_required
584 + local_block_ids0 = local_block_ids[0] if local_block_ids else []
585 + remote_block_ids0 = remote_block_ids[0]
586 + local_block_ids_mapped = self.get_mapped_blocks(
587 + np.asarray(local_block_ids0), block_size_ratio
588 + ).tolist()
589 + if len(local_block_ids_mapped) > len(remote_block_ids0):
590 + local_block_ids_mapped = local_block_ids_mapped[
591 + : len(remote_block_ids0)
592 + ]
593 + local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
594 + remote_block_ids = [remote_block_ids0]
595 +
596 + notif_id = f"{remote_request_id}:{self.world_size}".encode()
597 +
598 + if len(local_block_ids) == 0:
599 + logger.warning("No blocks to push for request %s", request_id)
600 + return
601 +
602 + # Align per-group block counts for push.
603 + local_block_ids = list(local_block_ids)
604 + remote_block_ids = list(remote_block_ids)
605 + for i in range(min(len(local_block_ids), len(remote_block_ids))):
606 + num_local = len(local_block_ids[i])
607 + num_remote = len(remote_block_ids[i])
608 + if num_local > num_remote:
609 + local_block_ids[i] = local_block_ids[i][:num_remote]
610 + elif num_local < num_remote:
611 + remote_block_ids[i] = remote_block_ids[i][:num_local]
612 +
613 + # Get descs ids.
614 + remote_block_descs_ids = self._compute_desc_ids(
615 + block_ids=remote_block_ids,
616 + dst_num_blocks=self.dst_num_blocks[dst_engine_id],
617 + block_size_ratio=None,
618 + physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
619 + )
620 + local_block_descs_ids = self._compute_desc_ids(
621 + block_ids=local_block_ids,
622 + dst_num_blocks=self.dst_num_blocks[self.engine_id],
623 + block_size_ratio=block_size_ratio,
624 + physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
625 + )
626 +
627 + assert len(local_block_descs_ids) == len(remote_block_descs_ids)
628 +
629 + handle = None
630 + try:
631 + handle = self.nixl_wrapper.make_prepped_xfer(
632 + "WRITE",
633 + local_xfer_side_handle,
634 + local_block_descs_ids,
635 + remote_xfer_side_handle,
636 + remote_block_descs_ids,
637 + notif_msg=notif_id,
638 + )
639 + self.nixl_wrapper.transfer(handle)
640 + # Track push WRITE handles so P can free blocks once done.
641 + with self._sending_transfers_lock:
642 + self._sending_transfers[request_id].append(handle)
643 + except Exception as e:
644 + self._log_failure(
645 + failure_type="transfer_setup_failed",
646 + req_id=request_id,
647 + msg="Push WRITE submission failed; releasing handle",
648 + error=e,
649 + dst_engine_id=dst_engine_id,
650 + remote_rank=remote_rank,
651 + )
652 + # On the P side this WRITE failure is purely outbound; we
653 + # don't have a ``_recving_metadata`` entry to invalidate, so
654 + # we just release the handle and let the engine reschedule
655 + # via the lease / watchdog.
656 + if handle is not None:
657 + self.nixl_wrapper.release_xfer_handle(handle)
658 + self.xfer_stats.record_failed_transfer()
659 +
660 + # --- Notification handling on engine main thread ------------------ #
661 +
662 + def _get_new_notifs(self) -> set[str]:
663 + """Drain HB / completion notifs forwarded by the writer thread.
664 +
665 + The writer owns ``nixl_wrapper.get_new_notifs`` for push; PUSH_REG
666 + notifs are handled there. Everything else is forwarded here for
667 + existing accounting.
668 + """
669 + assert self.transfer_topo is not None
670 + notified_req_ids: set[str] = set()
671 + while True:
672 + try:
673 + notif = self._pending_completion_notifs.get_nowait()
674 + except queue.Empty:
675 + break
676 +
677 + msg = notif.decode("utf-8")
678 + if msg.startswith("HB:"):
679 + self._handle_heartbeat(msg[3:])
680 + continue
681 +
682 + req_id, tp_size = msg.rsplit(":", 1)
683 +
684 + # Not tracked as a P-side send/process for this notif.
685 + if req_id not in self._reqs_to_send and req_id not in self._reqs_to_process:
686 + if req_id in self._recving_metadata:
687 + # D-side: P signalled push completion. The transfer was
688 + # driven entirely by P (we don't own a NIXL handle here),
689 + # so materialise an empty entry in ``_recving_transfers``
690 + # and let ``_pop_done_transfers`` report it done on the
691 + # next ``get_finished``.
692 + self._recving_transfers.setdefault(req_id, [])
693 + else:
694 + # Not tracked on either side (lease may have expired
695 + # before the notif arrived). Log and skip.
696 + logger.error(
697 + "Unrecognized request %s notif (may have expired).",
698 + req_id,
699 + )
700 + continue
701 +
702 + n_consumers = int(tp_size)
703 + tp_ratio = self.transfer_topo.tp_ratio(n_consumers)
704 + consumers_per_producer = -tp_ratio if n_consumers > self.world_size else 1
705 + self.consumer_notification_counts_by_req[req_id] += 1
706 + if (
707 + self.consumer_notification_counts_by_req[req_id]
708 + == consumers_per_producer
709 + ):
710 + notified_req_ids.add(req_id)
711 + del self.consumer_notification_counts_by_req[req_id]
712 + self._reqs_to_process.remove(req_id)
713 + self._reqs_to_send.pop(req_id, None)
714 + return notified_req_ids
715 +
716 + def get_finished(self) -> tuple[set[str], set[str]]:
717 + # Engine main thread asking for completions: also wake the writer
718 + # so it gets a chance to drain NIXL notifs (heartbeats, completion
719 + # notifs, late PUSH_REGs) even if it had been parked.
720 + self._push_writer_wake.set()
721 +
722 + done_sending, done_recving = super().get_finished()
723 +
724 + # ``_pop_done_transfers`` mutates ``_sending_transfers``; the
725 + # writer thread also appends to it, so guard the pop.
726 + with self._sending_transfers_lock:
727 + done_pushing = self._pop_done_transfers(self._sending_transfers)
728 + for req_id in done_pushing:
729 + self._reqs_to_send.pop(req_id, None)
730 + self._reqs_to_process.discard(req_id)
731 + self.consumer_notification_counts_by_req.pop(req_id, None)
732 + done_sending.add(req_id)
733 +
734 + # Tell the writer to drop any state it still holds for any
735 + # request that just finished (push completed) or expired
736 + # (lease ran out without a D registration ever arriving).
737 + for req_id in done_sending:
738 + self._evict_finished_inbox.put(req_id)
739 + if done_sending:
740 + self._push_writer_wake.set()
741 +
742 + return done_sending, done_recving
vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py
+6 -668
@@ -1,674 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Scheduler-side logic for the NIXL connector."""
3 +"""Backward-compatible re-export of NixlPullConnectorScheduler."""
4
5 -import threading
6 -import time
7 -from typing import TYPE_CHECKING, Any
8 -
9 -import msgspec
10 -import zmq
11 -
12 -from vllm import envs
13 -from vllm.distributed.kv_transfer.kv_connector.utils import (
14 - BlockIds,
15 - EngineId,
16 - yield_req_data,
17 -)
18 -from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19 - KVConnectorHandshakeMetadata,
20 - KVConnectorMetadata,
21 -)
22 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
23 - GET_META_MSG,
24 - HeartbeatInfo,
25 - NixlConnectorMetadata,
26 - NixlHandshakePayload,
27 - ReqId,
5 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
6 + NixlPullConnectorScheduler,
7 )
29 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx
30 -from vllm.logger import init_logger
31 -from vllm.platforms import current_platform
32 -from vllm.utils.math_utils import cdiv
33 -from vllm.utils.network_utils import make_zmq_path
34 -from vllm.v1.core.sched.output import SchedulerOutput
35 -from vllm.v1.kv_cache_interface import (
36 - FullAttentionSpec,
37 - MambaSpec,
38 - SlidingWindowSpec,
39 -)
40 -
41 -if TYPE_CHECKING:
42 - from vllm.config import VllmConfig
43 - from vllm.v1.core.kv_cache_manager import KVCacheBlocks
44 - from vllm.v1.kv_cache_interface import KVCacheConfig
45 - from vllm.v1.outputs import KVConnectorOutput
46 - from vllm.v1.request import Request
47 -
48 -logger = init_logger(__name__)
49 -
50 -
51 -class NixlConnectorScheduler:
52 - """Implementation of Scheduler side methods"""
53 -
54 - def __init__(
55 - self,
56 - vllm_config: "VllmConfig",
57 - engine_id: str,
58 - kv_cache_config: "KVCacheConfig",
59 - ):
60 - self.vllm_config = vllm_config
61 - self.block_size = vllm_config.cache_config.block_size
62 - self.engine_id: EngineId = engine_id
63 - self.kv_cache_config = kv_cache_config
64 - self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
65 - self.side_channel_port = (
66 - envs.VLLM_NIXL_SIDE_CHANNEL_PORT
67 - + vllm_config.parallel_config.data_parallel_index
68 - )
69 - assert vllm_config.kv_transfer_config is not None
70 - self._kv_lease_duration: int = (
71 - vllm_config.kv_transfer_config.get_from_extra_config(
72 - "kv_lease_duration", 30
73 - )
74 - )
75 - # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.
76 - self._heartbeat_interval = self._kv_lease_duration // 6
77 - if current_platform.device_type == "cpu":
78 - self.use_host_buffer = False
79 - else:
80 - self.use_host_buffer = (
81 - vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
82 - )
83 - self._is_hma_required = (
84 - not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
85 - # Also handle unlikely SW-only model case instead of checking num_groups>1.
86 - and any(
87 - not isinstance(g.kv_cache_spec, FullAttentionSpec)
88 - for g in kv_cache_config.kv_cache_groups
89 - )
90 - )
91 - self._has_mamba = any(
92 - isinstance(g.kv_cache_spec, MambaSpec)
93 - for g in kv_cache_config.kv_cache_groups
94 - )
95 -
96 - logger.info("Initializing NIXL Scheduler %s", engine_id)
97 - if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
98 - logger.info("Hybrid Memory Allocator is enabled with NIXL")
99 -
100 - # Background thread for handling new handshake requests.
101 - self._nixl_handshake_listener_t: threading.Thread | None = None
102 - self._stop_event = threading.Event()
103 -
104 - # Requests that need to start recv/send.
105 - # New requests are added by update_state_after_alloc in
106 - # the scheduler. Used to make metadata passed to Worker.
107 - self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
108 - self._reqs_need_save: dict[ReqId, Request] = {}
109 - # Reqs to send and their expiration time
110 - self._reqs_need_send: dict[ReqId, float] = {}
111 - self._reqs_in_batch: set[ReqId] = set()
112 - # Reqs to remove from processed set because they're not to send after
113 - # remote prefill or aborted.
114 - self._reqs_not_processed: set[ReqId] = set()
115 -
116 - # Heartbeat tracking: requests needing periodic lease-renewal heartbeats to
117 - # remote P-side, stored as ready-to-send HeartbeatInfo grouped by remote engine
118 - self._heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {}
119 - # Reverse lookup: local req_id -> (engine_id, remote_req_id) for O(1) removal
120 - self._heartbeat_req_engine: dict[ReqId, tuple[EngineId, ReqId]] = {}
121 - self._last_heartbeat_time: float = 0.0
122 -
123 - # Gather Sliding Window sizes for each kv cache group (if any) in number of
124 - # blocks per KV cache group. This is used to clip the local attention window.
125 - sw_sizes_tokens: list[tuple[int, int]] = [
126 - (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
127 - if isinstance(g.kv_cache_spec, SlidingWindowSpec)
128 - else (0, self.block_size)
129 - for g in kv_cache_config.kv_cache_groups
130 - ]
131 - # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively
132 - # account for boundary overlap eg window isn't fully aligned with blocks.
133 - self.blocks_per_sw = [
134 - cdiv(n_tokens, block_size) + 1 if n_tokens else 0
135 - for n_tokens, block_size in sw_sizes_tokens
136 - ]
137 -
138 - # Threshold to decide whether to compute kv cache locally
139 - # or pull from a remote node: minimum number of remote
140 - # tokens to amortize the xfer latencies
141 - self.kv_recompute_threshold: int = int(
142 - vllm_config.kv_transfer_config.get_from_extra_config(
143 - "kv_recompute_threshold", 64
144 - )
145 - )
146 -
147 - # Bi-directional KV transfer feature supports KV block
148 - # transfers from D node to P node
149 - self.is_bidirectional_kv_xfer_enabled = (
150 - vllm_config.kv_transfer_config.get_from_extra_config(
151 - "bidirectional_kv_xfer", False
152 - )
153 - )
154 - self.decoder_kv_blocks_ttl = (
155 - vllm_config.kv_transfer_config.get_from_extra_config(
156 - "decoder_kv_blocks_ttl", 480
157 - )
158 - )
159 -
160 - if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0:
161 - logger.info(
162 - "Bidirectional KV transfer is enabled and the kv "
163 - "recompute threshold is set to %d tokens."
164 - "KV blocks on D are released after a TTL of %d seconds.",
165 - self.kv_recompute_threshold,
166 - self.decoder_kv_blocks_ttl,
167 - )
168 -
169 - def shutdown(self):
170 - self._stop_event.set()
171 - if self._nixl_handshake_listener_t is not None:
172 - self._nixl_handshake_listener_t.join()
173 - self._nixl_handshake_listener_t = None
174 -
175 - def on_new_request(self, request: "Request") -> None:
176 - """Track a request that may need heartbeats."""
177 - params = request.kv_transfer_params
178 - # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are
179 - # effectively disabled for Bidirectional KV transfer.
180 - if params is None or not params.get("do_remote_prefill"):
181 - return
182 - # Only track if all required remote fields are present.
183 - remote_engine_id = params.get("remote_engine_id")
184 - remote_request_id = params.get("remote_request_id")
185 - host = params.get("remote_host")
186 - port = params.get("remote_port")
187 - tp_size = params.get("tp_size")
188 - if (
189 - remote_engine_id is None
190 - or remote_request_id is None
191 - or host is None
192 - or port is None
193 - or tp_size is None
194 - ):
195 - return
196 - if remote_engine_id not in self._heartbeat_by_engine:
197 - self._heartbeat_by_engine[remote_engine_id] = HeartbeatInfo(
198 - req_ids=set(),
199 - host=host,
200 - port=port,
201 - tp_size=tp_size,
202 - )
203 - self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id)
204 - self._heartbeat_req_engine[request.request_id] = (
205 - remote_engine_id,
206 - remote_request_id,
207 - )
208 -
209 - def _stop_heartbeat(self, req_id: ReqId) -> None:
210 - """Remove *req_id* from heartbeat tracking (if tracked)."""
211 - if key := self._heartbeat_req_engine.pop(req_id, None):
212 - engine_id, remote_id = key
213 - if info := self._heartbeat_by_engine.get(engine_id):
214 - info.req_ids.discard(remote_id)
215 - if not info.req_ids:
216 - # Clean up empty engines so we don't leak a key when remote dies.
217 - del self._heartbeat_by_engine[engine_id]
218 -
219 - def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
220 - """
221 - Clip the number of blocks to the sliding window size for each kv cache group
222 - that employs SWA.
223 - This is necessary because the KV Cache manager initially allocates blocks for
224 - the entire sequence length, and successively cleans up blocks that are outside
225 - the window prior to the `request_finished_all_groups` hook.
226 - """
227 - if len(block_ids) == 0 or not self._is_hma_required:
228 - # No blocks to clip eg Full prefix cache hit or not a hybrid model.
229 - return block_ids
230 - # NOTE (NickLucche) This logic is currently handled at the connector level
231 - # because offloading connectors might want to receive the whole sequence even
232 - # for SWA groups. We will abstract this logic once the interface is more stable
233 - assert len(block_ids) == len(self.blocks_per_sw), (
234 - "Number of KV cache groups must match"
235 - )
236 - # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged
237 - return tuple(
238 - [
239 - blocks[-self.blocks_per_sw[i] :]
240 - if self.blocks_per_sw[i] > 0
241 - else blocks
242 - for i, blocks in enumerate(block_ids)
243 - ]
244 - )
245 -
246 - def set_xfer_handshake_metadata(
247 - self, metadata: dict[int, KVConnectorHandshakeMetadata]
248 - ) -> None:
249 - """
250 - Set the KV connector handshake metadata for this connector.
251 -
252 - Args:
253 - metadata (dict): the handshake metadata to set.
254 - """
255 - encoded_data: dict[int, bytes] = {}
256 - encoder = msgspec.msgpack.Encoder()
257 - for tp_rank, rank_metadata in metadata.items():
258 - if not isinstance(rank_metadata, NixlHandshakePayload):
259 - raise ValueError(
260 - "NixlConnectorScheduler expects NixlHandshakePayload for "
261 - "handshake metadata."
262 - )
263 - encoded_data[tp_rank] = encoder.encode(rank_metadata)
264 - logger.debug(
265 - "Tp rank %d: encoded NixlHandshakePayload size: %s bytes",
266 - tp_rank,
267 - str(len(encoded_data[tp_rank])),
268 - )
269 -
270 - # Only start the listener when we have metadata to serve.
271 - if self._nixl_handshake_listener_t is None:
272 - ready_event = threading.Event()
273 - self._nixl_handshake_listener_t = threading.Thread(
274 - target=self._nixl_handshake_listener,
275 - args=(
276 - encoded_data,
277 - ready_event,
278 - self._stop_event,
279 - self.side_channel_host,
280 - self.side_channel_port,
281 - ),
282 - daemon=True,
283 - name="nixl_handshake_listener",
284 - )
285 - self._nixl_handshake_listener_t.start()
286 - ready_event.wait() # Wait for listener ZMQ socket to be ready.
287 -
288 - @staticmethod
289 - def _nixl_handshake_listener(
290 - encoded_data: dict[int, Any],
291 - ready_event: threading.Event,
292 - stop_event: threading.Event,
293 - host: str,
294 - port: int,
295 - ):
296 - """Background thread for getting new NIXL handshakes."""
297 - # NOTE(rob): this is a simple implementation. We will move
298 - # to a better approach via HTTP endpoint soon.
299 -
300 - # Listen for new requests for metadata.
301 - path = make_zmq_path("tcp", host, port)
302 - logger.debug("Starting listening on path: %s", path)
303 - with zmq_ctx(zmq.ROUTER, path) as sock:
304 - sock.setsockopt(zmq.RCVTIMEO, 1000)
305 - ready_event.set()
306 - while True:
307 - try:
308 - identity, _, msg = sock.recv_multipart()
309 - except zmq.Again:
310 - if stop_event.is_set():
311 - break
312 - continue
313 - # Decode the message which contains (GET_META_MSG, rank)
314 - msg, target_tp_rank = msgspec.msgpack.decode(msg)
315 - logger.debug(
316 - "Received message for tp rank %s",
317 - target_tp_rank,
318 - )
319 - if msg != GET_META_MSG:
320 - logger.warning("Connection listener got unexpected message %s", msg)
321 - sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))
322 -
323 - def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int:
324 - """D-side only. Returns N-1 for Mamba models since the decoder
325 - always recomputes the last token and must start from h(N-1)."""
326 - if self._has_mamba and num_prompt_tokens > 1:
327 - return num_prompt_tokens - 1
328 - return num_prompt_tokens
329 -
330 - def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
331 - """P-side only: drop the last prompt token so the prefiller computes
332 - h(N-1) instead of h(N). The decoder recomputes the last token to
333 - derive h(N) correctly.
334 -
335 - Guarded by ``_p_side_truncated`` to avoid repeated truncation if the
336 - request is preempted and rescheduled."""
337 - params = request.kv_transfer_params
338 - if (
339 - params is not None
340 - # Guard against repeated truncation after preemption/reschedule.
341 - and not params.get("_p_side_truncated")
342 - and request.num_prompt_tokens > 1
343 - ):
344 - if request.prompt_token_ids is not None:
345 - request.prompt_token_ids.pop()
346 - elif request.prompt_embeds is not None:
347 - request.prompt_embeds = request.prompt_embeds[:-1]
348 - else:
349 - return
350 -
351 - request._all_token_ids.pop()
352 - request.num_prompt_tokens -= 1
353 - request.max_tokens = 1
354 - params["_p_side_truncated"] = True
355 -
356 - def get_num_new_matched_tokens(
357 - self, request: "Request", num_computed_tokens: int
358 - ) -> tuple[int, bool]:
359 - """
360 - For remote prefill, pull all prompt blocks from remote
361 - asynchronously relative to engine execution.
362 -
363 - Args:
364 - request (Request): the request object.
365 - num_computed_tokens (int): the number of locally
366 - computed tokens for this request
367 - Returns:
368 - * the number of tokens that can be loaded from the
369 - external KV cache beyond what is already computed.
370 - * true if the external KV cache tokens will be loaded
371 - asynchronously (between scheduler steps).
372 - """
373 -
374 - params = request.kv_transfer_params
375 - logger.debug(
376 - "NIXLConnector get_num_new_matched_tokens: "
377 - "num_computed_tokens=%s, kv_transfer_params=%s",
378 - num_computed_tokens,
379 - params,
380 - )
381 -
382 - if params is not None and params.get("do_remote_prefill"):
383 - # Remote prefill: get all prompt blocks from remote.
384 - token_ids = request.prompt_token_ids or []
385 - actual = self._mamba_prefill_token_count(len(token_ids))
386 - count = actual - num_computed_tokens
387 - if count > 0:
388 - return count, True
389 -
390 - if params is not None and params.get("do_remote_decode") and self._has_mamba:
391 - self._truncate_mamba_request_for_prefill(request)
392 -
393 - if (
394 - params is not None
395 - and params.get("do_remote_decode")
396 - and params.get("remote_block_ids")
397 - and all(
398 - p in params
399 - for p in (
400 - "remote_engine_id",
401 - "remote_request_id",
402 - "remote_host",
403 - "remote_port",
404 - )
405 - )
406 - ):
407 - # Decode node has kv blocks for part of prefill request, so, provide them
408 - # as an external token count to scheduler.
409 - # The tokens will be loaded if not already present
410 - # in the prefill node local cache
411 - remote_num_tokens = params.get("remote_num_tokens") or 0
412 - count = (
413 - min(remote_num_tokens, request.num_prompt_tokens) - num_computed_tokens
414 - )
415 - if count > 0:
416 - # Check kv_recompute_threshold: skip pull if
417 - # remote tokens are below the threshold.
418 - if (
419 - self.kv_recompute_threshold > 0
420 - and count < self.kv_recompute_threshold
421 - ):
422 - logger.debug(
423 - "Skipping remote pull for %s: %d remote tokens < threshold %d",
424 - request.request_id,
425 - count,
426 - self.kv_recompute_threshold,
427 - )
428 - return 0, False
429 - return count, True
430 -
431 - # No remote prefill for this request.
432 - return 0, False
433 -
434 - def update_state_after_alloc(
435 - self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
436 - ):
437 - params = request.kv_transfer_params
438 - logger.debug(
439 - "NIXLConnector update_state_after_alloc: "
440 - "num_external_tokens=%s, kv_transfer_params=%s",
441 - num_external_tokens,
442 - params,
443 - )
444 -
445 - if not params:
446 - return
447 -
448 - if params.get("do_remote_decode") or (
449 - params.get("do_remote_prefill") and self.is_bidirectional_kv_xfer_enabled
450 - ):
451 - self._reqs_in_batch.add(request.request_id)
452 - if self.use_host_buffer and params.get("do_remote_decode"):
453 - # NOTE: when accelerator is not directly supported by Nixl,
454 - # prefilled blocks need to be saved to host memory before transfer.
455 - self._reqs_need_save[request.request_id] = request
456 - elif params.get("do_remote_prefill") or (
457 - params.get("do_remote_decode")
458 - and self.is_bidirectional_kv_xfer_enabled
459 - and not params.get("_remote_blocks_processed")
460 - ):
461 - if params.get("remote_block_ids"):
462 - if all(
463 - p in params
464 - for p in (
465 - "remote_engine_id",
466 - "remote_request_id",
467 - "remote_host",
468 - "remote_port",
469 - )
470 - ):
471 - # If remote_blocks and num_external_tokens = 0, we have
472 - # a full prefix cache hit on the local node. We need to call
473 - # send_notif in _read_blocks to free the memory on the remote node.
474 -
475 - unhashed_local_block_ids: BlockIds = (
476 - blocks.get_unhashed_block_ids_all_groups()
477 - if num_external_tokens > 0
478 - else ()
479 - )
480 - local_block_ids = self.get_sw_clipped_blocks(
481 - unhashed_local_block_ids
482 - )
483 -
484 - # Get unhashed blocks to pull from remote. Mind that a full prefix
485 - # cache hit is indicated with an empty list.
486 - self._reqs_need_recv[request.request_id] = (
487 - request,
488 - local_block_ids,
489 - )
490 -
491 - else:
492 - logger.warning(
493 - "Got invalid KVTransferParams: %s. This "
494 - "request will not utilize KVTransfer",
495 - params,
496 - )
497 - else:
498 - assert num_external_tokens == 0
499 - # Only trigger 1 KV transfer per request.
500 - params["do_remote_prefill"] = False
501 - params["_remote_blocks_processed"] = True
502 -
503 - def _build_save_meta(
504 - self,
505 - meta: NixlConnectorMetadata,
506 - scheduler_output: SchedulerOutput,
507 - ) -> None:
508 - # only called when use_host_buffer is True to build the save metadata
509 -
510 - # NOTE: For the prefill side, there might be a chance that an early added
511 - # request is a chunked prefill, so we need to check if new blocks are added
512 - for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output):
513 - req_to_save = self._reqs_need_save.get(req_id)
514 - if req_to_save is None or new_block_id_groups is None:
515 - continue
516 - req = req_to_save
517 -
518 - assert req.kv_transfer_params is not None
519 - clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups)
520 - meta.add_new_req_to_save(
521 - request_id=req_id,
522 - local_block_ids=clipped_block_id_groups,
523 - kv_transfer_params=req.kv_transfer_params,
524 - )
525 - assert scheduler_output.num_scheduled_tokens is not None
526 - num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
527 - is_partial = (
528 - req.num_computed_tokens + num_scheduled_tokens
529 - ) < req.num_prompt_tokens
530 - if not is_partial:
531 - # For non-partial prefills, once new req_meta is scheduled, it
532 - # can be removed from _reqs_need_save.
533 - # For partial prefill case, we will retain the request in
534 - # _reqs_need_save until all blocks are scheduled with req_meta.
535 - # Therefore, only pop if `not is_partial`.
536 - self._reqs_need_save.pop(req_id)
537 -
538 - def build_connector_meta(
539 - self,
540 - scheduler_output: SchedulerOutput,
541 - ) -> KVConnectorMetadata:
542 - meta = NixlConnectorMetadata()
543 -
544 - # Loop through scheduled reqs and convert to ReqMeta.
545 - for req_id, (req, block_ids) in self._reqs_need_recv.items():
546 - assert req.kv_transfer_params is not None
547 - meta.add_new_req_to_recv(
548 - request_id=req_id,
549 - local_block_ids=block_ids,
550 - kv_transfer_params=req.kv_transfer_params,
551 - )
552 -
553 - if self.use_host_buffer:
554 - self._build_save_meta(meta, scheduler_output)
555 -
556 - meta.reqs_to_send = self._reqs_need_send
557 - meta.reqs_in_batch = self._reqs_in_batch
558 - meta.reqs_not_processed = self._reqs_not_processed
559 -
560 - # Package heartbeats, throttled by heartbeat_interval.
561 - if self._heartbeat_by_engine:
562 - now = time.perf_counter()
563 - if now - self._last_heartbeat_time >= self._heartbeat_interval:
564 - self._last_heartbeat_time = now
565 - meta.heartbeat_by_engine = self._heartbeat_by_engine
566 -
567 - # Clear the list once workers start the transfers
568 - self._reqs_need_recv.clear()
569 - self._reqs_in_batch = set()
570 - self._reqs_not_processed = set()
571 - self._reqs_need_send = {}
572 -
573 - return meta
574 -
575 - def update_connector_output(self, connector_output: "KVConnectorOutput") -> None:
576 - """Stop heartbeating for requests whose KV transfer completed."""
577 - for req_id in connector_output.finished_recving or ():
578 - self._stop_heartbeat(req_id)
579 -
580 - def request_finished(
581 - self,
582 - request: "Request",
583 - block_ids: BlockIds,
584 - ) -> tuple[bool, dict[str, Any] | None]:
585 - """
586 - Once a request is finished, determine whether request blocks
587 - should be freed now or will be sent asynchronously and freed later.
588 - """
589 - from vllm.v1.request import RequestStatus
590 -
591 - params = request.kv_transfer_params
592 - logger.debug(
593 - "NIXLConnector request_finished(%s), request_status=%s, "
594 - "kv_transfer_params=%s",
595 - request.request_id,
596 - request.status,
597 - params,
598 - )
599 - if not params:
600 - return False, None
601 -
602 - is_p_node = bool(params.get("do_remote_decode"))
603 - is_d_node = not is_p_node
604 -
605 - # Stop heartbeating for aborted requests that never reached finished_recving:
606 - # normal path cleans up in update_connector_output.
607 - self._stop_heartbeat(request.request_id)
608 -
609 - if params.get("do_remote_prefill"):
610 - # If do_remote_prefill is still True when the request is finished,
611 - # update_state_after_alloc must not have been called (the request
612 - # must have been aborted before it was scheduled, e.g. via the
613 - # abort_immediately path used to clean up KV-transfer requests
614 - # rejected at the D-side serving layer).
615 - # To avoid stranding the prefill blocks in the prefill instance,
616 - # we must add empty block_ids to _reqs_need_recv so that our
617 - # worker side will notify and free blocks in the prefill instance.
618 - self._reqs_need_recv[request.request_id] = (request, [])
619 - params["do_remote_prefill"] = False
620 - return False, None
621 -
622 - if is_d_node and not self.is_bidirectional_kv_xfer_enabled:
623 - return False, None
624 -
625 - if request.status not in (
626 - RequestStatus.FINISHED_LENGTH_CAPPED,
627 - RequestStatus.FINISHED_STOPPED,
628 - ):
629 - # Also include the case of a P/D Prefill request with immediate
630 - # block free (eg abort). Stop tracking this request.
631 - self._reqs_not_processed.add(request.request_id)
632 - # Clear _reqs_need_save if a request is aborted as partial prefill.
633 - self._reqs_need_save.pop(request.request_id, None)
634 - return False, None
635 -
636 - # TODO: check whether block_ids actually ever be 0. If not we could
637 - # remove the conditional below
638 - delay_free_blocks = any(len(group) > 0 for group in block_ids)
639 - remote_num_tokens = 0
640 - if delay_free_blocks:
641 - # Prefill request on remote. It will be read from D upon completion
642 - request_kv_blocks_ttl = self._kv_lease_duration
643 - if is_d_node:
644 - # For blocks pinned on D, use a simpler timeout for now instead of a
645 - # lease mechanism as turn2 request is client-driven.
646 - request_kv_blocks_ttl = self.decoder_kv_blocks_ttl
647 - logger.debug(
648 - "NIXLConnector request_finished(%s) waiting for %d seconds "
649 - "before releasing blocks",
650 - request.request_id,
651 - request_kv_blocks_ttl,
652 - )
653 - self._reqs_need_send[request.request_id] = (
654 - time.perf_counter() + request_kv_blocks_ttl
655 - )
656 - # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
657 - # trimming down after allocating for the whole sequence length. Empty
658 - # blocks are always at the start of the list.
659 - # Here we "unpad" blocks to send the actual remote blocks to be read.
660 - block_ids = self.get_sw_clipped_blocks(block_ids)
8
662 - remote_num_tokens = request.num_computed_tokens
9 +# Backward compatibility: NixlConnectorScheduler is the pull-based scheduler.
10 +NixlConnectorScheduler = NixlPullConnectorScheduler
11
664 - return delay_free_blocks, dict(
665 - do_remote_prefill=is_p_node,
666 - do_remote_decode=is_d_node,
667 - remote_block_ids=block_ids,
668 - remote_engine_id=self.engine_id,
669 - remote_request_id=request.request_id,
670 - remote_host=self.side_channel_host,
671 - remote_port=self.side_channel_port,
672 - tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
673 - remote_num_tokens=remote_num_tokens,
674 - )
12 +__all__ = ["NixlConnectorScheduler", "NixlPullConnectorScheduler"]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py
+11
@@ -6,6 +6,7 @@ import contextlib
6 from collections.abc import Iterator
7 from typing import Any
8
9 +import regex as re
10 import zmq
11
12 from vllm.platforms import current_platform
@@ -55,3 +56,13 @@ def get_representative_spec_type(spec: KVCacheSpec) -> type[KVCacheSpec]:
56 inner = next(iter(spec.kv_cache_specs.values()))
57 return type(inner)
58 return type(spec)
59 +
60 +
61 +# Trailing 8-hex randomization suffix appended by
62 +# ``input_processor.assign_request_id`` as ``-{random_uuid():.8}``.
63 +_RANDOM_SUFFIX_RE = re.compile(r"-[0-9a-f]{8}$", re.IGNORECASE)
64 +
65 +
66 +def get_base_request_id(request_id: str) -> str:
67 + """Strip the per-request ``-<8 hex>`` randomization suffix, if present."""
68 + return _RANDOM_SUFFIX_RE.sub("", request_id)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+6 -2634
@@ -1,2641 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Worker-side logic for the NIXL connector."""
3 +"""Backward-compatible re-export of NixlPullConnectorWorker."""
4
5 -import logging
6 -import os
7 -import queue
8 -import threading
9 -import time
10 -import uuid
11 -from collections import defaultdict
12 -from collections.abc import Iterator
13 -from concurrent.futures import Future, ThreadPoolExecutor
14 -from typing import TYPE_CHECKING, Any, cast
15 -
16 -import msgspec
17 -import numpy as np
18 -import torch
19 -import zmq
20 -
21 -from vllm.distributed.kv_transfer.kv_connector.utils import (
22 - BlockIds,
23 - EngineId,
24 - EngineTransferInfo,
25 - TransferTopology,
26 - get_current_attn_backends,
27 - kv_postprocess_blksize_and_layout_on_receive,
28 - kv_postprocess_blksize_on_receive,
29 - kv_postprocess_layout_on_receive,
30 -)
31 -from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp
32 -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
33 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
34 - GET_META_MSG,
35 - NixlAgentMetadata,
36 - NixlConnectorMetadata,
37 - NixlHandshakePayload,
38 - ReqId,
39 - ReqMeta,
40 - TransferHandle,
41 - compute_nixl_compatibility_hash,
42 -)
43 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
44 - NixlKVConnectorStats,
45 -)
46 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
47 - ReadSpec,
48 - TPMapping,
49 - _is_attention_spec,
50 - _is_ssm_spec,
51 - compute_tp_mapping,
52 -)
53 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
54 - _NIXL_SUPPORTED_DEVICE,
55 - get_representative_spec_type,
56 - zmq_ctx,
57 -)
58 -from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
59 - MambaConvSplitInfo,
60 - derive_mamba_conv_split,
61 -)
62 -from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
63 -from vllm.distributed.parallel_state import (
64 - get_tensor_model_parallel_rank,
65 - get_tensor_model_parallel_world_size,
5 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import (
6 + NixlPullConnectorWorker,
7 )
67 -from vllm.logger import init_logger
68 -from vllm.platforms import current_platform
69 -from vllm.utils.network_utils import make_zmq_path
70 -from vllm.v1.attention.backends.utils import get_kv_cache_layout
71 -from vllm.v1.kv_cache_interface import (
72 - FullAttentionSpec,
73 - MambaSpec,
74 - MLAAttentionSpec,
75 - UniformTypeKVCacheSpecs,
76 -)
77 -from vllm.v1.worker.block_table import BlockTable
78 -from vllm.v1.worker.utils import select_common_block_size
79 -
80 -if TYPE_CHECKING:
81 - from vllm.config import VllmConfig
82 - from vllm.v1.kv_cache_interface import KVCacheConfig
83 -
84 -logger = init_logger(__name__)
85 -
86 -
87 -class NixlConnectorWorker:
88 - """Implementation of Worker side methods"""
89 -
90 - def _compute_desc_ids(
91 - self,
92 - block_ids: BlockIds,
93 - dst_num_blocks: int,
94 - block_size_ratio: float | None,
95 - physical_blocks_per_logical: int,
96 - ) -> np.ndarray:
97 - """Compute NIXL descriptor IDs for given block IDs."""
98 - num_fa_regions = self.num_regions
99 - num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0
100 -
101 - num_blocks = dst_num_blocks
102 - if block_size_ratio is not None:
103 - num_blocks = int(num_blocks * block_size_ratio)
104 - num_fa_descs = num_fa_regions * num_blocks
105 -
106 - # All-attention fast path: single vectorized broadcast.
107 - if num_ssm_regions == 0:
108 - # NOTE (NickLucche) With HMA, every kv group has the same number of layers
109 - # and layers from different groups share the same kv tensor.
110 - # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be
111 - # read across all regions, same for [3], but group0-group1 blocks will
112 - # always differ (different areas). Therefore we can just flatten the
113 - # block_ids and compute the descs ids for all groups at once.
114 - block_arr = np.concatenate(block_ids)[None, :]
115 - region_ids = np.arange(num_fa_regions)[:, None]
116 - return (region_ids * num_blocks + block_arr).flatten()
117 -
118 - # Compute desc ids per group using the right stride: FA descs have
119 - # num_blocks entries per region (kernel granularity), SSM descs have
120 - # logical_blocks entries per region (no kernel splitting).
121 - logical_blocks = num_blocks // physical_blocks_per_logical
122 - all_descs: list[np.ndarray] = []
123 - for i, group in enumerate(block_ids):
124 - group_arr = np.asarray(group)
125 - if _is_attention_spec(self._group_spec_types[i]):
126 - fa_region_ids = np.arange(num_fa_regions)[:, None]
127 - all_descs.append(
128 - (fa_region_ids * num_blocks + group_arr[None, :]).flatten()
129 - )
130 - elif _is_ssm_spec(self._group_spec_types[i]):
131 - # NOTE (NickLucche) SSM and Attention block regions can
132 - # be exchanged arbitrarily by manager. Therefore, descs
133 - # are laid out as:
134 - # [descs_fa (all regions) | descs_ssm (all regions)].
135 - # num_fa_descs offset must be computed per-engine since
136 - # P and D can have different num_blocks (and thus
137 - # different FA desc counts).
138 - ssm_region_ids = np.arange(num_ssm_regions)[:, None]
139 - all_descs.append(
140 - (
141 - ssm_region_ids * logical_blocks
142 - + group_arr[None, :]
143 - + num_fa_descs
144 - ).flatten()
145 - )
146 - else:
147 - raise ValueError(
148 - f"Unknown spec type {self._group_spec_types[i]} at index {i}"
149 - )
150 -
151 - return np.concatenate(all_descs)
152 -
153 - def _build_local_splits_from_plan(
154 - self,
155 - plan: TPMapping,
156 - src_blocks_data: list[tuple[int, int, int]],
157 - num_fa_descs: int,
158 - ) -> Iterator[list[tuple[int, int, int]]]:
159 - """Build split handle data for P_TP > D_TP scenario.
160 -
161 - num_fa_descs is the boundary between FA and SSM descriptors.
162 - Split counts are derived from source_ranks_per_group lengths.
163 - FA uses rank_to_attention_slot for the slot offset;
164 - SSM uses the rank's positional index.
165 - """
166 - fa_idx = next(
167 - i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
168 - )
169 - fa_num_splits = len(plan.source_ranks_per_group[fa_idx])
170 -
171 - has_ssm_descs = num_fa_descs < len(src_blocks_data)
172 - ssm_idx = next(
173 - (i for i, t in enumerate(self._group_spec_types) if _is_ssm_spec(t)),
174 - None,
175 - )
176 - ssm_num_splits = (
177 - len(plan.source_ranks_per_group[ssm_idx])
178 - if has_ssm_descs and ssm_idx is not None
179 - else 0
180 - )
181 -
182 - # Per-FA-descriptor replicate flag, in _build_fa_local emission order.
183 - fa_desc_replicated = self._fa_desc_replicated(num_fa_descs)
184 -
185 - for p_idx, p_rank in enumerate(plan.all_source_ranks):
186 - fa_slot = plan.rank_to_attention_slot.get(p_rank, 0)
187 -
188 - handle: list[tuple[int, int, int]] = []
189 - for j, (addr, local_len, dev) in enumerate(src_blocks_data):
190 - if j < num_fa_descs:
191 - if fa_desc_replicated[j]:
192 - # REPLICATE (MLA): whole block written on every rank.
193 - handle.append((addr, local_len, dev))
194 - else:
195 - # SPLIT (full-attn): this rank's head slice.
196 - chunk = local_len // fa_num_splits
197 - handle.append((addr + fa_slot * chunk, chunk, dev))
198 - else:
199 - chunk = local_len // ssm_num_splits
200 - handle.append((addr + p_idx * chunk, chunk, dev))
201 - yield handle
202 -
203 - def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]:
204 - """Per-FA-descriptor replicate flag, in _build_fa_local emission order
205 - (region-major; K then optional V per region). Length ``num_fa_descs``.
206 - """
207 - assert self.transfer_topo is not None
208 - n_regions = len(self.block_len_per_layer)
209 - # Unset only when the worker is built directly in unit tests; a real
210 - # model always registers regions (no-KV-cache crashes long before here).
211 - # Fall back to all-SPLIT to preserve the pre-per-region behavior.
212 - if n_regions == 0 or self.num_regions == 0:
213 - return [False] * num_fa_descs
214 - # Descriptors (blocks) per stream; all streams share the same count.
215 - nblk = num_fa_descs // self.num_regions
216 - virtually_split = self.transfer_topo.virtually_split_kv_in_blocks
217 - flags: list[bool] = []
218 - for i in range(n_regions):
219 - replicated = self._is_region_replicated(i)
220 - # REPLICATE (MLA) is key-only -> 1 stream; SPLIT emits K and V
221 - # (2 streams) under the virtually-split layout.
222 - num_streams = 1 if replicated or not virtually_split else 2
223 - flags.extend([replicated] * (num_streams * nblk))
224 - assert len(flags) == num_fa_descs, (
225 - f"FA desc flags {len(flags)} != num_fa_descs {num_fa_descs}"
226 - )
227 - return flags
228 -
229 - def _is_region_replicated(self, region_idx: int) -> bool:
230 - """Whether region ``region_idx`` is transferred REPLICATE vs SPLIT.
231 -
232 - REPLICATE (MLA): identical on every rank, whole block read from one
233 - rank at offset 0, key-only. SPLIT (full-attn): head-sharded across TP.
234 - Defaults to SPLIT when the per-region map is unset (e.g. tests that set
235 - block_len_per_layer without register_kv_caches).
236 - """
237 - return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx]
238 -
239 - def __init__(
240 - self,
241 - vllm_config: "VllmConfig",
242 - engine_id: str,
243 - kv_cache_config: "KVCacheConfig",
244 - ):
245 - nixl_wrapper_cls = NixlWrapper
246 - if nixl_wrapper_cls is None:
247 - logger.error("NIXL is not available")
248 - raise RuntimeError("NIXL is not available")
249 - logger.info("Initializing NIXL wrapper")
250 - logger.info("Initializing NIXL worker %s", engine_id)
251 -
252 - # Config.
253 - self.vllm_config = vllm_config
254 - # mypy will complain on re-assignment otherwise.
255 - self.block_size: int = cast(int, vllm_config.cache_config.block_size)
256 -
257 - if vllm_config.kv_transfer_config is None:
258 - raise ValueError("kv_transfer_config must be set for NixlConnector")
259 - self.kv_transfer_config = vllm_config.kv_transfer_config
260 -
261 - self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config(
262 - "backends", ["UCX"]
263 - )
264 - kv_lease_duration: int = vllm_config.kv_transfer_config.get_from_extra_config(
265 - "kv_lease_duration", 30
266 - )
267 - # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.
268 - self._lease_extension = kv_lease_duration * 2 // 3
269 -
270 - self._is_hma_required = (
271 - not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
272 - and any(
273 - not isinstance(g.kv_cache_spec, FullAttentionSpec)
274 - for g in kv_cache_config.kv_cache_groups
275 - )
276 - )
277 - self.kv_cache_config = kv_cache_config
278 - self._layer_specs = {
279 - layer: group.kv_cache_spec
280 - for group in kv_cache_config.kv_cache_groups
281 - for layer in group.layer_names
282 - }
283 - self.hma_group_size = len(kv_cache_config.kv_cache_tensors)
284 -
285 - # ---- Model state (derived from model config) ----
286 - mamba_ssm_size = (0, 0)
287 - # Conv state sub-projection decomposition (None when no Mamba).
288 - # The 3-read transfer requires DS (dim, state_len) conv layout so
289 - # that x/B/C sub-projections are contiguous in memory.
290 - self._conv_decomp: MambaConvSplitInfo | None = None
291 - self._has_mamba = any(
292 - isinstance(g.kv_cache_spec, MambaSpec)
293 - for g in kv_cache_config.kv_cache_groups
294 - )
295 - if self._has_mamba:
296 - assert self._is_hma_required
297 - from vllm.model_executor.layers.mamba.mamba_utils import (
298 - is_conv_state_dim_first,
299 - )
300 -
301 - assert is_conv_state_dim_first(), (
302 - "3-read Mamba conv transfer requires DS conv state layout. "
303 - "Set VLLM_SSM_CONV_STATE_LAYOUT=DS"
304 - )
305 - mamba_spec = next(
306 - spec
307 - for spec in self._layer_specs.values()
308 - if isinstance(spec, MambaSpec)
309 - )
310 - self._conv_decomp = derive_mamba_conv_split(
311 - mamba_spec,
312 - vllm_config.parallel_config.tensor_parallel_size,
313 - )
314 - mamba_ssm_size = self._conv_decomp.ssm_sizes
315 - self._mamba_ssm_size = mamba_ssm_size
316 -
317 - # Agent.
318 - non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
319 - # Configure NIXL num_threads to avoid UAR exhaustion on Mellanox NICs.
320 - # Each UCX thread allocates UARs (doorbell pages) via DevX, and
321 - # excessive NIXL UAR usage can exhaust NIC UAR space. This can cause
322 - # components like NVSHMEM (used by DeepEP kernels) to fail during RDMA
323 - # initialization with "mlx5dv_devx_alloc_uar" errors.
324 - # Ref: https://network.nvidia.com/files/doc-2020/ethernet-adapters-programming-manual.pdf#page=63
325 - num_threads = vllm_config.kv_transfer_config.get_from_extra_config(
326 - "num_threads", 4
327 - )
328 - if nixl_agent_config is None:
329 - config = None
330 - else:
331 - # Enable telemetry by default for NIXL 0.7.1 and above.
332 - config = (
333 - nixl_agent_config(backends=self.nixl_backends, capture_telemetry=True)
334 - if len(non_ucx_backends) > 0
335 - else nixl_agent_config(num_threads=num_threads, capture_telemetry=True)
336 - )
337 -
338 - self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config)
339 - # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.
340 - self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict)
341 -
342 - # Metadata.
343 - self.engine_id: EngineId = engine_id
344 - self.tp_rank = get_tensor_model_parallel_rank()
345 - self.world_size = get_tensor_model_parallel_world_size()
346 -
347 - self.num_blocks = kv_cache_config.num_blocks
348 - self.enable_permute_local_kv = False
349 - self.enable_heterogeneous_attn_post_process = False
350 -
351 - # KV Caches and nixl tracking data.
352 - self.device_type = current_platform.device_type
353 - self.kv_buffer_device: str = vllm_config.kv_transfer_config.kv_buffer_device
354 - if self.device_type not in _NIXL_SUPPORTED_DEVICE:
355 - raise RuntimeError(f"{self.device_type} is not supported.")
356 - elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[self.device_type]:
357 - raise RuntimeError(
358 - f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
359 - "is not supported."
360 - )
361 - self.device_kv_caches: dict[str, torch.Tensor] = {}
362 -
363 - # cpu kv buffer for xfer
364 - # used when device memory can not be registered under nixl
365 - self.host_xfer_buffers: dict[str, torch.Tensor] = {}
366 - if self.device_type == "cpu":
367 - self.use_host_buffer = False
368 - else:
369 - self.use_host_buffer = self.kv_buffer_device == "cpu"
370 -
371 - # reserve different cores for start_load_kv() from model_forward()
372 - if self.device_type == "cpu":
373 - numa_core_list = current_platform.discover_numa_topology()
374 - # setup one last core in each numa for kv transfer.
375 - rsv_cores_for_kv = [
376 - max(each_numa_core_list) for each_numa_core_list in numa_core_list
377 - ]
378 -
379 - if rsv_cores_for_kv:
380 - if not hasattr(os, "sched_setaffinity"):
381 - raise NotImplementedError(
382 - "os.sched_setaffinity is not available on this platform"
383 - )
384 - os.sched_setaffinity(0, rsv_cores_for_kv)
385 -
386 - # support for oot platform which can't register nixl memory
387 - # type based on kv_buffer_device
388 - nixl_memory_type = current_platform.get_nixl_memory_type()
389 - if nixl_memory_type is None:
390 - if self.kv_buffer_device in ["cuda", "xpu"]:
391 - nixl_memory_type = "VRAM"
392 - elif self.kv_buffer_device == "cpu":
393 - nixl_memory_type = "DRAM"
394 - if nixl_memory_type is None:
395 - raise RuntimeError(
396 - f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
397 - "is not supported."
398 - )
399 - self.nixl_memory_type = nixl_memory_type
400 -
401 - # Note: host xfer buffer ops when use_host_buffer is True
402 - self.copy_blocks: CopyBlocksOp | None = None
403 -
404 - # Map of engine_id -> kv_caches_base_addr. For TP case, each local
405 - self.device_id: int = 0
406 - # Current rank may pull from multiple remote TP workers.
407 - # EngineId, dict[int, list[int]] -> engine_id, tp_rank, base_addr_for_layer
408 - self.kv_caches_base_addr = defaultdict[EngineId, dict[int, list[int]]](dict)
409 -
410 - # Number of NIXL regions. Currently one region per cache
411 - # (so 1 per layer for MLA, otherwise 2 per layer)
412 - self.num_regions = 0
413 -
414 - # nixl_prepped_dlist_handle.
415 - self.src_xfer_handles_by_block_size: dict[int, int] = {}
416 - # Populated dynamically during handshake based on remote configuration.
417 - # Keep track of regions at different tp_ratio values. tp_ratio->handles
418 - self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {}
419 - # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}.
420 - self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict)
421 -
422 - # Map of engine_id -> num_blocks. All ranks in the same deployment will
423 - # have the same number of blocks.
424 - self.dst_num_blocks: dict[EngineId, int] = {}
425 - self._registered_descs: list[Any] = []
426 -
427 - # In progress transfers.
428 - # [req_id -> list[handle]]
429 - self._recving_metadata: dict[ReqId, ReqMeta] = {}
430 - self._recving_transfers = defaultdict[ReqId, list[TransferHandle]](list)
431 - # Track the expiration time of requests that are waiting to be sent.
432 - self._reqs_to_send: dict[ReqId, float] = {}
433 - # Set of requests that have been part of a batch, regardless of status.
434 - self._reqs_to_process: set[ReqId] = set()
435 -
436 - # Invalid blocks from failed NIXL operations (thread-safe queue of block ids)
437 - self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue()
438 - # requests that skipped transfer (handshake or transfer failures)
439 - # Uses Queue for thread-safe cross-thread coordination with the
440 - # background handshake thread, matching the _ready_requests pattern.
441 - self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
442 -
443 - # Handshake metadata of this worker for NIXL transfers.
444 - self.xfer_handshake_metadata: NixlHandshakePayload | None = None
445 - # Background thread for initializing new NIXL handshakes.
446 - self._handshake_initiation_executor = ThreadPoolExecutor(
447 - # NIXL is not guaranteed to be thread-safe, limit 1 worker.
448 - max_workers=1,
449 - thread_name_prefix="vllm-nixl-handshake-initiator",
450 - )
451 - self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]()
452 - self._handshake_futures: dict[EngineId, Future[dict[int, str]]] = {}
453 - # Protects _handshake_futures and _remote_agents.
454 - self._handshake_lock = threading.RLock()
455 -
456 - # TTL-based eviction of stale remote engine state.
457 - self._engine_last_active: dict[EngineId, float] = {}
458 - self._engine_ttl: float = vllm_config.kv_transfer_config.get_from_extra_config(
459 - "engine_ttl", 3600.0
460 - )
461 -
462 - self.block_size = vllm_config.cache_config.block_size
463 - self.model_config = vllm_config.model_config
464 -
465 - self.use_mla = self.model_config.use_mla
466 -
467 - # Get the attention backend from the first layer
468 - # NOTE (NickLucche) models with multiple backends are not supported yet
469 - self.attn_backends = get_current_attn_backends(vllm_config)
470 - self.backend_name = self.attn_backends[0].get_name()
471 -
472 - self.kv_cache_layout = get_kv_cache_layout()
473 - self.host_buffer_kv_cache_layout = self.kv_cache_layout
474 - logger.info(
475 - "Detected attention backend(s) %s",
476 - [backend.get_name() for backend in self.attn_backends],
477 - )
478 - logger.info("Detected kv cache layout %s", self.kv_cache_layout)
479 -
480 - # lazy initialized in register_kv_caches
481 - self.compat_hash: str | None = None
482 - self.transfer_topo: TransferTopology | None = None
483 -
484 - # With heterogeneous TP, P must wait for all assigned D TP workers to
485 - # finish reading before safely freeing the blocks.
486 - self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
487 - self.xfer_stats = NixlKVConnectorStats()
488 -
489 - self._physical_blocks_per_logical_kv_block = 1
490 - self._sync_block_size_with_kernel()
491 -
492 - # Unwrap UniformTypeKVCacheSpecs to get the representative spec type
493 - self._group_spec_types = tuple(
494 - get_representative_spec_type(g.kv_cache_spec)
495 - for g in self.kv_cache_config.kv_cache_groups
496 - )
497 -
498 - # Per-region MLA flag, 1:1 with block_len_per_layer. True -> REPLICATE
499 - # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models
500 - # combining both (e.g. GQA main + MLA Eagle-3 draft).
501 - self._region_is_mla = list[bool]()
502 -
503 - # Enable different block lengths for different layers *only* when MLA is used.
504 - # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.
505 - self.block_len_per_layer = list[int]()
506 -
507 - # Per-engine TP mappings. Generated during handshake.
508 - self.tp_mappings: dict[EngineId, TPMapping] = {}
509 -
510 - self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
511 - "enforce_handshake_compat", True
512 - )
513 -
514 - def _sync_block_size_with_kernel(self) -> None:
515 - backends = get_current_attn_backends(self.vllm_config)
516 - kernel_block_size = select_common_block_size(self.block_size, backends)
517 - # Number of blocks not accounting for kernel block mismatches
518 - self._logical_num_blocks = self.num_blocks
519 - if self.block_size != kernel_block_size:
520 - logger.info_once(
521 - "User-specified logical block size (%s) does not match"
522 - " physical kernel block size (%s). Using the latter.",
523 - self.block_size,
524 - kernel_block_size,
525 - )
526 - assert self.block_size > kernel_block_size
527 - self._physical_blocks_per_logical_kv_block = (
528 - self.block_size // kernel_block_size
529 - )
530 - self.block_size = kernel_block_size
531 - self.num_blocks *= self._physical_blocks_per_logical_kv_block
532 -
533 - def _nixl_handshake(
534 - self,
535 - host: str,
536 - port: int,
537 - remote_tp_size: int,
538 - expected_engine_id: str,
539 - ) -> dict[int, str]:
540 - """Do a NIXL handshake with a remote instance."""
541 -
542 - # the first time we connect to a remote agent.
543 - # be careful, the handshake happens in a background thread.
544 - # it does not have an active cuda context until any cuda runtime
545 - # call is made. when UCX fails to find a valid cuda context, it will
546 - # disable any cuda ipc communication, essentially disabling any NVLink
547 - # communication.
548 - # when we are using device buffers, we need to set the device
549 - # explicitly to make sure the handshake background thread has a valid
550 - # cuda context.
551 - if not self.use_host_buffer:
552 - current_platform.set_device(self.device_id)
553 -
554 - # When target instance TP > local TP, we need to perform multiple
555 - # handshakes. Do it in a single background job for simplicity.
556 - # Regardless, only handshake with the remote TP rank(s) that current
557 - # local rank will read from. Note that With homogeneous TP,
558 - # this happens to be the same single rank_i.
559 - assert self.transfer_topo is not None
560 - p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size)
561 - remote_rank_to_agent_name = {}
562 - path = make_zmq_path("tcp", host, port)
563 -
564 - with zmq_ctx(zmq.REQ, path) as sock:
565 - for remote_rank in p_remote_ranks:
566 - logger.debug(
567 - "Querying metadata on path: %s at remote tp rank %s",
568 - path,
569 - remote_rank,
570 - )
571 -
572 - start_time = time.perf_counter()
573 - # Send query for the request.
574 - msg = msgspec.msgpack.encode((GET_META_MSG, remote_rank))
575 - # Set receive timeout to 5 seconds to avoid hanging on dead server
576 - sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds
577 - sock.send(msg)
578 - handshake_bytes = sock.recv()
579 -
580 - # Decode handshake payload to get compatibility hash
581 - handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload)
582 - try:
583 - handshake_payload = handshake_decoder.decode(handshake_bytes)
584 - except (msgspec.DecodeError, msgspec.ValidationError) as e:
585 - raise RuntimeError(
586 - f"Failed to decode NixlHandshakePayload. This likely indicates "
587 - f"an incompatibility between connector version. Error: {e}"
588 - ) from e
589 -
590 - got_metadata_time = time.perf_counter()
591 - logger.debug(
592 - "NIXL handshake: get metadata took: %s",
593 - got_metadata_time - start_time,
594 - )
595 -
596 - # Check compatibility hash BEFORE decoding agent metadata
597 - assert self.compat_hash is not None
598 - if (
599 - self.enforce_compat_hash
600 - and handshake_payload.compatibility_hash != self.compat_hash
601 - ):
602 - raise RuntimeError(
603 - f"NIXL compatibility hash mismatch. "
604 - f"Local: {self.compat_hash}, "
605 - f"Remote: {handshake_payload.compatibility_hash}. "
606 - f"Prefill and decode instances have incompatible "
607 - f"configurations. This may be due to: different vLLM versions,"
608 - f" models, dtypes, KV cache layouts, attention backends, etc. "
609 - f"Both instances must use identical configurations."
610 - f"Disable this check using "
611 - f'--kv-transfer-config \'{{"kv_connector_extra_config": '
612 - f'{{"enforce_handshake_compat": false}}}}\''
613 - )
614 -
615 - logger.info(
616 - "NIXL compatibility check passed (hash: %s)",
617 - handshake_payload.compatibility_hash,
618 - )
619 -
620 - # Decode agent metadata
621 - metadata_decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
622 - try:
623 - metadata = metadata_decoder.decode(
624 - handshake_payload.agent_metadata_bytes
625 - )
626 - except (msgspec.DecodeError, msgspec.ValidationError) as e:
627 - # This should not happen if hash matched
628 - raise RuntimeError(
629 - f"Failed to decode NixlAgentMetadata. Error: {e}"
630 - ) from e
631 -
632 - # Ensure engine id matches.
633 - if metadata.engine_id != expected_engine_id:
634 - raise RuntimeError(
635 - f"Remote NIXL agent engine ID mismatch. "
636 - f"Expected {expected_engine_id},"
637 - f"received {metadata.engine_id}."
638 - )
639 -
640 - # Register Remote agent.
641 - remote_agent_name = self.add_remote_agent(
642 - metadata, remote_rank, remote_tp_size
643 - )
644 - setup_agent_time = time.perf_counter()
645 - logger.debug(
646 - "NIXL handshake: add agent took: %s",
647 - setup_agent_time - got_metadata_time,
648 - )
649 - remote_rank_to_agent_name[remote_rank] = remote_agent_name
650 - return remote_rank_to_agent_name
651 -
652 - def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> None:
653 - """
654 - Initialize transfer buffer in CPU mem for accelerators
655 - NOT directly supported by NIXL (e.g., tpu)
656 - """
657 - xfer_buffers: dict[str, torch.Tensor] = {}
658 - inv_order = [0, 1, 3, 2, 4]
659 - try:
660 - for layer_name, kv_cache in kv_caches.items():
661 - kv_shape = kv_cache.shape
662 - kv_dtype = kv_cache.dtype
663 - permute_shape = False
664 - if (
665 - self.kv_cache_layout == "NHD"
666 - and self.vllm_config.kv_transfer_config is not None
667 - and self.vllm_config.kv_transfer_config.enable_permute_local_kv
668 - ):
669 - logger.info_once(
670 - "'enable_permute_local_kv' flag is enabled while "
671 - "device KV Layout is NHD. Init host buffer with"
672 - " HND to better support Decode/Prefill TP_ratio > 1."
673 - )
674 - # Since NHD will not support Decode/Prefill TP_ratio > 1,
675 - # we can leverage host_buffer for permute
676 - self.host_buffer_kv_cache_layout = "HND"
677 - kv_shape = (
678 - tuple(kv_shape[i] for i in inv_order)
679 - if not self.use_mla
680 - else kv_shape
681 - )
682 - permute_shape = not self.use_mla
683 -
684 - xfer_buffers[layer_name] = torch.empty(
685 - kv_shape, dtype=kv_dtype, device="cpu"
686 - )
687 - if permute_shape:
688 - xfer_buffers[layer_name] = xfer_buffers[layer_name].permute(
689 - inv_order
690 - )
691 - except MemoryError as e:
692 - logger.error("NIXLConnectorWorker gets %s.", e)
693 - raise
694 -
695 - self.host_xfer_buffers = xfer_buffers
696 -
697 - def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
698 - """Assign copy (d2h, h2d) operations when host buffer is used."""
699 - # Set a no-op if the host buffer is not cpu.
700 - if self.kv_buffer_device != "cpu":
701 - return
702 - # Set a no-op if self.device_type is 'cpu'.
703 - if self.device_type == "cpu":
704 - return
705 - assert self.use_host_buffer
706 - self.copy_blocks = copy_operation
707 -
708 - def _log_failure(
709 - self,
710 - failure_type: str,
711 - req_id: str | None,
712 - msg: str = "",
713 - error: Exception | None = None,
714 - meta: ReqMeta | None = None,
715 - **extra_context,
716 - ):
717 - """Log transfer failure with structured context for easier debugging."""
718 - context: dict[str, Any] = {
719 - "failure_type": failure_type,
720 - "request_id": req_id,
721 - "engine_id": self.engine_id,
722 - }
723 - if meta is None and req_id is not None:
724 - # Try to get metadata from in progress transfers when not provided
725 - meta = self._recving_metadata.get(req_id)
726 -
727 - if meta and meta.remote:
728 - context.update(
729 - {
730 - "remote_engine_id": meta.remote.engine_id,
731 - "remote_request_id": meta.remote.request_id,
732 - "remote_host": meta.remote.host,
733 - "remote_port": meta.remote.port,
734 - "num_local_blocks": sum(
735 - len(group) for group in meta.local_block_ids
736 - ),
737 - "num_remote_blocks": sum(
738 - len(group) for group in meta.remote.block_ids
739 - ),
740 - "local_block_ids_sample": meta.local_block_ids[0][:10]
741 - if meta.local_block_ids
742 - else [],
743 - }
744 - )
745 -
746 - context.update(extra_context)
747 - if msg:
748 - failure_type = f"{failure_type}. {msg}"
749 -
750 - logger.error(
751 - "NIXL transfer failure: %s | Context: %s",
752 - failure_type,
753 - context,
754 - exc_info=error is not None,
755 - stacklevel=2,
756 - )
757 -
758 - def _ensure_handshake(
759 - self,
760 - engine_id: EngineId,
761 - host: str,
762 - port: int,
763 - tp_size: int,
764 - ) -> Future[dict[int, str]] | None:
765 - """
766 - Ensure a handshake is in-flight (or already done) for *engine_id*.
767 -
768 - Returns the ``Future`` if a handshake is pending (or was just
769 - started), or ``None`` if the handshake already completed
770 - successfully. Callers can attach per-request callbacks to the
771 - returned future.
772 - Failures to handshake are logged and the request is marked as failed.
773 - """
774 - self._evict_stale_engines()
775 - with self._handshake_lock:
776 - if engine_id in self._remote_agents:
777 - return None
778 - fut = self._handshake_futures.get(engine_id)
779 - if fut is not None:
780 - return fut
781 - fut = self._handshake_initiation_executor.submit(
782 - self._nixl_handshake,
783 - host,
784 - port,
785 - tp_size,
786 - engine_id,
787 - )
788 - self._handshake_futures[engine_id] = fut
789 -
790 - def done_callback(f: Future[dict[int, str]], eid=engine_id):
791 - with self._handshake_lock:
792 - del self._handshake_futures[eid]
793 - try:
794 - self._remote_agents[eid] = f.result()
795 - self._engine_last_active[eid] = time.perf_counter()
796 - except Exception as e:
797 - self._log_failure(
798 - failure_type="handshake_setup_failed",
799 - req_id=None,
800 - error=e,
801 - remote_engine_id=eid,
802 - )
803 -
804 - fut.add_done_callback(done_callback)
805 - return fut
806 -
807 - def _background_nixl_handshake(
808 - self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta
809 - ):
810 - # Do NIXL handshake in background and add to _ready_requests when done.
811 - assert meta.remote is not None
812 - fut = self._ensure_handshake(
813 - remote_engine_id,
814 - meta.remote.host,
815 - meta.remote.port,
816 - meta.tp_size,
817 - )
818 - if fut is None:
819 - # Already handshaked — only happens if caller does not pre-check.
820 - self._ready_requests.put((req_id, meta))
821 - return
822 -
823 - # Check handshake success before proceeding with request.
824 - def request_ready(f: Future[Any], entry=(req_id, meta)):
825 - try:
826 - f.result()
827 - self._ready_requests.put(entry)
828 - except Exception as e:
829 - self._log_failure(
830 - failure_type="handshake_failed",
831 - req_id=req_id,
832 - error=e,
833 - meta=meta,
834 - )
835 - self._handle_failed_transfer(req_id, None)
836 -
837 - fut.add_done_callback(request_ready)
838 -
839 - def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
840 - """Register a cross-layers KV cache tensor with NIXL.
841 -
842 - `use_uniform_kv_cache()` guarantees a single KV cache group whose
843 - layers all share the same `AttentionSpec`, so any layer name from
844 - `_layer_specs` yields the correct per-layer spec for `page_size_bytes`.
845 - """
846 - first_layer = next(iter(self._layer_specs))
847 - # Forwarding a real layer name rather than a synthetic key
848 - self.register_kv_caches({first_layer: kv_cache})
849 -
850 - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
851 - """Register the KV Cache data in nixl."""
852 - self.transfer_topo = TransferTopology(
853 - tp_rank=self.tp_rank,
854 - tp_size=self.world_size,
855 - block_size=self.block_size,
856 - engine_id=self.engine_id,
857 - is_mla=self.use_mla,
858 - total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
859 - attn_backends=self.attn_backends,
860 - # SSM States come in tuples (ssm, conv)
861 - tensor_shape=next(iter(kv_caches.values())).shape
862 - if not self._has_mamba
863 - else None,
864 - is_mamba=self._has_mamba,
865 - )
866 - self.compat_hash = compute_nixl_compatibility_hash(
867 - self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks
868 - )
869 -
870 - if self.use_host_buffer:
871 - self.initialize_host_xfer_buffer(kv_caches=kv_caches)
872 - assert len(self.host_xfer_buffers) == len(kv_caches), (
873 - f"host_buffer: {len(self.host_xfer_buffers)}, "
874 - f"kv_caches: {len(kv_caches)}"
875 - )
876 - xfer_buffers = self.host_xfer_buffers
877 - else:
878 - xfer_buffers = kv_caches
879 - assert not self.host_xfer_buffers, (
880 - "host_xfer_buffer should not be initialized when "
881 - f"kv_buffer_device is {self.kv_buffer_device}"
882 - )
883 -
884 - logger.info(
885 - "Registering KV_Caches. use_mla: %s, kv_buffer_device: %s, "
886 - "use_host_buffer: %s",
887 - self.use_mla,
888 - self.kv_buffer_device,
889 - self.use_host_buffer,
890 - )
891 -
892 - caches_data = []
893 - # With hybrid allocator, layers can share a kv cache tensor
894 - seen_base_addresses = []
895 -
896 - # Note(tms): I modified this from the original region setup code.
897 - # K and V are now in different regions. Advantage is that we can
898 - # elegantly support MLA and any cases where the K and V tensors
899 - # are non-contiguous (it's not locally guaranteed that they will be)
900 - # Disadvantage is that the encoded NixlAgentMetadata is now larger
901 - # (roughly 8KB vs 5KB).
902 - # Conversely for FlashInfer, K and V are registered in the same region
903 - # to better exploit the memory layout (ie num_blocks is the first dim).
904 - tensor_size_bytes = None
905 -
906 - for layer_name, cache_or_caches in xfer_buffers.items():
907 - # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
908 - # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
909 - # However, physical page_size may differ when kernel requires a specific
910 - # block size. This leads to SSM and FA layers having different num_blocks.
911 - # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.
912 - layer_spec = self._layer_specs.get(layer_name)
913 - if layer_spec is None:
914 - logger.debug(
915 - "Skipping layer %s as no KVCache spec is present. "
916 - "This is likely because the layer is sharing its KV cache",
917 - layer_name,
918 - )
919 - continue
920 - if isinstance(layer_spec, UniformTypeKVCacheSpecs):
921 - # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
922 - layer_spec = layer_spec.kv_cache_specs[layer_name]
923 - cache_list = self.transfer_topo.get_transfer_cache_regions(
924 - cache_or_caches, layer_spec
925 - )
926 - # `layer_spec.page_size_bytes` only accounts for logical page_size, that is
927 - # the page_size assuming constant `self._logical_num_blocks`.
928 - physical_page_size = (
929 - layer_spec.page_size_bytes
930 - if isinstance(layer_spec, MambaSpec)
931 - else layer_spec.page_size_bytes
932 - // self._physical_blocks_per_logical_kv_block
933 - )
934 - # For when registering multiple tensors eg K/V in separate regions.
935 - physical_page_size = physical_page_size // len(cache_list)
936 - if self.transfer_topo._cross_layers_blocks:
937 - # When cross-layers blocks are used, multiply by number of layers
938 - physical_page_size = physical_page_size * len(
939 - self.kv_cache_config.kv_cache_tensors
940 - )
941 - num_blocks = (
942 - self._logical_num_blocks
943 - if isinstance(layer_spec, MambaSpec)
944 - else self.num_blocks
945 - )
946 - # `page_size` accounts for physical blocks, st KVCache is always
947 - # [`num_blocks` * `page_size`]
948 - curr_tensor_size_bytes = num_blocks * physical_page_size
949 -
950 - # TODO (NickLucche) we could eventually unify how we handle FA/FI regions,
951 - # registering a single tensor for both K/V and splitting logically like FI.
952 - for cache in cache_list:
953 - base_addr = cache.data_ptr()
954 - if base_addr in seen_base_addresses:
955 - # NOTE (NickLucche) HMA employs memory pooling to share tensors
956 - # across groups. This results in skipping all tensors but the ones
957 - # pointed to by group0. Also, generally we will have more blocks
958 - # per tensor but fewer regions.
959 - logger.debug("Skipping %s because it's already seen", layer_name)
960 - continue
961 - logger.debug(
962 - "Registering layer %s with cache shape: %s", layer_name, cache.shape
963 - )
964 - seen_base_addresses.append(base_addr)
965 - # Only record non-Mamba page sizes.
966 - if isinstance(layer_spec, MambaSpec):
967 - self.block_len_per_layer.append(
968 - physical_page_size // self._physical_blocks_per_logical_kv_block
969 - )
970 - else:
971 - self.block_len_per_layer.append(physical_page_size)
972 - is_mla_region = isinstance(layer_spec, MLAAttentionSpec)
973 - self._region_is_mla.append(is_mla_region)
974 -
975 - # HeteroTP cannot transfer differently-sized regions, so every
976 - # non-MLA region in a group must share one tensor size (this also
977 - # holds for Mamba-like models). The sole exception is the DeepSeek
978 - # MLA indexer, which sits in a UniformTypeKVCacheSpecs group at a
979 - # different size; MLA regions are therefore exempt.
980 - if not is_mla_region:
981 - if tensor_size_bytes is None:
982 - tensor_size_bytes = curr_tensor_size_bytes
983 - assert tensor_size_bytes == curr_tensor_size_bytes, (
984 - "All non-MLA kv cache tensors must have the same size"
985 - )
986 -
987 - if cache.shape[0] != num_blocks:
988 - raise AssertionError(
989 - "All kv cache tensors must have the same number of "
990 - f"blocks; layer={layer_name}, "
991 - f"expected_num_blocks={num_blocks}, "
992 - f"cache_shape={tuple(cache.shape)}, "
993 - f"cache_stride={tuple(cache.stride())}, "
994 - f"layer_spec={type(layer_spec).__name__}, "
995 - f"backend={self.backend_name}, "
996 - "all_backends="
997 - f"{[backend.get_name() for backend in self.attn_backends]}, "
998 - f"kv_cache_layout={self.kv_cache_layout}, "
999 - "blocks_first="
1000 - f"{self.transfer_topo.is_kv_layout_blocks_first}"
1001 - )
1002 -
1003 - # Need to make sure the device ID is non-negative for NIXL,
1004 - # Torch uses -1 to indicate CPU tensors.
1005 - self.device_id = max(cache.get_device(), 0)
1006 - caches_data.append(
1007 - (base_addr, curr_tensor_size_bytes, self.device_id, "")
1008 - )
1009 -
1010 - logger.debug(
1011 - "Different block lengths collected: %s", set(self.block_len_per_layer)
1012 - )
1013 - assert (
1014 - len(self.block_len_per_layer)
1015 - == len(seen_base_addresses)
1016 - == len(self._region_is_mla)
1017 - )
1018 -
1019 - self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
1020 - self.num_regions = len(caches_data)
1021 -
1022 - if self.transfer_topo.virtually_split_kv_in_blocks:
1023 - # NOTE (NickLucche) When FlashInfer is used, memory is registered
1024 - # with joint KV for each block. This minimizes the overhead in
1025 - # registerMem allowing faster descs queries. In order to be able to
1026 - # split on kv_heads dim as required by heterogeneous TP, one must
1027 - # be able to index K/V separately. Hence we double the number
1028 - # of 'virtual' regions here and halve `block_len` below.
1029 - # Similarly for Mamba layers, we register SSM+Conv as a single region and
1030 - # then duplicate it logically to be able to index SSM/Conv separately.
1031 - # Exception: key-only REPLICATE regions (MLA) have no V half, so
1032 - # they contribute a single desc stream and are not doubled.
1033 - self.num_regions = sum(
1034 - 1 if self._is_region_replicated(i) else 2
1035 - for i in range(len(self._region_is_mla))
1036 - )
1037 -
1038 - # Total local FA descriptors (boundary between FA and mamba descs).
1039 - self.num_descs = self.num_regions * self.num_blocks
1040 -
1041 - descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
1042 - logger.debug("Registering descs: %s", caches_data)
1043 - self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
1044 - logger.debug("Done registering descs")
1045 - self._registered_descs.append(descs)
1046 -
1047 - self.device_kv_caches = kv_caches
1048 - self.dst_num_blocks[self.engine_id] = self.num_blocks
1049 -
1050 - if self._has_mamba:
1051 - logger.info(
1052 - "Hybrid SSM registration: num_blocks=%s, "
1053 - "logical_num_blocks=%s, ratio=%s, num_regions=%s, "
1054 - "num_descs=%s, mamba_ssm_size=%s, block_len_per_layer=%s",
1055 - self.num_blocks,
1056 - self._logical_num_blocks,
1057 - self._physical_blocks_per_logical_kv_block,
1058 - self.num_regions,
1059 - self.num_descs,
1060 - self._mamba_ssm_size,
1061 - set(self.block_len_per_layer),
1062 - )
1063 -
1064 - # Register local/src descr for NIXL xfer.
1065 - self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = (
1066 - self.register_local_xfer_handler(self.block_size)
1067 - )
1068 -
1069 - # After KV Caches registered, listen for new connections.
1070 - agent_metadata = NixlAgentMetadata(
1071 - engine_id=self.engine_id,
1072 - agent_metadata=self.nixl_wrapper.get_agent_metadata(),
1073 - device_id=self.device_id,
1074 - kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank],
1075 - num_blocks=self.num_blocks,
1076 - block_lens=self.block_len_per_layer,
1077 - kv_cache_layout=self.kv_cache_layout
1078 - if not self.use_host_buffer
1079 - else self.host_buffer_kv_cache_layout,
1080 - block_size=self.block_size,
1081 - ssm_sizes=self._mamba_ssm_size,
1082 - attn_backend_name=self.backend_name,
1083 - physical_blocks_per_logical_kv_block=(
1084 - self._physical_blocks_per_logical_kv_block
1085 - ),
1086 - )
1087 - # Wrap metadata in payload with hash for defensive decoding
1088 - assert self.compat_hash is not None
1089 - encoder = msgspec.msgpack.Encoder()
1090 - self.xfer_handshake_metadata = NixlHandshakePayload(
1091 - compatibility_hash=self.compat_hash,
1092 - agent_metadata_bytes=encoder.encode(agent_metadata),
1093 - )
1094 -
1095 - def _build_mamba_local(
1096 - self,
1097 - base_addresses: list[int],
1098 - block_size_ratio: int,
1099 - ) -> list[tuple[int, int, int]]:
1100 - """Build 4 desc regions (x, B, C, ssm) per layer for local mamba
1101 - blocks, enabling the 3-read transfer with DS conv layout."""
1102 - assert block_size_ratio == 1, (
1103 - "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
1104 - f"Got block_size_ratio={block_size_ratio}."
1105 - )
1106 - assert self._conv_decomp is not None
1107 - conv_offsets = self._conv_decomp.local_conv_offsets
1108 - conv_size, ssm_size = self._mamba_ssm_size
1109 - num_blocks = self._logical_num_blocks * block_size_ratio
1110 - physical_per_logical = self._physical_blocks_per_logical_kv_block
1111 -
1112 - result: list[tuple[int, int, int]] = []
1113 - for i, base_addr in enumerate(base_addresses):
1114 - # Jump one page_size, but ssm page_size may be bigger when kernel
1115 - # locks block size to a specific value (physical_per_logical scale).
1116 - page_stride = (
1117 - self.block_len_per_layer[i] // block_size_ratio * physical_per_logical
1118 - )
1119 - for off, sz in conv_offsets:
1120 - for blk in range(num_blocks):
1121 - result.append(
1122 - (base_addr + blk * page_stride + off, sz, self.device_id)
1123 - )
1124 - # SSM temporal state follows the conv state.
1125 - for blk in range(num_blocks):
1126 - result.append(
1127 - (
1128 - base_addr + blk * page_stride + conv_size,
1129 - ssm_size,
1130 - self.device_id,
1131 - )
1132 - )
1133 - return result
1134 -
1135 - def _build_mamba_remote(
1136 - self,
1137 - nixl_agent_meta: NixlAgentMetadata,
1138 - tp_ratio: int,
1139 - transfer_info: EngineTransferInfo,
1140 - ) -> list[tuple[int, int, int]]:
1141 - """Build 4 remote desc regions (proj0, proj1, proj2, ssm) per layer
1142 - for the 3-read transfer. For hetero-TP, each D rank reads only its
1143 - sub-projection slice from the P rank."""
1144 - assert self._conv_decomp is not None
1145 - effective_ratio = max(tp_ratio, 1)
1146 - # Mamba conv state is always TP-sharded, even when attention KV
1147 - # is replicated (num_kv_heads < tp_size).
1148 - local_offset = self.tp_rank % effective_ratio
1149 - conv_size_remote = nixl_agent_meta.ssm_sizes[0]
1150 -
1151 - conv_offsets = self._conv_decomp.remote_conv_offsets(local_offset, tp_ratio)
1152 - if tp_ratio >= 1:
1153 - ssm_read_size = self._mamba_ssm_size[1]
1154 - else:
1155 - ssm_read_size = nixl_agent_meta.ssm_sizes[1]
1156 -
1157 - remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical
1158 - num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical
1159 - device_id = nixl_agent_meta.device_id
1160 -
1161 - result: list[tuple[int, int, int]] = []
1162 - # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
1163 - # block lengths vary across layers (e.g. MLA).
1164 - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1165 - page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical
1166 - for off, sz in conv_offsets:
1167 - for blk in range(num_blocks):
1168 - result.append((base_addr + blk * page_stride + off, sz, device_id))
1169 - # SSM temporal state is also TP-sharded on the heads dimension.
1170 - for blk in range(num_blocks):
1171 - ssm_addr = (
1172 - base_addr
1173 - + blk * page_stride
1174 - + conv_size_remote
1175 - + local_offset * ssm_read_size
1176 - )
1177 - result.append((ssm_addr, ssm_read_size, device_id))
1178 - return result
1179 -
1180 - def _build_fa_local(
1181 - self,
1182 - base_addresses: list[int],
1183 - block_size_ratio: int,
1184 - ) -> list[tuple[int, int, int]]:
1185 - """Build local FA descriptors for all layers."""
1186 - assert self.transfer_topo is not None
1187 - num_blocks = self.num_blocks * block_size_ratio
1188 - result: list[tuple[int, int, int]] = []
1189 - for i, base_addr in enumerate(base_addresses):
1190 - kv_block_len = (
1191 - self.get_backend_aware_kv_block_len(
1192 - layer_idx=i, first_split=True, mamba_view=False
1193 - )
1194 - // block_size_ratio
1195 - )
1196 - page_stride = self.block_len_per_layer[i] // block_size_ratio
1197 - for block_id in range(num_blocks):
1198 - block_offset = block_id * page_stride
1199 - addr = base_addr + block_offset
1200 - result.append((addr, kv_block_len, self.device_id))
1201 -
1202 - if (
1203 - self.transfer_topo.virtually_split_kv_in_blocks
1204 - and not self._is_region_replicated(i)
1205 - ):
1206 - # Separate and interleave K/V regions to maintain the same
1207 - # descs ordering. This is needed for selecting contiguous heads
1208 - # when split across TP ranks. (Skipped for key-only REPLICATE.)
1209 - second_split = self.get_backend_aware_kv_block_len(
1210 - layer_idx=i, first_split=False, mamba_view=False
1211 - )
1212 - for block_id in range(num_blocks):
1213 - block_offset = block_id * page_stride
1214 - addr = base_addr + block_offset
1215 - v_addr = addr + kv_block_len
1216 - result.append((v_addr, second_split, self.device_id))
1217 - return result
1218 -
1219 - def _build_fa_remote(
1220 - self,
1221 - plan: TPMapping,
1222 - nixl_agent_meta: NixlAgentMetadata,
1223 - block_size_ratio: int,
1224 - ) -> list[tuple[int, int, int]]:
1225 - """Build remote FA descriptors for all layers."""
1226 - assert self.transfer_topo is not None
1227 - fa_group_idx = next(
1228 - i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
1229 - )
1230 - # SPLIT regions read their head slice from this many remote ranks at a
1231 - # per-rank offset; REPLICATE regions read the whole block once.
1232 - split_reads = len(plan.source_ranks_per_group[fa_group_idx])
1233 - num_blocks = nixl_agent_meta.num_blocks
1234 - result: list[tuple[int, int, int]] = []
1235 - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1236 - replicated = self._is_region_replicated(i)
1237 - # Read our whole local region size from remote..
1238 - local_block_len = self.get_backend_aware_kv_block_len(
1239 - layer_idx=i, first_split=True, mamba_view=False
1240 - )
1241 - remote_kv_block_len = local_block_len // block_size_ratio
1242 - if block_size_ratio > 1:
1243 - # ..using remote kv_block_len as transfer unit
1244 - local_block_len = remote_kv_block_len
1245 -
1246 - # REPLICATE reads the whole block once at offset 0; SPLIT gathers
1247 - # its head slice from `split_reads` remote ranks at a per-rank offset.
1248 - num_reads = 1 if replicated else split_reads
1249 - rank_offset = (
1250 - 0 if replicated else plan.rank_offset_factor * remote_kv_block_len
1251 - )
1252 - local_block_len = local_block_len // num_reads
1253 -
1254 - page_size = nixl_agent_meta.block_lens[i]
1255 - for block_id in range(num_blocks):
1256 - block_offset = block_id * page_size
1257 - # For each block, grab the kv heads chunk belonging to current local
1258 - # tp rank of size local_block_len.
1259 - addr = base_addr + block_offset + rank_offset
1260 - result.append((addr, local_block_len, nixl_agent_meta.device_id))
1261 -
1262 - emits_v = self.transfer_topo.virtually_split_kv_in_blocks and not replicated
1263 - if emits_v:
1264 - # With FlashInfer index V separately to allow head splitting.
1265 - second_split = self.get_backend_aware_kv_block_len(
1266 - layer_idx=i, first_split=False, mamba_view=False
1267 - )
1268 - second_split = second_split // num_reads
1269 - for block_id in range(num_blocks):
1270 - block_offset = block_id * page_size
1271 - addr = base_addr + block_offset + rank_offset
1272 - # Hop over the first split of remote page, K, to read V.
1273 - v_addr = addr + nixl_agent_meta.block_lens[i] // 2
1274 - result.append((v_addr, second_split, nixl_agent_meta.device_id))
1275 - return result
1276 -
1277 - def register_local_xfer_handler(
1278 - self,
1279 - block_size: int,
1280 - ) -> tuple[int, list[tuple[int, int, int]]]:
1281 - """
1282 - Function used for register local xfer handler with local block_size or
1283 - Remote block_size.
1284 -
1285 - When local block_size is same as remote block_size, we use local block_size
1286 - to register local_xfer_handler during init.
1287 -
1288 - When remote block size is less than local block size, we need to use
1289 - register another local_xfer_handler using remote block len to ensure
1290 - data copy correctness.
1291 - """
1292 - assert self.transfer_topo is not None
1293 - block_size_ratio = self.block_size // block_size
1294 - local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
1295 -
1296 - blocks_data = self._build_fa_local(local_base_addresses, block_size_ratio)
1297 - logger.debug(
1298 - "Created %s blocks for src engine %s and rank %s on device id %s",
1299 - len(blocks_data),
1300 - self.engine_id,
1301 - self.tp_rank,
1302 - self.device_id,
1303 - )
1304 - if self._has_mamba:
1305 - assert self.num_descs == len(blocks_data)
1306 - # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split
1307 - # is unnecessary — a single conv desc per block suffices. Consider
1308 - # adding a fast path that falls back to the standard 2-region
1309 - # registration (_build_fa_local mamba=True) when no hetero-TP
1310 - # remote has been seen. Currently we always register 4 regions
1311 - # because local descs are created before knowing the remote TP.
1312 - logger.debug("Registering local Mamba descriptors (4 regions/layer)")
1313 - blocks_data.extend(
1314 - self._build_mamba_local(local_base_addresses, block_size_ratio)
1315 - )
1316 -
1317 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
1318 - # NIXL_INIT_AGENT to be used for preparations of local descs.
1319 - return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data
1320 -
1321 - def add_remote_agent(
1322 - self,
1323 - nixl_agent_meta: NixlAgentMetadata,
1324 - remote_tp_rank: int = 0,
1325 - remote_tp_size: int = 1,
1326 - ) -> str:
1327 - """
1328 - Add the remote NIXL agent and prepare the descriptors for reading cache
1329 - blocks from remote.
1330 -
1331 - In particular, handle both homogeneous and heterogeneous TP. The former
1332 - requires local rank_i to read from remote rank_i.
1333 - The latter, in the case of D.world_size < P.world_size, requires that a
1334 - local (D) TP worker reads from multiple remote (P) TP workers.
1335 - Conversely, assuming D.world_size > P.world_size, two or more local TP
1336 - workers will read from a single remote TP worker.
1337 -
1338 - Here's an example for the last case described above (non-MLA):
1339 -
1340 - rank_offset p_remote_tp_rank
1341 - (kv split no)
1342 - --------------------------------
1343 - 0 0 Worker0 ---- 1st half of KV ----> Worker0 [ KV Cache ]
1344 - /
1345 - 1 0 Worker1 ---- 2nd half of KV -----/
1346 -
1347 - 0 1 Worker2 ---- 1st half of KV ----> Worker1 [ KV Cache ]
1348 - /
1349 - 1 1 Worker3 ---- 2nd half of KV -----/
1350 -
1351 -
1352 - Decoder TP workers Prefix TP workers
1353 - (world_size=4) (world_size=2)
1354 - tp_ratio = 4 // 2 = 2
1355 -
1356 - Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim]
1357 - then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "HND" layout format.
1358 - Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio
1359 - first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split
1360 - along the kv_heads dimension, and so forth until "tp_ratio" D TP workers have pulled from P-Worker0.
1361 -
1362 - Note that the above will also hold true for the homogeneous TP case, where tp_ratio evaluates to 1.
1363 -
1364 - Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0
1365 - so that the whole cache is shared by "tp_ratio" D TP workers.
1366 -
1367 - For Mamba hetero-TP, both tp_ratio > 0 (D_TP > P_TP) and
1368 - tp_ratio < 0 (P_TP > D_TP) are supported by the 3-read transfer.
1369 - """ # noqa: E501
1370 - engine_id = nixl_agent_meta.engine_id
1371 - # TODO re-evaluate refreshing for scaling/recovery
1372 - if remote_tp_rank in self._remote_agents.get(engine_id, {}):
1373 - logger.debug(
1374 - "Remote agent with engine_id %s and rank"
1375 - "%s already exchanged metadata, skip handshake.",
1376 - engine_id,
1377 - remote_tp_rank,
1378 - )
1379 - return self._remote_agents[engine_id][remote_tp_rank]
1380 -
1381 - ### Register remote engine in TransferTopology (idempotent).
1382 - assert self.transfer_topo is not None
1383 - transfer_topo = self.transfer_topo
1384 - physical_blocks_per_logical = (
1385 - nixl_agent_meta.physical_blocks_per_logical_kv_block
1386 - )
1387 - transfer_info = EngineTransferInfo(
1388 - remote_tp_size=remote_tp_size,
1389 - remote_block_size=nixl_agent_meta.block_size,
1390 - remote_block_len=nixl_agent_meta.block_lens[0],
1391 - remote_physical_blocks_per_logical=physical_blocks_per_logical,
1392 - )
1393 - transfer_topo.register_remote_engine(engine_id, transfer_info)
1394 - logger.info("Transfer plan: %s", transfer_topo.describe(engine_id))
1395 -
1396 - self.tp_mappings[engine_id] = compute_tp_mapping(
1397 - transfer_topology=transfer_topo,
1398 - remote_tp_size=remote_tp_size,
1399 - group_spec_types=self._group_spec_types,
1400 - )
1401 -
1402 - remote_agent_name = self.nixl_wrapper.add_remote_agent(
1403 - nixl_agent_meta.agent_metadata
1404 - )
1405 -
1406 - # Create dst descs and xfer side handles. TP workers have same #blocks
1407 - # so we only register once per engine_id.
1408 - # Example:
1409 - # block_size_ratio > 1:
1410 - # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
1411 - # local origin:| 0| 1| 8| 12|
1412 - # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
1413 - block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size)
1414 -
1415 - if engine_id not in self.dst_num_blocks:
1416 - self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
1417 -
1418 - # Keep track of remote agent kv caches base addresses.
1419 - self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
1420 - nixl_agent_meta.kv_caches_base_addr
1421 - )
1422 - self._validate_remote_agent_handshake(nixl_agent_meta, remote_tp_size)
1423 -
1424 - # This is 1 when P and D `--tensor-parallel-size` match. Otherwise,
1425 - # this is the ratio between the two sizes.
1426 - tp_ratio = transfer_topo.tp_ratio(remote_tp_size)
1427 -
1428 - logger.debug(
1429 - "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s",
1430 - engine_id,
1431 - remote_tp_rank,
1432 - tp_ratio,
1433 - )
1434 -
1435 - plan = self.tp_mappings[engine_id]
1436 -
1437 - ### (Optional) Register local agent memory regions. MLA is not split.
1438 - if (
1439 - tp_ratio < 0
1440 - and not self.use_mla
1441 - and tp_ratio not in self.src_xfer_handles_by_tp_ratio
1442 - ):
1443 - # Remote tp_size > local tp_size: read from multiple remote ranks.
1444 - # Logically "split" own regions into |tp_ratio| chunks. Mind that
1445 - # we only do this once per remote tp_size (replica-friendly).
1446 - self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
1447 -
1448 - for handle_data in self._build_local_splits_from_plan(
1449 - plan,
1450 - self.src_blocks_data,
1451 - self.num_descs,
1452 - ):
1453 - descs = self.nixl_wrapper.get_xfer_descs(
1454 - handle_data, self.nixl_memory_type
1455 - )
1456 - handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
1457 - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
1458 -
1459 - ### Register remote agent memory regions
1460 - # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With
1461 - # heterogeneous TP, prepare the descriptors by splitting the P KV cache along
1462 - # kv_head dim, of D worker's kv_head size (D>P).
1463 - # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].
1464 -
1465 - # Register all remote blocks, but only the corresponding kv heads.
1466 - blocks_data = self._build_fa_remote(
1467 - plan,
1468 - nixl_agent_meta,
1469 - block_size_ratio,
1470 - )
1471 - logger.debug(
1472 - "Created %s blocks for dst engine %s with remote rank %s and local rank %s",
1473 - len(blocks_data),
1474 - engine_id,
1475 - remote_tp_rank,
1476 - self.tp_rank,
1477 - )
1478 - if self._has_mamba:
1479 - logger.debug(
1480 - "Registering remote Mamba blocks for engine %s rank %s",
1481 - engine_id,
1482 - remote_tp_rank,
1483 - )
1484 - blocks_data.extend(
1485 - self._build_mamba_remote(
1486 - nixl_agent_meta,
1487 - tp_ratio,
1488 - transfer_info,
1489 - )
1490 - )
1491 -
1492 - # Register with NIXL.
1493 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
1494 - self.dst_xfer_side_handles[engine_id][remote_tp_rank] = (
1495 - self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs)
1496 - )
1497 -
1498 - if block_size_ratio > 1:
1499 - # when prefill with smaller block_size, we need to init a
1500 - # new handler with same block_len to match
1501 - self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = (
1502 - self.register_local_xfer_handler(nixl_agent_meta.block_size)[0]
1503 - )
1504 -
1505 - return remote_agent_name
1506 -
1507 - def _validate_remote_agent_handshake(
1508 - self, nixl_agent_meta: NixlAgentMetadata, remote_tp_size: int
1509 - ):
1510 - """
1511 - Validate the remote agent handshake metadata ensuring the
1512 - invariants hold true.
1513 - """
1514 - remote_engine_id = nixl_agent_meta.engine_id
1515 -
1516 - assert self.transfer_topo is not None
1517 - remote_info = self.transfer_topo.get_engine_info(remote_engine_id)
1518 - assert remote_info.remote_tp_size == remote_tp_size
1519 -
1520 - tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
1521 - block_size_ratio = self.transfer_topo.block_size_ratio(
1522 - nixl_agent_meta.block_size
1523 - )
1524 - # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba.
1525 - # Mamba models can have replicated FA KV with tp_ratio < 0.
1526 - # MLA models do not need to handle kv replication.
1527 - if not self.use_mla and not self._has_mamba:
1528 - assert not (
1529 - tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id)
1530 - )
1531 -
1532 - remote_physical_per_logical = (
1533 - nixl_agent_meta.physical_blocks_per_logical_kv_block
1534 - )
1535 - if (
1536 - self._has_mamba
1537 - and remote_physical_per_logical
1538 - != self._physical_blocks_per_logical_kv_block
1539 - and self.vllm_config.cache_config.enable_prefix_caching
1540 - ):
1541 - raise RuntimeError(
1542 - "Prefix caching with heterogeneous physical_blocks_per_logical "
1543 - "is not supported for Mamba hybrid models. "
1544 - f"Local: {self._physical_blocks_per_logical_kv_block}, "
1545 - f"Remote: {remote_physical_per_logical}. "
1546 - "Disable prefix caching with --no-enable-prefix-caching."
1547 - )
1548 -
1549 - if self._is_hma_required:
1550 - assert block_size_ratio == 1, (
1551 - "HMA does not support different remote block size yet"
1552 - )
1553 - kv_cache_layout = (
1554 - self.kv_cache_layout
1555 - if not self.use_host_buffer
1556 - else self.host_buffer_kv_cache_layout
1557 - )
1558 - if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout:
1559 - if (
1560 - self.kv_transfer_config.enable_permute_local_kv
1561 - and nixl_agent_meta.kv_cache_layout == "HND"
1562 - ):
1563 - logger.info(
1564 - "Remote is HND and local is NHD, enabled additional permute "
1565 - "on local device KV."
1566 - )
1567 - assert not self._is_hma_required, (
1568 - "HMA does not support block size post processing"
1569 - )
1570 - self.enable_permute_local_kv = True
1571 - else:
1572 - raise RuntimeError(
1573 - "Heterogeneous TP expects same kv_cache_layout. "
1574 - "Or enable experimental feature to use HND to NHD support by "
1575 - "setting 'enable_permute_local_kv'=True in --kv-transfer-config."
1576 - )
1577 - # if remote_agent used attn is not same as local,
1578 - # hint heterogenuous attn post process
1579 - if (
1580 - nixl_agent_meta.attn_backend_name != self.backend_name
1581 - and self.backend_name in ["CPU_ATTN"]
1582 - ):
1583 - if self._is_hma_required:
1584 - raise RuntimeError(
1585 - "heterogeneous attn post process is not supported with HMA"
1586 - )
1587 - logger.info(
1588 - "[Experimental] CPU_ATTN backend is used, "
1589 - "hint heterogeneous attn post process"
1590 - )
1591 - self.enable_heterogeneous_attn_post_process = True
1592 -
1593 - # Heterogeneous TP requires head-splitting, which only works with
1594 - # HND layout. MLA and replicated-KV cases don't split on heads.
1595 - # Mamba doesn't support heterogeneous TP.
1596 - if (
1597 - abs(tp_ratio) != 1
1598 - and not self.use_mla
1599 - and not self.transfer_topo.is_kv_replicated(remote_engine_id)
1600 - and kv_cache_layout != "HND"
1601 - and not self.enable_permute_local_kv
1602 - ):
1603 - raise RuntimeError(
1604 - "Heterogeneous TP head-dimension splitting requires contiguous heads. "
1605 - "Use HND layout on the prefill side."
1606 - )
1607 -
1608 - # Per-region block_len validation enforcing the P/D invariant.
1609 - # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer)
1610 - # only allow the number of blocks to differ; SPLIT regions scale with
1611 - # tp_ratio. Mamba uses the ssm_sizes counterpart, so skip block_len here.
1612 - if not self._has_mamba:
1613 - assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (
1614 - "Number of KV layers must match between prefill and decode"
1615 - )
1616 - model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated(
1617 - remote_engine_id
1618 - )
1619 - for i, local_len in enumerate(self.block_len_per_layer):
1620 - replicated = model_replicated or self._is_region_replicated(i)
1621 - remote_len = nixl_agent_meta.block_lens[i]
1622 - if replicated:
1623 - # Whole block copied; only the number of blocks may differ.
1624 - assert local_len // block_size_ratio == remote_len, (
1625 - "KV cache sizes must match between P and D when "
1626 - f"replicated (region {i}: local={local_len}, "
1627 - f"remote={remote_len}, bsr={block_size_ratio})."
1628 - )
1629 - elif tp_ratio > 0:
1630 - # D_TP >= P_TP: remote holds tp_ratio x local heads.
1631 - assert remote_len == (local_len * tp_ratio) // block_size_ratio, (
1632 - f"SPLIT region {i}: remote P KV block_len {remote_len} "
1633 - f"must equal local {local_len} * tp_ratio {tp_ratio} "
1634 - f"// block_size_ratio {block_size_ratio}."
1635 - )
1636 - else:
1637 - # P_TP > D_TP: local holds |tp_ratio| x remote heads.
1638 - assert block_size_ratio == 1, (
1639 - "Different local/remote block sizes are not supported "
1640 - "when P TP > D TP."
1641 - )
1642 - assert remote_len == local_len // (-tp_ratio), (
1643 - f"SPLIT region {i}: remote P KV block_len {remote_len} "
1644 - f"must equal local {local_len} // |tp_ratio| {-tp_ratio}."
1645 - )
1646 -
1647 - # TP workers that handhshake with same remote have same #blocks.
1648 - assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
1649 - # Same number of regions/~layers.
1650 - assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer)
1651 -
1652 - def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
1653 - """copy recved kv from host buffer to device."""
1654 - assert self.use_host_buffer
1655 - assert self.copy_blocks is not None
1656 -
1657 - local_block_ids = meta.local_physical_block_ids
1658 - # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups
1659 - for group_block_ids in local_block_ids:
1660 - self.copy_blocks(
1661 - self.host_xfer_buffers,
1662 - self.device_kv_caches,
1663 - group_block_ids,
1664 - group_block_ids,
1665 - "h2d",
1666 - )
1667 - if logger.isEnabledFor(logging.DEBUG):
1668 - logger.debug(
1669 - "synced recved kv of request[%s] to device kv buffer,"
1670 - "local_block_ids: %s. ",
1671 - req_id,
1672 - ",".join(map(str, local_block_ids)),
1673 - )
1674 -
1675 - def save_kv_to_host(self, metadata: NixlConnectorMetadata):
1676 - """copy kv from device to host buffer."""
1677 - assert self.use_host_buffer
1678 - assert self.copy_blocks is not None
1679 -
1680 - for req_id, meta in metadata.reqs_to_save.items():
1681 - meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
1682 - meta.local_block_ids
1683 - )
1684 - if logger.isEnabledFor(logging.DEBUG):
1685 - logger.debug(
1686 - "save_load_kv for request[%s] to host xfer buffer."
1687 - "local_block_ids: %s. ",
1688 - req_id,
1689 - ",".join(map(str, meta.local_physical_block_ids)),
1690 - )
1691 - # blocking
1692 - for group_block_ids in meta.local_physical_block_ids:
1693 - self.copy_blocks(
1694 - self.device_kv_caches,
1695 - self.host_xfer_buffers,
1696 - group_block_ids,
1697 - group_block_ids,
1698 - "d2h",
1699 - )
1700 -
1701 - def post_process_device_kv_on_receive(
1702 - self,
1703 - block_size_ratio: int,
1704 - block_ids_list: list[list[int]],
1705 - ):
1706 - """
1707 - Post process device kv cache after receiving from remote.
1708 -
1709 - 3 types of post processing supported:
1710 - * kv_cache_postprocess_layout => convert from HND to NHD
1711 - * kv_cache_postprocess_blksize => convert from small block size
1712 - to large block size
1713 - * kv_cache_postprocess_blksize_and_layout => convert from small
1714 - block size to large block size and convert from HND to NHD
1715 -
1716 - """
1717 - if len(self.device_kv_caches) == 0:
1718 - return
1719 - assert block_size_ratio >= 1, "Only nP < nD supported currently."
1720 - assert self.transfer_topo is not None
1721 - if self.enable_permute_local_kv and block_size_ratio > 1:
1722 - logger.debug(
1723 - "Post-processing device kv cache on receive by converting "
1724 - "block_size with %sx bigger and permuting layout from HND"
1725 - " to NHD.",
1726 - block_size_ratio,
1727 - )
1728 - elif self.enable_permute_local_kv:
1729 - logger.debug(
1730 - "Post-processing device kv cache on receive by permuting layout"
1731 - "from HND to NHD."
1732 - )
1733 - else:
1734 - logger.debug(
1735 - "Post-processing device kv cache on receive by converting "
1736 - "block_size with %sx bigger.",
1737 - block_size_ratio,
1738 - )
1739 -
1740 - split_k_and_v = self.transfer_topo.split_k_and_v
1741 -
1742 - for block_ids in block_ids_list:
1743 - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
1744 -
1745 - for _, cache_or_caches in self.device_kv_caches.items():
1746 - cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1747 - for cache in cache_list:
1748 - if self.enable_permute_local_kv and block_size_ratio > 1:
1749 - kv_postprocess_blksize_and_layout_on_receive(
1750 - cache, indices, block_size_ratio
1751 - )
1752 - elif self.enable_permute_local_kv:
1753 - kv_postprocess_layout_on_receive(cache, indices)
1754 - else:
1755 - kv_postprocess_blksize_on_receive(
1756 - cache, indices, block_size_ratio
1757 - )
1758 -
1759 - def post_process_device_kv_on_receive_heterogeneous_attn(
1760 - self, block_ids: list[int]
1761 - ):
1762 - """
1763 - Post process device kv cache after receiving from remote
1764 - for heterogeneous attention.
1765 - """
1766 - assert self.enable_heterogeneous_attn_post_process
1767 -
1768 - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
1769 -
1770 - for _, cache_or_caches in self.device_kv_caches.items():
1771 - blocks_to_update = cache_or_caches.index_select(1, indices)
1772 - current_platform.pack_kv_cache(
1773 - key=blocks_to_update[0],
1774 - value=blocks_to_update[1],
1775 - key_cache=cache_or_caches[0],
1776 - value_cache=cache_or_caches[1],
1777 - block_ids=block_ids,
1778 - indices=indices,
1779 - )
1780 -
1781 - def get_finished(self) -> tuple[set[str], set[str]]:
1782 - """
1783 - Get requests that are done sending or recving on this specific worker.
1784 - The scheduler process (via the MultiprocExecutor) will use this output
1785 - to track which workers are done.
1786 - """
1787 - assert self.transfer_topo is not None
1788 - done_sending = self._get_new_notifs()
1789 - done_recving = self._pop_done_transfers(self._recving_transfers)
1790 -
1791 - # Drain queue of requests where handshake or transfer setup failed.
1792 - failed_recv_reqs = set[ReqId]()
1793 - while not self._failed_recv_reqs.empty():
1794 - try:
1795 - failed_recv_reqs.add(self._failed_recv_reqs.get_nowait())
1796 - except queue.Empty:
1797 - break
1798 -
1799 - # Add failed requests to done_recving for scheduler tracking
1800 - # (blocks are already marked invalid, scheduler will handle recompute)
1801 - done_recving.update(failed_recv_reqs)
1802 -
1803 - if len(done_sending) > 0 or len(done_recving) > 0:
1804 - logger.debug(
1805 - "Rank %s, get_finished: %s requests done sending "
1806 - "and %s requests done recving (%s failed)",
1807 - self.tp_rank,
1808 - len(done_sending),
1809 - len(done_recving),
1810 - len(failed_recv_reqs),
1811 - )
1812 -
1813 - block_ids_for_blocksize_post_process = defaultdict(list)
1814 - block_ids_for_heterogeneous_attn_post_process = list[list[int]]()
1815 - for req_id in done_recving:
1816 - # clean up metadata for completed requests
1817 - meta = self._recving_metadata.pop(req_id, None)
1818 - assert meta is not None, f"{req_id} not found in recving_metadata list"
1819 -
1820 - # Skip KV sync and post-processing for failed requests
1821 - if req_id in failed_recv_reqs:
1822 - logger.warning(
1823 - "Skipping KV post-processing for failed request %s",
1824 - req_id,
1825 - )
1826 - continue
1827 -
1828 - assert meta.remote is not None
1829 - if self.use_host_buffer:
1830 - self.sync_recved_kv_to_device(req_id, meta)
1831 -
1832 - # post processing for heteroblocksize
1833 - remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id)
1834 - block_size_ratio = self.transfer_topo.block_size_ratio(
1835 - remote_info.remote_block_size
1836 - )
1837 - if not self.use_mla and (
1838 - block_size_ratio > 1 or self.enable_permute_local_kv
1839 - ):
1840 - assert not self._is_hma_required
1841 - block_ids_for_blocksize_post_process[block_size_ratio].append(
1842 - meta.local_physical_block_ids[0]
1843 - )
1844 - # post processing for heterogeneous attention
1845 - if self.enable_heterogeneous_attn_post_process:
1846 - block_ids_for_heterogeneous_attn_post_process.append(
1847 - meta.local_physical_block_ids[0]
1848 - )
1849 - for (
1850 - block_size_ratio,
1851 - block_ids_list,
1852 - ) in block_ids_for_blocksize_post_process.items():
1853 - self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list)
1854 -
1855 - for block_ids in block_ids_for_heterogeneous_attn_post_process:
1856 - self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
1857 -
1858 - # Handle timeout to avoid stranding blocks on remote.
1859 - now = time.perf_counter()
1860 - while self._reqs_to_send:
1861 - req_id, expires = next(iter(self._reqs_to_send.items()))
1862 - # Sorted dict, oldest requests are put first so we can exit early.
1863 - if now < expires:
1864 - break
1865 - count = self.consumer_notification_counts_by_req.pop(req_id, 0)
1866 - self.xfer_stats.record_kv_expired_req()
1867 - logger.warning(
1868 - "Releasing expired KV blocks for request %s which were "
1869 - "retrieved by %d remote worker(s) before lease expired.",
1870 - req_id,
1871 - count,
1872 - )
1873 - self._reqs_to_process.remove(req_id)
1874 - del self._reqs_to_send[req_id]
1875 - done_sending.add(req_id)
1876 -
1877 - return done_sending, done_recving
1878 -
1879 - def _get_new_notifs(self) -> set[str]:
1880 - """
1881 - Get req_ids which got a remote xfer message. When multiple consumers
1882 - are reading from the same producer (heterogeneous TP scenario), wait
1883 - for all consumers to be done pulling.
1884 -
1885 - Also handles heartbeat notifications ("HB:req1,req2,...") by
1886 - extending the lease on the referenced requests.
1887 - """
1888 - assert self.transfer_topo is not None
1889 - notified_req_ids: set[str] = set()
1890 - for notifs in self.nixl_wrapper.get_new_notifs().values():
1891 - for notif in notifs:
1892 - msg = notif.decode("utf-8")
1893 -
1894 - # Handle heartbeat messages from D-side.
1895 - if msg.startswith("HB:"):
1896 - self._handle_heartbeat(msg[3:])
1897 - continue
1898 -
1899 - req_id, tp_size = msg.rsplit(":", 1)
1900 - if (
1901 - req_id not in self._reqs_to_send
1902 - and req_id not in self._reqs_to_process
1903 - ):
1904 - logger.error(
1905 - "Potentially invalid KV blocks for "
1906 - "unrecognized request %s were retrieved by "
1907 - "a decode worker. They may have expired.",
1908 - req_id,
1909 - )
1910 - continue
1911 -
1912 - # NOTE: `tp_ratio` is the opposite when swapping local<>remote
1913 - n_consumers = int(tp_size)
1914 - tp_ratio = self.transfer_topo.tp_ratio(n_consumers)
1915 -
1916 - # Number of reads *per producer* to wait for.
1917 - # When remote D TP > local P TP we expect `tp_ratio` reads.
1918 - consumers_per_producer = (
1919 - -tp_ratio if n_consumers > self.world_size else 1
1920 - )
1921 -
1922 - self.consumer_notification_counts_by_req[req_id] += 1
1923 - # Wait all consumers (D) to be done reading before freeing.
1924 - if (
1925 - self.consumer_notification_counts_by_req[req_id]
1926 - == consumers_per_producer
1927 - ):
1928 - notified_req_ids.add(req_id)
1929 - del self.consumer_notification_counts_by_req[req_id]
1930 - self._reqs_to_process.remove(req_id)
1931 - self._reqs_to_send.pop(req_id, None)
1932 - return notified_req_ids
1933 -
1934 - def _handle_heartbeat(self, payload: str) -> None:
1935 - """Extend leases for requests referenced in a heartbeat.
1936 -
1937 - Args:
1938 - payload: comma-separated P-side request IDs, e.g.
1939 - "req_abc,req_def".
1940 - """
1941 - new_expiry = time.perf_counter() + self._lease_extension
1942 - for req_id in payload.split(","):
1943 - if req_id in self._reqs_to_send:
1944 - old = self._reqs_to_send[req_id]
1945 - self._reqs_to_send[req_id] = max(old, new_expiry)
1946 - logger.debug(
1947 - "Heartbeat extended lease for request %s "
1948 - "by %ds (old_expiry=%.1f, new_expiry=%.1f)",
1949 - req_id,
1950 - self._lease_extension,
1951 - old,
1952 - new_expiry,
1953 - )
1954 -
1955 - def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:
1956 - """
1957 - Pop completed xfers by checking for DONE state.
1958 - Args:
1959 - transfers: dict of req_id -> list[running_xfer]
1960 - Returns:
1961 - set of req_ids that have all done xfers
1962 - """
1963 - done_req_ids: set[str] = set()
1964 - for req_id, handles in list(transfers.items()):
1965 - in_progress = []
1966 - for handle in handles:
1967 - try:
1968 - xfer_state = self.nixl_wrapper.check_xfer_state(handle)
1969 - if xfer_state == "DONE":
1970 - # Get telemetry from NIXL
1971 - res = self.nixl_wrapper.get_xfer_telemetry(handle)
1972 - self.xfer_stats.record_transfer(res)
1973 - self.nixl_wrapper.release_xfer_handle(handle)
1974 - elif xfer_state == "PROC":
1975 - in_progress.append(handle)
1976 - continue
1977 - else:
1978 - self._log_failure(
1979 - failure_type="transfer_failed",
1980 - msg="Marking blocks as invalid",
1981 - req_id=req_id,
1982 - xfer_state=xfer_state,
1983 - )
1984 - self._handle_failed_transfer(req_id, handle)
1985 - except Exception as e:
1986 - self._log_failure(
1987 - failure_type="transfer_exception",
1988 - msg="Marking blocks as invalid",
1989 - req_id=req_id,
1990 - error=e,
1991 - )
1992 - self._handle_failed_transfer(req_id, handle)
1993 -
1994 - if not in_progress:
1995 - # Only report request as completed when all transfers are done.
1996 - done_req_ids.add(req_id)
1997 - del transfers[req_id]
1998 - else:
1999 - transfers[req_id] = in_progress
2000 - return done_req_ids
2001 -
2002 - def _handle_failed_transfer(self, req_id: str, handle: int | None):
2003 - """
2004 - Handle a failed transfer by marking all (logical) blocks as invalid and
2005 - recording the failure.
2006 -
2007 - Args:
2008 - req_id: The request ID.
2009 - handle: The transfer handle.
2010 - """
2011 - # Use .get() here as the metadata cleanup is handled by get_finished()
2012 - # TODO (NickLucche) handle failed transfer for HMA.
2013 - if (meta := self._recving_metadata.get(req_id)) and not self._is_hma_required:
2014 - self._invalid_block_ids.put(set(meta.local_block_ids[0]))
2015 - self._failed_recv_reqs.put(req_id)
2016 - if handle is not None:
2017 - self.nixl_wrapper.release_xfer_handle(handle)
2018 - self.xfer_stats.record_failed_transfer()
2019 -
2020 - def start_load_kv(self, metadata: NixlConnectorMetadata):
2021 - """
2022 - Start loading by triggering non-blocking nixl_xfer.
2023 - We check for these trnxs to complete in each step().
2024 - """
2025 - for req_id, meta in metadata.reqs_to_recv.items():
2026 - meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
2027 - meta.local_block_ids
2028 - )
2029 - assert meta.remote is not None
2030 - # Remote block IDs are kept logical here; expanded in
2031 - # _read_blocks_for_req using the remote engine's phys ratio.
2032 - remote_engine_id = meta.remote.engine_id
2033 - logger.debug(
2034 - "start_load_kv for request %s from remote engine %s. "
2035 - "Num local_block_ids: %s. Num remote_block_ids: %s. ",
2036 - req_id,
2037 - remote_engine_id,
2038 - len(meta.local_physical_block_ids),
2039 - len(meta.remote.block_ids),
2040 - )
2041 - # always store metadata for failure recovery
2042 - self._recving_metadata[req_id] = meta
2043 - if remote_engine_id not in self._remote_agents:
2044 - # Initiate handshake with remote engine to exchange metadata.
2045 - with self._handshake_lock:
2046 - if remote_engine_id not in self._remote_agents:
2047 - self._background_nixl_handshake(req_id, remote_engine_id, meta)
2048 - continue
2049 -
2050 - # Handshake already completed, start async read xfer.
2051 - self._read_blocks_for_req(req_id, meta)
2052 -
2053 - # Start transfers for requests whose handshakes have now finished.
2054 - while not self._ready_requests.empty():
2055 - self._read_blocks_for_req(*self._ready_requests.get_nowait())
2056 -
2057 - # Keep around the requests that have been part of a batch. This is
2058 - # needed because async scheduling pushes the misalignment between the
2059 - # moment in which requests expiration is set (P side) and the moment in
2060 - # which blocks are read from D. As P can now more easily lag behind D
2061 - # while processing the next batch, we make sure to only set an
2062 - # expiration for requests that have not been read from D yet.
2063 - for req_id in metadata.reqs_in_batch:
2064 - self._reqs_to_process.add(req_id)
2065 -
2066 - # Remove all requests that are not to be processed (eg aborted).
2067 - for req_id in metadata.reqs_not_processed:
2068 - self._reqs_to_process.discard(req_id)
2069 - # We should never get an abort after setting an expiry timer
2070 - assert req_id not in self._reqs_to_send
2071 -
2072 - # Add to requests that are waiting to be read and track expiration.
2073 - for req_id, expiration_time in metadata.reqs_to_send.items():
2074 - if req_id in self._reqs_to_process:
2075 - self._reqs_to_send[req_id] = expiration_time
2076 -
2077 - # Send heartbeats to P-side engines to keep KV blocks alive while
2078 - # requests sit in the D scheduler WAITING queue.
2079 - self._send_heartbeats(metadata)
2080 -
2081 - def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
2082 - """
2083 - Send heartbeat notifications to remote engines, extending lease on KV blocks.
2084 - """
2085 - for engine_id, hb_info in metadata.heartbeat_by_engine.items():
2086 - # Proactive handshake (this request may still be in waiting queue) so
2087 - # the **next** heartbeat for this remote can go through.
2088 - if (
2089 - self._ensure_handshake(
2090 - engine_id, hb_info.host, hb_info.port, hb_info.tp_size
2091 - )
2092 - is not None
2093 - ):
2094 - continue # handshake is still pending
2095 -
2096 - # Build the heartbeat message: "HB:req1,req2,..."
2097 - hb_msg = ("HB:" + ",".join(hb_info.req_ids)).encode()
2098 - for agent_name in self._remote_agents[engine_id].values():
2099 - try:
2100 - self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg)
2101 - except Exception:
2102 - logger.debug(
2103 - "Failed to send heartbeat to engine %s",
2104 - engine_id,
2105 - exc_info=True,
2106 - )
2107 -
2108 - def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
2109 - assert meta.remote is not None and self.transfer_topo is not None
2110 - engine_id = meta.remote.engine_id
2111 - # Update last activity from this remote. Mind that cleanup is done on main
2112 - # thread (this one), so we don't race on this structure.
2113 - self._engine_last_active[engine_id] = time.perf_counter()
2114 - plan = self.tp_mappings[engine_id]
2115 - remote_info = self.transfer_topo.get_engine_info(engine_id)
2116 - tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
2117 -
2118 - meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
2119 - meta.remote.block_ids,
2120 - remote_info.remote_physical_blocks_per_logical,
2121 - )
2122 - remote_block_ids = meta.remote.block_ids
2123 - local_block_ids = meta.local_physical_block_ids
2124 - num_groups = len(local_block_ids)
2125 - read_specs = [
2126 - ReadSpec(
2127 - remote_rank=rank,
2128 - local_block_ids=[
2129 - list(local_block_ids[g])
2130 - if rank in plan.source_ranks_per_group[g]
2131 - else []
2132 - for g in range(num_groups)
2133 - ],
2134 - remote_block_ids=[
2135 - list(remote_block_ids[g])
2136 - if rank in plan.source_ranks_per_group[g]
2137 - else []
2138 - for g in range(num_groups)
2139 - ],
2140 - )
2141 - for rank in plan.all_source_ranks
2142 - ]
2143 -
2144 - # D may have to perform multiple reads from different remote ranks.
2145 - # MLA opt: when P TP > D TP, only a single read is executed for
2146 - # the first remote rank (cache is duplicated)..
2147 - if self.use_mla and tp_ratio < 0:
2148 - assert len(read_specs) == 1
2149 -
2150 - for i, spec in enumerate(read_specs):
2151 - remote_block_size = remote_info.remote_block_size
2152 - logger.debug(
2153 - "Remote agent %s available, calling _read_blocks"
2154 - " on remote rank %s with remote block size %s for req %s",
2155 - meta.remote.engine_id,
2156 - spec.remote_rank,
2157 - remote_block_size,
2158 - req_id,
2159 - )
2160 - # Get side handles.
2161 - if tp_ratio < 0 and not self.use_mla:
2162 - assert remote_block_size == self.block_size
2163 - # Remote tp_size > local tp_size: we must perform multiple
2164 - # reads. Get the memory chunk onto which we will write to.
2165 - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
2166 - else:
2167 - # Single read from remote, we write to the whole memory region.
2168 - # Also handle remote block size different from local block size.
2169 - local_xfer_side_handle = self.src_xfer_handles_by_block_size[
2170 - remote_block_size
2171 - ]
2172 -
2173 - # Destination handle: remote_engine_id -> remote_rank -> handle.
2174 - remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
2175 - spec.remote_rank
2176 - ]
2177 -
2178 - self._read_blocks(
2179 - read_spec=spec,
2180 - request_id=req_id,
2181 - dst_engine_id=meta.remote.engine_id,
2182 - remote_request_id=meta.remote.request_id,
2183 - local_xfer_side_handle=local_xfer_side_handle,
2184 - remote_xfer_side_handle=remote_xfer_side_handle,
2185 - )
2186 -
2187 - if self.use_mla and tp_ratio < 0 and read_specs:
2188 - # ..but we still need to notify the other remote ranks that we
2189 - # have the blocks we need so they can update the request state.
2190 - notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()
2191 - remote_agents = self._remote_agents[meta.remote.engine_id]
2192 - for rank_to_notify, agent in remote_agents.items():
2193 - if rank_to_notify != read_specs[0].remote_rank:
2194 - self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
2195 -
2196 - def _read_blocks(
2197 - self,
2198 - read_spec: ReadSpec,
2199 - dst_engine_id: str,
2200 - request_id: str,
2201 - remote_request_id: str,
2202 - local_xfer_side_handle: int,
2203 - remote_xfer_side_handle: int,
2204 - ):
2205 - """
2206 - Post a READ point-to-point xfer request from a single local worker to
2207 - a single remote worker.
2208 - """
2209 - assert self.transfer_topo is not None
2210 - remote_rank = read_spec.remote_rank
2211 - local_block_ids = read_spec.local_block_ids
2212 - remote_block_ids = read_spec.remote_block_ids
2213 -
2214 - remote_info = self.transfer_topo.get_engine_info(dst_engine_id)
2215 - block_size_ratio = self.transfer_topo.block_size_ratio(
2216 - remote_info.remote_block_size
2217 - )
2218 - if block_size_ratio > 1:
2219 - # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.
2220 - assert not self._is_hma_required
2221 - local_block_ids0 = local_block_ids[0] if local_block_ids else []
2222 - remote_block_ids0 = remote_block_ids[0]
2223 - local_block_ids_mapped = self.get_mapped_blocks(
2224 - np.asarray(local_block_ids0), block_size_ratio
2225 - ).tolist()
2226 - if len(local_block_ids_mapped) > len(remote_block_ids0):
2227 - # NOTE:
2228 - # get_mapped_blocks will always expand block_ids for n times.
2229 - # ex:
2230 - # prefill block_ids with block_size as 4:
2231 - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
2232 - # Local decode block_ids with block_size as 16: [1, 2, 3]
2233 - # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to
2234 - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
2235 - # Then we clip local to align with prefill
2236 - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to
2237 - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
2238 - local_block_ids_mapped = local_block_ids_mapped[
2239 - : len(remote_block_ids0)
2240 - ]
2241 - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
2242 - remote_block_ids = [remote_block_ids0]
2243 - # NOTE(rob): having the staging blocks be on the READER side is
2244 - # not going to work well (since we will have to call rearrange tensors).
2245 - # after we detect the txn is complete (which means we cannot make the
2246 - # read trxn async easily). If we want to make "READ" happen cleanly,
2247 - # then we will need to have the staging blocks on the remote side.
2248 -
2249 - # NOTE(rob): according to nvidia the staging blocks are used to
2250 - # saturate IB with heterogeneous TP sizes.
2251 -
2252 - # Number of D TP workers that will read from dst P. Propagate info
2253 - # on notification so that dst worker can wait before freeing blocks.
2254 - notif_id = f"{remote_request_id}:{self.world_size}".encode()
2255 -
2256 - # Full prefix cache hit: do not need to read remote blocks,
2257 - # just notify P worker that we have the blocks we need.
2258 - if len(local_block_ids) == 0:
2259 - # A full prefix cache hit is indicated with an empty list.
2260 - agent_name = self._remote_agents[dst_engine_id][remote_rank]
2261 - try:
2262 - self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
2263 - except Exception as e:
2264 - self._log_failure(
2265 - failure_type="notification_failed",
2266 - msg="P worker blocks will be freed after timeout. "
2267 - "This may indicate network issues.",
2268 - req_id=request_id,
2269 - error=e,
2270 - dst_engine_id=dst_engine_id,
2271 - remote_rank=remote_rank,
2272 - remote_agent_name=agent_name,
2273 - )
2274 - self.xfer_stats.record_failed_notification()
2275 - return
2276 -
2277 - assert (
2278 - len(remote_block_ids)
2279 - == len(local_block_ids)
2280 - == len(self.kv_cache_config.kv_cache_groups)
2281 - )
2282 - remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical
2283 - local_block_ids, remote_block_ids = self._apply_prefix_caching(
2284 - local_block_ids, remote_block_ids, remote_physical_per_logical
2285 - )
2286 -
2287 - # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
2288 - # corresponding rank. With heterogeneous TP, fixing D>P, the D tp
2289 - # workers will issue xfers to parts of the P worker remote kv caches.
2290 -
2291 - # Get descs ids.
2292 - remote_block_descs_ids = self._compute_desc_ids(
2293 - block_ids=remote_block_ids,
2294 - dst_num_blocks=self.dst_num_blocks[dst_engine_id],
2295 - block_size_ratio=None,
2296 - physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
2297 - )
2298 - local_block_descs_ids = self._compute_desc_ids(
2299 - block_ids=local_block_ids,
2300 - dst_num_blocks=self.dst_num_blocks[self.engine_id],
2301 - block_size_ratio=block_size_ratio,
2302 - physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
2303 - )
2304 -
2305 - assert len(local_block_descs_ids) == len(remote_block_descs_ids)
2306 -
2307 - # Prepare transfer with Nixl.
2308 - handle = None
2309 - try:
2310 - handle = self.nixl_wrapper.make_prepped_xfer(
2311 - "READ",
2312 - local_xfer_side_handle,
2313 - local_block_descs_ids,
2314 - remote_xfer_side_handle,
2315 - remote_block_descs_ids,
2316 - notif_msg=notif_id,
2317 - )
2318 -
2319 - # Begin async xfer.
2320 - self.nixl_wrapper.transfer(handle)
2321 -
2322 - # Use handle to check completion in future step().
2323 - self._recving_transfers[request_id].append(handle)
2324 - except Exception as e:
2325 - # mark all (logical) blocks for this request as invalid
2326 - self._log_failure(
2327 - failure_type="transfer_setup_failed",
2328 - req_id=request_id,
2329 - msg="Marking blocks as invalid",
2330 - error=e,
2331 - dst_engine_id=dst_engine_id,
2332 - remote_rank=remote_rank,
2333 - )
2334 - self._handle_failed_transfer(request_id, handle)
2335 -
2336 - def get_mapped_blocks(
2337 - self, block_ids: np.ndarray, block_size_ratio: int
2338 - ) -> np.ndarray:
2339 - """
2340 - Calculates the new set of block IDs by mapping every element
2341 - in the (potentially sparse) input array.
2342 - Example: block_ids=[0, 2], block_size_ratio=2
2343 - get_mapped_blocks 0 1 [2 3] 4 5
2344 - # remote is |h0-b0|h1-b0||h0-b1|h1-b1||h0-b1|h1-b1||
2345 - # local is |h0-b0......||h1-b0......||h2-b0........
2346 - local_block_ids 0 [1] 2
2347 - """
2348 - if block_ids.size == 0:
2349 - return np.array([], dtype=np.int64)
2350 -
2351 - start_ids = block_ids * block_size_ratio
2352 - offsets = np.arange(block_size_ratio)
2353 - mapped_2d = start_ids[:, None] + offsets[None, :]
2354 -
2355 - return mapped_2d.flatten().astype(np.int64)
2356 -
2357 - def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:
2358 - """
2359 - Convert logical block ids to kernel physical block ids.
2360 - This is required when the logical block size (the one set by the user)
2361 - does not match the one required by the attn backend.
2362 - """
2363 - if self._physical_blocks_per_logical_kv_block == 1:
2364 - # Noop when physical and logical block sizes are the same
2365 - return block_ids
2366 - block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape(
2367 - 1, -1
2368 - )
2369 - # Mamba blocks have no logical<>physical discrepancy
2370 - group_specs = self.kv_cache_config.kv_cache_groups
2371 - return [
2372 - BlockTable.map_to_kernel_blocks(
2373 - np.array(group),
2374 - self._physical_blocks_per_logical_kv_block,
2375 - block_arange,
2376 - ).tolist()
2377 - if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
2378 - else group
2379 - for i, group in enumerate(block_ids)
2380 - ]
2381 -
2382 - def _apply_prefix_caching(
2383 - self,
2384 - local_block_ids: BlockIds,
2385 - remote_block_ids: BlockIds,
2386 - remote_physical_per_logical: int,
2387 - ) -> tuple[BlockIds, list]:
2388 - """Apply prefix caching by trimming local/remote block ID lists.
2389 -
2390 - For non-Mamba models: end-trim remote to match local count, so that
2391 - already-cached prefix blocks are skipped in the transfer.
2392 -
2393 - For Mamba hybrid (prefix caching not yet supported): front-trim both
2394 - to the minimum count to handle kernel block count discrepancies from
2395 - logical block rounding in heterogeneous TP.
2396 - """
2397 - # Partial prefix cache hit: just read uncomputed blocks.
2398 - # Skip mamba groups — their blocks represent full state (conv+ssm),
2399 - # not per-token data, so trimming would corrupt the transfer.
2400 - remote_block_ids = list(remote_block_ids)
2401 - if not self._has_mamba:
2402 - for i, remote_group in enumerate(remote_block_ids):
2403 - num_local_blocks = len(local_block_ids[i])
2404 - assert num_local_blocks <= len(remote_group)
2405 - if num_local_blocks < len(remote_group):
2406 - remote_block_ids[i] = remote_group[-num_local_blocks:]
2407 - else:
2408 - # (NOTE: ZhanqiuHu) Mamba hybrid: no prefix caching support so far.HeteroTP
2409 - # can cause different kernel block counts due to logical block rounding.
2410 - # Example: 640 prompt tokens, kernel_block_size=64
2411 - # remote physical_per_logical=10, local physical_per_logical=6
2412 - # remote logical ids from kv_transfer_params = [0]
2413 - # local logical ids allocated = [0, 1]
2414 - # remote kernel blocks: [0..9] (1*10=10)
2415 - # local kernel blocks: [0..11] (2*6=12)
2416 - # actual data blocks = ceil(640/64) = 10, trim both to 10
2417 - # Vice versa (remote physical_per_logical=6, local=10):
2418 - # remote logical ids = [0, 1], local logical ids = [0]
2419 - # remote kernel blocks: [0..11] (2*6=12)
2420 - # local kernel blocks: [0..9] (1*10=10)
2421 - # actual data blocks = ceil(640/64) = 10, trim both to 10
2422 - local_block_ids = list(local_block_ids)
2423 - for i, remote_group in enumerate(remote_block_ids):
2424 - num_local_blocks = len(local_block_ids[i])
2425 - num_remote_blocks = len(remote_group)
2426 - if (
2427 - _is_ssm_spec(self._group_spec_types[i])
2428 - and num_local_blocks < num_remote_blocks
2429 - ):
2430 - # NOTE (NickLucche): With prefix caching on SSM, (remote) blocks
2431 - # prior to the last one are placeholders (null blocks). Mind that
2432 - # this doesn't really impact transfer, as we only still care about
2433 - # the last "block", the full in-place state.
2434 - assert num_local_blocks == 1, "SSM can only have one local block"
2435 - remote_block_ids[i] = remote_group[-num_local_blocks:]
2436 - elif (
2437 - self._physical_blocks_per_logical_kv_block
2438 - == remote_physical_per_logical
2439 - and num_local_blocks < num_remote_blocks
2440 - ):
2441 - # Partial prefix cache hit for FA group.
2442 - remote_block_ids[i] = remote_group[-num_local_blocks:]
2443 - else:
2444 - # TODO Handle prefix caching with different block_sizes
2445 - max_padding = max(
2446 - self._physical_blocks_per_logical_kv_block,
2447 - remote_physical_per_logical,
2448 - )
2449 - assert abs(num_local_blocks - num_remote_blocks) < max_padding, (
2450 - f"Group {i}: |{num_local_blocks} - "
2451 - f"{num_remote_blocks}| >= {max_padding}"
2452 - )
2453 - num_blocks = min(num_local_blocks, num_remote_blocks)
2454 - local_block_ids[i] = local_block_ids[i][:num_blocks]
2455 - remote_block_ids[i] = remote_group[:num_blocks]
2456 - return local_block_ids, remote_block_ids
2457 -
2458 - def _logical_to_remote_kernel_block_ids(
2459 - self, block_ids: BlockIds, remote_physical_per_logical: int
2460 - ) -> BlockIds:
2461 - """Map logical block IDs to physical kernel block IDs on the remote.
2462 -
2463 - Args:
2464 - block_ids: per-group lists of logical block IDs.
2465 - remote_physical_per_logical: remote engine's physical blocks
2466 - per logical block.
2467 -
2468 - Returns:
2469 - Same structure with FA groups expanded (each logical block L
2470 - becomes kernel blocks [L*remote_physical_per_logical, ..
2471 - L*remote_physical_per_logical +
2472 - remote_physical_per_logical - 1]).
2473 - Mamba groups are passed through unchanged.
2474 - """
2475 - if remote_physical_per_logical == 1:
2476 - return block_ids
2477 - remote_arange = np.arange(remote_physical_per_logical).reshape(1, -1)
2478 - group_specs = self.kv_cache_config.kv_cache_groups
2479 - result = [
2480 - BlockTable.map_to_kernel_blocks(
2481 - np.array(group),
2482 - remote_physical_per_logical,
2483 - remote_arange,
2484 - ).tolist()
2485 - if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
2486 - else group
2487 - for i, group in enumerate(block_ids)
2488 - ]
2489 - return result
2490 -
2491 - def get_backend_aware_kv_block_len(
2492 - self, layer_idx: int, first_split: bool = True, mamba_view: bool = False
2493 - ) -> int:
2494 - """
2495 - Get the block length for one K/V element (K and V have the same size).
2496 -
2497 - For FA and other backends, this is equal to the length of the whole
2498 - block, as K and V are in separate regions.
2499 - For FlashInfer, this is half the length of the whole block, as K and V
2500 - share the same region.
2501 - Similarly, for SSM-based models, state and conv are interleaved, but crucially
2502 - the their size differs.
2503 - Reference diagram:
2504 - KVCacheTensor (Shared)
2505 - / \\
2506 - / \\
2507 - / \\
2508 - Attention (FlashInfer) View Mamba View
2509 - | |
2510 - | |
2511 - +-------------------+ +-------------------+
2512 - | KVCacheTensor | | KVCacheTensor |
2513 - | | | |
2514 - |<----- page ------>| |<----- page ------->|
2515 - | size | | size |
2516 - | Key 0 | Val 0 | |Conv 0 | SSM 0 |
2517 - | Key 1 | Val 1 | |Conv 1 | SSM 1 |
2518 - | ... | ... | | ... | ... |
2519 - | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 |
2520 - | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 |
2521 - +-------------------+ +--------------------+
2522 - |1st_split-2nd_split| |1st_split-2nd_split |
2523 - """
2524 - assert self.transfer_topo is not None
2525 - virtually_split = self.transfer_topo.virtually_split_kv_in_blocks
2526 - if virtually_split and mamba_view:
2527 - block_len = self._mamba_ssm_size[not first_split]
2528 - else:
2529 - # Per-descriptor block length: a SPLIT region (full-attn under the
2530 - # virtually-split layout) emits separate K and V and uses
2531 - # block_len//2; REPLICATE (MLA, key-only) and non-split layouts use
2532 - # the whole block.
2533 - half_block = virtually_split and not self._is_region_replicated(layer_idx)
2534 - block_len = self.block_len_per_layer[layer_idx] // (2 if half_block else 1)
2535 - return block_len
2536 -
2537 - def get_kv_connector_stats(self) -> KVConnectorStats | None:
2538 - """
2539 - Get the KV transfer stats for the connector.
2540 - """
2541 - # Clear stats for next iteration
2542 - if not self.xfer_stats.is_empty():
2543 - return self.xfer_stats.clone_and_reset()
2544 - return None
2545 -
2546 - def get_block_ids_with_load_errors(self) -> set[int]:
2547 - """
2548 - Return and clear the set of block IDs that failed to load.
2549 -
2550 - This is called by the scheduler to identify blocks that need
2551 - to be retried after a NIXL transfer failure.
2552 - """
2553 - # Drain the queue (thread-safe, no lock needed).
2554 - result: set[int] = set()
2555 - while not self._invalid_block_ids.empty():
2556 - try:
2557 - result.update(self._invalid_block_ids.get_nowait())
2558 - except queue.Empty:
2559 - break
2560 - return result
2561 -
2562 - def _evict_stale_engines(self) -> None:
2563 - """Scan for and evict remote engines that have exceeded their TTL.
2564 -
2565 - Called from the main thread in when a new remote engine appears.
2566 - We can only go OOM as we discover and register a new remote, therefore we make
2567 - sure we clean up stale engine data structures before then. This invariant
2568 - prevents us from using background threads, though memory usage is not guaranteed
2569 - to be "optimal" until a new handshake is performed.
2570 -
2571 - Engines with active transfers or pending handshakes cannot be stale:
2572 - - Active transfers touch _engine_last_active in start_load_kv.
2573 - - Pending handshakes don't have an _engine_last_active entry yet
2574 - """
2575 - # NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number
2576 - # of remote engines is registered all at once (adding a background cleanup
2577 - # thread wouldnt help either).
2578 - # If that scenario is plausible, we can follow up with an LRU eviction policy.
2579 - if self._engine_ttl <= 0:
2580 - return
2581 -
2582 - now = time.perf_counter()
2583 - for eid, last_active in list(self._engine_last_active.items()):
2584 - if now - last_active > self._engine_ttl:
2585 - self._cleanup_remote_engine(eid)
2586 -
2587 - def _cleanup_remote_engine(
2588 - self, engine_id: EngineId, *, log_eviction: bool = True
2589 - ) -> None:
2590 - """Remove all state for a single remote engine.
2591 -
2592 - Releases NIXL resources (dlist handles, remote agents) and clears
2593 - all per-engine data structures. Used by both TTL eviction and
2594 - shutdown.
2595 - """
2596 - assert engine_id in self._remote_agents
2597 -
2598 - for handle in self.dst_xfer_side_handles.pop(engine_id).values():
2599 - self.nixl_wrapper.release_dlist_handle(handle)
2600 - for agent_name in self._remote_agents.pop(engine_id).values():
2601 - self.nixl_wrapper.remove_remote_agent(agent_name)
2602 -
2603 - del self.kv_caches_base_addr[engine_id]
2604 - del self.dst_num_blocks[engine_id]
2605 - del self.tp_mappings[engine_id]
2606 - if self.transfer_topo is not None:
2607 - self.transfer_topo.unregister_remote_engine(engine_id)
8
2609 - last_active = self._engine_last_active.pop(engine_id)
2610 - if log_eviction:
2611 - logger.info(
2612 - "Evicted stale remote engine %s (inactive for %.1fs).",
2613 - engine_id,
2614 - time.perf_counter() - last_active,
2615 - )
9 +# Backward compatibility: NixlConnectorWorker is the pull-based worker.
10 +NixlConnectorWorker = NixlPullConnectorWorker
11
2617 - def __del__(self):
2618 - self.shutdown()
12
2620 - def shutdown(self):
2621 - """Shutdown the connector worker."""
2622 - if not hasattr(self, "_handshake_initiation_executor"):
2623 - # error happens during init, no need to shutdown
2624 - return
2625 - self._handshake_initiation_executor.shutdown(wait=False)
2626 - for handles in self._recving_transfers.values():
2627 - for handle in handles:
2628 - self.nixl_wrapper.release_xfer_handle(handle)
2629 - self._recving_transfers.clear()
2630 - for handle in self.src_xfer_handles_by_block_size.values():
2631 - self.nixl_wrapper.release_dlist_handle(handle)
2632 - self.src_xfer_handles_by_block_size.clear()
2633 - for handles in self.src_xfer_handles_by_tp_ratio.values():
2634 - for handle in handles:
2635 - self.nixl_wrapper.release_dlist_handle(handle)
2636 - self.src_xfer_handles_by_tp_ratio.clear()
2637 - for engine_id in list(self._remote_agents):
2638 - self._cleanup_remote_engine(engine_id, log_eviction=False)
2639 - for desc in self._registered_descs:
2640 - self.nixl_wrapper.deregister_memory(desc)
2641 - self._registered_descs.clear()
13 +__all__ = ["NixlConnectorWorker", "NixlPullConnectorWorker"]
vllm/v1/core/sched/scheduler.py
+13
@@ -2009,6 +2009,19 @@ class Scheduler(SchedulerInterface):
2009 )
2010 return len(self.requests) > num_in_queues
2011
2012 + def has_requests(self) -> bool:
2013 + # Override the interface default to also keep the engine alive while a
2014 + # connector still has pending push work (e.g. push-mode WRITE transfers
2015 + # in flight after all "live" requests have finished). Without this hook
2016 + # the engine would quiesce before the connector can drain completions.
2017 + # TODO: replace with a more general mechanism for connectors to keep
2018 + # the scheduler alive.
2019 + return (
2020 + self.has_unfinished_requests()
2021 + or self.has_finished_requests()
2022 + or (self.connector is not None and self.connector.has_pending_push_work())
2023 + )
2024 +
2025 def reset_prefix_cache(
2026 self, reset_running_requests: bool = False, reset_connector: bool = False
2027 ) -> bool: