feat(policies): route least_load by token-work expected-wait

bindings/python/src/lib.rs

@@ -377,6 +377,9 @@ struct Router {
     eviction_interval_secs: u64,
     max_tree_size: usize,
     block_size: usize,
+    least_load_kv_pressure_weight: f64,
+    least_load_default_throughput: f64,
+    least_load_mean_prefill_tokens: u32,
     max_idle_secs: u64,
     assignment_mode: String,
     max_payload_size: usize,
@@ -518,7 +521,9 @@ impl Router {
                 },
                 PolicyType::LeastLoad => ConfigPolicyConfig::LeastLoad {
                     load_check_interval_secs: 5,
-                    kv_pressure_weight: 1.5,
+                    kv_pressure_weight: self.least_load_kv_pressure_weight,
+                    mean_prefill_tokens: self.least_load_mean_prefill_tokens,
+                    default_throughput: self.least_load_default_throughput,
                 },
                 PolicyType::Bucket => ConfigPolicyConfig::Bucket {
                     balance_abs_threshold: self.balance_abs_threshold,
@@ -775,6 +780,9 @@ impl Router {
         eviction_interval_secs = 120,
         max_tree_size = 2usize.pow(26),
         block_size = 16,
+        least_load_kv_pressure_weight = 0.15,
+        least_load_default_throughput = 2000.0,
+        least_load_mean_prefill_tokens = 1024,
         max_idle_secs = 14400,
         assignment_mode = String::from("random"),
         max_payload_size = 512 * 1024 * 1024,
@@ -885,6 +893,9 @@ impl Router {
         eviction_interval_secs: u64,
         max_tree_size: usize,
         block_size: usize,
+        least_load_kv_pressure_weight: f64,
+        least_load_default_throughput: f64,
+        least_load_mean_prefill_tokens: u32,
         max_idle_secs: u64,
         assignment_mode: String,
         max_payload_size: usize,
@@ -1004,6 +1015,9 @@ impl Router {
             eviction_interval_secs,
             max_tree_size,
             block_size,
+            least_load_kv_pressure_weight,
+            least_load_default_throughput,
+            least_load_mean_prefill_tokens,
             max_idle_secs,
             assignment_mode,
             max_payload_size,

bindings/python/src/smg/router_args.py

@@ -51,6 +51,9 @@ class RouterArgs:
     eviction_interval_secs: int = 60
     max_tree_size: int = 2**26
     block_size: int = 16
+    least_load_kv_pressure_weight: float = 0.15
+    least_load_default_throughput: float = 2000.0
+    least_load_mean_prefill_tokens: int = 1024
     max_idle_secs: int = 4 * 3600
     assignment_mode: str = "random"  # Mode for manual policy new routing key assignment
     max_payload_size: int = 512 * 1024 * 1024  # 512MB default for large batches
@@ -332,6 +335,30 @@ def add_cli_args(
             default=RouterArgs.cache_threshold,
             help="Cache threshold (0.0-1.0) for cache-aware routing",
         )
+        routing_group.add_argument(
+            f"--{prefix}least-load-kv-pressure-weight",
+            type=float,
+            default=RouterArgs.least_load_kv_pressure_weight,
+            help="KV-pressure weight (seconds) for the least_load policy",
+        )
+        routing_group.add_argument(
+            f"--{prefix}least-load-default-throughput",
+            type=float,
+            default=RouterArgs.least_load_default_throughput,
+            help=(
+                "Fallback generation throughput (tokens/s) for least_load when a"
+                " backend reports no live throughput"
+            ),
+        )
+        routing_group.add_argument(
+            f"--{prefix}least-load-mean-prefill-tokens",
+            type=int,
+            default=RouterArgs.least_load_mean_prefill_tokens,
+            help=(
+                "Mean prefill tokens for least_load's in-flight estimate when a"
+                " request's token count is unknown at routing"
+            ),
+        )
         routing_group.add_argument(
             f"--{prefix}balance-abs-threshold",
             type=int,

crates/grpc_client/proto/sglang_scheduler.proto

@@ -518,6 +518,12 @@ message SchedulerLoad {
   double utilization = 10;
   int32 max_running_requests = 11;
 
+  // Queued token-work: waiting-queue tokens not yet served from cache
+  // (SGLang num_waiting_uncached_tokens). Field 17 is intentionally out of
+  // order — appended after the optional-section field numbers to preserve
+  // wire compatibility with existing readers.
+  int32 num_waiting_uncached_tokens = 17;
+
   // Optional sections
   optional MemoryMetrics memory = 12;
   optional SpeculativeMetrics speculative = 13;

crates/grpc_client/src/sglang_scheduler.rs

@@ -703,6 +703,7 @@ impl From<proto::SchedulerLoad> for openai_protocol::worker::SchedulerLoadSnapsh
             dp_rank: load.dp_rank,
             num_running_reqs: load.num_running_reqs,
             num_waiting_reqs: load.num_waiting_reqs,
+            num_waiting_uncached_tokens: load.num_waiting_uncached_tokens,
             num_total_reqs: load.num_total_reqs,
             num_used_tokens: load.num_used_tokens,
             max_total_num_tokens: load.max_total_num_tokens,

crates/grpc_client/src/tokenspeed_scheduler.rs

@@ -666,6 +666,8 @@ impl From<tokenspeed_proto::SchedulerLoad> for openai_protocol::worker::Schedule
             dp_rank: load.dp_rank,
             num_running_reqs: load.num_running_reqs,
             num_waiting_reqs: load.num_waiting_reqs,
+            // TokenSpeed does not report queued token-work; degrade to 0.
+            num_waiting_uncached_tokens: 0,
             num_total_reqs: load.num_total_reqs,
             num_used_tokens: load.num_used_tokens,
             max_total_num_tokens: load.max_total_num_tokens,

crates/protocols/src/worker.rs

@@ -1017,6 +1017,9 @@ pub struct SchedulerLoadSnapshot {
     pub dp_rank: i32,
     pub num_running_reqs: i32,
     pub num_waiting_reqs: i32,
+    /// Queued token-work: waiting-queue tokens not yet served from cache. 0 when
+    /// the backend does not report it — callers degrade gracefully.
+    pub num_waiting_uncached_tokens: i32,
     pub num_total_reqs: i32,
     pub num_used_tokens: i32,
     pub max_total_num_tokens: i32,
@@ -1051,6 +1054,19 @@ impl WorkerLoadResponse {
         self.loads.iter().map(|l| l.num_used_tokens as i64).sum()
     }
 
+    /// Total queued (waiting, uncached) tokens summed across all DP ranks.
+    pub fn total_waiting_uncached_tokens(&self) -> i64 {
+        self.loads
+            .iter()
+            .map(|l| l.num_waiting_uncached_tokens as i64)
+            .sum()
+    }
+
+    /// Total generation throughput (tokens/s) summed across all DP ranks.
+    pub fn total_gen_throughput(&self) -> f64 {
+        self.loads.iter().map(|l| l.gen_throughput).sum()
+    }
+
     pub fn dp_rank_loads(&self) -> HashMap<isize, isize> {
         let mut map = HashMap::new();
         for snapshot in &self.loads {

grpc_servicer/smg_grpc_servicer/sglang/servicer.py

@@ -90,6 +90,8 @@ def _convert_loads_to_protobuf(
         cache_hit_rate=result.cache_hit_rate,
         utilization=result.utilization,
         max_running_requests=result.max_running_requests,
+        # Queued token-work: waiting-queue tokens not served from cache.
+        num_waiting_uncached_tokens=result.num_waiting_uncached_tokens,
     )
 
     # Add optional sections using CopyFrom for proper protobuf assignment

model_gateway/src/config/types.rs

@@ -312,17 +312,27 @@ pub enum PolicyConfig {
     #[serde(rename = "power_of_two")]
     PowerOfTwo { load_check_interval_secs: u64 },
 
-    /// Least-load policy: routes to the worker minimizing
-    /// `in_flight + kv_pressure_weight * k/(1-k)` — the real-time in-flight request count plus
-    /// a convex KV-cache pressure term from the load monitor.
-    /// See `policies/least_load.rs`.
+    /// Least-(token-)work policy: routes to the worker minimizing the expected
+    /// wait `(queued_tokens + inflight_tokens) / throughput + kv_pressure_weight * k/(1-k)`
+    /// — token-work drain time plus a convex KV-cache pressure barrier, computed
+    /// from the load monitor with in-flight correction. See `policies/least_load.rs`.
     #[serde(rename = "least_load")]
     LeastLoad {
         #[serde(default = "default_least_load_interval")]
         load_check_interval_secs: u64,
-        /// KV-pressure weight (request-equivalents per unit of M/M/1 congestion).
+        /// KV-pressure weight `λ_t` (seconds): the time-cost of KV contention,
+        /// commensurate with the expected-queue-wait term.
         #[serde(default = "default_least_load_kv_pressure_weight")]
         kv_pressure_weight: f64,
+        /// Mean prefill length (tokens) used to estimate in-flight token-work
+        /// when a request's token count is unknown at routing time.
+        #[serde(default = "default_least_load_mean_prefill")]
+        mean_prefill_tokens: u32,
+        /// Fallback generation throughput (tokens/s) for the expected-wait term
+        /// when a backend reports no live `gen_throughput`. Set to the fleet's
+        /// per-replica generation rate; co-tunes with `kv_pressure_weight`.
+        #[serde(default = "default_least_load_throughput")]
+        default_throughput: f64,
     },
 
     #[serde(rename = "bucket")]
@@ -402,7 +412,15 @@ fn default_least_load_interval() -> u64 {
 }
 
 fn default_least_load_kv_pressure_weight() -> f64 {
-    1.5
+    0.15
+}
+
+fn default_least_load_mean_prefill() -> u32 {
+    1024
+}
+
+fn default_least_load_throughput() -> f64 {
+    2000.0
 }
 
 impl PolicyConfig {

model_gateway/src/config/validation.rs

@@ -303,6 +303,8 @@ impl ConfigValidator {
             PolicyConfig::LeastLoad {
                 load_check_interval_secs,
                 kv_pressure_weight,
+                mean_prefill_tokens,
+                default_throughput,
             } => {
                 if *load_check_interval_secs == 0 {
                     return Err(ConfigError::InvalidValue {
@@ -319,6 +321,22 @@ impl ConfigValidator {
                         reason: "Must be finite and >= 0.0".to_string(),
                     });
                 }
+
+                if *mean_prefill_tokens == 0 {
+                    return Err(ConfigError::InvalidValue {
+                        field: "mean_prefill_tokens".to_string(),
+                        value: mean_prefill_tokens.to_string(),
+                        reason: "Must be > 0".to_string(),
+                    });
+                }
+
+                if !default_throughput.is_finite() || *default_throughput <= 0.0 {
+                    return Err(ConfigError::InvalidValue {
+                        field: "default_throughput".to_string(),
+                        value: default_throughput.to_string(),
+                        reason: "Must be finite and > 0.0".to_string(),
+                    });
+                }
             }
             PolicyConfig::Bucket {
                 balance_abs_threshold: _,

model_gateway/src/main.rs

@@ -192,6 +192,20 @@ struct CliArgs {
     #[arg(long, default_value_t = 1.25, help_heading = "Routing Policy")]
     prefix_hash_load_factor: f64,
 
+    /// KV-pressure weight (seconds) for the least_load policy
+    #[arg(long, default_value_t = 0.15, help_heading = "Routing Policy")]
+    least_load_kv_pressure_weight: f64,
+
+    /// Fallback generation throughput (tokens/s) for least_load when a backend
+    /// reports no live throughput
+    #[arg(long, default_value_t = 2000.0, help_heading = "Routing Policy")]
+    least_load_default_throughput: f64,
+
+    /// Mean prefill tokens for least_load's in-flight estimate when a request's
+    /// token count is unknown at routing
+    #[arg(long, default_value_t = 1024, help_heading = "Routing Policy")]
+    least_load_mean_prefill_tokens: u32,
+
     /// Enable data parallelism aware scheduling
     #[arg(long, default_value_t = false, help_heading = "Routing Policy")]
     dp_aware: bool,
@@ -931,7 +945,9 @@ impl CliArgs {
             },
             "least_load" => PolicyConfig::LeastLoad {
                 load_check_interval_secs: 5,
-                kv_pressure_weight: 1.5,
+                kv_pressure_weight: self.least_load_kv_pressure_weight,
+                mean_prefill_tokens: self.least_load_mean_prefill_tokens,
+                default_throughput: self.least_load_default_throughput,
             },
             "prefix_hash" => PolicyConfig::PrefixHash {
                 prefix_token_count: self.prefix_token_count,

model_gateway/src/policies/factory.rs

@@ -20,9 +20,14 @@ impl PolicyFactory {
             PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()),
             PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()),
             PolicyConfig::LeastLoad {
-                kv_pressure_weight, ..
-            } => Arc::new(LeastLoadPolicy::with_kv_pressure_weight(
+                kv_pressure_weight,
+                mean_prefill_tokens,
+                default_throughput,
+                ..
+            } => Arc::new(LeastLoadPolicy::with_params(
                 *kv_pressure_weight,
+                *mean_prefill_tokens,
+                *default_throughput,
             )),
             PolicyConfig::CacheAware {
                 cache_threshold,