eredu_runtime/
speculative.rs

1//! Backend-neutral speculative request lifecycle and fair scheduling.
2
3use eredu_core::{
4    CompletedSpeculativeSchedule, PreparedSpeculativeLane, SpeculativeConstraint,
5    SpeculativeDriverError, SpeculativeExecutor, SpeculativeGenerationBatchOutput,
6    SpeculativeGenerationOutput, SpeculativeGenerationVisitor, SpeculativePublisher,
7    SpeculativeRequestTable, SpeculativeSampling,
8};
9
10/// Neutral owner of speculative request registration and fair scheduling.
11pub struct SpeculativeScheduler<'a, E, S, C, P>
12where
13    E: SpeculativeExecutor,
14    S: SpeculativeSampling<Logits = E::Logits, Error = E::Error>,
15    C: SpeculativeConstraint,
16    P: SpeculativePublisher<C>,
17{
18    executor: &'a mut E,
19    context: E::Context<'a>,
20    optimistic_execution_available: bool,
21    component_timings_collected: bool,
22    requests: SpeculativeRequestTable<'a, E, S, C, P>,
23}
24
25impl<'a, E, S, C, P> SpeculativeScheduler<'a, E, S, C, P>
26where
27    E: SpeculativeExecutor + 'a,
28    S: SpeculativeSampling<Logits = E::Logits, Error = E::Error, Context<'a> = E::Context<'a>> + 'a,
29    C: SpeculativeConstraint,
30    P: SpeculativePublisher<C>,
31{
32    /// Creates a scheduler for one prepared executor and placement.
33    #[allow(clippy::too_many_arguments)]
34    pub fn new(
35        executor: &'a mut E,
36        options: eredu_core::generation::SpeculativeSchedulerOptions,
37        topology: eredu_core::SpeculativeExecutionTopology,
38        optimistic_execution_available: bool,
39        component_timings_collected: bool,
40        context: E::Context<'a>,
41    ) -> Result<Self, SpeculativeDriverError<E::Error>> {
42        executor.set_telemetry_enabled(component_timings_collected);
43        Ok(Self {
44            executor,
45            context,
46            optimistic_execution_available,
47            component_timings_collected,
48            requests: SpeculativeRequestTable::new(options, topology)
49                .map_err(SpeculativeDriverError::Generation)?,
50        })
51    }
52
53    /// Registers and prefills one independently progressing lane.
54    pub fn submit(
55        &mut self,
56        mut lane: PreparedSpeculativeLane<'a, E, S, C, P>,
57    ) -> Result<eredu_core::generation::SpeculativeRequestId, SpeculativeDriverError<E::Error>>
58    {
59        self.requests.submit(
60            self.executor,
61            lane.take_cache(),
62            lane.take_input(),
63            lane.take_config(),
64            lane.take_runtime(),
65            lane.take_randomness(),
66            self.component_timings_collected,
67            self.context,
68        )
69    }
70
71    /// Performs one fairly selected lifecycle action.
72    pub fn step(&mut self) -> Result<bool, SpeculativeDriverError<E::Error>> {
73        self.requests.step(
74            self.executor,
75            self.optimistic_execution_available,
76            self.context,
77        )
78    }
79
80    /// Drives all registered lanes to terminal states.
81    pub fn run(&mut self) -> Result<(), SpeculativeDriverError<E::Error>> {
82        while self.step()? {}
83        Ok(())
84    }
85
86    /// Returns the current portable status for one lane.
87    pub fn status(
88        &self,
89        id: eredu_core::generation::SpeculativeRequestId,
90    ) -> Option<eredu_core::generation::SpeculativeRequestStatus> {
91        self.requests.status(id)
92    }
93
94    /// Requests cancellation at the next exact safe boundary.
95    pub fn cancel(
96        &mut self,
97        id: eredu_core::generation::SpeculativeRequestId,
98    ) -> Result<(), SpeculativeDriverError<E::Error>> {
99        self.requests.cancel(id)
100    }
101
102    /// Whether all registered lanes are terminal.
103    pub fn is_finished(&self) -> bool {
104        self.requests.is_finished()
105    }
106
107    /// Consumes a terminal scheduler into stable ordered results.
108    pub fn finish(
109        self,
110    ) -> Result<CompletedSpeculativeSchedule<S>, SpeculativeDriverError<E::Error>> {
111        self.requests.finish()
112    }
113}
114
115/// Facade-selected speculative generation driver.
116///
117/// Backends lend prepared native resources through
118/// [`SpeculativeGenerationVisitor`]. This driver alone registers lanes, runs
119/// the fair schedule, observes exact completions, and constructs public
120/// terminal outputs.
121#[derive(Debug, Clone, Copy, Default)]
122pub struct RunSpeculativeGeneration {
123    options: eredu_core::generation::SpeculativeSchedulerOptions,
124}
125
126impl RunSpeculativeGeneration {
127    /// Creates a driver with facade-selected scheduling and lookahead controls.
128    pub const fn new(options: eredu_core::generation::SpeculativeSchedulerOptions) -> Self {
129        Self { options }
130    }
131}
132
133impl SpeculativeGenerationVisitor for RunSpeculativeGeneration {
134    fn run<'a, E, S, C, P>(
135        self,
136        executor: &'a mut E,
137        lanes: Vec<PreparedSpeculativeLane<'a, E, S, C, P>>,
138        topology: eredu_core::SpeculativeExecutionTopology,
139        optimistic_execution_available: bool,
140        component_timings_collected: bool,
141        context: E::Context<'a>,
142    ) -> Result<SpeculativeGenerationBatchOutput, SpeculativeDriverError<E::Error>>
143    where
144        E: SpeculativeExecutor + 'a,
145        S: SpeculativeSampling<Logits = E::Logits, Error = E::Error, Context<'a> = E::Context<'a>>
146            + 'a,
147        C: SpeculativeConstraint,
148        P: SpeculativePublisher<C>,
149    {
150        let mut scheduler = SpeculativeScheduler::new(
151            executor,
152            self.options,
153            topology,
154            optimistic_execution_available,
155            component_timings_collected,
156            context,
157        )?;
158        for lane in lanes {
159            scheduler.submit(lane)?;
160        }
161        scheduler.run()?;
162        let mut completed = scheduler.finish()?;
163        let requests = completed
164            .take_requests()
165            .into_iter()
166            .map(|request| -> Result<_, SpeculativeDriverError<E::Error>> {
167                let finish_reason = request.finish_reason().ok_or_else(|| {
168                    SpeculativeDriverError::Generation(
169                        eredu_core::generation::GenerationError::MissingSpeculativeFinishReason {
170                            index: request.id().index(),
171                        },
172                    )
173                })?;
174                Ok(SpeculativeGenerationOutput::new(
175                    request.token_ids().to_vec(),
176                    finish_reason,
177                    request.stats().clone(),
178                ))
179            })
180            .collect::<Result<Vec<_>, _>>()?;
181        Ok(SpeculativeGenerationBatchOutput::new(
182            requests,
183            completed.take_scheduler(),
184        ))
185    }
186}
eredu_runtime/speculative.rs

eredu_runtime/
speculative.rs
