1use core::fmt;
41
42pub const TALKER_KV_VALUES_PER_TOKEN: u64 = 57_344;
47
48pub const MICRODECODER_KV_BYTES: u64 = 320 * 1024;
50
51pub const CODEC_DECODER_KV_BYTES: u64 = 2_359_296;
53
54pub const MAX_CONTEXT_TOKENS: u64 = 32_768;
59
60#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
62pub enum KvDtype {
63 Bf16,
65 F32,
67}
68
69impl KvDtype {
70 #[must_use]
72 pub const fn size_bytes(self) -> u64 {
73 match self {
74 Self::Bf16 => 2,
75 Self::F32 => 4,
76 }
77 }
78
79 #[must_use]
81 pub const fn as_str(self) -> &'static str {
82 match self {
83 Self::Bf16 => "bf16",
84 Self::F32 => "f32",
85 }
86 }
87}
88
89impl fmt::Display for KvDtype {
90 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
91 f.write_str(self.as_str())
92 }
93}
94
95#[derive(Clone, Copy, Debug, PartialEq, Eq)]
100pub enum BindingConstraint {
101 FrameCap,
103 ContextCeiling,
105 TextHeuristic,
112}
113
114impl BindingConstraint {
115 #[must_use]
117 pub const fn as_str(self) -> &'static str {
118 match self {
119 Self::FrameCap => "frame_cap",
120 Self::ContextCeiling => "context_ceiling",
121 Self::TextHeuristic => "text_heuristic",
122 }
123 }
124}
125
126pub const HEURISTIC_FRAMES_PER_PROMPT_TOKEN: u64 = 4;
132
133pub const HEURISTIC_FRAME_HEADROOM: u64 = 64;
135
136#[derive(Clone, Copy, Debug, PartialEq, Eq)]
138pub struct AdmissionRequest {
139 pub prompt_tokens: u64,
141 pub max_new_tokens: u64,
143 pub heuristic_eos_backstop: bool,
148 pub kv_dtype: KvDtype,
150 pub ring_buffer_bytes: u64,
152 pub weights_resident_bytes: u64,
154 pub budget_bytes: u64,
156}
157
158#[derive(Clone, Copy, Debug, PartialEq, Eq)]
163pub struct AdmissionPlan {
164 pub predicted_max_frames: u64,
166 pub binding_constraint: BindingConstraint,
168 pub kv_talker_bytes: u64,
170 pub bounded_state_bytes: u64,
172 pub weights_resident_bytes: u64,
174 pub predicted_peak_bytes: u64,
176 pub budget_bytes: u64,
178}
179
180impl AdmissionPlan {
181 #[must_use]
183 pub const fn shortfall_bytes(&self) -> u64 {
184 self.predicted_peak_bytes.saturating_sub(self.budget_bytes)
185 }
186
187 #[must_use]
189 pub const fn fits(&self) -> bool {
190 self.predicted_peak_bytes <= self.budget_bytes
191 }
192}
193
194#[derive(Clone, Copy, Debug, PartialEq, Eq)]
196pub enum AdmissionRejection {
197 PromptExceedsContext {
201 prompt_tokens: u64,
203 ceiling: u64,
205 },
206 NoFramesRequested,
208 BudgetExceeded {
210 plan: AdmissionPlan,
212 },
213 Overflow {
219 term: &'static str,
221 },
222}
223
224impl fmt::Display for AdmissionRejection {
225 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
226 match self {
227 Self::PromptExceedsContext {
228 prompt_tokens,
229 ceiling,
230 } => write!(
231 f,
232 "prompt is {prompt_tokens} tokens but the context ceiling is {ceiling}; \
233 no frames could be generated. Shorten the text or chunk it — raising the memory \
234 budget cannot help"
235 ),
236 Self::NoFramesRequested => {
237 f.write_str("max_new_tokens is 0; there is nothing to synthesize")
238 }
239 Self::BudgetExceeded { plan } => write!(
240 f,
241 "predicted peak {} bytes exceeds the {} byte budget by {} \
242 (talker KV {} over {} frames, bounded state {}, weights {}; binding constraint: {}). \
243 Rejected before allocating, so nothing was committed",
244 plan.predicted_peak_bytes,
245 plan.budget_bytes,
246 plan.shortfall_bytes(),
247 plan.kv_talker_bytes,
248 plan.predicted_max_frames,
249 plan.bounded_state_bytes,
250 plan.weights_resident_bytes,
251 plan.binding_constraint.as_str(),
252 ),
253 Self::Overflow { term } => write!(
254 f,
255 "admission arithmetic overflowed computing `{term}`; refusing rather than \
256 admitting on a wrapped total"
257 ),
258 }
259 }
260}
261
262impl core::error::Error for AdmissionRejection {}
263
264pub fn talker_kv_bytes(
270 prompt_tokens: u64,
271 frames: u64,
272 dtype: KvDtype,
273) -> Result<u64, AdmissionRejection> {
274 prompt_tokens
275 .checked_add(frames)
276 .and_then(|tokens| tokens.checked_mul(TALKER_KV_VALUES_PER_TOKEN))
277 .and_then(|values| values.checked_mul(dtype.size_bytes()))
278 .ok_or(AdmissionRejection::Overflow {
279 term: "talker_kv_bytes",
280 })
281}
282
283#[derive(Clone, Copy, Debug, PartialEq, Eq)]
290pub struct AdmissionPolicy {
291 pub budget_bytes: u64,
293 pub max_new_tokens: u64,
295 pub heuristic_eos_backstop: bool,
298 pub kv_dtype: KvDtype,
300 pub ring_buffer_bytes: u64,
302 pub weights_resident_bytes: u64,
304}
305
306pub const DEFAULT_BUDGET_BYTES: u64 = 2 * 1024 * 1024 * 1024;
311
312pub const DEFAULT_MAX_NEW_TOKENS: u64 = 8_192;
314
315impl Default for AdmissionPolicy {
316 fn default() -> Self {
317 Self {
318 budget_bytes: DEFAULT_BUDGET_BYTES,
319 max_new_tokens: DEFAULT_MAX_NEW_TOKENS,
320 heuristic_eos_backstop: true,
321 kv_dtype: KvDtype::Bf16,
322 ring_buffer_bytes: 0,
326 weights_resident_bytes: 0,
327 }
328 }
329}
330
331impl AdmissionPolicy {
332 #[must_use]
334 pub const fn request_for(&self, prompt_tokens: u64) -> AdmissionRequest {
335 AdmissionRequest {
336 prompt_tokens,
337 max_new_tokens: self.max_new_tokens,
338 heuristic_eos_backstop: self.heuristic_eos_backstop,
339 kv_dtype: self.kv_dtype,
340 ring_buffer_bytes: self.ring_buffer_bytes,
341 weights_resident_bytes: self.weights_resident_bytes,
342 budget_bytes: self.budget_bytes,
343 }
344 }
345
346 pub fn admit(&self, prompt_tokens: u64) -> Result<AdmissionPlan, AdmissionRejection> {
352 admit(&self.request_for(prompt_tokens))
353 }
354}
355
356pub fn admit(request: &AdmissionRequest) -> Result<AdmissionPlan, AdmissionRejection> {
362 if request.prompt_tokens >= MAX_CONTEXT_TOKENS {
363 return Err(AdmissionRejection::PromptExceedsContext {
364 prompt_tokens: request.prompt_tokens,
365 ceiling: MAX_CONTEXT_TOKENS,
366 });
367 }
368 if request.max_new_tokens == 0 {
369 return Err(AdmissionRejection::NoFramesRequested);
370 }
371
372 let headroom = MAX_CONTEXT_TOKENS - request.prompt_tokens;
373 let heuristic_cap = if request.heuristic_eos_backstop {
374 request
375 .prompt_tokens
376 .saturating_mul(HEURISTIC_FRAMES_PER_PROMPT_TOKEN)
377 .saturating_add(HEURISTIC_FRAME_HEADROOM)
378 } else {
379 u64::MAX
380 };
381 let predicted_max_frames = request.max_new_tokens.min(headroom).min(heuristic_cap);
382 let binding_constraint = if predicted_max_frames == heuristic_cap
383 && heuristic_cap < request.max_new_tokens.min(headroom)
384 {
385 BindingConstraint::TextHeuristic
386 } else if request.max_new_tokens <= headroom {
387 BindingConstraint::FrameCap
388 } else {
389 BindingConstraint::ContextCeiling
390 };
391
392 let kv_talker_bytes = talker_kv_bytes(
393 request.prompt_tokens,
394 predicted_max_frames,
395 request.kv_dtype,
396 )?;
397
398 let bounded_state_bytes = MICRODECODER_KV_BYTES
399 .checked_add(CODEC_DECODER_KV_BYTES)
400 .and_then(|sum| sum.checked_add(request.ring_buffer_bytes))
401 .ok_or(AdmissionRejection::Overflow {
402 term: "bounded_state_bytes",
403 })?;
404
405 let predicted_peak_bytes = kv_talker_bytes
406 .checked_add(bounded_state_bytes)
407 .and_then(|sum| sum.checked_add(request.weights_resident_bytes))
408 .ok_or(AdmissionRejection::Overflow {
409 term: "predicted_peak_bytes",
410 })?;
411
412 let plan = AdmissionPlan {
413 predicted_max_frames,
414 binding_constraint,
415 kv_talker_bytes,
416 bounded_state_bytes,
417 weights_resident_bytes: request.weights_resident_bytes,
418 predicted_peak_bytes,
419 budget_bytes: request.budget_bytes,
420 };
421
422 if plan.fits() {
423 Ok(plan)
424 } else {
425 Err(AdmissionRejection::BudgetExceeded { plan })
426 }
427}
428
429#[derive(Clone, Copy, Debug, PartialEq, Eq)]
434pub enum StopReason {
435 EndOfSpeech,
437 FrameCapReached,
439 DurationLimitReached,
441 Cancelled,
443}
444
445impl StopReason {
446 #[must_use]
448 pub const fn as_str(self) -> &'static str {
449 match self {
450 Self::EndOfSpeech => "end_of_speech",
451 Self::FrameCapReached => "frame_cap_reached",
452 Self::DurationLimitReached => "duration_limit_reached",
453 Self::Cancelled => "cancelled",
454 }
455 }
456
457 #[must_use]
462 pub const fn is_truncated(self) -> bool {
463 matches!(self, Self::FrameCapReached | Self::DurationLimitReached)
464 }
465
466 #[must_use]
471 pub const fn is_clean_completion(self) -> bool {
472 matches!(self, Self::EndOfSpeech)
473 }
474
475 #[must_use]
477 pub const fn remedy(self) -> Option<&'static str> {
478 match self {
479 Self::EndOfSpeech | Self::Cancelled => None,
480 Self::FrameCapReached => Some(
481 "the utterance hit the frame cap before the model finished speaking; raise \
482 --max-frames or split the text into shorter chunks",
483 ),
484 Self::DurationLimitReached => Some(
485 "the utterance hit the hard duration limit; raise it or split the text into \
486 shorter chunks",
487 ),
488 }
489 }
490}
491
492impl fmt::Display for StopReason {
493 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
494 f.write_str(self.as_str())
495 }
496}
497
498#[cfg(test)]
499mod tests {
500 use super::*;
501
502 const MIB: u64 = 1024 * 1024;
503 const GIB: u64 = 1024 * 1024 * 1024;
504
505 fn request(prompt_tokens: u64, max_new_tokens: u64, budget_bytes: u64) -> AdmissionRequest {
506 AdmissionRequest {
507 prompt_tokens,
508 max_new_tokens,
509 heuristic_eos_backstop: false,
511 kv_dtype: KvDtype::Bf16,
512 ring_buffer_bytes: 0,
513 weights_resident_bytes: 0,
514 budget_bytes,
515 }
516 }
517
518 #[test]
519 fn the_eos_backstop_binds_a_short_prompt_under_the_flat_default_cap() {
520 let mut with_backstop = request(28, DEFAULT_MAX_NEW_TOKENS, 2 * GIB);
521 with_backstop.heuristic_eos_backstop = true;
522 let plan = admit(&with_backstop).expect("fits easily");
523 assert_eq!(
524 plan.predicted_max_frames,
525 28 * HEURISTIC_FRAMES_PER_PROMPT_TOKEN + HEURISTIC_FRAME_HEADROOM
526 );
527 assert_eq!(plan.binding_constraint, BindingConstraint::TextHeuristic);
528 }
529
530 #[test]
531 fn an_explicit_cap_disables_the_eos_backstop_exactly() {
532 let explicit = request(28, 2_000, 2 * GIB);
535 let plan = admit(&explicit).expect("fits");
536 assert_eq!(plan.predicted_max_frames, 2_000);
537 assert_eq!(plan.binding_constraint, BindingConstraint::FrameCap);
538 }
539
540 #[test]
541 fn the_backstop_never_raises_a_smaller_explicit_cap() {
542 let mut small = request(1_000, 32, 2 * GIB);
543 small.heuristic_eos_backstop = true;
544 let plan = admit(&small).expect("fits");
545 assert_eq!(plan.predicted_max_frames, 32);
546 assert_eq!(plan.binding_constraint, BindingConstraint::FrameCap);
547 }
548
549 #[test]
552 fn talker_kv_matches_the_oq6_worked_points_exactly() {
553 assert_eq!(
554 talker_kv_bytes(512, 2048, KvDtype::Bf16).expect("no overflow"),
555 280 * MIB,
556 "512-token prompt + 2048-frame cap must be exactly 280 MiB"
557 );
558 assert_eq!(
559 talker_kv_bytes(512, 8192, KvDtype::Bf16).expect("no overflow"),
560 952 * MIB,
561 "512-token prompt + 8192-frame cap must be exactly 952 MiB"
562 );
563 assert_eq!(
565 talker_kv_bytes(0, MAX_CONTEXT_TOKENS, KvDtype::Bf16).expect("no overflow"),
566 7 * GIB / 2,
567 "the full 32768-token context must be exactly 3.50 GiB"
568 );
569 assert_eq!(
571 talker_kv_bytes(1, 0, KvDtype::Bf16).expect("no overflow"),
572 112 * 1024
573 );
574 assert_eq!(
576 talker_kv_bytes(512, 2048, KvDtype::F32).expect("no overflow"),
577 560 * MIB
578 );
579 }
580
581 #[test]
582 fn a_request_that_fits_is_admitted_with_its_full_prediction() {
583 let plan = admit(&request(512, 2048, 2 * GIB)).expect("must be admitted");
584 assert_eq!(plan.predicted_max_frames, 2048);
585 assert_eq!(plan.binding_constraint, BindingConstraint::FrameCap);
586 assert_eq!(plan.kv_talker_bytes, 280 * MIB);
587 assert_eq!(
588 plan.bounded_state_bytes,
589 MICRODECODER_KV_BYTES + CODEC_DECODER_KV_BYTES
590 );
591 assert!(plan.fits());
592 assert_eq!(plan.shortfall_bytes(), 0);
593 }
594
595 #[test]
597 fn an_over_budget_request_is_rejected_before_any_allocation_and_says_by_how_much() {
598 let error = admit(&request(512, 8192, 512 * MIB)).expect_err("must be rejected");
599 let AdmissionRejection::BudgetExceeded { plan } = error else {
600 panic!("expected a budget rejection, got {error}");
601 };
602 assert!(!plan.fits());
603 assert_eq!(plan.kv_talker_bytes, 952 * MIB);
604 assert!(plan.shortfall_bytes() > 0);
605
606 let rendered = error.to_string();
608 for expected in ["predicted peak", "budget", "exceeds", "before allocating"] {
609 assert!(
610 rendered.contains(expected),
611 "rejection is not actionable, missing `{expected}`: {rendered}"
612 );
613 }
614 }
615
616 #[test]
617 fn the_binding_constraint_is_reported_because_the_two_have_different_remedies() {
618 let plan = admit(&request(512, 8192, 8 * GIB)).expect("admitted");
620 assert_eq!(plan.binding_constraint, BindingConstraint::FrameCap);
621 assert_eq!(plan.predicted_max_frames, 8192);
622
623 let prompt = MAX_CONTEXT_TOKENS - 100;
625 let plan = admit(&request(prompt, 8192, 8 * GIB)).expect("admitted");
626 assert_eq!(plan.binding_constraint, BindingConstraint::ContextCeiling);
627 assert_eq!(plan.predicted_max_frames, 100);
628
629 let plan = admit(&request(24_000, 8192, 8 * GIB)).expect("admitted");
632 assert_eq!(plan.binding_constraint, BindingConstraint::FrameCap);
633 }
634
635 #[test]
636 fn a_prompt_at_or_past_the_ceiling_is_refused_as_unfixable_by_memory() {
637 for prompt in [MAX_CONTEXT_TOKENS, MAX_CONTEXT_TOKENS + 1, u64::MAX] {
638 let error = admit(&request(prompt, 1024, u64::MAX)).expect_err("must be rejected");
639 assert!(
640 matches!(error, AdmissionRejection::PromptExceedsContext { .. }),
641 "got {error}"
642 );
643 assert!(error.to_string().contains("cannot help"));
645 }
646 }
647
648 #[test]
649 fn zero_frames_is_refused_rather_than_admitted_as_a_no_op() {
650 let error = admit(&request(512, 0, u64::MAX)).expect_err("must be rejected");
651 assert_eq!(error, AdmissionRejection::NoFramesRequested);
652 }
653
654 #[test]
656 fn arithmetic_overflow_is_refused_never_wrapped_into_a_plausible_total() {
657 assert!(matches!(
658 talker_kv_bytes(u64::MAX, u64::MAX, KvDtype::F32),
659 Err(AdmissionRejection::Overflow { .. })
660 ));
661
662 let over = AdmissionRequest {
663 prompt_tokens: 512,
664 max_new_tokens: 2048,
665 heuristic_eos_backstop: false,
666 kv_dtype: KvDtype::Bf16,
667 ring_buffer_bytes: u64::MAX,
668 weights_resident_bytes: u64::MAX,
669 budget_bytes: u64::MAX,
670 };
671 let error = admit(&over).expect_err("overflow must not be admitted");
672 assert!(
673 matches!(error, AdmissionRejection::Overflow { .. }),
674 "a wrapped total is exactly the failure admission exists to prevent, got {error}"
675 );
676 }
677
678 #[test]
679 fn admission_is_exactly_at_the_boundary_not_off_by_one() {
680 let peak = admit(&request(512, 2048, u64::MAX))
681 .expect("admitted")
682 .predicted_peak_bytes;
683 assert!(admit(&request(512, 2048, peak)).is_ok());
685 assert!(admit(&request(512, 2048, peak - 1)).is_err());
686 }
687
688 #[test]
689 fn only_end_of_speech_counts_as_a_clean_completion() {
690 assert!(StopReason::EndOfSpeech.is_clean_completion());
691 assert!(!StopReason::EndOfSpeech.is_truncated());
692
693 for cut in [
695 StopReason::FrameCapReached,
696 StopReason::DurationLimitReached,
697 ] {
698 assert!(cut.is_truncated(), "{cut} must be reported as truncated");
699 assert!(
700 !cut.is_clean_completion(),
701 "{cut} must never be reported as an unqualified success — an agent cannot hear \
702 that the audio stopped mid-word"
703 );
704 assert!(
705 cut.remedy().is_some(),
706 "{cut} must tell the caller what to do"
707 );
708 }
709
710 assert!(!StopReason::Cancelled.is_clean_completion());
712 assert!(!StopReason::Cancelled.is_truncated());
713 }
714
715 #[test]
716 fn stop_reason_wire_strings_are_distinct_and_stable() {
717 let all = [
718 StopReason::EndOfSpeech,
719 StopReason::FrameCapReached,
720 StopReason::DurationLimitReached,
721 StopReason::Cancelled,
722 ];
723 let mut seen: Vec<&str> = all.iter().map(|reason| reason.as_str()).collect();
724 let count = seen.len();
725 seen.sort_unstable();
726 seen.dedup();
727 assert_eq!(seen.len(), count, "two stop reasons share a wire string");
728 assert_eq!(StopReason::FrameCapReached.as_str(), "frame_cap_reached");
729 }
730}