Skip to main content

eredu_core/
generation.rs

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