[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>
This commit is contained in:
wuyl1
2026-08-20 10:45:35 +08:00
committed by GitHub
co-authored by Wu, Yutong TianDi101 Zhangheng
parent 238ba40c27
commit 360d10d6bc
42 changed files with 6603 additions and 174 deletions
+11 -5
View File
@@ -108,6 +108,9 @@ jobs:
steps:
- uses: actions/checkout@v4
- name: Install protoc
run: sudo bash scripts/ci/utils/install_protoc.sh
- name: Verify dynamo-* SHA pin
run: |
set -euo pipefail
@@ -160,17 +163,17 @@ jobs:
- name: cargo check
working-directory: experimental/sgl-router
run: cargo check --all-targets
run: cargo check --workspace --all-targets
- name: cargo clippy
working-directory: experimental/sgl-router
run: cargo clippy --all-targets -- -D warnings
run: cargo clippy --workspace --all-targets -- -D warnings
- name: cargo fmt
working-directory: experimental/sgl-router
run: |
rustup toolchain install nightly --profile minimal --component rustfmt
cargo +nightly fmt -- --check
cargo +nightly fmt --all -- --check
# cargo-deny 0.19.6+ required for CVSS 4.0 parsing.
- name: Install cargo-deny
@@ -196,6 +199,9 @@ jobs:
steps:
- uses: actions/checkout@v4
- name: Install protoc
run: sudo bash scripts/ci/utils/install_protoc.sh
- name: Install OS deps (libssl, pkg-config)
run: |
sudo apt-get update
@@ -246,7 +252,7 @@ jobs:
# the Qwen3-0.6B tokenizer.json. Filter is matched against the
# full test path `tokenizer::parity::parity_matrix`; substring
# `parity_matrix` is unique to that test.
run: cargo test --release -- --skip parity_matrix
run: cargo test --release --workspace -- --skip parity_matrix
# Regenerate the cross-impl block-hash parity fixture and fail if it
# differs from the committed file. The Python script replicates
@@ -366,7 +372,7 @@ jobs:
working-directory: experimental/sgl-router
run: |
source "$HOME/.cargo/env"
cargo build --release
cargo build --release --workspace
# Install SGLang from the local checkout in editable mode (no PyPI
# version pin) — mirrors `pr-test-rust.yml` so the router e2e runs
+1 -1
View File
@@ -140,7 +140,7 @@ repos:
pass_filenames: false
- id: rustfmt-sgl-router
name: rustfmt experimental/sgl-router
entry: bash -c 'cd experimental/sgl-router && cargo fmt -- --check'
entry: bash -c 'cd experimental/sgl-router && cargo fmt --all -- --check'
language: system
files: ^experimental/sgl-router/.*\.rs$
pass_filenames: false
+25 -14
View File
@@ -34,21 +34,28 @@ ARG DEBIAN_VERSION=bookworm
######################## STAGE 1 — chef recipe ##########################
FROM rust:${RUST_VERSION}-${DEBIAN_VERSION} AS chef
RUN cargo install cargo-chef --locked --version ^0.1
WORKDIR /work
WORKDIR /work/sgl-router
COPY experimental/sgl-router/Cargo.toml ./
COPY experimental/sgl-router/rust-toolchain.toml ./
COPY experimental/sgl-router/sgl-kv-indexer/Cargo.toml sgl-kv-indexer/Cargo.toml
# Stub a minimal src tree so cargo can resolve the workspace, generate
# the lockfile (gitignored upstream), then prepare the chef recipe.
RUN mkdir -p src && echo "fn main() {}" > src/main.rs \
RUN mkdir -p src sgl-kv-indexer/src/bin \
&& echo "fn main() {}" > src/main.rs \
&& echo "" > src/lib.rs \
&& echo "" > sgl-kv-indexer/src/lib.rs \
&& echo "fn main() {}" > sgl-kv-indexer/src/bin/kv-indexer-server.rs \
&& echo "fn main() {}" > sgl-kv-indexer/src/bin/kv-indexer-bridge.rs \
&& cargo generate-lockfile \
&& cargo chef prepare --recipe-path recipe.json \
&& rm -rf src
&& rm -rf src sgl-kv-indexer/src
######################## STAGE 2 — builder ##############################
FROM rust:${RUST_VERSION}-${DEBIAN_VERSION} AS builder
RUN cargo install cargo-chef --locked --version ^0.1
WORKDIR /work
RUN apt-get update \
&& apt-get install -y --no-install-recommends protobuf-compiler \
&& rm -rf /var/lib/apt/lists/* \
&& cargo install cargo-chef --locked --version ^0.1
WORKDIR /work/sgl-router
# `dynamo-tokenizers` pulls in `pcre2-sys`, whose build.rs links the SYSTEM
# libpcre2-8 whenever pkg-config finds it (it does here — the rust:bookworm
@@ -58,29 +65,33 @@ WORKDIR /work
# it statically, keeping the runtime self-contained.
ENV PCRE2_SYS_STATIC=1
COPY --from=chef /work/recipe.json ./recipe.json
COPY --from=chef /work/Cargo.lock ./Cargo.lock
COPY experimental/sgl-router/rust-toolchain.toml ./
COPY --from=chef /work/sgl-router/recipe.json ./recipe.json
COPY --from=chef /work/sgl-router/Cargo.lock ./Cargo.lock
COPY experimental/sgl-router/sgl-kv-indexer/Cargo.toml sgl-kv-indexer/Cargo.toml
# Cook (compile + cache) the dep graph from the recipe. This layer's
# inputs are recipe.json + the toolchain — code changes in src/ do NOT
# invalidate it.
# Cook (compile + cache) the dep graph from the recipe. The recipe carries every
# workspace member's manifest, so chef recreates the Indexer's source stubs itself.
RUN cargo chef cook --release --recipe-path recipe.json
# Now bring in the real sources and the manifest they need.
COPY experimental/sgl-router/Cargo.toml ./
COPY experimental/sgl-router/src ./src
COPY experimental/sgl-router/sgl-kv-indexer/Cargo.toml sgl-kv-indexer/Cargo.toml
COPY experimental/sgl-router/sgl-kv-indexer/build.rs sgl-kv-indexer/build.rs
COPY experimental/sgl-router/sgl-kv-indexer/proto sgl-kv-indexer/proto
COPY experimental/sgl-router/sgl-kv-indexer/src sgl-kv-indexer/src
# --locked is intentionally omitted: the lockfile is generated in-container
# (gitignored upstream) and `cargo chef cook` may have mutated it during the
# dep-cook step, so a strict --locked check would spuriously fail.
RUN cargo build --release --bin sgl-router \
RUN touch sgl-kv-indexer/build.rs \
&& cargo build --release --bin sgl-router \
&& strip target/release/sgl-router
######################## STAGE 3 — runtime ##############################
FROM gcr.io/distroless/cc-debian12:nonroot AS runtime
COPY --from=builder /work/target/release/sgl-router /usr/local/bin/sgl-router
COPY --from=builder /work/sgl-router/target/release/sgl-router /usr/local/bin/sgl-router
# Default config path; mount your own via `-v <host-path>:/etc/sgl-router`.
ENV SGL_ROUTER_CONFIG=/etc/sgl-router/sgl-router.yaml
+4 -2
View File
@@ -1,6 +1,6 @@
[workspace]
resolver = "2"
members = ["."]
members = [".", "sgl-kv-indexer"]
[package]
name = "sgl-router"
@@ -33,6 +33,7 @@ axum = { version = "0.8", features = ["macros", "tracing"] }
tower = { version = "0.5", features = ["full"] }
tower-http = { version = "0.6", features = ["trace", "compression-gzip", "cors", "timeout", "request-id"] }
reqwest = { version = "0.12", features = ["stream", "json", "rustls-tls"], default-features = false }
sgl-kv-indexer = { path = "sgl-kv-indexer" }
# Serialization
serde = { version = "1", features = ["derive"] }
@@ -62,7 +63,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
futures = "0.3"
bytes = "1"
rand = "0.8"
tokio-stream = "0.1"
tokio-stream = { version = "0.1", features = ["net"] }
dashmap = "6"
kube = { version = "0.96", features = ["runtime", "derive"] }
k8s-openapi = { version = "0.23", features = ["v1_31"] }
@@ -86,6 +87,7 @@ serde_json = "1"
tempfile = "3"
tower = { version = "0.5", features = ["util"] }
tokio = { version = "1.42", features = ["test-util"] }
tonic = { version = "0.14.6", features = ["transport"] }
# Low-level msgpack encoder used to hand-construct wire bytes in
# kv_events golden-bytes tests (decode-only path uses rmp-serde).
rmp = "0.8"
+20
View File
@@ -50,6 +50,26 @@ Omit `--service-discovery-namespace` to watch all namespaces (requires
cluster-wide RBAC). For prefill/decode disaggregation, replace `--selector`
with `--prefill-selector` and `--decode-selector`.
External KV indexer as the cache-aware signal source:
```bash
sgl-router \
--model-id qwen3 \
--tokenizer-path /models/qwen3/tokenizer.json \
--worker-urls http://10.0.0.1:30000 http://10.0.0.2:30000 \
--policy cache_aware_zmq \
--kv-indexer-endpoint http://10.0.0.10:50051 \
--kv-indexer-query-timeout-ms 100 \
--kv-indexer-query-max-inflight 32
```
The existing cache-aware policy and thresholds are reused. When configured, the
Indexer replaces the Router-local radix tree as the cache signal. A successful
query with no usable match selects by minimum active load; connection failures,
timeouts, local admission rejection, and server rejection fail the Router
request with `503` rather than silently switching signals. The timeout and local
concurrency bound default to 100ms and 32 respectively.
## License
Apache-2.0.
@@ -0,0 +1,3 @@
target/
*.rs.bk
.DS_Store
@@ -0,0 +1,35 @@
[package]
name = "sgl-kv-indexer"
version = "0.1.0"
edition = "2021"
license = "Apache-2.0"
description = "In-memory KV cache indexer with a SGLang KV-event bridge"
publish = false # experimental crate; not published to crates.io
[lib]
name = "sgl_kv_indexer"
path = "src/lib.rs"
[[bin]]
name = "kv-indexer-server"
path = "src/bin/kv-indexer-server.rs"
[[bin]]
name = "kv-indexer-bridge"
path = "src/bin/kv-indexer-bridge.rs"
[dependencies]
bytes = "1"
prost = "0.14"
rmpv = "1.3.1"
tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "sync", "time"] }
tonic = { version = "0.14.6", features = ["transport"] }
tonic-prost = "0.14.6"
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }
# Must mirror the router's feature set: Cargo unifies features across the
# workspace, so enabling more transports here also enlarges the router binary.
zeromq = { version = "0.6", default-features = false, features = ["tokio-runtime", "tcp-transport"] }
[build-dependencies]
tonic-prost-build = "0.14.6"
@@ -0,0 +1,256 @@
# SGL KV Indexer (in-memory build)
`sgl-kv-indexer` is an experimental metadata service for SGLang KV-cache
blocks. It records which worker and storage tier currently holds each
content-addressed block, allowing a router to query likely cache hits without
moving KV data itself.
This build deliberately uses one process-local in-memory index. It has no
external storage dependency, but it is soft-state: restarting the Indexer loses
all placement metadata.
## Architecture
```text
SGLang worker ── ZMQ PUB ──> bridge ── gRPC ──> in-memory indexer
```
- SGLang publishes `BlockStored`, `BlockRemoved`, and `AllBlocksCleared` events.
- One `kv-indexer-bridge` follows each independent worker/rank event stream.
- One `kv-indexer-server` applies event batches and serves match queries.
- Placement, worker metadata, reverse holdings, and hit counters live in that
server process behind a single read/write lock.
Each apply RPC is ordered and atomic within the process, and a query sees a
consistent snapshot. The bridge splits larger worker event batches into
ordered RPCs of at most 16,384 hashes and 256 actions while keeping per-hash
metadata aligned. If a later RPC fails, earlier chunks may already be applied;
there is no rollback or replay. There is no persistence, replication, or state
sharing between Indexer servers.
## Operational contract
Run exactly one Indexer server for a deployment. Multiple bridge processes and
workers may report to it, but active-active Indexer servers have independent
state and must not be treated as replicas.
This build has no sequence gate, incarnation fencing, replay recovery, worker
liveness TTL, or restart recovery:
- An Indexer restart starts with an empty index.
- A worker death is not detected; its last placements remain until revoked or
until the Indexer restarts.
- Events published while a bridge is disconnected are not replayed.
- A publisher sequence gap is logged and otherwise ignored.
- Redelivered or reordered batches are applied again in arrival order.
Individual report, revoke, and clear mutations are idempotent. A future
Snapshot plus event-replay mechanism is required before production high
availability can rebuild state safely after restart or event loss.
## In-memory data model
The server stores:
- block hash → token count and `(worker, tier) → component mask`
- worker → router-facing address, `WorkerCacheSpec`, and reverse holdings by tier
- block hash → cumulative hit count
Component masks are `FULL=1`, `SWA=2`, `MAMBA=4`, and `0` for legacy
whole-block events. `REPORT` replaces the component snapshot for one
`(worker, tier, block)` placement. `REVOKE` removes that placement, and
`CLEAR_ALL_AT_TIER` removes every placement for the worker at that tier.
## Build
```bash
cd experimental/sgl-router
cargo build --release -p sgl-kv-indexer
```
This produces:
- `target/release/kv-indexer-server`
- `target/release/kv-indexer-bridge`
## End-to-end quickstart
Component-aware routing requires an SGLang engine build that supports
`component_types`. Start each command in its own terminal.
1. Start the single Indexer server:
```bash
KV_INDEXER_LISTEN_ADDR=127.0.0.1:50051 \
cargo run --release --bin kv-indexer-server
```
`KV_INDEXER_LISTEN_ADDR` defaults to `[::1]:50051`.
`KV_INDEXER_PREFIX_QUERY_MAX_INFLIGHT` sets the maximum number of prefix
queries executing concurrently and defaults to `32`. Requests above the limit
are rejected immediately with gRPC `RESOURCE_EXHAUSTED`.
There is no backend or storage configuration.
2. Start one bridge per worker event stream. This FULL+SWA example uses the
worker URL registered with the Router:
```bash
KV_INDEXER_WORKER_ID=worker-0 \
KV_INDEXER_WORKER_ADDRESS=http://127.0.0.1:30000 \
KV_INDEXER_ENDPOINT=http://127.0.0.1:50051 \
SGLANG_KV_EVENT_ENDPOINT=tcp://127.0.0.1:5567 \
SGLANG_KV_EVENT_TOPIC=kv-events \
KV_INDEXER_CACHE_COMPONENTS=full,swa \
KV_INDEXER_SWA_WINDOW_TOKENS=<model-window-tokens> \
KV_INDEXER_FULL_TIERS=HBM \
KV_INDEXER_SWA_TIERS=HBM \
cargo run --release --bin kv-indexer-bridge
```
For FULL+MAMBA, use:
```bash
KV_INDEXER_CACHE_COMPONENTS=full,mamba
KV_INDEXER_FULL_TIERS=HBM
KV_INDEXER_MAMBA_TIERS=HBM
```
3. Start the matching SGLang worker:
```bash
python -m sglang.launch_server \
--model-path <model> \
--port 30000 \
--kv-events-config \
'{"publisher":"zmq","endpoint":"tcp://*:5567","topic":"kv-events"}' \
--enable-kv-events-component-types
```
4. Start the Router with the Indexer as the authoritative cache signal:
```bash
sgl-router \
--model-id <model-id> \
--tokenizer-path <huggingface-repo-or-tokenizer> \
--worker-urls http://127.0.0.1:30000 \
--policy cache_aware_zmq \
--kv-indexer-endpoint http://127.0.0.1:50051 \
--kv-indexer-query-timeout-ms 100 \
--kv-indexer-query-max-inflight 32
```
For multiple workers, repeat steps 23 with unique worker IDs and ports.
`KV_INDEXER_WORKER_ADDRESS` must exactly match the corresponding Router URL.
The bridge sends its `WorkerCacheSpec` with every batch. Omitting
`KV_INDEXER_CACHE_COMPONENTS` clears any previously stored spec and uses legacy
whole-block matching.
## API
The protobuf service in `proto/kv_indexer.proto` provides:
- `ApplyExternalKvBatch`: ordered placement reports, revocations, and clears.
The request `seq` is carried for observability only.
- `MatchExternalKv`: workers and tiers holding requested block hashes.
- `MatchExternalKvPrefix`: per-worker longest contiguous reusable prefix.
- `GetExternalKvHitCounts`: per-block hit counters.
There is no gRPC health service in this build.
## Prefix routing semantics
For a legacy worker, `matched_prefix_blocks` is the largest `n` such that it
holds every block in `hashes[0..n)` without a gap. For a component-aware worker:
- FULL must be contiguous on every matched block.
- SWA must cover the trailing `swa_window_tokens` at the candidate boundary, or
form an unbroken run from the prompt head.
- MAMBA must be present on the candidate boundary block.
Component placements without a worker spec fail closed. Workers with an empty
router-facing address are excluded.
The Indexer returns every candidate sorted by prefix length; it does not choose
a worker. When configured, it replaces the Router's local radix tree as the
cache signal: the Router intersects Indexer results with its healthy candidates,
and a successful query with no usable match selects by minimum active load.
Indexer connection failures, timeouts, overload, and a prompt too long to fit one
gRPC message fall back to that same minimum-active-load selection, so an
unreachable Indexer costs cache affinity rather than availability; a rejected RPC
still fails the Router request with `503`, because it means the two sides
disagree on the request contract. An
endpoint the Router could never dial is rejected at startup instead of failing
every query later. The local radix tree is used only when no Indexer endpoint is
configured. The per-query deadline defaults to 100ms and can be changed with
`--kv-indexer-query-timeout-ms`. The Router-side admission bound defaults to 32
concurrent calls and can be changed with `--kv-indexer-query-max-inflight`.
Prefix queries carry no Indexer-imposed block cap beyond the caller's
`max_blocks` ceiling — unlike applies and `MatchExternalKv`, which reject above
16,384 hashes. The in-memory backend scans the request in one pass over a single
consistent snapshot, holding O(1) matching state per candidate worker and
considering only workers that hold the first block, so request length costs time
but not memory. Block hashes use packed `sfixed64` encoding, and the server
accepts decoded gRPC messages up to 8 MiB (roughly one million hashes).
Server work is bounded by `max_blocks` when the caller supplies one and by that
transport limit. The Router's per-query deadline bounds how long it waits for an
answer, but does not cancel a synchronous scan already in progress. A first-block
miss returns immediately with `blocks_read=1`.
Message decoding happens before a request reaches the service, so
`KV_INDEXER_PREFIX_QUERY_MAX_INFLIGHT` bounds the scan but not the bytes a peer
makes the server buffer. That is bounded instead by the HTTP/2 stream limit: each
connection is capped at 64 concurrent streams, bounding that connection to
64 × 8 MiB of undecoded requests. For a query past the 8 MiB ceiling, the Router
sends only the leading hashes that fit and still divides the returned prefix by
the full request's block count. This preserves a useful lower-bound cache signal
without overstating the match rate. If an Indexer has a lower ceiling and returns
gRPC `OUT_OF_RANGE`, the Router falls back to minimum active load.
Long scans hold the read lock throughout. Operators serving very long prompts
should set `max_blocks` instead of relying on the message-size limit.
## Overload behavior and observability
The Router and server apply separate admission bounds. The Router rejects a
query locally when its `--kv-indexer-query-max-inflight` permits are exhausted;
the server returns gRPC `RESOURCE_EXHAUSTED` when
`KV_INDEXER_PREFIX_QUERY_MAX_INFLIGHT` is exhausted. Both leave the request
routed by minimum active load, logged at `WARN` on the Router.
Every Router query publishes its timeout through the gRPC `grpc-timeout` header.
The server timestamps arrival and returns `DEADLINE_EXCEEDED` before backend work
when queueing has already consumed that budget. Apply/event RPCs are never shed,
because dropping one would permanently diverge the soft-state index.
Deadline shedding is logged at `INFO`; server admission rejection is logged at
`WARN`. Each rejection class reports totals 1, 2, 4, 8, and so on, making the
first overload visible at the default log level without log volume growing
linearly with sustained overload.
## Bridge configuration
Required or commonly used bridge variables:
- `KV_INDEXER_WORKER_ID`: unique ID for the worker event stream
- `KV_INDEXER_WORKER_ADDRESS`: Router-facing worker URL
- `KV_INDEXER_ENDPOINT`: Indexer endpoint, default `http://[::1]:50051`
- `SGLANG_KV_EVENT_ENDPOINT`: worker PUB endpoint
- `SGLANG_KV_EVENT_TOPIC`: ZMQ subscription topic
- `KV_INDEXER_CLEAR_TIERS`: tiers affected by clear, default `HBM,DRAM,SSD`
- `KV_INDEXER_CACHE_COMPONENTS`: optional `full,swa` or `full,mamba`
- `KV_INDEXER_SWA_WINDOW_TOKENS`: required when SWA is configured
- `KV_INDEXER_FULL_TIERS`, `KV_INDEXER_SWA_TIERS`,
`KV_INDEXER_MAMBA_TIERS`: servable component tiers
- `KV_INDEXER_CACHE_SPEC_VERSION`: component-rule version, default `1`
## Tests
No external service is needed:
```bash
cargo fmt --all -- --check
cargo clippy --all-targets -- -D warnings
cargo test
```
@@ -0,0 +1,10 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
fn main() -> Result<(), Box<dyn std::error::Error>> {
tonic_prost_build::configure()
.build_client(true)
.build_server(true)
.compile_protos(&["proto/kv_indexer.proto"], &["proto"])?;
Ok(())
}
@@ -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);
}
@@ -0,0 +1,170 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Deadline-based load shedding for the query path.
//!
//! A query that already waited out its caller's whole deadline can no longer be
//! answered usefully, so serving it only delays the rest of the backlog. The
//! budget is the caller's own `grpc-timeout`: no server-side threshold to tune,
//! and a caller that declared no deadline is never shed.
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tonic::metadata::MetadataMap;
use tonic::{Extensions, Request, Status};
/// When the request's headers were read off the connection.
#[derive(Clone, Copy)]
struct Arrival(Instant);
/// Counts rejections of one kind and reports on doubling totals, so the first
/// rejection is visible immediately and a sustained overload cannot flood the log.
pub(crate) struct RejectionLog(AtomicU64);
impl RejectionLog {
pub(crate) const fn new() -> Self {
Self(AtomicU64::new(0))
}
/// Records one rejection, returning the running total when it should be
/// logged and `None` when it should be absorbed.
pub(crate) fn record(&self) -> Option<u64> {
let total = self.0.fetch_add(1, Ordering::Relaxed) + 1;
total.is_power_of_two().then_some(total)
}
}
static DEADLINE_SHED_LOG: RejectionLog = RejectionLog::new();
/// Timestamps a request's arrival so the query path can measure how long it then
/// waited. Runs before the per-request task is spawned, so the stamp precedes
/// any scheduling delay. Without this interceptor nothing is ever shed.
pub fn stamp_arrival(mut request: Request<()>) -> Result<Request<()>, Status> {
request.extensions_mut().insert(Arrival(Instant::now()));
Ok(request)
}
/// Rejects a query that spent its caller's entire deadline waiting to be served.
///
/// Never applied to the apply path: dropping a KV event would leave the index
/// permanently diverged from the worker that reported it.
pub(crate) fn reject_if_deadline_passed(
metadata: &MetadataMap,
extensions: &Extensions,
) -> Result<(), Status> {
let (Some(arrival), Some(budget)) = (extensions.get::<Arrival>(), caller_deadline(metadata))
else {
return Ok(());
};
let waited = arrival.0.elapsed();
if waited < budget {
return Ok(());
}
if let Some(shed_total) = DEADLINE_SHED_LOG.record() {
tracing::info!(
shed_total,
waited_ms = waited.as_millis(),
budget_ms = budget.as_millis(),
"shedding prefix query whose caller deadline already passed"
);
}
Err(Status::deadline_exceeded(
"prefix query waited longer than its caller deadline",
))
}
/// The budget the caller declared in `grpc-timeout`, per the gRPC wire spec (up
/// to 8 digits followed by a unit). `None` for an absent or unparsable value,
/// which leaves the request unshed. Measured against the wait since arrival, so
/// transit time is ignored and shedding can only be late, never early.
fn caller_deadline(metadata: &MetadataMap) -> Option<Duration> {
let raw = metadata.get("grpc-timeout")?.to_str().ok()?;
let unit = *raw.as_bytes().last()?;
let value: u64 = raw.get(..raw.len() - 1)?.parse().ok()?;
match unit {
b'H' => value.checked_mul(60 * 60).map(Duration::from_secs),
b'M' => value.checked_mul(60).map(Duration::from_secs),
b'S' => Some(Duration::from_secs(value)),
b'm' => Some(Duration::from_millis(value)),
b'u' => Some(Duration::from_micros(value)),
b'n' => Some(Duration::from_nanos(value)),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn metadata(timeout: Option<&str>) -> MetadataMap {
let mut metadata = MetadataMap::new();
if let Some(timeout) = timeout {
metadata.insert("grpc-timeout", timeout.parse().unwrap());
}
metadata
}
fn extensions(arrival: Option<Instant>) -> Extensions {
let mut extensions = Extensions::new();
if let Some(arrival) = arrival {
extensions.insert(Arrival(arrival));
}
extensions
}
#[test]
fn parses_every_wire_unit() {
for (raw, expected) in [
("1H", Duration::from_secs(3600)),
("2M", Duration::from_secs(120)),
("3S", Duration::from_secs(3)),
("100m", Duration::from_millis(100)),
("250u", Duration::from_micros(250)),
("400n", Duration::from_nanos(400)),
] {
assert_eq!(caller_deadline(&metadata(Some(raw))), Some(expected));
}
}
#[test]
fn unparsable_deadline_is_ignored() {
for raw in ["", "m", "100", "100x", "abcm", "99999999999999999999H"] {
assert_eq!(caller_deadline(&metadata(Some(raw))), None, "{raw}");
}
assert_eq!(caller_deadline(&metadata(None)), None);
}
#[test]
fn sheds_only_once_the_caller_deadline_has_passed() {
let now = Instant::now();
let waited_past = now - Duration::from_millis(150);
let within = now - Duration::from_millis(10);
assert_eq!(
reject_if_deadline_passed(&metadata(Some("100m")), &extensions(Some(waited_past)))
.unwrap_err()
.code(),
tonic::Code::DeadlineExceeded
);
assert!(
reject_if_deadline_passed(&metadata(Some("100m")), &extensions(Some(within))).is_ok()
);
}
#[test]
fn rejections_are_reported_on_doubling_counts() {
let log = RejectionLog::new();
let reported: Vec<u64> = (0..16).filter_map(|_| log.record()).collect();
assert_eq!(reported, vec![1, 2, 4, 8, 16]);
}
#[test]
fn missing_deadline_or_arrival_never_sheds() {
let long_wait = Instant::now() - Duration::from_secs(60);
// A caller that declared no deadline keeps the pre-existing behaviour.
assert!(reject_if_deadline_passed(&metadata(None), &extensions(Some(long_wait))).is_ok());
// No interceptor installed: nothing to measure the wait against.
assert!(reject_if_deadline_passed(&metadata(Some("1m")), &extensions(None)).is_ok());
}
}
@@ -0,0 +1,18 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use sgl_kv_indexer::bridge::{run_bridge_until, BridgeConfig};
use sgl_kv_indexer::shutdown_signal;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()),
)
.init();
let config = BridgeConfig::from_env()?;
run_bridge_until(config, shutdown_signal()).await?;
Ok(())
}
@@ -0,0 +1,94 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use std::net::SocketAddr;
use std::sync::Arc;
use std::{env, io};
use sgl_kv_indexer::{
server_builder, shutdown_signal, stamp_arrival, InMemoryKvIndexerBackend, KvIndexerBackend,
KvIndexerService, DEFAULT_PREFIX_QUERY_MAX_INFLIGHT, MAX_CONCURRENT_STREAMS,
};
use tonic::service::interceptor::InterceptedService;
use tracing::info;
const PREFIX_QUERY_MAX_INFLIGHT_ENV: &str = "KV_INDEXER_PREFIX_QUERY_MAX_INFLIGHT";
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()),
)
.init();
let addr = std::env::var("KV_INDEXER_LISTEN_ADDR")
.unwrap_or_else(|_| "[::1]:50051".to_string())
.parse::<SocketAddr>()?;
let prefix_query_max_inflight = prefix_query_max_inflight_from_env()?;
let backend: Arc<dyn KvIndexerBackend> = Arc::new(InMemoryKvIndexerBackend::new());
// The interceptor timestamps each request before its own task is queued,
// which is what lets the query path shed work whose deadline expired.
let service = InterceptedService::new(
KvIndexerService::with_prefix_query_max_inflight(backend, prefix_query_max_inflight)
.into_server(),
stamp_arrival,
);
info!(
%addr,
prefix_query_max_inflight,
max_concurrent_streams = MAX_CONCURRENT_STREAMS,
"starting single-server in-memory SGLang KV Indexer"
);
server_builder()
.add_service(service)
.serve_with_shutdown(addr, shutdown_signal())
.await?;
Ok(())
}
fn prefix_query_max_inflight_from_env() -> io::Result<usize> {
match env::var(PREFIX_QUERY_MAX_INFLIGHT_ENV) {
Ok(raw) => parse_prefix_query_max_inflight(&raw),
Err(env::VarError::NotPresent) => Ok(DEFAULT_PREFIX_QUERY_MAX_INFLIGHT),
Err(env::VarError::NotUnicode(_)) => Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("{PREFIX_QUERY_MAX_INFLIGHT_ENV} must be valid UTF-8"),
)),
}
}
fn parse_prefix_query_max_inflight(raw: &str) -> io::Result<usize> {
let value = raw.parse::<usize>().map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
format!("{PREFIX_QUERY_MAX_INFLIGHT_ENV} must be a positive integer, got {raw:?}"),
)
})?;
if value == 0 {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("{PREFIX_QUERY_MAX_INFLIGHT_ENV} must be greater than zero"),
));
}
Ok(value)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_positive_prefix_query_limit() {
assert_eq!(parse_prefix_query_max_inflight("64").unwrap(), 64);
}
#[test]
fn rejects_invalid_prefix_query_limit() {
assert!(parse_prefix_query_max_inflight("0").is_err());
assert!(parse_prefix_query_max_inflight("many").is_err());
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,382 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Router-facing client for the prefix-match query.
//!
//! A successful query distinguishes a real match from an empty result. Transport
//! failures, deadlines, and server rejections stay distinct errors so the caller
//! chooses between degrading and failing the request, instead of silently using
//! a different signal.
use std::time::Duration;
use tokio::sync::{Semaphore, SemaphorePermit};
use tonic::transport::{Channel, Endpoint};
use crate::pb::kv_indexer_client::KvIndexerClient;
use crate::pb::MatchExternalKvPrefixRequest;
use crate::service::MAX_GRPC_DECODING_MESSAGE_SIZE;
/// Default per-query deadline. Indexer failures are request failures, so this
/// absorbs normal cross-host jitter without stalling a request indefinitely.
pub const DEFAULT_QUERY_DEADLINE: Duration = Duration::from_millis(100);
/// Default process-local bound on prefix-query RPCs issued by one client.
pub const DEFAULT_QUERY_MAX_INFLIGHT: usize = 32;
// Leave room for the packed field tag, length prefix, and future scalar fields.
const PREFIX_QUERY_ENCODING_HEADROOM: usize = 16;
const MAX_PREFIX_HASHES_PER_QUERY: usize =
(MAX_GRPC_DECODING_MESSAGE_SIZE - PREFIX_QUERY_ENCODING_HEADROOM) / 8;
/// One worker's contiguous prefix hit.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PrefixMatch {
/// Router-facing routing identity; intersect byte-for-byte with registered
/// worker URLs. Never empty (the indexer drops unroutable workers).
pub address: String,
/// Length of the contiguous request prefix this worker holds.
pub matched_prefix_blocks: u32,
/// Opaque worker id, for the caller's logs only.
pub worker_id: String,
}
/// A failed prefix query.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PrefixIndexError {
/// The endpoint could not be reached.
Unreachable,
/// The query exceeded its deadline.
Timeout,
/// The client or Indexer shed the query because its in-flight limit was hit.
Overloaded,
/// The query exceeded the Indexer's gRPC message-size limit, so no worker's
/// prefix was scanned. Bounded by prompt length, not by load: retrying the
/// same prompt cannot succeed.
QueryTooLarge,
/// The server rejected the request.
Rejected(tonic::Code),
}
impl std::fmt::Display for PrefixIndexError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Unreachable => f.write_str("KV Indexer is unreachable"),
Self::Timeout => f.write_str("KV Indexer query timed out"),
Self::Overloaded => f.write_str("KV Indexer is overloaded"),
Self::QueryTooLarge => {
f.write_str("KV Indexer query exceeded the gRPC message-size limit")
}
Self::Rejected(code) => write!(f, "KV Indexer rejected the query with {code}"),
}
}
}
impl std::error::Error for PrefixIndexError {}
/// A configured endpoint that is not a usable gRPC target.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InvalidEndpoint {
endpoint: String,
reason: &'static str,
}
impl std::fmt::Display for InvalidEndpoint {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"invalid KV Indexer endpoint `{}`: {}",
self.endpoint, self.reason
)
}
}
impl std::error::Error for InvalidEndpoint {}
/// Result of a successful prefix query.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PrefixOutcome {
Matched {
/// Sorted by `matched_prefix_blocks`, descending.
matches: Vec<PrefixMatch>,
/// Longest contiguous prefix held by any single worker.
best_prefix_blocks: u32,
},
/// No worker holds a prefix (or the request had no hashes).
Empty,
}
/// Client configuration.
#[derive(Debug, Clone)]
pub struct PrefixIndexConfig {
/// gRPC endpoint of the indexer, e.g. `http://10.0.0.1:50051`.
pub endpoint: String,
/// Per-query deadline.
pub query_deadline: Duration,
/// Maximum prefix-query RPCs issued concurrently by this client.
pub max_inflight: usize,
}
impl PrefixIndexConfig {
/// Config with the default query deadline ([`DEFAULT_QUERY_DEADLINE`]).
pub fn new(endpoint: impl Into<String>) -> Self {
Self {
endpoint: endpoint.into(),
query_deadline: DEFAULT_QUERY_DEADLINE,
max_inflight: DEFAULT_QUERY_MAX_INFLIGHT,
}
}
}
/// The prefix-match query the router links against.
#[tonic::async_trait]
pub trait PrefixIndex: Send + Sync {
/// Queries the longest contiguous prefix each worker holds for `hashes`
/// (prompt order, `hashes[0]` first).
async fn match_prefix(&self, hashes: Vec<i64>) -> Result<PrefixOutcome, PrefixIndexError>;
}
/// tonic-backed [`PrefixIndex`] with a lazily-established connection.
pub struct GrpcPrefixIndex {
channel: Channel,
deadline: Duration,
prefix_query_semaphore: Semaphore,
}
impl GrpcPrefixIndex {
/// Rejects an unusable endpoint instead of building a client that can only
/// fail, so a misconfigured address stops startup rather than silently
/// costing every request its cache affinity.
pub fn new(config: PrefixIndexConfig) -> Result<Self, InvalidEndpoint> {
assert!(
config.max_inflight > 0,
"prefix query max inflight must be greater than zero"
);
Ok(Self {
channel: parse_endpoint(&config.endpoint)?.connect_lazy(),
deadline: config.query_deadline,
prefix_query_semaphore: Semaphore::new(config.max_inflight),
})
}
fn try_acquire_prefix_query(&self) -> Result<SemaphorePermit<'_>, PrefixIndexError> {
self.prefix_query_semaphore
.try_acquire()
.map_err(|_| PrefixIndexError::Overloaded)
}
}
fn truncate_prefix_query(hashes: &mut Vec<i64>) -> Option<usize> {
let total = hashes.len();
hashes.truncate(MAX_PREFIX_HASHES_PER_QUERY);
(hashes.len() != total).then_some(total)
}
#[tonic::async_trait]
impl PrefixIndex for GrpcPrefixIndex {
async fn match_prefix(&self, mut hashes: Vec<i64>) -> Result<PrefixOutcome, PrefixIndexError> {
if hashes.is_empty() {
return Ok(PrefixOutcome::Empty);
}
if let Some(total_hashes) = truncate_prefix_query(&mut hashes) {
tracing::warn!(
total_hashes,
queried_hashes = hashes.len(),
"KV Indexer query truncated to the gRPC message-size limit"
);
}
let _permit = self.try_acquire_prefix_query()?;
let mut client = KvIndexerClient::new(self.channel.clone());
let mut request = tonic::Request::new(MatchExternalKvPrefixRequest {
hashes,
// The policy retains the full query length as its denominator, so a
// transport-limited prefix cannot turn a partial scan into a perfect
// hit.
max_blocks: 0,
});
// On the wire so the indexer can drop a query this caller already stopped
// waiting for. The local timeout below stays the hard stop, since it also
// covers a stall before the channel applies its own deadline.
request.set_timeout(self.deadline);
match tokio::time::timeout(self.deadline, client.match_external_kv_prefix(request)).await {
Err(_) => Err(PrefixIndexError::Timeout),
Ok(Err(status)) => Err(classify(status.code())),
Ok(Ok(response)) => {
let response = response.into_inner();
if response.matches.is_empty() {
return Ok(PrefixOutcome::Empty);
}
let matches = response
.matches
.into_iter()
.map(|m| PrefixMatch {
address: m.worker_address,
matched_prefix_blocks: m.matched_prefix_blocks,
worker_id: m.worker_id,
})
.collect();
Ok(PrefixOutcome::Matched {
matches,
best_prefix_blocks: response.best_prefix_blocks,
})
}
}
}
}
/// Validates the endpoint the operator configured. tonic itself only checks URI
/// syntax, which accepts a host:port with no scheme and then fails on every
/// connect, so the scheme and host are checked here.
fn parse_endpoint(endpoint: &str) -> Result<Endpoint, InvalidEndpoint> {
let reject = |reason: &'static str| InvalidEndpoint {
endpoint: endpoint.to_string(),
reason,
};
let parsed =
Endpoint::from_shared(endpoint.to_string()).map_err(|_| reject("not a valid URI"))?;
// A `unix:` endpoint is fully specified by its socket path.
if endpoint.starts_with("unix:") {
return Ok(parsed);
}
match parsed.uri().scheme_str() {
None => Err(reject("missing scheme, expected http:// or https://")),
Some("http" | "https") => {
if parsed.uri().host().unwrap_or_default().is_empty() {
Err(reject("missing host"))
} else {
Ok(parsed)
}
}
Some(_) => Err(reject("unsupported scheme, expected http:// or https://")),
}
}
fn classify(code: tonic::Code) -> PrefixIndexError {
match code {
tonic::Code::Unavailable => PrefixIndexError::Unreachable,
// The indexer sheds an expired query as DEADLINE_EXCEEDED, while tonic
// reports its own enforcement of the same `grpc-timeout` as CANCELLED.
// This client cancels a query for no other reason.
tonic::Code::DeadlineExceeded | tonic::Code::Cancelled => PrefixIndexError::Timeout,
tonic::Code::ResourceExhausted => PrefixIndexError::Overloaded,
// The indexer's decoder refuses a message past its size limit with
// OUT_OF_RANGE. A prompt too long to carry is not a disagreement about
// the request contract, so it stays separable from `Rejected` and the
// caller can degrade instead of failing the request.
tonic::Code::OutOfRange => PrefixIndexError::QueryTooLarge,
_ => PrefixIndexError::Rejected(code),
}
}
#[cfg(test)]
mod tests {
use prost::Message;
use super::*;
#[test]
fn classifies_resource_exhausted_as_overload() {
assert_eq!(
classify(tonic::Code::ResourceExhausted),
PrefixIndexError::Overloaded
);
}
/// An over-limit query must stay distinguishable from a contract rejection:
/// the caller degrades on the former and fails the request on the latter.
#[test]
fn classifies_over_limit_message_as_too_large() {
assert_eq!(
classify(tonic::Code::OutOfRange),
PrefixIndexError::QueryTooLarge
);
assert_eq!(
classify(tonic::Code::InvalidArgument),
PrefixIndexError::Rejected(tonic::Code::InvalidArgument)
);
}
#[test]
fn oversized_query_keeps_a_prefix_within_the_transport_limit() {
let total = MAX_PREFIX_HASHES_PER_QUERY + 1;
let mut hashes = vec![-1; total];
assert_eq!(truncate_prefix_query(&mut hashes), Some(total));
assert_eq!(hashes.len(), MAX_PREFIX_HASHES_PER_QUERY);
assert!(
MatchExternalKvPrefixRequest {
hashes,
max_blocks: 0,
}
.encoded_len()
<= MAX_GRPC_DECODING_MESSAGE_SIZE
);
}
#[test]
fn classifies_both_deadline_signals_as_timeout() {
for code in [tonic::Code::DeadlineExceeded, tonic::Code::Cancelled] {
assert_eq!(classify(code), PrefixIndexError::Timeout);
}
}
#[test]
fn accepts_endpoints_the_client_can_actually_dial() {
for endpoint in [
"http://10.0.0.1:50051",
"https://indexer.svc:443",
"unix:/tmp/i",
] {
assert!(
parse_endpoint(endpoint).is_ok(),
"{endpoint} should be accepted"
);
}
}
/// A host:port with no scheme parses as a URI but can never connect, which
/// is the misconfiguration that otherwise only shows up under traffic.
#[test]
fn rejects_endpoints_that_could_only_fail_at_query_time() {
for endpoint in [
"10.0.0.1:50051",
"indexer.svc",
"grpc://10.0.0.1:50051",
"http://",
] {
let error = parse_endpoint(endpoint)
.expect_err(&format!("{endpoint} should be rejected"))
.to_string();
assert!(
error.contains(endpoint),
"error should name the endpoint: {error}"
);
}
}
#[test]
fn construction_fails_on_an_invalid_endpoint() {
assert!(GrpcPrefixIndex::new(PrefixIndexConfig::new("10.0.0.1:50051")).is_err());
}
#[tokio::test]
async fn local_admission_rejects_without_queueing() {
let index = GrpcPrefixIndex::new(PrefixIndexConfig {
endpoint: "http://127.0.0.1:1".to_string(),
query_deadline: DEFAULT_QUERY_DEADLINE,
max_inflight: 1,
})
.expect("valid endpoint");
let permit = index.try_acquire_prefix_query().unwrap();
assert_eq!(
index.try_acquire_prefix_query().unwrap_err(),
PrefixIndexError::Overloaded
);
drop(permit);
assert!(index.try_acquire_prefix_query().is_ok());
}
}
@@ -0,0 +1,34 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! SGLang KV Indexer: a gRPC service that tracks externally-managed KV cache
//! block placements (as reported by inference engines such as SGLang HiCache)
//! and answers placement-match queries for KV-aware routing.
pub mod bridge;
pub mod client;
pub mod pb {
tonic::include_proto!("kv_indexer.v1");
}
mod admission;
mod memory_backend;
mod service;
mod shutdown;
pub use admission::stamp_arrival;
pub use client::{
GrpcPrefixIndex, InvalidEndpoint, PrefixIndex, PrefixIndexConfig, PrefixIndexError,
PrefixMatch, PrefixOutcome, DEFAULT_QUERY_MAX_INFLIGHT,
};
pub use memory_backend::InMemoryKvIndexerBackend;
pub use service::{
component_bit, server_builder, BlockComponents, KvIndexerBackend, KvIndexerService,
WorkerPrefixInput, COMPONENT_FULL, COMPONENT_MAMBA, COMPONENT_SWA,
DEFAULT_PREFIX_QUERY_MAX_INFLIGHT, MAX_CONCURRENT_STREAMS, MAX_GRPC_DECODING_MESSAGE_SIZE,
};
pub use shutdown::shutdown_signal;
/// Re-exported because [`PrefixIndexError::Rejected`] carries it, so callers can
/// match on a rejection without depending on tonic.
pub use tonic::Code as RpcCode;
@@ -0,0 +1,500 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Process-local storage backend for the KV Indexer.
//!
//! The complete placement view lives behind one [`RwLock`], making an apply batch
//! atomic and every query a consistent snapshot. The state is soft: not shared
//! with another server, and lost when the process exits.
use std::collections::{BTreeMap, HashMap, HashSet};
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use tonic::Status;
use crate::pb::{
ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvActionType,
ExternalKvNodeMatch, GetExternalKvHitCountsRequest, GetExternalKvHitCountsResponse,
HitCountEntry, MatchExternalKvPrefixRequest, MatchExternalKvPrefixResponse,
MatchExternalKvRequest, MatchExternalKvResponse, TierHashes, WorkerCacheSpec,
};
use crate::service::{assemble_prefix_response, prefix_limit, WorkerPrefixScanner};
use crate::{BlockComponents, KvIndexerBackend, WorkerPrefixInput};
#[derive(Debug, Default)]
struct BlockRecord {
/// Shared block token count. A zero value means legacy/unspecified.
token_count: u32,
/// Resident component snapshot for each `(worker, tier)`.
placements: HashMap<(String, i32), u32>,
}
#[derive(Debug, Default)]
struct WorkerRecord {
address: String,
spec: Option<WorkerCacheSpec>,
/// Reverse index used by CLEAR_ALL_AT_TIER.
holdings: HashMap<i32, HashSet<i64>>,
}
#[derive(Debug, Default)]
struct State {
blocks: HashMap<i64, BlockRecord>,
workers: HashMap<String, WorkerRecord>,
hit_counts: HashMap<i64, u64>,
}
struct WorkerView {
worker_id: String,
address: String,
spec: Option<WorkerCacheSpec>,
hashes_by_tier: BTreeMap<i32, Vec<(i64, u32, u32)>>,
blocks: Vec<Option<BlockComponents>>,
}
#[derive(Debug)]
struct PrefixCandidate {
worker_id: String,
address: String,
scanner: WorkerPrefixScanner,
}
/// Single-process, soft-state KV placement index.
#[derive(Debug, Default)]
pub struct InMemoryKvIndexerBackend {
state: RwLock<State>,
}
impl InMemoryKvIndexerBackend {
pub fn new() -> Self {
Self::default()
}
fn read_state(&self) -> Result<RwLockReadGuard<'_, State>, Status> {
self.state
.read()
.map_err(|_| Status::internal("in-memory backend lock poisoned"))
}
fn write_state(&self) -> Result<RwLockWriteGuard<'_, State>, Status> {
self.state
.write()
.map_err(|_| Status::internal("in-memory backend lock poisoned"))
}
fn apply(
&self,
req: ApplyExternalKvBatchRequest,
) -> Result<ApplyExternalKvBatchResponse, Status> {
let mut state = self.write_state()?;
let worker_id = req.worker_id;
// Address and spec are snapshots carried on every batch. Empty address
// makes the worker unroutable; absent spec returns it to legacy mode.
{
let worker = state.workers.entry(worker_id.clone()).or_default();
worker.address = req.worker_address;
worker.spec = req.cache_spec;
}
for action in req.actions {
match ExternalKvActionType::try_from(action.r#type) {
Ok(ExternalKvActionType::ActionReport) => {
let has_masks = !action.component_masks.is_empty();
let has_sizes = !action.block_sizes.is_empty();
// REPORT is a REPLACE snapshot. Keep the final occurrence
// when a coalesced action repeats one hash.
let mut last_by_hash: HashMap<i64, (u32, u32)> = HashMap::new();
for (index, hash) in action.hashes.into_iter().enumerate() {
let mask = if has_masks {
action.component_masks[index]
} else {
0
};
let token_count = if has_sizes {
action.block_sizes[index]
} else {
0
};
last_by_hash.insert(hash, (mask, token_count));
}
for (hash, (mask, token_count)) in last_by_hash {
let block = state.blocks.entry(hash).or_default();
block
.placements
.insert((worker_id.clone(), action.tier), mask);
// A legacy report carries no size, so 0 means
// "unknown" and must not erase a known count.
if token_count > 0 {
block.token_count = token_count;
}
state
.workers
.entry(worker_id.clone())
.or_default()
.holdings
.entry(action.tier)
.or_default()
.insert(hash);
}
}
Ok(ExternalKvActionType::ActionRevoke) => {
for hash in action.hashes {
revoke_one(&mut state, &worker_id, &hash, action.tier);
}
}
Ok(ExternalKvActionType::ActionClearAllAtTier) => {
let hashes = state
.workers
.get(&worker_id)
.and_then(|worker| worker.holdings.get(&action.tier))
.cloned()
.unwrap_or_default();
for hash in hashes {
revoke_one(&mut state, &worker_id, &hash, action.tier);
}
}
Ok(ExternalKvActionType::ActionUnknown) | Err(_) => {
return Err(Status::invalid_argument("unsupported action type"));
}
}
}
Ok(ApplyExternalKvBatchResponse {})
}
fn do_match(&self, req: MatchExternalKvRequest) -> Result<MatchExternalKvResponse, Status> {
let hashes = dedup_preserve_order(&req.hashes);
let workers = if req.count_as_hit {
let mut state = self.write_state()?;
let (workers, matched_hashes) = Self::collect_worker_views(&state, &hashes, false);
for hash in matched_hashes {
let count = state.hit_counts.entry(hash).or_default();
*count = count.saturating_add(1);
}
workers
} else {
let state = self.read_state()?;
Self::collect_worker_views(&state, &hashes, false).0
};
let matches = workers
.into_iter()
.map(|worker| ExternalKvNodeMatch {
worker_id: worker.worker_id,
address: worker.address,
hashes_by_tier: worker
.hashes_by_tier
.into_iter()
.map(|(tier, placements)| TierHashes {
tier,
hashes: placements.iter().map(|(hash, _, _)| *hash).collect(),
component_masks: placements.iter().map(|(_, mask, _)| *mask).collect(),
block_sizes: placements
.into_iter()
.map(|(_, _, block_size)| block_size)
.collect(),
})
.collect(),
})
.collect();
Ok(MatchExternalKvResponse { matches })
}
fn collect_worker_views(
state: &State,
hashes: &[i64],
with_blocks: bool,
) -> (Vec<WorkerView>, Vec<i64>) {
// Keyed by a borrow of the stored worker id, so it is copied once per
// worker in the result rather than once per scanned placement.
let mut worker_order: Vec<&str> = Vec::new();
let mut by_worker: HashMap<&str, WorkerView> = HashMap::new();
let mut matched_hashes = Vec::new();
for (index, hash) in hashes.iter().enumerate() {
let Some(block) = state.blocks.get(hash) else {
continue;
};
if block.placements.is_empty() {
continue;
}
matched_hashes.push(*hash);
for ((worker, tier), mask) in &block.placements {
let view = by_worker.entry(worker.as_str()).or_insert_with(|| {
worker_order.push(worker.as_str());
let metadata = state.workers.get(worker);
WorkerView {
worker_id: worker.clone(),
address: metadata
.map(|worker| worker.address.clone())
.unwrap_or_default(),
spec: metadata.and_then(|worker| worker.spec),
hashes_by_tier: BTreeMap::new(),
blocks: if with_blocks {
vec![None; hashes.len()]
} else {
Vec::new()
},
}
});
if with_blocks {
let components = view.blocks[index].get_or_insert_with(|| BlockComponents {
token_count: block.token_count,
tier_masks: Vec::new(),
});
components.tier_masks.push((*tier, *mask));
} else {
view.hashes_by_tier.entry(*tier).or_default().push((
*hash,
*mask,
block.token_count,
));
}
}
}
let workers = worker_order
.into_iter()
.filter_map(|worker| by_worker.remove(worker))
.collect();
(workers, matched_hashes)
}
fn collect_prefix_inputs_locked(state: &State, hashes: &[i64]) -> Vec<WorkerPrefixInput> {
Self::collect_worker_views(state, hashes, true)
.0
.into_iter()
.map(|worker| WorkerPrefixInput {
worker_id: worker.worker_id,
address: worker.address,
spec: worker.spec,
blocks: worker.blocks,
})
.collect()
}
fn do_match_prefix(
&self,
req: MatchExternalKvPrefixRequest,
) -> Result<MatchExternalKvPrefixResponse, Status> {
let limit = prefix_limit(req.hashes.len(), req.max_blocks);
let hashes = &req.hashes[..limit];
if hashes.is_empty() {
return Ok(MatchExternalKvPrefixResponse::default());
}
// Only workers holding block zero can own a non-empty prefix, so the
// candidate set is fixed up front and each candidate reuses one scanner and
// one block view. Allocation is O(first-block holders) whatever the request
// length, which is why there is no scan cap: length costs time, not memory.
let state = self.read_state()?;
let Some(first) = state
.blocks
.get(&hashes[0])
.filter(|block| !block.placements.is_empty())
else {
return Ok(MatchExternalKvPrefixResponse {
matches: Vec::new(),
best_prefix_blocks: 0,
blocks_read: 1,
});
};
let mut seen = HashSet::new();
let mut candidates: Vec<PrefixCandidate> = first
.placements
.keys()
.filter(|(worker, _)| seen.insert(worker.as_str()))
.map(|(worker, _)| {
let metadata = state.workers.get(worker);
PrefixCandidate {
address: metadata
.map(|worker| worker.address.clone())
.unwrap_or_default(),
scanner: WorkerPrefixScanner::new(
metadata.and_then(|worker| worker.spec.as_ref()),
),
worker_id: worker.clone(),
}
})
.collect();
let candidate_by_id: HashMap<String, usize> = candidates
.iter()
.enumerate()
.map(|(index, candidate)| (candidate.worker_id.clone(), index))
.collect();
let mut present = vec![false; candidates.len()];
let mut block_views: Vec<BlockComponents> = (0..candidates.len())
.map(|_| BlockComponents {
token_count: 0,
tier_masks: Vec::new(),
})
.collect();
for hash in hashes {
present.fill(false);
for block in &mut block_views {
block.token_count = 0;
block.tier_masks.clear();
}
if let Some(block) = state.blocks.get(hash) {
for ((worker, tier), mask) in &block.placements {
let Some(&index) = candidate_by_id.get(worker) else {
continue;
};
present[index] = true;
block_views[index].token_count = block.token_count;
block_views[index].tier_masks.push((*tier, *mask));
}
}
for (index, candidate) in candidates.iter_mut().enumerate() {
candidate
.scanner
.push(present[index].then_some(&block_views[index]));
}
}
let entries = candidates
.into_iter()
.filter_map(|candidate| {
let prefix = candidate.scanner.prefix();
(!candidate.address.is_empty() && prefix > 0).then_some((
candidate.worker_id,
candidate.address,
prefix,
))
})
.collect();
Ok(assemble_prefix_response(entries, limit as u32))
}
fn do_hit_counts(
&self,
req: GetExternalKvHitCountsRequest,
) -> Result<GetExternalKvHitCountsResponse, Status> {
let state = self.read_state()?;
let entries = dedup_preserve_order(&req.hashes)
.into_iter()
.filter_map(|hash| {
state
.hit_counts
.get(&hash)
.copied()
.map(|hit_count_total| HitCountEntry {
hash,
hit_count_total,
})
})
.collect();
Ok(GetExternalKvHitCountsResponse { entries })
}
}
fn revoke_one(state: &mut State, worker_id: &str, hash: &i64, tier: i32) {
let mut remove_block = false;
if let Some(block) = state.blocks.get_mut(hash) {
block.placements.remove(&(worker_id.to_string(), tier));
remove_block = block.placements.is_empty();
}
if let Some(worker) = state.workers.get_mut(worker_id) {
if let Some(hashes) = worker.holdings.get_mut(&tier) {
hashes.remove(hash);
if hashes.is_empty() {
worker.holdings.remove(&tier);
}
}
}
if remove_block {
state.blocks.remove(hash);
state.hit_counts.remove(hash);
}
}
fn dedup_preserve_order(hashes: &[i64]) -> Vec<i64> {
let mut seen = HashSet::new();
hashes
.iter()
.filter(|hash| seen.insert(**hash))
.copied()
.collect()
}
#[tonic::async_trait]
impl KvIndexerBackend for InMemoryKvIndexerBackend {
async fn apply_external_kv_batch(
&self,
request: ApplyExternalKvBatchRequest,
) -> Result<ApplyExternalKvBatchResponse, Status> {
self.apply(request)
}
async fn match_external_kv(
&self,
request: MatchExternalKvRequest,
) -> Result<MatchExternalKvResponse, Status> {
self.do_match(request)
}
async fn collect_worker_prefix_inputs(
&self,
hashes: &[i64],
) -> Result<Vec<WorkerPrefixInput>, Status> {
let state = self.read_state()?;
Ok(Self::collect_prefix_inputs_locked(&state, hashes))
}
async fn match_external_kv_prefix(
&self,
request: MatchExternalKvPrefixRequest,
) -> Result<MatchExternalKvPrefixResponse, Status> {
self.do_match_prefix(request)
}
async fn get_external_kv_hit_counts(
&self,
request: GetExternalKvHitCountsRequest,
) -> Result<GetExternalKvHitCountsResponse, Status> {
self.do_hit_counts(request)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{mpsc, Arc};
use std::time::Duration;
#[test]
fn dedup_keeps_first_seen_order() {
let hashes = vec![1, -2, 1, 3];
assert_eq!(dedup_preserve_order(&hashes), vec![1, -2, 3]);
}
#[test]
fn match_without_hit_count_uses_shared_lock() {
let backend = Arc::new(InMemoryKvIndexerBackend::new());
let read_guard = backend.read_state().unwrap();
let query_backend = Arc::clone(&backend);
let (result_tx, result_rx) = mpsc::channel();
let query = std::thread::spawn(move || {
result_tx
.send(query_backend.do_match(MatchExternalKvRequest {
hashes: vec![-1],
count_as_hit: false,
}))
.unwrap();
});
result_rx
.recv_timeout(Duration::from_secs(1))
.expect("read-only match should not wait for another reader")
.unwrap();
drop(read_guard);
query.join().unwrap();
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,42 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Termination signalling shared by the server and bridge binaries.
use tracing::{info, warn};
/// Resolves once the process has been asked to terminate, on `SIGTERM` (which
/// container runtimes send before escalating to `SIGKILL`) or `SIGINT`.
///
/// A handler that cannot be installed never resolves: firing immediately would
/// shut the process down at startup instead of leaving it running unsupervised.
pub async fn shutdown_signal() {
let interrupt = async {
if let Err(error) = tokio::signal::ctrl_c().await {
warn!(%error, "cannot handle SIGINT");
std::future::pending::<()>().await;
}
};
#[cfg(unix)]
let terminate = async {
use tokio::signal::unix::{signal, SignalKind};
match signal(SignalKind::terminate()) {
Ok(mut stream) => {
stream.recv().await;
}
Err(error) => {
warn!(%error, "cannot handle SIGTERM");
std::future::pending::<()>().await;
}
}
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
() = interrupt => info!("received SIGINT; shutting down"),
() = terminate => info!("received SIGTERM; shutting down"),
}
}
@@ -0,0 +1,11 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use std::time::{SystemTime, UNIX_EPOCH};
pub fn nanos() -> u128 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
}
@@ -0,0 +1,57 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use sgl_kv_indexer::pb::{
ApplyExternalKvBatchRequest, ExternalKvAction, ExternalKvActionType, TierType,
};
pub fn hbm() -> i32 {
TierType::TierHbm as i32
}
pub fn dram() -> i32 {
TierType::TierDram as i32
}
pub fn action(kind: ExternalKvActionType, tier: i32, hashes: &[i64]) -> ExternalKvAction {
ExternalKvAction {
r#type: kind as i32,
tier,
hashes: hashes.to_vec(),
component_masks: Vec::new(),
block_sizes: Vec::new(),
}
}
/// A component-aware REPORT action: each hash carries its component bitmask and
/// token count, index-aligned with `hashes`.
#[allow(dead_code)] // used by memory_integration, not grpc_contract
pub fn component_report(
tier: i32,
hashes: &[i64],
masks: &[u32],
block_sizes: &[u32],
) -> ExternalKvAction {
ExternalKvAction {
r#type: ExternalKvActionType::ActionReport as i32,
tier,
hashes: hashes.to_vec(),
component_masks: masks.to_vec(),
block_sizes: block_sizes.to_vec(),
}
}
pub fn apply_request(
worker: &str,
address: &str,
seq: u64,
actions: Vec<ExternalKvAction>,
) -> ApplyExternalKvBatchRequest {
ApplyExternalKvBatchRequest {
worker_id: worker.to_string(),
seq,
actions,
worker_address: address.to_string(),
cache_spec: None,
}
}
@@ -0,0 +1,13 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
use std::net::SocketAddr;
/// Reserves an ephemeral loopback port. The listener is dropped immediately,
/// so callers that spawn a server should retain their connect-retry loop.
pub fn free_addr() -> SocketAddr {
std::net::TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
}
@@ -0,0 +1,517 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! gRPC contract tests: exercise all four RPCs of the `KVIndexer` service
//! over the wire (real tonic server + client), not just the backend trait.
#[path = "common/id.rs"]
mod test_id;
#[allow(dead_code)]
#[path = "common/kv.rs"]
mod test_kv;
#[path = "common/net.rs"]
mod test_net;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use prost::Message;
use tokio::sync::Semaphore;
use tonic::transport::Server;
use tonic::{Code, Status};
use sgl_kv_indexer::pb::kv_indexer_client::KvIndexerClient;
use sgl_kv_indexer::pb::kv_indexer_server::KvIndexerServer;
use sgl_kv_indexer::pb::{
ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvAction,
ExternalKvActionType, GetExternalKvHitCountsRequest, GetExternalKvHitCountsResponse,
MatchExternalKvPrefixRequest, MatchExternalKvPrefixResponse, MatchExternalKvRequest,
MatchExternalKvResponse,
};
use sgl_kv_indexer::{
server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerBackend, KvIndexerService,
PrefixIndex, PrefixIndexConfig, MAX_GRPC_DECODING_MESSAGE_SIZE,
};
use test_id::nanos;
use test_kv::{action, apply_request, hbm};
use test_net::free_addr;
async fn start_backend(
backend: InMemoryKvIndexerBackend,
) -> KvIndexerClient<tonic::transport::Channel> {
let svc = KvIndexerService::new(backend).into_server();
let addr = free_addr();
tokio::spawn(async move {
server_builder()
.add_service(svc)
.serve(addr)
.await
.expect("server serve");
});
let endpoint = format!("http://{addr}");
for _ in 0..50 {
if let Ok(c) = KvIndexerClient::connect(endpoint.clone()).await {
return c;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
panic!("client failed to connect to {endpoint}");
}
#[derive(Clone)]
struct BlockingPrefixBackend {
entered: Arc<AtomicUsize>,
release: Arc<Semaphore>,
}
#[tonic::async_trait]
impl KvIndexerBackend for BlockingPrefixBackend {
async fn apply_external_kv_batch(
&self,
_request: ApplyExternalKvBatchRequest,
) -> Result<ApplyExternalKvBatchResponse, Status> {
Ok(ApplyExternalKvBatchResponse::default())
}
async fn match_external_kv(
&self,
_request: MatchExternalKvRequest,
) -> Result<MatchExternalKvResponse, Status> {
Ok(MatchExternalKvResponse::default())
}
async fn match_external_kv_prefix(
&self,
_request: MatchExternalKvPrefixRequest,
) -> Result<MatchExternalKvPrefixResponse, Status> {
self.entered.fetch_add(1, Ordering::SeqCst);
let _permit = self.release.acquire().await.expect("semaphore open");
Ok(MatchExternalKvPrefixResponse::default())
}
async fn get_external_kv_hit_counts(
&self,
_request: GetExternalKvHitCountsRequest,
) -> Result<GetExternalKvHitCountsResponse, Status> {
Ok(GetExternalKvHitCountsResponse::default())
}
}
async fn start_blocking_backend(
backend: BlockingPrefixBackend,
) -> KvIndexerClient<tonic::transport::Channel> {
let svc = KvIndexerService::with_prefix_query_max_inflight(backend, 2).into_server();
let addr = free_addr();
tokio::spawn(async move {
server_builder()
.add_service(svc)
.serve(addr)
.await
.expect("server serve");
});
let endpoint = format!("http://{addr}");
for _ in 0..50 {
if let Ok(client) = KvIndexerClient::connect(endpoint.clone()).await {
return client;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
panic!("client failed to connect to {endpoint}");
}
/// Starts a real gRPC server with isolated process-local state.
async fn start() -> KvIndexerClient<tonic::transport::Channel> {
start_backend(InMemoryKvIndexerBackend::new()).await
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn prefix_limit_rejects_over_real_grpc_without_blocking_writes() {
let entered = Arc::new(AtomicUsize::new(0));
let release = Arc::new(Semaphore::new(0));
let backend = BlockingPrefixBackend {
entered: Arc::clone(&entered),
release: Arc::clone(&release),
};
let client = start_blocking_backend(backend).await;
let request = || MatchExternalKvPrefixRequest {
hashes: vec![-1],
max_blocks: 0,
};
let mut first_client = client.clone();
let first = tokio::spawn(async move { first_client.match_external_kv_prefix(request()).await });
let mut second_client = client.clone();
let second =
tokio::spawn(async move { second_client.match_external_kv_prefix(request()).await });
tokio::time::timeout(Duration::from_secs(1), async {
while entered.load(Ordering::SeqCst) != 2 {
tokio::task::yield_now().await;
}
})
.await
.expect("two prefix queries should enter the backend");
let mut rejected_client = client.clone();
let rejected = tokio::time::timeout(
Duration::from_secs(1),
rejected_client.match_external_kv_prefix(request()),
)
.await
.expect("overload response should be immediate")
.expect_err("third prefix query should be rejected");
assert_eq!(rejected.code(), Code::ResourceExhausted);
assert_eq!(entered.load(Ordering::SeqCst), 2);
let mut write_client = client.clone();
tokio::time::timeout(
Duration::from_secs(1),
write_client.apply_external_kv_batch(ApplyExternalKvBatchRequest {
worker_id: "worker".into(),
..Default::default()
}),
)
.await
.expect("writes should not share the prefix-query limit")
.expect("write should succeed");
release.add_permits(2);
first.await.expect("first task").expect("first response");
second.await.expect("second task").expect("second response");
}
fn apply(
worker: &str,
addr: &str,
seq: u64,
action_type: ExternalKvActionType,
tier: i32,
hashes: &[i64],
) -> ApplyExternalKvBatchRequest {
apply_request(worker, addr, seq, vec![action(action_type, tier, hashes)])
}
fn apply_report(
worker: &str,
addr: &str,
seq: u64,
tier: i32,
hashes: &[i64],
) -> ApplyExternalKvBatchRequest {
apply(
worker,
addr,
seq,
ExternalKvActionType::ActionReport,
tier,
hashes,
)
}
#[tokio::test]
async fn multiple_workers_share_one_indexer_server() {
let mut indexer = start().await;
let suffix = nanos();
let worker_0 = format!("worker-0-{suffix}");
let worker_1 = format!("worker-1-{suffix}");
let (hash_0, hash_1, shared_hash) = (1, 2, 3);
indexer
.apply_external_kv_batch(apply_report(
&worker_0,
"10.0.0.1:9000",
1,
hbm(),
&[hash_0, shared_hash],
))
.await
.expect("apply worker-0");
indexer
.apply_external_kv_batch(apply_report(
&worker_1,
"10.0.0.2:9000",
1,
hbm(),
&[hash_1, shared_hash],
))
.await
.expect("apply worker-1");
let response = indexer
.match_external_kv(MatchExternalKvRequest {
hashes: vec![hash_0, hash_1, shared_hash],
count_as_hit: false,
})
.await
.expect("query indexer")
.into_inner();
assert!(response
.matches
.iter()
.any(|entry| entry.worker_id == worker_0));
assert!(response
.matches
.iter()
.any(|entry| entry.worker_id == worker_1));
// Keep one wire-level smoke check for hit counting; detailed counter
// semantics live in memory_integration.rs.
indexer
.match_external_kv(MatchExternalKvRequest {
hashes: vec![hash_0],
count_as_hit: true,
})
.await
.expect("counting match over gRPC");
let miss = 4;
let counts = indexer
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest {
hashes: vec![hash_0, miss],
})
.await
.expect("hit counts over gRPC")
.into_inner();
let count = |hash: i64| {
counts
.entries
.iter()
.find(|entry| entry.hash == hash)
.map(|entry| entry.hit_count_total)
.unwrap_or(0)
};
assert!(count(hash_0) >= 1, "matched hash should have a hit");
assert_eq!(count(miss), 0, "unmatched hash must not be counted");
}
#[tokio::test]
async fn validation_errors_map_to_invalid_argument_over_grpc() {
let mut c = start().await;
let err = c
.apply_external_kv_batch(apply_report("", "addr", 1, hbm(), &[1]))
.await
.expect_err("empty worker_id must be rejected");
assert_eq!(err.code(), Code::InvalidArgument);
// An action type outside the enum can only arrive over the wire; the
// in-process tests cover the mapped `ActionUnknown` variant instead.
let unmapped_action_type = ApplyExternalKvBatchRequest {
worker_id: "w".into(),
seq: 1,
worker_address: String::new(),
cache_spec: None,
actions: vec![ExternalKvAction {
r#type: 999,
tier: hbm(),
hashes: vec![1],
component_masks: Vec::new(),
block_sizes: Vec::new(),
}],
};
let err = c
.apply_external_kv_batch(unmapped_action_type)
.await
.expect_err("unknown action type must be rejected");
assert_eq!(err.code(), Code::InvalidArgument);
}
#[tokio::test]
async fn match_prefix_over_grpc() {
let mut c = start().await;
let (w_long, w_short) = (format!("long-{}", nanos()), format!("short-{}", nanos()));
let (a, b, d) = (1, 2, 3);
c.apply_external_kv_batch(apply_report(&w_long, "10.0.0.1:9000", 1, hbm(), &[a, b, d]))
.await
.expect("apply long");
c.apply_external_kv_batch(apply_report(&w_short, "10.0.0.2:9000", 1, hbm(), &[a]))
.await
.expect("apply short");
let resp = c
.match_external_kv_prefix(MatchExternalKvPrefixRequest {
hashes: vec![a, b, d],
max_blocks: 0,
})
.await
.expect("prefix ok")
.into_inner();
assert_eq!(resp.best_prefix_blocks, 3);
assert_eq!(resp.blocks_read, 3);
// Descending by prefix length: long (3) before short (1).
assert_eq!(resp.matches.len(), 2);
assert_eq!(resp.matches[0].worker_id, w_long);
assert_eq!(resp.matches[0].matched_prefix_blocks, 3);
assert_eq!(resp.matches[0].worker_address, "10.0.0.1:9000");
assert_eq!(resp.matches[1].worker_id, w_short);
assert_eq!(resp.matches[1].matched_prefix_blocks, 1);
}
#[tokio::test]
async fn prefix_query_scans_more_than_one_apply_chunk_over_grpc() {
const APPLY_CHUNK_SIZE: usize = 16_384;
let mut indexer = start().await;
let hashes: Vec<i64> = (0..=APPLY_CHUNK_SIZE as i64).collect();
for (seq, chunk) in hashes.chunks(APPLY_CHUNK_SIZE).enumerate() {
indexer
.apply_external_kv_batch(apply_report(
"large-prefix-worker",
"10.0.0.1:9000",
seq as u64,
hbm(),
chunk,
))
.await
.expect("bounded apply chunk");
}
let response = indexer
.match_external_kv_prefix(MatchExternalKvPrefixRequest {
hashes,
max_blocks: 0,
})
.await
.expect("prefix request larger than one apply chunk")
.into_inner();
assert_eq!(response.best_prefix_blocks as usize, APPLY_CHUNK_SIZE + 1);
assert_eq!(response.blocks_read as usize, APPLY_CHUNK_SIZE + 1);
}
#[tokio::test]
async fn packed_signed_hash_query_can_exceed_tonics_default_receive_limit() {
const TONIC_DEFAULT_RECEIVE_LIMIT: usize = 4 * 1024 * 1024;
const HASH_COUNT: usize = 600_000;
let mut indexer = start().await;
indexer
.apply_external_kv_batch(apply_report(
"large-wire-worker",
"10.0.0.1:9000",
1,
hbm(),
&[-1],
))
.await
.expect("store the signed first hash");
let mut hashes = Vec::with_capacity(HASH_COUNT);
hashes.push(-1);
hashes.extend((1..HASH_COUNT).map(|value| value as i64));
let request = MatchExternalKvPrefixRequest {
hashes,
max_blocks: 0,
};
assert!(request.encoded_len() > TONIC_DEFAULT_RECEIVE_LIMIT);
assert!(request.encoded_len() < MAX_GRPC_DECODING_MESSAGE_SIZE);
let response = indexer
.match_external_kv_prefix(request)
.await
.expect("configured server accepts a packed request larger than 4 MiB")
.into_inner();
assert_eq!(response.best_prefix_blocks, 1);
}
/// Past the configured ceiling the server must answer OUT_OF_RANGE, because that
/// is the code the router maps to a degraded (cache-affinity-free) route rather
/// than to a failed request. A different code there would fail the request.
#[tokio::test]
async fn query_past_the_configured_limit_is_refused_as_out_of_range() {
let hash_count = MAX_GRPC_DECODING_MESSAGE_SIZE / std::mem::size_of::<i64>() + 1_024;
let request = MatchExternalKvPrefixRequest {
hashes: (0..hash_count).map(|value| value as i64).collect(),
max_blocks: 0,
};
assert!(request.encoded_len() > MAX_GRPC_DECODING_MESSAGE_SIZE);
let status = start()
.await
.match_external_kv_prefix(request)
.await
.expect_err("a request past the ceiling must be refused");
assert_eq!(status.code(), Code::OutOfRange);
}
/// Serves an empty backend behind an interceptor that records the `grpc-timeout`
/// of every request, and returns the router-facing client alongside the capture.
async fn start_recording_deadlines(
query_deadline: Duration,
) -> (GrpcPrefixIndex, Arc<Mutex<Vec<String>>>) {
let seen = Arc::new(Mutex::new(Vec::new()));
let recorder = Arc::clone(&seen);
let svc = KvIndexerServer::with_interceptor(
KvIndexerService::new(InMemoryKvIndexerBackend::new()),
move |request: tonic::Request<()>| {
if let Some(timeout) = request.metadata().get("grpc-timeout") {
recorder
.lock()
.expect("deadline recorder")
.push(timeout.to_str().expect("ascii timeout").to_string());
}
Ok(request)
},
);
let addr = free_addr();
tokio::spawn(async move {
Server::builder()
.add_service(svc)
.serve(addr)
.await
.expect("server serve");
});
let endpoint = format!("http://{addr}");
for _ in 0..50 {
if KvIndexerClient::connect(endpoint.clone()).await.is_ok() {
break;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
let index = GrpcPrefixIndex::new(PrefixIndexConfig {
endpoint,
query_deadline,
max_inflight: sgl_kv_indexer::DEFAULT_QUERY_MAX_INFLIGHT,
})
.expect("test endpoint is valid");
(index, seen)
}
/// The router-facing client must publish its deadline on the wire: that header is
/// the only thing letting the indexer shed a query whose caller gave up.
#[tokio::test]
async fn router_client_publishes_its_deadline_on_the_wire() {
let (index, seen) = start_recording_deadlines(Duration::from_millis(250)).await;
index
.match_prefix(vec![1, 2, 3])
.await
.expect("query reaches the indexer");
let seen = seen.lock().expect("deadline recorder").clone();
assert_eq!(
seen.len(),
1,
"exactly one query reached the server: {seen:?}"
);
let raw = &seen[0];
// Asserted structurally, not byte-for-byte: the wire spec lets the sender
// pick any unit that fits, so pinning tonic's choice would fail on a
// change that is still correct.
let (digits, unit) = raw.split_at(raw.len() - 1);
assert!(
matches!(unit, "H" | "M" | "S" | "m" | "u" | "n"),
"unit is one the wire spec defines: {raw:?}"
);
let value: u64 = digits.parse().expect("timeout value is numeric");
assert!(
value > 0,
"a budget of zero would shed every query: {raw:?}"
);
}
@@ -0,0 +1,886 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Integration tests for the process-local in-memory backend.
#[path = "common/kv.rs"]
mod test_kv;
use std::sync::Arc;
use sgl_kv_indexer::pb::{
ApplyExternalKvBatchRequest, ApplyExternalKvBatchResponse, ExternalKvActionType,
GetExternalKvHitCountsRequest, GetExternalKvHitCountsResponse, MatchExternalKvPrefixRequest,
MatchExternalKvPrefixResponse, MatchExternalKvRequest, MatchExternalKvResponse,
WorkerCacheSpec,
};
use sgl_kv_indexer::{
InMemoryKvIndexerBackend, KvIndexerBackend, WorkerPrefixInput, COMPONENT_FULL, COMPONENT_SWA,
};
use test_kv::{action, apply_request as apply_req, component_report, dram, hbm};
use tonic::Status;
fn backend() -> InMemoryKvIndexerBackend {
InMemoryKvIndexerBackend::new()
}
fn match_req(hs: &[i64], count_as_hit: bool) -> MatchExternalKvRequest {
MatchExternalKvRequest {
hashes: hs.to_vec(),
count_as_hit,
}
}
/// Returns the tiers a worker holds a hash at, per the match response.
fn tiers_for(resp: &MatchExternalKvResponse, worker: &str, hash: i64) -> Vec<i32> {
let mut tiers = Vec::new();
for m in &resp.matches {
if m.worker_id != worker {
continue;
}
for th in &m.hashes_by_tier {
if th.hashes.contains(&hash) {
tiers.push(th.tier);
}
}
}
tiers.sort_unstable();
tiers
}
macro_rules! itest {
($name:ident, $b:ident, $body:block) => {
#[tokio::test]
async fn $name() {
let $b = backend();
$body
}
};
}
itest!(report_then_match_returns_worker_and_address, b, {
b.apply_external_kv_batch(apply_req(
"w1",
"10.0.0.1:9000",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1, 2])],
))
.await
.unwrap();
let resp = b
.match_external_kv(match_req(&[1, 2, 3], false))
.await
.unwrap();
assert_eq!(resp.matches.len(), 1);
let m = &resp.matches[0];
assert_eq!(m.worker_id, "w1");
assert_eq!(m.address, "10.0.0.1:9000");
assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]);
assert_eq!(tiers_for(&resp, "w1", 2), vec![hbm()]);
assert!(tiers_for(&resp, "w1", 3).is_empty());
});
itest!(large_request_preserves_complete_ordered_results, b, {
// Exercise a large write and read while preserving complete ordered results.
let expected_hashes: Vec<i64> = (0..300).collect();
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(
ExternalKvActionType::ActionReport,
hbm(),
&expected_hashes,
)],
))
.await
.unwrap();
let resp = b
.match_external_kv(match_req(&expected_hashes, false))
.await
.unwrap();
let worker = resp
.matches
.iter()
.find(|m| m.worker_id == "w1")
.expect("worker must match");
let tier = worker
.hashes_by_tier
.iter()
.find(|t| t.tier == hbm())
.expect("HBM tier must match");
assert_eq!(tier.hashes, expected_hashes);
});
itest!(duplicate_report_is_idempotent, b, {
for _ in 0..3 {
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
))
.await
.unwrap();
}
let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap();
assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]);
});
itest!(identical_batch_replay_is_idempotent, b, {
// Stores, removes, then stores the same hash again; the net state is
// "stored". Re-delivering the identical batch must not change it.
let batch = apply_req(
"w1",
"a",
7,
vec![
action(ExternalKvActionType::ActionReport, hbm(), &[9]),
action(ExternalKvActionType::ActionRevoke, hbm(), &[9]),
action(ExternalKvActionType::ActionReport, hbm(), &[9]),
],
);
b.apply_external_kv_batch(batch.clone()).await.unwrap();
let first = b.match_external_kv(match_req(&[9], false)).await.unwrap();
b.apply_external_kv_batch(batch).await.unwrap();
let second = b.match_external_kv(match_req(&[9], false)).await.unwrap();
assert_eq!(tiers_for(&first, "w1", 9), vec![hbm()]);
assert_eq!(tiers_for(&second, "w1", 9), vec![hbm()]);
});
itest!(recomputed_full_node_restores_hbm_placement, b, {
// HiRadixCache lifecycle for an exact-match recomputation:
// BlockStored(GPU) -> BlockRemoved(GPU) -> BlockStored(GPU).
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
))
.await
.unwrap();
b.apply_external_kv_batch(apply_req(
"w1",
"a",
2,
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])],
))
.await
.unwrap();
let evicted = b.match_external_kv(match_req(&[1], false)).await.unwrap();
assert!(tiers_for(&evicted, "w1", 1).is_empty());
b.apply_external_kv_batch(apply_req(
"w1",
"a",
3,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
))
.await
.unwrap();
let restored = b.match_external_kv(match_req(&[1], false)).await.unwrap();
assert_eq!(tiers_for(&restored, "w1", 1), vec![hbm()]);
});
itest!(recomputed_split_reports_only_materialized_hashes, b, {
// An evicted [prefix -> old suffix] is partially recomputed as
// [prefix -> new suffix]. The old suffix must remain absent.
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1, 2])],
))
.await
.unwrap();
b.apply_external_kv_batch(apply_req(
"w1",
"a",
2,
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1, 2])],
))
.await
.unwrap();
b.apply_external_kv_batch(apply_req(
"w1",
"a",
3,
vec![
action(ExternalKvActionType::ActionReport, hbm(), &[1]),
action(ExternalKvActionType::ActionReport, hbm(), &[3]),
],
))
.await
.unwrap();
let result = b
.match_external_kv(match_req(&[1, 2, 3], false))
.await
.unwrap();
assert_eq!(tiers_for(&result, "w1", 1), vec![hbm()]);
assert!(tiers_for(&result, "w1", 2).is_empty());
assert_eq!(tiers_for(&result, "w1", 3), vec![hbm()]);
});
itest!(recomputed_batch_replay_keeps_cpu_copy, b, {
// Re-materializing on GPU must not revoke the existing host backup, and
// re-delivering the same batch must leave both tiers unchanged.
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(ExternalKvActionType::ActionReport, dram(), &[1])],
))
.await
.unwrap();
let recomputed = apply_req(
"w1",
"a",
2,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
);
b.apply_external_kv_batch(recomputed.clone()).await.unwrap();
b.apply_external_kv_batch(recomputed).await.unwrap();
let result = b.match_external_kv(match_req(&[1], false)).await.unwrap();
assert_eq!(tiers_for(&result, "w1", 1), vec![hbm(), dram()]);
});
itest!(revoke_partial_tier_keeps_other_tier, b, {
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![
action(ExternalKvActionType::ActionReport, hbm(), &[1]),
action(ExternalKvActionType::ActionReport, dram(), &[1]),
action(ExternalKvActionType::ActionRevoke, hbm(), &[1]),
],
))
.await
.unwrap();
let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap();
assert_eq!(tiers_for(&resp, "w1", 1), vec![dram()]);
});
itest!(revoke_missing_hash_is_idempotent, b, {
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[404])],
))
.await
.unwrap();
let resp = b.match_external_kv(match_req(&[404], false)).await.unwrap();
assert!(resp.matches.is_empty());
});
itest!(multi_worker_multi_tier, b, {
b.apply_external_kv_batch(apply_req(
"w1",
"a1",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
))
.await
.unwrap();
b.apply_external_kv_batch(apply_req(
"w2",
"a2",
1,
vec![action(ExternalKvActionType::ActionReport, dram(), &[1])],
))
.await
.unwrap();
let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap();
assert_eq!(resp.matches.len(), 2);
assert_eq!(tiers_for(&resp, "w1", 1), vec![hbm()]);
assert_eq!(tiers_for(&resp, "w2", 1), vec![dram()]);
});
itest!(clear_all_at_tier_removes_only_that_tier, b, {
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![
action(ExternalKvActionType::ActionReport, hbm(), &[1, 2, 3]),
action(ExternalKvActionType::ActionReport, dram(), &[1]),
],
))
.await
.unwrap();
b.apply_external_kv_batch(apply_req(
"w1",
"a",
2,
vec![action(
ExternalKvActionType::ActionClearAllAtTier,
hbm(),
&[],
)],
))
.await
.unwrap();
let resp = b
.match_external_kv(match_req(&[1, 2, 3], false))
.await
.unwrap();
assert_eq!(tiers_for(&resp, "w1", 1), vec![dram()]);
assert!(tiers_for(&resp, "w1", 2).is_empty());
assert!(tiers_for(&resp, "w1", 3).is_empty());
});
itest!(
count_as_hit_only_counts_matched_and_replay_does_not_double,
b,
{
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
))
.await
.unwrap();
// Diagnostic match (count_as_hit=false) must not count.
b.match_external_kv(match_req(&[1, 2], false))
.await
.unwrap();
let counts = b
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1, 2] })
.await
.unwrap();
assert!(counts.entries.is_empty());
// Counting match: only the matched hash "1" is counted, "2" (a miss) is not.
b.match_external_kv(match_req(&[1, 2], true)).await.unwrap();
let counts = b
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1, 2] })
.await
.unwrap();
assert_eq!(counts.entries.len(), 1);
assert_eq!(counts.entries[0].hash, 1);
assert_eq!(counts.entries[0].hit_count_total, 1);
// Replaying the apply batch must not touch hit counts.
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
))
.await
.unwrap();
let counts = b
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] })
.await
.unwrap();
assert_eq!(counts.entries[0].hit_count_total, 1);
}
);
itest!(full_revoke_drops_hit_key, b, {
// Report a block, count a hit (creates the co-located :h key), then fully
// revoke it. The hit key must go with the placement, or a
// matched-then-evicted block leaks its counter forever.
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![action(ExternalKvActionType::ActionReport, hbm(), &[1])],
))
.await
.unwrap();
// Counting match creates the hit key with c=1.
b.match_external_kv(match_req(&[1], true)).await.unwrap();
let counts = b
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] })
.await
.unwrap();
assert_eq!(counts.entries.len(), 1);
assert_eq!(counts.entries[0].hit_count_total, 1);
// Fully revoke the block: placement empties, so the hit key must go too.
b.apply_external_kv_batch(apply_req(
"w1",
"a",
2,
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])],
))
.await
.unwrap();
let resp = b.match_external_kv(match_req(&[1], false)).await.unwrap();
assert!(resp.matches.is_empty());
// Hit key is gone too: a leaked :h would still report a count here.
let counts = b
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] })
.await
.unwrap();
assert!(
counts.entries.is_empty(),
"hit key leaked after full revoke: {:?}",
counts.entries
);
});
itest!(partial_revoke_keeps_hit_key, b, {
// Block present at two tiers; count a hit, then revoke only one tier. Placement
// is still non-empty, so the hit key must survive (guard against over-deletion).
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![
action(ExternalKvActionType::ActionReport, hbm(), &[1]),
action(ExternalKvActionType::ActionReport, dram(), &[1]),
],
))
.await
.unwrap();
b.match_external_kv(match_req(&[1], true)).await.unwrap();
b.apply_external_kv_batch(apply_req(
"w1",
"a",
2,
vec![action(ExternalKvActionType::ActionRevoke, hbm(), &[1])],
))
.await
.unwrap();
let counts = b
.get_external_kv_hit_counts(GetExternalKvHitCountsRequest { hashes: vec![1] })
.await
.unwrap();
assert_eq!(counts.entries.len(), 1);
assert_eq!(counts.entries[0].hit_count_total, 1);
});
itest!(batch_action_order_is_preserved, b, {
// revoke-then-report on the same hash within one batch must net to "stored".
b.apply_external_kv_batch(apply_req(
"w1",
"a",
1,
vec![
action(ExternalKvActionType::ActionRevoke, hbm(), &[5]),
action(ExternalKvActionType::ActionReport, hbm(), &[5]),
],
))
.await
.unwrap();
let resp = b.match_external_kv(match_req(&[5], false)).await.unwrap();
assert_eq!(tiers_for(&resp, "w1", 5), vec![hbm()]);
// report-then-revoke on the same hash must net to "absent".
b.apply_external_kv_batch(apply_req(
"w1",
"a",
2,
vec![
action(ExternalKvActionType::ActionReport, hbm(), &[6]),
action(ExternalKvActionType::ActionRevoke, hbm(), &[6]),
],
))
.await
.unwrap();
let resp = b.match_external_kv(match_req(&[6], false)).await.unwrap();
assert!(tiers_for(&resp, "w1", 6).is_empty());
});
// --- prefix query: backend override vs. the trait's default implementation ---
//
// The trait default is the written semantics and the backend override is a read
// optimization, so they must agree field-for-field on the parts that ARE the
// contract (per-worker prefix set and best_prefix_blocks). `blocks_read` is
// observability and legitimately differs, so it is not compared.
/// Delegates every RPC to an in-memory backend EXCEPT `match_external_kv_prefix`,
/// which it leaves to the trait default — giving a reference answer computed from
/// the same state the optimized path reads.
struct DefaultViaMemory(Arc<InMemoryKvIndexerBackend>);
#[tonic::async_trait]
impl KvIndexerBackend for DefaultViaMemory {
async fn apply_external_kv_batch(
&self,
request: ApplyExternalKvBatchRequest,
) -> Result<ApplyExternalKvBatchResponse, Status> {
self.0.apply_external_kv_batch(request).await
}
async fn match_external_kv(
&self,
request: MatchExternalKvRequest,
) -> Result<MatchExternalKvResponse, Status> {
self.0.match_external_kv(request).await
}
// Delegate the component-aware read to the same backend so the trait
// default computes over the same placement and specs the fast path sees.
async fn collect_worker_prefix_inputs(
&self,
hashes: &[i64],
) -> Result<Vec<WorkerPrefixInput>, Status> {
self.0.collect_worker_prefix_inputs(hashes).await
}
async fn get_external_kv_hit_counts(
&self,
request: GetExternalKvHitCountsRequest,
) -> Result<GetExternalKvHitCountsResponse, Status> {
self.0.get_external_kv_hit_counts(request).await
}
}
fn shared_state_pair() -> (Arc<InMemoryKvIndexerBackend>, DefaultViaMemory) {
let backend = Arc::new(InMemoryKvIndexerBackend::new());
let reference = DefaultViaMemory(Arc::clone(&backend));
(backend, reference)
}
fn prefix_req(hs: &[i64]) -> MatchExternalKvPrefixRequest {
MatchExternalKvPrefixRequest {
hashes: hs.to_vec(),
max_blocks: 0,
}
}
/// Sorted `(worker_id, matched_prefix_blocks)` — the semantic content of a
/// prefix response, independent of `blocks_read`.
fn prefix_pairs(resp: &MatchExternalKvPrefixResponse) -> Vec<(String, u32)> {
let mut pairs: Vec<(String, u32)> = resp
.matches
.iter()
.map(|m| (m.worker_id.clone(), m.matched_prefix_blocks))
.collect();
pairs.sort();
pairs
}
fn report(worker: &str, addr: &str, seq: u64, hs: &[i64]) -> ApplyExternalKvBatchRequest {
apply_req(
worker,
addr,
seq,
vec![action(ExternalKvActionType::ActionReport, hbm(), hs)],
)
}
#[tokio::test]
async fn prefix_fast_path_matches_default_impl() {
let (fast, reference) = shared_state_pair();
// Nested prefixes (hole-free), a diverging branch, and a hole.
fast.apply_external_kv_batch(report("w-long", "10.0.0.1:1", 1, &[1, 2, 3, 4]))
.await
.unwrap();
fast.apply_external_kv_batch(report("w-short", "10.0.0.2:1", 1, &[1, 2]))
.await
.unwrap();
// w-hole holds 1, 3, 4 but not 2: strict prefix must be 1.
fast.apply_external_kv_batch(report("w-hole", "10.0.0.3:1", 1, &[1, 3, 4]))
.await
.unwrap();
// w-noaddr is unroutable and must be excluded by both paths.
fast.apply_external_kv_batch(report("w-noaddr", "", 1, &[1, 2]))
.await
.unwrap();
let query = [1, 2, 3, 4];
let fast_resp = fast
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
let ref_resp = reference
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
assert_eq!(prefix_pairs(&fast_resp), prefix_pairs(&ref_resp));
assert_eq!(fast_resp.best_prefix_blocks, ref_resp.best_prefix_blocks);
assert_eq!(fast_resp.best_prefix_blocks, 4);
assert_eq!(
prefix_pairs(&fast_resp),
vec![
("w-hole".to_string(), 1),
("w-long".to_string(), 4),
("w-short".to_string(), 2),
]
);
assert!(fast_resp
.matches
.iter()
.all(|m| !m.worker_address.is_empty()));
// Descending order and first-block read are part of the response contract.
assert_eq!(fast_resp.matches[0].matched_prefix_blocks, 4);
assert!(fast_resp.blocks_read >= 1);
}
#[tokio::test]
async fn prefix_first_block_miss_reads_one_block() {
let b = backend();
// No worker holds the first queried block; the scan stops after one read.
b.apply_external_kv_batch(report("w1", "10.0.0.1:1", 1, &[2, 3]))
.await
.unwrap();
let resp = b
.match_external_kv_prefix(prefix_req(&[1, 2, 3]))
.await
.unwrap();
assert!(resp.matches.is_empty());
assert_eq!(resp.best_prefix_blocks, 0);
assert_eq!(resp.blocks_read, 1);
}
#[tokio::test]
async fn prefix_max_blocks_caps_the_scan() {
let b = backend();
b.apply_external_kv_batch(report("w1", "10.0.0.1:1", 1, &[1, 2, 3, 4]))
.await
.unwrap();
let resp = b
.match_external_kv_prefix(MatchExternalKvPrefixRequest {
hashes: vec![1, 2, 3, 4],
max_blocks: 2,
})
.await
.unwrap();
// Capped at 2 even though the worker holds all four.
assert_eq!(resp.best_prefix_blocks, 2);
assert_eq!(resp.blocks_read, 2);
assert_eq!(resp.matches.len(), 1);
assert_eq!(resp.matches[0].matched_prefix_blocks, 2);
}
// --- component-aware placement & prefix -------------------------------------
/// A hybrid-SWA spec: full servable from HBM+DRAM, swa a 100-token trailing
/// window servable from HBM.
fn swa_spec() -> WorkerCacheSpec {
WorkerCacheSpec {
version: 1,
components: COMPONENT_FULL | COMPONENT_SWA,
swa_window_tokens: 100,
full_tier_mask: (1 << hbm()) | (1 << dram()),
swa_tier_mask: 1 << hbm(),
mamba_tier_mask: 0,
}
}
fn apply_with_spec(
worker: &str,
addr: &str,
seq: u64,
spec: WorkerCacheSpec,
actions: Vec<sgl_kv_indexer::pb::ExternalKvAction>,
) -> ApplyExternalKvBatchRequest {
let mut req = apply_req(worker, addr, seq, actions);
req.cache_spec = Some(spec);
req
}
#[tokio::test]
async fn component_prefix_matches_default_impl() {
let (fast, reference) = shared_state_pair();
// Four full blocks (50 tokens each); swa present on all but the 4th, so the
// largest boundary with an unbroken 100-token swa window is 3.
let report = component_report(
hbm(),
&[1, 2, 3, 4],
&[
COMPONENT_FULL | COMPONENT_SWA,
COMPONENT_FULL | COMPONENT_SWA,
COMPONENT_FULL | COMPONENT_SWA,
COMPONENT_FULL,
],
&[50, 50, 50, 50],
);
fast.apply_external_kv_batch(apply_with_spec(
"w-swa",
"10.0.0.1:1",
1,
swa_spec(),
vec![report],
))
.await
.unwrap();
let query = [1, 2, 3, 4];
let fast_resp = fast
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
let ref_resp = reference
.match_external_kv_prefix(prefix_req(&query))
.await
.unwrap();
assert_eq!(prefix_pairs(&fast_resp), prefix_pairs(&ref_resp));
assert_eq!(fast_resp.best_prefix_blocks, ref_resp.best_prefix_blocks);
assert_eq!(prefix_pairs(&fast_resp), vec![("w-swa".to_string(), 3)]);
}
#[tokio::test]
async fn partial_eviction_replace_shrinks_component_set() {
let b = backend();
// Store full+swa, then restate to full-only (partial swa eviction) via a
// REPLACE snapshot for the same (hash, tier). No BlockRemoved is involved.
b.apply_external_kv_batch(apply_with_spec(
"w1",
"10.0.0.1:1",
1,
swa_spec(),
vec![component_report(
hbm(),
&[1, 2],
&[
COMPONENT_FULL | COMPONENT_SWA,
COMPONENT_FULL | COMPONENT_SWA,
],
&[80, 80],
)],
))
.await
.unwrap();
// Both blocks reusable (window 100 met by 2x80 tokens; head rule also holds).
let before = b
.match_external_kv_prefix(prefix_req(&[1, 2]))
.await
.unwrap();
assert_eq!(before.best_prefix_blocks, 2);
// Restate the second block to full only: swa gone there.
b.apply_external_kv_batch(apply_with_spec(
"w1",
"10.0.0.1:1",
2,
swa_spec(),
vec![component_report(hbm(), &[2], &[COMPONENT_FULL], &[80])],
))
.await
.unwrap();
// That block has no swa now, and its trailing window (only 80 < 100) is not
// headed, so the largest valid boundary drops to 1.
let after = b
.match_external_kv_prefix(prefix_req(&[1, 2]))
.await
.unwrap();
assert_eq!(after.best_prefix_blocks, 1);
let snapshot = b
.match_external_kv(match_req(&[1, 2], false))
.await
.unwrap();
let tier = &snapshot.matches[0].hashes_by_tier[0];
assert_eq!(tier.hashes, vec![1, 2]);
assert_eq!(
tier.component_masks,
vec![COMPONENT_FULL | COMPONENT_SWA, COMPONENT_FULL]
);
assert_eq!(tier.block_sizes, vec![80, 80]);
}
#[tokio::test]
async fn component_aware_worker_without_spec_is_excluded() {
let b = backend();
// Report component-aware placement but never send a spec: the worker cannot
// be interpreted and must be excluded (NoSignal-safe), never over-reported.
b.apply_external_kv_batch(apply_req(
"w1",
"10.0.0.1:1",
1,
vec![component_report(
hbm(),
&[1, 2],
&[COMPONENT_FULL, COMPONENT_FULL],
&[16, 16],
)],
))
.await
.unwrap();
let resp = b
.match_external_kv_prefix(prefix_req(&[1, 2]))
.await
.unwrap();
assert!(resp.matches.is_empty());
assert_eq!(resp.best_prefix_blocks, 0);
}
#[tokio::test]
async fn duplicate_hash_in_one_report_keeps_last_snapshot() {
let b = backend();
// A single REPORT action naming the same hash twice (a coalesced
// store+restate): the LAST snapshot must win deterministically, never a race.
b.apply_external_kv_batch(apply_with_spec(
"w1",
"10.0.0.1:1",
1,
swa_spec(),
vec![component_report(
hbm(),
&[1, 1],
&[COMPONENT_FULL | COMPONENT_SWA, COMPONENT_FULL],
&[80, 80],
)],
))
.await
.unwrap();
// The hash ends as full-only (last snapshot); with swa required and a lone 80-token
// block that is not a full head window, the boundary requiring swa fails,
// so no reusable prefix.
let resp = b.match_external_kv_prefix(prefix_req(&[1])).await.unwrap();
assert_eq!(resp.best_prefix_blocks, 0);
}
#[tokio::test]
async fn absent_spec_batch_clears_stored_spec() {
let b = backend();
// First a component-aware batch establishes a spec + a reusable block.
b.apply_external_kv_batch(apply_with_spec(
"w1",
"10.0.0.1:1",
1,
swa_spec(),
vec![component_report(
hbm(),
&[1],
&[COMPONENT_FULL | COMPONENT_SWA],
&[200],
)],
))
.await
.unwrap();
assert_eq!(
b.match_external_kv_prefix(prefix_req(&[1]))
.await
.unwrap()
.best_prefix_blocks,
1
);
// A later batch with NO spec (worker reverted to legacy) must clear the old
// spec. The still-component-aware placement can then no longer be interpreted
// (component data but no spec) -> fail closed, never scored on stale rules.
b.apply_external_kv_batch(apply_req(
"w1",
"10.0.0.1:1",
2,
vec![action(ExternalKvActionType::ActionReport, dram(), &[2])],
))
.await
.unwrap();
let resp = b.match_external_kv_prefix(prefix_req(&[1])).await.unwrap();
assert!(resp.matches.is_empty());
assert_eq!(resp.best_prefix_blocks, 0);
}
+157 -6
View File
@@ -12,10 +12,14 @@ use std::num::NonZeroU32;
use crate::config::{
default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs,
resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config,
DiscoveryBackend, K8sDiscoveryConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, StickyConfig,
DiscoveryBackend, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig,
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
StickyConfig,
};
const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100;
const DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT: usize = sgl_kv_indexer::DEFAULT_QUERY_MAX_INFLIGHT;
/// `sgl-router` — slim KV-aware OpenAI-compatible router for SGLang workers.
///
/// Discovery is mutually exclusive: pass `--worker-urls` for a static
@@ -69,6 +73,18 @@ pub struct Cli {
/// Multiplicative load spread gating the absolute balance check.
#[arg(long)]
pub balance_rel_threshold: Option<f32>,
/// External KV indexer gRPC endpoint used as the authoritative cache signal.
/// Needs an explicit scheme, e.g. `http://10.0.0.1:50051`.
#[arg(long)]
pub kv_indexer_endpoint: Option<String>,
/// KV Indexer query timeout in milliseconds. Requires
/// `--kv-indexer-endpoint`; defaults to 100.
#[arg(long)]
pub kv_indexer_query_timeout_ms: Option<u64>,
/// Maximum concurrent KV Indexer queries issued by this Router. Requires
/// `--kv-indexer-endpoint`; defaults to 32.
#[arg(long)]
pub kv_indexer_query_max_inflight: Option<usize>,
// ---- sticky-session policy (only used by `--policy sticky`) ----
/// Request header carrying the routing key for sticky-session routing.
@@ -157,14 +173,35 @@ impl Cli {
}
let tuned_cache_aware = self.cache_threshold.is_some()
|| self.balance_abs_threshold.is_some()
|| self.balance_rel_threshold.is_some();
|| self.balance_rel_threshold.is_some()
|| self.kv_indexer_endpoint.is_some()
|| self.kv_indexer_query_timeout_ms.is_some()
|| self.kv_indexer_query_max_inflight.is_some();
if tuned_cache_aware && self.policy != PolicyKind::CacheAwareZmq {
return Err(anyhow!(
"--cache-threshold / --balance-abs-threshold / --balance-rel-threshold \
require --policy cache_aware_zmq"
"cache-aware tuning flags require --policy cache_aware_zmq"
));
}
if self.kv_indexer_query_timeout_ms == Some(0) {
return Err(anyhow!(
"--kv-indexer-query-timeout-ms must be greater than zero"
));
}
if self.kv_indexer_query_timeout_ms.is_some() && self.kv_indexer_endpoint.is_none() {
return Err(anyhow!(
"--kv-indexer-query-timeout-ms requires --kv-indexer-endpoint"
));
}
if self.kv_indexer_query_max_inflight == Some(0) {
return Err(anyhow!(
"--kv-indexer-query-max-inflight must be greater than zero"
));
}
if self.kv_indexer_query_max_inflight.is_some() && self.kv_indexer_endpoint.is_none() {
return Err(anyhow!(
"--kv-indexer-query-max-inflight requires --kv-indexer-endpoint"
));
}
let tuned_sticky = self.routing_key_header.is_some()
|| self.sticky_fallback_policy.is_some()
|| self.sticky_idle_secs.is_some()
@@ -234,6 +271,12 @@ impl Cli {
// Only build a CacheAwareConfig when the operator tuned at least
// one knob; otherwise leave it None so the policy uses its own
// defaults. Unset knobs fall back to the per-field defaults.
let kv_indexer_query_timeout_ms = self
.kv_indexer_query_timeout_ms
.unwrap_or(DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS);
let kv_indexer_query_max_inflight = self
.kv_indexer_query_max_inflight
.unwrap_or(DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT);
let cache_aware = if tuned_cache_aware {
let d = CacheAwareConfig::default();
Some(CacheAwareConfig {
@@ -244,6 +287,11 @@ impl Cli {
balance_rel_threshold: self
.balance_rel_threshold
.unwrap_or(d.balance_rel_threshold),
kv_indexer_endpoint: self.kv_indexer_endpoint.map(|url| KvIndexerEndpointConfig {
url,
query_timeout_ms: kv_indexer_query_timeout_ms,
query_max_inflight: kv_indexer_query_max_inflight,
}),
})
} else {
None
@@ -744,6 +792,109 @@ mod tests {
assert_eq!(ca.balance_abs_threshold, 32);
}
#[test]
fn kv_indexer_reuses_cache_aware_policy_config() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--policy",
"cache_aware_zmq",
"--kv-indexer-endpoint",
"http://indexer:50051",
"--kv-indexer-query-timeout-ms",
"75",
"--kv-indexer-query-max-inflight",
"17",
]))
.unwrap();
let cache = c.model.cache_aware.expect("cache-aware config");
let indexer = cache.kv_indexer_endpoint.expect("Indexer config");
assert_eq!(indexer.url, "http://indexer:50051");
assert_eq!(indexer.query_timeout_ms, 75);
assert_eq!(indexer.query_max_inflight, 17);
}
#[test]
fn kv_indexer_uses_safe_query_defaults() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--policy",
"cache_aware_zmq",
"--kv-indexer-endpoint",
"http://indexer:50051",
]))
.unwrap();
let indexer = c
.model
.cache_aware
.expect("cache-aware config")
.kv_indexer_endpoint
.expect("Indexer config");
assert_eq!(indexer.query_timeout_ms, 100);
assert_eq!(indexer.query_max_inflight, 32);
}
#[test]
fn kv_indexer_requires_cache_aware_policy() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--kv-indexer-endpoint",
"http://indexer:50051",
]))
.unwrap_err()
.to_string();
assert!(err.contains("require --policy cache_aware_zmq"));
}
#[test]
fn kv_indexer_timeout_requires_endpoint() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--policy",
"cache_aware_zmq",
"--kv-indexer-query-timeout-ms",
"75",
]))
.unwrap_err()
.to_string();
assert!(err.contains("requires --kv-indexer-endpoint"));
}
#[test]
fn kv_indexer_max_inflight_requires_endpoint() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--policy",
"cache_aware_zmq",
"--kv-indexer-query-max-inflight",
"17",
]))
.unwrap_err()
.to_string();
assert!(err.contains("requires --kv-indexer-endpoint"));
}
#[test]
fn kv_indexer_max_inflight_must_be_positive() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--policy",
"cache_aware_zmq",
"--kv-indexer-endpoint",
"http://indexer:50051",
"--kv-indexer-query-max-inflight",
"0",
]))
.unwrap_err()
.to_string();
assert!(err.contains("must be greater than zero"));
}
#[test]
fn no_cache_aware_flags_leaves_none() {
let c = into_config_owned(with_model(&[
+15 -3
View File
@@ -146,7 +146,7 @@ pub struct ModelConfig {
pub tokenizer_path: String,
pub policy: PolicyKind,
pub circuit_breaker: Option<CircuitBreakerConfig>,
/// Tuning for the cache-aware ZMQ policy. Ignored unless
/// Tuning for cache-aware routing. Ignored unless
/// `policy = "cache_aware_zmq"`. `None` falls back to defaults at
/// policy construction time.
pub cache_aware: Option<CacheAwareConfig>,
@@ -157,8 +157,16 @@ pub struct ModelConfig {
pub sticky: Option<StickyConfig>,
}
/// Per-model cache-aware-ZMQ tuning.
#[derive(Debug, Clone, Copy)]
/// External KV Indexer client settings.
#[derive(Debug, Clone)]
pub struct KvIndexerEndpointConfig {
pub url: String,
pub query_timeout_ms: u64,
pub query_max_inflight: usize,
}
/// Per-model cache-aware tuning.
#[derive(Debug, Clone)]
pub struct CacheAwareConfig {
/// Lower bound on `matched_blocks / total_blocks` for the tree match
/// to win the selection. Below this, the policy falls back to
@@ -174,6 +182,9 @@ pub struct CacheAwareConfig {
/// that the absolute check is gated on. Default 1.1 — 10 % relative
/// difference triggers re-balancing.
pub balance_rel_threshold: f32,
/// Optional external KV Indexer client configuration. When configured, it
/// replaces the local ZMQ radix tree as the cache signal.
pub kv_indexer_endpoint: Option<KvIndexerEndpointConfig>,
}
impl Default for CacheAwareConfig {
@@ -182,6 +193,7 @@ impl Default for CacheAwareConfig {
cache_threshold: default_cache_threshold(),
balance_abs_threshold: default_balance_abs(),
balance_rel_threshold: default_balance_rel(),
kv_indexer_endpoint: None,
}
}
}
+44 -19
View File
@@ -98,19 +98,43 @@ async fn main() -> Result<()> {
);
let registry = Arc::new(sgl_router::workers::WorkerRegistry::default());
let prefix_index = cfg
.model
.cache_aware
.as_ref()
.and_then(|cache| cache.kv_indexer_endpoint.as_ref())
.map(|indexer| {
let config = sgl_kv_indexer::PrefixIndexConfig {
endpoint: indexer.url.clone(),
query_deadline: std::time::Duration::from_millis(indexer.query_timeout_ms),
max_inflight: indexer.query_max_inflight,
};
sgl_kv_indexer::GrpcPrefixIndex::new(config)
.map(Arc::new)
.context("configure KV Indexer client")
})
.transpose()?;
// Build the KV-event index up front so the cache-aware-zmq policy can
// share its `HashTree` handle + `BlockSizeOracle`. When no model uses
// `cache_aware_zmq`, the index is still constructed (cheap) but no
// subscribers are ever added.
// share its `HashTree` handle + `BlockSizeOracle`. An external Indexer makes
// the local tree irrelevant to routing, so only discover hash metadata rather
// than duplicating every KV event.
let block_size_oracle = sgl_router::policies::kv_events::BlockSizeOracle::new();
let kv_index = sgl_router::policies::kv_events::KvEventIndex::new_with_http_and_oracle(
reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(2))
.build()
.expect("default http client builds"),
Arc::clone(&block_size_oracle),
);
let kv_event_http = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(2))
.build()
.expect("default http client builds");
let kv_index = if prefix_index.is_some() {
sgl_router::policies::kv_events::KvEventIndex::new_metadata_only_with_http_and_oracle(
kv_event_http,
Arc::clone(&block_size_oracle),
)
} else {
sgl_router::policies::kv_events::KvEventIndex::new_with_http_and_oracle(
kv_event_http,
Arc::clone(&block_size_oracle),
)
};
let policies = Arc::new(
sgl_router::policies::factory::build_registry(
&cfg,
@@ -163,16 +187,17 @@ async fn main() -> Result<()> {
.context("build proxy client")?,
);
let ctx = Arc::new(
sgl_router::server::app_context::AppContext::with_active_load(
cfg.clone(),
tokenizers,
proxy,
registry,
policies,
active_load,
),
let mut app_ctx = sgl_router::server::app_context::AppContext::with_active_load(
cfg.clone(),
tokenizers,
proxy,
registry,
policies,
active_load,
);
app_ctx.prefix_index = prefix_index;
app_ctx.block_size_oracle = block_size_oracle;
let ctx = Arc::new(app_ctx);
ctx.mark_ready();
let app = sgl_router::server::app::build_router(ctx.clone());
@@ -130,6 +130,46 @@ impl CacheAwareZmqPolicy {
let rel_threshold = (min_load as f32 * self.config.balance_rel_threshold) as usize;
abs_diff > self.config.balance_abs_threshold && max_load > rel_threshold
}
fn select_external(
&self,
workers: &[Arc<Worker>],
ctx: &SelectionContext<'_>,
signal: &crate::policies::ExternalPrefixSignal,
) -> Option<Arc<Worker>> {
let sgl_kv_indexer::PrefixOutcome::Matched { matches, .. } = &signal.outcome else {
return None;
};
if signal.query_blocks == 0 {
return None;
}
// The index may include unhealthy workers or workers in another pool.
let best_routable_blocks = matches
.iter()
.filter(|m| workers.iter().any(|worker| worker.url == m.address))
.map(|m| m.matched_prefix_blocks)
.max()
.unwrap_or(0);
let match_rate = best_routable_blocks as f32 / signal.query_blocks as f32;
if let Some(metrics) = self.metrics.get() {
metrics.observe_overlap_blocks(ctx.model().0.as_str(), best_routable_blocks as u64);
}
if match_rate <= self.config.cache_threshold {
return None;
}
workers
.iter()
.filter(|worker| {
matches.iter().any(|m| {
m.matched_prefix_blocks == best_routable_blocks && m.address == worker.url
})
})
.min_by_key(|worker| worker.active_load())
.cloned()
}
}
impl Policy for CacheAwareZmqPolicy {
@@ -144,6 +184,14 @@ impl Policy for CacheAwareZmqPolicy {
return Self::pick_min_load(workers);
}
// An external signal is authoritative: an empty/unusable result
// degrades only to min-load and never consults the local radix tree.
if let Some(signal) = ctx.external_prefix() {
return self
.select_external(workers, ctx, signal)
.or_else(|| Self::pick_min_load(workers));
}
// 2. Routing tokens. Prefer the ids computed once at ingress; fall
// back to tokenizing the body here so the policy stays usable for
// callers that don't pre-tokenize (e.g. unit tests). In production
@@ -263,6 +311,7 @@ mod tests {
cache_threshold: 0.5,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
}
}
@@ -350,6 +399,70 @@ mod tests {
assert_eq!(chosen.url, "http://w1:30000");
}
#[test]
fn external_prefix_signal_skips_unroutable_best_match() {
let mut config = cfg_default();
config.cache_threshold = 0.0;
let policy = CacheAwareZmqPolicy::new(
config,
Arc::new(HashTree::new()),
tokenizer_registry_with_tiny(),
oracle_for_tests(4),
);
let w0 = worker("http://w0:30000", "tiny");
let w1 = worker("http://w1:30000", "tiny");
let _load = w0.load_guard();
let workers = vec![Arc::clone(&w0), Arc::clone(&w1)];
let signal = crate::policies::ExternalPrefixSignal {
outcome: sgl_kv_indexer::PrefixOutcome::Matched {
matches: vec![
sgl_kv_indexer::PrefixMatch {
address: "http://gone:30000".into(),
matched_prefix_blocks: 4,
worker_id: "gone".into(),
},
sgl_kv_indexer::PrefixMatch {
address: w0.url.clone(),
matched_prefix_blocks: 2,
worker_id: "w0".into(),
},
],
best_prefix_blocks: 4,
},
query_blocks: 4,
};
let model = ModelId("tiny".into());
let ctx = SelectionContext::new(&model, None).with_external_prefix(Some(&signal));
let chosen = policy.select(&workers, &ctx).expect("must pick");
assert_eq!(chosen.url, w0.url);
}
#[test]
fn external_empty_result_uses_min_load_without_local_tree() {
let tree = Arc::new(HashTree::new());
let registry = tokenizer_registry_with_tiny();
let text = "hello world hello world hello world";
let tok = registry.get("tiny").unwrap();
let ids = adapter::encode(&tok, text).unwrap();
let hashes = compute_block_hashes(&ids, 4);
tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes);
let policy = CacheAwareZmqPolicy::new(cfg_default(), tree, registry, oracle_for_tests(4));
let w0 = worker("http://w0:30000", "tiny");
let w1 = worker("http://w1:30000", "tiny");
let _load = w0.load_guard();
let workers = vec![Arc::clone(&w0), Arc::clone(&w1)];
let signal = crate::policies::ExternalPrefixSignal {
outcome: sgl_kv_indexer::PrefixOutcome::Empty,
query_blocks: 4,
};
let model = ModelId("tiny".into());
let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap();
let ctx = SelectionContext::new(&model, Some(&body)).with_external_prefix(Some(&signal));
let chosen = policy.select(&workers, &ctx).expect("must pick");
assert_eq!(chosen.url, w1.url);
}
/// Tree contains w0's prefix; cache-aware selection picks w0 even
/// though w1 has lower load (the load skew is below the imbalance
/// threshold, so cache wins).
@@ -377,6 +490,7 @@ mod tests {
cache_threshold: 0.0, // any match counts
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
registry,
@@ -414,6 +528,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
registry,
@@ -459,6 +574,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
toks,
@@ -510,6 +626,7 @@ mod tests {
cache_threshold: 1.0, // match_rate <= 1.0 always -> always fall back
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
toks,
@@ -593,6 +710,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
Arc::clone(&registry),
@@ -632,6 +750,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
Arc::clone(&registry),
@@ -690,6 +809,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
registry,
@@ -762,6 +882,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
registry,
@@ -800,6 +921,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
registry,
@@ -903,6 +1025,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
registry,
@@ -938,6 +1061,7 @@ mod tests {
cache_threshold: 0.0, // would normally always match
balance_abs_threshold: 5,
balance_rel_threshold: 2.0,
kv_indexer_endpoint: None,
},
tree,
registry,
@@ -1063,6 +1187,7 @@ mod tests {
cache_threshold: 0.99,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
tokenizer_registry_with_tiny(),
@@ -1146,6 +1271,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree.clone(),
registry,
@@ -1243,6 +1369,7 @@ mod tests {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
tree,
registry,
@@ -65,7 +65,7 @@ pub fn build_policy(
PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()),
PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()),
PolicyKind::CacheAwareZmq => {
let cache_cfg = model.cache_aware.unwrap_or_default();
let cache_cfg = model.cache_aware.clone().unwrap_or_default();
Arc::new(CacheAwareZmqPolicy::new(
cache_cfg,
tree,
@@ -65,6 +65,7 @@ struct WorkerEntry {
/// routing path entirely.
pub struct KvEventIndex {
tree: Arc<HashTree>,
maintain_tree: bool,
subscribers: Arc<KvEventSubscriberRegistry>,
pump: Mutex<Option<JoinHandle<()>>>,
pump_cancel: CancellationToken,
@@ -114,6 +115,24 @@ impl KvEventIndex {
pub fn new_with_http_and_oracle(
http: reqwest::Client,
block_size_oracle: Arc<BlockSizeOracle>,
) -> Arc<Self> {
Self::new_with_mode(http, block_size_oracle, true)
}
/// Discovers worker hash metadata only: seeds the shared [`BlockSizeOracle`]
/// but neither subscribes to KV events nor maintains the local tree, because
/// an external Indexer is the routing signal.
pub fn new_metadata_only_with_http_and_oracle(
http: reqwest::Client,
block_size_oracle: Arc<BlockSizeOracle>,
) -> Arc<Self> {
Self::new_with_mode(http, block_size_oracle, false)
}
fn new_with_mode(
http: reqwest::Client,
block_size_oracle: Arc<BlockSizeOracle>,
maintain_tree: bool,
) -> Arc<Self> {
let tree = Arc::new(HashTree::new());
let (tx, rx) = mpsc::channel::<WorkerEvent>(EVENT_CHANNEL_BUFFER);
@@ -130,6 +149,7 @@ impl KvEventIndex {
));
Arc::new(Self {
tree,
maintain_tree,
subscribers,
pump: Mutex::new(Some(pump)),
pump_cancel,
@@ -208,6 +228,15 @@ impl KvEventIndex {
// hash KV blocks over token bigrams, so the policy must use the bigram
// hasher for its query hashes to match the worker's stored hashes.
self.block_size_oracle.set_bigram(cfg.is_bigram);
if !self.maintain_tree {
info!(
worker_url = %worker_url,
block_size = cfg.block_size,
is_bigram = cfg.is_bigram,
"kv-events: external Indexer configured; discovered hash metadata without subscribing"
);
return;
}
info!(
worker_url = %worker_url,
dp_size = cfg.dp_size,
@@ -710,4 +739,28 @@ mod tests {
);
index.shutdown().await;
}
#[tokio::test]
async fn metadata_only_mode_seeds_oracle_without_registering_subscribers() {
let oracle = BlockSizeOracle::new();
let index = KvEventIndex::new_metadata_only_with_http_and_oracle(
reqwest::Client::new(),
Arc::clone(&oracle),
);
let cfg = EventConfig {
host: "127.0.0.1".into(),
port_base: 30400,
topic: "kv-events".into(),
block_size: 64,
dp_size: 2,
is_bigram: true,
};
index.add_worker("http://127.0.0.1:30400", Some(cfg)).await;
assert_eq!(oracle.get(), Some(64));
assert!(oracle.is_bigram());
assert_eq!(index.known_worker_count(), 0);
index.shutdown().await;
}
}
@@ -34,6 +34,13 @@ pub struct RequestTokens {
pub engine_equivalent: bool,
}
/// External indexer answer prepared by the async ingress path for the
/// synchronous cache-aware policy.
pub struct ExternalPrefixSignal {
pub outcome: sgl_kv_indexer::PrefixOutcome,
pub query_blocks: usize,
}
/// Produce the routing tokens — and whether they are engine-equivalent —
/// from an already-parsed request body, using the shared tokenizer registry.
///
@@ -181,6 +188,7 @@ pub struct SelectionContext<'a> {
request_body: Option<&'a [u8]>,
routing_key: Option<&'a str>,
request_tokens: Option<&'a [u32]>,
external_prefix: Option<&'a ExternalPrefixSignal>,
}
impl<'a> SelectionContext<'a> {
@@ -190,6 +198,7 @@ impl<'a> SelectionContext<'a> {
request_body,
routing_key: None,
request_tokens: None,
external_prefix: None,
}
}
@@ -203,6 +212,7 @@ impl<'a> SelectionContext<'a> {
request_body,
routing_key,
request_tokens: None,
external_prefix: None,
}
}
@@ -214,6 +224,14 @@ impl<'a> SelectionContext<'a> {
self
}
pub fn with_external_prefix(
mut self,
external_prefix: Option<&'a ExternalPrefixSignal>,
) -> Self {
self.external_prefix = external_prefix;
self
}
pub fn model(&self) -> &ModelId {
self.model
}
@@ -231,6 +249,10 @@ impl<'a> SelectionContext<'a> {
pub fn request_tokens(&self) -> Option<&[u32]> {
self.request_tokens
}
pub fn external_prefix(&self) -> Option<&ExternalPrefixSignal> {
self.external_prefix
}
}
pub trait Policy: Send + Sync + std::fmt::Debug {
@@ -4,6 +4,7 @@
use crate::config::Config;
use crate::policies::active_load::ActiveLoadRegistry;
use crate::policies::kv_events::BlockSizeOracle;
use crate::policies::PolicyRegistry;
use crate::proxy::Proxy;
use crate::server::metrics::MetricsRegistry;
@@ -12,7 +13,6 @@ use crate::workers::WorkerRegistry;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
#[derive(Debug)]
pub struct AppContext {
pub config: Config,
pub tokenizers: Arc<TokenizerRegistry>,
@@ -30,6 +30,8 @@ pub struct AppContext {
/// (active_load gauge + stale_requests_total), and PD resolver
/// (decode_affinity_total).
pub metrics: Arc<MetricsRegistry>,
pub prefix_index: Option<Arc<sgl_kv_indexer::GrpcPrefixIndex>>,
pub block_size_oracle: Arc<BlockSizeOracle>,
ready: AtomicBool,
}
@@ -82,6 +84,8 @@ impl AppContext {
policies,
active_load,
metrics,
prefix_index: None,
block_size_oracle: BlockSizeOracle::new(),
ready: AtomicBool::new(false),
}
}
@@ -127,6 +131,8 @@ impl AppContext {
policies: Arc::new(PolicyRegistry::default()),
active_load: ActiveLoadRegistry::with_defaults(),
metrics: MetricsRegistry::new(),
prefix_index: None,
block_size_oracle: BlockSizeOracle::new(),
ready: AtomicBool::new(false),
}
}
@@ -2,8 +2,9 @@
// SPDX-License-Identifier: Apache-2.0
use crate::discovery::{ModelId, WorkerMode};
use crate::policies::kv_events::{compute_block_hashes, compute_block_hashes_bigram};
use crate::policies::registry::{PdPoolResolver, PdResolveError};
use crate::policies::{request_tokens_for, RequestTokens, SelectionContext};
use crate::policies::{request_tokens_for, ExternalPrefixSignal, RequestTokens, SelectionContext};
use crate::server::app_context::AppContext;
use crate::server::error::ApiError;
use crate::server::metrics::{
@@ -16,6 +17,7 @@ use axum::http::{HeaderMap, HeaderName, HeaderValue, Response};
use bytes::Bytes;
use serde::de::IgnoredAny;
use serde::Deserialize;
use sgl_kv_indexer::PrefixIndex;
use std::collections::HashMap;
use std::sync::Arc;
@@ -167,6 +169,34 @@ pub async fn chat_completions(
let request_tokens = request_value
.as_ref()
.and_then(|v| request_tokens_for(&ctx.tokenizers, &model_id, v));
let external_prefix = match (
ctx.prefix_index.as_ref(),
request_tokens.as_ref(),
ctx.block_size_oracle.get(),
) {
(Some(index), Some(tokens), Some(block_size)) => {
let hashes = if ctx.block_size_oracle.is_bigram() {
compute_block_hashes_bigram(&tokens.ids, block_size as usize)
} else {
compute_block_hashes(&tokens.ids, block_size as usize)
};
let query_blocks = hashes.len();
let outcome = if hashes.is_empty() {
sgl_kv_indexer::PrefixOutcome::Empty
} else {
resolve_prefix_query(index.match_prefix(hashes).await, &model_str)?
};
Some(ExternalPrefixSignal {
outcome,
query_blocks,
})
}
(Some(_), _, _) => Some(ExternalPrefixSignal {
outcome: sgl_kv_indexer::PrefixOutcome::Empty,
query_blocks: 0,
}),
_ => None,
};
// Sticky-session routing key. When the sticky policy is configured,
// read the routing key from the operator-chosen header into the
@@ -181,7 +211,8 @@ pub async fn chat_completions(
.and_then(|v| v.to_str().ok())
.filter(|s| !s.is_empty());
let selection_ctx = SelectionContext::with_routing_key(&model_id, Some(&body), routing_key)
.with_request_tokens(request_tokens.as_ref().map(|t| t.ids.as_slice()));
.with_request_tokens(request_tokens.as_ref().map(|t| t.ids.as_slice()))
.with_external_prefix(external_prefix.as_ref());
let worker =
policy
.select(&workers, &selection_ctx)
@@ -620,6 +651,42 @@ pub async fn chat_completions(
}
}
fn resolve_prefix_query(
result: Result<sgl_kv_indexer::PrefixOutcome, sgl_kv_indexer::PrefixIndexError>,
model: &str,
) -> Result<sgl_kv_indexer::PrefixOutcome, ApiError> {
use sgl_kv_indexer::PrefixIndexError;
match result {
Ok(outcome) => Ok(outcome),
// The prefix hit only improves worker choice, so an indexer that is
// shedding, slow, or down costs cache affinity — not availability.
Err(
error @ (PrefixIndexError::Overloaded
| PrefixIndexError::Timeout
| PrefixIndexError::Unreachable),
) => {
tracing::warn!(%model, error = %error, "KV Indexer unavailable; falling back to min-load routing");
Ok(sgl_kv_indexer::PrefixOutcome::Empty)
}
// A prompt too long to fit one gRPC message is still a prompt a worker
// can serve, so it costs cache affinity like the cases above. Logged
// separately because the remedy is operational — raise the indexer's
// message limit — rather than waiting for the indexer to recover.
Err(error @ PrefixIndexError::QueryTooLarge) => {
tracing::warn!(%model, error = %error, "prompt exceeds the KV Indexer query size limit; falling back to min-load routing");
Ok(sgl_kv_indexer::PrefixOutcome::Empty)
}
// A rejection means the router and the indexer disagree on the request
// contract; degrading would hide that from every request.
Err(error) => {
tracing::warn!(%model, error = %error, "KV Indexer rejected the query");
Err(ApiError::PolicySelectionFailed {
model: model.to_string(),
})
}
}
}
/// Estimate prefill-token count from the raw request body for use as
/// the active-load `prefill_load` counter. Returns 1 at minimum so
/// a registered request always shows up as "load > 0" — under-counting
@@ -911,6 +978,38 @@ fn parse_probe(body: &Bytes) -> Result<RequestProbe, ApiError> {
mod tests {
use super::*;
/// An unavailable indexer must never fail a request that min-load routing
/// can still serve. `QueryTooLarge` belongs here too: a prompt that outgrows
/// the query's message limit loses cache affinity, not availability.
#[test]
fn unavailable_indexer_degrades_to_empty_prefix_signal() {
for error in [
sgl_kv_indexer::PrefixIndexError::Overloaded,
sgl_kv_indexer::PrefixIndexError::Timeout,
sgl_kv_indexer::PrefixIndexError::Unreachable,
sgl_kv_indexer::PrefixIndexError::QueryTooLarge,
] {
assert_eq!(
resolve_prefix_query(Err(error.clone()), "tiny").unwrap(),
sgl_kv_indexer::PrefixOutcome::Empty,
"{error} should degrade"
);
}
}
#[test]
fn rejected_indexer_query_still_fails_selection() {
assert!(matches!(
resolve_prefix_query(
Err(sgl_kv_indexer::PrefixIndexError::Rejected(
sgl_kv_indexer::RpcCode::InvalidArgument
)),
"tiny"
),
Err(ApiError::PolicySelectionFailed { .. })
));
}
/// `generate_room_id` MUST return values in `[0, i64::MAX]`. The
/// SGLang prefill stores `bootstrap_room` as `torch.int64`; a u64
/// with the top bit set would wrap negative on the engine side.
@@ -109,6 +109,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
kv_indexer_endpoint: None,
},
kv_index.tree(),
Arc::clone(&tokenizers),
@@ -1,44 +1,28 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
# SPDX-License-Identifier: Apache-2.0
"""Content-based cross-router routing test for cache-aware-zmq.
"""Content-based routing test for both cache-aware-zmq index backends.
Two routers + two SGLang workers + one shared model. Each router runs an
independent ``cache_aware_zmq`` policy whose ``KvEventIndex`` subscribes
to **both** workers' KV publishers.
Two SGLang workers publish KV events to two routers at once: one runs the
local ``KvEventIndex`` (SUB straight to the workers) and one runs against an
external KV Indexer fed by a ``kv-indexer-bridge`` per worker. Every
subscriber attaches before the single warmup, so one pair of disjoint
prefixes exercises both index backends without a second model load.
The test warms each worker with a DIFFERENT prefix DIRECTLY (bypassing
both routers), then sends those prefixes through each router and
asserts that routing follows the prefix CONTENT: ``PREFIX_X`` lands on
the worker holding X, ``PREFIX_Y`` lands on the worker holding Y, on
both routers.
# Why content-based, not convergence
An earlier version of this test asserted that both routers converged on
the *same dominant worker* after a one-prefix warmup. That property
sounds like it pins the ZMQ-fan-out contract, but it doesn't: when the
KV-event path is broken (subscribers never opened, e.g. a worker's
``/server_info`` lacks the ``kv_events`` block), ``cache_aware_zmq``
silently degrades to **min-load** — which, with sequential requests
holding ``active_load`` at zero, picks the same worker deterministically
on every call within a router. Both routers' min-load picks happened to
agree often enough (about half the time, modulo HashSet seed) to make
the convergence assertion pass even when no event ever flowed.
Content-based routing is uniquely sensitive to the KV-event path. Two
disjoint prefixes warmed on two different workers can only be routed
correctly if the router knows *which worker holds which content* — the
only mechanism that supplies that information is the ``BlockStored``
event stream. Under min-load fallback, both prefixes route to the same
default worker on each router, so the ``PREFIX_Y → worker_y`` assertion
fails regardless of which worker min-load defaults to.
Assert on content, not on convergence: a broken event path degrades
``cache_aware_zmq`` to content-blind min-load, which routes both prefixes to
one worker and so fails at least one assertion below.
"""
from __future__ import annotations
import os
import re
import socket
import subprocess
import time
from contextlib import contextmanager
from pathlib import Path
import httpx
import pytest
@@ -79,6 +63,69 @@ _REQ_TOTAL_RE = re.compile(
_LABEL_RE = re.compile(r'(\w+)="([^"]*)"')
def _open_port() -> int:
with socket.socket() as sock:
sock.bind(("127.0.0.1", 0))
return sock.getsockname()[1]
@contextmanager
def _run(binary: Path, env: dict[str, str], log_path: Path):
with log_path.open("w") as log:
process = subprocess.Popen(
[str(binary)],
env={**os.environ, **env},
stdout=log,
stderr=subprocess.STDOUT,
start_new_session=True,
)
try:
yield process
finally:
process.terminate()
try:
process.wait(timeout=10)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout=5)
def _wait_for_indexer(process: subprocess.Popen, port: int, log_path: Path) -> None:
deadline = time.time() + 10
while time.time() < deadline:
if process.poll() is not None:
raise RuntimeError(
f"KV Indexer exited during startup:\n{log_path.read_text()}"
)
try:
with socket.create_connection(("127.0.0.1", port), timeout=0.2):
return
except OSError:
time.sleep(0.1)
raise RuntimeError("timed out waiting for KV Indexer")
def _wait_for_bridge(process: subprocess.Popen, log_path: Path) -> None:
deadline = time.time() + 10
while time.time() < deadline:
output = log_path.read_text(errors="replace")
if "bridge session established" in output:
# ZMQ connect is asynchronous; let the subscription reach the PUB.
time.sleep(0.5)
return
if process.poll() is not None:
raise RuntimeError(f"KV Indexer Bridge exited during startup:\n{output}")
time.sleep(0.1)
raise RuntimeError(f"timed out waiting for KV Indexer Bridge:\n{output}")
def _dump_logs(logs: dict[str, Path]) -> None:
"""Print the tail of each Indexer/Bridge log so a routing failure is debuggable."""
for name, path in logs.items():
tail = path.read_text(errors="replace")[-4000:] if path.exists() else "<no log>"
print(f"\n----- {name} -----\n{tail}")
def _success_counts_by_worker(router_url: str) -> dict[str, int]:
"""Scrape ``/metrics`` and return ``{worker_url: success_count}``."""
r = httpx.get(f"{router_url}/metrics", timeout=5.0)
@@ -173,26 +220,21 @@ def _route_through(router_url: str, model_id: str, prompt: str) -> str:
@pytest.mark.real_gpu
@pytest.mark.slow
def test_two_routers_route_by_prefix_content(
router_binary, # noqa: ARG001 — fixture forces release-binary presence
def test_routers_route_by_prefix_content(
router_binary,
gpu_allocator,
tmp_path,
):
"""Each router must route by prefix CONTENT, agreeing across routers.
With each worker direct-warmed by a different disjoint prefix, the
only way a router can route ``PREFIX_X → worker_x`` AND
``PREFIX_Y → worker_y`` is by consulting a HashTree populated from
the BlockStored events the workers emit. Min-load fallback (the
failure mode when no SUB socket opened) is content-blind and would
route both prefixes to whichever worker its tiebreaker prefers.
"""
"""Both the local ZMQ index and the external Indexer must route by content."""
spec = get_model_spec("qwen3-0.6b")
gpus = gpu_allocator.acquire(2)
# Workers run with the model's REAL chat template (no override): the engine
# caches chat-templated tokens, and the router renders the same template
# (loaded from the model's tokenizer_config.json) before hashing. This
# exercises the production chat-template tokenization path, which aligns
# router query hashes with the engine's templated blocks.
indexer_port = _open_port()
indexer_endpoint = f"http://127.0.0.1:{indexer_port}"
indexer_binary = router_binary.parent / "kv-indexer-server"
bridge_binary = router_binary.parent / "kv-indexer-bridge"
logs = {
name: tmp_path / f"{name}.log" for name in ("indexer", "bridge-x", "bridge-y")
}
try:
with (
spawn_worker(
@@ -205,54 +247,79 @@ def test_two_routers_route_by_prefix_content(
gpu_ids=[gpus[1]],
enable_kv_events=True,
) as worker_y,
Gateway() as router_a,
Gateway() as router_b,
_run(
indexer_binary,
{"KV_INDEXER_LISTEN_ADDR": f"127.0.0.1:{indexer_port}"},
logs["indexer"],
) as indexer,
):
_wait_for_indexer(indexer, indexer_port, logs["indexer"])
worker_urls = [worker_x.url, worker_y.url]
for gw in (router_a, router_b):
gw.start_regular(
def bridge_env(worker, worker_id: str) -> dict[str, str]:
assert worker.kv_events_endpoint is not None
return {
"KV_INDEXER_WORKER_ID": worker_id,
"KV_INDEXER_WORKER_ADDRESS": worker.url,
"KV_INDEXER_ENDPOINT": indexer_endpoint,
"SGLANG_KV_EVENT_ENDPOINT": worker.kv_events_endpoint.replace(
"*", "127.0.0.1"
),
"SGLANG_KV_EVENT_TOPIC": "kv",
}
with (
_run(
bridge_binary, bridge_env(worker_x, "worker-x"), logs["bridge-x"]
) as bridge_x,
_run(
bridge_binary, bridge_env(worker_y, "worker-y"), logs["bridge-y"]
) as bridge_y,
Gateway() as local,
Gateway() as external,
):
local.start_regular(
model_id=spec["model"],
tokenizer_path=spec["model"],
worker_urls=worker_urls,
policy="cache_aware_zmq",
timeout=120.0,
)
# 1. Direct-warm each worker with its own prefix. Must happen
# AFTER both routers have started — ZMQ PUB/SUB doesn't
# replay messages emitted before SUB attaches, so any
# BlockStored event predating subscription is lost and
# the HashTree never sees it.
_direct_warm(worker_x.url, spec["model"], PREFIX_X)
_direct_warm(worker_y.url, spec["model"], PREFIX_Y)
# 2. Drain the SUB mpsc + pump-apply path. Sub-second under
# loopback ZMQ; 2 s leaves comfortable headroom.
time.sleep(2.0)
# 3. Content-routing assertion (×4): each prefix must land
# on the worker that holds it, on either router.
#
# The four assertions below are independently strong:
# min-load fallback routes both prefixes on a given
# router to a single default worker, so for ANY broken-
# fan-out scenario at least one of the four fails.
for router, label in ((router_a, "A"), (router_b, "B")):
landed = _route_through(router.base_url, spec["model"], PREFIX_X)
assert landed == worker_x.url, (
f"router {label}: PREFIX_X must route to worker_x "
f"({worker_x.url}); landed on {landed}. "
f"Likely cause: HashTree is empty — KV-event "
f"subscriber never opened, or BlockStored events "
f"never reached the pump."
)
landed = _route_through(router.base_url, spec["model"], PREFIX_Y)
assert landed == worker_y.url, (
f"router {label}: PREFIX_Y must route to worker_y "
f"({worker_y.url}); landed on {landed}. "
f"Likely cause: HashTree is empty — KV-event "
f"subscriber never opened, or BlockStored events "
f"never reached the pump."
external.start_regular(
model_id=spec["model"],
tokenizer_path=spec["model"],
worker_urls=worker_urls,
policy="cache_aware_zmq",
kv_indexer_endpoint=indexer_endpoint,
timeout=120.0,
)
_wait_for_bridge(bridge_x, logs["bridge-x"])
_wait_for_bridge(bridge_y, logs["bridge-y"])
_direct_warm(worker_x.url, spec["model"], PREFIX_X)
_direct_warm(worker_y.url, spec["model"], PREFIX_Y)
time.sleep(2.0)
try:
for router, label in (
(local, "local-index"),
(external, "external-indexer"),
):
landed = _route_through(
router.base_url, spec["model"], PREFIX_X
)
assert (
landed == worker_x.url
), f"router {label}: PREFIX_X must route to {worker_x.url}; landed on {landed}"
landed = _route_through(
router.base_url, spec["model"], PREFIX_Y
)
assert (
landed == worker_y.url
), f"router {label}: PREFIX_Y must route to {worker_y.url}; landed on {landed}"
except Exception:
_dump_logs(logs)
raise
finally:
gpu_allocator.release(gpus)
@@ -172,6 +172,7 @@ class Gateway:
tokenizer_path: str,
worker_urls: list[str],
policy: str = "round_robin",
kv_indexer_endpoint: str | None = None,
timeout: float = 60.0,
) -> None:
"""Start the router in regular (non-PD) mode.
@@ -186,6 +187,7 @@ class Gateway:
metadata are learned from ``/server_info``.
policy: Policy kind — ``round_robin``, ``random``, ``power_of_two``,
or ``cache_aware_zmq``.
kv_indexer_endpoint: Optional external KV Indexer gRPC endpoint.
timeout: How long to wait for ``/readyz`` before giving up.
"""
self._launch(
@@ -194,6 +196,7 @@ class Gateway:
tokenizer_path=tokenizer_path,
urls=list(worker_urls),
policy=policy,
kv_indexer_endpoint=kv_indexer_endpoint,
),
timeout=timeout,
)
@@ -293,6 +296,7 @@ class Gateway:
tokenizer_path: str,
urls: list[str],
policy: str,
kv_indexer_endpoint: str | None = None,
) -> list[str]:
resolved_tokenizer = _resolve_tokenizer_path(tokenizer_path)
@@ -317,6 +321,8 @@ class Gateway:
"--stale-request-timeout-secs",
str(self.stale_request_timeout_secs),
]
if kv_indexer_endpoint is not None:
args += ["--kv-indexer-endpoint", kv_indexer_endpoint]
# `--worker-urls` is multi-valued; keep it last so clap doesn't
# absorb a following flag as a URL.
args += ["--worker-urls", *urls]
@@ -10,7 +10,11 @@ FROM rust:1.90-bookworm AS builder
# `channel = "1.90"`.
ENV RUSTUP_TOOLCHAIN=1.90.0
# libssl-dev + pkg-config ship with rust:1.90-bookworm already; no apt-get needed.
# libssl-dev + pkg-config ship with rust:1.90-bookworm already; protoc does not,
# and the Indexer's build script needs it to compile the KV-indexer protos.
RUN apt-get update \
&& apt-get install -y --no-install-recommends protobuf-compiler \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /build
@@ -12,18 +12,10 @@
//! doesn't render tool schemas, so its ids would diverge from the engine).
//! * A request with multimodal (array) content → `input_ids` omitted (a text
//! tokenizer can't represent image content).
//!
//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches
//! the built-in V4 chat encoder — the engine-equivalent path — without a
//! template fixture.
use axum::body::Body;
use axum::http::{Request, StatusCode};
use serde_json::{json, Value};
use sgl_router::config::{
ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{BlockSizeOracle, HashTree};
@@ -36,33 +28,9 @@ use std::sync::Arc;
use std::time::Duration;
use tower::ServiceExt;
use crate::common::cache_aware_fixture::{config, MODEL};
use crate::common::mock_worker::MockWorker;
const MODEL: &str = "deepseek-v4-tiny";
fn config() -> Config {
Config {
server: ServerConfig {
host: "0".into(),
port: 0,
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
id: MODEL.into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::CacheAwareZmq,
circuit_breaker: None,
cache_aware: Some(CacheAwareConfig::default()),
sticky: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
}
fn build_ctx(url: String) -> Arc<AppContext> {
let cfg = config();
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
@@ -0,0 +1,39 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Shared router config for the cache-aware proxy tests.
//!
//! The model id contains `deepseek-v4` so the tokenizer registry auto-attaches the
//! built-in V4 chat encoder — the engine-equivalent path — with no template fixture.
use sgl_router::config::{
ActiveLoadConfig, CacheAwareConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
pub const MODEL: &str = "deepseek-v4-tiny";
/// A single-model `cache_aware_zmq` router. Discovery is a placeholder because
/// every caller installs its own `WorkerRegistry`.
pub fn config() -> Config {
Config {
server: ServerConfig {
host: "0".into(),
port: 0,
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
id: MODEL.into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::CacheAwareZmq,
circuit_breaker: None,
cache_aware: Some(CacheAwareConfig::default()),
sticky: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
}
@@ -3,5 +3,6 @@
//! Shared test harness re-exports.
pub mod cache_aware_fixture;
pub mod mock_worker;
pub mod streaming;
@@ -0,0 +1,135 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Full HTTP routing path backed by a real in-memory Indexer gRPC server.
use std::sync::Arc;
use std::time::Duration;
use axum::body::Body;
use axum::http::{Request, StatusCode};
use serde_json::json;
use sgl_kv_indexer::pb::kv_indexer_client::KvIndexerClient;
use sgl_kv_indexer::pb::{
ApplyExternalKvBatchRequest, ExternalKvAction, ExternalKvActionType, TierType,
};
use sgl_kv_indexer::{
server_builder, GrpcPrefixIndex, InMemoryKvIndexerBackend, KvIndexerService, PrefixIndexConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry;
use sgl_router::policies::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
use sgl_router::policies::request_tokens_for;
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::WorkerRegistry;
use tokio_stream::wrappers::TcpListenerStream;
use tower::ServiceExt;
use crate::common::cache_aware_fixture::{config, MODEL};
use crate::common::mock_worker::MockWorker;
#[tokio::test]
async fn external_indexer_routes_to_the_cached_worker() {
let cached = MockWorker::start(vec![]).await;
let uncached = MockWorker::start(vec![]).await;
let cfg = config();
let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
let body = json!({
"model": MODEL,
"messages": [{"role": "user", "content": "hello there friend"}],
});
let tokens = request_tokens_for(&tokenizers, &ModelId(MODEL.into()), &body)
.expect("test prompt tokenizes");
let hashes = compute_block_hashes(&tokens.ids, 1);
assert!(!hashes.is_empty());
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let endpoint = format!("http://{}", listener.local_addr().unwrap());
let server = tokio::spawn(async move {
server_builder()
.add_service(KvIndexerService::new(InMemoryKvIndexerBackend::new()).into_server())
.serve_with_incoming(TcpListenerStream::new(listener))
.await
.unwrap();
});
let mut indexer = KvIndexerClient::connect(endpoint.clone()).await.unwrap();
indexer
.apply_external_kv_batch(ApplyExternalKvBatchRequest {
worker_id: "cached-worker".into(),
seq: 1,
actions: vec![ExternalKvAction {
r#type: ExternalKvActionType::ActionReport as i32,
tier: TierType::TierHbm as i32,
hashes: hashes.clone(),
component_masks: Vec::new(),
block_sizes: Vec::new(),
}],
worker_address: cached.url.clone(),
cache_spec: None,
})
.await
.unwrap();
let registry = Arc::new(WorkerRegistry::default());
for url in [&cached.url, &uncached.url] {
registry
.add(WorkerSpec {
id: WorkerId(url.clone()),
url: url.clone(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId(MODEL.into())],
bootstrap_port: None,
})
.unwrap();
}
let oracle = BlockSizeOracle::new();
oracle.try_set(1).unwrap();
let policies = Arc::new(
build_registry(
&cfg,
Arc::new(HashTree::new()),
Arc::clone(&tokenizers),
Arc::clone(&oracle),
)
.unwrap(),
);
let mut ctx = AppContext::new(
cfg,
tokenizers,
Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()),
registry,
policies,
);
ctx.prefix_index = Some(Arc::new(
GrpcPrefixIndex::new(PrefixIndexConfig {
endpoint,
query_deadline: Duration::from_secs(1),
max_inflight: 4,
})
.unwrap(),
));
ctx.block_size_oracle = oracle;
let app = build_router(Arc::new(ctx));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert!(cached.captured.lock().unwrap().last_body.is_some());
assert!(uncached.captured.lock().unwrap().last_body.is_none());
server.abort();
}
@@ -12,6 +12,7 @@ mod common;
mod cache_aware_input_ids;
mod chat_routing;
mod external_indexer_routing;
mod failover;
mod graceful_shutdown;
mod header_forwarding;