Skip to content

Commit 44ec239

Browse files
voipmonitorclaudelukealonso
authored
Remove repeated work from scheduler planning (#188)
* Remove repeated work from scheduler planning Funded acceptance learning compared every probe candidate with the ordinary winner slot by slot, and each comparison of a packed round re-estimated the whole round's prefix lengths, including lane values, while the allocation loop re-evaluated every lane's acceptance curve for each extra row. At concurrency 32 one plan call took 10-35 ms whenever acceptance credit was full. The worker does not poll completions while it plans, so a finished round waited for the plan before its successor or a draining probe could start: on DeepSeek V4.1 TP4 every few seconds, and for about 0.7 s after 32 requests arrived. Planning also repeated linear work per lane and per shape: - Each speculative-token bound scanned every lane of the submitted steps to find the request's pending advance. The ceiling is now tracked on the request where commit and completion already track its pending count. - Packed rounds rebuilt rankings, gains, and per-lane curves for every shape and priorities for every ordering. Shapes sharing draft and width bounds now share one ranking, lane curves are evaluated once per call, priorities once per shape, and only the leading cohort of each ordering is sorted. - Padding pruning compared all candidate pairs and re-estimated round lengths for each pair. Only candidates with the same kind, cohort, and prefill row are compared, and each candidate's lengths are estimated once. - Service rates and next-progress values evaluated the forecast curve for every batch and width, and every speculative duration repeated a hashed lookup of the full graph, which the packed round table already includes. Selection, round lengths, and their arithmetic are unchanged. Seeded scheduler traces over 1056 scenarios, with 1-64 requests, host and device prefix allocation, fixed widths, mixed graphs, grammar and host-stop lanes, queued rounds, retirements, and timing and acceptance probes, record every plan, lane value, pruned candidate list, probe choice, and round length: about 210 000 plans compare byte for byte with main; the first quarter also in a debug build. Median plan time with one round queued (release test build): requests allocation main after acceptance probe: main after 8 device 140 us 60 us 2563 us 111 us 16 device 289 us 96 us 2153 us 201 us 32 device 669 us 196 us 5312 us 413 us 64 device 1773 us 436 us 13380 us 943 us 32 host 754 us 250 us 1363 us 497 us 64 host 1957 us 615 us 3546 us 1215 us DeepSeek V4.1 Flash TP4 on switched GPUs 0-3 of frank1 and frank2, concurrency 32, seven drafts, verification budget at most 128, alpha 0.6, Engram tables in RAM, with the exposed-handoff cost change in both builds (tokens/s, mean of ten 30 s cells, worst cell in parentheses): host/package before after frank1 main 2471 (2428) 2524 (2490) frank1 B 2474 (2432) 2529 (2495) frank2 main 2478 (2398) 2537 (2503) frank2 B 2450 (2313) 2540 (2515) all 40 cells 2468 2532 (+2.6%) In a diagnostic run of the same service, the GPU busy share of the steady window rose from 98.4-98.9% to 99.8-99.9%, and no worker iteration exceeded 10 ms. Validation: Helios library tests (runtime-package and server features) pass in release and debug builds except five frontend-template, package-cache, and image-hint tests that fail identically on main for missing local assets. Debug builds check the tracked pending ceiling against the submitted lanes on every use. A new test covers queued rounds, completion, and retirement. Limitations: while a winning probe waits for the queue to drain, the worker still plans again on every iteration; each plan is now short. Signed-off-by: Martin Vit <martin@voipmonitor.org> Design-Signed-Off-By: Claude (Opus 5.5); design-review=1c01 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Reuse scheduling cursors and rank speculative prefixes Derive pending advance from the scheduled and consumed cursors, retaining submission scans for completion reconciliation. Stream acceptance candidates into winner selection without storing all candidate length arrays. Rank monotone prefix survival once with depth and slot tie-breaking. Share the ranking with host cohort allocation and preserve request limits, fixed widths, and GPU-owned allocation ceilings. Validation: 346 scheduler and runtime phase tests pass in both debug and release builds; two tests are ignored by default. Differential tests cover host/GPU allocation, MTP/copy proposals, ties, tails, and cohort growth. A release CPU benchmark measures planning with one queued round and funded or unfunded acceptance probes. GPU serving measurements were not run. Design-Signed-Off-By: Codex (GPT-6); design-review=1c01 Signed-off-by: Luke Alonso <lalonso@gmail.com> --------- Signed-off-by: Martin Vit <martin@voipmonitor.org> Signed-off-by: Luke Alonso <lalonso@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Luke Alonso <lalonso@gmail.com>
1 parent cc696a6 commit 44ec239

7 files changed

Lines changed: 391 additions & 80 deletions

File tree

‎helios/src/scheduler.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2087,7 +2087,7 @@ impl RequestScheduler {
20872087
fn speculative_tokens(&self, request: &LiveRequest, width: usize) -> usize {
20882088
let context = self.limits.max_sequence_tokens - request.scheduled_consumed as usize;
20892089
let output = (request.max_new_tokens - request.generated_tokens) as usize;
2090-
let output = output.saturating_sub(self.pending_token_ceiling(request.handle));
2090+
let output = output.saturating_sub(self.pending_ceiling(request));
20912091
let proposal = if request.host_stops
20922092
|| (request.grammar_constrained
20932093
&& self.proposal_source == speculation::ProposalSource::Mtp)

‎helios/src/scheduler/learning.rs‎

Lines changed: 28 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -18,40 +18,44 @@ impl RequestScheduler {
1818
self.packed_round_candidates(&mut candidates, values, age_price, true);
1919
self.mixed_candidates(&mut candidates, values, age_price, true);
2020
let count = self.prune_padding(&mut candidates);
21-
let observed_width = |candidate: Candidate, slot: usize, optimistic: bool| {
22-
if candidate.mask & (1 << slot) == 0 || candidate.prefill_slot == Some(slot)
23-
|| !matches!(candidate.kind, GraphKind::Speculative | GraphKind::Mixed) {
24-
1
25-
} else {
26-
if let Some(lengths) = candidate.host_lengths {
21+
// Round lengths depend only on the candidate, so each is estimated once.
22+
let widths = |candidate: Candidate, optimistic: bool| -> [usize; REQUEST_LANES] {
23+
let speculative = matches!(candidate.kind, GraphKind::Speculative | GraphKind::Mixed);
24+
let estimated = (speculative && candidate.host_lengths.is_none() && candidate.packed_token_capacity != 0)
25+
.then(|| self.prefix_lengths_with_values(candidate.mask, candidate.batch, candidate.draft_tokens,
26+
candidate.capacity, candidate.packed_token_capacity, false, values, optimistic));
27+
std::array::from_fn(|slot| {
28+
if candidate.mask & (1 << slot) == 0 || candidate.prefill_slot == Some(slot) || !speculative {
29+
1
30+
} else if let Some(lengths) = candidate.host_lengths {
2731
usize::from(lengths[slot])
28-
} else if candidate.packed_token_capacity != 0 {
29-
self.estimated_round_lengths(candidate.mask, candidate.batch, candidate.draft_tokens,
30-
candidate.capacity, candidate.packed_token_capacity, optimistic)[slot]
32+
} else if let Some(lengths) = estimated {
33+
lengths[slot]
3134
} else {
3235
self.speculative_tokens(self.slots[slot].request.as_ref().unwrap(), candidate.capacity)
3336
}
34-
}
37+
})
3538
};
36-
candidates[..count].iter().copied().filter(|candidate| {
37-
candidate.rate.duration_ns != 0 && candidate.rate.compare(ordinary.rate).is_gt()
38-
&& bits(candidate.mask).any(|slot| {
39-
let request = self.slots[slot].request.as_ref().unwrap();
40-
let width = observed_width(*candidate, slot, true);
41-
width > observed_width(ordinary, slot, false)
42-
&& request.max_new_tokens - request.generated_tokens > width as u32
43-
&& self.limits.max_sequence_tokens - request.scheduled_consumed as usize > width
44-
})
45-
}).max_by(|a, b| {
39+
let ordinary_widths = widths(ordinary, false);
40+
candidates[..count].iter().copied()
41+
.filter(|candidate| candidate.rate.duration_ns != 0 && candidate.rate.compare(ordinary.rate).is_gt())
42+
.map(|candidate| (candidate, widths(candidate, true)))
43+
.filter(|(candidate, widths)| bits(candidate.mask).any(|slot| {
44+
let request = self.slots[slot].request.as_ref().unwrap();
45+
let width = widths[slot];
46+
width > ordinary_widths[slot]
47+
&& request.max_new_tokens - request.generated_tokens > width as u32
48+
&& self.limits.max_sequence_tokens - request.scheduled_consumed as usize > width
49+
})).max_by(|(a, a_widths), (b, b_widths)| {
4650
let gain = |candidate: &Candidate| (candidate.rate.score - ordinary.rate.score) / candidate.rate.duration_ns as f64;
47-
let information = |candidate: &Candidate| bits(candidate.mask).flat_map(|slot| {
48-
(observed_width(ordinary, slot, false) - 1..observed_width(*candidate, slot, true) - 1)
51+
let information = |candidate: &Candidate, widths: &[usize; REQUEST_LANES]| bits(candidate.mask).flat_map(|slot| {
52+
(ordinary_widths[slot] - 1..widths[slot] - 1)
4953
.map(move |depth| self.speculative_information(self.slots[slot].request.as_ref().unwrap(), depth))
5054
}).fold(f64::INFINITY, f64::min);
5155
gain(a).total_cmp(&gain(b))
52-
.then_with(|| information(b).total_cmp(&information(a)))
56+
.then_with(|| information(b, b_widths).total_cmp(&information(a, a_widths)))
5357
.then(b.rate.duration_ns.cmp(&a.rate.duration_ns))
5458
.then(b.batch.cmp(&a.batch)).then(b.capacity.cmp(&a.capacity)).then(b.mask.cmp(&a.mask))
55-
})
59+
}).map(|(candidate, _)| candidate)
5660
}
5761
}
Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
/*
2+
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
use super::*;
7+
use crate::execution::PrefixAllocation;
8+
9+
fn planning_scheduler(active: usize, allocation: PrefixAllocation) -> RequestScheduler {
10+
let mut scheduler = engine(active, true);
11+
let mut graphs = super::super::tests::duration_graphs(active);
12+
let rectangular = graphs.clone();
13+
for graph in rectangular.iter().filter(|g| g.kind == GraphKind::Speculative) {
14+
for rows in [1, 2, 4] {
15+
let packed = graph.batch_capacity * rows;
16+
if packed >= graph.packed_token_capacity { continue; }
17+
let mut graph = graph.clone();
18+
graph.name = format!("{}-p{packed}", graph.name);
19+
graph.packed_token_capacity = packed;
20+
graphs.push(graph);
21+
}
22+
}
23+
for graph in &mut graphs {
24+
graph.estimated_duration_ns = if graph.kind == GraphKind::Speculative {
25+
100_000 + graph.batch_capacity as u64 * 10_000
26+
+ graph.draft_tokens as u64 * 15_000 + graph.packed_token_capacity as u64 * 2_000
27+
} else { MS };
28+
if graph.kind == GraphKind::Speculative {
29+
graph.confidence_width = if allocation == PrefixAllocation::Gpu { graph.draft_tokens } else { 0 };
30+
scheduler.capabilities.speculative_allocations.push((graph.batch_capacity,
31+
graph.draft_tokens, graph.token_capacity, graph.packed_token_capacity, allocation));
32+
}
33+
}
34+
scheduler.measured_graph_durations(&graphs).unwrap();
35+
for id in 0..active {
36+
let handle = add(&mut scheduler, id as u32, 64, 256);
37+
let request = scheduler.request_mut(handle).unwrap();
38+
request.phase = RequestPhase::Decode;
39+
request.consumed_tokens = 80;
40+
request.scheduled_consumed = 80;
41+
request.generated_tokens = 17;
42+
for sample in 0..8 {
43+
let accepted = (id + sample) % 8;
44+
request.acceptance.observe(speculation::ProposalSource::Mtp, 17,
45+
7, accepted, accepted < 7, None);
46+
}
47+
}
48+
scheduler
49+
}
50+
51+
fn greedy_lengths(scheduler: &RequestScheduler, mask: u64, batch: usize, drafts: usize,
52+
width: usize, packed: usize, preserve: bool, values: &[utility::LaneValue; REQUEST_LANES],
53+
optimistic: bool) -> [usize; REQUEST_LANES] {
54+
let ceilings: [_; REQUEST_LANES] = std::array::from_fn(|slot| {
55+
if mask & (1 << slot) == 0 { 0 }
56+
else { scheduler.speculative_tokens(scheduler.slots[slot].request.as_ref().unwrap(), width.min(drafts + 1)) }
57+
});
58+
let gpu = scheduler.gpu_prefix_allocation(batch, drafts, width, packed);
59+
if (gpu && preserve) || (!gpu && scheduler.speculative_fixed_width.is_some())
60+
|| ceilings.iter().sum::<usize>() <= packed { return ceilings; }
61+
let probabilities: [[f64; 7]; REQUEST_LANES] = std::array::from_fn(|slot| {
62+
let Some(request) = scheduler.slots.get(slot).and_then(|slot| slot.request.as_ref()) else { return [0.0; 7]; };
63+
if gpu { scheduler.speculative_curve(request, optimistic) }
64+
else { request.acceptance.curve(scheduler.proposal_source, request.generated_tokens, optimistic)
65+
.map(|p| p.powf(scheduler.policy.adaptive_verification_alpha) * values[slot].output_token) }
66+
});
67+
let mut lengths = [0; REQUEST_LANES];
68+
for slot in policy::bits(mask) { lengths[slot] = 1; }
69+
for _ in 0..packed.saturating_sub(mask.count_ones() as usize) {
70+
let Some(slot) = policy::bits(mask).filter(|&slot| lengths[slot] < ceilings[slot])
71+
.max_by(|&a, &b| probabilities[a][lengths[a] - 1].total_cmp(&probabilities[b][lengths[b] - 1])
72+
.then_with(|| lengths[b].cmp(&lengths[a])).then_with(|| b.cmp(&a))) else { break; };
73+
lengths[slot] += 1;
74+
}
75+
lengths
76+
}
77+
78+
#[test]
79+
fn ranked_prefix_allocation_matches_greedy_frontiers() {
80+
for active in [1, 3, 8, 32, 64] {
81+
for allocation in [PrefixAllocation::Host, PrefixAllocation::Gpu] {
82+
for source in [speculation::ProposalSource::Mtp, speculation::ProposalSource::Copy] {
83+
let mut scheduler = planning_scheduler(active, allocation);
84+
scheduler.proposal_source = source;
85+
let round = ScheduledStep {
86+
kind: GraphKind::Speculative, batch_capacity: active.next_power_of_two(),
87+
draft_tokens: 7, token_capacity: 8, packed_token_capacity: active.next_power_of_two() * 8,
88+
mixed: None, requires_feedback: false,
89+
lanes: scheduler.slots.iter().filter_map(|slot| slot.request.as_ref()).map(|request|
90+
ScheduledLane { request: request.handle, start_position: request.consumed_tokens, tokens: 8 }).collect(),
91+
};
92+
let mut records = vec![CompletionRecord::default(); scheduler.limits.request_capacity()];
93+
let confidence: Vec<_> = round.lanes.iter().enumerate().map(|(row, lane)| {
94+
records[lane.request.slot as usize] = CompletionRecord {
95+
status: 1, generated_tokens: 18, vector: Some(TokenVector {
96+
emitted_count: 1, offered_drafts: 7, verification_width: 8, ..Default::default()
97+
}), ..Default::default()
98+
};
99+
std::array::from_fn(|depth| [0.0, 0.25, 0.75, 1.0][(row + depth) % 4])
100+
}).collect();
101+
scheduler.observe_packed_forecasts(&round, &records, &confidence);
102+
for seed in 0..4 {
103+
scheduler.policy.adaptive_verification_alpha = [0.5, 1.0, 2.0, 4.0][seed];
104+
let mut mask = 0u64;
105+
let mut values = [utility::LaneValue::default(); REQUEST_LANES];
106+
for (slot, entry) in scheduler.slots.iter_mut().enumerate() {
107+
let Some(request) = entry.request.as_mut() else { continue; };
108+
if (slot + seed) % 3 != 1 { mask |= 1 << slot; }
109+
request.max_new_tokens = request.generated_tokens + 1 + ((slot + seed) % 8) as u32;
110+
request.host_stops = seed == 2 && slot % 7 == 0;
111+
request.grammar_constrained = seed == 3 && slot % 11 == 0;
112+
request.copy_proposal = Some(speculation::CopyProposal {
113+
consumed: request.consumed_tokens, tokens: [9; 7], count: (slot + seed) % 8,
114+
grammar_masks: vec![],
115+
});
116+
if seed == 0 { request.acceptance = Default::default(); }
117+
else { request.acceptance.observe(source, request.generated_tokens,
118+
7, slot % 8, slot % 8 < 7, None); }
119+
values[slot].output_token = [0.0, 1.0, 1.0, 4.0][(slot + seed) % 4];
120+
}
121+
scheduler.limits.max_sequence_tokens = if seed == 3 { 83 } else { 8192 };
122+
scheduler.forecasts.population = [0.9, 0.9, 0.5, 0.5, 0.25, 0.0, 0.0];
123+
let lanes = mask.count_ones() as usize;
124+
for width in [1, 2, 4, 8] {
125+
for drafts in [width - 1, 7] {
126+
for packed in lanes..=lanes * width {
127+
scheduler.capabilities.speculative_allocations = vec![
128+
(active.next_power_of_two(), drafts, width, packed, allocation)];
129+
for (fixed, preserve, optimistic) in [
130+
(false, false, false), (false, false, true),
131+
(false, true, false), (true, false, false),
132+
] {
133+
scheduler.speculative_fixed_width = fixed.then_some(width);
134+
let expected = greedy_lengths(&scheduler, mask, active.next_power_of_two(),
135+
drafts, width, packed, preserve, &values, optimistic);
136+
let actual = scheduler.prefix_lengths_with_values(mask, active.next_power_of_two(),
137+
drafts, width, packed, preserve, &values, optimistic);
138+
assert_eq!(actual, expected,
139+
"active={active} allocation={allocation:?} source={source:?} seed={seed} width={width} drafts={drafts} packed={packed} fixed={fixed} preserve={preserve} optimistic={optimistic}");
140+
}
141+
}
142+
}
143+
}
144+
}
145+
}
146+
}
147+
}
148+
}
149+
150+
#[test]
151+
fn host_cohort_ranking_matches_greedy_reallocation() {
152+
for active in [3, 8, 32, 64] {
153+
let scheduler = planning_scheduler(active, PrefixAllocation::Host);
154+
let (values, _, _) = scheduler.lane_values(scheduler.window_remaining().unwrap());
155+
let available: Vec<_> = (0..active).rev().collect();
156+
for optimistic in [false, true] {
157+
let ranking = scheduler.host_prefix_ranking(&available, 7, 8, &values, optimistic);
158+
for packed in [active, active * 2, active * 4] {
159+
let mut lengths = [0; REQUEST_LANES];
160+
let mut frontier = std::collections::BinaryHeap::new();
161+
let mut mask = 0;
162+
for (count, &slot) in available.iter().enumerate() {
163+
mask |= 1 << slot;
164+
ranking.extend_cohort(slot, &mut lengths, &mut frontier, packed - count - 1);
165+
assert_eq!(lengths, greedy_lengths(&scheduler, mask, active.next_power_of_two(),
166+
7, 8, packed, false, &values, optimistic));
167+
}
168+
}
169+
}
170+
}
171+
}
172+
173+
#[test]
174+
#[ignore = "CPU planning benchmark; run in release mode with --ignored --nocapture --test-threads=1"]
175+
fn benchmark_packed_planning() {
176+
assert!(!cfg!(debug_assertions), "planning measurements require a release build");
177+
for active in [8, 16, 32, 64] {
178+
for allocation in [PrefixAllocation::Host, PrefixAllocation::Gpu] {
179+
let mut scheduler = planning_scheduler(active, allocation);
180+
let first = scheduler.plan().unwrap();
181+
assert!(!first.requires_feedback);
182+
scheduler.commit(&first, Fence(scheduler.submitted.0 + 1)).unwrap();
183+
for probe in [false, true] {
184+
scheduler.refresh.acceptance_credit_ns = if probe { refresh::CREDIT_CAP_NS } else { 0.0 };
185+
for _ in 0..10 { std::hint::black_box(scheduler.plan()); }
186+
let samples: Vec<_> = (0..9).map(|_| {
187+
let start = std::time::Instant::now();
188+
for _ in 0..50 { std::hint::black_box(scheduler.plan()); }
189+
start.elapsed().as_nanos() / 50
190+
}).collect();
191+
println!("planning requests={active} allocation={allocation:?} queued=1 probe={probe} ns={samples:?}");
192+
}
193+
}
194+
}
195+
}

0 commit comments

Comments
 (0)