feat(zmq): piggyback a scheduler-load snapshot on slim output batches

python/tokenspeed/runtime/engine/event_loop.py

@@ -553,6 +553,20 @@ def __init__(
             ),
             metrics=self.metrics,
         )
+
+        # msgpack wire only: piggyback a load snapshot on every output batch.
+        # The socket predates the scheduler (it is built during the startup
+        # handshake), so the sampler late-binds here — the first point where
+        # scheduler, cache geometry, and output processor all exist. Same
+        # sources as the per-iteration Prometheus snapshot in
+        # _record_scheduler_iteration_metrics.
+        if self.server_args.zmq_msgpack and self.attn_tp_rank == 0:
+            self.send_to_tokenizer.load_fn = lambda: (
+                len(self.output_processor.rid_to_state),
+                self.scheduler.waiting_size(),
+                self.scheduler.active_kv_pages(),
+                self._scheduler_cache_geometry.num_usable_pages,
+            )
         if server_args.disaggregation_mode != "null":
             kv_args = get_kv_args(
                 global_rank,

python/tokenspeed/runtime/engine/io_struct.py

@@ -712,10 +712,26 @@ class BatchTokenIDOutSlim(BaseBatchReq, kw_only=True):
     # DP the batch itself names its rank. Appended field: defaults to 0 so
     # older peers on either side stay compatible.
     engine_index: int = 0
+    # Piggybacked scheduler-load snapshot from the producing rank, sampled at
+    # send time. The pickle-mode GetLoad poll has no msgpack transport (control
+    # replies are dropped on this wire), so the output batch is the only
+    # in-band load channel; an external frontend uses these for least-loaded
+    # routing across a DP group. Same sources as the per-iteration Prometheus
+    # snapshot: running = resident request states, waiting = scheduler queue
+    # depth, and the KV ratio's numerator/denominator carried as exact page
+    # counts. Appended fields: all default to 0 so older peers stay
+    # compatible; a 0 kv_total_pages means "no snapshot" to the frontend.
+    num_running: int = 0
+    num_waiting: int = 0
+    kv_active_pages: int = 0
+    kv_total_pages: int = 0
 
     @classmethod
     def from_full(
-        cls, out: BatchTokenIDOut, engine_index: int = 0
+        cls,
+        out: BatchTokenIDOut,
+        engine_index: int = 0,
+        load: tuple[int, int, int, int] | None = None,
     ) -> "BatchTokenIDOutSlim":
         # Token source: ``out.output_ids`` — the not-yet-sent slice of each
         # request's generated ids. NOT ``out.decode_ids``: that is the
@@ -729,8 +745,13 @@ def from_full(
                 "BatchTokenIDOut.output_ids is None; the msgpack wire needs "
                 "the per-request generated token ids"
             )
+        num_running, num_waiting, kv_active_pages, kv_total_pages = load or (0, 0, 0, 0)
         return cls(
             engine_index=engine_index,
+            num_running=num_running,
+            num_waiting=num_waiting,
+            kv_active_pages=kv_active_pages,
+            kv_total_pages=kv_total_pages,
             rids=list(out.rids),
             output_ids=[list(ids) for ids in out.output_ids],
             finished_reasons=[_finish_type(fr) for fr in out.finished_reasons],

python/tokenspeed/runtime/engine/zmq_msgpack.py

@@ -144,13 +144,25 @@ def __init__(self, socket: zmq.Socket, engine_index: int = 0) -> None:
         self._socket = socket
         self._engine_index = engine_index
         self._encoder = MsgpackEncoder()
+        # Late-bound scheduler-load sampler, () -> (num_running, num_waiting,
+        # kv_active_pages, kv_total_pages). The socket is created during the
+        # startup handshake, before the scheduler exists, so the event loop
+        # binds this once the scheduler is up. None means "no snapshot":
+        # the slim batch's load fields stay at their 0 defaults.
+        self.load_fn = None
 
     def send_pyobj(self, obj) -> None:
         if isinstance(obj, BatchTokenIDOut):
             # The PULL side carries no routing identity, so the batch itself
             # names its producing rank; the frontend attributes per-rank
-            # outputs and load by this index under DP.
-            slim = BatchTokenIDOutSlim.from_full(obj, engine_index=self._engine_index)
+            # outputs and load by this index under DP. The load snapshot rides
+            # every batch because this wire has no control-reply channel for
+            # the GetLoad poll (dropped below).
+            slim = BatchTokenIDOutSlim.from_full(
+                obj,
+                engine_index=self._engine_index,
+                load=self.load_fn() if self.load_fn is not None else None,
+            )
             self._socket.send_multipart(self._encoder.encode(slim), copy=False)
         else:
             logger.warning(

test/runtime/test_zmq_msgpack.py

@@ -271,6 +271,56 @@ def test_slim_out_engine_index_defaults_for_older_senders():
     assert rt.engine_index == 0
 
 
+def test_slim_out_piggybacks_the_load_snapshot():
+    # The msgpack wire has no control-reply channel for the GetLoad poll, so
+    # the output batch carries the producing rank's scheduler-load snapshot:
+    # (num_running, num_waiting, kv_active_pages, kv_total_pages).
+    slim = BatchTokenIDOutSlim.from_full(_make_batch_out(), load=(2, 5, 100, 400))
+    rt = msgspec.msgpack.Decoder(BatchTokenIDOutSlim).decode(_encode_payload(slim))
+    assert rt.num_running == 2
+    assert rt.num_waiting == 5
+    assert rt.kv_active_pages == 100
+    assert rt.kv_total_pages == 400
+
+
+def test_slim_out_load_defaults_for_older_senders():
+    # Appended fields: a 10-element array (engine_index era, pre-load) must
+    # still decode with the zero "no snapshot" defaults — kv_total_pages == 0
+    # is the frontend's signal that no load rode this batch.
+    slim = BatchTokenIDOutSlim.from_full(_make_batch_out(), load=(2, 5, 100, 400))
+    raw = msgspec.msgpack.decode(_encode_payload(slim))
+    rt = msgspec.msgpack.Decoder(BatchTokenIDOutSlim).decode(
+        msgspec.msgpack.encode(raw[:10])
+    )
+    assert rt.num_running == 0
+    assert rt.num_waiting == 0
+    assert rt.kv_active_pages == 0
+    assert rt.kv_total_pages == 0
+
+
+def test_send_socket_samples_load_fn_per_batch():
+    # MsgpackSendSocket late-binds a load sampler (the socket exists before
+    # the scheduler); every BatchTokenIDOut send must sample it fresh.
+    sent = []
+
+    class _Sock:
+        def send_multipart(self, frames, copy=False):
+            sent.append(list(frames))
+
+    sender = zmq_msgpack.MsgpackSendSocket(_Sock(), engine_index=1)
+    samples = iter([(1, 0, 10, 400), (0, 0, 0, 400)])
+    sender.load_fn = lambda: next(samples)
+    sender.send_pyobj(_make_batch_out())
+    sender.send_pyobj(_make_batch_out())
+
+    decoder = msgspec.msgpack.Decoder(BatchTokenIDOutSlim)
+    first = decoder.decode(sent[0][0])
+    second = decoder.decode(sent[1][0])
+    assert (first.num_running, first.kv_active_pages) == (1, 10)
+    assert (second.num_running, second.kv_active_pages) == (0, 0)
+    assert first.engine_index == second.engine_index == 1
+
+
 def test_slim_out_sources_output_ids_not_the_detok_window():
     # ``decode_ids`` on the io_struct is the incremental-detokenization window,
     # which starts at the prompt tail for context; ``output_ids`` is the
@@ -305,22 +355,24 @@ def test_slim_out_carries_logprob_columns():
 
 
 def test_slim_out_is_tagged_positional_tuple():
-    """The wire form must be a 10-element tagged array (the frontend's codec
-    depends on the tag and column order; engine_index is the appended tail)."""
+    """The wire form must be a 14-element tagged array (the frontend's codec
+    depends on the tag and column order; engine_index and the load snapshot
+    are the appended tail)."""
     out = _make_batch_out(
         output_token_logprobs_val=[[-0.5]], output_token_logprobs_idx=[[10]]
     )
-    slim = BatchTokenIDOutSlim.from_full(out, engine_index=1)
+    slim = BatchTokenIDOutSlim.from_full(out, engine_index=1, load=(2, 5, 100, 400))
     raw = msgspec.msgpack.decode(_encode_payload(slim))
     assert isinstance(raw, list)
     assert raw[0] == "BatchTokenIDOutSlim"
-    assert len(raw) == 10
+    assert len(raw) == 14
     assert raw[1] == ["r1"]  # rids
     assert raw[2] == [[10, 11]]  # output_ids
     assert raw[3] == ["length"]  # finished_reasons
     assert raw[7] == [[pytest.approx(-0.5)]]
     assert raw[8] == [[10]]
-    assert raw[9] == 1  # engine_index (appended)
+    assert raw[9] == 1  # engine_index (appended, #1046)
+    assert raw[10:14] == [2, 5, 100, 400]  # load snapshot (appended)
 
 
 def test_slim_out_finish_reason_none_maps_to_empty():