Charge scheduler host handoff only where the queue drains

docs/scheduling-runtime-design.md

@@ -83,7 +83,11 @@ never the proposals or allocation of the completed program.
 Costs use the full graph, including proposal work, allocation, verification,
 commit, collectives, and readiness stalls. GPU replay and serial host handoff
 have separate estimates; queued completion wall time does not become an isolated
-handoff sample. A trace attaches physical B/D/W/P, selected cost, observation
+handoff sample. Handoff is exposed only when the next submission waits for a
+completion: the median handoff of a graph's drained executions is charged in
+proportion to its ordinary executions that require completion feedback. Host
+delays around pipelined executions and probes that drain only to measure do not
+change a graph's cost. A trace attaches physical B/D/W/P, selected cost, observation
 age, selection reason, and exploration credit to each committed plan.
 `flare_speculative_*` counters cover complete speculative programs rather than
 separate producer and target calls.

helios/README.md

@@ -337,8 +337,11 @@ lane's logical token limit even when the completion reports physical padding.
 update it with the median of up to nine replay observations. After 30 seconds without a valid observation,
 the cost becomes stale but remains available for selection. The next valid
 observation replaces the stale estimate. Tensor-parallel observations use the
-maximum positive rank interval. Isolated executions also measure host handoff;
-queued executions exclude queue wait from the timing updates. Unmeasured
+maximum positive rank interval. Executions that drain the queue also measure host
+handoff, which is charged in proportion to the graph's ordinary executions that
+must deliver completion feedback before another submission. Pipelined executions
+charge no handoff, and queued executions exclude queue wait from the timing
+updates. Unmeasured
 candidates compete by useful work and can beat measured candidates. Both
 unmeasured and stale physical dispatches require an empty queue and completion
 before another submission.

helios/src/scheduler.rs

@@ -563,6 +563,7 @@ impl ScheduledPlan {
 
 struct Submission {
     step: ScheduledStep,
+    delivers_feedback: bool,
     expected_finish_ns: u64,
     estimated_duration_ns: u64,
     participants: [Option<service::ServiceParticipant>; REQUEST_LANES],
@@ -1026,7 +1027,7 @@ impl RequestScheduler {
         step: &ScheduledStep,
         timing: crate::execution::ExecutionTiming,
     ) {
-        self.costs.observe(step, timing);
+        self.costs.observe(step, timing, Some(self.step_delivers_feedback(step)));
     }
 
     pub fn submitted_timing(&self, fence: Fence) -> Result<SubmittedTiming> {
@@ -1051,16 +1052,22 @@ impl RequestScheduler {
     }
 
     pub fn plan_requires_feedback(&self, plan: &ScheduledPlan) -> bool {
-        let geometry = (plan.kind, plan.batch_capacity, plan.draft_tokens, plan.token_capacity,
-            plan.packed_token_capacity, plan.prefill_capacity());
-        plan.requires_feedback || !matches!(plan.kind, GraphKind::Decode | GraphKind::Speculative)
+        plan.requires_feedback || self.step_delivers_feedback(plan)
+            || !self.costs.is_fresh(cost::Key::from_step(plan))
+            || plan.origin.is_probe()
+    }
+
+    /// Whether the host must consume this step's completion before another submission
+    /// for reasons other than measuring its cost.
+    fn step_delivers_feedback(&self, step: &ScheduledStep) -> bool {
+        let geometry = (step.kind, step.batch_capacity, step.draft_tokens, step.token_capacity,
+            step.packed_token_capacity, step.prefill_capacity());
+        !matches!(step.kind, GraphKind::Decode | GraphKind::Speculative)
             || self.capabilities.completion_inputs.iter().any(|(key, inputs)| *key == geometry && !inputs.is_empty())
-            || plan.lanes.iter().any(|lane| self.request(lane.request).is_ok_and(|request|
+            || step.lanes.iter().any(|lane| self.request(lane.request).is_ok_and(|request|
                 request.grammar_constrained || request.host_stops
                     || self.proposal_source == speculation::ProposalSource::Copy
                     || self.decoded_history.is_some()))
-            || !self.costs.is_fresh(cost::Key::from_step(plan))
-            || plan.origin.is_probe()
     }
 
     fn selected_cost(&self, key: cost::Key) -> SelectedCost {
@@ -2133,6 +2140,7 @@ impl RequestScheduler {
             }
         }
         let feedback = self.plan_requires_feedback(plan);
+        let delivers_feedback = self.step_delivers_feedback(step);
         if self.submissions.len() >= 2
             || (feedback && !self.submissions.is_empty())
             || self.submissions.values().any(|step| step.requires_feedback)
@@ -2223,6 +2231,7 @@ impl RequestScheduler {
             fence,
             Submission {
                 step: ScheduledStep { requires_feedback: feedback, ..step.clone() },
+                delivers_feedback,
                 expected_finish_ns,
                 estimated_duration_ns: plan.estimated_duration_ns,
                 participants,
@@ -2440,7 +2449,8 @@ impl RequestScheduler {
         self.refresh.validate_completion(step.origin, timing)?;
         let submission = self.submissions.remove(&fence).unwrap();
         self.refresh.complete(submission.origin, timing);
-        self.observe_execution(&submission.step, timing);
+        self.costs.observe(&submission.step, timing,
+            (!submission.origin.is_probe()).then_some(submission.delivers_feedback));
         let step = submission.step;
         let now_ns = self.now_ns;
         let proposal_source = self.proposal_source;

helios/src/scheduler/cost.rs

@@ -50,6 +50,7 @@ pub(super) struct PackedRoundShape {
 struct Estimate {
     nanos: u64,
     replays: ReplayTimes,
+    handoff: HostHandoff,
     handoff_ns: u64,
     observed_at_ns: u64,
     last_attempt_ns: u64,
@@ -81,6 +82,39 @@ impl ReplayTimes {
     }
 }
 
+/// Host time a graph leaves the GPU idle between executions. A pipelined execution
+/// overlaps its handoff with the queued one, and host delays around it belong to the
+/// host rather than to the graph. Only an execution that must deliver completion
+/// feedback before the next submission drains the queue and exposes its handoff, so
+/// the median drained handoff is charged in proportion to such ordinary executions.
+#[derive(Default)]
+struct HostHandoff {
+    drained: ReplayTimes,
+    drained_ns: u64,
+    feedback: u16,
+    executions: u32,
+}
+
+impl HostHandoff {
+    /// `drained_ns` is the handoff of an execution that ran with an empty queue
+    /// because it required feedback. `feedback` records whether an ordinary
+    /// execution required it; probes drain the queue only to measure and omit it.
+    fn observe(&mut self, drained_ns: Option<u64>, feedback: Option<bool>) -> u64 {
+        if let Some(ns) = drained_ns {
+            self.drained_ns = self.drained.observe(ns, true);
+        }
+        if let Some(feedback) = feedback {
+            self.feedback = self.feedback << 1 | u16::from(feedback);
+            self.executions = (self.executions + 1).min(u16::BITS);
+        }
+        if self.executions == 0 {
+            return 0;
+        }
+        let executions = u128::from(self.executions);
+        ((u128::from(self.drained_ns) * u128::from(self.feedback.count_ones()) + executions / 2) / executions) as u64
+    }
+}
+
 impl Estimate {
     fn fresh(&self, now_ns: u64, ttl_ns: u64) -> bool {
         self.nanos != 0 && now_ns.saturating_sub(self.observed_at_ns) < ttl_ns
@@ -114,6 +148,7 @@ impl ActionCosts {
         let estimate = &mut self.entries[self.index[&key]].1;
         estimate.nanos = nanos;
         estimate.replays.observe(nanos, false);
+        estimate.handoff = HostHandoff::default();
         estimate.handoff_ns = 0;
         estimate.observed_at_ns = self.now_ns;
         self.refresh_packed_durations();
@@ -139,6 +174,7 @@ impl ActionCosts {
                 Estimate {
                     nanos: 0,
                     replays: ReplayTimes::default(),
+                    handoff: HostHandoff::default(),
                     handoff_ns: 0,
                     observed_at_ns: 0,
                     last_attempt_ns: 0,
@@ -246,7 +282,9 @@ impl ActionCosts {
         self.duration_key(Key::from_step(step))
     }
 
-    pub fn observe(&mut self, step: &ScheduledStep, timing: ExecutionTiming) {
+    /// `feedback` tells whether an ordinary execution had to deliver completion
+    /// feedback before the next submission; it is absent for probes.
+    pub fn observe(&mut self, step: &ScheduledStep, timing: ExecutionTiming, feedback: Option<bool>) {
         let Some(raw_sample) = timing.gpu_ns.into_iter().flatten().max().filter(|n| *n > 0) else {
             return;
         };
@@ -265,14 +303,12 @@ impl ActionCosts {
         if let Some(&index) = self.index.get(&key) {
             let estimate = &mut self.entries[index].1;
             let fresh = estimate.fresh(self.now_ns, self.ttl_ns);
-            if !timing.queued {
-                let handoff = timing.wall_ns.saturating_sub(raw_sample);
-                estimate.handoff_ns = if fresh {
-                    ((7 * u128::from(estimate.handoff_ns) + u128::from(handoff) + 4) / 8) as u64
-                } else { handoff };
-            } else if !fresh {
-                estimate.handoff_ns = 0;
+            if !fresh {
+                estimate.handoff = HostHandoff::default();
             }
+            let drained = (step.requires_feedback && !timing.queued)
+                .then(|| timing.wall_ns.saturating_sub(raw_sample));
+            estimate.handoff_ns = estimate.handoff.observe(drained, feedback);
             estimate.nanos = estimate.replays.observe(sample, fresh);
             estimate.observed_at_ns = self.now_ns;
             self.refresh_packed_durations();
@@ -282,7 +318,7 @@ impl ActionCosts {
 
 #[cfg(test)]
 mod tests {
-    use super::ReplayTimes;
+    use super::{HostHandoff, ReplayTimes};
 
     #[test]
     fn isolated_slow_replay_does_not_displace_the_steady_cost() {
@@ -310,6 +346,32 @@ mod tests {
         assert_eq!(replays.observe(16, true), 24);
     }
 
+    #[test]
+    fn pipelined_executions_expose_no_handoff() {
+        let mut handoff = HostHandoff::default();
+        assert_eq!(handoff.observe(Some(90_000), None), 0);
+        for _ in 0..16 {
+            assert_eq!(handoff.observe(None, Some(false)), 0);
+        }
+    }
+
+    #[test]
+    fn feedback_executions_charge_their_median_drained_handoff() {
+        let mut handoff = HostHandoff::default();
+        for _ in 0..8 {
+            assert_eq!(handoff.observe(Some(100_000), Some(true)), 100_000);
+        }
+        assert_eq!(handoff.observe(Some(30_000_000), Some(true)), 100_000);
+        for _ in 0..8 {
+            handoff.observe(None, Some(false));
+        }
+        assert_eq!(handoff.observe(None, Some(false)), 43_750);
+        for _ in 0..7 {
+            handoff.observe(None, Some(false));
+        }
+        assert_eq!(handoff.observe(None, None), 0);
+    }
+
     #[test]
     fn even_sample_median_does_not_overflow() {
         let mut replays = ReplayTimes::default();

helios/src/scheduler/tests.rs

@@ -43,6 +43,7 @@ pub(super) fn measured_costs(graphs: &[crate::execution::GraphDescriptor]) -> co
                 ],
                 ..Default::default()
             },
+            None,
         );
     }
     costs
@@ -1731,6 +1732,7 @@ fn graph_duration_median_uses_runtime_observations_and_excludes_queue_wait() {
     s.admit(spec(65, 8)).unwrap();
     assert!(s.measured_graph_durations(&graphs).is_err());
     let plan = s.plan().unwrap();
+    assert_eq!(plan.kind, GraphKind::Init);
     assert_eq!(s.estimated_duration_ns(&plan), Some(1000));
     for queued in [false, true] {
         s.observe_execution(
@@ -1760,11 +1762,14 @@ fn graph_duration_median_uses_runtime_observations_and_excludes_queue_wait() {
             },
         );
     }
-    assert_eq!(s.estimated_duration_ns(&plan), Some(14075));
+    // This state operation must deliver its completion before the next submission,
+    // so the 98200 ns handoff of its drained execution is exposed. The queued
+    // execution adds only its GPU interval; its wall time includes the queue wait.
+    assert_eq!(s.estimated_duration_ns(&plan), Some(1800 + 98_200));
     let mut same_graph = (*plan).clone();
     same_graph.lanes[0].start_position = 32769;
     same_graph.lanes[0].tokens = 1;
-    assert_eq!(s.estimated_duration_ns(&same_graph), Some(14075));
+    assert_eq!(s.estimated_duration_ns(&same_graph), Some(1800 + 98_200));
     same_graph.batch_capacity = 2;
     assert_eq!(s.estimated_duration_ns(&same_graph), Some(1000));
     s.observe_execution(
@@ -1776,27 +1781,48 @@ fn graph_duration_median_uses_runtime_observations_and_excludes_queue_wait() {
             queued: false,
         },
     );
-    assert_eq!(s.estimated_duration_ns(&plan), Some(14075));
+    assert_eq!(s.estimated_duration_ns(&plan), Some(1800 + 98_200));
+}
+
+#[test]
+fn host_delays_around_pipelined_steps_do_not_change_their_cost() {
+    use crate::execution::ExecutionTiming;
+    let mut costs = cost::ActionCosts::prepared(&duration_graphs(1)).unwrap();
+    let step = ScheduledStep { draft_tokens: 0, kind: GraphKind::Decode, batch_capacity: 1, token_capacity: 1,
+        packed_token_capacity: 0, mixed: None, lanes: vec![], requires_feedback: false };
+    let mut gpu_ns = [None; 16];
+    gpu_ns[0] = Some(30_000);
+    let timing = |wall_ns, queued| ExecutionTiming { gpu_ns, readiness: None, wall_ns, queued };
+    // A refresh drains the queue to measure the step and records its handoff.
+    costs.observe(&ScheduledStep { requires_feedback: true, ..step.clone() }, timing(30_100, false), None);
+    assert_eq!(costs.get(&step), Some(30_000));
+    // A late host leaves the following ordinary steps without a queued successor.
+    for _ in 0..5 {
+        costs.observe(&step, timing(60_000, false), Some(false));
+    }
+    assert_eq!(costs.get(&step), Some(30_000));
+    costs.observe(&step, timing(60_000, true), Some(false));
+    assert_eq!(costs.get(&step), Some(30_000));
 }
 
 #[test]
 fn readiness_correction_changes_shape_cost_without_changing_handoff_or_service_time() {
     use crate::execution::{ExecutionTiming, ReadinessIntervals};
     let mut costs = cost::ActionCosts::prepared(&duration_graphs(1)).unwrap();
     let step = ScheduledStep { draft_tokens: 0, kind: GraphKind::Decode, batch_capacity: 1, token_capacity: 1,
-        packed_token_capacity: 0, mixed: None, lanes: vec![], requires_feedback: false };
+        packed_token_capacity: 0, mixed: None, lanes: vec![], requires_feedback: true };
     let timestamps = [1000u64, 3000, 3100].into_iter().flat_map(u64::to_ne_bytes).collect::<Vec<_>>();
     let mut gpu_ns = [None; 16];
     gpu_ns[0] = Some(10_000);
     gpu_ns[1] = Some(9_000);
     let mut timing = ExecutionTiming { gpu_ns, readiness: Some(ReadinessIntervals::decode(&timestamps).unwrap()),
         wall_ns: 12_000, queued: false };
-    costs.observe(&step, timing);
+    costs.observe(&step, timing, Some(true));
     assert_eq!(costs.get(&step), Some(10_000));
     assert_eq!(cost::observed_duration_ns(timing), Some(12_000));
     timing.queued = true;
     timing.wall_ns = 100_000;
-    costs.observe(&step, timing);
+    costs.observe(&step, timing, Some(true));
     assert_eq!(costs.get(&step), Some(10_000));
     assert_eq!(cost::observed_duration_ns(timing), Some(10_000));
     assert_eq!(timing.gpu_ns[0], Some(10_000));

helios/src/server/runtime/phase_tests.rs

@@ -1109,7 +1109,10 @@ fn worker_observes_prefill_decode_and_prerequisite_timings() {
         ];
         frame.timing.wall_ns = 350;
         worker.complete(kind, frame).unwrap();
-        assert_eq!(worker.scheduler.estimated_duration_ns(&step), Some(619));
+        // The median of the prepared 1000 ns and the observed 200 ns, plus the 150 ns
+        // handoff of steps whose completion must be consumed before the next submission.
+        let handoff = if kind == GraphKind::Decode { 0 } else { 150 };
+        assert_eq!(worker.scheduler.estimated_duration_ns(&step), Some(600 + handoff));
         let metrics = fixture.metrics.snapshot(false);
         assert_eq!(metrics.verification_drafts, 0);
         assert_eq!(
@@ -1600,7 +1603,9 @@ fn speculative_dashboard_counts_live_rounds_and_batched_forward_passes() {
         let source = fixture.config.speculative_proposal_source;
         assert!(profile.prefix_probability(source, 0).unwrap() > 0.5);
         assert!(profile.prefix_probability(source, 2).unwrap() < 0.125);
-        assert_eq!(worker.scheduler.estimated_duration_ns(&step), Some(1063));
+        // Copied proposals need each completion before the next round and expose its
+        // 500 ns handoff; MTP rounds overlap it with the queued round.
+        assert_eq!(worker.scheduler.estimated_duration_ns(&step), Some(if copy { 1500 } else { 1000 }));
         let metrics = fixture.metrics.snapshot(!copy);
         assert_eq!((metrics.prefill_tokens, metrics.generated_tokens), (13, 4));
         assert_eq!((metrics.accepted_drafts, metrics.draft_rounds), (2, 1));