[AMD-yanfeiwang]

Motivation

Follow-up to #42027, which this PR is stacked on (its commit is included until it merges).

With DP attention every decode step waits for the busiest DP rank, so what matters is how many requests each rank holds. p2c on num_total_tokens evens out KV across decode ranks but leaves their request counts as uneven as round_robin (tables below).

Modifications

SGLANG_ROUTER_P2C_LOAD_METRIC selects what power_of_two compares:

value	compares	ties
tokens (default)	polled num_total_tokens, two random candidates (unchanged)	in-flight requests
requests	the router's in-flight requests, two random candidates	polled tokens
least_requests	the router's in-flight requests, every candidate	polled tokens, then random

In-flight requests are counted from dispatch to the end of the response (#42027). For a decode rank that is its running batch, every queue (waiting, prealloc, transfer, retracted) and requests still in prefill: the set /v1/loads reports as num_running_reqs + num_waiting_reqs, but live instead of polled. least_requests skips workers the load monitor has no entry for unless none has one, and reads the counters without locking, so concurrent selections may pick the same worker.

The switch is an env var to keep the policy config and CLI unchanged; happy to turn it into a --p2c-load-metric flag instead.

Results (MI355X)

DeepSeek-V4-Pro AgentX trace replay, 1P1D, prefill TP8/EP8 and decode TP8, both with DP attention (8 decode ranks); router --policy consistent_hashing --dp-aware --worker-startup-check-interval 1 (the load monitor polls at that interval). Router built from this branch on lmsysorg/sglang-rocm:v0.5.21-rocm724-mi35x-20261007; each point is a 1 h measurement.

Absolute values; in parentheses, the change against round_robin. All power_of_two columns poll loads every 1 s.

Output throughput (tok/s/GPU)

conc	round_robin	tokens	requests	least_requests
192	369.9	366.7 (-0.9%)	375.5 (+1.5%)	372.8 (+0.8%)
256	439.4	438.4 (-0.2%)	446.3 (+1.6%)	444.3 (+1.1%)
384	563.5	568.4 (+0.9%)	570.1 (+1.2%)	570.0 (+1.2%)
512	590.7	601.1 (+1.8%)	602.6 (+2.0%)	594.3 (+0.6%)

TPOT p50 (ms)

conc	round_robin	tokens	requests	least_requests
192	12.45	12.63 (+1.4%)	11.90 (-4.4%)	11.80 (-5.2%)
256	13.92	14.07 (+1.1%)	13.57 (-2.5%)	13.08 (-6.0%)
384	17.34	16.67 (-3.9%)	15.92 (-8.2%)	15.65 (-9.7%)
512	18.45	18.72 (+1.5%)	17.74 (-3.8%)	17.99 (-2.5%)

TPOT p90 (ms)

conc	round_robin	tokens	requests	least_requests
192	14.82	14.85 (+0.2%)	13.81 (-6.8%)	13.75 (-7.2%)
256	16.35	16.58 (+1.4%)	15.99 (-2.2%)	15.48 (-5.3%)
384	20.67	19.82 (-4.1%)	18.57 (-10.2%)	18.56 (-10.2%)
512	22.22	22.21 (0.0%)	21.59 (-2.8%)	21.70 (-2.3%)

Decode-rank balance over the measured hour, per-second samples across the 8 ranks (round_robin / tokens / requests / least_requests); outstanding = running + all queues:

	c192	c256	c384	c512
outstanding requests, CV	0.17 / 0.22 / 0.09 / 0.05	0.14 / 0.18 / 0.06 / 0.03	0.10 / 0.15 / 0.04 / 0.02	0.08 / 0.12 / 0.02 / 0.01
running batch, CV	0.22 / 0.24 / 0.17 / 0.14	0.19 / 0.20 / 0.14 / 0.13	0.16 / 0.17 / 0.13 / 0.13	0.17 / 0.15 / 0.14 / 0.14
busiest rank's batch	14.8 / 15.2 / 13.1 / 12.7	19.0 / 19.4 / 17.6 / 17.0	28.8 / 28.3 / 25.9 / 25.3	32.5 / 32.5 / 30.6 / 30.5
busiest / mean rank KV tokens	1.43 / 1.23 / 1.35 / 1.32	1.38 / 1.17 / 1.28 / 1.27	1.31 / 1.12 / 1.22 / 1.21	1.24 / 1.08 / 1.19 / 1.18

requests lowers the busiest rank's batch by 6-12% and least_requests by 6-14%, which is where the TPOT gain comes from. Throughput moves less because output is under 1% of the tokens in this input-heavy workload. TTFT stays within run-to-run noise except c384 (p50 +12% with requests, +15% with least_requests), to be rechecked.

Concurrency 512 over the hour

p2c token and p2c in-flight requests below are tokens and requests.

requests keeps outstanding requests flat across decode ranks for the whole hour (bottom) and least_requests flatter still, while tokens is less even than round_robin. At 512 more than half (55-57%) of each rank's in-flight requests are still waiting on KV transfer, so equal in-flight counts no longer mean equal running batches: the busiest rank's batch (top) improves less than at lower concurrency, and evening out in-flight requests further with least_requests (CV 0.01 vs 0.02) leaves it unchanged (30.5 vs 30.6).

With requests every rank's hourly mean running batch is within 2% of the 8-rank mean, and no rank is the busiest more than 17% of the time (round_robin: DP6 at +6.3% and 21%; 12.5% if even).

Test

Two new unit tests: requests prefers fewer in-flight requests over a lower token load, and least_requests picks the global minimum while skipping workers without a polled load. cargo test --lib (398 passed), cargo clippy --all-targets --all-features -- -D warnings and cargo +nightly fmt -- --check pass.

---

CI States

Latest PR Test (Base): ✅ Run #38356273395
Latest PR Test (Extra): ❌ Run #38356273155
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

[inkcherry]

hi, @ShangmingCai , could you please help on review this PR and #42027. This PR improves load balancing in the PD (prefill–decode disaggregation) setup.

We will try to land these optimizations in the experimental router later. That will take some time.

[ShangmingCai]

LGTM

[ShangmingCai]

/tag-and-rerun-ci