diff --git a/examples/sglang_2p1d_glm5.2/README.md b/examples/sglang_2p1d_glm5.2/README.md index 80d9ea07e..0bf8df7e7 100644 --- a/examples/sglang_2p1d_glm5.2/README.md +++ b/examples/sglang_2p1d_glm5.2/README.md @@ -139,6 +139,41 @@ CONC=176 bash engine/agentx/agentx.sh # DURATION=3600 WARMUP_PER_LANE=10 by bash engine/down.sh # then wait for VRAM to drain on both nodes ``` +## Decode request balancing + +`--decode-input-demand on` balances the decode ranks by in-flight input tokens. +That evens their KV, but not their request counts: one rank can hold many short +requests for another's one long request, and the busiest rank sets every step. +`--decode-request-weight W` adds `W` tokens per request a rank already holds to +its demand (the shape of Dynamo's `router-decode-active-request-weight`); 0 is +token-only demand. The kit passes `DECODE_REQUEST_WEIGHT` (env.sh, default +1,000,000); `DECODE_REQUEST_WEIGHT=0` turns it off. + +1P1D C88 on two AUS MI355X nodes (smci355-ccs-aus-n04-25 / n06-33), prefills +with the EAGLE flags (note 10), one deployment, prefill cache flushed and router +restarted before each point, 10 warm-up requests per lane, **1800 s measured**: + +| weight | tok/s/GPU | requests (per s) | interactivity p50 / p90 | TTFT p50 / p90 / mean | ITL p50 / p90 | busiest rank running, median / p90 | +|---|---|---|---|---|---|---| +| 0 | 26,462 | 5,160 (2.87) | 75.8 / 55.7 | 5.5 / 11.1 / 8.3 s | 13.2 / 18.0 ms | 8 / 11 | +| 1,000,000 | 27,186 | 5,184 (2.88) | 86.5 / 79.6 | 5.6 / 12.0 / 9.5 s | 11.6 / 12.6 ms | 6 / 8 | + +- p90 interactivity +43% at the same throughput. Both points take the same time + per step at the same busiest-rank load (40.5 / 40.8 ms at 6 requests, 43.6 / + 43.5 ms at 8): the gain is the load distribution, the busiest rank no longer + reaching 10-14 requests. +- TTFT rises (mean +15%, p90 +8%): the ranks' context is less even (busiest / + mean 1.23 → 1.52 at the median), so more KV lands on fewer ranks. +- On this point one more request on the busiest rank cost as much as about + 1.05M more context tokens there (least-squares fit of step time against the + busiest rank's requests and tokens and the totals over 3 s bins of the decode + log), which is why the weight is a million; on 2P1D C176 with the v0.5.21 + image the same fit gave 0.17-0.19M. A 180,000 weight on another node pair + gave 27,577 tok/s/GPU at 85.0 / 74.4. +- 1800 s windows understate tok/s/GPU against the canonical 3600 s by about 6% + (sessions reach their longer turns later), and leave p90 interactivity within + 2%. + ## Notes **1. Forced MTP acceptance produces wrong text.** `SIM_ACCEPT_LEN` makes the diff --git a/examples/sglang_2p1d_glm5.2/engine/common.sh b/examples/sglang_2p1d_glm5.2/engine/common.sh index 8002714da..5bdec2d08 100755 --- a/examples/sglang_2p1d_glm5.2/engine/common.sh +++ b/examples/sglang_2p1d_glm5.2/engine/common.sh @@ -25,7 +25,7 @@ derive_names # run, and forwarded to the decode node. KIT_VARS=(PREFILL_NODE_IP DECODE_NODE_IP DECODE_SSH MODEL SERVED_MODEL IMAGE WORK_DIR RDMA_DEVICES MC_GID_INDEX MOONCAKE_DISABLE_HIP_DMABUF PREFILL_WORKERS PREFILL0_GPUS PREFILL1_GPUS DECODE_GPUS - SIM_ACCEPT_LEN PREFILL_EXTRA_ARGS DECODE_EXTRA_ARGS ROUTER_EXTRA_ARGS + SIM_ACCEPT_LEN DECODE_REQUEST_WEIGHT PREFILL_EXTRA_ARGS DECODE_EXTRA_ARGS ROUTER_EXTRA_ARGS ETCD_PORT ETCD_PEER_PORT ROUTER_PORT PREFILL_PORT DECODE_PORT BOOTSTRAP_PORT KV_EVENTS_PORT KV_SNAPSHOT_PORT ETCD_IMAGE CTR) diff --git a/examples/sglang_2p1d_glm5.2/engine/up.sh b/examples/sglang_2p1d_glm5.2/engine/up.sh index 1171ac6eb..aa1c83da0 100755 --- a/examples/sglang_2p1d_glm5.2/engine/up.sh +++ b/examples/sglang_2p1d_glm5.2/engine/up.sh @@ -49,7 +49,7 @@ docker run -d --name "$CTR-router" --network host -v "$MODEL:$MODEL:ro" "$IMAGE" --request-transport http --kv-event-transport zmq --router-tokenizer-path "$MODEL" \ --router-policy kv-aware --kv-prefill-overlap-weight 20 --kv-decode-overlap-weight 2 \ --session-affinity prefill --pd-prefill-guard-release completion --decode-input-demand on \ - "${router_extra[@]}" >/dev/null + --decode-request-weight "$DECODE_REQUEST_WEIGHT" "${router_extra[@]}" >/dev/null wait_health "$ROUTER_URL/health" "$CTR-router" # A prefill registers only after infera's start-up probe has moved KV to every # decode rank, a minute or more after its /health answers. diff --git a/examples/sglang_2p1d_glm5.2/env.sh b/examples/sglang_2p1d_glm5.2/env.sh index c96765769..8e0e66bf8 100755 --- a/examples/sglang_2p1d_glm5.2/env.sh +++ b/examples/sglang_2p1d_glm5.2/env.sh @@ -58,6 +58,11 @@ export DECODE_GPUS="${DECODE_GPUS:-0,1,2,3,4,5,6,7}" # reproduces the AgentX reference numbers (InferenceX's GLM-5.2 value). export SIM_ACCEPT_LEN="${SIM_ACCEPT_LEN:-}" +# Tokens one held request is worth when the router balances the decode ranks +# (--decode-request-weight); 0 balances by input tokens only (README "Decode +# request balancing"). +export DECODE_REQUEST_WEIGHT="${DECODE_REQUEST_WEIGHT:-1000000}" + # Extra flags, space-separated, appended last. export PREFILL_EXTRA_ARGS="${PREFILL_EXTRA_ARGS:-}" export DECODE_EXTRA_ARGS="${DECODE_EXTRA_ARGS:-}" diff --git a/infera/server/args.py b/infera/server/args.py index 600b789a4..8ab090ff8 100644 --- a/infera/server/args.py +++ b/infera/server/args.py @@ -19,6 +19,13 @@ def _session_ttl(raw: str) -> int: return ttl +def _non_negative_int(raw: str) -> int: + value = int(raw) + if value < 0: + raise argparse.ArgumentTypeError(f"must be >= 0, got {value}") + return value + + def parse_server_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="Infera server") parser.add_argument("--host", default="0.0.0.0") @@ -134,6 +141,15 @@ def parse_server_args(argv: list[str] | None = None) -> argparse.Namespace: "input tokens. 'shadow' books demand and reports a suggestion without " "changing selection. Unset defers to $INFERA_DECODE_INPUT_DEMAND, else off.", ) + parser.add_argument( + "--decode-request-weight", + type=_non_negative_int, + default=None, + help="Rust backend with --decode-input-demand: input tokens charged per " + "request a Decode target already holds, so DP ranks are balanced by " + "request count as well as tokens. Unset defers to " + "$INFERA_DECODE_REQUEST_WEIGHT, else 0 (tokens only).", + ) parser.add_argument( "--pd-prefill-guard-release", choices=("decode", "completion"), @@ -342,6 +358,7 @@ def parse_server_args(argv: list[str] | None = None) -> argparse.Namespace: flag for flag, value in ( ("--decode-input-demand", args.decode_input_demand), + ("--decode-request-weight", args.decode_request_weight), ("--pd-prefill-guard-release", args.pd_prefill_guard_release), ) if value is not None diff --git a/infera/server/launch_rust.py b/infera/server/launch_rust.py index ee62f61f4..0c1db9ddb 100644 --- a/infera/server/launch_rust.py +++ b/infera/server/launch_rust.py @@ -132,6 +132,7 @@ def exec_rust(args: argparse.Namespace) -> None: ("--nats-req-max-duration-s", args.nats_req_max_duration), ("--nats-req-max-pending", args.nats_req_max_pending), ("--decode-input-demand", args.decode_input_demand), + ("--decode-request-weight", args.decode_request_weight), ("--pd-prefill-guard-release", args.pd_prefill_guard_release), ): if value is not None: diff --git a/manual/features/kv_aware_routing.md b/manual/features/kv_aware_routing.md index 8fe902beb..b15ee59cf 100644 --- a/manual/features/kv_aware_routing.md +++ b/manual/features/kv_aware_routing.md @@ -277,6 +277,32 @@ the mean Decode allocation wait from 10.5 s to 3.3 s. That run also moved to different nodes, so its throughput change is not attributable to demand routing alone. +### Request-count weight + +Under DP attention every decode step waits for the busiest rank, and that rank's +step time grows with its request count as well as its context. Token-only demand +evens the context but can leave one rank with many short requests. +`--decode-request-weight W` (or `INFERA_DECODE_REQUEST_WEIGHT`; requires +`--decode-input-demand shadow|on`) changes a candidate's cost to + + in-flight input tokens + this request's tokens + W × requests the target holds + +where held requests are the reservations above (booked at the pick, released at +response EOF, cancellation or failure). `W = 0`, the default, is token-only +demand. Dynamo's `router-decode-active-request-weight` scores Decode workers the +same way (`decode blocks + weight × active requests`). + +Choose `W` as the context, in tokens, that costs one decode step as much as one +more request on the busiest rank. A least-squares fit of step time against the +busiest rank's requests and tokens (3 s bins of the Decode log) gave about 1.05M +for GLM-5.2 1P1D C88 on SGLang v0.5.22 and 0.17-0.19M for 2P1D C176 on v0.5.21. +On that 1P1D C88 deployment, `W = 1000000` raised p90 interactivity from 55.7 to +79.6 at the same throughput and raised TTFT (mean +15%), because context spreads +less evenly; see `examples/sglang_2p1d_glm5.2/README.md`. + +`/metrics` adds `infera_router_decode_demand_requests{target}`, the held +requests per Decode route key. + ### Future: one tokenization end to end The router and each engine still tokenize the same prompt separately. Sending the diff --git a/manual/reference/cli.md b/manual/reference/cli.md index cd2180731..e3f8e375c 100644 --- a/manual/reference/cli.md +++ b/manual/reference/cli.md @@ -21,6 +21,7 @@ the **same** on the server and every worker. See | `--session-affinity` | `off` | `off` \| `prefill` \| `both`: pin `X-Session-ID` sessions on top of the policy | | `--session-affinity-ttl-seconds` | `3600` | idle seconds before a session binding is dropped, 1–86400 | | `--decode-input-demand` | (unset: `$INFERA_DECODE_INPUT_DEMAND`, else `off`) | `off` \| `shadow` \| `on`: pick Decode targets by in-flight input tokens (Rust backend, `kv-aware`) | +| `--decode-request-weight` | (unset: `$INFERA_DECODE_REQUEST_WEIGHT`, else `0`) | with `--decode-input-demand`: input tokens charged per request a Decode target already holds, so ranks balance by request count as well as tokens | | `--pd-prefill-guard-release` | (unset: `$INFERA_PD_PREFILL_GUARD_RELEASE`, else `decode`) | `decode` \| `completion`: release PD Prefill load when Decode ends or when Prefill completes (Rust backend) | | `--kv-overlap-weight` | `1.0` | KV-aware: cache locality vs load balance | | `--kv-prefill-overlap-weight` | (unset) | KV-aware PD: prefill-side weight (typical `20.0`); overrides the global | diff --git a/rust/router/src/config.rs b/rust/router/src/config.rs index 0f79e9649..a73211c65 100644 --- a/rust/router/src/config.rs +++ b/rust/router/src/config.rs @@ -79,6 +79,14 @@ pub struct Config { )] pub decode_input_demand: crate::decode_demand::Mode, + /// Decode input demand: the input tokens one in-flight request is worth on a + /// decode target, added per request it already holds. Under DP attention a + /// step waits for the busiest rank, whose cost grows with its request count + /// as well as its context, so token-only balancing can leave one rank with + /// many short requests. 0 keeps token-only balancing. + #[arg(long, default_value_t = 0, env = "INFERA_DECODE_REQUEST_WEIGHT")] + pub decode_request_weight: u64, + /// `etcd` (external) or `kubernetes` (workers publish into their own Pod /// annotation and the API server is watched). #[arg(long, default_value = "etcd")] @@ -258,6 +266,11 @@ impl Config { { anyhow::bail!("--decode-input-demand requires --router-policy kv-aware"); } + if self.decode_request_weight > 0 + && self.decode_input_demand == crate::decode_demand::Mode::Off + { + anyhow::bail!("--decode-request-weight requires --decode-input-demand shadow or on"); + } if self.router_policy == "kv-aware" && self.kv_tokenizer_path.is_none() { tracing::warn!( "--router-policy kv-aware without --kv-tokenizer-path: block hashes \ @@ -502,4 +515,25 @@ mod tests { } assert!(Config::try_parse_from(["router", "--decode-input-demand", "enabled"]).is_err()); } + + #[test] + fn decode_request_weight_requires_decode_demand() { + let weighted = |mode: &str| { + Config::try_parse_from([ + "router", + "--router-policy", + "kv-aware", + "--decode-input-demand", + mode, + "--decode-request-weight", + "180000", + ]) + .unwrap() + }; + assert!(weighted("off").validate().is_err()); + assert!(weighted("shadow").validate().is_ok()); + assert!(weighted("on").validate().is_ok()); + let default = Config::try_parse_from(["router"]).unwrap(); + assert_eq!(default.decode_request_weight, 0); + } } diff --git a/rust/router/src/decode_demand.rs b/rust/router/src/decode_demand.rs index b971a47f2..21e6ab6e1 100644 --- a/rust/router/src/decode_demand.rs +++ b/rust/router/src/decode_demand.rs @@ -64,6 +64,7 @@ impl Reservation { entry.tokens += u128::from(tokens); entry.requests += 1; crate::metrics::set_decode_demand_tokens(&key, Some(entry.tokens)); + crate::metrics::set_decode_demand_requests(&key, Some(entry.requests)); Self { ledger: ledger.clone(), key, @@ -85,8 +86,10 @@ impl Drop for Reservation { if entry.requests == 0 { state.remove(&self.key); crate::metrics::set_decode_demand_tokens(&self.key, None); + crate::metrics::set_decode_demand_requests(&self.key, None); } else { crate::metrics::set_decode_demand_tokens(&self.key, Some(entry.tokens)); + crate::metrics::set_decode_demand_requests(&self.key, Some(entry.requests)); } } } diff --git a/rust/router/src/main.rs b/rust/router/src/main.rs index 199daef3d..6fab98481 100644 --- a/rust/router/src/main.rs +++ b/rust/router/src/main.rs @@ -92,6 +92,7 @@ async fn main() -> anyhow::Result<()> { cfg.kv_per_worker_template_kwargs, )) .with_decode_input_demand(cfg.decode_input_demand) + .with_decode_request_weight(cfg.decode_request_weight) // A worker that served traffic before the router subscribed emits // its one rooted event to nobody, and no later event can rebuild // the chain. Asking it to flush is the only repair, so the router diff --git a/rust/router/src/metrics.rs b/rust/router/src/metrics.rs index 7381e8eff..3f231eb91 100644 --- a/rust/router/src/metrics.rs +++ b/rust/router/src/metrics.rs @@ -36,6 +36,7 @@ struct Metrics { prefix_blocks_hit: CounterVec, prefix_blocks_total: CounterVec, decode_demand_tokens: GaugeVec, + decode_demand_requests: GaugeVec, decode_demand_decisions: CounterVec, engine_kv_cache_usage: GaugeVec, engine_prefix_cache_hit_rate: GaugeVec, @@ -148,6 +149,12 @@ fn m() -> &'static Metrics { &["target"] ) .expect("register decode_demand_tokens"), + decode_demand_requests: register_gauge_vec!( + "infera_router_decode_demand_requests", + "In-flight requests booked by decode input-demand routing, by Decode route key", + &["target"] + ) + .expect("register decode_demand_requests"), decode_demand_decisions: register_counter_vec!( "infera_router_decode_demand_decisions_total", "Decode input-demand picks by mode, exact or estimated input, and whether the \ @@ -245,6 +252,16 @@ pub fn set_decode_demand_tokens(target: &str, tokens: Option) { } } +pub fn set_decode_demand_requests(target: &str, requests: Option) { + let gauge = &m().decode_demand_requests; + match requests { + Some(requests) => gauge.with_label_values(&[target]).set(requests as f64), + None => { + let _ = gauge.remove_label_values(&[target]); + } + } +} + pub fn record_decode_demand_decision(mode: &str, exact: bool, agrees: bool) { let input = if exact { "exact" } else { "estimated" }; let agrees = if agrees { "true" } else { "false" }; diff --git a/rust/router/src/policy.rs b/rust/router/src/policy.rs index 9a32b0e33..7e1a9554b 100644 --- a/rust/router/src/policy.rs +++ b/rust/router/src/policy.rs @@ -336,6 +336,9 @@ const ZERO_HIT_ALARM_REPEAT: u64 = 1024; /// it. Both halves of the load term are needed -- see [`RECENT_DECAY`]. pub struct KvEventAwarePolicy { decode_input_demand: DecodeDemandMode, + /// Input tokens charged per request a Decode target already holds; see + /// `Config::decode_request_weight`. + decode_request_weight: u64, demand: Ledger, kv: Arc, hasher: Arc, @@ -381,6 +384,7 @@ impl KvEventAwarePolicy { ) -> Self { KvEventAwarePolicy { decode_input_demand: DecodeDemandMode::Off, + decode_request_weight: 0, demand: Ledger::default(), kv, hasher: Arc::new(hasher), @@ -402,6 +406,11 @@ impl KvEventAwarePolicy { self } + pub fn with_decode_request_weight(mut self, tokens_per_request: u64) -> Self { + self.decode_request_weight = tokens_per_request; + self + } + #[cfg(test)] pub(crate) fn demand_snapshot(&self) -> HashMap { self.demand.lock().unwrap().clone() @@ -773,11 +782,15 @@ impl Policy for KvEventAwarePolicy { let (picked, suggested, reservation, costs) = { // Only snapshot, selection and reservation hold the ledger lock. let mut ledger = self.demand.lock().unwrap_or_else(PoisonError::into_inner); + let weight = u128::from(self.decode_request_weight); let costs: Vec = keys .iter() .zip(&work) .map(|(key, (tokens, _))| { - ledger.get(key).map_or(0, |d| d.tokens) + u128::from(*tokens) + let booked = ledger.get(key); + booked.map_or(0, |d| d.tokens) + + u128::from(*tokens) + + weight * booked.map_or(0, |d| d.requests as u128) }) .collect(); let suggested = (0..targets.len()) @@ -804,6 +817,7 @@ impl Policy for KvEventAwarePolicy { && !self.parity.diverged(&targets[picked].worker.worker_id); tracing::debug!(target: "infera_router::decode_demand", mode=?self.decode_input_demand, exact, candidates=?keys, projected_tokens=?costs, + request_weight=self.decode_request_weight, legacy_selected=%keys[legacy_pick], selected=%keys[picked], suggested=%keys[suggested], "decode input demand"); crate::metrics::record_decode_demand_decision( @@ -1437,6 +1451,58 @@ mod tests { assert!(policy.demand_snapshot().is_empty()); } + #[test] + fn decode_request_weight_balances_requests_held() { + // "a" holds one long request, "b" three short ones: token-only demand + // sends the next request to "b"; at 1000 tokens per held request "a" + // costs 2000 + 1000 + 32 against "b"'s 300 + 3000 + 32. + for (weight, want) in [(0, "b"), (1000, "a")] { + let policy = demand_policy(DecodeDemandMode::On).with_decode_request_weight(weight); + let workers = vec![worker("a", 0, None), worker("b", 0, None)]; + let held = { + let mut ledger = policy.demand.lock().unwrap(); + let mut held = vec![Reservation::book( + &policy.demand, + &mut ledger, + "a".into(), + 2000, + )]; + for _ in 0..3 { + held.push(Reservation::book( + &policy.demand, + &mut ledger, + "b".into(), + 100, + )); + } + held + }; + let pick = policy.pick(&workers, &json!({"prompt": vec![7;32]}), Role::Decode); + assert_eq!(pick.target.route_key(), want, "weight {weight}"); + let snapshot = policy.demand_snapshot(); + assert_eq!(snapshot[want].requests, if want == "a" { 2 } else { 4 }); + drop((held, pick)); + assert!(policy.demand_snapshot().is_empty()); + } + } + + #[test] + fn decode_request_weight_in_shadow_keeps_the_legacy_pick() { + let legacy = demand_policy(DecodeDemandMode::Off); + let shadow = demand_policy(DecodeDemandMode::Shadow).with_decode_request_weight(1_000_000); + let workers = vec![worker("a", 16, None), worker("b", 16, None)]; + let mut picks = Vec::new(); + for n in 1..12 { + let request = json!({"prompt": vec![7; n * 16]}); + let a = legacy.pick(&workers, &request, Role::Decode); + let b = shadow.pick(&workers, &request, Role::Decode); + assert_eq!(a.target.route_key(), b.target.route_key()); + picks.push((a, b)); + } + drop(picks); + assert!(shadow.demand_snapshot().is_empty()); + } + #[test] fn decode_demand_estimated_inputs_still_steer() { let policy = demand_policy(DecodeDemandMode::On); diff --git a/tests/unit/server/test_launch_rust.py b/tests/unit/server/test_launch_rust.py index 69e41f26f..69cf02903 100644 --- a/tests/unit/server/test_launch_rust.py +++ b/tests/unit/server/test_launch_rust.py @@ -54,6 +54,7 @@ def _args(**overrides) -> argparse.Namespace: session_affinity="off", session_affinity_ttl_seconds=3600, decode_input_demand=None, + decode_request_weight=None, pd_prefill_guard_release=None, ) base.update(overrides) @@ -128,19 +129,27 @@ def test_rust_only_routing_switches_are_forwarded_only_when_given(monkeypatch): """Unset defers to the INFERA_* variables the Rust side reads itself.""" argv = _argv(monkeypatch) assert "--decode-input-demand" not in argv + assert "--decode-request-weight" not in argv assert "--pd-prefill-guard-release" not in argv argv = _argv( monkeypatch, router_policy="kv-aware", decode_input_demand="shadow", + decode_request_weight=0, pd_prefill_guard_release="completion", ) assert _value_of(argv, "--decode-input-demand") == "shadow" + assert _value_of(argv, "--decode-request-weight") == "0" assert _value_of(argv, "--pd-prefill-guard-release") == "completion" @pytest.mark.parametrize( - "flags", [["--decode-input-demand", "on"], ["--pd-prefill-guard-release", "completion"]] + "flags", + [ + ["--decode-input-demand", "on"], + ["--decode-request-weight", "180000"], + ["--pd-prefill-guard-release", "completion"], + ], ) def test_rust_only_routing_switches_need_the_rust_backend(flags, capsys): argv = ["--router-tokenizer-path", "tokenizer", *flags] @@ -148,7 +157,14 @@ def test_rust_only_routing_switches_need_the_rust_backend(flags, capsys): parse_server_args(["--router-backend", "python", *argv]) assert "requires --router-backend rust" in capsys.readouterr().err args = parse_server_args(["--router-backend", "rust", *argv]) - assert flags[1] in (args.decode_input_demand, args.pd_prefill_guard_release) + given = (args.decode_input_demand, args.decode_request_weight, args.pd_prefill_guard_release) + assert flags[1] in [str(v) for v in given if v is not None] + + +def test_decode_request_weight_rejects_negative(capsys): + with pytest.raises(SystemExit): + parse_server_args(["--router-backend", "rust", "--decode-request-weight", "-1"]) + assert "must be >= 0" in capsys.readouterr().err @pytest.mark.parametrize(