1use crate::backend::{
4 BoundedCompletionWait, BoundedCompletionWaitError, CompletionCancellationMode,
5};
6use serde::{Deserialize, Serialize};
7use std::sync::{
8 atomic::{AtomicBool, Ordering},
9 Arc,
10};
11
12#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum FinishReason {
16 Eos,
18 StopSequence,
20 GrammarComplete,
22 MaxTokens,
24 Cancelled,
26}
27
28#[derive(Debug, Clone, Default)]
30pub struct GenerationCancellationToken {
31 cancelled: Arc<AtomicBool>,
32}
33
34impl GenerationCancellationToken {
35 pub fn new() -> Self {
37 Self::default()
38 }
39
40 pub fn cancel(&self) {
42 self.cancelled.store(true, Ordering::Release);
43 }
44
45 pub fn is_cancelled(&self) -> bool {
47 self.cancelled.load(Ordering::Acquire)
48 }
49}
50
51#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)]
53pub struct TokenTerminalSignals {
54 pub stop_sequence: bool,
56 pub grammar_complete: bool,
58}
59
60#[derive(Debug, Clone, Copy, Eq, PartialEq)]
62pub struct TokenCommit {
63 pub token_id: u32,
65 pub position: usize,
67 pub finish_reason: Option<FinishReason>,
69}
70
71#[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 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 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 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 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 pub fn observe_cancellation(&mut self, cancellation: &GenerationCancellationToken) -> bool {
156 cancellation.is_cancelled() && self.cancel()
157 }
158
159 pub fn tokens(&self) -> &[u32] {
161 &self.tokens
162 }
163
164 pub fn into_tokens(self) -> Vec<u32> {
166 self.tokens
167 }
168
169 pub fn remaining(&self) -> usize {
171 self.max_tokens.saturating_sub(self.tokens.len())
172 }
173
174 pub const fn max_tokens(&self) -> usize {
176 self.max_tokens
177 }
178
179 pub const fn finish_reason(&self) -> Option<FinishReason> {
181 self.finish_reason
182 }
183
184 pub const fn is_finished(&self) -> bool {
186 self.finish_reason.is_some()
187 }
188}
189
190#[derive(Debug, Clone, Copy, Eq, PartialEq)]
192pub enum SpeculativeTail {
193 Replacement,
195 Bonus,
197}
198
199#[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 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 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 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 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 pub const fn is_full_acceptance(&self) -> bool {
259 self.accepted == self.proposal_count
260 }
261
262 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#[derive(Debug, Clone, Copy, Eq, PartialEq)]
284pub struct SpeculativeCommitPlan<'a> {
285 pub accepted_proposals: usize,
287 pub committed_tokens: &'a [u32],
289 pub verified_inputs: usize,
291 pub full_acceptance: bool,
293 pub tail: Option<SpeculativeTail>,
295 pub terminal: bool,
297}
298
299#[derive(Debug, Clone, Copy, Eq, PartialEq)]
301pub enum OptimisticReuseDecision {
302 DiscardTerminal,
304 DiscardMismatch,
306 MatchedConsumed,
308 MatchedRetained,
310}
311
312pub 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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
341pub struct SpeculativeConfig {
342 pub max_tokens: usize,
344 pub max_draft_tokens: usize,
346 pub temperature: f32,
348 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
378#[serde(default)]
379pub struct SpeculativeSchedulerOptions {
380 pub max_in_flight_verifications: usize,
382 pub max_optimistic_branches: usize,
384 pub lookahead_blocks: usize,
386 pub adaptive_lookahead: bool,
388 pub adaptive_lookahead_min_blocks: usize,
390 pub completion_timeout_milliseconds: u64,
392 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 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 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 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 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#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
473pub struct SpeculativeRequestId(usize);
474
475impl SpeculativeRequestId {
476 pub const fn new(index: usize) -> Self {
478 Self(index)
479 }
480
481 pub const fn index(self) -> usize {
483 self.0
484 }
485}
486
487#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
489#[serde(rename_all = "snake_case")]
490#[non_exhaustive]
491pub enum SpeculativeRequestStatus {
492 Prefill,
494 ReadyToDraft,
496 ReadyToSubmitVerification,
498 TargetVerificationInFlight,
500 OptimisticDraftRunning,
502 OptimisticDraftReady,
504 VerificationResolution,
506 Completed,
508 Cancelled,
510}
511
512#[derive(Debug, Clone, Copy, Eq, PartialEq)]
514pub enum SpeculativeCancellationDisposition {
515 AlreadyTerminal,
517 CancelNow,
519 Deferred,
521}
522
523#[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 pub const fn new() -> Self {
539 Self {
540 status: SpeculativeRequestStatus::Prefill,
541 cancellation_pending: false,
542 }
543 }
544
545 pub const fn completed() -> Self {
547 Self {
548 status: SpeculativeRequestStatus::Completed,
549 cancellation_pending: false,
550 }
551 }
552
553 pub const fn cancelled() -> Self {
555 Self {
556 status: SpeculativeRequestStatus::Cancelled,
557 cancellation_pending: false,
558 }
559 }
560
561 pub const fn status(&self) -> SpeculativeRequestStatus {
563 self.status
564 }
565
566 pub const fn cancellation_pending(&self) -> bool {
568 self.cancellation_pending
569 }
570
571 pub const fn is_terminal(&self) -> bool {
573 matches!(
574 self.status,
575 SpeculativeRequestStatus::Completed | SpeculativeRequestStatus::Cancelled
576 )
577 }
578
579 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 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#[derive(Debug, Clone, Default, PartialEq, Deserialize, Serialize)]
663pub struct CheckpointGenerationConfig {
664 #[serde(default)]
666 pub do_sample: Option<bool>,
667 #[serde(default)]
669 pub temperature: Option<f32>,
670 #[serde(default)]
672 pub top_k: Option<i32>,
673 #[serde(default)]
675 pub top_p: Option<f32>,
676 #[serde(default)]
678 pub min_p: Option<f32>,
679 #[serde(default)]
681 pub repetition_penalty: Option<f32>,
682 #[serde(default)]
684 pub repeat_last_n: Option<i32>,
685 #[serde(default)]
687 pub frequency_penalty: Option<f32>,
688 #[serde(default)]
690 pub presence_penalty: Option<f32>,
691 #[serde(default)]
693 pub max_new_tokens: Option<usize>,
694}
695
696#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
698pub struct GenerationConfigOverrides {
699 pub do_sample: Option<bool>,
701 pub temperature: Option<f32>,
703 pub top_k: Option<i32>,
705 pub top_p: Option<f32>,
707 pub min_p: Option<f32>,
709 pub repetition_penalty: Option<f32>,
711 pub repeat_last_n: Option<i32>,
713 pub frequency_penalty: Option<f32>,
715 pub presence_penalty: Option<f32>,
717 pub max_new_tokens: Option<usize>,
719}
720
721#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
723pub struct ResolvedGenerationConfig {
724 pub do_sample: bool,
726 pub temperature: f32,
728 pub top_k: i32,
730 pub top_p: f32,
732 pub min_p: f32,
734 pub repetition_penalty: f32,
736 pub repeat_last_n: i32,
738 pub frequency_penalty: f32,
740 pub presence_penalty: f32,
742 pub max_new_tokens: Option<usize>,
744}
745
746pub 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#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
843pub enum SemanticEvent {
844 ReasoningDelta(String),
846 TextDelta(String),
848 ToolCallStart {
850 index: usize,
852 id: String,
854 name: String,
856 },
857 ToolArgumentsDelta {
859 index: usize,
861 json_fragment: String,
863 },
864 ToolCallEnd,
866 Finished {
868 reason: FinishReason,
870 },
871}
872
873#[derive(Debug, Clone, PartialEq, thiserror::Error)]
875pub enum GenerationError {
876 #[error("generation has already finished")]
878 AlreadyFinished,
879 #[error("generation is fenced after a failed prediction or delivery")]
881 FailedGeneration,
882 #[error("speculative verification requires at least one proposal")]
884 EmptyProposalBlock,
885 #[error("invalid speculative verification state transition")]
887 InvalidSpeculativeTransition,
888 #[error("speculative verification round is not resolved")]
890 IncompleteSpeculativeRound,
891 #[error("optimistic proposal prefix diverged from the canonical committed prefix")]
893 OptimisticPrefixDiverged,
894 #[error("optimistic proposal branch is empty")]
896 EmptyOptimisticBranch,
897 #[error("speculative verification already retains an optimistic branch")]
899 OptimisticBranchAlreadyPresent,
900 #[error("unknown speculative request id {index}")]
902 UnknownSpeculativeRequest {
903 index: usize,
905 },
906 #[error(
908 "promoted speculative block has {proposed} proposals but canonical capacity is {capacity}"
909 )]
910 ProposalCapacityExceeded {
911 proposed: usize,
913 capacity: usize,
915 },
916 #[error("cannot finish a speculative scheduler with active requests")]
918 ActiveSpeculativeRequests,
919 #[error("completed speculative request {index} has no finish reason")]
921 MissingSpeculativeFinishReason {
922 index: usize,
924 },
925 #[error("speculative max_draft_tokens must be positive")]
927 ZeroDraftTokens,
928 #[error("speculative backend does not permit any draft tokens")]
930 NoBackendDraftCapacity,
931 #[error("temperature must be finite and non-negative, got {0}")]
933 InvalidTemperature(f32),
934 #[error("do_sample=true requires a temperature greater than zero")]
936 StochasticZeroTemperature,
937 #[error("Mirostat V2 tau must be finite and positive, got {0}")]
939 InvalidMirostatTau(f32),
940 #[error("Mirostat V2 eta must be finite and positive, got {0}")]
942 InvalidMirostatEta(f32),
943 #[error("top_k must be non-negative, got {0}")]
945 InvalidTopK(i32),
946 #[error("top_p must be between zero and one, got {0}")]
948 InvalidTopP(f32),
949 #[error("min_p must be between zero and one, got {0}")]
951 InvalidMinP(f32),
952 #[error("repetition_penalty must be finite and positive, got {0}")]
954 InvalidRepetitionPenalty(f32),
955 #[error("frequency_penalty must be finite, got {0}")]
957 InvalidFrequencyPenalty(f32),
958 #[error("presence_penalty must be finite, got {0}")]
960 InvalidPresencePenalty(f32),
961 #[error("max_new_tokens must be positive when supplied")]
963 ZeroTokenBudget,
964 #[error("speculative max_in_flight_verifications must be positive")]
966 ZeroInFlightVerifications,
967 #[error("speculative scheduler currently supports at most one lookahead block")]
969 TooManyLookaheadBlocks,
970 #[error("speculative lookahead requires at least one optimistic branch slot")]
972 LookaheadWithoutBranchCapacity,
973 #[error("speculative adaptive_lookahead_min_blocks must be positive")]
975 ZeroAdaptiveLookaheadWindow,
976 #[error("speculative completion timeout must be positive")]
978 ZeroSpeculativeCompletionTimeout,
979 #[error("speculative completion timeout is too large")]
981 SpeculativeCompletionTimeoutTooLarge,
982 #[error("speculative scheduler reached a non-terminal state with no eligible operation")]
984 StalledSpeculativeSchedule,
985 #[error("invalid speculative request status transition from {from:?} to {to:?}")]
987 InvalidSpeculativeStatusTransition {
988 from: SpeculativeRequestStatus,
990 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}