Skip to content

feat(router): query and feed the shared prefix-cache index from the policy layer - #2438

Open
slin1237 wants to merge 2 commits into
feat/radix-indexfrom
feat/radix-gateway-integration
Open

slin1237 wants to merge 2 commits into
feat/radix-indexfrom
feat/radix-gateway-integration

Conversation

@slin1237

@slin1237 slin1237 commented Sep 7, 2026

Copy link
Copy Markdown
Member

Note

3 of 3 in the shared prefix-cache index stack. Base: #2437 (the index service). Flag-gated: with --kv-indexer-url unset every routing path is byte-identical to today.

Description

Problem

The index service (#2437) is useless until the gateway queries it at routing time and feeds it after each request — and that must not be re-plumbed by hand in every router.

Solution

One seam on PolicyRegistry: resolve_remote_overlap (the routing-time query, 2 ms deadline, None on every skip case; gated on any cache_aware policy that could consume the scores, PD/EPD legs included) and publish_placement (the post-dispatch publish, once per request, never at select time — a shed or retry must not advertise a phantom holder). A router opts in with that pair; the registry owns the handle. With the shared index, follow-up cache hit stays flat at 0.95 from 1→8 gateways vs. 0.94→0.66 for per-gateway state.

The shared index replaces local state, it does not sit beside it. With --kv-indexer-url set the gateway subscribes to no worker KV events and, when a request gets nothing usable from the index (outage, timeout, nothing to hash), cache_aware picks on load alone through a new select_worker_remote_only hook (logged as branch = remote_none) instead of reading or populating its local prefix trees. A local tree fed only by remote misses is a partial, per-gateway view — exactly the state this index exists to remove. With the flag unset that hook is never reached.

routing path status
gRPC Regular built — byte-equivalent refactor of the previous inline prefetch
gRPC PD/EPD built — overlap steers the prefill leg only; publish targets the prefill worker (decode/encode never hold the prompt KV)
HTTP regular, tokens built — owned-token hoist out of the request lease, publish on dispatch success
HTTP regular, string prompts (Bytes keyspace) built — needs design sign-off: byte-prefix affinity is coarser than token-prefix; BYTE_BLOCK is a fixed constant
HTTP PD · HTTP streamed pass-through · worker-event forwarding not built (see Status)

Changes

  • policies/remote_index.rs (new), policies/registry.rs — the seam; policies/mod.rs.
  • routers/grpc/common/stages/worker_selection.rs, routers/grpc/pipeline.rs, routers/grpc/context.rs — gRPC opt-in (Regular + PD/EPD).
  • routers/http/router.rs — HTTP opt-in (token tree, string-mode fallback).
  • workflow/steps/** — worker add/drop lifecycle signals to the index; service_discovery.rs, app_context.rs, config/*, main.rs--kv-indexer-url / --kv-indexer-block-size wiring (default shared with the bridge — a mismatch would silently split the fleet into two keyspaces).
  • observability/metrics.rs — query/publish counters and latency.
  • crates/kv_index — byte-chunk content hasher for string mode (token hasher untouched).

Suggested reading order: policies/remote_index.rspolicies/registry.rs (the two calls) → one router (routers/grpc/common/stages/worker_selection.rs) → workflow/steps.

Test Plan

  • Flag off: byte-identical (registry/selection/pipeline suites unchanged: policies 260, gRPC 295, HTTP 33, workflow 87).
  • Flag on: same skip conditions, query args, published chains, and x-smg-index-* echo headers as the pre-refactor inline code; PD/EPD is a pure superset. Placement carries a three-state RemoteLookup (not attempted / missed / hit): only the two call sites that query the index build one, so HTTP PD, transcription, the streamed pass-through and retry re-selection place exactly as before the index existed. New locks: the remote-only pick leaves the local token tree empty, remote_miss is the labeled path for an answer with no local holder, --kv-indexer-block-size 0 is rejected at parse time and the two flags round-trip into RouterConfig.
  • End to end via the harness (feat(sim): no-GPU simulation harness for cache-aware routing and the shared index #2439): agentic workload sprayed across gateways — cache hit 0.95 flat vs. 0.94→0.66; routing precision 0.97 flat vs. 0.91→0.12.

Validation (end-to-end campaign, 2026-09-08)

Harness (#2439): 120 mock workers with prefix caching and KV events, 305 sessions/s (~405 req/s), 150 s windows, sessions sprayed across gateways; one seed per leg; 68,392 requests and 0 errors in every leg. Follow-up (turn-2) cached tokens, sum/sum:

regime 1 gw 2 gw 4 gw 8 gw
shared index, event feed (--kv-indexer-url + bridge) 0.815 0.882 0.893 0.875
shared index, placement feed (eventless engine) 0.826 0.884 0.896 0.890
per-gateway event trees (each gateway consumes all workers' events) 0.824 0.889 0.870 0.707
no sharing (per-gateway approximate tree) 0.598 0.473 0.328 0.249
  • Follow-ups landing on the worker that served turn 1: placement feed 0.90 / 0.97 / 0.98 / 0.99; per-gateway trees 0.85 / 0.93 / 0.91 / 0.74; no sharing 0.52 → 0.11.
  • Fleet load imbalance (CoV) at 8 gateways: shared index 0.063, per-gateway trees 0.472. Peak gateway RSS at 8 gateways: 605 MiB (index) vs 787 (event trees) vs 1056 (no sharing).
  • Index lookups: remote_hit ≥ 98.6%, degraded (2 ms deadline missed or disconnected) ≤ 1.4%; event feed predicted-vs-actual cached tokens p95 error 0; placement feed p95 4.1–4.9k tokens (the output tail the worker caches after the routing-time chain).
  • 240 workers, 611 sessions/s (~823 req/s), 8 gateways, 180 s, over h2 (with fix(http): do not relay empty upstream chunks as empty h2 DATA frames #2488 applied): shared index 0.873 / 0.878, per-gateway trees 0.727, mesh TreeSync 0.633; CoV 0.045 vs 0.43 / 0.46; gateway RSS 560 vs 954 / 1022 MiB; 163,958 requests, 0 errors per leg.
  • Half load (150 sessions/s): 0.939 follow-up cached, 0.98 same-worker, prediction p95 error 0 on the placement feed. Points at ~2× engine capacity (60 workers @305, 120 @610) saturate the engines (e2e p50 16–20 s) and the 2 ms index deadline fails open to load-only routing, as designed.
  • Flag off: 737 tests across the policies, gRPC, HTTP, placement, workflow and indexer suites unchanged.

Side-by-side evaluation (2026-09-09, 240 workers, 8 gateways, 2 seeds, mean ± 95% CI)

Same workload for every column; only where prefix knowledge lives differs. Sessions are sprayed across gateways at random (no sticky load balancer in front), so the incumbent's per-gateway sticky routing key only helps when a follow-up happens to land on the gateway that served turn 1 (1 in 8). With a sticky LB the incumbent behaves like the single-gateway column below (0.95).

305 sessions/s (~410 req/s) production (hash + sticky) per-gateway kv_index event trees shared index, event feed shared index, placement feed
follow-up cached tokens 0.266 0.922 0.931 0.932
follow-up on its turn-1 worker 0.13 0.964 0.969 0.971
fleet load CoV 0.047 0.823 0.077 0.077
max gateway RSS (MiB) 196 693 409 391
max gateway CPU (mean %) 9 27 14 15
611 sessions/s (~820 req/s) production event trees index, events index, placements
follow-up cached tokens 0.248 0.715 0.874 0.875
follow-up on its turn-1 worker 0.13 0.747 0.912 0.968
fleet load CoV 0.034 0.462 0.044 0.052
max gateway RSS (MiB) 206 924 547 551
max gateway CPU (mean %) 16 55 28 28
900 sessions/s (~1210 req/s, ~1.5× the calibrated engine load) production event trees index, events index, placements
follow-up cached tokens 0.221 0.299 0.633 0.636
follow-up on its turn-1 worker 0.13 0.169 0.650 0.965
fleet load CoV 0.026 0.114 0.031 0.040
max gateway RSS (MiB) 216 1190 727 697
max gateway CPU (mean %) 24 88 44 42

Reading it: the per-gateway event trees (every gateway ingesting every worker's events) fall apart as load rises — CPU 88%, imbalance CoV 0.46–0.82, hit rate 0.30 at 1210 req/s — while the shared index holds routing precision (placement feed on its turn-1 worker 0.97 at every rate) and the fleet balanced at CoV ≤ 0.08, with 40–60% of the event trees' gateway memory and half the CPU. At 900 sessions/s every regime loses cache to engine KV eviction (the mock's 1.2M-token KV at 1.5× the calibrated load); the event feed's lower same-worker share there is it correctly routing away from blocks the engine evicted (prediction exact share 0.98), while the placement feed keeps routing to the original worker (exact share 0.86, p95 error 11k tokens). Zero request errors in all 24 eight-gateway legs (82k–242k requests each).

Index path latency (8 gateways): service-side query engine time p50 20 µs, p99 0.08–0.23 ms; apply p50 5 µs (events) / 11–21 µs (placements), p99 up to 0.49 ms at 1210 req/s. Gateway-side lookup, as the 2 ms deadline sees it: p50 0.52 ms, p90 0.94 ms, p99 3.7–4.1 ms — the p99 tail is what the 2.1–2.6% degraded (fail-open) share is; with one gateway p99 is 1.0 ms and degraded 0.04%. The 0.5 ms floor over a 20 µs engine query is client/transport overhead (one subscribe stream per gateway) and is the next performance target. Index replica: 71–200 MiB RSS, 6–30% of a core, 0.9996/0.977 hit share (1/8 gateways).

Input/output shapes (8 gateways, 240 workers, 2 seeds; session rate scaled per shape to hold ~35 concurrent requests per worker; zero request errors except 28 ± 32 on the event trees' long-short legs):

follow-up cached tokens production event trees index, events index, placements
short prompts (1–4k), short outputs (200–800), 1320 req/s 0.737 0.776 0.775 0.775
long prompts (12–24k), short outputs, 1320 req/s 0.243 0.183 0.932 0.934
long prompts, long outputs (2–5k), 536 req/s 0.238 0.742 0.935 0.938

Short prompts are mostly the shared 2k system prefix every worker holds, so every regime scores it and the index has nothing to add (same-worker 0.38–0.40 — overlap ties, load decides). Long prompts are where the index matters: at 1320 req/s of 20k-token prompts the per-gateway event trees collapse (0.18, e2e p50 20 s vs 1.8 s, CoV 0.12, RSS 1.2 GiB peak) while the shared index holds 0.93 with same-worker 0.94–0.97, CoV 0.02–0.06, and 375–466 MiB of gateway memory. Placement apply p99 grows with chain length (0.65 ms on long-short); lookup p99 stays 3.7–3.9 ms, degraded share 2.2–2.7%.

30-minute soaks (8 gateways, 120 workers, 305 sessions/s, 822,748 requests each, zero errors): follow-up cache flat at 0.87–0.88 for the whole half hour on both feeds (minute 1: 0.90 / 0.88; last full minute: 0.89 / 0.90), index timeouts 0.7–1.0%, routing precision 0.98 (placements) / 0.91 (events, routing away from evicted blocks), gateway RSS flat at ~350–360 MiB after warm-up. Index replica: event feed flat at 61 MiB and ~490k blocks; placement feed held ~850k blocks but its RSS grew from 120 to 460 MiB over the 30 minutes at a constant block count — a memory growth in the placement path under sustained capacity cuts, under investigation (see #2437).

Single gateway (all regimes): the one gateway process saturates at ~240 req/s on this host, so 611 and 900 sessions/s are gateway-bound (queue 503s, e2e 40–70 s) and all four regimes read alike; at 305 they all sit at 0.93–0.95 because one gateway sees every turn.

Chaos (2026-09-09/10, 8 gateways, 120 mock workers, 305 sessions/s, ~68k requests per leg, 2 seeds; convergence reruns 2026-09-14 on the anti-entropy + recency-eviction build)

leg follow-up cached follow-up on its turn-1 worker request errors notes
20 workers removed then 20 added mid-run 0.836 0.964 0 replicas converged (0 holders differing)
one gateway cold-restarted at 60 s, shared index 0.895 0.985 748 (in-flight at the kill) the restarted gateway's first 300 follow-ups: 0.88, the others 0.89
one gateway cold-restarted at 60 s, per-gateway event trees 0.718 0.752 764 the restarted gateway's first 300 follow-ups: 0.68 vs 0.73 elsewhere
rolling restart of both index replicas 0.789 0.862 0 degraded lookups 10.8% during the restarts
one replica removed permanently at 60 s 0.890 0.986 0
every gateway cut from its replica for 45 s 0.615 0.625 0 fail-open, 31% degraded; last minute back at 0.87
half the gateways cut from their replica for 45 s 0.681 0.719 0
inter-replica partition 45 s (event feed) 0.838 0.879 0 routing unaffected

The cold-gateway pair is the thesis in one row: a gateway that restarts with the shared index is at the fleet's hit rate from its first requests; one that has to rebuild its own event trees is 7 points behind for the whole run. Every leg ran with zero request errors except the two that kill a gateway with requests in flight.

Replica consistency. The first chaos pass found the inter-replica partition left the two replicas permanently diverged (120/120 holders, one side still holding blocks the worker had evicted) and the rolling restart 78 and 27 holders apart. #2437 gained anti-entropy (peers exchange per-holder digests every 15 s and pull holders where a peer is provably ahead). Rerun on that build: the partition leg converges exactly (0/120, identical 555,196 blocks on both replicas, 292 holders pulled by anti-entropy); the placement-feed legs still show 33 and 24 holders differing, every one of them a worker above its capacity on both replicas, i.e. the two replicas ran their local capacity cut at different moments. The harness now reports that class separately from a placement holder differing under capacity (a lost update); none occurred.

Memory over the 30-minute placement soak, after #2500's recency-ordered cut: index replica RSS 43 → 108 MiB by minute 10, then flat at 102 MiB through minute 30 (before: 486 MiB at minute 30 and climbing). Gateway RSS was flat in both.

Caveat. The 2026-09-14 reruns ran on a laptop shared with other active agent sessions; a host-wide sampler caught a Go build at 460% CPU, rustfmt and clippy-driver at 95–99%, and a Codex process at 30–90% mid-run. Their follow-up cache ratios dip, to the minute, wherever such a burst pushed index lookups past the 2 ms deadline on all eight gateways at once (the service's own query time stayed under 0.1 ms p99 throughout), so the hit-rate figures in this PR remain those of the quiet-host campaign above; the rerun contributes the convergence and memory verdicts, which are end-state measurements.

Status

  • HTTP PD: its prefill workers are a separate pool, so steering only pays off once it also publishes — an all-or-nothing change across its streaming dispatch internals; own PR.
  • HTTP streamed pass-through selects under UNKNOWN_MODEL_ID, so joining the index would split the keyspace against the buffered path.
  • Worker-event forwarding (bridge in the gateway) is unwired.

@coderabbitai

coderabbitai Bot commented Sep 7, 2026

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: f7dbfb2d-10c2-4221-a436-5ad25ee05cab

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Comment @coderabbitai help to get the list of available commands.

Comment thread model_gateway/src/policies/registry.rs Outdated
Comment thread model_gateway/src/policies/cache_aware.rs Outdated
Comment thread model_gateway/src/routers/http/router.rs Outdated
Comment thread model_gateway/src/policies/registry.rs Outdated
Comment thread model_gateway/src/routers/grpc/pipeline.rs Outdated
Comment thread model_gateway/src/workflow/steps/shared/update_policies.rs
Comment thread model_gateway/src/main.rs
@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch from 95bdcd3 to c045476 Compare September 7, 2026 11:30
@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch 3 times, most recently from e69039e to 667cbf9 Compare September 7, 2026 22:52
Comment thread model_gateway/src/routers/common/placement.rs
Comment thread model_gateway/src/routers/http/router.rs Outdated
@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch from 667cbf9 to 051c22e Compare September 8, 2026 21:39
///
/// This trait provides a unified interface for implementing routing algorithms
/// that can work with both regular single-worker selection and PD dual-worker selection.
pub(crate) mod remote_index;

@claude claude Bot Sep 14, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: The mod declaration was inserted inside LoadBalancingPolicy's doc comment, so it steals it.

Lines 48–51 (/// Core trait for load balancing policies … regular single-worker selection and PD dual-worker selection.) are an outer doc comment; whatever item follows takes them. That is now pub(crate) mod remote_index;, so rustdoc renders "Core trait for load balancing policies" as the module's summary (on top of the module's own //! header), and the trait this PR extends with two new methods ends up undocumented.

Fix is to move this line up with the other mod declarations above the pub use block, leaving lines 48–51 attached to pub trait LoadBalancingPolicy as before.

.filter_map(|&idx| {
let url = workers[idx].url();
remote
.scores

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: Quadratic string-compare scan on the synchronous routing path, at the exact fleet size this PR benchmarks.

For each healthy worker this walks the whole scores vec comparing URLs, so it is O(workers × holders) String/&str comparisons. QueryOutcome::Scores is not top-k capped anywhere in engine.rs — it returns every holder that matched — so on the validated 240-worker fleet a well-shared prefix gives ~240 × 240 ≈ 58k comparisons per request, inside the policy call that holds up selection, at 1210 req/s.

The local paths don't pay this: they get a per-worker score out of the indexer rather than scanning a list. One pass to build a lookup keyed on &str and then one probe per worker keeps it O(W + S):

        let by_holder: std::collections::HashMap<&str, u32> = remote
            .scores
            .iter()
            .filter(|(_, matched)| *matched > 0)
            .map(|(holder, matched)| (holder.as_str(), *matched))
            .collect();
        let mut candidates: Vec<OverlapCandidate> = healthy_indices
            .iter()
            .filter_map(|&idx| {
                by_holder
                    .get(workers[idx].url())
                    .map(|matched| OverlapCandidate {
                        idx,
                        effective_score: f64::from(*matched),
                    })
            })
            .collect();

(A duplicate holder in scores would now resolve to the last rather than the first entry; if the index can emit duplicates, and_modify(|v| *v = (*v).max(matched)) covers it.)

@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch from f69af37 to 794febe Compare September 14, 2026 10:27
@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch from 794febe to 93e2643 Compare September 14, 2026 11:54
@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch 2 times, most recently from d96499d to 39e5da0 Compare September 15, 2026 14:33
}
})
};
let mut p = limit - limit % alignment;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Important: The candidate descent is linear in the request's length for every answering worker, on the synchronous selection path.

accepts runs once per aligned candidate from limit down to the first acceptor, and alignment is the LCM of the lanes' blocks in keyspace units — which is 1 whenever a lane's engine block equals the keyspace block (the common case: --kv-indexer-block-size is meant to match the bridge/engine page). So the step is one unit, and limit is the request's whole prefix: a 20k-token prompt at a 16-token block gives limit = 1250.

The workers that scan farthest are the typical ones, not a corner case: a holder of just the shared 2k system prefix probes ~1125 candidates before accepting ~125. QueryOutcome::Scores is not top-k capped in engine.rs, so on the 240-worker fleet this PR validates that is ~240 × ~1100 accepts calls, each re-walking every lane's intervals — covered_from_zero is recomputed per candidate inside the closure, and a rejecting Checkpoint lane does a full .any() over its runs. Order 10^6 interval visits per request, inside the resolve_remote_overlap the router awaits before selection, at the 1210 req/s this stack benchmarks.

A cheap bound exists because Full acceptance is monotone in p (p <= covered_from_zero): start the descent at min(limit, min over Full lanes of covered_from_zero(..)) truncated to alignment, which collapses the common case to a handful of probes. Worth hoisting the per-candidate covered_from_zero and block_units recomputation out of the closure at the same time (block_units is already computed once into lookbacks), and for checkpoint lanes the acceptable positions are enumerable directly from the intervals rather than probed.

Only lane-published (hybrid) models reach this path — a plain holder takes the lanes.len() == 1 fast path above — but that is exactly the configuration this commit adds.

Comment thread model_gateway/src/policies/reuse.rs Outdated
let mut by_worker: Vec<(String, Vec<(Option<LaneMeta>, Vec<(u32, u32)>)>, bool)> = Vec::new();
for answer in answers {
let (worker, lane) = worker_of(&answer.holder);
let entry = match by_worker.iter_mut().find(|(w, _, _)| w == worker) {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: Grouping the answers by worker with a linear find is O(answers × workers) string comparisons per request.

find_matches sorts its answers by matched_blocks descending, then holder name (crates/radix_index/src/engine.rs), so a worker's lanes are not adjacent in the vec — every answer walks the whole accumulated by_worker comparing String == &str. At the fleet size this PR validates (240 workers × ~3 lanes ≈ 720 answers, up to 240 entries each) that is ~10^5 string compares per request on the same synchronous routing path as reusable_units.

A HashMap<&str, usize> from the borrowed worker slice to its index in by_worker makes it one pass and keeps the single to_string() per worker (the answers outlive the fold, so borrowing answer.holder needs the loop restructured to take the intervals out by index, or just key the map on String — still one hash per answer instead of a scan).

Metrics::record_remote_index_query(label, started.elapsed());
let scores = match outcome {
radix_index::client::QueryOutcome::Scores(answers) => {
super::reuse::aggregate(answers, tokens.len(), block)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: remote_hit is now recorded for answers that aggregate to nothing, so an unreadable lane fleet is indistinguishable from a healthy one.

outcome_label is computed from the raw QueryOutcome two lines up, but reuse::aggregate can return an empty vec for a perfectly non-empty answer: any worker with one lane whose lane_meta the consumer cannot read is dropped whole (reuse.rs:190-193), and LaneMeta::parse accepts only full_attention | mla_attention | sliding_window | mamba. chunked_local — which reuse.rs's own test pins as unscorable, and which is what an interleaved-local-attention model reports — drops every worker in the fleet. Empty scores then becomes RemoteLookup::Missedselect_worker_remote_only, i.e. load-only routing with the local trees deliberately left empty.

The visible result is a deployment routing purely on load while smg_remote_index_query_total{outcome="remote_hit"} reads ~100% — the exact metric the PR's validation table uses to argue the index is answering (remote_hit ≥ 98.6%). Worth either re-labeling when the aggregate comes back empty (a distinct remote_unscorable, so the fail-open is countable) or a rate-limited warn! the first time a lane_meta fails to parse. Same at line 398 for the bytes path.

@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch from 39e5da0 to de89c8f Compare September 15, 2026 14:49
@slin1237
slin1237 requested a review from njhill as a code owner September 15, 2026 14:50
@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch from de89c8f to d194fb7 Compare September 15, 2026 14:50
match lane.map(|_| LaneMeta::parse(&answer.lane_meta)) {
// A lane whose description the consumer cannot read makes
// the worker unscorable: better no claim than a wrong one.
Some(None) => entry.2 = true,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Important: An index replica that restarts makes every lane-published worker permanently unscorable, because the bridge announces lane metadata exactly once per worker connection — not once per replica.

The chain:

  1. convert_batch_lanes attaches proto::Added { metadata: info.metadata() } only when lanes.announced.get(&group) != Some(&info) (crates/radix_index/src/bridge.rs:312-318), and LaneBook is created once per worker_loop (bridge.rs:499) — its lifetime is the worker connection, so it is not reset when the bridge's index client reconnects or fails over.
  2. On the replica side, a Publish for an unknown holder creates it inferred, with metadata: Vec::new() (engine.rs:622, engine.rs:1102). So after a replica restart the lane holders come back from the ongoing event stream with empty lane_meta.
  3. Here, lane.map(|_| LaneMeta::parse(&answer.lane_meta)) on a lane-suffixed holder with empty metadata yields Some(None)entry.2 = true → the worker is dropped whole at line 205, including the lanes whose metadata did survive and its plain placement holder.

Every lane worker dropping out makes scores empty, which resolve_remote_overlap turns into RemoteLookup::Missedselect_worker_remote_only → load-only routing, with the local trees deliberately left empty. Unlike the unknown-kind case, this does not self-heal: the bridge has no trigger to re-announce until the worker↔bridge connection itself drops, so a replica restart costs prefix affinity for the remaining life of every worker loop. AddedControl's own doc says empty metadata means "no claim: a bare lifecycle re-announce never clears standing metadata" — i.e. empty is explicitly not a negative claim, but this treats it as one.

Distinguishing "lane has no metadata yet" from "lane has metadata I can't parse" would fix the restart case without weakening the unknown-kind guard — an empty lane_meta could keep the worker's readable lanes (or fall back to the contiguous depth of its plain holder) rather than voiding it. If the fail-closed read is deliberate, the replica needs to either persist/anti-entropy the metadata before serving, or the bridge needs to re-announce on index reconnect.

…olicy layer

Wires the gateway to smg-radix-index behind --kv-indexer-url; with the
flag unset every routing path is byte-identical to before.

One seam on PolicyRegistry so no router re-plumbs the index by hand:
resolve_remote_overlap runs the routing-time query (2 ms deadline, None
on every skip case: flag off, non-cache_aware policy, no tokens, sticky
override) and publish_placement records the served worker after
dispatch, never at select time, so a shed or retry cannot advertise a
phantom holder. The registry owns the client handle; worker add/drop
flows through the workflow steps as lifecycle signals.

Opted in: gRPC Regular (byte-equivalent refactor of the previous inline
prefetch), gRPC PD/EPD (overlap steers the prefill leg only; decode and
encode never hold the prompt KV), HTTP regular in token mode, and HTTP
string prompts via a byte-chunk keyspace (kv_index gains the byte-chunk
content hasher; the token hasher is untouched). HTTP PD, the streamed
pass-through (selects under UNKNOWN_MODEL_ID) and worker-event
forwarding stay on plain selection and are called out in the PR.

The remote overlap rides alongside the request's cache namespace from

Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The shared index now answers per holder with the covered runs of the
request and the publisher's opaque lane metadata, and a worker with
several independently evicted caches (full attention, sliding window,
recurrent state) is published as one holder per lane. The gateway is
where those answers become one number per worker: the largest prefix
position every lane accepts at that position, which is what vLLM's
hybrid coordinator computes natively.

policies/reuse.rs holds the rules: full and MLA attention need coverage
from 0, a sliding window needs cdiv(W - 1, block) contiguous engine
blocks (at least one) before the candidate, a recurrent lane needs a
saved state exactly there (any block end inside a run), and candidates
walk down from the request bound aligned to the least common multiple
of the lanes' engine blocks. A holder without lane metadata keeps its
depth from 0, so single-cache workers score exactly as before. A lane
the gateway cannot read disqualifies its worker rather than guessing.

Checked against the native-lookup fixture recorded in #2497 (vLLM
find_longest_cache_hit expectations, 160 cases): all 125 whole-block
cases agree; the 35 partial-hash-hit cases are not representable in the
event feed, whose partial blocks the converter drops, and are skipped
with that reason.

Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
@slin1237
slin1237 force-pushed the feat/radix-gateway-integration branch from 18430a7 to c849649 Compare September 15, 2026 15:44
if unreadable {
return None;
}
let has_lanes = lanes.iter().any(|(m, _)| m.is_some());

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Important: A lane with no coverage of the query is absent from the answers, and this treats absence as "no constraint" instead of "accepts only 0" — so the joint score is overstated for exactly the hybrid models the rules exist for.

Engine::find_matches builds its answers from tree.coverage(chain, ..) and only emits holders that produced at least one run on the query path (crates/radix_index/src/engine.rs:1039-1067) — a lane holder that is registered but holds nothing on this chain simply never appears. aggregate therefore sees only the lanes that matched, and reusable_units computes the fixed point over that subset.

Failure scenario: a hybrid worker w1 with a full-attention lane w1#0 (blocks [0,8) cached for this prompt) and a Mamba lane w1#1 with no saved state on this path. The index answers one holder. has_lanes is true, views has a single Full lane, and the worker scores 8 units. The engine's own find_longest_cache_hit would return 0 — the recurrent group has nothing to resume from, so the whole prefix is recomputed. The router then places the request on w1 believing 8 blocks are reusable, i.e. the precision this commit is meant to add is inverted for the one worker shape it targets.

The module cannot distinguish "lane absent because it holds nothing" from "worker has only these lanes" with the data it currently gets. Options: carry the worker's lane count (or the full lane roster) in the announce so aggregate can detect a short answer and score 0; or have the index answer every live lane of a worker that matched at all, with an empty interval list for the ones that did not.

(Latent today — nothing in-tree publishes lane metadata any more — but it is the core rule of this commit, so worth settling before a producer lands.)

match key {
"kind" => {
kind = Some(match value {
"full_attention" | "mla_attention" => LaneKind::Full,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: This vocabulary has no producer in the tree, and the only in-tree examples of the encoding spell the kinds differently — so the contract is unpinned in a way that fails silently and fleet-wide.

The doc comment above (lines 29–31) says these bytes are what "the bridge writes as kind=..;window=..;block=..", but crates/radix_index/src/bridge.rs contains no metadata/lane code at all (this push removed convert_batch_lanes), and client.rs:360 publishes metadata: Vec::new(). The only places the encoding appears outside this file are the engine's own tests, which write kind=full;block=1, kind=swa;block=1;window=3, kind=mamba;block=4 (engine.rs:1951-1953, 2002-2004) — and the fixture loaded by native_lookup_fixture_agrees_with_the_rules uses full / swa / mamba too (line 403-407 maps them). Of those, only mamba parses here; full and swa hit the _ => return None arm.

The failure mode is the worst possible one: an unparseable lane sets entry.2 = true in aggregate, which drops the entire worker from the scores — so a fleet-wide spelling drift yields remote_hit with an empty score vec and load-only routing, with nothing in the metrics to say why.

Worth pinning the vocabulary in one place both sides compile against (e.g. a LaneKind/lane_meta encoder+decoder pair in radix_index, with the engine tests and the fixture mapper using it), rather than three independent string literals.

// A checkpoint lane's runs are whole engine blocks, so a
// state exists at every block end inside a run, not only
// at the run's end (adjacent blocks merge into one run).
LaneKind::Checkpoint => lane.intervals.iter().any(|&(start, end)| {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: The checkpoint rule assumes the publisher marks a saved state as a whole engine block, and the only in-tree example of a mamba lane uses the other encoding.

(p - start).is_multiple_of(block_units) anchors the block grid at each run's start, which is only correct if runs are whole engine blocks (as the comment states, and as the fixture mapper constructs them at line 421-424: (end_units - block_units, end_units)).

engine.rs:2004 publishes the same concept as single-unit marks at the state boundaries: ("w1#2", 0, &[(3, 4), (7, 8)][..], "kind=mamba;block=4") — states at positions 4 and 8, with block=4 tokens over a block_size=1 keyspace, i.e. block_units == 4. Feed those intervals here and every candidate is rejected: at p = 4, start = 3, p - start = 1, not a multiple of 4 → the lane accepts nothing but 0, so the worker scores 0 despite holding two usable states.

Since there is no producer yet, either encoding is still on the table — but the two in-tree examples already disagree, so it's worth stating the granularity as part of the lane_meta contract (see the note on the kind vocabulary above) instead of leaving it to a doc comment on the consumer. An alternative that is robust to both: require p to be a block end relative to the keyspace origin (p.is_multiple_of(block_units)) plus p inside a run, which the fixture encoding also satisfies whenever runs are block-aligned.

/// async pipeline stage before the (synchronous) policy call.
#[derive(Debug, Clone, Default)]
pub struct RemoteOverlap {
/// Per-holder (worker url, matched prefix blocks), descending.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: This doc is now wrong in the one word that matters — these are per-worker reuse scores, not per-holder matched blocks.

As of this push both resolve_remote_overlap and resolve_remote_overlap_bytes run the raw answers through reuse::aggregate (registry.rs:337-339, 395-397), which folds worker#lane holders into one entry per worker and replaces matched_blocks with the jointly-reusable unit count from the reuse rules (bounded below the request's last token, zero scores dropped). "Holder" is now a specific term in this stack — radix_index::engine::LANE_SEPARATOR splits it into worker + lane — so "per-holder" reads as the opposite of what the field contains.

IndexPrediction::scores in remote_index.rs:72 has the same stale wording ("Per-holder (url, matched blocks) as answered"), and it is the one that feeds predicted_tokens_for → the x-smg-index-* echo headers the harness uses to separate index error from policy spill, so the distinction is load-bearing for reading those numbers.

Suggestion for this line:

Suggested change
/// Per-holder (worker url, matched prefix blocks), descending.
/// Per-worker (worker url, jointly reusable prefix blocks) as
/// resolved by `reuse::aggregate`, descending, zero scores dropped.

.await;
let label = outcome_label(&outcome);
Metrics::record_remote_index_query(label, started.elapsed());
let scores = match outcome {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: String mode passes a byte unit into rules whose lane metadata is in tokens, so the two silently collapse instead of disagreeing loudly.

aggregate(answers, text.len(), BYTE_BLOCK) makes unit_tokens = 256 bytes, but LaneMeta::block_tokens is the engine block in tokens (reuse.rs:37-38). reusable_units then computes block_units = meta.block_tokens as usize / unit_tokens, which is 0 for every realistic engine block (16–512 tokens ÷ 256 bytes), and returns None at reuse.rs:125-127 — i.e. any lane-published holder in the Bytes keyspace is unscorable and aggregate's filter_map drops that worker from the scores entirely.

Latent today: the Bytes keyspace is fed only by string-mode publish_placement, which writes plain worker holders with no lane_meta, so the lane branch is never entered there. But the two paths call the same function with incompatible units, and the failure is a silent drop rather than an error — the same unit hazard the block_size comment right below this line already had to work around for the decay math. Worth either rejecting lanes explicitly in the byte keyspace, or carrying the unit kind (tokens vs bytes) into reusable_units so a future lane-publishing bridge can't be silently unscorable.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

dependencies Dependency updates grpc gRPC client and router changes kv-index KV index crate changes model-gateway Model gateway crate changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant