@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3179e53135
Fork 0

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 11:15 UTC

[P/D] Prefill compute optimizations with bi-directional KV cache transfers between P and D nodes (#32553)

Signed-off-by: Sunita Nadampalli <nadampal@amazon.com>

snadampal committed Apr 30, 2026 at 03:14 UTC 3179e53135dbc0bbb9d845fb56e5380a6b88157e
4 files changed +1570 -9
benchmarks/multi_turn/benchmark_serving_multi_turn.py
+5
@@ -217,6 +217,7 @@ async def send_request(
217 min_tokens: int | None = None,
218 max_tokens: int | None = None,
219 timeout_sec: int = 120,
220 + conversation_id: str | None = None,
221 ) -> ServerResponse:
222 payload = {
223 "model": model,
@@ -225,6 +226,9 @@ async def send_request(
226 "temperature": 0.0,
227 }
228
229 + if conversation_id is not None:
230 + payload["conversation_id"] = conversation_id
231 +
232 if stream:
233 payload["stream"] = True
234 payload["stream_options"] = {"include_usage": False}
@@ -419,6 +423,7 @@ async def send_turn(
423 min_tokens,
424 max_tokens,
425 req_args.timeout_sec,
426 + conversation_id=conv_id,
427 )
428
429 if response.valid is False:
examples/online_serving/disaggregated_serving/disagg_proxy_multiturn.py new
+562
@@ -0,0 +1,562 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Disaggregated Prefill/Decode Proxy with Bidirectional KV Transfer
5 +
6 +This proxy sits between clients and a vLLM Prefill/Decode (P/D) deployment,
7 +routing multi-turn chat requests so that each turn reuses KV cache blocks
8 +from the previous turn's Decode node via bidirectional KV transfer.
9 +
10 +Architecture:
11 + Client ──► Proxy ──► Prefill (P) ──► Decode (D)
12 + │ │ │
13 + │ kv_transfer_params flow: │
14 + │ D finish ──► proxy caches │
15 + │ next turn ──► proxy sends │
16 + │ cached D blocks to P ──► │
17 + │ P reads D blocks (bidir) │
18 + │ P sends its blocks to D │
19 +
20 +Per-request flow:
21 + 1. Client sends chat/completions request to proxy.
22 + 2. Proxy looks up cached D block info from the previous turn
23 + (keyed by conversation_id).
24 + 3. If cache hit, proxy attaches D's block info to the request
25 + so P can read D's KV blocks instead of recomputing.
26 + 4. Proxy sends request to P (max_tokens=1, non-streaming).
27 + 5. P returns kv_transfer_params with its own block info.
28 + 6. Proxy forwards request + P's block info to D (streaming).
29 + 7. D streams the response. The final chunk includes D's
30 + kv_transfer_params, which the proxy caches for the next turn.
31 + 8. Proxy returns D's response to the client.
32 +
33 +Conversation isolation:
34 + Each request must include a ``conversation_id`` field (top-level in
35 + the JSON body) to scope the KV cache across turns. Without it, the
36 + proxy cannot link turns and falls back to no-cache behavior.
37 +
38 +Usage:
39 + python disagg_proxy_multiturn.py \\
40 + --host 0.0.0.0 --port 8000 \\
41 + --prefiller-host 10.0.0.1 --prefiller-port 8100 \\
42 + --decoder-host 10.0.0.2 --decoder-port 8200
43 +
44 +Dependencies:
45 + pip install fastapi uvicorn httpx
46 +"""
47 +
48 +from __future__ import annotations
49 +
50 +import argparse
51 +import itertools
52 +import json
53 +import logging
54 +import os
55 +import time
56 +import uuid
57 +from contextlib import asynccontextmanager
58 +from dataclasses import dataclass, field
59 +from typing import Any
60 +
61 +import httpx
62 +from fastapi import FastAPI, Request
63 +from fastapi.responses import JSONResponse, StreamingResponse
64 +
65 +# Logging
66 +logging.basicConfig(
67 + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
68 + datefmt="%Y-%m-%d %H:%M:%S",
69 + level=logging.INFO,
70 +)
71 +logger = logging.getLogger("disagg_proxy")
72 +
73 +
74 +# Data structures
75 +@dataclass
76 +class CachedKVEntry:
77 + """KV transfer parameters cached from D's response for one turn."""
78 +
79 + kv_transfer_params: dict[str, Any]
80 + timestamp: float = field(default_factory=time.time)
81 +
82 +
83 +class ConversationKVCache:
84 + """Per-conversation KV block cache.
85 +
86 + Each conversation is identified by a ``conversation_id`` supplied by
87 + the client. After D finishes a turn, its ``kv_transfer_params`` are
88 + stored here. On the next turn, the proxy retrieves them so P can
89 + read D's blocks via bidirectional KV transfer.
90 + """
91 +
92 + def __init__(self, ttl_seconds: float = 600.0) -> None:
93 + self._store: dict[str, CachedKVEntry] = {}
94 + self._ttl = ttl_seconds
95 +
96 + def get(self, conversation_id: str) -> dict[str, Any] | None:
97 + """Retrieve and consume cached KV params for a conversation.
98 +
99 + Returns a *copy* of the kv_transfer_params dict, or None.
100 + The entry is removed after retrieval (single-use).
101 + """
102 + entry = self._store.pop(conversation_id, None)
103 + if entry is None:
104 + return None
105 + age = time.time() - entry.timestamp
106 + if age > self._ttl:
107 + logger.info(
108 + "conv=%s: stale cache entry (age=%.1fs > ttl=%.1fs), discarding",
109 + conversation_id,
110 + age,
111 + self._ttl,
112 + )
113 + return None
114 + logger.info(
115 + "conv=%s: cache HIT (age=%.1fs)",
116 + conversation_id,
117 + age,
118 + )
119 + return dict(entry.kv_transfer_params)
120 +
121 + def put(self, conversation_id: str, kv_params: dict[str, Any]) -> None:
122 + """Store D's kv_transfer_params for a conversation."""
123 + self._store[conversation_id] = CachedKVEntry(
124 + kv_transfer_params=dict(kv_params), # defensive copy
125 + )
126 + logger.info(
127 + "conv=%s: cached D blocks (remote_request_id=%s, blocks=%d)",
128 + conversation_id,
129 + kv_params.get("remote_request_id", "?"),
130 + len(kv_params.get("remote_block_ids", [[]])[0])
131 + if kv_params.get("remote_block_ids")
132 + else 0,
133 + )
134 +
135 + def evict_stale(self) -> int:
136 + """Remove entries older than TTL. Returns count of evicted entries."""
137 + now = time.time()
138 + stale = [
139 + cid
140 + for cid, entry in self._store.items()
141 + if now - entry.timestamp > self._ttl
142 + ]
143 + for cid in stale:
144 + del self._store[cid]
145 + return len(stale)
146 +
147 + @property
148 + def size(self) -> int:
149 + return len(self._store)
150 +
151 +
152 +# Global state
153 +kv_cache = ConversationKVCache(
154 + ttl_seconds=450.0
155 +) # Must be < VLLM_NIXL_ABORT_REQUEST_TIMEOUT (480s)
156 +
157 +
158 +# Service client helpers
159 +@dataclass
160 +class ServiceClient:
161 + """Wrapper around an httpx.AsyncClient for a P or D instance."""
162 +
163 + client: httpx.AsyncClient
164 + host: str
165 + port: int
166 + id: int
167 +
168 +
169 +def _make_headers(request_id: str) -> dict[str, str]:
170 + """Build HTTP headers for upstream requests."""
171 + headers = {"X-Request-Id": request_id}
172 + api_key = os.environ.get("OPENAI_API_KEY")
173 + if api_key:
174 + headers["Authorization"] = f"Bearer {api_key}"
175 + return headers
176 +
177 +
178 +async def _send_to_prefill(
179 + client: ServiceClient,
180 + endpoint: str,
181 + req_data: dict[str, Any],
182 + request_id: str,
183 +) -> dict[str, Any]:
184 + """Send a non-streaming prefill request (max_tokens=1).
185 +
186 + Returns the JSON response from P, which includes kv_transfer_params.
187 + """
188 + payload = req_data.copy()
189 + payload["stream"] = False
190 + payload["max_tokens"] = 1
191 + payload.pop("max_completion_tokens", None)
192 + payload.pop("min_tokens", None)
193 + payload.pop("stream_options", None)
194 +
195 + resp = await client.client.post(
196 + endpoint,
197 + json=payload,
198 + headers=_make_headers(request_id),
199 + )
200 + resp.raise_for_status()
201 + return resp.json()
202 +
203 +
204 +async def _stream_from_decode(
205 + client: ServiceClient,
206 + endpoint: str,
207 + req_data: dict[str, Any],
208 + request_id: str,
209 + conversation_id: str,
210 +) -> tuple[str, str | None, dict[str, Any] | None, str, str | None, int | None]:
211 + """Stream response from D, capturing text and kv_transfer_params.
212 +
213 + Returns (collected_text, finish_reason, kv_params, response_id, created).
214 + Also stores kv_params in the conversation cache.
215 + """
216 + payload = req_data.copy()
217 + payload["stream"] = True
218 +
219 + collected_text = ""
220 + finish_reason: str | None = None
221 + response_id: str | None = None
222 + model_name: str | None = None
223 + created: int | None = None
224 + captured_kv: dict[str, Any] | None = None
225 +
226 + async with client.client.stream(
227 + "POST",
228 + endpoint,
229 + json=payload,
230 + headers=_make_headers(request_id),
231 + ) as resp:
232 + resp.raise_for_status()
233 + async for line in resp.aiter_lines():
234 + if not line or not line.startswith("data: "):
235 + continue
236 + if line == "data: [DONE]":
237 + break
238 + try:
239 + chunk = json.loads(line[6:])
240 + except json.JSONDecodeError:
241 + continue
242 +
243 + if response_id is None:
244 + response_id = chunk.get("id")
245 + model_name = chunk.get("model")
246 + created = chunk.get("created")
247 +
248 + for choice in chunk.get("choices", []):
249 + collected_text += choice.get("text", "")
250 + delta = choice.get("delta", {})
251 + collected_text += delta.get("content", "")
252 + if choice.get("finish_reason"):
253 + finish_reason = choice["finish_reason"]
254 +
255 + kv_params = chunk.get("kv_transfer_params")
256 + if kv_params:
257 + kv_params["remote_host"] = client.host
258 + captured_kv = kv_params
259 + if conversation_id:
260 + kv_cache.put(conversation_id, kv_params)
261 +
262 + return (
263 + collected_text,
264 + finish_reason,
265 + captured_kv,
266 + response_id or request_id,
267 + model_name,
268 + created,
269 + )
270 +
271 +
272 +async def _stream_from_decode_sse(
273 + client: ServiceClient,
274 + endpoint: str,
275 + req_data: dict[str, Any],
276 + request_id: str,
277 + conversation_id: str,
278 +):
279 + """Yield SSE chunks from D to the client, capturing kv_transfer_params."""
280 + payload = req_data.copy()
281 + payload["stream"] = True
282 +
283 + async with client.client.stream(
284 + "POST",
285 + endpoint,
286 + json=payload,
287 + headers=_make_headers(request_id),
288 + ) as resp:
289 + resp.raise_for_status()
290 + async for line in resp.aiter_lines():
291 + if not line:
292 + yield "\n"
293 + continue
294 +
295 + if line.startswith("data: ") and line != "data: [DONE]":
296 + try:
297 + chunk = json.loads(line[6:])
298 + kv_params = chunk.get("kv_transfer_params")
299 + if kv_params and conversation_id:
300 + kv_params["remote_host"] = client.host
301 + kv_cache.put(conversation_id, kv_params)
302 + except json.JSONDecodeError:
303 + pass
304 +
305 + yield line + "\n"
306 +
307 +
308 +# FastAPI application
309 +@asynccontextmanager
310 +async def lifespan(app: FastAPI):
311 + """Initialize HTTP clients for P and D instances."""
312 + app.state.prefill_clients: list[ServiceClient] = []
313 + app.state.decode_clients: list[ServiceClient] = []
314 +
315 + for i, (host, port) in enumerate(global_args.prefiller_instances):
316 + app.state.prefill_clients.append(
317 + ServiceClient(
318 + client=httpx.AsyncClient(
319 + timeout=None,
320 + base_url=f"http://{host}:{port}/v1",
321 + ),
322 + host=host,
323 + port=port,
324 + id=i,
325 + )
326 + )
327 +
328 + for i, (host, port) in enumerate(global_args.decoder_instances):
329 + app.state.decode_clients.append(
330 + ServiceClient(
331 + client=httpx.AsyncClient(
332 + timeout=None,
333 + base_url=f"http://{host}:{port}/v1",
334 + ),
335 + host=host,
336 + port=port,
337 + id=i,
338 + )
339 + )
340 +
341 + app.state.prefill_iter = itertools.cycle(range(len(app.state.prefill_clients)))
342 + app.state.decode_iter = itertools.cycle(range(len(app.state.decode_clients)))
343 +
344 + logger.info(
345 + "Ready: %d prefill, %d decode instances",
346 + len(app.state.prefill_clients),
347 + len(app.state.decode_clients),
348 + )
349 + yield
350 +
351 + for sc in app.state.prefill_clients + app.state.decode_clients:
352 + await sc.client.aclose()
353 +
354 +
355 +app = FastAPI(title="Disaggregated P/D Proxy (Multi-turn)", lifespan=lifespan)
356 +
357 +
358 +def _next_client(app_state, role: str) -> ServiceClient:
359 + if role == "prefill":
360 + return app_state.prefill_clients[next(app_state.prefill_iter)]
361 + return app_state.decode_clients[next(app_state.decode_iter)]
362 +
363 +
364 +# Request handler
365 +async def _handle_request(api_path: str, request: Request):
366 + """Core request handler for both /v1/chat/completions and /v1/completions."""
367 + req_data = await request.json()
368 + request_id = str(uuid.uuid4())
369 + conversation_id: str = req_data.pop("conversation_id", "")
370 + client_wants_stream = req_data.get("stream", False)
371 +
372 + if not conversation_id:
373 + logger.warning(
374 + "[%s] No conversation_id provided — KV cache reuse disabled "
375 + "for this request. Add a 'conversation_id' field to enable "
376 + "cross-turn KV sharing.",
377 + request_id,
378 + )
379 +
380 + # Step 1: Look up cached D blocks from the previous turn
381 + cached_kv = kv_cache.get(conversation_id) if conversation_id else None
382 +
383 + if cached_kv:
384 + # Tell P to read D's blocks (bidirectional transfer)
385 + cached_kv["do_remote_decode"] = True
386 + cached_kv["do_remote_prefill"] = False
387 + req_data["kv_transfer_params"] = cached_kv
388 + logger.info(
389 + "[%s] conv=%s: sending D's cached blocks to P (remote_request_id=%s)",
390 + request_id,
391 + conversation_id,
392 + cached_kv.get("remote_request_id"),
393 + )
394 + else:
395 + # No cached blocks — P recomputes from scratch
396 + req_data["kv_transfer_params"] = {
397 + "do_remote_decode": True,
398 + "do_remote_prefill": False,
399 + "remote_engine_id": None,
400 + "remote_block_ids": None,
401 + "remote_host": None,
402 + "remote_port": None,
403 + }
404 + logger.info("[%s] conv=%s: cache MISS", request_id, conversation_id)
405 +
406 + # Step 2: Send to Prefill node (non-streaming, max_tokens=1)
407 + prefill_client = _next_client(request.app.state, "prefill")
408 + t0 = time.time()
409 + prefill_resp = await _send_to_prefill(
410 + prefill_client,
411 + api_path,
412 + req_data,
413 + request_id,
414 + )
415 + logger.info(
416 + "[%s] Prefill done in %.0fms",
417 + request_id,
418 + (time.time() - t0) * 1000,
419 + )
420 +
421 + # Attach P's kv_transfer_params for D to read P's blocks
422 + p_kv_params = prefill_resp.get("kv_transfer_params", {})
423 + if p_kv_params:
424 + p_kv_params["remote_host"] = prefill_client.host
425 + req_data["kv_transfer_params"] = p_kv_params
426 +
427 + # Step 3: Stream from Decode node, capturing kv_transfer_params
428 + decode_client = _next_client(request.app.state, "decode")
429 +
430 + if client_wants_stream:
431 + return StreamingResponse(
432 + _stream_from_decode_sse(
433 + decode_client,
434 + api_path,
435 + req_data,
436 + request_id,
437 + conversation_id,
438 + ),
439 + media_type="text/event-stream",
440 + )
441 +
442 + text, finish_reason, _, resp_id, model, created = await _stream_from_decode(
443 + decode_client,
444 + api_path,
445 + req_data,
446 + request_id,
447 + conversation_id,
448 + )
449 +
450 + # Build OpenAI-compatible response
451 + is_chat = "messages" in req_data
452 + if is_chat:
453 + body = {
454 + "id": resp_id,
455 + "object": "chat.completion",
456 + "created": created or int(time.time()),
457 + "model": model or req_data.get("model", ""),
458 + "choices": [
459 + {
460 + "index": 0,
461 + "message": {"role": "assistant", "content": text},
462 + "finish_reason": finish_reason,
463 + }
464 + ],
465 + "usage": None,
466 + }
467 + else:
468 + body = {
469 + "id": resp_id,
470 + "object": "text_completion",
471 + "created": created or int(time.time()),
472 + "model": model or req_data.get("model", ""),
473 + "choices": [
474 + {
475 + "index": 0,
476 + "text": text,
477 + "logprobs": None,
478 + "finish_reason": finish_reason,
479 + }
480 + ],
481 + "usage": None,
482 + }
483 + return JSONResponse(content=body)
484 +
485 +
486 +# Routes
487 +@app.post("/v1/chat/completions")
488 +async def chat_completions(request: Request):
489 + return await _handle_request("/chat/completions", request)
490 +
491 +
492 +@app.post("/v1/completions")
493 +async def completions(request: Request):
494 + return await _handle_request("/completions", request)
495 +
496 +
497 +@app.get("/health")
498 +async def health():
499 + evicted = kv_cache.evict_stale()
500 + return {
501 + "status": "ok",
502 + "cached_conversations": kv_cache.size,
503 + "evicted_stale": evicted,
504 + }
505 +
506 +
507 +# CLI
508 +def parse_args() -> argparse.Namespace:
509 + p = argparse.ArgumentParser(
510 + description="Disaggregated P/D proxy with bidirectional KV transfer",
511 + )
512 + p.add_argument("--host", default="0.0.0.0")
513 + p.add_argument("--port", type=int, default=8000)
514 + p.add_argument(
515 + "--prefiller-host",
516 + "--prefiller-hosts",
517 + dest="prefiller_hosts",
518 + nargs="+",
519 + default=["localhost"],
520 + )
521 + p.add_argument(
522 + "--prefiller-port",
523 + "--prefiller-ports",
524 + dest="prefiller_ports",
525 + type=int,
526 + nargs="+",
527 + default=[8100],
528 + )
529 + p.add_argument(
530 + "--decoder-host",
531 + "--decoder-hosts",
532 + dest="decoder_hosts",
533 + nargs="+",
534 + default=["localhost"],
535 + )
536 + p.add_argument(
537 + "--decoder-port",
538 + "--decoder-ports",
539 + dest="decoder_ports",
540 + type=int,
541 + nargs="+",
542 + default=[8200],
543 + )
544 + args = p.parse_args()
545 +
546 + if len(args.prefiller_hosts) != len(args.prefiller_ports):
547 + p.error("Number of prefiller hosts must match ports")
548 + if len(args.decoder_hosts) != len(args.decoder_ports):
549 + p.error("Number of decoder hosts must match ports")
550 +
551 + args.prefiller_instances = list(zip(args.prefiller_hosts, args.prefiller_ports))
552 + args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
553 + return args
554 +
555 +
556 +if __name__ == "__main__":
557 + global global_args
558 + global_args = parse_args()
559 +
560 + import uvicorn
561 +
562 + uvicorn.run(app, host=global_args.host, port=global_args.port)
tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py new
+915
@@ -0,0 +1,915 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for bi-directional KV cache transfer between P and D nodes.
4 +
5 +Tests cover the new behaviors added by the bi-directional KV transfer PR:
6 +1. P-node scheduler lifecycle: P pulls KV from D using remote_block_ids,
7 + eliminating redundant prefill computation in multi-turn conversations.
8 +2. P-node metadata: NixlConnectorMetadata correctly populates recv metadata
9 + when P pulls KV from D (do_remote_decode=True + remote_block_ids).
10 +3. P-node worker: start_load_kv processes reqs_to_recv for KV pull from D.
11 +4. D-node request_finished: returns kv_transfer_params with remote_block_ids
12 + and remote_num_tokens so P can pull KV in future turns.
13 +5. Edge cases:
14 + - No double read after reschedule (_remote_blocks_processed flag)
15 + - remote_num_tokens bounded by block capacity (num_computed_tokens)
16 + - kv_recompute_threshold skips small transfers
17 + - P-node holds blocks for D after finishing
18 + - Cache MISS first turn falls back to local prefill
19 + - Partial remote coverage: P pulls partial, computes the rest
20 + - _remote_blocks_processed flag persists across reschedules
21 +
22 +P-node flags: do_remote_prefill=False (prefill locally),
23 +do_remote_decode=True (don't decode locally, send KV to D).
24 +P pulls KV from D when remote_block_ids is not None and
25 +external tokens > 0.
26 +"""
27 +
28 +import copy
29 +import time
30 +from unittest.mock import patch
31 +
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 (
36 + NixlConnector,
37 + NixlConnectorMetadata,
38 +)
39 +from vllm.forward_context import ForwardContext
40 +from vllm.v1.outputs import (
41 + EMPTY_MODEL_RUNNER_OUTPUT,
42 + KVConnectorOutput,
43 +)
44 +from vllm.v1.request import RequestStatus
45 +
46 +from .test_nixl_connector import FakeNixlConnectorWorker, FakeNixlWrapper
47 +from .utils import (
48 + assert_scheduler_empty,
49 + create_model_runner_output,
50 + create_request,
51 + create_scheduler,
52 + create_vllm_config,
53 + make_kv_cache_config,
54 +)
55 +
56 +pytestmark = pytest.mark.cpu_test
57 +
58 +# Common extra config for all bi-directional KV transfer tests.
59 +BIDIR_KV_EXTRA_CONFIG = {"bidirectional_kv_xfer": True, "kv_recompute_threshold": 0}
60 +
61 +
62 +# Helpers
63 +
64 +
65 +def _make_p_node_turn2_request(
66 + request_id, block_size, num_tokens, num_remote_blocks=3, remote_num_tokens=None
67 +):
68 + """Create a P-node Turn 2 request with remote_block_ids from D."""
69 + request = create_request(
70 + request_id=request_id,
71 + block_size=block_size,
72 + num_tokens=num_tokens,
73 + do_remote_decode=True,
74 + )
75 + if remote_num_tokens is None:
76 + remote_num_tokens = num_remote_blocks * block_size
77 + request.kv_transfer_params["remote_block_ids"] = [list(range(num_remote_blocks))]
78 + request.kv_transfer_params["remote_num_tokens"] = remote_num_tokens
79 + request.kv_transfer_params["remote_engine_id"] = "decode-engine"
80 + request.kv_transfer_params["remote_request_id"] = f"decode-{request_id}"
81 + request.kv_transfer_params["remote_host"] = "decode-host"
82 + request.kv_transfer_params["remote_port"] = 5678
83 + return request
84 +
85 +
86 +def _make_connector_with_fake_worker(
87 + hand_shake_latency=0, cycles_before_done=0, do_handshake=True
88 +):
89 + """Create a NixlConnector with FakeNixlConnectorWorker."""
90 + vllm_config = create_vllm_config()
91 + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
92 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
93 + connector.connector_worker = FakeNixlConnectorWorker(
94 + vllm_config,
95 + connector.engine_id,
96 + hand_shake_latency=hand_shake_latency,
97 + kv_cache_config=kv_cache_config,
98 + )
99 + worker = connector.connector_worker
100 + assert isinstance(worker.nixl_wrapper, FakeNixlWrapper)
101 + worker.nixl_wrapper.set_cycles_before_xfer_done(cycles_before_done)
102 + worker.kv_cache_layout = "HND"
103 + if do_handshake:
104 + remote_agents = worker._nixl_handshake(
105 + host="localhost",
106 + port=1234,
107 + remote_tp_size=1,
108 + expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
109 + )
110 + worker._remote_agents[FakeNixlConnectorWorker.REMOTE_ENGINE_ID] = remote_agents
111 + return connector, worker
112 +
113 +
114 +def _make_p_node_recv_metadata(request_id, local_blocks, remote_blocks):
115 + """Build NixlConnectorMetadata for P-node pulling KV from D."""
116 + meta = NixlConnectorMetadata()
117 + meta.add_new_req_to_recv(
118 + request_id=request_id,
119 + local_block_ids=(local_blocks,),
120 + kv_transfer_params={
121 + "do_remote_prefill": False,
122 + "do_remote_decode": True,
123 + "remote_block_ids": (remote_blocks,),
124 + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
125 + "remote_request_id": f"decode-{request_id}",
126 + "remote_host": "localhost",
127 + "remote_port": 1234,
128 + "remote_tp_size": 1,
129 + },
130 + )
131 + return meta
132 +
133 +
134 +def _do_load_kv(connector, metadata):
135 + """Bind metadata and call start_load_kv."""
136 + connector.bind_connector_metadata(metadata)
137 + ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
138 + connector.start_load_kv(ctx)
139 +
140 +
141 +# 1. P-node scheduler lifecycle tests
142 +
143 +
144 +def test_multiturn_lifecycle():
145 + """Full two-turn lifecycle on the P node:
146 + Turn 1: P prefills locally (do_remote_prefill=False), sends KV to D
147 + (do_remote_decode=True). Finishes LENGTH_CAPPED with remote_block_ids.
148 + Turn 2: P receives remote_block_ids from D. P pulls KV from D because
149 + remote_block_ids is not None and external tokens > 0. Computes only
150 + new tokens, finishes LENGTH_CAPPED."""
151 + vllm_config = create_vllm_config(
152 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
153 + )
154 + scheduler = create_scheduler(vllm_config)
155 + BS = vllm_config.cache_config.block_size
156 +
157 + t1 = create_request(
158 + request_id=100, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
159 + )
160 + scheduler.add_request(t1)
161 + t1_id = t1.request_id
162 + so = scheduler.schedule()
163 + mro = create_model_runner_output(reqs=[t1])
164 + eco = scheduler.update_from_output(so, mro)
165 + assert t1.status == RequestStatus.FINISHED_LENGTH_CAPPED
166 + kv = eco[0].outputs[0].kv_transfer_params
167 + assert kv and sum(len(g) for g in kv["remote_block_ids"]) > 0
168 + so = scheduler.schedule()
169 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
170 +
171 + t2 = _make_p_node_turn2_request(200, BS, int(BS * 2.5))
172 + scheduler.add_request(t2)
173 + t2_id = t2.request_id
174 + so = scheduler.schedule()
175 + assert t2.status == RequestStatus.WAITING_FOR_REMOTE_KVS
176 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
177 + so = scheduler.schedule()
178 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
179 + mro.kv_connector_output = KVConnectorOutput(finished_recving={t2_id})
180 + scheduler.update_from_output(so, mro)
181 + so = scheduler.schedule()
182 + mro = create_model_runner_output(reqs=[t2])
183 + scheduler.update_from_output(so, mro)
184 + assert t2.status == RequestStatus.FINISHED_LENGTH_CAPPED
185 + so = scheduler.schedule()
186 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
187 + so = scheduler.schedule()
188 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
189 + mro.kv_connector_output = KVConnectorOutput(finished_sending={t1_id, t2_id})
190 + scheduler.update_from_output(so, mro)
191 + assert_scheduler_empty(scheduler)
192 +
193 +
194 +def test_first_turn_no_remote_blocks():
195 + """First turn: P has no remote_block_ids from D yet.
196 + Standard local prefill, returns kv_transfer_params for future turns."""
197 + vllm_config = create_vllm_config(
198 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
199 + )
200 + scheduler = create_scheduler(vllm_config)
201 + BS = vllm_config.cache_config.block_size
202 + req = create_request(
203 + request_id=3, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
204 + )
205 + scheduler.add_request(req)
206 + req_id = req.request_id
207 + so = scheduler.schedule()
208 + assert req.status != RequestStatus.WAITING_FOR_REMOTE_KVS
209 + mro = create_model_runner_output(reqs=[req])
210 + eco = scheduler.update_from_output(so, mro)
211 + assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
212 + assert eco[0].outputs[0].kv_transfer_params is not None
213 + so = scheduler.schedule()
214 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
215 + so = scheduler.schedule()
216 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
217 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
218 + scheduler.update_from_output(so, mro)
219 + assert_scheduler_empty(scheduler)
220 +
221 +
222 +def test_abort_p_side_during_send():
223 + """P-side do_remote_decode=True: blocks held until finished_sending."""
224 + vllm_config = create_vllm_config(
225 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
226 + )
227 + scheduler = create_scheduler(vllm_config)
228 + BS = vllm_config.cache_config.block_size
229 + req = create_request(
230 + request_id=42, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
231 + )
232 + scheduler.add_request(req)
233 + req_id = req.request_id
234 + so = scheduler.schedule()
235 + mro = create_model_runner_output(reqs=[req])
236 + scheduler.update_from_output(so, mro)
237 + assert req_id in scheduler.requests
238 + so = scheduler.schedule()
239 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
240 + assert req_id in scheduler.requests
241 + so = scheduler.schedule()
242 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
243 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
244 + scheduler.update_from_output(so, mro)
245 + assert_scheduler_empty(scheduler)
246 +
247 +
248 +def test_abort_p_side_non_length_capped():
249 + """P-side abort with non-LENGTH_CAPPED → immediate block free."""
250 + vllm_config = create_vllm_config(
251 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
252 + )
253 + scheduler = create_scheduler(vllm_config)
254 + BS = vllm_config.cache_config.block_size
255 + req = create_request(
256 + request_id=44, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
257 + )
258 + req.sampling_params.max_tokens = 100
259 + req.max_tokens = 100
260 + scheduler.add_request(req)
261 + req_id = req.request_id
262 + so = scheduler.schedule()
263 + mro = create_model_runner_output(reqs=[req])
264 + scheduler.update_from_output(so, mro)
265 + scheduler.finish_requests([req_id], RequestStatus.FINISHED_ABORTED)
266 + conn = scheduler.connector.connector_scheduler
267 + assert req_id in conn._reqs_not_processed
268 + assert req_id not in scheduler.requests
269 + so = scheduler.schedule()
270 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
271 + assert_scheduler_empty(scheduler)
272 +
273 +
274 +def test_remote_blocks_exceed_prompt_tokens():
275 + """D provides more remote tokens than P's prompt needs.
276 + P caps external tokens to prompt length."""
277 + vllm_config = create_vllm_config(
278 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
279 + )
280 + scheduler = create_scheduler(vllm_config)
281 + BS = vllm_config.cache_config.block_size
282 + NUM_TOKENS = int(BS * 2.5)
283 + req = _make_p_node_turn2_request(
284 + 300, BS, NUM_TOKENS, num_remote_blocks=5, remote_num_tokens=5 * BS
285 + )
286 + scheduler.add_request(req)
287 + req_id = req.request_id
288 + so = scheduler.schedule()
289 + assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
290 + assert req.num_computed_tokens == NUM_TOKENS
291 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
292 + so = scheduler.schedule()
293 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
294 + mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
295 + scheduler.update_from_output(so, mro)
296 + so = scheduler.schedule()
297 + mro = create_model_runner_output(reqs=[req])
298 + scheduler.update_from_output(so, mro)
299 + assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
300 + so = scheduler.schedule()
301 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
302 + so = scheduler.schedule()
303 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
304 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
305 + scheduler.update_from_output(so, mro)
306 + assert_scheduler_empty(scheduler)
307 +
308 +
309 +def test_p_node_pulls_partial_last_block_from_d():
310 + """D sends remote_block_ids with partially filled last block.
311 + remote_num_tokens < len(remote_block_ids) * block_size.
312 + P pulls only remote_num_tokens worth of external tokens."""
313 + vllm_config = create_vllm_config(
314 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
315 + )
316 + scheduler = create_scheduler(vllm_config)
317 + BS = vllm_config.cache_config.block_size
318 + num_remote_blocks = 3
319 + remote_num_tokens = int(BS * 2.5)
320 + assert remote_num_tokens < num_remote_blocks * BS
321 + NUM_TOKENS = int(BS * 3.5)
322 + req = _make_p_node_turn2_request(
323 + 400,
324 + BS,
325 + NUM_TOKENS,
326 + num_remote_blocks=num_remote_blocks,
327 + remote_num_tokens=remote_num_tokens,
328 + )
329 + scheduler.add_request(req)
330 + req_id = req.request_id
331 + so = scheduler.schedule()
332 + assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
333 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
334 + so = scheduler.schedule()
335 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
336 + mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
337 + scheduler.update_from_output(so, mro)
338 + so = scheduler.schedule()
339 + assert len(scheduler.running) == 1
340 + mro = create_model_runner_output(reqs=[req])
341 + scheduler.update_from_output(so, mro)
342 + assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
343 + so = scheduler.schedule()
344 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
345 + so = scheduler.schedule()
346 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
347 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
348 + scheduler.update_from_output(so, mro)
349 + assert_scheduler_empty(scheduler)
350 +
351 +
352 +# 2. P-node metadata tests
353 +
354 +
355 +def test_add_new_req_to_recv_populates_remote_meta():
356 + """add_new_req_to_recv correctly populates RemoteMeta for P-node
357 + bi-directional KV pull from D."""
358 + meta = NixlConnectorMetadata()
359 + kv_params = {
360 + "remote_block_ids": [[0, 1, 2]],
361 + "remote_engine_id": "decode-engine",
362 + "remote_request_id": "decode-req-123",
363 + "remote_host": "decode-host",
364 + "remote_port": 5678,
365 + }
366 + local_block_ids = ([10, 11, 12],)
367 + meta.add_new_req_to_recv(
368 + request_id="test-req",
369 + local_block_ids=local_block_ids,
370 + kv_transfer_params=kv_params,
371 + )
372 + assert "test-req" in meta.reqs_to_recv
373 + rm = meta.reqs_to_recv["test-req"]
374 + assert rm.remote is not None
375 + assert rm.remote.block_ids == kv_params["remote_block_ids"]
376 + assert rm.remote.engine_id == "decode-engine"
377 + assert rm.remote.request_id == "decode-req-123"
378 + assert rm.remote.host == "decode-host"
379 + assert rm.remote.port == 5678
380 + assert rm.local_block_ids == local_block_ids
381 +
382 +
383 +def test_build_connector_meta_recv_entries():
384 + """P-node scheduler: do_remote_decode=True + remote_block_ids →
385 + _reqs_need_recv populated, build_connector_meta produces reqs_to_recv."""
386 + vllm_config = create_vllm_config(
387 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
388 + )
389 + scheduler = create_scheduler(vllm_config)
390 + BS = vllm_config.cache_config.block_size
391 + req = _make_p_node_turn2_request(1, BS, int(BS * 2.5))
392 + scheduler.add_request(req)
393 + req_id = req.request_id
394 + so = scheduler.schedule()
395 + assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
396 + meta = so.kv_connector_metadata
397 + assert isinstance(meta, NixlConnectorMetadata)
398 + assert req_id in meta.reqs_to_recv
399 + rm = meta.reqs_to_recv[req_id]
400 + assert rm.remote is not None
401 + assert rm.remote.engine_id == "decode-engine"
402 +
403 +
404 +def test_build_connector_meta_clears_reqs_need_recv():
405 + """After build_connector_meta, _reqs_need_recv is cleared."""
406 + vllm_config = create_vllm_config(
407 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
408 + )
409 + scheduler = create_scheduler(vllm_config)
410 + BS = vllm_config.cache_config.block_size
411 + req = _make_p_node_turn2_request(2, BS, int(BS * 2.5))
412 + scheduler.add_request(req)
413 + conn = scheduler.connector.connector_scheduler
414 + scheduler.schedule()
415 + assert len(conn._reqs_need_recv) == 0
416 +
417 +
418 +def test_build_connector_meta_multiple_requests():
419 + """Multiple P-node requests all included in reqs_to_recv."""
420 + vllm_config = create_vllm_config(
421 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
422 + )
423 + scheduler = create_scheduler(vllm_config)
424 + BS = vllm_config.cache_config.block_size
425 + reqs = [_make_p_node_turn2_request(10 + i, BS, int(BS * 2.5)) for i in range(3)]
426 + for r in reqs:
427 + scheduler.add_request(r)
428 + so = scheduler.schedule()
429 + meta = so.kv_connector_metadata
430 + assert isinstance(meta, NixlConnectorMetadata)
431 + assert len(meta.reqs_to_recv) == 3
432 + for r in reqs:
433 + assert r.request_id in meta.reqs_to_recv
434 +
435 +
436 +# 3. P-node worker tests (FakeNixlWrapper)
437 +
438 +
439 +@patch(
440 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
441 + FakeNixlWrapper,
442 +)
443 +def test_p_node_pull_kv_from_d(dist_init):
444 + """P node pulls KV from D via start_load_kv with reqs_to_recv."""
445 + connector, worker = _make_connector_with_fake_worker()
446 + meta = _make_p_node_recv_metadata("req-p1", [10, 11, 12], [20, 21, 22])
447 + _do_load_kv(connector, meta)
448 + assert "req-p1" in worker._recving_metadata
449 + _, done_recving = connector.get_finished(finished_req_ids=set())
450 + assert "req-p1" in done_recving
451 +
452 +
453 +@patch(
454 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
455 + FakeNixlWrapper,
456 +)
457 +def test_p_node_pull_then_send_kv(dist_init):
458 + """Full P-node bi-directional: pull KV from D → prefill →
459 + send KV back to D via notification."""
460 + connector, worker = _make_connector_with_fake_worker()
461 + meta = _make_p_node_recv_metadata("req-p2", [10, 11], [20, 21])
462 + _do_load_kv(connector, meta)
463 + _, done_recving = connector.get_finished(finished_req_ids=set())
464 + assert "req-p2" in done_recving
465 + worker._reqs_to_send["req-p2"] = time.perf_counter() + 60
466 + worker._reqs_to_process.add("req-p2")
467 + notif = f"req-p2:{worker.world_size}".encode()
468 + orig = worker.nixl_wrapper.get_new_notifs
469 + worker.nixl_wrapper.get_new_notifs = lambda: {"agent": [notif]}
470 + done_sending, _ = connector.get_finished(finished_req_ids=set())
471 + assert "req-p2" in done_sending
472 + worker.nixl_wrapper.get_new_notifs = orig
473 +
474 +
475 +@patch(
476 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
477 + FakeNixlWrapper,
478 +)
479 +def test_p_node_deferred_pull_on_no_handshake(dist_init):
480 + """P defers KV pull when no prior handshake exists."""
481 + connector, worker = _make_connector_with_fake_worker(
482 + hand_shake_latency=0, do_handshake=False
483 + )
484 + meta = _make_p_node_recv_metadata("req-p3", [10, 11], [20, 21])
485 + _do_load_kv(connector, meta)
486 + assert "req-p3" in worker._recving_metadata
487 + timeout = 3.0
488 + start = time.perf_counter()
489 + while time.perf_counter() - start < timeout:
490 + connector.bind_connector_metadata(NixlConnectorMetadata())
491 + ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
492 + connector.start_load_kv(ctx)
493 + _, done = connector.get_finished(finished_req_ids=set())
494 + if "req-p3" in done:
495 + return
496 + time.sleep(0.2)
497 + raise AssertionError("Transfer did not complete after async handshake")
498 +
499 +
500 +# 4. D-node request_finished returns kv_transfer_params (new behavior)
501 +
502 +
503 +def test_d_node_request_finished_returns_kv_params():
504 + """D-node request_finished returns kv_transfer_params with
505 + do_remote_decode=True, remote_block_ids, remote_num_tokens
506 + for P to pull. These params go directly to P node."""
507 + vllm_config = create_vllm_config(
508 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
509 + )
510 + scheduler = create_scheduler(vllm_config)
511 + BS = vllm_config.cache_config.block_size
512 + req = create_request(
513 + request_id=1, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True
514 + )
515 + scheduler.add_request(req)
516 + req_id = req.request_id
517 + so = scheduler.schedule()
518 + scheduler.update_from_output(
519 + so, create_model_runner_output(reqs=[], finished_recving={req_id})
520 + )
521 + so = scheduler.schedule()
522 + eco = scheduler.update_from_output(
523 + so, create_model_runner_output(reqs=[req], use_eos=True)
524 + )
525 + assert req.status == RequestStatus.FINISHED_STOPPED
526 + kv = eco[0].outputs[0].kv_transfer_params
527 + assert kv is not None
528 + assert kv["do_remote_decode"] is True
529 + assert kv["do_remote_prefill"] is False
530 + assert "remote_block_ids" in kv
531 + assert "remote_num_tokens" in kv
532 + assert kv["remote_num_tokens"] > 0
533 +
534 +
535 +def test_d_node_request_finished_delays_block_free():
536 + """D-node holds blocks (delay_free=True) until P reads them."""
537 + vllm_config = create_vllm_config(
538 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
539 + )
540 + scheduler = create_scheduler(vllm_config)
541 + BS = vllm_config.cache_config.block_size
542 + req = create_request(
543 + request_id=2, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True
544 + )
545 + scheduler.add_request(req)
546 + req_id = req.request_id
547 + so = scheduler.schedule()
548 + scheduler.update_from_output(
549 + so, create_model_runner_output(reqs=[], finished_recving={req_id})
550 + )
551 + so = scheduler.schedule()
552 + scheduler.update_from_output(
553 + so, create_model_runner_output(reqs=[req], use_eos=True)
554 + )
555 + assert req_id in scheduler.requests
556 + conn = scheduler.connector.connector_scheduler
557 + assert req_id in conn._reqs_need_send
558 +
559 +
560 +def test_d_node_request_finished_remote_num_tokens():
561 + """D-node kv_transfer_params includes correct remote_num_tokens."""
562 + vllm_config = create_vllm_config(
563 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
564 + )
565 + scheduler = create_scheduler(vllm_config)
566 + BS = vllm_config.cache_config.block_size
567 + req = create_request(
568 + request_id=3, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True
569 + )
570 + scheduler.add_request(req)
571 + req_id = req.request_id
572 + so = scheduler.schedule()
573 + scheduler.update_from_output(
574 + so, create_model_runner_output(reqs=[], finished_recving={req_id})
575 + )
576 + so = scheduler.schedule()
577 + eco = scheduler.update_from_output(
578 + so, create_model_runner_output(reqs=[req], use_eos=True)
579 + )
580 + kv = eco[0].outputs[0].kv_transfer_params
581 + assert kv["remote_num_tokens"] > 0
582 + assert sum(len(g) for g in kv["remote_block_ids"]) > 0
583 +
584 +
585 +def test_d_node_partial_last_block_remote_num_tokens():
586 + """D-node: remote_num_tokens < len(remote_block_ids) * block_size
587 + when last block is partially filled."""
588 + vllm_config = create_vllm_config(
589 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
590 + )
591 + scheduler = create_scheduler(vllm_config)
592 + BS = vllm_config.cache_config.block_size
593 + req = create_request(
594 + request_id=5, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True
595 + )
596 + scheduler.add_request(req)
597 + req_id = req.request_id
598 + so = scheduler.schedule()
599 + scheduler.update_from_output(
600 + so, create_model_runner_output(reqs=[], finished_recving={req_id})
601 + )
602 + so = scheduler.schedule()
603 + eco = scheduler.update_from_output(
604 + so, create_model_runner_output(reqs=[req], use_eos=True)
605 + )
606 + kv = eco[0].outputs[0].kv_transfer_params
607 + total_blocks = sum(len(g) for g in kv["remote_block_ids"])
608 + assert total_blocks == 3
609 + assert kv["remote_num_tokens"] < total_blocks * BS
610 + assert kv["remote_num_tokens"] > 0
611 +
612 +
613 +# 5. Edge case tests
614 +
615 +
616 +def test_no_double_read_blocks_after_reschedule():
617 + """Edge case 1: update_state_after_alloc called twice for the same
618 + bidirectional request (once on initial schedule, once after
619 + WAITING_FOR_REMOTE_KVS → reschedule). The _remote_blocks_processed
620 + flag must prevent the request from being added to _reqs_need_recv
621 + twice, which would cause P to read D's blocks twice."""
622 + vllm_config = create_vllm_config(
623 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
624 + )
625 + scheduler = create_scheduler(vllm_config)
626 + BS = vllm_config.cache_config.block_size
627 + req = _make_p_node_turn2_request(500, BS, int(BS * 2.5))
628 + scheduler.add_request(req)
629 + req_id = req.request_id
630 + conn = scheduler.connector.connector_scheduler
631 +
632 + # First schedule: request enters WAITING_FOR_REMOTE_KVS,
633 + # _reqs_need_recv populated then cleared by build_connector_meta.
634 + so = scheduler.schedule()
635 + assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
636 + meta = so.kv_connector_metadata
637 + assert isinstance(meta, NixlConnectorMetadata)
638 + assert req_id in meta.reqs_to_recv
639 + # _reqs_need_recv should be cleared after build_connector_meta
640 + assert len(conn._reqs_need_recv) == 0
641 +
642 + # Simulate recv completion
643 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
644 + so = scheduler.schedule()
645 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
646 + mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
647 + scheduler.update_from_output(so, mro)
648 +
649 + # Second schedule after recv: update_state_after_alloc called again.
650 + # The _remote_blocks_processed flag should prevent re-entry.
651 + so = scheduler.schedule()
652 + meta2 = so.kv_connector_metadata
653 + assert isinstance(meta2, NixlConnectorMetadata)
654 + # Must NOT be in reqs_to_recv again
655 + assert req_id not in meta2.reqs_to_recv
656 +
657 + # Clean up
658 + mro = create_model_runner_output(reqs=[req])
659 + scheduler.update_from_output(so, mro)
660 + assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
661 + so = scheduler.schedule()
662 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
663 + so = scheduler.schedule()
664 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
665 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
666 + scheduler.update_from_output(so, mro)
667 + assert_scheduler_empty(scheduler)
668 +
669 +
670 +def test_remote_num_tokens_bounded_by_blocks():
671 + """Edge case 2: D-node request_finished must return
672 + remote_num_tokens <= len(remote_block_ids) * block_size.
673 + request.num_tokens includes the last sampled token which has no KV
674 + in the cache, so remote_num_tokens must use num_computed_tokens."""
675 + vllm_config = create_vllm_config(
676 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
677 + )
678 + scheduler = create_scheduler(vllm_config)
679 + BS = vllm_config.cache_config.block_size
680 + req = create_request(
681 + request_id=501,
682 + block_size=BS,
683 + num_tokens=int(BS * 2.5),
684 + do_remote_prefill=True,
685 + )
686 + scheduler.add_request(req)
687 + req_id = req.request_id
688 + so = scheduler.schedule()
689 + scheduler.update_from_output(
690 + so, create_model_runner_output(reqs=[], finished_recving={req_id})
691 + )
692 + so = scheduler.schedule()
693 + eco = scheduler.update_from_output(
694 + so, create_model_runner_output(reqs=[req], use_eos=True)
695 + )
696 + kv = eco[0].outputs[0].kv_transfer_params
697 + assert kv is not None
698 + total_blocks = sum(len(g) for g in kv["remote_block_ids"])
699 + max_tokens_in_blocks = total_blocks * BS
700 + assert kv["remote_num_tokens"] <= max_tokens_in_blocks, (
701 + f"remote_num_tokens ({kv['remote_num_tokens']}) exceeds "
702 + f"block capacity ({max_tokens_in_blocks})"
703 + )
704 + assert kv["remote_num_tokens"] > 0
705 +
706 +
707 +def test_kv_recompute_threshold_skips_small_transfer():
708 + """Edge case 3: When remote tokens are below kv_recompute_threshold,
709 + P should skip the remote pull and compute locally instead of
710 + entering WAITING_FOR_REMOTE_KVS."""
711 + threshold = 256
712 + vllm_config = create_vllm_config(
713 + kv_connector_extra_config={
714 + "bidirectional_kv_xfer": True,
715 + "kv_recompute_threshold": threshold,
716 + },
717 + )
718 + scheduler = create_scheduler(vllm_config)
719 + BS = vllm_config.cache_config.block_size
720 +
721 + # Create request where remote tokens (48) < threshold (256)
722 + req = _make_p_node_turn2_request(
723 + 502,
724 + BS,
725 + int(BS * 2.5),
726 + num_remote_blocks=3,
727 + remote_num_tokens=3 * BS,
728 + )
729 + scheduler.add_request(req)
730 + so = scheduler.schedule()
731 + # Should NOT enter WAITING_FOR_REMOTE_KVS — threshold not met
732 + assert req.status != RequestStatus.WAITING_FOR_REMOTE_KVS
733 + assert req.status == RequestStatus.RUNNING
734 +
735 + # Clean up
736 + mro = create_model_runner_output(reqs=[req])
737 + scheduler.update_from_output(so, mro)
738 + assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
739 + so = scheduler.schedule()
740 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
741 + so = scheduler.schedule()
742 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
743 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req.request_id})
744 + scheduler.update_from_output(so, mro)
745 + assert_scheduler_empty(scheduler)
746 +
747 +
748 +def test_p_node_finished_holds_blocks_for_d():
749 + """Edge case 4: P-node finishes with FINISHED_LENGTH_CAPPED and
750 + do_remote_decode=True. P must hold blocks (delay_free=True) and
751 + return kv_transfer_params with do_remote_prefill=True so D can
752 + read P's blocks."""
753 + vllm_config = create_vllm_config(
754 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
755 + )
756 + scheduler = create_scheduler(vllm_config)
757 + BS = vllm_config.cache_config.block_size
758 + req = create_request(
759 + request_id=503,
760 + block_size=BS,
761 + num_tokens=int(BS * 2.5),
762 + do_remote_decode=True,
763 + )
764 + scheduler.add_request(req)
765 + req_id = req.request_id
766 + so = scheduler.schedule()
767 + mro = create_model_runner_output(reqs=[req])
768 + eco = scheduler.update_from_output(so, mro)
769 + assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
770 + kv = eco[0].outputs[0].kv_transfer_params
771 + assert kv is not None
772 + # P-node finished: should tell D to pull (do_remote_prefill=True)
773 + assert kv["do_remote_prefill"] is True
774 + assert kv["do_remote_decode"] is False
775 + assert "remote_block_ids" in kv
776 + assert sum(len(g) for g in kv["remote_block_ids"]) > 0
777 + # Blocks should be held (request still tracked)
778 + assert req_id in scheduler.requests
779 +
780 + # Clean up: simulate D reading and notifying
781 + so = scheduler.schedule()
782 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
783 + so = scheduler.schedule()
784 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
785 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
786 + scheduler.update_from_output(so, mro)
787 + assert_scheduler_empty(scheduler)
788 +
789 +
790 +def test_cache_miss_first_turn_no_remote_pull():
791 + """Edge case 5: First turn with do_remote_decode=True but no
792 + remote_block_ids (cache MISS). P should prefill locally with
793 + num_external_tokens=0 and not enter WAITING_FOR_REMOTE_KVS."""
794 + vllm_config = create_vllm_config(
795 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
796 + )
797 + scheduler = create_scheduler(vllm_config)
798 + BS = vllm_config.cache_config.block_size
799 + req = create_request(
800 + request_id=504,
801 + block_size=BS,
802 + num_tokens=int(BS * 2.5),
803 + do_remote_decode=True,
804 + )
805 + # No remote_block_ids set — this is a cache MISS
806 + assert req.kv_transfer_params.get("remote_block_ids") is None
807 + scheduler.add_request(req)
808 + so = scheduler.schedule()
809 + # Should NOT wait for remote KVs
810 + assert req.status != RequestStatus.WAITING_FOR_REMOTE_KVS
811 + assert req.status == RequestStatus.RUNNING
812 +
813 + # Clean up
814 + mro = create_model_runner_output(reqs=[req])
815 + scheduler.update_from_output(so, mro)
816 + so = scheduler.schedule()
817 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
818 + so = scheduler.schedule()
819 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
820 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req.request_id})
821 + scheduler.update_from_output(so, mro)
822 + assert_scheduler_empty(scheduler)
823 +
824 +
825 +def test_partial_remote_tokens_less_than_prompt():
826 + """Edge case 6: D's remote_num_tokens covers only part of P's
827 + prompt. P should pull remote_num_tokens worth of external tokens
828 + and compute the rest locally."""
829 + vllm_config = create_vllm_config(
830 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
831 + )
832 + scheduler = create_scheduler(vllm_config)
833 + BS = vllm_config.cache_config.block_size
834 + NUM_TOKENS = int(BS * 4.5) # 72 tokens
835 + # D provides only 2 blocks (32 tokens) out of 72
836 + req = _make_p_node_turn2_request(
837 + 505,
838 + BS,
839 + NUM_TOKENS,
840 + num_remote_blocks=2,
841 + remote_num_tokens=2 * BS,
842 + )
843 + scheduler.add_request(req)
844 + req_id = req.request_id
845 + so = scheduler.schedule()
846 + assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
847 + # num_computed_tokens should reflect the external tokens pulled
848 + # (capped to remote_num_tokens, not full prompt)
849 + assert req.num_computed_tokens < NUM_TOKENS
850 +
851 + # Complete the transfer and finish
852 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
853 + so = scheduler.schedule()
854 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
855 + mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
856 + scheduler.update_from_output(so, mro)
857 + so = scheduler.schedule()
858 + mro = create_model_runner_output(reqs=[req])
859 + scheduler.update_from_output(so, mro)
860 + assert req.status == RequestStatus.FINISHED_LENGTH_CAPPED
861 + so = scheduler.schedule()
862 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
863 + so = scheduler.schedule()
864 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
865 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
866 + scheduler.update_from_output(so, mro)
867 + assert_scheduler_empty(scheduler)
868 +
869 +
870 +def test_remote_blocks_processed_flag_persists():
871 + """Edge case 7: After recv completes and request is rescheduled,
872 + the _remote_blocks_processed flag in kv_transfer_params prevents
873 + the bidirectional path from re-entering _reqs_need_recv."""
874 + vllm_config = create_vllm_config(
875 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG,
876 + )
877 + scheduler = create_scheduler(vllm_config)
878 + BS = vllm_config.cache_config.block_size
879 + req = _make_p_node_turn2_request(506, BS, int(BS * 2.5))
880 + scheduler.add_request(req)
881 + req_id = req.request_id
882 + conn = scheduler.connector.connector_scheduler
883 +
884 + # First schedule → WAITING_FOR_REMOTE_KVS
885 + so = scheduler.schedule()
886 + assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
887 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
888 +
889 + # Recv completes
890 + so = scheduler.schedule()
891 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
892 + mro.kv_connector_output = KVConnectorOutput(finished_recving={req_id})
893 + scheduler.update_from_output(so, mro)
894 +
895 + # Verify the flag is set
896 + assert req.kv_transfer_params.get("_remote_blocks_processed") is True
897 +
898 + # Next schedule: update_state_after_alloc is called again.
899 + # _reqs_need_recv must NOT contain this request.
900 + so = scheduler.schedule()
901 + assert req_id not in conn._reqs_need_recv
902 + meta = so.kv_connector_metadata
903 + assert isinstance(meta, NixlConnectorMetadata)
904 + assert req_id not in meta.reqs_to_recv
905 +
906 + # Clean up
907 + mro = create_model_runner_output(reqs=[req])
908 + scheduler.update_from_output(so, mro)
909 + so = scheduler.schedule()
910 + scheduler.update_from_output(so, EMPTY_MODEL_RUNNER_OUTPUT)
911 + so = scheduler.schedule()
912 + mro = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
913 + mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})
914 + scheduler.update_from_output(so, mro)
915 + assert_scheduler_empty(scheduler)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py
+88 -9
@@ -119,6 +119,30 @@ class NixlConnectorScheduler:
119 for n_tokens, block_size in sw_sizes_tokens
120 ]
121
122 + # Threshold to decide whether to compute kv cache locally
123 + # or pull from a remote node: minimum number of remote
124 + # tokens to amortize the xfer latencies
125 + self.kv_recompute_threshold: int = int(
126 + vllm_config.kv_transfer_config.get_from_extra_config(
127 + "kv_recompute_threshold", 64
128 + )
129 + )
130 +
131 + # Bi-directional KV transfer feature supports KV block
132 + # transfers from D node to P node
133 + self.is_bidirectional_kv_xfer_enabled = (
134 + vllm_config.kv_transfer_config.get_from_extra_config(
135 + "bidirectional_kv_xfer", False
136 + )
137 + )
138 +
139 + if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0:
140 + logger.info(
141 + "Bidirectional KV transfer is enabled and the kv "
142 + "recompute threshold is set to %d tokens",
143 + self.kv_recompute_threshold,
144 + )
145 +
146 def shutdown(self):
147 self._stop_event.set()
148 if self._nixl_handshake_listener_t is not None:
@@ -298,6 +322,44 @@ class NixlConnectorScheduler:
322 if params is not None and params.get("do_remote_decode") and self._has_mamba:
323 self._truncate_mamba_request_for_prefill(request)
324
325 + if (
326 + params is not None
327 + and params.get("do_remote_decode")
328 + and params.get("remote_block_ids")
329 + and all(
330 + p in params
331 + for p in (
332 + "remote_engine_id",
333 + "remote_request_id",
334 + "remote_host",
335 + "remote_port",
336 + )
337 + )
338 + ):
339 + # Decode node has kv blocks for part of prefill request, so, provide them
340 + # as an external token count to scheduler.
341 + # The tokens will be loaded if not already present
342 + # in the prefill node local cache
343 + remote_num_tokens = params.get("remote_num_tokens") or 0
344 + count = (
345 + min(remote_num_tokens, request.num_prompt_tokens) - num_computed_tokens
346 + )
347 + if count > 0:
348 + # Check kv_recompute_threshold: skip pull if
349 + # remote tokens are below the threshold.
350 + if (
351 + self.kv_recompute_threshold > 0
352 + and count < self.kv_recompute_threshold
353 + ):
354 + logger.debug(
355 + "Skipping remote pull for %s: %d remote tokens < threshold %d",
356 + request.request_id,
357 + count,
358 + self.kv_recompute_threshold,
359 + )
360 + return 0, False
361 + return count, True
362 +
363 # No remote prefill for this request.
364 return 0, False
365
@@ -315,13 +377,19 @@ class NixlConnectorScheduler:
377 if not params:
378 return
379
318 - if params.get("do_remote_decode"):
380 + if params.get("do_remote_decode") or (
381 + params.get("do_remote_prefill") and self.is_bidirectional_kv_xfer_enabled
382 + ):
383 self._reqs_in_batch.add(request.request_id)
384 if self.use_host_buffer and params.get("do_remote_decode"):
385 # NOTE: when accelerator is not directly supported by Nixl,
386 # prefilled blocks need to be saved to host memory before transfer.
387 self._reqs_need_save[request.request_id] = request
324 - elif params.get("do_remote_prefill"):
388 + elif params.get("do_remote_prefill") or (
389 + params.get("do_remote_decode")
390 + and self.is_bidirectional_kv_xfer_enabled
391 + and not params.get("_remote_blocks_processed")
392 + ):
393 if params.get("remote_block_ids"):
394 if all(
395 p in params
@@ -333,8 +401,8 @@ class NixlConnectorScheduler:
401 )
402 ):
403 # If remote_blocks and num_external_tokens = 0, we have
336 - # a full prefix cache hit on the D worker. We need to call
337 - # send_notif in _read_blocks to free the memory on the P.
404 + # a full prefix cache hit on the local node. We need to call
405 + # send_notif in _read_blocks to free the memory on the remote node.
406
407 unhashed_local_block_ids: BlockIds = (
408 blocks.get_unhashed_block_ids_all_groups()
@@ -362,6 +430,7 @@ class NixlConnectorScheduler:
430 assert num_external_tokens == 0
431 # Only trigger 1 KV transfer per request.
432 params["do_remote_prefill"] = False
433 + params["_remote_blocks_processed"] = True
434
435 def _build_save_meta(
436 self,
@@ -450,6 +519,9 @@ class NixlConnectorScheduler:
519 if not params:
520 return False, None
521
522 + is_p_node = bool(params.get("do_remote_decode"))
523 + is_d_node = not is_p_node
524 +
525 if params.get("do_remote_prefill"):
526 # If do_remote_prefill is still True when the request is finished,
527 # update_state_after_alloc must not have been called (the request
@@ -461,9 +533,13 @@ class NixlConnectorScheduler:
533 params["do_remote_prefill"] = False
534 return False, None
535
464 - if not params.get("do_remote_decode"):
536 + if is_d_node and not self.is_bidirectional_kv_xfer_enabled:
537 return False, None
466 - if request.status != RequestStatus.FINISHED_LENGTH_CAPPED:
538 +
539 + if request.status not in (
540 + RequestStatus.FINISHED_LENGTH_CAPPED,
541 + RequestStatus.FINISHED_STOPPED,
542 + ):
543 # Also include the case of a P/D Prefill request with immediate
544 # block free (eg abort). Stop tracking this request.
545 self._reqs_not_processed.add(request.request_id)
@@ -474,7 +550,7 @@ class NixlConnectorScheduler:
550 # TODO: check whether block_ids actually ever be 0. If not we could
551 # remove the conditional below
552 delay_free_blocks = any(len(group) > 0 for group in block_ids)
477 -
553 + remote_num_tokens = 0
554 if delay_free_blocks:
555 # Prefill request on remote. It will be read from D upon completion
556 logger.debug(
@@ -492,13 +568,16 @@ class NixlConnectorScheduler:
568 # Here we "unpad" blocks to send the actual remote blocks to be read.
569 block_ids = self.get_sw_clipped_blocks(block_ids)
570
571 + remote_num_tokens = request.num_computed_tokens
572 +
573 return delay_free_blocks, dict(
496 - do_remote_prefill=True,
497 - do_remote_decode=False,
574 + do_remote_prefill=is_p_node,
575 + do_remote_decode=is_d_node,
576 remote_block_ids=block_ids,
577 remote_engine_id=self.engine_id,
578 remote_request_id=request.request_id,
579 remote_host=self.side_channel_host,
580 remote_port=self.side_channel_port,
581 tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
582 + remote_num_tokens=remote_num_tokens,
583 )