[router] cache_aware_zmq: load-aware selection from engine LoadSnapshot

[Kangyan-Zhou]

Motivation

The cache_aware_zmq policy prices workers by a router-side in-flight counter. That counter reflects what this router dispatched, not what the engine has queued: it cannot see traffic from other router replicas or from direct clients, and for streaming responses it stays held for the full response rather than for the time the request occupies the scheduler.

With #28599 the engine publishes its LoadSnapshot on a dedicated socket. This subscribes to it and routes on real queue depth, keeping the in-flight counter as the fallback for any worker whose reported picture is missing or incomplete.

Depends on #28599 for the publisher and the /server_info advertisement.

Modifications

Discovery. Workers advertise load_endpoint_port_base in /server_info's kv_events block. It is Option<u16>: an engine that does not publish load simply omits it, and such a worker opens no load sockets and keeps using the router-side signal. Both the introspector and the discovery fallback thread the field.

Subscription. A SubKind::Load registry opens one SUB per advertised DP rank on load_port_base + dp_rank, subscribe-all, sharing the existing pump channel. It is a separate registry from the KV one, each with its own handle map, so the two never collide despite sharing the (url, dp_rank) key.

Decoding. The payload is the engine's LoadSnapshot, a msgpack map keyed by field name. Every field defaults to zero — the Python side is omit_defaults — and unknown keys are ignored, because the snapshot carries considerably more than this policy reads (memory, speculative, LoRA, disaggregation and queue sub-structs) and will keep growing. Tolerating both is what stops an engine-side addition from becoming a router-side decode failure.

A worker is trusted only when its picture is complete and fresh. EngineLoadTable::snapshot_fresh sums a worker's ranks and omits the worker entirely if any rank is stale or if fewer ranks have reported than the worker advertised. Freshness alone is not sufficient: an entry that does not exist cannot be judged stale, so a rank whose publisher failed to bind at startup would leave that worker reporting a partial sum indefinitely — looking idle and drawing more traffic precisely because it is broken. Workers are recorded with the rank count they advertised, which gives the read path its denominator; a worker that never advertised a load topic has no denominator and is taken at face value.

Selection. WorkerLoads is built once per select() so every comparison in one decision sees a consistent view, and load_of returns engine-reported depth where a fresh snapshot exists and Worker::active_load() otherwise. Both the imbalance fast-path and the matched-set tiebreak go through it.

Degradation

Every failure path lands on the prior behavior rather than an error:

• no advertised load port → no load sockets opened;
• an undecodable frame → logged and dropped, the subscriber survives;
• entries older than the freshness window, or an incomplete rank set → the worker is omitted from the table;
• worker removal clears its entries so a re-added worker cannot inherit stale load.

Known gap

load_of can mix measurement bases within a single selection: engine queue depth for workers with a complete fresh snapshot, and the router-side in-flight counter for the rest. Those quantities are not interchangeable, so a fleet where only some workers report load can produce a fabricated spread in the imbalance check. Making the choice all-or-nothing per selection is a follow-up.

Accuracy Tests

Not applicable — routing policy only, no change to model execution or output.

Speed Tests and Profiling

No benchmark run. snapshot_fresh is computed once per selection so per-worker lookups are O(1); it does allocate per worker on the routing path, which is a candidate for Arc<str> keys if it shows up in profiles.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Tests cover the wire contract and each degradation path: decoding a load frame into WorkerEvent::Load; ignoring unknown fields and defaulting absent ones; rank sums, staleness, partial freshness, and a worker missing a rank entirely; engine load overriding active_load, the matched-set tiebreak using it, and stale load falling back; clearing a removed worker; and the /server_info introspection path both with and without the load port advertised.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!