Skip to main content

eredu_core/
generation.rs

1//! Portable generation configuration, lifecycle, and speculative bookkeeping.
2
3use crate::backend::{
4    BoundedCompletionWait, BoundedCompletionWaitError, CompletionCancellationMode,
5};
6use serde::{Deserialize, Serialize};
7use std::sync::{
8    atomic::{AtomicBool, Ordering},
9    Arc,
10};
11
12/// Why generation reached a terminal state.
13#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum FinishReason {
16    /// A checkpoint end-of-sequence token was committed.
17    Eos,
18    /// A caller or protocol stop sequence matched.
19    StopSequence,
20    /// The committed generation grammar reached an accepting state.
21    GrammarComplete,
22    /// The configured output-token budget was exhausted.
23    MaxTokens,
24    /// The caller cooperatively cancelled generation.
25    Cancelled,
26}
27
28/// Cheap thread-safe cooperative cancellation observed between submissions.
29#[derive(Debug, Clone, Default)]
30pub struct GenerationCancellationToken {
31    cancelled: Arc<AtomicBool>,
32}
33
34impl GenerationCancellationToken {
35    /// Creates an active token.
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    /// Permanently requests cancellation for every clone.
41    pub fn cancel(&self) {
42        self.cancelled.store(true, Ordering::Release);
43    }
44
45    /// Returns whether cancellation has been requested.
46    pub fn is_cancelled(&self) -> bool {
47        self.cancelled.load(Ordering::Acquire)
48    }
49}
50
51/// Terminal signals observed while committing one token.
52#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)]
53pub struct TokenTerminalSignals {
54    /// Decoded output matched a stop sequence.
55    pub stop_sequence: bool,
56    /// The committed grammar is complete.
57    pub grammar_complete: bool,
58}
59
60/// Result of one canonical token commit.
61#[derive(Debug, Clone, Copy, Eq, PartialEq)]
62pub struct TokenCommit {
63    /// Token id committed to the logical output.
64    pub token_id: u32,
65    /// Zero-based generated-token position.
66    pub position: usize,
67    /// Terminal reason selected after this commit, if any.
68    pub finish_reason: Option<FinishReason>,
69}
70
71/// Canonical committed-token sequence and terminal-condition precedence.
72#[derive(Debug, Clone, Eq, PartialEq)]
73pub struct GenerationSequence {
74    max_tokens: usize,
75    eos_token_ids: Vec<u32>,
76    tokens: Vec<u32>,
77    finish_reason: Option<FinishReason>,
78}
79
80impl GenerationSequence {
81    /// Logical data copied with the complete sequence, including EOS policy.
82    /// Excludes allocator capacity and overhead.
83    pub fn snapshot_storage_bytes(&self) -> Option<u64> {
84        let Self {
85            max_tokens: _,
86            eos_token_ids,
87            tokens,
88            finish_reason: _,
89        } = self;
90        (std::mem::size_of::<Self>() as u64)
91            .checked_add((eos_token_ids.len() as u64).checked_mul(4)?)?
92            .checked_add((tokens.len() as u64).checked_mul(4)?)
93    }
94    /// Creates an empty output sequence with a fixed token budget.
95    pub fn new(max_tokens: usize, eos_token_ids: impl IntoIterator<Item = u32>) -> Self {
96        let mut eos_token_ids = eos_token_ids.into_iter().collect::<Vec<_>>();
97        eos_token_ids.sort_unstable();
98        eos_token_ids.dedup();
99        Self {
100            max_tokens,
101            eos_token_ids,
102            tokens: Vec::with_capacity(max_tokens),
103            finish_reason: (max_tokens == 0).then_some(FinishReason::MaxTokens),
104        }
105    }
106
107    /// Commits one token and applies stop, grammar, EOS, and budget precedence.
108    pub fn commit(
109        &mut self,
110        token_id: u32,
111        signals: TokenTerminalSignals,
112    ) -> Result<TokenCommit, GenerationError> {
113        if self.finish_reason.is_some() {
114            return Err(GenerationError::AlreadyFinished);
115        }
116        let position = self.tokens.len();
117        self.tokens.push(token_id);
118        let finish_reason = signals
119            .stop_sequence
120            .then_some(FinishReason::StopSequence)
121            .or_else(|| {
122                signals
123                    .grammar_complete
124                    .then_some(FinishReason::GrammarComplete)
125            })
126            .or_else(|| {
127                self.eos_token_ids
128                    .binary_search(&token_id)
129                    .is_ok()
130                    .then_some(FinishReason::Eos)
131            })
132            .or_else(|| (self.tokens.len() == self.max_tokens).then_some(FinishReason::MaxTokens));
133        self.finish_reason = finish_reason;
134        Ok(TokenCommit {
135            token_id,
136            position,
137            finish_reason,
138        })
139    }
140
141    /// Applies cancellation if the sequence has not already terminated.
142    pub fn cancel(&mut self) -> bool {
143        if self.tokens.is_empty() && self.finish_reason == Some(FinishReason::MaxTokens) {
144            self.finish_reason = Some(FinishReason::Cancelled);
145            true
146        } else if self.finish_reason.is_some() {
147            false
148        } else {
149            self.finish_reason = Some(FinishReason::Cancelled);
150            true
151        }
152    }
153
154    /// Observes a cooperative token at a submission boundary.
155    pub fn observe_cancellation(&mut self, cancellation: &GenerationCancellationToken) -> bool {
156        cancellation.is_cancelled() && self.cancel()
157    }
158
159    /// Committed tokenizer ids in canonical order.
160    pub fn tokens(&self) -> &[u32] {
161        &self.tokens
162    }
163
164    /// Consumes the state into committed tokenizer ids.
165    pub fn into_tokens(self) -> Vec<u32> {
166        self.tokens
167    }
168
169    /// Number of remaining token slots.
170    pub fn remaining(&self) -> usize {
171        self.max_tokens.saturating_sub(self.tokens.len())
172    }
173
174    /// Configured output-token budget.
175    pub const fn max_tokens(&self) -> usize {
176        self.max_tokens
177    }
178
179    /// Selected terminal reason.
180    pub const fn finish_reason(&self) -> Option<FinishReason> {
181        self.finish_reason
182    }
183
184    /// Whether no further token may be committed.
185    pub const fn is_finished(&self) -> bool {
186        self.finish_reason.is_some()
187    }
188}
189
190/// Kind of non-proposal token emitted by speculative verification.
191#[derive(Debug, Clone, Copy, Eq, PartialEq)]
192pub enum SpeculativeTail {
193    /// A rejected proposal was replaced from the target distribution.
194    Replacement,
195    /// Every proposal was accepted and the target emitted its bonus token.
196    Bonus,
197}
198
199/// Canonical bookkeeping for one proposal/verification transaction.
200#[derive(Debug, Clone, Eq, PartialEq)]
201pub struct SpeculativeRound {
202    proposal_count: usize,
203    accepted: usize,
204    committed_tokens: Vec<u32>,
205    tail: Option<SpeculativeTail>,
206    terminal: bool,
207}
208
209impl SpeculativeRound {
210    /// Starts a verification round for a non-empty proposal block.
211    pub fn new(proposal_count: usize) -> Result<Self, GenerationError> {
212        if proposal_count == 0 {
213            return Err(GenerationError::EmptyProposalBlock);
214        }
215        Ok(Self {
216            proposal_count,
217            accepted: 0,
218            committed_tokens: Vec::with_capacity(proposal_count + 1),
219            tail: None,
220            terminal: false,
221        })
222    }
223
224    /// Records the next accepted proposal token.
225    pub fn accept(&mut self, token: u32, terminal: bool) -> Result<(), GenerationError> {
226        if self.tail.is_some() || self.accepted == self.proposal_count || self.terminal {
227            return Err(GenerationError::InvalidSpeculativeTransition);
228        }
229        self.accepted += 1;
230        self.committed_tokens.push(token);
231        self.terminal = terminal;
232        Ok(())
233    }
234
235    /// Records the target replacement for the first rejected proposal.
236    pub fn reject_with(&mut self, token: u32, terminal: bool) -> Result<(), GenerationError> {
237        if self.tail.is_some() || self.accepted == self.proposal_count || self.terminal {
238            return Err(GenerationError::InvalidSpeculativeTransition);
239        }
240        self.tail = Some(SpeculativeTail::Replacement);
241        self.committed_tokens.push(token);
242        self.terminal = terminal;
243        Ok(())
244    }
245
246    /// Records the target bonus after complete proposal acceptance.
247    pub fn bonus(&mut self, token: u32, terminal: bool) -> Result<(), GenerationError> {
248        if self.tail.is_some() || self.accepted != self.proposal_count || self.terminal {
249            return Err(GenerationError::InvalidSpeculativeTransition);
250        }
251        self.tail = Some(SpeculativeTail::Bonus);
252        self.committed_tokens.push(token);
253        self.terminal = terminal;
254        Ok(())
255    }
256
257    /// Whether every proposal was accepted.
258    pub const fn is_full_acceptance(&self) -> bool {
259        self.accepted == self.proposal_count
260    }
261
262    /// Produces the exact cache-retention and output-publication plan.
263    pub fn commit_plan(&self) -> Result<SpeculativeCommitPlan<'_>, GenerationError> {
264        if !self.terminal && self.tail.is_none() {
265            return Err(GenerationError::IncompleteSpeculativeRound);
266        }
267        Ok(SpeculativeCommitPlan {
268            accepted_proposals: self.accepted,
269            committed_tokens: &self.committed_tokens,
270            verified_inputs: if self.tail.is_some() {
271                1 + self.accepted
272            } else {
273                self.accepted
274            },
275            full_acceptance: self.accepted == self.proposal_count,
276            tail: self.tail,
277            terminal: self.terminal,
278        })
279    }
280}
281
282/// Borrowed resolution plan for a speculative transaction.
283#[derive(Debug, Clone, Copy, Eq, PartialEq)]
284pub struct SpeculativeCommitPlan<'a> {
285    /// Number of accepted assistant proposals.
286    pub accepted_proposals: usize,
287    /// Tokens that become visible after the backend cache commit succeeds.
288    pub committed_tokens: &'a [u32],
289    /// Verification inputs retained in the target cache.
290    pub verified_inputs: usize,
291    /// Whether every assistant proposal was accepted.
292    pub full_acceptance: bool,
293    /// Optional replacement or target bonus.
294    pub tail: Option<SpeculativeTail>,
295    /// Whether the last committed token terminated generation.
296    pub terminal: bool,
297}
298
299/// Reuse decision for a proposal block drafted against an assumed prefix.
300#[derive(Debug, Clone, Copy, Eq, PartialEq)]
301pub enum OptimisticReuseDecision {
302    /// Terminal output makes the complete branch unusable.
303    DiscardTerminal,
304    /// The target bonus differed from the optimistic first token.
305    DiscardMismatch,
306    /// The matching first token consumed the complete optimistic block.
307    MatchedConsumed,
308    /// The first optimistic token matched and later proposals remain reusable.
309    MatchedRetained,
310}
311
312/// Resolves optimistic proposal reuse without inspecting backend state.
313pub fn resolve_optimistic_reuse(
314    assumed_prefix: &[u32],
315    canonical_prefix: &[u32],
316    optimistic_tokens: &[u32],
317    bonus: u32,
318    terminal: bool,
319) -> Result<OptimisticReuseDecision, GenerationError> {
320    if assumed_prefix != canonical_prefix {
321        return Err(GenerationError::OptimisticPrefixDiverged);
322    }
323    let first = optimistic_tokens
324        .first()
325        .ok_or(GenerationError::EmptyOptimisticBranch)?;
326    if terminal {
327        return Ok(OptimisticReuseDecision::DiscardTerminal);
328    }
329    if *first != bonus {
330        return Ok(OptimisticReuseDecision::DiscardMismatch);
331    }
332    Ok(if optimistic_tokens.len() == 1 {
333        OptimisticReuseDecision::MatchedConsumed
334    } else {
335        OptimisticReuseDecision::MatchedRetained
336    })
337}
338
339/// Options shared by speculative multi-token backends.
340#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
341pub struct SpeculativeConfig {
342    /// Maximum number of output tokens, including terminal tokens.
343    pub max_tokens: usize,
344    /// Maximum assistant proposals per verification round.
345    pub max_draft_tokens: usize,
346    /// Sampling temperature. Zero selects greedy verification.
347    pub temperature: f32,
348    /// Token ids that terminate a sequence.
349    pub eos_token_ids: Vec<u32>,
350}
351
352impl Default for SpeculativeConfig {
353    fn default() -> Self {
354        Self {
355            max_tokens: 256,
356            max_draft_tokens: 4,
357            temperature: 0.0,
358            eos_token_ids: Vec::new(),
359        }
360    }
361}
362
363impl SpeculativeConfig {
364    /// Validates backend-independent speculative settings.
365    pub fn validate(&self) -> Result<(), GenerationError> {
366        if self.max_draft_tokens == 0 {
367            return Err(GenerationError::ZeroDraftTokens);
368        }
369        if !self.temperature.is_finite() || self.temperature < 0.0 {
370            return Err(GenerationError::InvalidTemperature(self.temperature));
371        }
372        Ok(())
373    }
374}
375
376/// Bounded fair-scheduler settings for speculative requests.
377#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
378#[serde(default)]
379pub struct SpeculativeSchedulerOptions {
380    /// Maximum retained target verification transactions.
381    pub max_in_flight_verifications: usize,
382    /// Maximum retained optimistic branches.
383    pub max_optimistic_branches: usize,
384    /// Proposal blocks drafted ahead per request; currently zero or one.
385    pub lookahead_blocks: usize,
386    /// Disable lookahead adaptively when discarded work dominates reuse.
387    pub adaptive_lookahead: bool,
388    /// Resolved branches required before adaptive disabling.
389    pub adaptive_lookahead_min_blocks: usize,
390    /// Maximum milliseconds spent awaiting one exact verification completion.
391    pub completion_timeout_milliseconds: u64,
392    /// Safe disposition required if verification remains live at its deadline.
393    pub completion_cancellation: CompletionCancellationMode,
394}
395
396impl Default for SpeculativeSchedulerOptions {
397    fn default() -> Self {
398        Self {
399            max_in_flight_verifications: 1,
400            max_optimistic_branches: 1,
401            lookahead_blocks: 1,
402            adaptive_lookahead: true,
403            adaptive_lookahead_min_blocks: 4,
404            completion_timeout_milliseconds: 30_000,
405            completion_cancellation: CompletionCancellationMode::QuarantineUntilComplete,
406        }
407    }
408}
409
410impl SpeculativeSchedulerOptions {
411    /// Enables or disables same-request optimistic lookahead.
412    pub fn with_lookahead(mut self, enabled: bool) -> Self {
413        self.lookahead_blocks = usize::from(enabled);
414        if enabled {
415            self.max_optimistic_branches = self.max_optimistic_branches.max(1);
416        }
417        self
418    }
419
420    /// Selects the exact-completion deadline and safe timeout disposition.
421    pub fn with_completion_wait(
422        mut self,
423        timeout: std::time::Duration,
424        cancellation: CompletionCancellationMode,
425    ) -> Result<Self, GenerationError> {
426        let timeout = u64::try_from(timeout.as_millis())
427            .map_err(|_| GenerationError::SpeculativeCompletionTimeoutTooLarge)?;
428        if timeout == 0 {
429            return Err(GenerationError::ZeroSpeculativeCompletionTimeout);
430        }
431        self.completion_timeout_milliseconds = timeout;
432        self.completion_cancellation = cancellation;
433        Ok(self)
434    }
435
436    /// Resolves the portable fields into the bounded backend completion policy.
437    pub fn completion_wait(self) -> Result<BoundedCompletionWait, GenerationError> {
438        BoundedCompletionWait::new(
439            std::time::Duration::from_millis(self.completion_timeout_milliseconds),
440            self.completion_cancellation,
441        )
442        .map_err(|error| match error {
443            BoundedCompletionWaitError::ZeroTimeout => {
444                GenerationError::ZeroSpeculativeCompletionTimeout
445            }
446        })
447    }
448
449    /// Validates scheduler capacity and lookahead invariants.
450    pub fn validate(self) -> Result<Self, GenerationError> {
451        if self.max_in_flight_verifications == 0 {
452            return Err(GenerationError::ZeroInFlightVerifications);
453        }
454        if self.lookahead_blocks > 1 {
455            return Err(GenerationError::TooManyLookaheadBlocks);
456        }
457        if self.lookahead_blocks > 0 && self.max_optimistic_branches == 0 {
458            return Err(GenerationError::LookaheadWithoutBranchCapacity);
459        }
460        if self.lookahead_blocks > 0
461            && self.adaptive_lookahead
462            && self.adaptive_lookahead_min_blocks == 0
463        {
464            return Err(GenerationError::ZeroAdaptiveLookaheadWindow);
465        }
466        self.completion_wait()?;
467        Ok(self)
468    }
469}
470
471/// Stable speculative scheduler request identifier.
472#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
473pub struct SpeculativeRequestId(usize);
474
475impl SpeculativeRequestId {
476    /// Creates an identifier from its stable scheduler insertion index.
477    pub const fn new(index: usize) -> Self {
478        Self(index)
479    }
480
481    /// Returns the scheduler insertion index.
482    pub const fn index(self) -> usize {
483        self.0
484    }
485}
486
487/// Explicit speculative request/round state.
488#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
489#[serde(rename_all = "snake_case")]
490#[non_exhaustive]
491pub enum SpeculativeRequestStatus {
492    /// Target prompt prefill and first-token sampling.
493    Prefill,
494    /// Committed target state is ready to seed proposals.
495    ReadyToDraft,
496    /// A proposal block is ready for target submission.
497    ReadyToSubmitVerification,
498    /// Target verification is submitted and unresolved.
499    TargetVerificationInFlight,
500    /// Same-request continuation is being drafted optimistically.
501    OptimisticDraftRunning,
502    /// Verification is in flight and its optimistic branch is ready.
503    OptimisticDraftReady,
504    /// Target results are being accepted/rejected and committed.
505    VerificationResolution,
506    /// The request reached a normal terminal condition.
507    Completed,
508    /// The request was cancelled.
509    Cancelled,
510}
511
512/// Result of requesting cancellation at a speculative submission boundary.
513#[derive(Debug, Clone, Copy, Eq, PartialEq)]
514pub enum SpeculativeCancellationDisposition {
515    /// The request was already terminal.
516    AlreadyTerminal,
517    /// No backend submission is retained, so cancellation completed now.
518    CancelNow,
519    /// A backend submission must reach its exact safe boundary first.
520    Deferred,
521}
522
523/// Validated backend-neutral lifecycle of one speculative request.
524#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
525pub struct SpeculativeRequestLifecycle {
526    status: SpeculativeRequestStatus,
527    cancellation_pending: bool,
528}
529
530impl Default for SpeculativeRequestLifecycle {
531    fn default() -> Self {
532        Self::new()
533    }
534}
535
536impl SpeculativeRequestLifecycle {
537    /// Starts a request at prompt prefill.
538    pub const fn new() -> Self {
539        Self {
540            status: SpeculativeRequestStatus::Prefill,
541            cancellation_pending: false,
542        }
543    }
544
545    /// Creates a request completed before backend submission.
546    pub const fn completed() -> Self {
547        Self {
548            status: SpeculativeRequestStatus::Completed,
549            cancellation_pending: false,
550        }
551    }
552
553    /// Creates a request cancelled before backend submission.
554    pub const fn cancelled() -> Self {
555        Self {
556            status: SpeculativeRequestStatus::Cancelled,
557            cancellation_pending: false,
558        }
559    }
560
561    /// Current lifecycle status.
562    pub const fn status(&self) -> SpeculativeRequestStatus {
563        self.status
564    }
565
566    /// Whether cancellation must be applied after an exact backend boundary.
567    pub const fn cancellation_pending(&self) -> bool {
568        self.cancellation_pending
569    }
570
571    /// Whether the request can own no further submissions.
572    pub const fn is_terminal(&self) -> bool {
573        matches!(
574            self.status,
575            SpeculativeRequestStatus::Completed | SpeculativeRequestStatus::Cancelled
576        )
577    }
578
579    /// Requests cancellation without discarding an in-flight backend transaction.
580    pub fn request_cancellation(
581        &mut self,
582        submission_retained: bool,
583    ) -> Result<SpeculativeCancellationDisposition, GenerationError> {
584        if self.is_terminal() {
585            return Ok(SpeculativeCancellationDisposition::AlreadyTerminal);
586        }
587        if submission_retained {
588            self.cancellation_pending = true;
589            Ok(SpeculativeCancellationDisposition::Deferred)
590        } else {
591            self.transition(SpeculativeRequestStatus::Cancelled)?;
592            Ok(SpeculativeCancellationDisposition::CancelNow)
593        }
594    }
595
596    /// Applies one legal lifecycle transition.
597    pub fn transition(&mut self, next: SpeculativeRequestStatus) -> Result<(), GenerationError> {
598        let allowed = matches!(
599            (self.status, next),
600            (
601                SpeculativeRequestStatus::Prefill,
602                SpeculativeRequestStatus::ReadyToDraft
603            ) | (
604                SpeculativeRequestStatus::Prefill,
605                SpeculativeRequestStatus::Completed
606            ) | (
607                SpeculativeRequestStatus::Prefill,
608                SpeculativeRequestStatus::Cancelled
609            ) | (
610                SpeculativeRequestStatus::ReadyToDraft,
611                SpeculativeRequestStatus::ReadyToSubmitVerification
612            ) | (
613                SpeculativeRequestStatus::ReadyToDraft,
614                SpeculativeRequestStatus::Completed
615            ) | (
616                SpeculativeRequestStatus::ReadyToDraft,
617                SpeculativeRequestStatus::Cancelled
618            ) | (
619                SpeculativeRequestStatus::ReadyToSubmitVerification,
620                SpeculativeRequestStatus::TargetVerificationInFlight
621            ) | (
622                SpeculativeRequestStatus::ReadyToSubmitVerification,
623                SpeculativeRequestStatus::Cancelled
624            ) | (
625                SpeculativeRequestStatus::TargetVerificationInFlight,
626                SpeculativeRequestStatus::OptimisticDraftRunning
627            ) | (
628                SpeculativeRequestStatus::TargetVerificationInFlight,
629                SpeculativeRequestStatus::VerificationResolution
630            ) | (
631                SpeculativeRequestStatus::OptimisticDraftRunning,
632                SpeculativeRequestStatus::OptimisticDraftReady
633            ) | (
634                SpeculativeRequestStatus::OptimisticDraftReady,
635                SpeculativeRequestStatus::VerificationResolution
636            ) | (
637                SpeculativeRequestStatus::VerificationResolution,
638                SpeculativeRequestStatus::ReadyToDraft
639            ) | (
640                SpeculativeRequestStatus::VerificationResolution,
641                SpeculativeRequestStatus::Completed
642            ) | (
643                SpeculativeRequestStatus::VerificationResolution,
644                SpeculativeRequestStatus::Cancelled
645            )
646        );
647        if !allowed {
648            return Err(GenerationError::InvalidSpeculativeStatusTransition {
649                from: self.status,
650                to: next,
651            });
652        }
653        self.status = next;
654        if self.is_terminal() {
655            self.cancellation_pending = false;
656        }
657        Ok(())
658    }
659}
660
661/// Sampling values declared by a checkpoint generation configuration.
662#[derive(Debug, Clone, Default, PartialEq, Deserialize, Serialize)]
663pub struct CheckpointGenerationConfig {
664    /// Whether stochastic sampling is recommended.
665    #[serde(default)]
666    pub do_sample: Option<bool>,
667    /// Recommended temperature.
668    #[serde(default)]
669    pub temperature: Option<f32>,
670    /// Recommended top-k cutoff.
671    #[serde(default)]
672    pub top_k: Option<i32>,
673    /// Recommended nucleus probability.
674    #[serde(default)]
675    pub top_p: Option<f32>,
676    /// Recommended minimum-probability filter.
677    #[serde(default)]
678    pub min_p: Option<f32>,
679    /// Recommended multiplicative repetition penalty.
680    #[serde(default)]
681    pub repetition_penalty: Option<f32>,
682    /// Recommended generated-history window for repetition penalties.
683    #[serde(default)]
684    pub repeat_last_n: Option<i32>,
685    /// Recommended additive frequency penalty.
686    #[serde(default)]
687    pub frequency_penalty: Option<f32>,
688    /// Recommended additive presence penalty.
689    #[serde(default)]
690    pub presence_penalty: Option<f32>,
691    /// Recommended maximum number of new tokens.
692    #[serde(default)]
693    pub max_new_tokens: Option<usize>,
694}
695
696/// Per-request overrides layered over checkpoint generation settings.
697#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
698pub struct GenerationConfigOverrides {
699    /// Overrides stochastic-versus-greedy selection.
700    pub do_sample: Option<bool>,
701    /// Overrides temperature.
702    pub temperature: Option<f32>,
703    /// Overrides top-k filtering.
704    pub top_k: Option<i32>,
705    /// Overrides top-p filtering.
706    pub top_p: Option<f32>,
707    /// Overrides min-p filtering.
708    pub min_p: Option<f32>,
709    /// Overrides the multiplicative repetition penalty.
710    pub repetition_penalty: Option<f32>,
711    /// Overrides the generated-history window used by repetition penalties.
712    pub repeat_last_n: Option<i32>,
713    /// Overrides the additive frequency penalty.
714    pub frequency_penalty: Option<f32>,
715    /// Overrides the additive presence penalty.
716    pub presence_penalty: Option<f32>,
717    /// Overrides the output-token budget.
718    pub max_new_tokens: Option<usize>,
719}
720
721/// Fully resolved and validated generation settings.
722#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
723pub struct ResolvedGenerationConfig {
724    /// Whether sampling is stochastic.
725    pub do_sample: bool,
726    /// Effective temperature.
727    pub temperature: f32,
728    /// Effective top-k cutoff.
729    pub top_k: i32,
730    /// Effective top-p probability.
731    pub top_p: f32,
732    /// Effective min-p probability.
733    pub min_p: f32,
734    /// Effective multiplicative repetition penalty.
735    pub repetition_penalty: f32,
736    /// Effective generated-history window for repetition penalties.
737    pub repeat_last_n: i32,
738    /// Effective additive frequency penalty.
739    pub frequency_penalty: f32,
740    /// Effective additive presence penalty.
741    pub presence_penalty: f32,
742    /// Effective token budget, when declared.
743    pub max_new_tokens: Option<usize>,
744}
745
746/// Resolves checkpoint and request sampling settings without a tensor runtime.
747pub fn resolve_generation_config(
748    checkpoint: Option<&CheckpointGenerationConfig>,
749    overrides: GenerationConfigOverrides,
750) -> Result<ResolvedGenerationConfig, GenerationError> {
751    let checkpoint_present = checkpoint.is_some();
752    let checkpoint = checkpoint.cloned().unwrap_or_default();
753    let (do_sample, temperature) = if let Some(do_sample) = overrides.do_sample {
754        if do_sample {
755            (
756                true,
757                overrides
758                    .temperature
759                    .or(checkpoint.temperature)
760                    .unwrap_or(1.0),
761            )
762        } else {
763            (false, 0.0)
764        }
765    } else if let Some(temperature) = overrides.temperature {
766        (temperature > 0.0, temperature)
767    } else if checkpoint.do_sample.unwrap_or(false) {
768        (true, checkpoint.temperature.unwrap_or(1.0))
769    } else {
770        (false, 0.0)
771    };
772    let resolved = ResolvedGenerationConfig {
773        do_sample,
774        temperature,
775        top_k: overrides
776            .top_k
777            .or(checkpoint.top_k)
778            .unwrap_or(if checkpoint_present { 50 } else { 40 }),
779        top_p: overrides
780            .top_p
781            .or(checkpoint.top_p)
782            .unwrap_or(if checkpoint_present { 1.0 } else { 0.95 }),
783        min_p: overrides
784            .min_p
785            .or(checkpoint.min_p)
786            .unwrap_or(if checkpoint_present { 0.0 } else { 0.05 }),
787        repetition_penalty: overrides
788            .repetition_penalty
789            .or(checkpoint.repetition_penalty)
790            .unwrap_or(1.0),
791        repeat_last_n: overrides
792            .repeat_last_n
793            .or(checkpoint.repeat_last_n)
794            .unwrap_or(64),
795        frequency_penalty: overrides
796            .frequency_penalty
797            .or(checkpoint.frequency_penalty)
798            .unwrap_or(0.0),
799        presence_penalty: overrides
800            .presence_penalty
801            .or(checkpoint.presence_penalty)
802            .unwrap_or(0.0),
803        max_new_tokens: overrides.max_new_tokens.or(checkpoint.max_new_tokens),
804    };
805    if !resolved.temperature.is_finite() || resolved.temperature < 0.0 {
806        return Err(GenerationError::InvalidTemperature(resolved.temperature));
807    }
808    if resolved.do_sample && resolved.temperature == 0.0 {
809        return Err(GenerationError::StochasticZeroTemperature);
810    }
811    if resolved.top_k < 0 {
812        return Err(GenerationError::InvalidTopK(resolved.top_k));
813    }
814    if !resolved.top_p.is_finite() || !(0.0..=1.0).contains(&resolved.top_p) {
815        return Err(GenerationError::InvalidTopP(resolved.top_p));
816    }
817    if !resolved.min_p.is_finite() || !(0.0..=1.0).contains(&resolved.min_p) {
818        return Err(GenerationError::InvalidMinP(resolved.min_p));
819    }
820    if !resolved.repetition_penalty.is_finite() || resolved.repetition_penalty <= 0.0 {
821        return Err(GenerationError::InvalidRepetitionPenalty(
822            resolved.repetition_penalty,
823        ));
824    }
825    if !resolved.frequency_penalty.is_finite() {
826        return Err(GenerationError::InvalidFrequencyPenalty(
827            resolved.frequency_penalty,
828        ));
829    }
830    if !resolved.presence_penalty.is_finite() {
831        return Err(GenerationError::InvalidPresencePenalty(
832            resolved.presence_penalty,
833        ));
834    }
835    if resolved.max_new_tokens == Some(0) {
836        return Err(GenerationError::ZeroTokenBudget);
837    }
838    Ok(resolved)
839}
840
841/// Semantic event emitted by generation orchestration.
842#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
843pub enum SemanticEvent {
844    /// Incremental reasoning content.
845    ReasoningDelta(String),
846    /// Incremental user-visible text.
847    TextDelta(String),
848    /// A structured tool call began.
849    ToolCallStart {
850        /// Zero-based tool-call position in the assistant turn.
851        index: usize,
852        /// Stable tool-call identifier.
853        id: String,
854        /// Tool name selected by the model.
855        name: String,
856    },
857    /// Incremental structured tool arguments.
858    ToolArgumentsDelta {
859        /// Zero-based tool-call position in the assistant turn.
860        index: usize,
861        /// A fragment of the tool call's JSON arguments.
862        json_fragment: String,
863    },
864    /// A structured tool call ended.
865    ToolCallEnd,
866    /// Generation terminated.
867    Finished {
868        /// The condition that ended the stream.
869        reason: FinishReason,
870    },
871}
872
873/// Invalid backend-neutral generation configuration or state transition.
874#[derive(Debug, Clone, PartialEq, thiserror::Error)]
875pub enum GenerationError {
876    /// A token was committed after termination.
877    #[error("generation has already finished")]
878    AlreadyFinished,
879    /// A prior prediction or its semantic delivery failed and requires recovery.
880    #[error("generation is fenced after a failed prediction or delivery")]
881    FailedGeneration,
882    /// A proposal transaction cannot be empty.
883    #[error("speculative verification requires at least one proposal")]
884    EmptyProposalBlock,
885    /// Proposal acceptance/rejection order was invalid.
886    #[error("invalid speculative verification state transition")]
887    InvalidSpeculativeTransition,
888    /// No terminal outcome or replacement/bonus tail resolved the round.
889    #[error("speculative verification round is not resolved")]
890    IncompleteSpeculativeRound,
891    /// Optimistic state was forked from a different prefix.
892    #[error("optimistic proposal prefix diverged from the canonical committed prefix")]
893    OptimisticPrefixDiverged,
894    /// An optimistic branch contained no proposals.
895    #[error("optimistic proposal branch is empty")]
896    EmptyOptimisticBranch,
897    /// A second optimistic branch was installed for one target transaction.
898    #[error("speculative verification already retains an optimistic branch")]
899    OptimisticBranchAlreadyPresent,
900    /// A speculative request identifier does not belong to this table.
901    #[error("unknown speculative request id {index}")]
902    UnknownSpeculativeRequest {
903        /// Requested stable insertion index.
904        index: usize,
905    },
906    /// A promoted proposal block no longer fits the canonical request budget.
907    #[error(
908        "promoted speculative block has {proposed} proposals but canonical capacity is {capacity}"
909    )]
910    ProposalCapacityExceeded {
911        /// Proposals retained by the promoted block.
912        proposed: usize,
913        /// Current canonical proposal capacity.
914        capacity: usize,
915    },
916    /// A speculative request table was consumed before reaching terminal state.
917    #[error("cannot finish a speculative scheduler with active requests")]
918    ActiveSpeculativeRequests,
919    /// A terminal speculative request did not retain its canonical finish reason.
920    #[error("completed speculative request {index} has no finish reason")]
921    MissingSpeculativeFinishReason {
922        /// Stable request insertion index.
923        index: usize,
924    },
925    /// No assistant proposal may be generated per round.
926    #[error("speculative max_draft_tokens must be positive")]
927    ZeroDraftTokens,
928    /// The selected backend cannot submit an assistant proposal.
929    #[error("speculative backend does not permit any draft tokens")]
930    NoBackendDraftCapacity,
931    /// Temperature is NaN, infinite, or negative.
932    #[error("temperature must be finite and non-negative, got {0}")]
933    InvalidTemperature(f32),
934    /// Stochastic sampling cannot use zero temperature.
935    #[error("do_sample=true requires a temperature greater than zero")]
936    StochasticZeroTemperature,
937    /// Mirostat V2 target surprise is non-finite or non-positive.
938    #[error("Mirostat V2 tau must be finite and positive, got {0}")]
939    InvalidMirostatTau(f32),
940    /// Mirostat V2 adaptation rate is non-finite or non-positive.
941    #[error("Mirostat V2 eta must be finite and positive, got {0}")]
942    InvalidMirostatEta(f32),
943    /// Top-k is negative.
944    #[error("top_k must be non-negative, got {0}")]
945    InvalidTopK(i32),
946    /// Top-p lies outside zero through one.
947    #[error("top_p must be between zero and one, got {0}")]
948    InvalidTopP(f32),
949    /// Min-p lies outside zero through one.
950    #[error("min_p must be between zero and one, got {0}")]
951    InvalidMinP(f32),
952    /// Repetition penalty is non-finite or non-positive.
953    #[error("repetition_penalty must be finite and positive, got {0}")]
954    InvalidRepetitionPenalty(f32),
955    /// Frequency penalty is non-finite.
956    #[error("frequency_penalty must be finite, got {0}")]
957    InvalidFrequencyPenalty(f32),
958    /// Presence penalty is non-finite.
959    #[error("presence_penalty must be finite, got {0}")]
960    InvalidPresencePenalty(f32),
961    /// An explicit token budget was zero.
962    #[error("max_new_tokens must be positive when supplied")]
963    ZeroTokenBudget,
964    /// Scheduler cannot retain any target transaction.
965    #[error("speculative max_in_flight_verifications must be positive")]
966    ZeroInFlightVerifications,
967    /// Current scheduler supports no more than one lookahead block.
968    #[error("speculative scheduler currently supports at most one lookahead block")]
969    TooManyLookaheadBlocks,
970    /// Lookahead was enabled with no branch capacity.
971    #[error("speculative lookahead requires at least one optimistic branch slot")]
972    LookaheadWithoutBranchCapacity,
973    /// Adaptive lookahead needs a non-zero observation window.
974    #[error("speculative adaptive_lookahead_min_blocks must be positive")]
975    ZeroAdaptiveLookaheadWindow,
976    /// Exact verification must always have a positive completion deadline.
977    #[error("speculative completion timeout must be positive")]
978    ZeroSpeculativeCompletionTimeout,
979    /// The host duration cannot be represented by the portable millisecond field.
980    #[error("speculative completion timeout is too large")]
981    SpeculativeCompletionTimeoutTooLarge,
982    /// Active requests expose no legal scheduler action.
983    #[error("speculative scheduler reached a non-terminal state with no eligible operation")]
984    StalledSpeculativeSchedule,
985    /// The requested speculative lifecycle edge is invalid.
986    #[error("invalid speculative request status transition from {from:?} to {to:?}")]
987    InvalidSpeculativeStatusTransition {
988        /// Current status.
989        from: SpeculativeRequestStatus,
990        /// Requested status.
991        to: SpeculativeRequestStatus,
992    },
993}
994
995#[cfg(test)]
996mod tests {
997    use super::*;
998
999    #[test]
1000    fn terminal_precedence_and_cancellation_are_canonical() {
1001        let mut sequence = GenerationSequence::new(2, [7]);
1002        let first = sequence.commit(1, TokenTerminalSignals::default()).unwrap();
1003        assert_eq!(first.finish_reason, None);
1004        let second = sequence
1005            .commit(
1006                7,
1007                TokenTerminalSignals {
1008                    stop_sequence: true,
1009                    grammar_complete: true,
1010                },
1011            )
1012            .unwrap();
1013        assert_eq!(second.finish_reason, Some(FinishReason::StopSequence));
1014        assert!(!sequence.cancel());
1015
1016        let token = GenerationCancellationToken::new();
1017        let mut active = GenerationSequence::new(3, []);
1018        token.cancel();
1019        assert!(active.observe_cancellation(&token));
1020        assert_eq!(active.finish_reason(), Some(FinishReason::Cancelled));
1021    }
1022
1023    #[test]
1024    fn speculative_commit_plan_preserves_trailing_token_cache_semantics() {
1025        let mut rejected = SpeculativeRound::new(3).unwrap();
1026        rejected.accept(10, false).unwrap();
1027        rejected.reject_with(99, false).unwrap();
1028        let plan = rejected.commit_plan().unwrap();
1029        assert_eq!(plan.accepted_proposals, 1);
1030        assert_eq!(plan.committed_tokens, &[10, 99]);
1031        assert_eq!(plan.verified_inputs, 2);
1032        assert_eq!(plan.tail, Some(SpeculativeTail::Replacement));
1033
1034        let mut terminal_accept = SpeculativeRound::new(2).unwrap();
1035        terminal_accept.accept(10, false).unwrap();
1036        terminal_accept.accept(11, true).unwrap();
1037        let plan = terminal_accept.commit_plan().unwrap();
1038        assert!(plan.full_acceptance);
1039        assert_eq!(plan.verified_inputs, 2);
1040        assert_eq!(plan.tail, None);
1041
1042        let mut bonus = SpeculativeRound::new(2).unwrap();
1043        bonus.accept(10, false).unwrap();
1044        bonus.accept(11, false).unwrap();
1045        bonus.bonus(12, false).unwrap();
1046        assert_eq!(bonus.commit_plan().unwrap().verified_inputs, 3);
1047
1048        let mut incomplete = SpeculativeRound::new(2).unwrap();
1049        incomplete.accept(10, false).unwrap();
1050        assert!(matches!(
1051            incomplete.commit_plan(),
1052            Err(GenerationError::IncompleteSpeculativeRound)
1053        ));
1054    }
1055
1056    #[test]
1057    fn optimistic_reuse_is_pure_and_fail_closed() {
1058        assert_eq!(
1059            resolve_optimistic_reuse(&[1], &[1], &[2, 3], 2, false).unwrap(),
1060            OptimisticReuseDecision::MatchedRetained
1061        );
1062        assert_eq!(
1063            resolve_optimistic_reuse(&[1], &[1], &[2], 2, false).unwrap(),
1064            OptimisticReuseDecision::MatchedConsumed
1065        );
1066        assert_eq!(
1067            resolve_optimistic_reuse(&[1], &[1], &[2], 9, false).unwrap(),
1068            OptimisticReuseDecision::DiscardMismatch
1069        );
1070        assert!(matches!(
1071            resolve_optimistic_reuse(&[1], &[9], &[2], 2, false),
1072            Err(GenerationError::OptimisticPrefixDiverged)
1073        ));
1074    }
1075
1076    #[test]
1077    fn sampler_and_scheduler_configuration_validate_without_a_backend() {
1078        let checkpoint = CheckpointGenerationConfig {
1079            do_sample: Some(true),
1080            temperature: Some(0.8),
1081            top_k: Some(64),
1082            repetition_penalty: Some(1.1),
1083            ..CheckpointGenerationConfig::default()
1084        };
1085        let resolved =
1086            resolve_generation_config(Some(&checkpoint), GenerationConfigOverrides::default())
1087                .unwrap();
1088        assert!(resolved.do_sample);
1089        assert_eq!(resolved.top_k, 64);
1090        assert_eq!(resolved.repetition_penalty, 1.1);
1091        assert_eq!(resolved.repeat_last_n, 64);
1092        assert!(matches!(
1093            resolve_generation_config(
1094                None,
1095                GenerationConfigOverrides {
1096                    frequency_penalty: Some(f32::NAN),
1097                    ..GenerationConfigOverrides::default()
1098                }
1099            ),
1100            Err(GenerationError::InvalidFrequencyPenalty(value)) if value.is_nan()
1101        ));
1102        assert!(SpeculativeConfig::default().validate().is_ok());
1103        assert!(SpeculativeSchedulerOptions::default().validate().is_ok());
1104        assert!(matches!(
1105            SpeculativeSchedulerOptions {
1106                max_in_flight_verifications: 0,
1107                ..SpeculativeSchedulerOptions::default()
1108            }
1109            .validate(),
1110            Err(GenerationError::ZeroInFlightVerifications)
1111        ));
1112
1113        let config_json = serde_json::to_string(&resolved).unwrap();
1114        assert_eq!(
1115            serde_json::from_str::<ResolvedGenerationConfig>(&config_json).unwrap(),
1116            resolved
1117        );
1118        let options = SpeculativeSchedulerOptions::default();
1119        let options_json = serde_json::to_string(&options).unwrap();
1120        assert_eq!(
1121            serde_json::from_str::<SpeculativeSchedulerOptions>(&options_json).unwrap(),
1122            options
1123        );
1124    }
1125
1126    #[test]
1127    fn semantic_events_round_trip_without_a_backend() {
1128        let event = SemanticEvent::ToolCallStart {
1129            index: 2,
1130            id: "call_2".into(),
1131            name: "lookup".into(),
1132        };
1133        let json = serde_json::to_string(&event).unwrap();
1134        assert_eq!(serde_json::from_str::<SemanticEvent>(&json).unwrap(), event);
1135
1136        let mut zero_budget = GenerationSequence::new(0, []);
1137        assert_eq!(zero_budget.finish_reason(), Some(FinishReason::MaxTokens));
1138        assert!(zero_budget.cancel());
1139        assert_eq!(zero_budget.finish_reason(), Some(FinishReason::Cancelled));
1140    }
1141
1142    #[test]
1143    fn speculative_request_lifecycle_defers_cancellation_exactly() {
1144        let mut lifecycle = SpeculativeRequestLifecycle::new();
1145        lifecycle
1146            .transition(SpeculativeRequestStatus::ReadyToDraft)
1147            .unwrap();
1148        lifecycle
1149            .transition(SpeculativeRequestStatus::ReadyToSubmitVerification)
1150            .unwrap();
1151        lifecycle
1152            .transition(SpeculativeRequestStatus::TargetVerificationInFlight)
1153            .unwrap();
1154        assert_eq!(
1155            lifecycle.request_cancellation(true).unwrap(),
1156            SpeculativeCancellationDisposition::Deferred
1157        );
1158        assert!(lifecycle.cancellation_pending());
1159        lifecycle
1160            .transition(SpeculativeRequestStatus::VerificationResolution)
1161            .unwrap();
1162        lifecycle
1163            .transition(SpeculativeRequestStatus::Cancelled)
1164            .unwrap();
1165        assert!(lifecycle.is_terminal());
1166        assert!(!lifecycle.cancellation_pending());
1167        assert!(matches!(
1168            lifecycle.transition(SpeculativeRequestStatus::ReadyToDraft),
1169            Err(GenerationError::InvalidSpeculativeStatusTransition { .. })
1170        ));
1171    }
1172}