1use serde::{Deserialize, Serialize};
4use std::sync::{
5 atomic::{AtomicBool, Ordering},
6 Arc,
7};
8
9#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum FinishReason {
13 Eos,
15 StopSequence,
17 GrammarComplete,
19 MaxTokens,
21 Cancelled,
23}
24
25#[derive(Debug, Clone, Default)]
27pub struct GenerationCancellationToken {
28 cancelled: Arc<AtomicBool>,
29}
30
31impl GenerationCancellationToken {
32 pub fn new() -> Self {
34 Self::default()
35 }
36
37 pub fn cancel(&self) {
39 self.cancelled.store(true, Ordering::Release);
40 }
41
42 pub fn is_cancelled(&self) -> bool {
44 self.cancelled.load(Ordering::Acquire)
45 }
46}
47
48#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)]
50pub struct TokenTerminalSignals {
51 pub stop_sequence: bool,
53 pub grammar_complete: bool,
55}
56
57#[derive(Debug, Clone, Copy, Eq, PartialEq)]
59pub struct TokenCommit {
60 pub token_id: u32,
62 pub position: usize,
64 pub finish_reason: Option<FinishReason>,
66}
67
68#[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 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 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 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 pub fn observe_cancellation(&mut self, cancellation: &GenerationCancellationToken) -> bool {
140 cancellation.is_cancelled() && self.cancel()
141 }
142
143 pub fn tokens(&self) -> &[u32] {
145 &self.tokens
146 }
147
148 pub fn into_tokens(self) -> Vec<u32> {
150 self.tokens
151 }
152
153 pub fn remaining(&self) -> usize {
155 self.max_tokens.saturating_sub(self.tokens.len())
156 }
157
158 pub const fn max_tokens(&self) -> usize {
160 self.max_tokens
161 }
162
163 pub const fn finish_reason(&self) -> Option<FinishReason> {
165 self.finish_reason
166 }
167
168 pub const fn is_finished(&self) -> bool {
170 self.finish_reason.is_some()
171 }
172}
173
174#[derive(Debug, Clone, Copy, Eq, PartialEq)]
176pub enum SpeculativeTail {
177 Replacement,
179 Bonus,
181}
182
183#[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 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 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 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 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 pub const fn is_full_acceptance(&self) -> bool {
243 self.accepted == self.proposal_count
244 }
245
246 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#[derive(Debug, Clone, Copy, Eq, PartialEq)]
268pub struct SpeculativeCommitPlan<'a> {
269 pub accepted_proposals: usize,
271 pub committed_tokens: &'a [u32],
273 pub verified_inputs: usize,
275 pub full_acceptance: bool,
277 pub tail: Option<SpeculativeTail>,
279 pub terminal: bool,
281}
282
283#[derive(Debug, Clone, Copy, Eq, PartialEq)]
285pub enum OptimisticReuseDecision {
286 DiscardTerminal,
288 DiscardMismatch,
290 MatchedConsumed,
292 MatchedRetained,
294}
295
296pub 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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
325pub struct SpeculativeConfig {
326 pub max_tokens: usize,
328 pub max_draft_tokens: usize,
330 pub temperature: f32,
332 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
362pub struct SpeculativeSchedulerOptions {
363 pub max_in_flight_verifications: usize,
365 pub max_optimistic_branches: usize,
367 pub lookahead_blocks: usize,
369 pub adaptive_lookahead: bool,
371 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 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 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#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
420pub struct SpeculativeRequestId(usize);
421
422impl SpeculativeRequestId {
423 pub const fn new(index: usize) -> Self {
425 Self(index)
426 }
427
428 pub const fn index(self) -> usize {
430 self.0
431 }
432}
433
434#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
436#[serde(rename_all = "snake_case")]
437#[non_exhaustive]
438pub enum SpeculativeRequestStatus {
439 Prefill,
441 ReadyToDraft,
443 ReadyToSubmitVerification,
445 TargetVerificationInFlight,
447 OptimisticDraftRunning,
449 OptimisticDraftReady,
451 VerificationResolution,
453 Completed,
455 Cancelled,
457}
458
459#[derive(Debug, Clone, Copy, Eq, PartialEq)]
461pub enum SpeculativeCancellationDisposition {
462 AlreadyTerminal,
464 CancelNow,
466 Deferred,
468}
469
470#[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 pub const fn new() -> Self {
486 Self {
487 status: SpeculativeRequestStatus::Prefill,
488 cancellation_pending: false,
489 }
490 }
491
492 pub const fn completed() -> Self {
494 Self {
495 status: SpeculativeRequestStatus::Completed,
496 cancellation_pending: false,
497 }
498 }
499
500 pub const fn cancelled() -> Self {
502 Self {
503 status: SpeculativeRequestStatus::Cancelled,
504 cancellation_pending: false,
505 }
506 }
507
508 pub const fn status(&self) -> SpeculativeRequestStatus {
510 self.status
511 }
512
513 pub const fn cancellation_pending(&self) -> bool {
515 self.cancellation_pending
516 }
517
518 pub const fn is_terminal(&self) -> bool {
520 matches!(
521 self.status,
522 SpeculativeRequestStatus::Completed | SpeculativeRequestStatus::Cancelled
523 )
524 }
525
526 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 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#[derive(Debug, Clone, Default, PartialEq, Deserialize, Serialize)]
610pub struct CheckpointGenerationConfig {
611 #[serde(default)]
613 pub do_sample: Option<bool>,
614 #[serde(default)]
616 pub temperature: Option<f32>,
617 #[serde(default)]
619 pub top_k: Option<i32>,
620 #[serde(default)]
622 pub top_p: Option<f32>,
623 #[serde(default)]
625 pub min_p: Option<f32>,
626 #[serde(default)]
628 pub repetition_penalty: Option<f32>,
629 #[serde(default)]
631 pub repeat_last_n: Option<i32>,
632 #[serde(default)]
634 pub frequency_penalty: Option<f32>,
635 #[serde(default)]
637 pub presence_penalty: Option<f32>,
638 #[serde(default)]
640 pub max_new_tokens: Option<usize>,
641}
642
643#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
645pub struct GenerationConfigOverrides {
646 pub do_sample: Option<bool>,
648 pub temperature: Option<f32>,
650 pub top_k: Option<i32>,
652 pub top_p: Option<f32>,
654 pub min_p: Option<f32>,
656 pub repetition_penalty: Option<f32>,
658 pub repeat_last_n: Option<i32>,
660 pub frequency_penalty: Option<f32>,
662 pub presence_penalty: Option<f32>,
664 pub max_new_tokens: Option<usize>,
666}
667
668#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
670pub struct ResolvedGenerationConfig {
671 pub do_sample: bool,
673 pub temperature: f32,
675 pub top_k: i32,
677 pub top_p: f32,
679 pub min_p: f32,
681 pub repetition_penalty: f32,
683 pub repeat_last_n: i32,
685 pub frequency_penalty: f32,
687 pub presence_penalty: f32,
689 pub max_new_tokens: Option<usize>,
691}
692
693pub 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#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
790pub enum SemanticEvent {
791 ReasoningDelta(String),
793 TextDelta(String),
795 ToolCallStart {
797 index: usize,
799 id: String,
801 name: String,
803 },
804 ToolArgumentsDelta {
806 index: usize,
808 json_fragment: String,
810 },
811 ToolCallEnd,
813 Finished {
815 reason: FinishReason,
817 },
818}
819
820#[derive(Debug, Clone, PartialEq, thiserror::Error)]
822pub enum GenerationError {
823 #[error("generation has already finished")]
825 AlreadyFinished,
826 #[error("speculative verification requires at least one proposal")]
828 EmptyProposalBlock,
829 #[error("invalid speculative verification state transition")]
831 InvalidSpeculativeTransition,
832 #[error("speculative verification round is not resolved")]
834 IncompleteSpeculativeRound,
835 #[error("optimistic proposal prefix diverged from the canonical committed prefix")]
837 OptimisticPrefixDiverged,
838 #[error("optimistic proposal branch is empty")]
840 EmptyOptimisticBranch,
841 #[error("speculative verification already retains an optimistic branch")]
843 OptimisticBranchAlreadyPresent,
844 #[error("unknown speculative request id {index}")]
846 UnknownSpeculativeRequest {
847 index: usize,
849 },
850 #[error(
852 "promoted speculative block has {proposed} proposals but canonical capacity is {capacity}"
853 )]
854 ProposalCapacityExceeded {
855 proposed: usize,
857 capacity: usize,
859 },
860 #[error("cannot finish a speculative scheduler with active requests")]
862 ActiveSpeculativeRequests,
863 #[error("completed speculative request {index} has no finish reason")]
865 MissingSpeculativeFinishReason {
866 index: usize,
868 },
869 #[error("speculative max_draft_tokens must be positive")]
871 ZeroDraftTokens,
872 #[error("speculative backend does not permit any draft tokens")]
874 NoBackendDraftCapacity,
875 #[error("temperature must be finite and non-negative, got {0}")]
877 InvalidTemperature(f32),
878 #[error("do_sample=true requires a temperature greater than zero")]
880 StochasticZeroTemperature,
881 #[error("Mirostat V2 tau must be finite and positive, got {0}")]
883 InvalidMirostatTau(f32),
884 #[error("Mirostat V2 eta must be finite and positive, got {0}")]
886 InvalidMirostatEta(f32),
887 #[error("top_k must be non-negative, got {0}")]
889 InvalidTopK(i32),
890 #[error("top_p must be between zero and one, got {0}")]
892 InvalidTopP(f32),
893 #[error("min_p must be between zero and one, got {0}")]
895 InvalidMinP(f32),
896 #[error("repetition_penalty must be finite and positive, got {0}")]
898 InvalidRepetitionPenalty(f32),
899 #[error("frequency_penalty must be finite, got {0}")]
901 InvalidFrequencyPenalty(f32),
902 #[error("presence_penalty must be finite, got {0}")]
904 InvalidPresencePenalty(f32),
905 #[error("max_new_tokens must be positive when supplied")]
907 ZeroTokenBudget,
908 #[error("speculative max_in_flight_verifications must be positive")]
910 ZeroInFlightVerifications,
911 #[error("speculative scheduler currently supports at most one lookahead block")]
913 TooManyLookaheadBlocks,
914 #[error("speculative lookahead requires at least one optimistic branch slot")]
916 LookaheadWithoutBranchCapacity,
917 #[error("speculative adaptive_lookahead_min_blocks must be positive")]
919 ZeroAdaptiveLookaheadWindow,
920 #[error("speculative scheduler reached a non-terminal state with no eligible operation")]
922 StalledSpeculativeSchedule,
923 #[error("invalid speculative request status transition from {from:?} to {to:?}")]
925 InvalidSpeculativeStatusTransition {
926 from: SpeculativeRequestStatus,
928 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}