router: --decode-request-weight balances decode DP ranks by held requests

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

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)
 

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.

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:-}"

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

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:

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

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 |

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);
+    }
 }

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));
         }
     }
 }

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

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<u128>) {
     }
 }
 
+pub fn set_decode_demand_requests(target: &str, requests: Option<usize>) {
+    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" };

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<KvEventClient>,
     hasher: Arc<BlockHasher>,
@@ -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<String, crate::decode_demand::Demand> {
         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<u128> = 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);