veloxquant_mlx.routing package

Rate-aware session admission and shard placement for a shared KV pool (issue #278).

Adapts “CacheRoute: Planned Prefix-Affinity Routing for Large-Scale LLM Serving” (Cheng, Meta, Aug 2026, arXiv:2608.19677) — originally a routing plan across many servers in a distributed serving fleet — to a single VeloxQuant-MLX process: many concurrent session owners sharing one BlockPoolAllocator, optionally partitioned into shards.

Core pieces:

  • RateEstimator — tracks a smoothed per-owner request rate from live traffic (or skip this and supply your own {owner: rate} source, e.g. offline aggregate statistics).

  • CacheRoutePlanner — builds a RoutingTable once per control interval: sizes each hot owner’s shard count from its rate (Eq. 1 in the paper), admits owners into a warm set up to a slot budget in decreasing-rate order, and places admitted owners with longest-processing-time-first (LPT) list scheduling so shard load stays balanced (Algorithm 1).

  • RoutingTable — the resulting fixed plan: which owners are admitted, which shard(s) each should prefer, and the expected load per shard the plan was built from.

This module only plans admission and placement — it does not allocate or evict any block itself. Combine it with BlockPoolAllocator / PooledKVCache: use RoutingTable.is_admitted(owner) to decide which owners’ blocks to protect from eviction under pressure, and RoutingTable.preferred_shard to pick which pool partition a session’s blocks should live in.

Typical usage:

from veloxquant_mlx.routing import CacheRoutePlanner, RateEstimator

estimator = RateEstimator()
planner = CacheRoutePlanner(n_shards=4, qcap=50.0, warm_slots_per_shard=200.0)

# on every incoming request:
estimator.record(owner=session_id)

# once per control interval:
table = planner.plan(estimator.rates())
shard = table.preferred_shard(session_id, fallback=session_id % planner.n_shards)
if table.is_admitted(session_id):
    ...  # protect this session's blocks from eviction
class veloxquant_mlx.routing.SessionRate(owner, rate)[source]

Bases: object

A session’s estimated request rate, in requests per unit time.

Variables:

  • owner – Opaque session/owner id, matching BlockPoolAllocator owner ids.

  • rate – Estimated request rate (lambda_b in the paper). Any consistent unit works as long as qcap uses the same unit.

Parameters:
owner: int
rate: float
class veloxquant_mlx.routing.RoutingTable(shards, expected_load, n_shards, admitted_rate_total, unadmitted_rate_total)[source]

Bases: object

A fixed admission-and-placement plan for one control interval.

Variables:

  • shards – Maps each admitted owner to the tuple of shard ids it may use, ordered by preference (most-loaded-eligible-first is not guaranteed; callers needing a single preferred shard should use shards[owner][0]). An owner with more than one shard should split its load evenly across them, per the paper’s kb replication (Eq. 1) — an owner needs more than one shard only when its rate exceeds qcap on its own.

  • expected_load – Per-shard expected load after this plan’s placement, indexed by shard id. Sums admitted owners’ rate / kb over every shard each owner was placed on, plus each shard’s cold-tail share of unadmitted traffic. For diagnostics and shadow-replay comparisons only — not enforced at dispatch time.

  • n_shards – Total number of shards this plan was computed for.

  • admitted_rate_total – Sum of rate over admitted owners, for computing what fraction of traffic the warm set actually covers.

  • unadmitted_rate_total – Sum of rate over owners considered but not admitted (the plan’s cold-tail fallback traffic).

Parameters:
shards: dict[int, tuple[int, ...]]
expected_load: list[float]
n_shards: int
admitted_rate_total: float
unadmitted_rate_total: float
is_admitted(owner)[source]

Whether owner has a warm-set placement in this plan.

Return type:

bool

Parameters:

owner (int)

preferred_shard(owner, fallback=0)[source]

The single best shard for owner, or fallback if not admitted.

For an owner with kb > 1 shards, returns the currently least-loaded of its assigned shards (ties broken by lowest shard id), so repeated calls track load drift within the same interval even though the set of eligible shards stays fixed.

Return type:

int

Parameters:
imbalance()[source]

Ratio of the most-loaded shard’s expected load to the mean.

Matches the “imbalance” metric in the paper’s Table 4/5 (1.0x is perfectly balanced; higher is worse). Returns 1.0 for zero total load (nothing to balance).

Return type:

float

class veloxquant_mlx.routing.CacheRoutePlanner(n_shards, qcap, warm_slots_per_shard)[source]

Bases: object

Builds a periodic RoutingTable from measured session rates.

Implements Algorithm 1 from the paper: size each admitted session’s shard count from its rate relative to qcap (Eq. 1), admit sessions in decreasing-rate order up to the warm-slot budget, then place admitted sessions with longest-processing-time-first (LPT) list scheduling so no shard’s expected load is far from the mean.

This planner is stateless between calls — call plan() once per control interval with fresh rates, and hold onto the returned RoutingTable for the whole interval rather than recomputing per request. Recomputing every interval, versus reusing a stale table, is a real tradeoff the paper measures directly (Appendix E): a fresh plan changes most owners’ shard assignments (high “cache churn”) even when only a few owners’ rates actually moved, which costs a rewarming dip; a stale plan drifts and loses some hit rate/latency instead. This class always recomputes fresh — callers wanting hysteresis should diff two RoutingTable instances themselves before switching.

Parameters:

  • n_shards (int) – Number of pool shards to plan across (R in the paper). Use 1 if the pool isn’t partitioned — every admitted owner then simply gets warm-set priority with no placement decision to make.

  • qcap (float) – Load cap per shard, in the same rate unit as SessionRate.rate (Section 3.1). Calibrate from a single shard’s measured latency/load knee, as the paper recommends — this module does not calibrate it for you.

  • warm_slots_per_shard (float) – Warm-set capacity per shard, W in the paper. Total admission budget is n_shards * warm_slots_per_shard, counted in the same units as rate/qcap (equal-slot admission model; see the paper’s Section 3.1 and 6 for the heterogeneous-prefix-size limitation this inherits).

Example:

planner = CacheRoutePlanner(n_shards=4, qcap=50.0, warm_slots_per_shard=200.0)
table = planner.plan(estimator.rates())
shard = table.preferred_shard(owner, fallback=hash(owner) % planner.n_shards)
plan(rates)[source]

Compute a fresh routing table from the given session rates.

Parameters:

rates (list[SessionRate]) – Estimated per-session rates for this control interval. Order does not matter; sessions are re-sorted internally by decreasing rate as the paper’s Algorithm 1 requires.

Return type:

RoutingTable

Returns:

A RoutingTable valid for the caller’s next control interval.

class veloxquant_mlx.routing.RateEstimator(half_life=20.0, max_owners=None, _counts=<factory>, _last_seen=<factory>, _clock=0)[source]

Bases: object

Tracks a per-owner exponentially-weighted request rate estimate.

A minimal helper for callers who don’t already track per-session rates: call record() on every incoming request, then rates() to get the SessionRate list a CacheRoutePlanner needs. Not required — any source of {owner: rate} works, including offline aggregate statistics as in the paper’s telemetry-derived workload.

Memory is bounded by the number of distinct owners tracked (one float in _counts and one int in _last_seen each), not by how many requests have been recorded — an owner’s whole contribution to memory is released by forget(). That bound is only enforced if the caller calls forget() on session close; nothing here does it automatically. A caller that can’t guarantee that — e.g. a long-running server seeing many short-lived session IDs — should set max_owners so _counts self-bounds instead of growing without limit.

Parameters:

  • half_life (float) – Number of record() calls (from any owner) for an owner’s rate contribution to decay by half. Smaller values track bursts faster but are noisier; larger values are stable but slow to react to a session going cold or hot.

  • max_owners (int | None) – Maximum distinct owners to track at once. None (the default) preserves the original unbounded behavior. When set, recording a request from a not-yet-tracked owner that would push the owner count past this cap first evicts the single currently-lowest-rate owner (an approximation of “the coldest session,” not a full LRU) to make room. Does not add a branch to the already-tracked-owner path, so the hot case (an owner already in _counts) is unaffected.

  • _counts (dict[int, float])

  • _last_seen (dict[int, int])

  • _clock (int)

half_life: float = 20.0
max_owners: int | None = None
record(owner)[source]

Record one request arrival for owner.

The estimate is an exponentially weighted count of the owner’s arrivals on a global clock that ticks once per recorded request (from any owner), so it is proportional to the owner’s arrival rate and decays while the owner is silent. Cost is one dict lookup, one pow and two dict writes for an already-tracked owner; the eviction scan only runs for a new owner when max_owners is set.

Return type:

None

Parameters:

owner (int)

rate(owner)[source]

Current smoothed rate estimate for owner (0.0 if never seen).

Return type:

float

Parameters:

owner (int)

rates()[source]

Snapshot every tracked owner’s current rate as a SessionRate list.

Return type:

list[SessionRate]

forget(owner)[source]

Drop all tracked state for owner (e.g. on session close).

Return type:

None

Parameters:

owner (int)

Submodules