@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7e4bc2cecb
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 3, 06:45 UTC

[Docs][PD][NIXL] Lease extension mechanism for blocks on P (#43099)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed May 20, 2026 at 08:58 UTC 7e4bc2cecb3a8aede2d10c86a3a1a4bd98e26100
1 file changed +136
docs/design/nixl_kv_cache_lease.md new
+136
@@ -0,0 +1,136 @@
1 +# NIXL KV Cache Lease Renewal
2 +
3 +In disaggregated prefill/decode deployments, the Prefill instance (P) must hold KV cache blocks in GPU memory after completing a prefill, waiting for the Decode instance (D) to read them via RDMA. A mechanism is needed to determine when those blocks can safely be freed when D isn't able to retrieve them. This mechanism was introduced in [PR #41383](https://github.com/vllm-project/vllm/pull/41383).
4 +
5 +## Motivation
6 +
7 +### The single-timeout problem
8 +
9 +The original design used a single, large timeout (`VLLM_NIXL_ABORT_REQUEST_TIMEOUT`, default 480s) to control how long P retains KV blocks. When D crashed or disconnected, P would hold onto potentially several GBs of "dead" blocks for up to 8 minutes before reclaiming them. During this window, subsequent requests hitting P would find reduced cache capacity and experience degraded performance.
10 +
11 +### The overloading problem
12 +
13 +Simply lowering the timeout introduces a different failure mode. Under traffic surges, requests can sit in D's waiting queue for a long time before being scheduled. If the fixed timeout on P is too short, blocks get freed before D ever has a chance to read them --- causing unnecessary recomputation and wasted prefill work.
14 +
15 +### Solution: lease renewal via heartbeats
16 +
17 +The lease renewal mechanism addresses both problems simultaneously. P grants a **short initial lease** (default 30s) when prefill completes. While a request is **queued or in-flight** on D, D **periodically sends heartbeats** to P extending the lease. If D crashes and stops heartbeating, P reclaims blocks within seconds of the last heartbeat rather than waiting minutes. If D is merely overloaded, the heartbeats keep the blocks alive for as long as needed.
18 +
19 +## How It Works
20 +
21 +### Lease lifecycle
22 +
23 +When P finishes a prefill, it pins the KV blocks with an initial lease duration (`kv_lease_duration`, default 30s). From that point, the blocks are held until either:
24 +
25 +1. **D completes the KV transfer** --- P receives a read-completion notification and frees the blocks immediately.
26 +2. **D keeps heartbeating** --- each heartbeat extends the lease by `lease_duration * 2/3` (~20s), keeping blocks alive indefinitely while D is healthy.
27 +3. **No heartbeat arrives** --- the lease expires and P reclaims the blocks.
28 +
29 +### Piggybacking on NIXL notifications
30 +
31 +Rather than introducing a new transport channel, heartbeats reuse NIXL's existing notification system (`send_notif` / `get_new_notifs`). The notification medium is backend-specific, with automatic fallback from IB/RoCE to TCP already handled by NIXL. Each single heartbeat message sent from D to a particular P renews all requests pinned in P on behalf of that D --- in other words, a single batched message per iteration renews the lease of multiple requests.
32 +
33 +### Scheduler-side tracking (D)
34 +
35 +A critical insight is that heartbeating must start **as soon as a request enters D's scheduler** --- not when it gets scheduled for execution. Under heavy load, a request may sit in the waiting queue for much longer than the initial lease duration, and the gap between arrival and scheduling is unbounded.
36 +
37 +To achieve this, D's connector (`NixlConnectorScheduler`) hooks into the scheduler via `on_new_request()`. When a request with `do_remote_prefill=True` arrives, the connector immediately starts tracking it for heartbeats. Requests are grouped by `remote_engine_id` for efficient batching. On each scheduler step, heartbeat metadata is packaged into `NixlConnectorMetadata` and sent to the worker, throttled by a heartbeat interval of `lease_duration // 6` (~5s).
38 +
39 +Tracking stops when either the KV transfer completes (via `update_connector_output`) or the request finishes/aborts (via `request_finished`).
40 +
41 +### Timing and simplicity
42 +
43 +Heartbeat sending and processing happen **in the forward loop**, not in a background thread. This means timing is not millisecond-precise --- a long model forward pass will delay heartbeats. However, the lease durations are configured with sufficient margin: with default settings, the heartbeat interval (~5s) and lease extension (~20s) are at least an order of magnitude larger than a typical forward pass. This avoids lock complexity between threads while keeping the design simple and extensible.
44 +
45 +## Happy Path
46 +
47 +```mermaid
48 +sequenceDiagram
49 + participant R as Routing Proxy
50 + participant P as Prefill Instance
51 + participant D as Decode Instance
52 +
53 + R->>P: Request (do_remote_decode=True)
54 + P->>P: Run prefill
55 + P->>P: Grant lease (30s)
56 + P->>R: Response (with kv_transfer_params)
57 +
58 + R->>D: Request (do_remote_prefill=True)
59 + note over D: Request enters waiting queue
60 + D->>D: on_new_request() starts tracking
61 +
62 + loop Every ~5s (heartbeat interval)
63 + D->>P: Heartbeat (extend lease)
64 + P->>P: Lease extended by ~20s
65 + end
66 +
67 + note over D: Request scheduled for execution
68 + D->>P: KV transfer (RDMA read)
69 + P-->D: Transfer complete
70 + D->>D: Stop heartbeating
71 + P->>P: Free KV blocks
72 +```
73 +
74 +## Decode Instance Crash
75 +
76 +```mermaid
77 +sequenceDiagram
78 + participant R as Routing Proxy
79 + participant P as Prefill Instance
80 + participant D as Decode Instance
81 +
82 + R->>P: Request (do_remote_decode=True)
83 + P->>P: Run prefill (holds onto KVs with lease)
84 + P->>R: Response
85 +
86 + R->>D: Request (do_remote_prefill=True)
87 + D->>P: Heartbeat (extend lease)
88 + D->>P: Heartbeat (extend lease)
89 + note over D: D crashes
90 + note over P: No heartbeat received
91 + P->>P: Lease expires (~20s, not 480s)
92 + P->>P: Free KV blocks
93 +```
94 +
95 +### Worker-side sending and receiving
96 +
97 +**On D (sending):** During `start_load_kv()` (called every forward pass), the worker reads `metadata.heartbeat_by_engine` and sends batched heartbeat notifications to each remote P engine. If D hasn't yet handshaked with P for a given engine (common for requests still in the waiting queue), it triggers a **proactive handshake** in a background thread.
98 +The heartbeat is deferred to the next step once the handshake completes --- the early handshake also **speeds up the eventual KV transfer.**
99 +
100 +**On P (receiving):** In `_get_new_notifs()`, P's worker checks incoming NIXL notifications. Messages starting with `"HB:"` are routed to `_handle_heartbeat()`, which extends the lease expiry for each referenced request using `max(old_expiry, now + lease_extension)`. This ensures leases are never accidentally shortened.
101 +
102 +## Bidirectional KV Transfer
103 +
104 +For multi-turn conversations, [bidirectional KV transfer](../features/disagg_prefill.md) allows D to cache KV blocks that P can pull from on subsequent turns. Since the timing of the next conversational turn is **client-dependent** (not controlled by the system), the heartbeat-based lease mechanism does not apply here. Instead, a separate `decoder_kv_blocks_ttl` (default 480s) provides a simple fixed timeout for blocks cached on D. If the client takes too long to continue the conversation, the blocks expire and P recomputes. Future work may extend a symmetric heartbeat mechanism to this case.
105 +
106 +## Key Design Decisions
107 +
108 +- **Per-request leasing, not per-instance.** P has no notion of which D its KV blocks belong to --- block ownership is only resolved after prefill completes and the router selects a D. Leasing at the request level avoids coupling P/D selection in the load balancer. In practice, D batches lease extensions toward the same P by grouping requests with the same `remote_engine_id`.
109 +
110 +- **NIXL notifications as transport.** Heartbeats reuse the existing `send_notif`/`get_new_notifs` system rather than adding ZMQ connections or API changes. The notification medium is backend-specific with IB/RoCE-to-TCP fallback already handled, making heartbeats work across any NIXL-supported transport.
111 +
112 +- **No background thread.** Heartbeat sending and processing happen in the forward loop (`start_load_kv` / `get_finished`). This avoids lock complexity between threads. Lease durations provide sufficient margin over forward-pass latency (seconds vs. milliseconds).
113 +
114 +- **Proactive handshake.** When D needs to heartbeat a P engine it hasn't connected to yet (common for requests still in the waiting queue), it triggers an early handshake in a background thread. This also speeds up the eventual KV transfer.
115 +
116 +- **Heterogeneous TP support.** When P TP > D TP (e.g., P TP=4, D TP=2), a single D worker pulls from multiple P workers. Heartbeats must be sent to all P workers for a given engine. Conversely, when D TP > P TP, a single P receives notifications from multiple Ds, which simply refreshes the TTL multiple times with no downside.
117 +
118 +## Configuration
119 +
120 +The lease mechanism is controlled through `kv_connector_extra_config` in `--kv-transfer-config`:
121 +
122 +| Parameter | Default | Description |
123 +|-------------------------|---------|---------------------------------------------------------------------------------------------------------------------------------------------------------------|
124 +| `kv_lease_duration` | 30s | Initial lease duration on P. Heartbeat interval and extension amount are derived automatically (`interval = duration // 6`, `extension = duration * 2 // 3`). |
125 +| `decoder_kv_blocks_ttl` | 480s | TTL for KV blocks cached on D in bidirectional transfer mode. Simple fixed timeout, not renewed via heartbeats. |
126 +
127 +```bash
128 +vllm serve <MODEL> \
129 + --kv-transfer-config '{
130 + "kv_connector": "NixlConnector",
131 + "kv_role": "kv_both",
132 + "kv_connector_extra_config": {"kv_lease_duration": 60}
133 + }'
134 +```
135 +
136 +For full NixlConnector configuration details, see the [NixlConnector Usage Guide](../features/nixl_connector_usage.md).