metrale_scheduler/
plan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: What the core asks the device to run ([`StepPlan`]) and what comes back
4//! ([`StepResult`]). Generic over the device's logits handle `P`.
5//!
6//! Owner: scheduler.
7//! Invariants: none beyond the types.
8
9/// 2026-09-25: A short label for traces.
10pub trait Label {
11    fn label(&self) -> String;
12}
13
14/// 2026-09-25: Where a fed decode row's input token comes from.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub enum FeedSource {
17    /// 2026-09-25: The previous step's argmax for row `from_row`, read from the device
18    /// feed cell that step wrote; the host does not know the id at launch.
19    Feed { from_row: u32 },
20    /// 2026-09-25: A token id the host supplies.
21    Host(u32),
22}
23
24/// 2026-09-25: The two token ids the host would re-pick around for a row on the
25/// argmax fast path (`u32::MAX` = none): the closed-thinking mask.
26pub type RowMask = [u32; 2];
27
28/// 2026-09-25: How the logits of a decode step come back to the host.
29pub enum Readback<'a> {
30    /// 2026-09-25: One argmax id per row, computed on the device.
31    Argmax,
32    /// 2026-09-25: The device argmax with the per-row two-id mask applied exactly as
33    /// the host re-pick would (see `argmax_feed.cu`); the ids also land
34    /// in the device feed cells for a following fed step.
35    ArgmaxMasked { masks: Vec<RowMask> },
36    /// 2026-09-25: The whole logits block, copied into `into` (resized by the router to
37    /// `rows * vocab * elem_bytes`; the core samples over it).
38    HostLogits { into: &'a mut Vec<u8> },
39}
40
41impl std::fmt::Debug for Readback<'_> {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        match self {
44            Self::Argmax => f.write_str("Argmax"),
45            Self::ArgmaxMasked { masks } => write!(f, "ArgmaxMasked{{n={}}}", masks.len()),
46            Self::HostLogits { .. } => f.write_str("HostLogits"),
47        }
48    }
49}
50
51/// 2026-09-25: The DFlash context commit a decode step performs between its forward and
52/// its readback (the drafter's view of the serial tokens).
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub enum DecodeCtxCommit {
55    None,
56    /// 2026-09-25: `commit_ctx(row, 1, seq_len - 1, row)` for every row.
57    Unified,
58    /// 2026-09-25: `dflash_serial_ctx_append(row 0)` (single-row batches only).
59    SerialAppend,
60}
61
62/// 2026-09-25: One device step. The sequence rows it advances travel beside it as
63/// `&mut S`.
64#[derive(Debug)]
65pub enum StepPlan<'a> {
66    /// 2026-09-25: Batched decode of one token per row, the context commit, then the
67    /// readback.
68    Decode {
69        /// 2026-09-25: The pending decode input of each row, in row order.
70        tokens: Vec<u32>,
71        ctx: DecodeCtxCommit,
72        readback: Readback<'a>,
73    },
74    /// 2026-09-25: A decode step launched while the previous one is still in flight:
75    /// each row's input is resolved on the device from `sources`, the
76    /// readback is the masked feed argmax, and the host bookkeeping the
77    /// plain step performs at launch (`tokens.push`, `seq_len += 1`) is
78    /// deferred to the core, which applies it once the previous step is
79    /// committed. The context commit runs between the forward and the
80    /// readback, as it does for the plain step.
81    DecodeFed {
82        sources: Vec<FeedSource>,
83        masks: Vec<RowMask>,
84        ctx: DecodeCtxCommit,
85    },
86}
87
88impl Label for StepPlan<'_> {
89    fn label(&self) -> String {
90        match self {
91            Self::Decode {
92                tokens,
93                ctx,
94                readback,
95            } => format!(
96                "Decode{{n={}, ctx={ctx:?}, readback={readback:?}}}",
97                tokens.len()
98            ),
99            Self::DecodeFed { sources, ctx, .. } => {
100                let fed = sources
101                    .iter()
102                    .filter(|s| matches!(s, FeedSource::Feed { .. }))
103                    .count();
104                format!("DecodeFed{{n={}, fed={fed}, ctx={ctx:?}}}", sources.len())
105            }
106        }
107    }
108}
109
110/// 2026-09-25: The rows of a decode step after the readback.
111#[derive(Debug)]
112pub enum DecodeRows {
113    /// 2026-09-25: The device argmax per row.
114    Tokens(Vec<u32>),
115    /// 2026-09-25: The logits block was copied into the plan's buffer; `elem_bytes` is
116    /// the element width the model writes (2 = BF16, 4 = FP32).
117    HostLogits { elem_bytes: usize },
118}
119
120#[derive(Debug)]
121pub enum StepOutcome<P> {
122    Decode {
123        /// 2026-09-25: The device logits the step produced, for a later
124        /// [`crate::Effect::ReadLogits`].
125        logits: P,
126        rows: DecodeRows,
127    },
128}
129
130#[derive(Debug)]
131pub struct StepResult<P> {
132    pub ticket: crate::Ticket,
133    pub outcome: StepOutcome<P>,
134}
metrale_scheduler/plan.rs

metrale_scheduler/
plan.rs
