feat(router): add per-worker PD Prefill admission

[junliu-mde]

Description

PD and EPD routing has no way to bound how many Prefill requests one SMG process
assigns to a single Prefill worker. When a Prefill worker is saturated, SMG keeps
dispatching to it and the extra requests queue inside the inference engine, where
SMG can neither observe nor bound their wait.

This PR adds an optional per-worker Prefill in-flight limit and a bounded,
Router-wide FIFO admission queue in front of it.

Behavior

Options

Option	Description	Default
--prefill-max-inflight-requests-per-worker	Maximum in-flight Prefill requests per registered Prefill worker. A non-positive value disables admission.	-1
--prefill-queue-size	Router-wide FIFO queue depth. 0 disables waiting.	100 when enabled
--prefill-queue-timeout-secs	Maximum time a request may wait for Prefill capacity.	60 when enabled

A positive per-worker limit enables admission and the queue. Queue options
without a positive limit are rejected at startup, as is combining admission with
--priority-scheduler-enabled or a non-PD/EPD routing mode.

Accounting

• One client request occupies one Prefill admission slot. Batch inputs and
  parallel samples requested through n do not multiply slots: every backend
  sub-request of one client request shares the same reservation.
• A slot is held from worker selection until the Prefill phase reaches EOF, fails,
  or is cancelled. It is not held for the Decode phase.
• Worker load accounting is unchanged. Decode load, and Prefill load when
  admission is disabled, remain one unit per backend sub-request.
• The limit and the queue are local to one SMG process. With --dp-aware, every
  registered rank is a separate worker and gets its own limit.

Selection and queueing

• Waiting requests do not bind to a worker. Only the queue head can select one.
• At the head, SMG re-reads worker health, capacity, runtime pairing, and cache
  state, filters workers that are already at the limit, and only then runs the
  Prefill policy. Candidate filtering, policy selection, and slot reservation run
  under one lock, so a stateful policy such as cache_aware never commits a
  routing decision that then fails to reserve capacity.
• New requests cannot bypass existing waiters.
• With consistent_hashing, X-SMG-Target-Worker stays strict: the request waits
  for that worker instead of moving to another one.
• The queue head is woken when a Prefill slot is released and when a worker is
  registered, removed, replaced, or changes health state.
• The queue head also rechecks once per second. A worker whose circuit breaker
  moves from open to half-open does so lazily on its next state check and emits
  no registry event, so recovery would otherwise be invisible to a waiting
  request. Only the head polls, so the rest of the queue stays asleep.
• Cancelling a request removes its queue ticket and wakes the next waiter.

Errors

• A full queue returns 429 with pd_prefill_queue_full.
• An expired queue wait returns 429 with pd_prefill_queue_timeout.
• Both are local admission errors and are not retried inside SMG.
• No matching, healthy worker remains a 503 availability error and is not queued.

Scope

• Decode workers keep their existing load tracking but have no admission limit.
  Decode scheduling stays with the inference engine.
• HTTP PD, gRPC PD, and gRPC EPD share one admission controller.
• Regular gRPC PD now uses separate Prefill and Decode policies, matching HTTP and
  EPD.
• Decode load is acquired only after the Prefill reservation succeeds, so a
  rejected candidate no longer briefly inflates Decode load.
• The streaming /v1/responses handlers now wait for the backend request to
  start before they return. They previously returned 200 and an SSE stream
  immediately, which would have turned an admission 429 into a 200 followed
  by an SSE error event. Any failure to start the backend request now reaches
  the client as its real HTTP status.

Metrics

Metric	Type	Labels
smg_pd_prefill_admission_inflight	Gauge	worker
smg_pd_prefill_admission_queued	Gauge	—
smg_pd_prefill_admission_wait_seconds	Histogram	—
smg_pd_prefill_admission_rejections_total	Counter	reason

These measure SMG admission only. Engine-side queue metrics are unchanged.

Test Plan

Run against this branch:

Command	Result
cargo fmt --all -- --check	pass
cargo clippy --all-targets --all-features -- -D warnings	pass
cargo check -p smg -p smg-golang -p smg-python	pass
cargo test -p smg --lib	1378 passed, 0 failed, 5 ignored
cargo test -p smg --test routing_tests	102 passed, 0 failed
pytest bindings/python/tests on a Python 3.9 abi3 wheel built from this commit	248 passed, 4 skipped

New unit tests cover strict FIFO order, no queue bypass, queue full, queue
timeout, cancellation cleanup, re-selection without worker binding, one slot per
client request, zero-size queue rejection, release visibility after a drop, and
the availability path. Configuration tests cover the startup validation rules.

End-to-end run

Started the router from the built wheel against two mock workers: a Prefill
worker with a 3 s generation delay and an immediate Decode worker.

With --prefill-max-inflight-requests-per-worker 1 --prefill-queue-size 1 --prefill-queue-timeout-secs 2 and three overlapping requests:

Request	Result
1	200
2	429 pd_prefill_queue_timeout after 2.0 s
3	429 pd_prefill_queue_full

With --prefill-max-inflight-requests-per-worker 1 --prefill-queue-size 4 --prefill-queue-timeout-secs 30 and three overlapping requests: all returned
200, finished in arrival order, and took 9 s in total for three 3 s Prefills.
A sample taken while they were in flight reported
smg_pd_prefill_admission_inflight = 1 and smg_pd_prefill_admission_queued = 2;
the two queued requests waited 8.1 s in total.

Startup validation, each rejected with its own message:

Configuration	Result
Admission plus --priority-scheduler-enabled	startup fails
--prefill-queue-size without a positive per-worker limit	startup fails
A positive per-worker limit outside PD/EPD mode	startup fails

With no admission options the router starts and serves a PD request as before,
and with a per-worker limit above the offered load it serves the same request.

Checklist

• Formatting passes
• Checks and tests pass
• Documentation updated
• (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

Important

Review skipped

Draft detected.

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: ASSERTIVE

Plan: Pro Plus

Run ID: 2ad272a9-4d36-4ab2-a642-19cd7d4fb85d

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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