NIXL KV 캐시 임대 갱신

NIXL KV 캐시 임대 갱신 (NIXL KV Cache Lease Renewal)

disaggregated prefill/decode 배포에서 Prefill 인스턴스(P)는 prefill이 끝난 뒤에도 KV 캐시 블록을 GPU 메모리에 붙잡아 두고, Decode 인스턴스(D)가 RDMA로 그것을 읽어 갈 때까지 기다려야 합니다. D가 블록을 가져가지 못할 때 안전하게 해제하는 시점을 판단하는 메커니즘이 필요한데, 이 문서는 그 lease 갱신(임대 갱신) 메커니즘을 설명합니다.

출처: 문서

본문

disaggregated prefill/decode 배포에서 Prefill 인스턴스(P)는 prefill을 마친 뒤 KV 캐시 블록을 GPU 메모리에 보관한 채, Decode 인스턴스(D)가 RDMA로 읽어 갈 때까지 기다려야 합니다. D가 블록을 가져가지 못할 때 언제 안전하게 해제할지를 결정할 메커니즘이 필요한데, 이 메커니즘이 PR #41383에서 도입되었습니다.

동기 (Motivation)

단일 타임아웃의 문제 (The single-timeout problem)

원래 설계는 P가 KV 블록을 얼마나 오래 보유할지를 제어하기 위해 단일한 큰 타임아웃(VLLM_NIXL_ABORT_REQUEST_TIMEOUT, 기본 480초)을 사용했습니다. D가 크래시하거나 연결이 끊기면, P는 수 GB에 달하는 "죽은" 블록을 회수하기 전까지 최대 8분 동안 붙잡아 두었습니다. 이 시간 동안 P에 도착하는 후속 요청들은 줄어든 캐시 용량을 마주하게 되어 성능이 저하되었습니다.

과부하의 문제 (The overloading problem)

단순히 타임아웃을 낮추면 다른 실패 모드가 나타납니다. 트래픽이 급증하면 요청이 스케줄링되기 전까지 D의 대기 큐에 오래 머물 수 있습니다. P의 고정 타임아웃이 너무 짧으면, D가 읽어 갈 기회조차 얻기 전에 블록이 해제되어 불필요한 재계산과 낭비된 prefill 작업이 발생합니다.

해법: 하트비트를 통한 임대 갱신 (Solution: lease renewal via heartbeats)

임대 갱신 메커니즘은 위 두 문제를 동시에 해결합니다. P는 prefill이 완료되면 짧은 초기 임대(기본 30초)를 부여합니다. 요청이 D에서 큐 대기 중이거나 실행 중인 동안 D는 주기적으로 P에 하트비트를 보내 임대를 연장합니다. D가 크래시하여 하트비트를 멈추면, P는 수 분을 기다리는 대신 마지막 하트비트로부터 수 초 내에 블록을 회수합니다. D가 단지 과부하 상태라면, 하트비트가 필요한 만큼 블록을 살려 둡니다.

동작 방식 (How It Works)

임대 수명 주기 (Lease lifecycle)

P가 prefill을 마치면 초기 임대 기간(kv_lease_duration, 기본 30초)으로 KV 블록을 핀(pin)합니다. 이후 블록은 다음 중 하나가 될 때까지 보유됩니다.

D가 KV 전송을 완료 — P는 read-completion 알림을 받고 즉시 블록을 해제합니다.
D가 계속 하트비트 — 각 하트비트는 임대를 lease_duration * 2/3(~20초) 연장하며, D가 정상인 동안 블록을 무기한 살려 둡니다.
하트비트가 도착하지 않음 — 임대가 만료되고 P가 블록을 회수합니다.

NIXL 알림에 편승하기 (Piggybacking on NIXL notifications)

새로운 전송 채널을 도입하는 대신, 하트비트는 NIXL의 기존 알림 시스템(send_notif / get_new_notifs)을 재사용합니다. 알림 매체는 백엔드별로 다르며, IB/RoCE에서 TCP로의 자동 폴백은 이미 NIXL이 처리합니다. D에서 특정 P로 보내는 각 하트비트 메시지 하나는 해당 D를 대신해 P에 핀된 모든 요청의 임대를 갱신합니다. 즉, iteration마다 단일 배치 메시지 하나가 여러 요청의 임대를 갱신합니다.

스케줄러 측 추적 (D) (Scheduler-side tracking (D))

핵심 통찰은 하트비팅이 요청이 실행을 위해 스케줄링될 때가 아니라 D의 스케줄러에 들어가는 즉시 시작되어야 한다는 점입니다. 부하가 심할 때 요청은 초기 임대 기간보다 훨씬 오래 대기 큐에 머물 수 있는데, 도착과 스케줄링 사이의 간격에는 상한이 없기 때문입니다.

이를 위해 D의 커넥터(NixlConnectorScheduler)는 on_new_request()를 통해 스케줄러에 연결됩니다. do_remote_prefill=True인 요청이 도착하면 커넥터는 즉시 하트비트용 추적을 시작합니다. 요청은 효율적인 배치를 위해 remote_engine_id로 그룹화됩니다. 각 스케줄러 스텝마다 하트비트 메타데이터가 NixlConnectorMetadata에 담겨 워커로 보내지며, lease_duration // 6(~5초)의 하트비트 간격으로 제한(throttle)됩니다.

추적은 KV 전송이 완료되거나(update_connector_output) 요청이 끝나거나 중단될 때(request_finished) 멈춥니다.

타이밍과 단순성 (Timing and simplicity)

하트비트 전송·처리는 백그라운드 스레드가 아니라 forward 루프에서 일어납니다. 즉 타이밍이 밀리초 단위로 정밀하지 않습니다 — 긴 모델 forward pass는 하트비트를 지연시킵니다. 하지만 임대 기간은 충분한 여유를 두고 설정됩니다. 기본 설정에서 하트비트 간격(~5초)과 임대 연장(~20초)은 전형적인 forward pass보다 최소 한 자릿수는 큽니다. 이렇게 해서 스레드 간 락 복잡성을 피하면서 설계를 단순하고 확장 가능하게 유지합니다.

정상 경로 (Happy Path)

sequenceDiagram
    participant R as Routing Proxy
    participant P as Prefill Instance
    participant D as Decode Instance

    R->>P: Request (do_remote_decode=True)
    P->>P: Run prefill
    P->>P: Grant lease (30s)
    P->>R: Response (with kv_transfer_params)

    R->>D: Request (do_remote_prefill=True)
    note over D: Request enters waiting queue
    D->>D: on_new_request() starts tracking

    loop Every ~5s (heartbeat interval)
        D->>P: Heartbeat (extend lease)
        P->>P: Lease extended by ~20s
    end

    note over D: Request scheduled for execution
    D->>P: KV transfer (RDMA read)
    P-->D: Transfer complete
    D->>D: Stop heartbeating
    P->>P: Free KV blocks

Decode 인스턴스 크래시 (Decode Instance Crash)

sequenceDiagram
    participant R as Routing Proxy
    participant P as Prefill Instance
    participant D as Decode Instance

    R->>P: Request (do_remote_decode=True)
    P->>P: Run prefill (holds onto KVs with lease)
    P->>R: Response

    R->>D: Request (do_remote_prefill=True)
    D->>P: Heartbeat (extend lease)
    D->>P: Heartbeat (extend lease)
    note over D: D crashes
    note over P: No heartbeat received
    P->>P: Lease expires (~20s, not 480s)
    P->>P: Free KV blocks

워커 측 송·수신 (Worker-side sending and receiving)

D(송신)에서: start_load_kv()(forward pass마다 호출) 동안 워커는 metadata.heartbeat_by_engine를 읽고 각 원격 P 엔진에 배치된 하트비트 알림을 보냅니다. D가 특정 엔진에 대해 아직 P와 핸드셰이크하지 않았다면(대기 큐에 있는 요청에서 흔한 경우), 백그라운드 스레드에서 선제적 핸드셰이크를 트리거합니다. 핸드셰이크가 끝나면 하트비트는 다음 스텝으로 미뤄집니다 — 조기 핸드셰이크는 결과적으로 KV 전송도 빨라지게 합니다.

P(수신)에서: _get_new_notifs()에서 P의 워커는 들어오는 NIXL 알림을 확인합니다. "HB:"로 시작하는 메시지는 _handle_heartbeat()로 라우팅되어, 각 참조 요청의 임대 만료 시각을 max(old_expiry, now + lease_extension)으로 연장합니다. 이로써 임대가 실수로 짧아지는 일이 없습니다.

양방향 KV 전송 (Bidirectional KV Transfer)

멀티턴 대화에서 양방향 KV 전송을 사용하면 D가 KV 블록을 캐시해 두고 이후 턴에서 P가 끌어다 쓸 수 있습니다. 다음 대화 턴의 시점은 클라이언트에 달려 있어(시스템이 제어하지 못함) 하트비트 기반 임대 메커니즘은 여기에 적용되지 않습니다. 대신 별도의 decoder_kv_blocks_ttl(기본 480초)이 D에 캐시된 블록의 단순 고정 타임아웃을 제공합니다. 클라이언트가 대화를 이어가는 데 너무 오래 걸리면 블록이 만료됩니다. D는 만료 시각을 P에 전달해 P가 블록 만료를 알고 재계산할 수 있게 합니다. 마감 시한은 D에서 만들어진 perf_counter 값이고 두 엔진은 서로 다른 프로세스(시계가 무관함)에서 돌기 때문에, P는 핸드셰이크 왕복으로 D에 대한 시계 오프셋(clock offset)을 추정해, 그 마감 시한을 자신의 perf_counter와 비교하기 전에 적용합니다. 향후 작업으로 이 경우에도 대칭 하트비트 메커니즘을 확장할 수 있습니다.

핵심 설계 결정 (Key Design Decisions)

요청 단위 임대, 인스턴스 단위 아님. P는 자신의 KV 블록이 어느 D에 속하는지 알 수 없습니다 — 블록 소유권은 prefill이 끝나고 라우터가 D를 고른 뒤에야 결정됩니다. 요청 레벨에서 임대하면 로드 밸런서에서 P/D 선택이 커플링되는 것을 피할 수 있습니다. 실제로 D는 같은 remote_engine_id의 요청을 그룹화해 같은 P로의 임대 연장을 배치합니다.
전송은 NIXL 알림 사용. 하트비트는 ZMQ 연결이나 API 변경을 추가하는 대신 기존 send_notif/get_new_notifs 시스템을 재사용합니다. 알림 매체는 백엔드별이며 IB/RoCE→TCP 폴백이 이미 처리되어, NIXL이 지원하는 어떤 전송 위에서도 하트비트가 동작합니다.
백그라운드 스레드 없음. 하트비트 송·수신은 forward 루프(start_load_kv / get_finished)에서 일어나 스레드 간 락 복잡성을 피합니다. 임대 기간은 forward-pass 지연(밀리초)보다 충분한 여유(초)를 제공합니다.
선제적 핸드셰이크. D가 아직 연결하지 않은 P 엔진에 하트비트를 보내야 할 때(대기 큐의 요청에서 흔함), 백그라운드 스레드에서 조기 핸드셰이크를 트리거합니다. 이는 결과적으로 KV 전송도 빨라지게 합니다.
이종 TP 지원. P TP > D TP(예: P TP=4, D TP=2)일 때 단일 D 워커가 여러 P 워커에서 가져옵니다. 특정 엔진의 모든 P 워커에 하트비트를 보내야 합니다. 반대로 D TP > P TP면 단일 P가 여러 D로부터 알림을 받으며, 단순히 TTL을 여러 번 갱신할 뿐 부작용이 없습니다.

설정 (Configuration)

임대 메커니즘은 --kv-transfer-config의 kv_connector_extra_config로 제어합니다.

Parameter	Default	Description
`kv_lease_duration`	30s	P의 초기 임대 기간. 하트비트 간격과 연장량은 자동으로 파생됩니다(`interval = duration // 6`, `extension = duration * 2 // 3`).
`decoder_kv_blocks_ttl`	480s	양방향 전송 모드에서 D에 캐시된 KV 블록의 TTL. 하트비트로 갱신되지 않는 단순 고정 타임아웃입니다.

vllm serve <MODEL> \
  --kv-transfer-config '{
    "kv_connector": "NixlConnector",
    "kv_role": "kv_producer",
    "kv_connector_extra_config": {"kv_lease_duration": 60}
  }'

NixlConnector 전체 설정 세부사항은 NixlConnector 사용 가이드를 참고하세요.

더 알아보기 (Learn more)

NixlConnector 사용 가이드 — --kv-transfer-config 전체 옵션
NIXL push-mode KV 전송