fix(router): account for all output blocks in response chunks

lib/llm/src/kv_router/routing_host/request_guard.rs

@@ -373,6 +373,7 @@ impl RequestObservability {
 }
 
 struct OutputBlockUpdate {
+    num_blocks: usize,
     decay_fraction: Option<f64>,
 }
 
@@ -547,12 +548,16 @@ impl OutputBlockTracker {
             return None;
         }
 
+        let num_blocks = new_total_blocks - self.current_total_blocks;
         // Advance before returning so a failed scheduler update preserves existing no-retry behavior.
         self.current_total_blocks = new_total_blocks;
         let decay_fraction = self
             .expected_output_tokens
             .map(|expected| (1.0 - cumulative_osl as f64 / expected.max(1) as f64).max(0.0));
-        Some(OutputBlockUpdate { decay_fraction })
+        Some(OutputBlockUpdate {
+            num_blocks,
+            decay_fraction,
+        })
     }
 }
 
@@ -573,6 +578,16 @@ pub(super) struct RequestGuard {
 }
 
 impl RequestGuard {
+    #[cfg(test)]
+    pub(super) fn booking_for_test(
+        &self,
+    ) -> &dynamo_kv_router::scheduling::queue::SchedulerBookingDescriptor {
+        self.cleanup
+            .lifecycle()
+            .expect("expected a tracked booking")
+            .booking()
+    }
+
     pub(super) fn new_kv(
         chooser: Arc<KvRouter>,
         request_metrics: Arc<RouterRequestMetrics>,
@@ -803,16 +818,21 @@ impl RequestGuard {
         if let RequestCleanup::Kv(cleanup) = &self.cleanup
             && let Some(lifecycle) = cleanup.lifecycle()
             && lifecycle.is_active()
-            && let Err(error) = cleanup
-                .chooser
-                .enqueue_output_block_if_booking(lifecycle.booking(), update.decay_fraction)
-                .await
         {
-            tracing::warn!(
-                request_id = %cleanup.context_id,
-                %error,
-                "Failed to add output block"
-            );
+            for _ in 0..update.num_blocks {
+                if let Err(error) = cleanup
+                    .chooser
+                    .enqueue_output_block_if_booking(lifecycle.booking(), update.decay_fraction)
+                    .await
+                {
+                    tracing::warn!(
+                        request_id = %cleanup.context_id,
+                        %error,
+                        "Failed to add output block"
+                    );
+                    break;
+                }
+            }
         }
 
         self.observability.observe_output_block_boundary();
@@ -861,6 +881,57 @@ impl Drop for RequestGuard {
     }
 }
 
+#[cfg(test)]
+mod output_block_tests {
+    use super::OutputBlockTracker;
+
+    #[test]
+    fn counts_every_crossed_boundary_once() {
+        let mut tracker = OutputBlockTracker::new(true, 16, 16, None);
+        assert!(tracker.observe(0).is_none());
+        let first = tracker.observe(1).unwrap();
+        assert_eq!(first.num_blocks, 1);
+        assert_eq!(first.decay_fraction, None);
+        assert!(tracker.observe(16).is_none());
+        assert_eq!(tracker.observe(33).unwrap().num_blocks, 2);
+        assert!(tracker.observe(33).is_none());
+        assert!(tracker.observe(48).is_none());
+        assert_eq!(tracker.observe(49).unwrap().num_blocks, 1);
+    }
+
+    #[test]
+    fn accounts_for_a_partial_prompt_block() {
+        let mut tracker = OutputBlockTracker::new(true, 15, 16, None);
+        assert!(tracker.observe(1).is_none());
+        assert_eq!(tracker.observe(34).unwrap().num_blocks, 3);
+        assert!(tracker.observe(49).is_none());
+    }
+
+    #[test]
+    fn preserves_decay_at_the_observed_output_length() {
+        let mut tracker = OutputBlockTracker::new(true, 16, 16, Some(64));
+        let update = tracker.observe(32).unwrap();
+        assert_eq!(update.num_blocks, 2);
+        assert_eq!(update.decay_fraction, Some(0.5));
+        let update = tracker.observe(64).unwrap();
+        assert_eq!(update.num_blocks, 2);
+        assert_eq!(update.decay_fraction, Some(0.0));
+        let update = tracker.observe(96).unwrap();
+        assert_eq!(update.num_blocks, 2);
+        assert_eq!(update.decay_fraction, Some(0.0));
+
+        let mut zero_expected = OutputBlockTracker::new(true, 16, 16, Some(0));
+        assert_eq!(zero_expected.observe(32).unwrap().decay_fraction, Some(0.0));
+    }
+
+    #[test]
+    fn disabled_tracking_emits_no_updates() {
+        let mut tracker = OutputBlockTracker::new(false, 16, 16, None);
+        assert!(tracker.observe(48).is_none());
+        assert_eq!(tracker.current_total_blocks, 1);
+    }
+}
+
 #[cfg(test)]
 mod output_hash_tests {
     use super::*;

lib/llm/src/kv_router/routing_host/tests.rs

@@ -1144,6 +1144,116 @@ async fn stream_failure_releases_booking_before_error_is_observable() {
     runtime.shutdown();
 }
 
+#[tokio::test]
+#[serial_test::serial]
+async fn output_block_accounting_tracks_grouped_chunks() {
+    for (track_output_blocks, expected_output_blocks) in [(true, 3), (false, 0)] {
+        let config = KvRouterConfig {
+            skip_initial_worker_wait: true,
+            use_kv_events: false,
+            router_track_active_blocks: true,
+            router_track_output_blocks: track_output_blocks,
+            ..Default::default()
+        };
+        let (router, runtime) = router_with_config(
+            None,
+            HashMap::from([(7, ModelRuntimeConfig::default())]),
+            config,
+        )
+        .await;
+        let prompt_tokens = 16_usize;
+        let chunks = [1_usize, 32, 15];
+        let mut input = request();
+        input.token_ids = (1..=prompt_tokens as u32).collect::<Vec<_>>().into();
+        let input = Context::new(input);
+        let (mut selection, _) = router
+            .select_with_affinity(
+                &input,
+                RequestPhase::Aggregated,
+                false,
+                &CleanupBudget::default(),
+            )
+            .await
+            .unwrap();
+        let mut guard = router
+            .track_selection(
+                &input,
+                &mut selection,
+                RequestPhase::Aggregated,
+                false,
+                &CleanupBudget::default(),
+            )
+            .await
+            .unwrap();
+        let mut next_token = prompt_tokens as u32 + 1;
+        let initial_loads = router
+            .kv_router()
+            .get_potential_loads(&[], None, None, None, None)
+            .await
+            .unwrap();
+        let initial_blocks = initial_loads
+            .iter()
+            .find(|load| load.worker_id == 7 && load.dp_rank == 0)
+            .unwrap()
+            .potential_decode_blocks;
+        for size in &chunks {
+            let token_ids = (next_token..next_token + *size as u32).collect();
+            next_token += *size as u32;
+            guard
+                .on_item(&Annotated::from_data(LLMEngineOutput {
+                    token_ids,
+                    index: Some(0),
+                    ..Default::default()
+                }))
+                .await;
+        }
+        // Output updates are enqueued without waiting for their application.
+        // This idempotent, acknowledged command on the same actor is a FIFO
+        // barrier, so the load observation cannot race the output updates.
+        router
+            .kv_router()
+            .mark_prefill_completed_if_booking(guard.booking_for_test())
+            .await
+            .unwrap();
+        let loads = router
+            .kv_router()
+            .get_potential_loads(&[], None, None, None, None)
+            .await
+            .unwrap();
+        let load = loads
+            .iter()
+            .find(|load| load.worker_id == 7 && load.dp_rank == 0)
+            .unwrap();
+        assert_eq!(load.active_requests, 1);
+        // Keep the scheduler's existing prompt accounting unchanged;
+        // this regression checks only the growth caused by output.
+        let expected_blocks = initial_blocks + expected_output_blocks;
+        println!(
+            "track_output_blocks={track_output_blocks} prompt_tokens={prompt_tokens} output_tokens=48 block_size=16 chunks={chunks:?} initial_blocks={initial_blocks} observed_blocks={} expected_blocks={expected_blocks}",
+            load.potential_decode_blocks
+        );
+        let observed_blocks = load.potential_decode_blocks;
+        guard.finish().await;
+        let loads = router
+            .kv_router()
+            .get_potential_loads(&[], None, None, None, None)
+            .await
+            .unwrap();
+        let load = loads
+            .iter()
+            .find(|load| load.worker_id == 7 && load.dp_rank == 0)
+            .unwrap();
+        assert_eq!(load.active_requests, 0);
+        assert_eq!(load.potential_decode_blocks, 0);
+        drop(router);
+        runtime.shutdown();
+        assert_eq!(
+            observed_blocks, expected_blocks,
+            "incorrect grouped-output accounting with tracking={track_output_blocks}"
+        );
+    }
+}
+
 async fn router(session_affinity_ttl: Option<Duration>) -> (RoutingHost, Runtime) {
     router_with_workers(session_affinity_ttl, &[7]).await
 }
@@ -1189,6 +1299,41 @@ async fn router_with_worker_configs_and_classifier(
     session_affinity_ttl: Option<Duration>,
     workers: HashMap<u64, ModelRuntimeConfig>,
     classifier: Option<impl RequestClassifier>,
+) -> (RoutingHost, Runtime) {
+    let config = KvRouterConfig {
+        skip_initial_worker_wait: true,
+        use_kv_events: false,
+        router_track_active_blocks: false,
+        ..Default::default()
+    };
+    Box::pin(router_with_config_and_classifier(
+        session_affinity_ttl,
+        workers,
+        config,
+        classifier,
+    ))
+    .await
+}
+
+async fn router_with_config(
+    session_affinity_ttl: Option<Duration>,
+    workers: HashMap<u64, ModelRuntimeConfig>,
+    config: KvRouterConfig,
+) -> (RoutingHost, Runtime) {
+    Box::pin(router_with_config_and_classifier(
+        session_affinity_ttl,
+        workers,
+        config,
+        None::<RecordingClassifier>,
+    ))
+    .await
+}
+
+async fn router_with_config_and_classifier(
+    session_affinity_ttl: Option<Duration>,
+    workers: HashMap<u64, ModelRuntimeConfig>,
+    config: KvRouterConfig,
+    classifier: Option<impl RequestClassifier>,
 ) -> (RoutingHost, Runtime) {
     let runtime = Runtime::from_current().unwrap();
     // Each runtime has its own in-memory discovery store, so fixture namespaces can repeat.
@@ -1204,12 +1349,6 @@ async fn router_with_worker_configs_and_classifier(
     let client = endpoint.client().await.unwrap();
     let worker_ids = workers.keys().copied().collect::<Vec<_>>();
     let (_tx, workers) = watch::channel(workers);
-    let config = KvRouterConfig {
-        skip_initial_worker_wait: true,
-        use_kv_events: false,
-        router_track_active_blocks: false,
-        ..Default::default()
-    };
     let mut chooser = KvRouter::new(
         endpoint,
         client.clone(),