feat(grpc): add vLLM KV cache event support (SubscribeKvEvents bridge)

docs/getting-started/kv-events-cache-aware.md

@@ -91,6 +91,30 @@ All fields and defaults match SGLang's `KVEventsConfig` (see `python/sglang/srt/
 
 For data-parallel deployments, the actual TCP port becomes `endpoint_port + dp_rank` (rank 0 keeps the configured port).
 
+### Alternative: launch a vLLM worker
+
+vLLM publishes KV cache events on a ZMQ socket; enable them with `--kv-events-config` and run the worker in SMG gRPC mode:
+
+```bash
+pip install "smg-grpc-servicer[vllm]"
+
+# `--grpc` runs vLLM in SMG gRPC mode (loads smg-grpc-servicer);
+# --kv-events-config turns on KV-event publishing:
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+  --grpc \
+  --kv-events-config '{"enable_kv_cache_events": true, "publisher": "zmq", "endpoint": "tcp://*:5557", "topic": "kv-events"}'
+```
+
+Event-driven routing needs the worker in **SMG gRPC mode** (`--grpc`) — KV events stream over the `SubscribeKvEvents` RPC, so an HTTP worker can't participate. See [gRPC Workers](grpc-workers.md) for additional launch flags (host, port, TP size).
+
+| Field | Why |
+|---|---|
+| `enable_kv_cache_events: true` | vLLM-specific master switch. Without it no events are published even if a publisher is set. |
+| `publisher: "zmq"` | Selects the ZMQ publisher the servicer bridges. |
+| `endpoint` / `topic` | ZMQ `PUB` address and topic prefix. For data-parallel, the port is `endpoint_port + dp_rank`; SMG currently consumes rank 0. |
+
+Everything downstream — SMG flags, block-size learning, and the verification logs — is unchanged; `KvEventMonitor` consumes the events the same way for any gRPC worker.
+
 ---
 
 ## Step 2 — Launch SMG

grpc_servicer/pyproject.toml

@@ -30,7 +30,7 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
-vllm = ["vllm>=0.19.0"]
+vllm = ["vllm>=0.19.0", "pyzmq>=25.0.0", "msgspec>=0.18.0"]
 sglang = ["sglang>=0.5.10"]
 # smg-grpc-proto>=0.4.7 is the first release that ships mlx_engine_pb2;
 # without this floor, installing [mlx] against an older proto build would

grpc_servicer/smg_grpc_servicer/vllm/kv_events.py

@@ -0,0 +1,178 @@
+"""vLLM KV-cache-event → proto conversion and ZMQ streaming.
+
+Engine-free: imports only stdlib + the generated proto, and dispatches vLLM
+events by class name (BlockStored / BlockRemoved / AllBlocksCleared), so it
+needs no vLLM import and is unit-testable without vLLM installed.
+"""
+
+import logging
+from collections.abc import AsyncIterator, Awaitable, Callable
+
+from smg_grpc_proto.generated import common_pb2
+
+logger = logging.getLogger(__name__)
+
+_U64_MASK = 0xFFFFFFFFFFFFFFFF
+_I64_SIGN_BIT = 0x8000000000000000
+_U64_MODULUS = 0x10000000000000000
+
+
+def to_int64(value: int | bytes) -> int:
+    """Reduce a vLLM block hash to a signed int64 for the proto block_hash field.
+
+    vLLM's block hash is ``int | bytes`` (sha256 bytes when int hashes are
+    disabled); bytes are read big-endian. SMG uses the hash only as a node
+    identity, so the 64-bit reduction is safe as long as it stays deterministic.
+    """
+    if isinstance(value, (bytes, bytearray)):
+        value = int.from_bytes(value, "big")
+    masked = value & _U64_MASK
+    if masked >= _I64_SIGN_BIT:
+        masked -= _U64_MODULUS
+    return masked
+
+
+def endpoint_for_rank(endpoint: str, dp_rank: int) -> str:
+    """Resolve a vLLM KV-events PUB endpoint to a connectable SUB address.
+
+    Bind wildcards (``*``, ``0.0.0.0``) are rewritten to ``127.0.0.1`` (the
+    latter is not connectable on macOS/Windows). For data-parallel deployments
+    each rank publishes on ``base_port + dp_rank``; non-tcp endpoints (ipc://,
+    inproc://) get the wildcard substituted but no port arithmetic.
+    """
+    resolved = endpoint.replace("*", "127.0.0.1").replace("0.0.0.0", "127.0.0.1")
+    if resolved.startswith("tcp://") and dp_rank:
+        host, sep, port = resolved.rpartition(":")
+        if sep and port.isdigit():
+            return f"{host}:{int(port) + dp_rank}"
+    return resolved
+
+
+def resolve_kv_events_config(engine: object):
+    """Return the vLLM KVEventsConfig iff ZMQ event publishing is enabled, else None.
+
+    Reads ``engine.vllm_config.kv_events_config`` via getattr so it tolerates any
+    vLLM version and is testable with a fake engine.
+    """
+    vllm_config = getattr(engine, "vllm_config", None)
+    cfg = getattr(vllm_config, "kv_events_config", None)
+    if cfg is None:
+        return None
+    if not getattr(cfg, "enable_kv_cache_events", False):
+        logger.info("vLLM KV cache events not enabled; SubscribeKvEvents disabled")
+        return None
+    if getattr(cfg, "publisher", None) != "zmq":
+        logger.info(
+            "vLLM KV events publisher is %r, not 'zmq'; SubscribeKvEvents disabled",
+            getattr(cfg, "publisher", None),
+        )
+        return None
+    logger.info("vLLM KV events enabled: endpoint=%s", getattr(cfg, "endpoint", "?"))
+    return cfg
+
+
+def convert_event(event: object, event_id: int) -> common_pb2.KvCacheEvent | None:
+    """Convert one decoded vLLM event to a proto KvCacheEvent (or None if unknown)."""
+    name = type(event).__name__
+
+    if name == "BlockStored":
+        block_size = int(event.block_size)
+        blocks = []
+        for i, block_hash in enumerate(event.block_hashes):
+            start = i * block_size
+            end = start + block_size
+            block = common_pb2.KvBlock(
+                block_hash=to_int64(block_hash),
+                token_ids=list(event.token_ids[start:end]),
+                block_size=block_size,
+            )
+            lora_id = getattr(event, "lora_id", None)
+            if lora_id is not None:
+                block.lora_id = to_int64(lora_id)
+            blocks.append(block)
+        stored = common_pb2.KvBlocksStored(blocks=blocks)
+        parent = getattr(event, "parent_block_hash", None)
+        if parent is not None:
+            stored.parent_block_hash = to_int64(parent)
+        return common_pb2.KvCacheEvent(event_id=event_id, stored=stored)
+
+    if name == "BlockRemoved":
+        return common_pb2.KvCacheEvent(
+            event_id=event_id,
+            removed=common_pb2.KvBlocksRemoved(
+                block_hashes=[to_int64(h) for h in event.block_hashes]
+            ),
+        )
+
+    if name == "AllBlocksCleared":
+        return common_pb2.KvCacheEvent(event_id=event_id, cleared=common_pb2.KvCacheCleared())
+
+    logger.debug("Unknown vLLM KV event type %r, skipping", name)
+    return None
+
+
+def convert_batch(
+    raw_batch: object, seq_num: int, event_id_start: int
+) -> tuple[common_pb2.KvEventBatch, int]:
+    """Convert a decoded vLLM KVEventBatch to a proto KvEventBatch.
+
+    Returns the proto batch and the new event-id counter. The counter advances
+    once per input event (even if unconvertible) so ids stay monotonic.
+    """
+    proto = common_pb2.KvEventBatch(sequence_number=seq_num, timestamp=raw_batch.ts)
+    dp_rank = getattr(raw_batch, "data_parallel_rank", None)
+    if dp_rank is not None:
+        proto.dp_rank = dp_rank
+
+    event_id = event_id_start
+    for event in raw_batch.events:
+        event_id += 1
+        proto_event = convert_event(event, event_id)
+        if proto_event is not None:
+            proto.events.append(proto_event)
+    return proto, event_id
+
+
+async def stream_kv_events(
+    sub_socket: object,
+    decode: Callable[[bytes], object],
+    send_initial_metadata: Callable[[], Awaitable[None]],
+    is_cancelled: Callable[[], bool],
+    *,
+    recv_timeout: float = 1.0,
+) -> AsyncIterator[common_pb2.KvEventBatch]:
+    """Core ZMQ→proto streaming loop, decoupled from vLLM and gRPC types.
+
+    Args:
+        sub_socket: a connected ``zmq.asyncio`` SUB socket (duck-typed; only
+            ``poll()`` and ``recv_multipart()`` are used). The caller owns the
+            socket lifecycle (this function never closes it).
+        decode: bytes → decoded vLLM batch (e.g. ``msgspec.msgpack.Decoder(KVEventBatch).decode``).
+        send_initial_metadata: awaitable called once before the first recv so the
+            gRPC client's ``subscribe_kv_events().await`` resolves promptly.
+        is_cancelled: returns True when the RPC is cancelled; loop then exits.
+        recv_timeout: poll timeout so cancellation is observed even when idle.
+
+    Yields proto KvEventBatch using the ZMQ publisher's native sequence numbers.
+    """
+    await send_initial_metadata()
+    event_id = 0
+    while not is_cancelled():
+        # poll() before recv: cancelling a zmq.asyncio recv future does not
+        # cancel the in-flight ZMQ recv and can drop an already-dequeued message.
+        if not await sub_socket.poll(timeout=int(recv_timeout * 1000)):
+            continue
+        frames = await sub_socket.recv_multipart()
+
+        # ZMQ multipart: [topic, 8-byte big-endian seq, msgpack payload].
+        if len(frames) < 3:
+            continue
+        zmq_seq = int.from_bytes(frames[1], "big")
+        try:
+            raw_batch = decode(frames[2])
+        except Exception as e:  # noqa: BLE001 - one bad frame must not kill the stream
+            logger.warning("Failed to decode vLLM KV event batch: %s", e)
+            continue
+
+        proto_batch, event_id = convert_batch(raw_batch, zmq_seq, event_id)
+        yield proto_batch

grpc_servicer/smg_grpc_servicer/vllm/servicer.py

@@ -5,6 +5,7 @@
 Implements the VllmEngine gRPC service on top of vLLM's EngineClient.
 """
 
+import asyncio
 import hashlib
 import itertools
 import json
@@ -14,11 +15,15 @@
 from pathlib import Path
 
 import grpc
+import msgspec
 import torch
+import zmq
+import zmq.asyncio
 from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
 from smg_grpc_proto.generated import common_pb2
 from transformers import BatchFeature
 from vllm import PoolingParams, SamplingParams, TokensPrompt
+from vllm.distributed.kv_events import KVEventBatch
 from vllm.engine.protocol import EngineClient
 from vllm.inputs.engine import MultiModalInput as VllmMultiModalInput
 from vllm.inputs.engine import mm_input, tokens_input
@@ -33,6 +38,11 @@
 from vllm.sampling_params import RequestOutputKind, StructuredOutputsParams
 
 from smg_grpc_servicer.tokenizer_bundle import CHUNK_SIZE, build_tokenizer_zip
+from smg_grpc_servicer.vllm.kv_events import (
+    endpoint_for_rank,
+    resolve_kv_events_config,
+    stream_kv_events,
+)
 from smg_grpc_servicer.vllm.kv_transfer import params_from_request, params_to_response_fields
 
 logger = init_logger(__name__)
@@ -136,6 +146,9 @@ def __init__(self, async_llm: EngineClient, start_time: float):
         """
         self.engine = async_llm
         self.start_time = start_time
+        # Resolve KV-event publishing config from the engine. Non-None only when
+        # vLLM was started with --kv-events-config enabling the ZMQ publisher.
+        self._kv_events_config = resolve_kv_events_config(async_llm)
         logger.info("VllmEngineServicer initialized")
 
     async def Generate(
@@ -993,3 +1006,54 @@ def _complete_response(
                 **stop_kwargs,
             ),
         )
+
+    async def SubscribeKvEvents(
+        self,
+        request: common_pb2.SubscribeKvEventsRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> AsyncIterator[common_pb2.KvEventBatch]:
+        """Bridge vLLM's in-process ZMQ KV cache events to a gRPC stream.
+
+        The ZMQ publisher's sequence numbers are used directly as the gRPC
+        batch sequence numbers.
+        """
+        if self._kv_events_config is None:
+            await context.abort(
+                grpc.StatusCode.UNIMPLEMENTED,
+                "KV cache events not enabled. Start vLLM with "
+                "--kv-events-config "
+                '\'{"enable_kv_cache_events": true, "publisher": "zmq"}\'',
+            )
+
+        config = self._kv_events_config
+
+        # For DP attention each rank publishes on port + rank with independent
+        # sequence counters; subscribing to several on one socket interleaves
+        # them and breaks gap detection. Subscribe to rank 0 only for now.
+        # TODO(phase3): per-rank virtual workers or merged renumbering.
+        pub_endpoint = endpoint_for_rank(config.endpoint, 0)
+
+        zmq_ctx = zmq.asyncio.Context.instance()
+        sub_socket = zmq_ctx.socket(zmq.SUB)
+        sub_socket.subscribe(config.topic.encode("utf-8"))
+        sub_socket.connect(pub_endpoint)
+        logger.info("SubscribeKvEvents: connected to ZMQ endpoint %s", pub_endpoint)
+
+        decoder = msgspec.msgpack.Decoder(KVEventBatch)
+
+        try:
+            async for proto_batch in stream_kv_events(
+                sub_socket,
+                decoder.decode,
+                lambda: context.send_initial_metadata(()),
+                context.cancelled,
+            ):
+                yield proto_batch
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            logger.exception("SubscribeKvEvents failed")
+            await context.abort(grpc.StatusCode.INTERNAL, str(e))
+        finally:
+            sub_socket.close(linger=0)
+            logger.info("SubscribeKvEvents: stream closed")