@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 136 lines 9.9 KB
Rendered Raw
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. D communicates back the expiry time so P can know when blocks are expired and recompute. Because the deadline is a `perf_counter` value produced on D and the two engines run in separate processes (with unrelated clocks), P estimates the clock offset to D from the handshake round-trip and applies it before comparing the deadline against its own `perf_counter`. 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_producer",
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).