diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md index 1bf73f077663..fe837f0d2774 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/configuration-and-tuning.md @@ -201,17 +201,21 @@ every eligible worker is busy for that class, but a new arrival cannot bypass an existing backlog in the same class. Queue limits are configured per discovered worker endpoint with -`request_queue_limit_per_worker`, `raw_isl_token_queue_limit_per_worker`, and -`cached_token_queue_limit_per_worker`. The effective class-local limit is the -configured value multiplied by the current number of discovered endpoints. +`request_queue_limit_per_worker`, `raw_isl_token_queue_limit_per_worker`, +`cached_token_queue_limit_per_worker`, and +`uncached_token_queue_limit_per_worker`. The effective class-local limit is +each configured value multiplied by the current number of discovered endpoints. Limits are checked against current usage before adding the incoming request, so the request that crosses a limit is accepted and the next queued request is -rejected with HTTP 529 and the effective total. Worker removal does not evict +rejected with HTTP 503 and the effective total. Worker removal does not evict queued requests; new arrivals reject until usage drains or capacity returns. -DRR charges the uncached-token snapshot captured at enqueue, while raw, cached, -and uncached snapshots remain unchanged for limits, WSPT, counters, and later -dispatch. For the ring cursor, deficit charging, weighted bursts, and bounded -bulk-credit behavior, see +The uncached-token limit counts only the queued ISL portion that still needs +prefill compute (raw ISL minus the best cached tokens across eligible +workers), so cache-heavy requests do not consume the queue budget that +protects prefill capacity. DRR charges the uncached-token snapshot captured +at enqueue, while raw, cached, and uncached snapshots remain unchanged for +limits, WSPT, counters, and later dispatch. For the ring cursor, deficit +charging, weighted bursts, and bounded bulk-credit behavior, see [Deficit Round Robin Queue Scheduling](deficit-round-robin.md). Every matrix class must identify both `policy_family` and `cache_bucket`; a diff --git a/lib/kv-router/src/scheduling/policy_config.rs b/lib/kv-router/src/scheduling/policy_config.rs index 6d65f37a7237..3c80810ac5fd 100644 --- a/lib/kv-router/src/scheduling/policy_config.rs +++ b/lib/kv-router/src/scheduling/policy_config.rs @@ -42,6 +42,7 @@ pub struct PolicyClassConfig { pub request_queue_limit_per_worker: Option, pub raw_isl_token_queue_limit_per_worker: Option, pub cached_token_queue_limit_per_worker: Option, + pub uncached_token_queue_limit_per_worker: Option, } impl PolicyClassConfig { @@ -123,6 +124,7 @@ impl PolicyProfile { request_queue_limit_per_worker: None, raw_isl_token_queue_limit_per_worker: None, cached_token_queue_limit_per_worker: None, + uncached_token_queue_limit_per_worker: None, }; Self { classes: vec![class], @@ -316,6 +318,7 @@ struct RawPolicyClassConfig { request_queue_limit_per_worker: Option, raw_isl_token_queue_limit_per_worker: Option, cached_token_queue_limit_per_worker: Option, + uncached_token_queue_limit_per_worker: Option, } fn resolve_profile( @@ -504,6 +507,7 @@ fn resolve_policy_class( request_queue_limit_per_worker: raw.request_queue_limit_per_worker, raw_isl_token_queue_limit_per_worker: raw.raw_isl_token_queue_limit_per_worker, cached_token_queue_limit_per_worker: raw.cached_token_queue_limit_per_worker, + uncached_token_queue_limit_per_worker: raw.uncached_token_queue_limit_per_worker, }, binding, }) @@ -774,6 +778,53 @@ models: assert_eq!(unmatched.default_class().prefill_busy_threshold_frac, None); } + #[test] + fn uncached_token_queue_limit_parses_and_defaults_to_none() { + let config = RouterPolicyConfig::from_yaml( + r#" +default_policy_family: standard +uncached_isl_buckets: + - min_tokens: 0 + bucket: all +policy_classes: + - name: capped + policy_family: standard + cache_bucket: all + quantum: 1 + uncached_token_queue_limit_per_worker: 666667 + - name: uncapped + policy_family: uncapped + cache_bucket: all + quantum: 1 +"#, + ) + .unwrap(); + + let profile = config.resolve_profile(None, None, RouterQueuePolicy::Fcfs); + assert_eq!( + profile + .class(profile.resolve_class_index(Some("capped"), usize::MAX)) + .uncached_token_queue_limit_per_worker, + Some(666667) + ); + // Classes that omit the knob keep the disabled-by-default contract and + // existing YAML stays accepted without the field. + assert_eq!( + profile + .class(profile.resolve_class_index(Some("uncapped"), usize::MAX)) + .uncached_token_queue_limit_per_worker, + None + ); + + let synthetic = PolicyProfile::synthetic(None, RouterQueuePolicy::Fcfs); + assert_eq!( + synthetic + .default_class() + .uncached_token_queue_limit_per_worker, + None + ); + } + #[test] fn rootless_model_config_falls_back_for_unmatched_model() { let config = RouterPolicyConfig::from_yaml( diff --git a/lib/kv-router/src/scheduling/policy_queue.rs b/lib/kv-router/src/scheduling/policy_queue.rs index 1e51dcaae561..f609be333d5f 100644 --- a/lib/kv-router/src/scheduling/policy_queue.rs +++ b/lib/kv-router/src/scheduling/policy_queue.rs @@ -41,6 +41,7 @@ pub enum QueueLimitKind { Requests, RawIslTokens, CachedTokens, + UncachedTokens, } impl std::fmt::Display for QueueLimitKind { @@ -49,6 +50,7 @@ impl std::fmt::Display for QueueLimitKind { Self::Requests => formatter.write_str("requests"), Self::RawIslTokens => formatter.write_str("raw_isl_tokens"), Self::CachedTokens => formatter.write_str("cached_tokens"), + Self::UncachedTokens => formatter.write_str("uncached_tokens"), } } } @@ -70,6 +72,7 @@ pub struct PolicyQueueStats { pub requests: usize, pub raw_isl_tokens: usize, pub cached_tokens: usize, + pub uncached_tokens: usize, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -799,6 +802,11 @@ fn queue_rejection(class: &PolicyClassQueue, worker_count: usize) -> Optio class.stats.cached_tokens, class.config.cached_token_queue_limit_per_worker, ), + ( + QueueLimitKind::UncachedTokens, + class.stats.uncached_tokens, + class.config.uncached_token_queue_limit_per_worker, + ), ] { let limit = limit_per_worker.map(|limit| limit.saturating_mul(worker_count)); if limit.is_some_and(|limit| current >= limit) { @@ -818,12 +826,18 @@ fn add_stats(stats: &mut PolicyQueueStats, snapshot: QueueSnapshot) { stats.requests += 1; stats.raw_isl_tokens = stats.raw_isl_tokens.saturating_add(snapshot.raw_isl_tokens); stats.cached_tokens = stats.cached_tokens.saturating_add(snapshot.cached_tokens); + stats.uncached_tokens = stats + .uncached_tokens + .saturating_add(snapshot.uncached_tokens); } fn subtract_stats(stats: &mut PolicyQueueStats, snapshot: QueueSnapshot) { stats.requests = stats.requests.saturating_sub(1); stats.raw_isl_tokens = stats.raw_isl_tokens.saturating_sub(snapshot.raw_isl_tokens); stats.cached_tokens = stats.cached_tokens.saturating_sub(snapshot.cached_tokens); + stats.uncached_tokens = stats + .uncached_tokens + .saturating_sub(snapshot.uncached_tokens); } #[cfg(test)] @@ -1350,6 +1364,228 @@ policy_classes: assert_eq!(no_workers_rejection.limit, 0); } + #[test] + fn uncached_token_cap_charges_uncached_work_only() { + let mut queue = PolicyQueue::new(profile( + r#" +default_policy_family: capped +uncached_isl_buckets: + - min_tokens: 0 + bucket: all +policy_classes: + - name: capped + policy_family: capped + cache_bucket: all + quantum: 10 + uncached_token_queue_limit_per_worker: 3 +"#, + )); + queue + .enqueue( + 0, + 2, + QueueSnapshot::new(10, 8), + 0.0, + 0.0, + 0, + WorkerPlacement::Any, + "uncached-queued", + ) + .unwrap(); + // Cache-rich work consumes none of the uncached budget. + queue + .enqueue( + 0, + 2, + QueueSnapshot::new(100, 100), + 1.0, + 0.0, + 0, + WorkerPlacement::Any, + "cache-rich", + ) + .unwrap(); + assert_eq!(queue.class_stats(0).uncached_tokens, 2); + + // Pre-add: the entry crossing the cap is admitted... + queue + .enqueue( + 0, + 2, + QueueSnapshot::new(4, 0), + 2.0, + 0.0, + 0, + WorkerPlacement::Any, + "crossing", + ) + .unwrap(); + assert_eq!(queue.class_stats(0).uncached_tokens, 6); + + // ...and the next arrival is rejected against existing usage only. + let (rejection, payload) = queue + .enqueue( + 0, + 2, + QueueSnapshot::new(1, 0), + 3.0, + 0.0, + 0, + WorkerPlacement::Any, + "rejected", + ) + .unwrap_err(); + assert_eq!(payload, "rejected"); + assert_eq!(rejection.limit_kind, QueueLimitKind::UncachedTokens); + assert_eq!(rejection.current, 6); + assert_eq!(rejection.limit, 6); + assert_eq!(queue.class_stats(0).uncached_tokens, 6); + + // Draining the earliest entry returns its 2 uncached tokens, so the + // same request shape is admitted again. + assert_eq!( + queue.pop_next(|_, _, _| true).unwrap().into_payload(), + "uncached-queued" + ); + assert_eq!(queue.class_stats(0).uncached_tokens, 4); + queue + .enqueue( + 0, + 2, + QueueSnapshot::new(1, 0), + 4.0, + 0.0, + 0, + WorkerPlacement::Any, + "readmitted", + ) + .unwrap(); + assert_eq!(queue.class_stats(0).uncached_tokens, 5); + } + + #[test] + fn uncached_token_cap_scales_with_worker_count() { + let mut queue = PolicyQueue::new(profile( + r#" +default_policy_family: capped +uncached_isl_buckets: + - min_tokens: 0 + bucket: all +policy_classes: + - name: capped + policy_family: capped + cache_bucket: all + quantum: 10 + uncached_token_queue_limit_per_worker: 5 +"#, + )); + queue + .enqueue( + 0, + 2, + QueueSnapshot::new(10, 0), + 0.0, + 0.0, + 0, + WorkerPlacement::Any, + "at-cap", + ) + .unwrap(); + let (shrunk, _) = queue + .enqueue( + 0, + 1, + QueueSnapshot::new(1, 0), + 1.0, + 0.0, + 0, + WorkerPlacement::Any, + "shrunk", + ) + .unwrap_err(); + assert_eq!(shrunk.limit_kind, QueueLimitKind::UncachedTokens); + assert_eq!(shrunk.current, 10); + assert_eq!(shrunk.limit, 5); + + // One more discovered worker grows the cap without evicting, and the + // crossing entry is still admitted under the pre-add comparison. + queue + .enqueue( + 0, + 3, + QueueSnapshot::new(4, 0), + 2.0, + 0.0, + 0, + WorkerPlacement::Any, + "grown-cap", + ) + .unwrap(); + queue + .enqueue( + 0, + 3, + QueueSnapshot::new(1, 0), + 3.0, + 0.0, + 0, + WorkerPlacement::Any, + "crossing-grown", + ) + .unwrap(); + assert_eq!(queue.class_stats(0).uncached_tokens, 15); + let (at_grown, _) = queue + .enqueue( + 0, + 3, + QueueSnapshot::new(1, 0), + 4.0, + 0.0, + 0, + WorkerPlacement::Any, + "at-grown", + ) + .unwrap_err(); + assert_eq!(at_grown.limit_kind, QueueLimitKind::UncachedTokens); + assert_eq!(at_grown.current, 15); + assert_eq!(at_grown.limit, 15); + } + + #[test] + fn uncached_token_zero_limit_sheds_immediately() { + let mut queue = PolicyQueue::new(profile( + r#" +default_policy_family: shed +uncached_isl_buckets: + - min_tokens: 0 + bucket: all +policy_classes: + - name: shed + policy_family: shed + cache_bucket: all + quantum: 10 + uncached_token_queue_limit_per_worker: 0 +"#, + )); + let (rejection, payload) = queue + .enqueue( + 0, + 4, + QueueSnapshot::new(1, 0), + 0.0, + 0.0, + 0, + WorkerPlacement::Any, + "shed", + ) + .unwrap_err(); + assert_eq!(payload, "shed"); + assert_eq!(rejection.limit_kind, QueueLimitKind::UncachedTokens); + assert_eq!(rejection.current, 0); + assert_eq!(rejection.limit, 0); + assert_eq!(queue.pending_count(), 0); + } + #[test] fn per_worker_limit_multiplication_saturates() { let mut queue = PolicyQueue::new(profile(&format!( diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index f48840008412..e0076beb8a85 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -3251,6 +3251,130 @@ policy_classes: assert_eq!(queue.pending_count(), 2); } + #[tokio::test(flavor = "multi_thread")] + async fn class_local_uncached_limit_rejects_by_pending_compute() { + let profile = policy_profile( + r#" +default_policy_family: compute-capped +uncached_isl_buckets: + - min_tokens: 0 + bucket: all +policy_classes: + - name: compute-capped + policy_family: compute-capped + cache_bucket: all + quantum: 1 + prefill_busy_threshold: 0 + uncached_token_queue_limit_per_worker: 100 +"#, + ); + let (queue, _slots) = make_queue_with_profile(1, 16, 64, profile); + + let (active, active_rx) = make_request("active", 64); + queue.enqueue(active).await; + active_rx.await.unwrap().unwrap(); + + let (first, _first_rx) = make_request("first", 64); + queue.enqueue(first).await; + + let (second, _second_rx) = make_request("second", 64); + queue.enqueue(second).await; + assert_eq!(queue.pending_count(), 2); + + let (rejected, rejected_rx) = make_request("rejected", 64); + queue.enqueue(rejected).await; + let error = rejected_rx.await.unwrap().unwrap_err(); + let KvSchedulerError::QueueRejected(rejection) = &error else { + panic!("expected queue rejection, got {error:?}"); + }; + assert_eq!(rejection.policy_class, "compute-capped"); + assert_eq!( + rejection.limit_kind, + super::super::QueueLimitKind::UncachedTokens + ); + assert_eq!(rejection.current, 128); + assert_eq!(rejection.limit, 100); + assert!(!error.is_overload()); + assert_eq!( + rejection.to_string(), + "router policy class \"compute-capped\" queue uncached_tokens limit reached \ + (current=128, limit=100)" + ); + + assert_eq!( + queue.class_queue_stats(0), + Some(ClassQueueStats { + pending_count: 2, + pending_isl_tokens: 128, + pending_cached_tokens: 0, + }) + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn class_local_uncached_limit_accounts_cached_overlap() { + let profile = policy_profile( + r#" +default_policy_family: compute-capped +uncached_isl_buckets: + - min_tokens: 0 + bucket: all +policy_classes: + - name: compute-capped + policy_family: compute-capped + cache_bucket: all + quantum: 1 + prefill_busy_threshold: 0 + uncached_token_queue_limit_per_worker: 100 +"#, + ); + let (queue, _slots) = make_queue_with_profile(1, 16, 64, profile); + + let (active, active_rx) = make_request("active", 100); + queue.enqueue(active).await; + active_rx.await.unwrap().unwrap(); + + // Two queued requests with 100 ISL but 50 cached tokens each: only + // 50 uncached tokens per request are charged. Under raw-ISL + // accounting the second enqueue would already be rejected. + let half_overlap = OverlapSignals { + effective_cached_tokens: HashMap::from([(WorkerWithDpRank::new(0, 0), 50)]), + ..OverlapSignals::default() + }; + let (mut first, _first_rx) = make_request("first", 100); + first.overlap = half_overlap.clone(); + queue.enqueue(first).await; + + let (mut second, _second_rx) = make_request("second", 100); + second.overlap = half_overlap; + queue.enqueue(second).await; + assert_eq!(queue.pending_count(), 2); + + let (rejected, rejected_rx) = make_request("rejected", 100); + queue.enqueue(rejected).await; + let error = rejected_rx.await.unwrap().unwrap_err(); + let KvSchedulerError::QueueRejected(rejection) = &error else { + panic!("expected queue rejection, got {error:?}"); + }; + assert_eq!( + rejection.limit_kind, + super::super::QueueLimitKind::UncachedTokens + ); + // 100 queued uncached tokens (2 x 50), not the 200 raw ISL tokens. + assert_eq!(rejection.current, 100); + assert_eq!(rejection.limit, 100); + + // Raw and cached gauges still charge the full prompt picture. + assert_eq!( + queue.class_queue_stats(0), + Some(ClassQueueStats { + pending_count: 2, + pending_isl_tokens: 200, + pending_cached_tokens: 100, + }) + ); + } + #[tokio::test(start_paused = true)] async fn test_queue_update_uses_decayed_oldest_prefill_load() { let estimator: Arc = Arc::new(FixedPrefillLoadEstimator { diff --git a/lib/llm/src/kv_router/embedded.rs b/lib/llm/src/kv_router/embedded.rs index e60b18eb4e19..14127f7fe162 100644 --- a/lib/llm/src/kv_router/embedded.rs +++ b/lib/llm/src/kv_router/embedded.rs @@ -84,6 +84,7 @@ fn record_queue_rejection( QueueLimitKind::Requests => handles.request_limit_rejections.inc(), QueueLimitKind::RawIslTokens => handles.raw_isl_limit_rejections.inc(), QueueLimitKind::CachedTokens => handles.cached_token_limit_rejections.inc(), + QueueLimitKind::UncachedTokens => handles.uncached_token_limit_rejections.inc(), } } diff --git a/lib/llm/src/kv_router/metrics.rs b/lib/llm/src/kv_router/metrics.rs index 02c584de6c62..6bcd7879acf1 100644 --- a/lib/llm/src/kv_router/metrics.rs +++ b/lib/llm/src/kv_router/metrics.rs @@ -702,6 +702,7 @@ pub struct RouterQueueMetricHandles { pub request_limit_rejections: IntCounter, pub raw_isl_limit_rejections: IntCounter, pub cached_token_limit_rejections: IntCounter, + pub uncached_token_limit_rejections: IntCounter, } pub static ROUTER_QUEUE_METRICS: LazyLock = @@ -766,6 +767,7 @@ impl RouterQueueMetrics { request_limit_rejections: rejection("request_limit"), raw_isl_limit_rejections: rejection("raw_isl_token_limit"), cached_token_limit_rejections: rejection("cached_token_limit"), + uncached_token_limit_rejections: rejection("uncached_token_limit"), } } } @@ -1587,6 +1589,7 @@ dynamo_frontend_worker_active_prefill_tokens{dp_rank=\"0\",worker_id=\"123\",wor dynamo_frontend_router_queue_backpressure_total{model=\"model\",policy_class=\"default\",reason=\"cached_token_limit\",worker_type=\"decode\"} 0 dynamo_frontend_router_queue_backpressure_total{model=\"model\",policy_class=\"default\",reason=\"raw_isl_token_limit\",worker_type=\"decode\"} 0 dynamo_frontend_router_queue_backpressure_total{model=\"model\",policy_class=\"default\",reason=\"request_limit\",worker_type=\"decode\"} 0 +dynamo_frontend_router_queue_backpressure_total{model=\"model\",policy_class=\"default\",reason=\"uncached_token_limit\",worker_type=\"decode\"} 0 # HELP dynamo_frontend_router_queue_pending_cached_tokens Estimated cached tokens for requests pending in the router scheduler queue # TYPE dynamo_frontend_router_queue_pending_cached_tokens gauge dynamo_frontend_router_queue_pending_cached_tokens{model=\"model\",policy_class=\"default\",worker_type=\"decode\"} 512