[Feature] Add process-local in-memory KV indexer and Router integration (#33370)

Co-authored-by: Wu, Yutong <yutong.wu@amd.com> Co-authored-by: TianDi101 <ditian12@amd.com> Co-authored-by: Zhangheng <hzh0425@apache.org>
2026-08-20 10:45:35 +08:00
parent 238ba40c27
commit 360d10d6bc
42 changed files with 6603 additions and 174 deletions
@@ -0,0 +1,185 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+syntax = "proto3";
+
+package kv_indexer.v1;
+
+// Storage tier for externally managed KV cache blocks.
+enum TierType {
+  TIER_UNKNOWN = 0;
+  TIER_HBM = 1;
+  TIER_DRAM = 2;
+  TIER_SSD = 3;
+}
+
+// KV components are a fixed set; each component's rule is a property of its
+// type, not a per-worker choice:
+//   FULL  (bit 0, value 1) -- resident on every prefix block.
+//   SWA   (bit 1, value 2) -- resident contiguously over at least
+//                             `swa_window_tokens` tokens ending at the boundary,
+//                             or in an unbroken run from the head.
+//   MAMBA (bit 2, value 4) -- resident on the boundary block only.
+// Carried as a bitmask; FULL is always present on a stored block.
+
+// A worker's versioned cache spec: which components its cache holds and the
+// parameters the fixed rules need. Absent means legacy / full-only behaviour.
+message WorkerCacheSpec {
+  uint32 version = 1;
+  // Bitmask of components this worker's cache holds (FULL|SWA|MAMBA); selects
+  // which fixed rules gate the prefix.
+  uint32 components = 2;
+  // Sliding-window size in tokens; used only when SWA is present.
+  uint32 swa_window_tokens = 3;
+  // Per-component servable tiers, each a bitmask of (1 << TierType), intersected
+  // with the indexer's servable set (V1: HBM+DRAM) at query time.
+  uint32 full_tier_mask = 4;
+  uint32 swa_tier_mask = 5;
+  uint32 mamba_tier_mask = 6;
+}
+
+message MatchExternalKvRequest {
+  // Block-level hashes matched against placement metadata. This RPC returns
+  // placement matches only; it does not compute the longest reusable prefix.
+  repeated sfixed64 hashes = 1 [packed = true];
+
+  // When true, only hashes that are actually matched should increase hit
+  // counters. Diagnostic callers should leave this false.
+  bool count_as_hit = 2;
+}
+
+// The kind of mutation a single ExternalKvAction carries. A whole SGLang
+// KVEventBatch is applied in one call while preserving exact action order.
+enum ExternalKvActionType {
+  ACTION_UNKNOWN = 0;
+  ACTION_REPORT = 1;
+  ACTION_REVOKE = 2;
+  ACTION_CLEAR_ALL_AT_TIER = 3;
+}
+
+message ExternalKvAction {
+  ExternalKvActionType type = 1;
+  TierType tier = 2;
+  // Non-empty for REPORT/REVOKE; ignored for CLEAR_ALL_AT_TIER.
+  repeated sfixed64 hashes = 3 [packed = true];
+
+  // REPORT only. Per-hash resident component bitmask at this tier, a REPLACE
+  // snapshot index-aligned with `hashes`. Empty means every hash is a legacy
+  // whole-block store (no component detail). REVOKE/CLEAR ignore it.
+  repeated uint32 component_masks = 4;
+
+  // REPORT only. Per-hash token count (block_size), index-aligned with `hashes`,
+  // used to accumulate SWA trailing windows. Empty when not supplied (legacy).
+  repeated uint32 block_sizes = 5;
+}
+
+message ApplyExternalKvBatchRequest {
+  string worker_id = 1;
+  // The SGLang batch sequence number, monotonic per worker. Observability only:
+  // every batch is applied, with no deduplication, fencing, or checkpointing.
+  uint64 seq = 2;
+  repeated ExternalKvAction actions = 3;
+  // The worker's KV-transfer address, used to populate MatchExternalKvResponse.
+  // Supplied by the bridge; may be empty, and then matches carry no address.
+  string worker_address = 4;
+
+  // Field 5 held the worker incarnation token used for restart fencing.
+  reserved 5;
+
+  // The worker's component cache spec, carried on every batch so it self-heals
+  // across reconnects and indexer restarts. Absent means legacy / full-only.
+  WorkerCacheSpec cache_spec = 6;
+}
+
+// Deliberately empty: applies are unconditional, so there is nothing to report
+// back. Kept as a message so the RPC signature stays stable.
+message ApplyExternalKvBatchResponse {
+  // Fields 1-3 held the durable seq checkpoint and duplicate flag.
+  reserved 1, 2, 3;
+}
+
+message TierHashes {
+  TierType tier = 1;
+  repeated sfixed64 hashes = 2 [packed = true];
+  // Diagnostic snapshots aligned with `hashes`. Empty only for backends that
+  // cannot expose component detail.
+  repeated uint32 component_masks = 3;
+  repeated uint32 block_sizes = 4;
+}
+
+message ExternalKvNodeMatch {
+  string worker_id = 1;
+  string address = 2;
+  repeated TierHashes hashes_by_tier = 3;
+}
+
+message MatchExternalKvResponse {
+  repeated ExternalKvNodeMatch matches = 1;
+}
+
+message MatchExternalKvPrefixRequest {
+  // Block hashes in prompt order. hashes[0] MUST be the request's first block:
+  // prefix matching starts there and stops at the first block a worker is
+  // missing, so a misordered list silently truncates every match.
+  repeated sfixed64 hashes = 1 [packed = true];
+
+  // Caller-supplied ceiling on how many leading blocks to consider. 0 means no
+  // caller ceiling: the server considers every block sent. Unlike the mutating
+  // RPCs, a prefix query is not rejected for length; its scan holds O(1) state
+  // per candidate worker, so length costs time rather than memory.
+  uint32 max_blocks = 2;
+}
+
+message ExternalKvPrefixMatch {
+  // The worker's router-facing routing identity, NOT its KV-transfer address.
+  // The router intersects this byte-for-byte with the worker URLs it registered,
+  // so a mismatch silently disables cache-aware routing. Workers with an empty
+  // address are excluded from this response entirely.
+  string worker_address = 1;
+
+  // Largest n such that this worker holds hashes[0..n) contiguously.
+  uint32 matched_prefix_blocks = 2;
+
+  // Opaque worker id, carried for the caller's logs only; not a routing key.
+  string worker_id = 3;
+}
+
+message MatchExternalKvPrefixResponse {
+  // Matches sorted by matched_prefix_blocks, descending.
+  repeated ExternalKvPrefixMatch matches = 1;
+
+  // The longest contiguous prefix held by any single worker (0 when no match).
+  uint32 best_prefix_blocks = 2;
+
+  // How many blocks the server actually read placement for, so early termination
+  // and truncation are observable. NOT part of the prefix semantics: backends may
+  // report different values for the same matches.
+  uint32 blocks_read = 3;
+}
+
+message HitCountEntry {
+  sfixed64 hash = 1;
+  uint64 hit_count_total = 2;
+}
+
+message GetExternalKvHitCountsRequest {
+  repeated sfixed64 hashes = 1 [packed = true];
+}
+
+message GetExternalKvHitCountsResponse {
+  repeated HitCountEntry entries = 1;
+}
+
+service KVIndexer {
+  // The sole mutation API: applies an ordered SGLang KVEventBatch.
+  rpc ApplyExternalKvBatch(ApplyExternalKvBatchRequest) returns (ApplyExternalKvBatchResponse);
+
+  rpc MatchExternalKv(MatchExternalKvRequest) returns (MatchExternalKvResponse);
+
+  // Answers, per worker, how long a contiguous prefix of the request it holds.
+  // The indexer never picks a worker: it cannot see the router's health, load, or
+  // pool split, so the final choice stays with the router.
+  rpc MatchExternalKvPrefix(MatchExternalKvPrefixRequest) returns (MatchExternalKvPrefixResponse);
+
+  rpc GetExternalKvHitCounts(GetExternalKvHitCountsRequest) returns (GetExternalKvHitCountsResponse);
+}