Ilker ManapandClaude Opus 5.5
117f417ef5
Plan §7.5 tier 2, opt-in per profile with kv_events. The agent starts vLLM publishing KV cache events to a loopback port it listens on (pure-Go zmq4, msgpack in both of vLLM's encodings), with sha256_cbor hashing and a fixed PYTHONHASHSEED; it mirrors each engine's blocks and sends net changes up the existing stream, with a full copy on every new one and a reset on gaps, restarts and clears. The control plane hands them to the router, for the sending node's own replicas only. The router tokenizes through a replica's /tokenize, hashes the tokens as vLLM does (checked against cbor2/hashlib golden values), and scores replicas by the share of leading blocks held, falling back to the approximate tree. New metrics beside llmgw_prefix_decisions_total. Untested against a running vLLM. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
# Publishes KV cache events the way vLLM's ZmqEventPublisher does, with
|
|
# libzmq (pyzmq) and msgspec and vLLM's own struct definitions, for
|
|
# TestInteropWithPyZMQ. Not vLLM itself: its publisher thread reduced to the
|
|
# calls that put bytes on the wire.
|
|
#
|
|
# python pypub.py tcp://127.0.0.1:PORT
|
|
import sys
|
|
import time
|
|
from typing import Any
|
|
|
|
import msgspec
|
|
import zmq
|
|
|
|
|
|
class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
|
|
ts: float
|
|
events: list[Any]
|
|
data_parallel_rank: int | None = None
|
|
|
|
|
|
class KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True):
|
|
pass
|
|
|
|
|
|
class BlockStored(KVCacheEvent):
|
|
block_hashes: list[bytes | int]
|
|
parent_block_hash: bytes | int | None
|
|
token_ids: list[int]
|
|
block_size: int
|
|
lora_id: int | None
|
|
medium: str | None
|
|
lora_name: str | None
|
|
extra_keys: list[tuple[Any, ...] | None] | None = None
|
|
group_idx: int | None = None
|
|
|
|
|
|
class BlockRemoved(KVCacheEvent):
|
|
block_hashes: list[bytes | int]
|
|
medium: str | None
|
|
group_idx: int | None = None
|
|
|
|
|
|
class AllBlocksCleared(KVCacheEvent):
|
|
pass
|
|
|
|
|
|
class KVEventBatch(EventBatch):
|
|
events: list[BlockStored | BlockRemoved | AllBlocksCleared]
|
|
|
|
|
|
# The chain TestBlockHashesAreVLLMsSha256CborChain checks.
|
|
TOKENS = [1, 23, 24, 255, 256, 65535, 65536, 151643, 4294967296] + list(range(100, 140))
|
|
HASHES = [2557436066646899323, 14514142130713257781, 480514406790448453]
|
|
|
|
ctx = zmq.Context.instance()
|
|
pub = ctx.socket(zmq.PUB)
|
|
pub.connect(sys.argv[1]) # a host, not "*": vLLM connects
|
|
enc = msgspec.msgpack.Encoder()
|
|
seq = 0
|
|
|
|
|
|
def publish(batch):
|
|
global seq
|
|
pub.send_multipart((b"kv-events", seq.to_bytes(8, "big"), enc.encode(batch)))
|
|
seq += 1
|
|
|
|
|
|
stored = KVEventBatch(ts=time.time(), events=[BlockStored(
|
|
block_hashes=HASHES, parent_block_hash=None, token_ids=TOKENS[:48],
|
|
block_size=16, lora_id=None, medium="GPU", lora_name=None, group_idx=0)])
|
|
# Until the subscriber has said what it wants, PUB drops; the Go side waits
|
|
# for three blocks, then for one, which the removal leaves.
|
|
for _ in range(40):
|
|
publish(stored)
|
|
time.sleep(0.05)
|
|
publish(KVEventBatch(ts=time.time(), events=[BlockRemoved(block_hashes=HASHES[1:], medium="GPU")]))
|
|
time.sleep(0.5)
|