perf(router): reduce scheduler and dispatch overhead

lib/kv-router/benches/worker_selection.rs

@@ -154,8 +154,10 @@ fn fixture_with_preferred_taints(
     const PREFERRED_TAINTS: [&str; 4] = ["rack-a", "zone-a", "gpu-a", "node-a"];
     assert!(preferred_taint_count <= PREFERRED_TAINTS.len());
     let mut workers = HashMap::with_capacity(worker_count);
-    let mut effective_overlap_blocks = FxHashMap::default();
-    let mut effective_cached_tokens = FxHashMap::default();
+    let mut effective_overlap_blocks =
+        FxHashMap::with_capacity_and_hasher(worker_count, Default::default());
+    let mut effective_cached_tokens =
+        FxHashMap::with_capacity_and_hasher(worker_count, Default::default());
     let mut worker_loads = FxHashMap::with_capacity_and_hasher(worker_count, Default::default());
 
     for worker_id in 0..worker_count as WorkerId {
@@ -234,7 +236,7 @@ fn worker_selection(c: &mut Criterion) {
             group.measurement_time(Duration::from_secs(5));
             group.sample_size(50);
 
-            for worker_count in [2, 32, 1_024, 10_000] {
+            for worker_count in [2, 32, 1_024, 2_048, 10_000] {
                 let (workers, request) = fixture(worker_count);
                 let selector = DefaultWorkerSelector::new(
                     Some(KvRouterConfig {
@@ -425,11 +427,49 @@ fn default_policy_wrapper(c: &mut Criterion) {
     group.finish();
 }
 
+fn worker_projection(c: &mut Criterion) {
+    let mut group = c.benchmark_group("worker_projection");
+    group.warm_up_time(Duration::from_secs(2));
+    group.measurement_time(Duration::from_secs(5));
+    group.sample_size(50);
+    for count in [32, 2_048] {
+        let (workers, request) = fixture(count);
+        let slots = dynamo_kv_router::ActiveSequencesMultiWorker::new(
+            dynamo_kv_router::NoopSequencePublisher,
+            16,
+            workers.keys().map(|&id| (id, (0, 1))).collect(),
+            false,
+            0,
+            "bench",
+        );
+        let now = tokio::time::Instant::now();
+        group.bench_function(BenchmarkId::new("allocate", count), |b| {
+            b.iter(|| black_box(slots.project_worker_loads(request.token_seq.as_deref(), now)))
+        });
+        #[cfg(feature = "bench")]
+        {
+            let mut projections = FxHashMap::default();
+            group.bench_function(BenchmarkId::new("reuse", count), |b| {
+                b.iter(|| {
+                    slots.bench_project_worker_loads_into(
+                        request.token_seq.as_deref(),
+                        now,
+                        &mut projections,
+                    );
+                    black_box(&projections);
+                })
+            });
+        }
+    }
+    group.finish();
+}
+
 criterion_group!(
     benches,
     worker_selection,
     custom_worker_selection,
     unused_preferred_taint_metadata,
-    default_policy_wrapper
+    default_policy_wrapper,
+    worker_projection
 );
 criterion_main!(benches);

lib/kv-router/src/scheduling/filter.rs

@@ -322,7 +322,7 @@ impl<'a> RoutingEligibility<'a> {
     }
 }
 
-fn worker_config_for_rank<C: WorkerConfigLike>(
+pub(super) fn worker_config_for_rank<C: WorkerConfigLike>(
     workers: &HashMap<WorkerId, C>,
     worker: WorkerWithDpRank,
 ) -> Result<&C, WorkerEligibilityError> {

lib/kv-router/src/scheduling/local.rs

@@ -623,6 +623,7 @@ where
         }))
     }
 
+    /// Apply completion and notify pending admission before returning.
     pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> {
         let request_id = request_id.to_string();
         let worker = self.slots.request_worker(&request_id);
@@ -634,14 +635,12 @@ where
         }
         self.slots.publish_prefill_completed(&request_id);
         if outcome.is_applied() {
-            match worker {
-                Some(worker) => self.queue.update_worker(worker).await,
-                None => self.queue.update().await,
-            }
+            self.queue.capacity_changed(worker);
         }
         Ok(())
     }
 
+    /// Release state and notify pending admission before returning.
     pub async fn free(&self, request_id: &str) -> Result<(), SequenceError> {
         let request_id = request_id.to_string();
         let worker = self.slots.request_worker(&request_id);
@@ -650,10 +649,7 @@ where
             return Err(SequenceError::RequestNotFound { request_id });
         }
         if outcome.is_applied() {
-            match worker {
-                Some(worker) => self.queue.update_worker(worker).await,
-                None => self.queue.update().await,
-            }
+            self.queue.capacity_changed(worker);
         }
         Ok(())
     }
@@ -672,7 +668,7 @@ where
             .slots
             .free_if_worker(&request_id, worker, Instant::now())?;
         if outcome.is_applied() {
-            self.queue.update_worker(worker).await;
+            self.queue.capacity_changed(Some(worker));
         }
         Ok(())
     }
@@ -698,7 +694,7 @@ where
     }
 
     /// Complete owner cleanup after a successful release (including a stale
-    /// booking's `NoChange`), before the cancellable queue-progress wait.
+    /// booking's `NoChange`), then notify pending admission before returning.
     pub(crate) async fn free_if_booking_with_cleanup(
         &self,
         booking: &SchedulerBookingDescriptor,
@@ -712,7 +708,7 @@ where
         )?;
         cleanup();
         if outcome.is_applied() {
-            self.queue.update_worker(booking.worker).await;
+            self.queue.capacity_changed(Some(booking.worker));
         }
         Ok(outcome)
     }
@@ -723,15 +719,27 @@ where
     }
 
     /// `NoChange` when the booking no longer matches or its prefill was
-    /// already marked complete.
+    /// already marked complete. Success confirms the state mutation and capacity
+    /// notification, not completion of pending admission.
     #[doc(hidden)]
     pub async fn mark_prefill_completed_if_booking(
         &self,
         booking: &SchedulerBookingDescriptor,
     ) -> Result<LifecycleMutationOutcome, KvSchedulerError> {
-        self.queue
-            .mark_prefill_completed_if_booking(booking.clone())
-            .await
+        self.queue.ensure_running()?;
+        let outcome = self
+            .slots
+            .mark_prefill_completed_if_booking(
+                &booking.request_id,
+                booking.worker,
+                booking.attempt_id,
+                Instant::now(),
+            )
+            .map_err(|error| KvSchedulerError::BookingFailed(error.to_string()))?;
+        if outcome.is_applied() {
+            self.queue.capacity_changed(Some(booking.worker));
+        }
+        Ok(outcome)
     }
 
     /// Republish the ordered prefill-completion event while `booking` is live.
@@ -743,6 +751,13 @@ where
         self.slots.publish_prefill_completed_if_booking(booking)
     }
 
+    /// Wait for an admission recheck when queueing is enabled and the actor is running.
+    /// This does not wait for blocked requests to finish or for the queue to empty.
+    /// With queueing disabled, return immediately without an admission barrier.
+    pub async fn update_queue(&self) {
+        self.queue.update().await;
+    }
+
     pub fn pending_count(&self) -> usize {
         self.queue.pending_count()
     }
@@ -778,9 +793,10 @@ where
         booking: &SchedulerBookingDescriptor,
         decay_fraction: Option<f64>,
     ) -> Result<(), KvSchedulerError> {
-        self.queue
-            .add_output_block_if_booking(booking.clone(), decay_fraction)
-            .await
+        self.queue.ensure_running()?;
+        self.add_output_block_if_booking_sync(booking, decay_fraction)
+            .map(|_| ())
+            .map_err(|error| KvSchedulerError::BookingFailed(error.to_string()))
     }
 
     /// `add_output_block_if_booking` applied inline, like `add_output_block`,
@@ -799,14 +815,14 @@ where
         )
     }
 
+    /// Apply an output update before returning, without waiting for admission.
     #[doc(hidden)]
     pub async fn enqueue_output_block_if_booking(
         &self,
         booking: &SchedulerBookingDescriptor,
         decay_fraction: Option<f64>,
     ) -> Result<(), KvSchedulerError> {
-        self.queue
-            .enqueue_output_block_if_booking(booking.clone(), decay_fraction)
+        self.add_output_block_if_booking(booking, decay_fraction)
             .await
     }
 
@@ -892,6 +908,70 @@ mod tests {
         duration: Duration,
     }
 
+    fn poll_once<F: std::future::Future>(future: F) -> std::task::Poll<F::Output> {
+        std::pin::pin!(future).poll(&mut std::task::Context::from_waker(std::task::Waker::noop()))
+    }
+
+    #[tokio::test]
+    async fn lifecycle_updates_do_not_wait_for_the_actor_in_either_queueing_mode() {
+        let worker = WorkerWithDpRank::new(0, 0);
+        for threshold in [None, Some(0.5)] {
+            let (scheduler, slots, _configs, cancellation) = make_scheduler(
+                HashMap::from([(0, SimpleWorkerConfig::default())]),
+                threshold,
+                false,
+                None,
+            );
+            let booking = scheduler
+                .add_request_if_registered_guarded(SequenceRequest {
+                    request_id: "output".into(),
+                    token_sequence: None,
+                    track_prefill_tokens: true,
+                    expected_output_tokens: None,
+                    prefill_load_hint: Some(crate::protocols::PrefillLoadHint {
+                        initial_effective_prefill_tokens: 64,
+                        expected_prefill_duration: None,
+                    }),
+                    worker,
+                    lora_name: None,
+                })
+                .unwrap()
+                .commit();
+            let before = slots.active_blocks()[&worker];
+            // A single poll on this current-thread runtime cannot run the actor.
+            assert!(matches!(
+                poll_once(scheduler.enqueue_output_block_if_booking(&booking, None)),
+                std::task::Poll::Ready(Ok(()))
+            ));
+            assert!(slots.active_blocks()[&worker] > before);
+            assert!(matches!(
+                poll_once(scheduler.mark_prefill_completed_if_booking(&booking)),
+                std::task::Poll::Ready(Ok(LifecycleMutationOutcome::Applied))
+            ));
+            assert_eq!(slots.active_tokens(Instant::now())[&worker], 0);
+            assert!(matches!(
+                poll_once(scheduler.mark_prefill_completed_if_booking(&booking)),
+                std::task::Poll::Ready(Ok(LifecycleMutationOutcome::NoChange))
+            ));
+            assert!(matches!(
+                poll_once(scheduler.free_if_booking(&booking)),
+                std::task::Poll::Ready(Ok(LifecycleMutationOutcome::Applied))
+            ));
+            assert!(matches!(
+                poll_once(scheduler.add_output_block_if_booking(&booking, None)),
+                std::task::Poll::Ready(Ok(()))
+            ));
+            slots.assert_completely_drained(Instant::now());
+            if threshold.is_some() {
+                assert!(poll_once(scheduler.update_queue()).is_pending());
+                scheduler.update_queue().await;
+            } else {
+                assert!(poll_once(scheduler.update_queue()).is_ready());
+            }
+            cancellation.cancel();
+        }
+    }
+
     impl PrefillLoadEstimator for FixedPrefillLoadEstimator {
         fn predict_prefill_duration(
             &self,