KV 이벤트

KV 이벤트 (KV Events)

KV 이벤트(KV Events)는 vLLM V1 엔진이 KV 캐시에 일어난 일(블록 저장·제거·전체 클리어)을 외부 구독자에게 배포하는 메커니즘입니다. 이 예제는 ZMQ 구독/재생(replay) 소켓을 사용해 kv-events 토픽을 듣고 이벤트 배치를 디코딩하는 구독자(subscriber)를 보여줍니다.

출처: 문서

본문

소스: https://github.com/vllm-project/vllm/tree/main/examples/features/kv_events

KV 이벤트 구독자 (Kv Events Subscriber)

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any

import msgspec
import zmq
from msgspec.msgpack import Decoder

from vllm.v1.core.kv_cache_utils import ExternalBlockHash

#
# Types copied from vllm.distributed.kv_events
#
class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
    ts: float
    events: list[Any]

class KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True):
    """Base class for all KV cache-related events."""

class BlockStored(KVCacheEvent):
    block_hashes: list[ExternalBlockHash]
    parent_block_hash: ExternalBlockHash | None
    token_ids: list[int]
    block_size: int

    lora_id: int | None
    """Deprecated: use `lora_name` for KV block key hash.
    Retained for backward compatibility.
    """

    medium: str | None
    lora_name: str | None

    extra_keys: list[tuple[Any, ...] | None] | None = None
    """Extra keys used in block hash computation, one entry per block in
    block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,
    prompt embeddings data, etc. for that specific block.
    """

    group_idx: int | None = None
    kv_cache_spec_kind: str | None = None
    kv_cache_spec_sliding_window: int | None = None
    locality: str | None = None
    session_id: str | None = None

class BlockRemoved(KVCacheEvent):
    block_hashes: list[ExternalBlockHash]
    medium: str | None
    group_idx: int | None = None
    locality: str | None = None

class AllBlocksCleared(KVCacheEvent):
    pass

class KVEventBatch(EventBatch):
    events: list[BlockStored | BlockRemoved | AllBlocksCleared]

def process_event(event_batch):
    print(f"Received event batch at {event_batch.ts}:")
    for event in event_batch.events:
        print(f"  - {event}")

def main():
    decoder = Decoder(type=KVEventBatch)
    last_seq = -1

    context = zmq.Context()

    # Set up the main subscription socket
    sub = context.socket(zmq.SUB)
    sub.connect("tcp://localhost:5557")
    topic = "kv-events"
    sub.setsockopt_string(zmq.SUBSCRIBE, topic)

    # Initialize replay socket
    replay = context.socket(zmq.REQ)
    replay.connect("tcp://localhost:5558")
    poller = zmq.Poller()
    poller.register(replay, zmq.POLLIN)

    print("Listening for KV cache events on topic:", topic)

    while True:
        try:
            if sub.poll(50):
                _, seq_bytes, payload = sub.recv_multipart()
                seq = int.from_bytes(seq_bytes, "big")

                if last_seq >= 0 and seq > last_seq + 1:
                    missed = seq - last_seq - 1
                    print(
                        f"Missed {missed} messages (last: {last_seq}, current: {seq})"
                    )

                    replay.send((last_seq + 1).to_bytes(8, "big"))

                    while poller.poll(timeout=200):
                        _, seq_bytes, replay_payload = replay.recv_multipart()
                        if not replay_payload:
                            # End of replay marker is sent as an empty frame
                            # for the payload
                            break

                        replay_seq = int.from_bytes(seq_bytes, "big")

                        if replay_seq > last_seq:
                            event_batch = decoder.decode(replay_payload)
                            process_event(event_batch)
                            last_seq = replay_seq
                            if replay_seq >= seq - 1:
                                break

                event_batch = decoder.decode(payload)
                process_event(event_batch)

            # ... do other periodic work or check for shutdown ...

        except KeyboardInterrupt:
            print("Interrupted")
            break
        except Exception as e:
            print("Error decoding message:", e)

if __name__ == "__main__":
    main()

동작 요약:

• 이벤트 타입(KVCacheEvent의 파생)은 BlockStored(블록 해시·부모 블록 해시·토큰 id·블록 크기·LoRA·medium·extra_keys·session_id 등), BlockRemoved, AllBlocksCleared입니다. 메시지는 msgspec msgpack으로 인코딩되며 Decoder(type=KVEventBatch)로 디코딩합니다.
• 주 소켓은 tcp://localhost:5557의 kv-events 토픽을 SUBSCRIBE합니다. 시퀀스 번호(seq)가 점프하면(메시지 누락) 누락분을 재생 소켓(tcp://localhost:5558)으로 요청해 받아 옵니다.
• 재생 응답의 페이로드가 빈 프레임이면 재생 종료 마커로 처리합니다.

더 알아보기 (Learn more)

• vllm.distributed.kv_events — 원본 이벤트 타입 정의
• ExternalBlockHash — 블록 해시 값
• Disaggregated Serving