1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
// SPDX-License-Identifier: MIT OR Apache-2.0
//! 2026-09-25: The launch half of the plain decode step: the KV-exhaustion retry loop
//! (`launch_with_preemption`), which `decode_batch_with_preemption` awaits at
//! once and the pipelined lane keeps in flight, the decode-preemption victim
//! policy, and the readback flavour and DFlash context commit a launch plans.
//!
//! Owner: scheduler.
//! Invariants: none beyond the types.
use super::*;
use crate::scheduler::io::{DecodeCtxCommit, DeviceError, Readback, SchedIo, SpillIo, StepPlan};
use metrale_scheduler::Ticket;
/// 2026-09-25: How a plain decode step's argmax readback is planned.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum FeedReadback {
/// 2026-09-25: The readback `decode_readback_plan` picks, unchanged.
Plain,
/// 2026-09-25: The masked feed argmax: the router runs ahead of the host, so a masked
/// id must be re-picked on the device rather than by a host round trip.
Masked { think_end_token: Option<u32> },
}
/// 2026-09-25: A resumed sequence is passed over as a decode-preemption victim until it
/// has generated this many more tokens, unless every candidate is immune (see
/// [`choose_decode_victim`]).
pub(super) const PREEMPT_IMMUNITY_TOKENS: usize = 64;
/// 2026-09-25: Pick the decode-preemption victim.
///
/// Policy: fewest generated tokens, ties broken by the lower slot.
///
/// Exclusions:
/// * grammar-active sequences;
/// * without spill, sequences whose tokens contain vision pads (their KV
/// came from image embeddings a token re-prefill cannot reproduce);
/// * immune sequences (see [`PREEMPT_IMMUNITY_TOKENS`]), unless every
/// eligible candidate is immune.
pub(super) fn choose_decode_victim(
model: &dyn Model,
active: &[ActiveSeq],
can_spill: bool,
) -> Option<usize> {
let eligible = |a: &ActiveSeq| {
a.grammar_state.is_none() && (can_spill || !model.tokens_contain_vision_pad(&a.seq.tokens))
};
let pick = |immune_ok: bool| {
active
.iter()
.enumerate()
.filter(|(_, a)| {
eligible(a) && (immune_ok || a.output_tokens.len() >= a.preempt_immune_until_tokens)
})
.min_by_key(|(_, a)| (a.output_tokens.len(), a.seq.slot_idx))
.map(|(i, _)| i)
};
pick(false).or_else(|| pick(true))
}
/// 2026-09-25: Send `e` to every row and drain them. Unlike the forward-error path in
/// [`launch_with_preemption`], this does not check the fault latch.
pub(super) fn fail_readback(io: &SchedIo, active: &mut Vec<ActiveSeq>, e: &anyhow::Error) {
tracing::error!("argmax_batch error: {e:#}");
for mut a in active.drain(..) {
send_error(io, &mut a, &format!("{e:#}"));
}
}
/// 2026-09-25: Which DFlash context commit a plain decode step of `n` rows performs.
pub(super) fn decode_ctx_commit(n: usize, levers: &super::levers::SchedLevers) -> DecodeCtxCommit {
if levers.dflash_unified_ctx {
DecodeCtxCommit::Unified
} else if n == 1 && levers.dflash_serial_append {
DecodeCtxCommit::SerialAppend
} else {
DecodeCtxCommit::None
}
}
/// 2026-09-25: Launch a plain decode step, preempting a victim and retrying while the
/// pool reports exhaustion. Returns the ticket unawaited, so a pipelining
/// caller can keep the step in flight; `None` when every row was failed
/// instead. `feed` decides the argmax readback's flavour.
pub(super) fn launch_with_preemption(
sched: &crate::scheduler::sched_ctx::SchedCtx,
active: &mut Vec<ActiveSeq>,
spill: Option<&dyn SpillIo>,
swapped: &mut Vec<SwappedSeq>,
preempted: &mut Vec<PreemptedSeq>,
staging: &mut Vec<u8>,
feed: FeedReadback,
) -> Option<Ticket> {
let io = &sched.io;
let model = io.dev.model();
loop {
let tokens: Vec<u32> = active.iter().map(|a| a.last_token).collect();
let readback = match (
super::decode_logits_step::decode_readback_plan(active, sched, staging),
feed,
) {
(Readback::Argmax, FeedReadback::Masked { think_end_token }) => {
Readback::ArgmaxMasked {
masks: active
.iter()
.map(|a| super::core::pipeline::row_mask(a, think_end_token))
.collect(),
}
}
(planned, _) => planned,
};
let ctx = decode_ctx_commit(active.len(), &sched.levers);
let mut refs: Vec<&mut SequenceState> = active.iter_mut().map(|a| &mut a.seq).collect();
let launched = io.dev.launch(
StepPlan::Decode {
tokens,
ctx,
readback,
},
&mut refs,
);
match launched {
Ok(t) => return Some(t),
Err(DeviceError::Readback(e)) | Err(DeviceError::Effect(e)) => {
drop(refs);
fail_readback(io, active, &e);
return None;
}
Err(err) => {
drop(refs);
let e = err.into_inner();
let victim = if format!("{e:#}").contains("KV cache exhausted") && active.len() > 1
{
choose_decode_victim(model, active, spill.is_some())
} else {
None
};
let Some(vi) = victim else {
tracing::error!("decode_batch error: {e:#}");
for mut a in active.drain(..) {
send_error(io, &mut a, &format!("{e:#}"));
}
// 2026-09-25: If the process-wide fault latch reports a lost CUDA
// context, request shutdown: every later request would fail
// the same way. `request` is idempotent, so repeated failures
// do not repeat it.
if let Some(reason) = metrale_core::fault::global().fault() {
crate::tui::shutdown::request(reason);
}
return None;
};
// 2026-09-25: `remove`, not `swap_remove`, keeps the ascending-slot order
// the caller sorted `active` into.
let v = active.remove(vi);
tracing::warn!(
"KV cache exhausted during decode: preempting slot={} \
({} blocks, {} tokens generated) for later RESUME so the \
other {} sequence(s) can continue",
v.seq.slot_idx,
v.seq.block_table.len(),
v.output_tokens.len(),
active.len(),
);
match spill {
Some(sp) => match super::preempt::spill_out_sequence(io, v, sp) {
Ok(s) => swapped.push(s),
Err((v, spill_err)) => {
tracing::warn!(
"decode-preempt spill failed ({spill_err:#}); \
requeuing victim for re-prefill instead"
);
preempted.push(super::preempt::preempt_requeue(io, v));
}
},
None => preempted.push(super::preempt::preempt_requeue(io, v)),
}
}
}
}
}