1use crate::{
8 artifact::{
9 plan_model_preparation, ArtifactFormat, ArtifactInspection, ModelConfigurationResolver,
10 PreparationPolicy,
11 },
12 backend::{BackendProvider, ModelLoadingBackend, ModelRuntime, SessionCapabilities},
13 execution::{
14 DevicePlan, DraftingPlan, ExecutionPlan, ExpertCachePlan, ResidencyPlan,
15 DEFAULT_MAX_CACHED_SHARDS,
16 },
17 speculative::SpeculativeDraft,
18};
19use serde::{Deserialize, Serialize};
20use std::{
21 path::PathBuf,
22 sync::atomic::{AtomicU64, Ordering},
23 time::Duration,
24};
25
26static NEXT_EXECUTION_PLAN_TARGET_ID: AtomicU64 = AtomicU64::new(1);
27
28fn next_execution_plan_target_id() -> Result<u64, AutomaticPlanningError> {
29 NEXT_EXECUTION_PLAN_TARGET_ID
30 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
31 current.checked_add(1)
32 })
33 .map_err(|_| {
34 AutomaticPlanningError::Invalid(
35 "execution-plan target identity space is exhausted".into(),
36 )
37 })
38}
39
40pub const AUTOMATIC_SCHEMA_VERSION: u32 = 6;
42
43#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
45#[serde(rename_all = "snake_case")]
46pub enum ObservationKind {
47 Exact,
49 Conservative,
51 Observational,
53 Estimated,
55}
56
57#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
59#[serde(tag = "status", rename_all = "snake_case")]
60pub enum Observed<T> {
61 Available {
63 value: T,
65 kind: ObservationKind,
67 source: String,
69 },
70 Unsupported {
72 reason: String,
74 },
75 Unavailable {
77 reason: String,
79 },
80}
81
82impl<T> Observed<T> {
83 pub fn exact(value: T, source: impl Into<String>) -> Self {
85 Self::Available {
86 value,
87 kind: ObservationKind::Exact,
88 source: source.into(),
89 }
90 }
91
92 pub fn unavailable(reason: impl Into<String>) -> Self {
94 Self::Unavailable {
95 reason: reason.into(),
96 }
97 }
98
99 pub fn unsupported(reason: impl Into<String>) -> Self {
101 Self::Unsupported {
102 reason: reason.into(),
103 }
104 }
105
106 pub const fn value(&self) -> Option<&T> {
108 match self {
109 Self::Available { value, .. } => Some(value),
110 Self::Unsupported { .. } | Self::Unavailable { .. } => None,
111 }
112 }
113}
114
115fn unobserved_embedded_draft_layers() -> Observed<usize> {
116 Observed::unavailable("embedded drafting requires normalized architecture inspection")
117}
118
119#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
121pub struct ModelResourceProfile {
122 pub schema_version: u32,
124 pub path: PathBuf,
126 pub artifact_format: ArtifactFormat,
128 #[serde(skip_serializing_if = "Option::is_none")]
130 pub model_family: Option<String>,
131 #[serde(skip_serializing_if = "Option::is_none")]
133 pub architecture: Option<String>,
134 #[serde(skip_serializing_if = "Option::is_none")]
136 pub tensor_count: Option<usize>,
137 #[serde(skip_serializing_if = "Option::is_none")]
139 pub checkpoint_shards: Option<usize>,
140 #[serde(default = "unobserved_embedded_draft_layers")]
142 pub embedded_draft_layers: Observed<usize>,
143 pub stored_tensor_bytes: Observed<u64>,
145 pub largest_stored_tensor_bytes: Observed<u64>,
147 pub materialized_parameter_bytes: Observed<u64>,
149 pub pinned_parameter_bytes: Observed<u64>,
151 pub largest_execution_group_bytes: Observed<u64>,
153 pub largest_adjacent_execution_groups_bytes: Observed<u64>,
155 pub expert_parameter_bytes: Observed<u64>,
157}
158
159impl ModelResourceProfile {
160 pub fn unmeasured(path: PathBuf, artifact_format: ArtifactFormat) -> Self {
162 let unavailable = || {
163 Observed::unavailable("resource value requires a validated checkpoint parameter plan")
164 };
165 Self {
166 schema_version: AUTOMATIC_SCHEMA_VERSION,
167 path,
168 artifact_format,
169 model_family: None,
170 architecture: None,
171 tensor_count: None,
172 checkpoint_shards: None,
173 embedded_draft_layers: unobserved_embedded_draft_layers(),
174 stored_tensor_bytes: Observed::unavailable(
175 "checkpoint tensor catalog was not established",
176 ),
177 largest_stored_tensor_bytes: Observed::unavailable(
178 "checkpoint tensor catalog was not established",
179 ),
180 materialized_parameter_bytes: unavailable(),
181 pinned_parameter_bytes: unavailable(),
182 largest_execution_group_bytes: unavailable(),
183 largest_adjacent_execution_groups_bytes: unavailable(),
184 expert_parameter_bytes: unavailable(),
185 }
186 }
187}
188
189#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
191pub struct HardwareDeviceProfile {
192 pub id: String,
194 pub family: String,
196 pub index: usize,
198 pub total_memory_bytes: Observed<u64>,
200 pub available_memory_bytes: Observed<u64>,
202}
203
204#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
206pub struct HardwareBackendProfile {
207 pub backend: crate::execution::BackendId,
209 pub available: bool,
211 #[serde(skip_serializing_if = "Option::is_none")]
213 pub detail: Option<String>,
214 pub devices: Vec<HardwareDeviceProfile>,
216}
217
218#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
220pub struct HardwareProfile {
221 pub schema_version: u32,
223 pub operating_system: String,
225 pub architecture: String,
227 pub logical_cpu_count: Observed<u64>,
229 pub physical_memory_bytes: Observed<u64>,
231 pub available_memory_bytes: Observed<u64>,
233 pub physical_memory_semantics: HardwareMemorySemantics,
235 pub backends: Vec<HardwareBackendProfile>,
237}
238
239impl HardwareProfile {
240 pub fn observe_host(
246 physical_memory_bytes: Observed<u64>,
247 available_memory_bytes: Observed<u64>,
248 physical_memory_semantics: HardwareMemorySemantics,
249 backends: Vec<HardwareBackendProfile>,
250 ) -> Self {
251 let logical_cpu_count = std::thread::available_parallelism().map_or_else(
252 |error| Observed::unavailable(error.to_string()),
253 |count| Observed::exact(count.get() as u64, "std::thread::available_parallelism"),
254 );
255 Self {
256 schema_version: AUTOMATIC_SCHEMA_VERSION,
257 operating_system: std::env::consts::OS.into(),
258 architecture: std::env::consts::ARCH.into(),
259 logical_cpu_count,
260 physical_memory_bytes,
261 available_memory_bytes,
262 physical_memory_semantics,
263 backends,
264 }
265 }
266}
267
268#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
270#[serde(rename_all = "snake_case")]
271pub enum HardwareMemorySemantics {
272 Unified,
274 SeparateTiers,
276 Unknown,
278}
279
280impl From<crate::capability::PhysicalMemorySemantics> for HardwareMemorySemantics {
281 fn from(value: crate::capability::PhysicalMemorySemantics) -> Self {
282 match value {
283 crate::capability::PhysicalMemorySemantics::Unified => Self::Unified,
284 crate::capability::PhysicalMemorySemantics::SeparateTiers => Self::SeparateTiers,
285 crate::capability::PhysicalMemorySemantics::Unknown => Self::Unknown,
286 }
287 }
288}
289
290#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
292#[serde(rename_all = "snake_case")]
293pub enum PlanExplanationLevel {
294 Decision,
296 Warning,
298 Rejection,
300}
301
302#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
304pub struct PlanExplanationEntry {
305 pub level: PlanExplanationLevel,
307 pub code: String,
309 pub detail: String,
311}
312
313#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
315pub struct PlanExplanation {
316 pub summary: String,
318 pub entries: Vec<PlanExplanationEntry>,
320}
321
322#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
324pub struct ExecutionPlanReport {
325 pub schema_version: u32,
327 pub hardware: HardwareProfile,
329 pub resources: ModelResourceProfile,
331 pub plan: ExecutionPlan,
333 pub explanation: PlanExplanation,
335}
336
337#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
339#[serde(default)]
340#[non_exhaustive]
341pub struct AutomaticPlannerPolicy {
342 pub device_memory_fallback_bytes: u64,
344 pub host_memory_fallback_bytes: u64,
346 pub memory_headroom_percent: u8,
348 pub expert_cache_share_percent: u8,
350 pub device_layer_window: usize,
352 pub max_cached_shards: usize,
354 pub embedded_mtp_draft_tokens: usize,
356 pub minimum_feedback_tokens: usize,
358}
359
360impl Default for AutomaticPlannerPolicy {
361 fn default() -> Self {
362 Self {
363 device_memory_fallback_bytes: 4 << 30,
364 host_memory_fallback_bytes: 16 << 30,
365 memory_headroom_percent: 30,
366 expert_cache_share_percent: 40,
367 device_layer_window: 1,
368 max_cached_shards: DEFAULT_MAX_CACHED_SHARDS,
369 embedded_mtp_draft_tokens: 3,
370 minimum_feedback_tokens: 1,
371 }
372 }
373}
374
375#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
377pub struct TimingTelemetry {
378 pub load_seconds: f64,
380 pub generation_seconds: f64,
382 #[serde(skip_serializing_if = "Option::is_none")]
384 pub time_to_first_token_seconds: Option<f64>,
385 pub total_seconds: f64,
387 pub token_rate: f64,
389 #[serde(skip_serializing_if = "Option::is_none")]
391 pub decode_token_rate: Option<f64>,
392}
393
394impl TimingTelemetry {
395 pub fn new(
397 load: Duration,
398 generation: Duration,
399 time_to_first_token: Option<Duration>,
400 generated_tokens: usize,
401 total: Duration,
402 ) -> Self {
403 fn rate(tokens: usize, elapsed: Duration) -> f64 {
404 if elapsed.is_zero() {
405 0.0
406 } else {
407 tokens as f64 / elapsed.as_secs_f64()
408 }
409 }
410 Self {
411 load_seconds: load.as_secs_f64(),
412 generation_seconds: generation.as_secs_f64(),
413 time_to_first_token_seconds: time_to_first_token.map(|value| value.as_secs_f64()),
414 total_seconds: total.as_secs_f64(),
415 token_rate: rate(generated_tokens, generation),
416 decode_token_rate: time_to_first_token.map(|first| {
417 rate(
418 generated_tokens.saturating_sub(1),
419 generation.saturating_sub(first),
420 )
421 }),
422 }
423 }
424}
425
426#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
428pub struct AllocatorTelemetry {
429 pub peak_bytes: u64,
431 pub active_bytes: u64,
433 pub cache_bytes: u64,
435}
436
437#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
439pub struct ResidencyTelemetry {
440 pub planned_disk_bytes: u64,
442 pub planned_host_bytes: u64,
444 pub planned_device_bytes: u64,
446 pub current_host_bytes: u64,
448 pub current_device_bytes: u64,
450 pub peak_host_bytes: u64,
452 pub peak_device_bytes: u64,
454 pub transfers: Vec<TransferTelemetry>,
456}
457
458#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
460pub struct TransferTelemetry {
461 pub direction: String,
463 pub count: u64,
465 pub bytes: u64,
467 pub seconds: DurationSeconds,
469}
470
471#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
473#[serde(transparent)]
474pub struct DurationSeconds(pub f64);
475
476impl PartialEq for DurationSeconds {
477 fn eq(&self, other: &Self) -> bool {
478 self.0.to_bits() == other.0.to_bits()
479 }
480}
481impl Eq for DurationSeconds {}
482
483#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
485pub struct ExpertCacheTelemetry {
486 pub owned_experts: usize,
488 pub owned_bytes: u64,
490 pub host_resident_experts: usize,
492 pub device_resident_experts: usize,
494 pub host_resident_bytes: u64,
496 pub device_resident_bytes: u64,
498 pub peak_host_resident_bytes: u64,
500 pub peak_device_resident_bytes: u64,
502}
503
504#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
506pub struct SpeculativeDecodingTelemetry {
507 pub execution_topology: String,
509 pub target_tokens: usize,
511 pub draft_tokens: usize,
513 pub accepted_tokens: usize,
515 pub accept_rate: f64,
517 pub rounds: usize,
519 pub accept_lens: Vec<usize>,
521 pub emitted_tokens: usize,
523 pub optimistic_draft_tokens: usize,
525 pub reused_optimistic_tokens: usize,
527 pub discarded_optimistic_tokens: usize,
529 pub adaptive_lookahead_disabled: bool,
531 pub optimistic_draft_seconds: f64,
533 pub verification_in_flight_seconds: f64,
535}
536
537pub fn speculative_decoding_telemetry(
539 stats: &crate::speculative::SpeculativeStats,
540) -> SpeculativeDecodingTelemetry {
541 SpeculativeDecodingTelemetry {
542 execution_topology: stats.execution_topology().to_string(),
543 target_tokens: stats.target_tokens(),
544 draft_tokens: stats.draft_tokens(),
545 accepted_tokens: stats.accepted_tokens(),
546 accept_rate: stats.accept_rate(),
547 rounds: stats.rounds(),
548 accept_lens: stats.accept_lens().to_vec(),
549 emitted_tokens: stats.emitted_tokens(),
550 optimistic_draft_tokens: stats.optimistic_draft_tokens(),
551 reused_optimistic_tokens: stats.reused_optimistic_tokens(),
552 discarded_optimistic_tokens: stats.discarded_optimistic_tokens(),
553 adaptive_lookahead_disabled: stats.adaptive_lookahead_disabled(),
554 optimistic_draft_seconds: stats.optimistic_draft_time().as_secs_f64(),
555 verification_in_flight_seconds: stats.verification_in_flight_time().as_secs_f64(),
556 }
557}
558
559#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
561pub struct ExecutionTelemetry {
562 pub schema_version: u32,
564 pub effective_model_type: String,
566 #[serde(skip_serializing_if = "Option::is_none")]
568 pub plan: Option<ExecutionPlan>,
569 #[serde(skip_serializing_if = "Option::is_none")]
571 pub plan_explanation: Option<PlanExplanation>,
572 #[serde(skip_serializing_if = "Option::is_none")]
574 pub hardware: Option<HardwareProfile>,
575 #[serde(skip_serializing_if = "Option::is_none")]
577 pub resources: Option<ModelResourceProfile>,
578 pub prompt_tokens: usize,
580 pub generated_tokens: usize,
582 pub stop_reason: String,
584 pub timing: TimingTelemetry,
586 #[serde(skip_serializing_if = "Option::is_none")]
588 pub allocator: Option<AllocatorTelemetry>,
589 #[serde(skip_serializing_if = "Option::is_none")]
591 pub residency: Option<ResidencyTelemetry>,
592 #[serde(skip_serializing_if = "Option::is_none")]
594 pub expert_cache: Option<ExpertCacheTelemetry>,
595 #[serde(skip_serializing_if = "Option::is_none")]
597 pub speculative: Option<SpeculativeDecodingTelemetry>,
598}
599
600#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
602#[non_exhaustive]
603pub struct AutomaticPlanRequest {
604 pub schema_version: u32,
606 pub model_path: PathBuf,
608 pub device: DevicePlan,
610 #[serde(default, skip_serializing_if = "Vec::is_empty")]
612 pub prior_telemetry: Vec<ExecutionTelemetry>,
613}
614
615impl AutomaticPlanRequest {
616 pub fn new(model_path: impl Into<PathBuf>, device: DevicePlan) -> Self {
618 Self {
619 schema_version: AUTOMATIC_SCHEMA_VERSION,
620 model_path: model_path.into(),
621 device,
622 prior_telemetry: Vec::new(),
623 }
624 }
625
626 pub fn with_prior_telemetry(
628 mut self,
629 telemetry: impl IntoIterator<Item = ExecutionTelemetry>,
630 ) -> Self {
631 self.prior_telemetry.extend(telemetry);
632 self
633 }
634}
635
636#[derive(Debug, Clone, Eq, PartialEq)]
638pub struct CandidateAdmission {
639 pub supported: bool,
641 pub rejection: Option<String>,
643}
644
645#[derive(Debug, Clone, Copy, Eq, PartialEq)]
647pub struct BoundedResidencyRequirement {
648 pub static_bytes: u64,
650 pub window_bytes: u64,
652 pub required_bytes: u64,
654 pub depth: usize,
656}
657
658pub trait AutomaticPlanningBackend {
660 type Inspection;
662 fn backend_id(&self) -> crate::execution::BackendId;
664 fn discover_hardware(&self) -> Result<HardwareProfile, AutomaticPlanningError>;
666 fn inspect_resources(
668 &self,
669 model_path: &std::path::Path,
670 ) -> Result<(ModelResourceProfile, Self::Inspection), AutomaticPlanningError>;
671 fn admit_candidate(
673 &self,
674 inspection: &Self::Inspection,
675 plan: &ExecutionPlan,
676 ) -> Result<CandidateAdmission, AutomaticPlanningError>;
677 fn bounded_residency_requirement(
679 &self,
680 inspection: &Self::Inspection,
681 plan: &ExecutionPlan,
682 ) -> Result<BoundedResidencyRequirement, AutomaticPlanningError>;
683}
684
685pub struct RetainedAutomaticPlan<I> {
687 report: ExecutionPlanReport,
688 inspection: I,
689}
690
691impl<I> RetainedAutomaticPlan<I> {
692 pub const fn report(&self) -> &ExecutionPlanReport {
694 &self.report
695 }
696
697 pub fn into_parts(self) -> (ExecutionPlanReport, I) {
699 (self.report, self.inspection)
700 }
701}
702
703pub struct ExecutionPlanTargetSelection<B: ModelLoadingBackend> {
705 policy: PreparationPolicy,
706 selected: B::SelectedPreparation,
707 capabilities: SessionCapabilities,
708}
709
710impl<B: ModelLoadingBackend> ExecutionPlanTargetSelection<B> {
711 pub fn new(
713 policy: PreparationPolicy,
714 selected: B::SelectedPreparation,
715 capabilities: SessionCapabilities,
716 ) -> Self {
717 Self {
718 policy,
719 selected,
720 capabilities,
721 }
722 }
723}
724
725pub struct SelectedExecutionPlanTarget<B: ModelLoadingBackend> {
727 execution_plan: ExecutionPlan,
728 preparation: crate::backend::SelectedModelPreparation<B>,
729 target_id: u64,
730}
731
732impl<B: ModelLoadingBackend> SelectedExecutionPlanTarget<B> {
733 fn into_preparation(self) -> crate::backend::SelectedModelPreparation<B> {
734 self.preparation
735 }
736
737 pub fn inspection(
739 &self,
740 ) -> &ArtifactInspection<<B::ConfigurationResolver as ModelConfigurationResolver>::ArtifactPlan>
741 {
742 self.preparation.plan().inspection()
743 }
744
745 pub const fn execution_plan(&self) -> &ExecutionPlan {
747 &self.execution_plan
748 }
749}
750
751pub struct ExecutionPlanTarget<B: ModelLoadingBackend> {
757 backend: B,
758 selected: SelectedExecutionPlanTarget<B>,
759}
760
761pub type PreparedExecutionPlanTarget<B> = ModelRuntime<B>;
763
764pub type ExecutionPlanTargetLoadError<B> =
766 crate::backend::ModelLoadError<<B as BackendProvider>::Error>;
767
768impl<B: ModelLoadingBackend> ExecutionPlanTarget<B> {
769 pub fn new(backend: B, selected: SelectedExecutionPlanTarget<B>) -> Self {
775 Self { backend, selected }
776 }
777
778 pub const fn backend(&self) -> &B {
780 &self.backend
781 }
782
783 pub fn into_runtime(
785 self,
786 ) -> Result<PreparedExecutionPlanTarget<B>, ExecutionPlanTargetLoadError<B>> {
787 let target_id = self.selected.target_id;
788 let preparation = self.selected.into_preparation();
789 let prepared = crate::backend::prepare_selected_model(&self.backend, preparation)?;
790 ModelRuntime::from_prepared_execution_plan_target(self.backend, prepared, target_id)
791 .map_err(crate::backend::ModelLoadError::Backend)
792 }
793}
794
795#[derive(Debug, Clone, Copy, Eq, PartialEq)]
801pub struct TokenizerCompatibilityProof {
802 fingerprint: [u8; 32],
803}
804
805impl TokenizerCompatibilityProof {
806 pub fn prove(
808 target_fingerprint: [u8; 32],
809 assistant_fingerprint: [u8; 32],
810 ) -> Result<Self, TokenizerCompatibilityError> {
811 if target_fingerprint != assistant_fingerprint {
812 return Err(TokenizerCompatibilityError);
813 }
814 Ok(Self {
815 fingerprint: target_fingerprint,
816 })
817 }
818
819 pub const fn fingerprint(self) -> [u8; 32] {
821 self.fingerprint
822 }
823
824 pub fn validate_target(
826 self,
827 target_fingerprint: [u8; 32],
828 ) -> Result<(), TokenizerCompatibilityError> {
829 if self.fingerprint != target_fingerprint {
830 return Err(TokenizerCompatibilityError);
831 }
832 Ok(())
833 }
834}
835
836#[derive(Debug, Clone, Copy, Eq, PartialEq, thiserror::Error)]
838#[error("assistant token-id vocabulary mapping does not match the target")]
839pub struct TokenizerCompatibilityError;
840
841#[derive(Debug, Clone, Eq, PartialEq)]
843pub struct ExternalDraftArtifact<P> {
844 pub preparation: P,
846 pub tokenizer_compatibility: TokenizerCompatibilityProof,
848}
849
850pub struct SelectedExecutionPlanDrafting<P> {
857 execution_plan: ExecutionPlan,
858 target_id: u64,
859 external_artifact: Option<ExternalDraftArtifact<P>>,
860}
861
862impl<P> std::fmt::Debug for SelectedExecutionPlanDrafting<P> {
863 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
864 formatter
865 .debug_struct("SelectedExecutionPlanDrafting")
866 .field("execution_plan", &self.execution_plan)
867 .field("has_external_artifact", &self.external_artifact.is_some())
868 .finish()
869 }
870}
871
872impl<P> SelectedExecutionPlanDrafting<P> {
873 pub fn into_external_artifact<B: BackendProvider>(
879 self,
880 plan: &ExecutionPlan,
881 target: &ModelRuntime<B>,
882 ) -> Result<Option<ExternalDraftArtifact<P>>, AutomaticPlanningError> {
883 if self.execution_plan != *plan {
884 return Err(AutomaticPlanningError::Invalid(
885 "selected drafting was established for a different execution plan".into(),
886 ));
887 }
888 if target.execution_plan_target_id() != Some(self.target_id) {
889 return Err(AutomaticPlanningError::Invalid(
890 "selected drafting was established for a different realized target".into(),
891 ));
892 }
893 Ok(self.external_artifact)
894 }
895}
896
897pub enum RealizedDrafting<D> {
899 Disabled,
901 Embedded,
903 External(D),
905}
906
907impl<D> RealizedDrafting<D> {
908 pub fn as_speculative_draft(&mut self) -> Option<SpeculativeDraft<'_, D>> {
910 match self {
911 Self::Disabled => None,
912 Self::Embedded => Some(SpeculativeDraft::Embedded),
913 Self::External(drafter) => Some(SpeculativeDraft::External(drafter)),
914 }
915 }
916
917 pub const fn is_external(&self) -> bool {
919 matches!(self, Self::External(_))
920 }
921}
922
923pub trait ExecutionPlanBackendFactory: AutomaticPlanningBackend {
930 type Backend: ModelLoadingBackend;
932 type DrafterPreparation;
934 type SelectedDrafterPreparation;
936 type Drafter;
938
939 fn select_target(
941 &self,
942 inspection: &ArtifactInspection<
943 <<Self::Backend as ModelLoadingBackend>::ConfigurationResolver as ModelConfigurationResolver>::ArtifactPlan,
944 >,
945 plan: &ExecutionPlan,
946 ) -> Result<ExecutionPlanTargetSelection<Self::Backend>, AutomaticPlanningError>;
947
948 fn select_drafting(
953 &self,
954 plan: &ExecutionPlan,
955 target: &SelectedExecutionPlanTarget<Self::Backend>,
956 external_artifact: Option<ExternalDraftArtifact<Self::DrafterPreparation>>,
957 ) -> Result<
958 Option<ExternalDraftArtifact<Self::SelectedDrafterPreparation>>,
959 AutomaticPlanningError,
960 >;
961
962 fn realize_target(
967 &self,
968 selected: SelectedExecutionPlanTarget<Self::Backend>,
969 ) -> Result<ExecutionPlanTarget<Self::Backend>, AutomaticPlanningError>;
970
971 fn realize_drafting(
979 &self,
980 plan: &ExecutionPlan,
981 target: &ModelRuntime<Self::Backend>,
982 selected: SelectedExecutionPlanDrafting<Self::SelectedDrafterPreparation>,
983 ) -> Result<RealizedDrafting<Self::Drafter>, AutomaticPlanningError>;
984}
985
986pub fn select_execution_plan_target<F: ExecutionPlanBackendFactory>(
988 factory: &F,
989 plan: &ExecutionPlan,
990 inspection: ArtifactInspection<
991 <<F::Backend as ModelLoadingBackend>::ConfigurationResolver as ModelConfigurationResolver>::ArtifactPlan,
992 >,
993) -> Result<SelectedExecutionPlanTarget<F::Backend>, AutomaticPlanningError> {
994 let expected_backend = factory.backend_id();
995 if plan.device.backend != expected_backend {
996 return Err(AutomaticPlanningError::Invalid(format!(
997 "execution plan selects backend {} but factory owns {}",
998 plan.device.backend, expected_backend
999 )));
1000 }
1001 plan.validate_structure()
1002 .map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
1003
1004 let selection = factory.select_target(&inspection, plan)?;
1005 selection
1006 .policy
1007 .validate_session_capabilities(&selection.capabilities)
1008 .map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
1009 let preparation = plan_model_preparation(inspection, selection.policy, selection.capabilities)
1010 .map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
1011 Ok(SelectedExecutionPlanTarget {
1012 execution_plan: plan.clone(),
1013 preparation: crate::backend::SelectedModelPreparation::new(preparation, selection.selected),
1014 target_id: next_execution_plan_target_id()?,
1015 })
1016}
1017
1018pub fn realize_execution_plan_target<F: ExecutionPlanBackendFactory>(
1020 factory: &F,
1021 plan: &ExecutionPlan,
1022 selected: SelectedExecutionPlanTarget<F::Backend>,
1023) -> Result<ExecutionPlanTarget<F::Backend>, AutomaticPlanningError> {
1024 let expected_backend = factory.backend_id();
1025 if plan.device.backend != expected_backend {
1026 return Err(AutomaticPlanningError::Invalid(format!(
1027 "execution plan selects backend {} but factory owns {}",
1028 plan.device.backend, expected_backend
1029 )));
1030 }
1031 plan.validate_structure()
1032 .map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
1033 if selected.execution_plan != *plan {
1034 return Err(AutomaticPlanningError::Invalid(
1035 "selected target was established for a different execution plan".into(),
1036 ));
1037 }
1038 let realization = factory.realize_target(selected)?;
1039 let descriptor = realization.backend().descriptor();
1040 if descriptor.name() != expected_backend.as_str() {
1041 return Err(AutomaticPlanningError::Invalid(format!(
1042 "factory identity {} does not match realized backend {}",
1043 expected_backend,
1044 descriptor.name()
1045 )));
1046 }
1047 let devices =
1048 realization
1049 .backend()
1050 .devices()
1051 .map_err(|error| AutomaticPlanningError::Backend {
1052 operation: "realize_execution_plan_devices",
1053 message: error.to_string(),
1054 })?;
1055 let capabilities = devices
1056 .iter()
1057 .find_map(|(device, capabilities)| {
1058 (device.id() == plan.device.device).then_some(capabilities)
1059 })
1060 .ok_or_else(|| {
1061 AutomaticPlanningError::Invalid(format!(
1062 "realized backend {} does not expose selected device {}",
1063 expected_backend, plan.device.device
1064 ))
1065 })?;
1066 plan.validate_device_capabilities(capabilities)
1067 .map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
1068 Ok(realization)
1069}
1070
1071pub fn realize_execution_plan_drafting<F: ExecutionPlanBackendFactory>(
1073 factory: &F,
1074 plan: &ExecutionPlan,
1075 target: &ModelRuntime<F::Backend>,
1076 selected: SelectedExecutionPlanDrafting<F::SelectedDrafterPreparation>,
1077) -> Result<RealizedDrafting<F::Drafter>, AutomaticPlanningError> {
1078 if selected.execution_plan != *plan {
1079 return Err(AutomaticPlanningError::Invalid(
1080 "selected drafting was established for a different execution plan".into(),
1081 ));
1082 }
1083 if target.execution_plan_target_id() != Some(selected.target_id) {
1084 return Err(AutomaticPlanningError::Invalid(
1085 "selected drafting was established for a different realized target".into(),
1086 ));
1087 }
1088 match (&plan.drafting, selected.external_artifact.as_ref()) {
1089 (DraftingPlan::External { .. }, None) => {
1090 return Err(AutomaticPlanningError::Invalid(
1091 "external drafting requires proven tokenizer compatibility".into(),
1092 ));
1093 }
1094 (DraftingPlan::Disabled | DraftingPlan::Embedded { .. }, Some(_)) => {
1095 return Err(AutomaticPlanningError::Invalid(
1096 "tokenizer compatibility was supplied for a plan without an external assistant"
1097 .into(),
1098 ));
1099 }
1100 _ => {}
1101 }
1102 let drafting = factory.realize_drafting(plan, target, selected)?;
1103 let matches_plan = matches!(
1104 (&plan.drafting, &drafting),
1105 (DraftingPlan::Disabled, RealizedDrafting::Disabled)
1106 | (DraftingPlan::Embedded { .. }, RealizedDrafting::Embedded)
1107 | (DraftingPlan::External { .. }, RealizedDrafting::External(_))
1108 );
1109 if !matches_plan {
1110 return Err(AutomaticPlanningError::Invalid(
1111 "backend factory realized a drafting mode different from the execution plan".into(),
1112 ));
1113 }
1114 Ok(drafting)
1115}
1116
1117pub fn select_execution_plan_drafting<F: ExecutionPlanBackendFactory>(
1119 factory: &F,
1120 plan: &ExecutionPlan,
1121 target: &SelectedExecutionPlanTarget<F::Backend>,
1122 external_artifact: Option<ExternalDraftArtifact<F::DrafterPreparation>>,
1123) -> Result<SelectedExecutionPlanDrafting<F::SelectedDrafterPreparation>, AutomaticPlanningError> {
1124 if target.execution_plan != *plan {
1125 return Err(AutomaticPlanningError::Invalid(
1126 "selected target was established for a different execution plan".into(),
1127 ));
1128 }
1129 match (&plan.drafting, external_artifact.as_ref()) {
1130 (DraftingPlan::External { .. }, None) => {
1131 return Err(AutomaticPlanningError::Invalid(
1132 "external drafting requires proven tokenizer compatibility".into(),
1133 ));
1134 }
1135 (DraftingPlan::Disabled | DraftingPlan::Embedded { .. }, Some(_)) => {
1136 return Err(AutomaticPlanningError::Invalid(
1137 "tokenizer compatibility was supplied for a plan without an external assistant"
1138 .into(),
1139 ));
1140 }
1141 _ => {}
1142 }
1143 let external_artifact = factory.select_drafting(plan, target, external_artifact)?;
1144 Ok(SelectedExecutionPlanDrafting {
1145 execution_plan: plan.clone(),
1146 target_id: target.target_id,
1147 external_artifact,
1148 })
1149}
1150
1151#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1153pub enum AutomaticPlanningError {
1154 #[error("automatic planning error: {0}")]
1156 Invalid(String),
1157 #[error("automatic planning backend failed during {operation}: {message}")]
1159 Backend {
1160 operation: &'static str,
1162 message: String,
1164 },
1165}
1166
1167#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Default)]
1169pub struct AutomaticPlanner {
1170 policy: AutomaticPlannerPolicy,
1171}
1172
1173impl AutomaticPlanner {
1174 pub fn new(policy: AutomaticPlannerPolicy) -> Self {
1176 Self { policy }
1177 }
1178
1179 pub fn policy(&self) -> &AutomaticPlannerPolicy {
1181 &self.policy
1182 }
1183
1184 pub fn plan<B: AutomaticPlanningBackend>(
1186 &self,
1187 backend: &B,
1188 request: &AutomaticPlanRequest,
1189 ) -> Result<ExecutionPlanReport, AutomaticPlanningError> {
1190 Ok(self.plan_retained(backend, request)?.into_parts().0)
1191 }
1192
1193 pub fn plan_retained<B: AutomaticPlanningBackend>(
1195 &self,
1196 backend: &B,
1197 request: &AutomaticPlanRequest,
1198 ) -> Result<RetainedAutomaticPlan<B::Inspection>, AutomaticPlanningError> {
1199 validate_request(request, &self.policy)?;
1200 let backend_id = backend.backend_id();
1201 if request.device.backend != backend_id {
1202 return Err(AutomaticPlanningError::Invalid(format!(
1203 "selected planning backend {} cannot plan device owned by {}",
1204 backend_id, request.device.backend
1205 )));
1206 }
1207 let hardware = backend.discover_hardware()?;
1208 validate_device(&hardware, &request.device)?;
1209 let (mut resources, inspection) = backend.inspect_resources(&request.model_path)?;
1210 let selected_device =
1211 selected_device(&hardware, &request.device).expect("validated device is present");
1212 let device_capacity = memory_basis(
1213 observed_u64(&selected_device.available_memory_bytes),
1214 observed_u64(&selected_device.total_memory_bytes)
1215 .or_else(|| observed_u64(&hardware.physical_memory_bytes)),
1216 hardware.physical_memory_semantics,
1217 );
1218 let host_capacity = memory_basis(
1219 observed_u64(&hardware.available_memory_bytes),
1220 observed_u64(&hardware.physical_memory_bytes),
1221 hardware.physical_memory_semantics,
1222 );
1223 let device_budget = budget(
1224 device_capacity,
1225 self.policy.device_memory_fallback_bytes,
1226 self.policy.memory_headroom_percent,
1227 );
1228 let host_budget = budget(
1229 host_capacity,
1230 self.policy.host_memory_fallback_bytes,
1231 self.policy.memory_headroom_percent,
1232 );
1233 let model_bytes = observed_u64(&resources.materialized_parameter_bytes)
1234 .or_else(|| observed_u64(&resources.stored_tensor_bytes));
1235 let candidates = base_candidates(
1236 request.device.clone(),
1237 device_budget,
1238 host_budget,
1239 &self.policy,
1240 );
1241 let resident = backend.admit_candidate(&inspection, &candidates[0])?;
1242 let mut layerwise = backend.admit_candidate(&inspection, &candidates[1])?;
1243 let mut disk = backend.admit_candidate(&inspection, &candidates[2])?;
1244 let resident_fits = model_bytes.is_some_and(|bytes| bytes <= device_budget);
1245 let layerwise_host_fits = model_bytes.is_some_and(|bytes| {
1246 if hardware.physical_memory_semantics == HardwareMemorySemantics::Unified {
1247 bytes <= host_budget.saturating_mul(2)
1248 } else {
1249 bytes <= host_budget
1250 }
1251 });
1252 if !resident_fits || !resident.supported {
1253 apply_bounded_probe(
1254 backend,
1255 &inspection,
1256 &candidates[1],
1257 device_budget,
1258 &mut layerwise,
1259 &mut resources,
1260 false,
1261 )?;
1262 apply_bounded_probe(
1263 backend,
1264 &inspection,
1265 &candidates[2],
1266 device_budget,
1267 &mut disk,
1268 &mut resources,
1269 true,
1270 )?;
1271 }
1272 let selected =
1273 if resident_fits && resident.supported {
1274 0
1275 } else if layerwise_host_fits && layerwise.supported {
1276 1
1277 } else if disk.supported {
1278 2
1279 } else {
1280 return Err(AutomaticPlanningError::Invalid(format!(
1281 "no loadable single-device policy: resident: {}; layerwise: {}; disk-streamed: {}",
1282 rejection(&resident), rejection(&layerwise), rejection(&disk)
1283 )));
1284 };
1285 let mut plan = candidates[selected].clone();
1286 let mut entries = vec![PlanExplanationEntry {
1287 level: PlanExplanationLevel::Decision,
1288 code: "single_device_scope".into(),
1289 detail: format!(
1290 "automatic planning is restricted to {}:{} with {}% memory headroom",
1291 request.device.backend, request.device.device, self.policy.memory_headroom_percent
1292 ),
1293 }];
1294 if selected > 0 {
1295 entries.push(PlanExplanationEntry {
1296 level: PlanExplanationLevel::Rejection,
1297 code: "fully_resident_not_admitted".into(),
1298 detail: resident
1299 .rejection
1300 .unwrap_or_else(|| "the model exceeds the device memory budget".into()),
1301 });
1302 }
1303 if selected > 1 {
1304 entries.push(PlanExplanationEntry {
1305 level: PlanExplanationLevel::Rejection,
1306 code: "layerwise_not_admitted".into(),
1307 detail: layerwise
1308 .rejection
1309 .unwrap_or_else(|| "the model exceeds the host-backed admission budget".into()),
1310 });
1311 }
1312 let mut summary = match selected {
1313 0 => "selected fully resident execution for the lowest expected latency".to_string(),
1314 1 => "selected host-backed layerwise execution with a validated bounded device window"
1315 .to_string(),
1316 _ => "selected bounded dense disk streaming because resident and layerwise admission failed"
1317 .to_string(),
1318 };
1319
1320 if selected > 0 {
1321 let expert_plan = with_expert_cache(plan.clone(), &self.policy);
1322 let expert = backend.admit_candidate(&inspection, &expert_plan)?;
1323 if expert.supported {
1324 plan = expert_plan;
1325 entries.push(PlanExplanationEntry {
1326 level: PlanExplanationLevel::Decision,
1327 code: "expert_cache_selected".into(),
1328 detail: "the backend admitted independent routed-expert caching".into(),
1329 });
1330 }
1331 }
1332
1333 let embedded_layers = resources.embedded_draft_layers.value().copied();
1334 if embedded_layers.is_some_and(|layers| layers > 0) {
1335 plan.drafting = DraftingPlan::Embedded {
1336 max_draft_tokens: self.policy.embedded_mtp_draft_tokens,
1337 lookahead: true,
1338 adaptive_lookahead: true,
1339 };
1340 entries.push(PlanExplanationEntry {
1341 level: PlanExplanationLevel::Decision,
1342 code: "embedded_mtp_selected".into(),
1343 detail: "checkpoint metadata advertises embedded prediction layers".into(),
1344 });
1345 }
1346
1347 if let Some((feedback, samples, median)) = select_feedback_plan(
1348 backend,
1349 &inspection,
1350 request,
1351 &hardware,
1352 &resources,
1353 &self.policy,
1354 embedded_layers,
1355 )? {
1356 plan = feedback;
1357 summary = format!(
1358 "selected a previously observed plan at {median:.2} median decode tokens/s"
1359 );
1360 entries.push(PlanExplanationEntry {
1361 level: PlanExplanationLevel::Decision,
1362 code: "prior_telemetry_selected".into(),
1363 detail: format!("selected using {samples} matching runtime sample(s)"),
1364 });
1365 }
1366
1367 let final_admission = backend.admit_candidate(&inspection, &plan)?;
1368 if !final_admission.supported {
1369 return Err(AutomaticPlanningError::Invalid(format!(
1370 "selected final plan is not loadable: {}",
1371 rejection(&final_admission)
1372 )));
1373 }
1374 if !matches!(plan.residency(), ResidencyPlan::FullyResident) {
1375 let final_budget = match plan.residency() {
1376 ResidencyPlan::LayerwiseHost {
1377 device_budget_bytes,
1378 ..
1379 } => device_budget_bytes.unwrap_or(device_budget),
1380 ResidencyPlan::DenseDiskStream {
1381 device_budget_bytes,
1382 ..
1383 } => *device_budget_bytes,
1384 ResidencyPlan::FullyResident => unreachable!(),
1385 };
1386 let mut final_probe = final_admission;
1387 apply_bounded_probe(
1388 backend,
1389 &inspection,
1390 &plan,
1391 final_budget,
1392 &mut final_probe,
1393 &mut resources,
1394 matches!(plan.residency(), ResidencyPlan::DenseDiskStream { .. }),
1395 )?;
1396 if !final_probe.supported {
1397 return Err(AutomaticPlanningError::Invalid(format!(
1398 "selected final plan exceeds its exact bounded residency: {}",
1399 rejection(&final_probe)
1400 )));
1401 }
1402 }
1403 let report = ExecutionPlanReport {
1404 schema_version: AUTOMATIC_SCHEMA_VERSION,
1405 hardware,
1406 resources,
1407 plan,
1408 explanation: PlanExplanation { summary, entries },
1409 };
1410 Ok(RetainedAutomaticPlan { report, inspection })
1411 }
1412}
1413
1414fn observed_u64(value: &Observed<u64>) -> Option<u64> {
1415 value.value().copied()
1416}
1417
1418fn validate_request(
1419 request: &AutomaticPlanRequest,
1420 policy: &AutomaticPlannerPolicy,
1421) -> Result<(), AutomaticPlanningError> {
1422 if request.schema_version != AUTOMATIC_SCHEMA_VERSION {
1423 return Err(AutomaticPlanningError::Invalid(format!(
1424 "automatic request schema {} does not match supported schema {}",
1425 request.schema_version, AUTOMATIC_SCHEMA_VERSION
1426 )));
1427 }
1428 if policy.device_memory_fallback_bytes == 0 || policy.host_memory_fallback_bytes == 0 {
1429 return Err(AutomaticPlanningError::Invalid(
1430 "automatic fallback memory budgets must be greater than zero".into(),
1431 ));
1432 }
1433 if policy.memory_headroom_percent >= 100
1434 || policy.expert_cache_share_percent == 0
1435 || policy.expert_cache_share_percent >= 100
1436 || policy.device_layer_window == 0
1437 || policy.max_cached_shards == 0
1438 || policy.embedded_mtp_draft_tokens == 0
1439 || policy.minimum_feedback_tokens == 0
1440 {
1441 return Err(AutomaticPlanningError::Invalid(
1442 "automatic percentage and count policy values are outside their valid ranges".into(),
1443 ));
1444 }
1445 Ok(())
1446}
1447
1448fn selected_device<'a>(
1449 hardware: &'a HardwareProfile,
1450 device: &DevicePlan,
1451) -> Option<&'a HardwareDeviceProfile> {
1452 hardware
1453 .backends
1454 .iter()
1455 .find(|backend| backend.backend == device.backend && backend.available)
1456 .and_then(|backend| backend.devices.iter().find(|item| item.id == device.device))
1457}
1458
1459fn validate_device(
1460 hardware: &HardwareProfile,
1461 device: &DevicePlan,
1462) -> Result<(), AutomaticPlanningError> {
1463 selected_device(hardware, device)
1464 .map(|_| ())
1465 .ok_or_else(|| {
1466 AutomaticPlanningError::Invalid(format!(
1467 "hardware discovery did not report available {} device {}",
1468 device.backend, device.device
1469 ))
1470 })
1471}
1472
1473fn memory_basis(
1474 available: Option<u64>,
1475 physical: Option<u64>,
1476 semantics: HardwareMemorySemantics,
1477) -> Option<u64> {
1478 available.or_else(|| {
1479 (semantics == HardwareMemorySemantics::Unified)
1480 .then_some(physical)
1481 .flatten()
1482 })
1483}
1484
1485fn budget(available: Option<u64>, fallback: u64, headroom_percent: u8) -> u64 {
1486 available
1487 .map(|bytes| bytes.saturating_mul(u64::from(100 - headroom_percent)) / 100)
1488 .unwrap_or(fallback)
1489 .max(1)
1490}
1491
1492fn base_candidates(
1493 device: DevicePlan,
1494 device_budget: u64,
1495 host_budget: u64,
1496 policy: &AutomaticPlannerPolicy,
1497) -> [ExecutionPlan; 3] {
1498 let mut resident = ExecutionPlan::fully_resident(device);
1499 resident.max_cached_shards = policy.max_cached_shards;
1500 let mut layerwise = resident.clone();
1501 layerwise.residency = ResidencyPlan::LayerwiseHost {
1502 device_layer_window: policy.device_layer_window,
1503 device_budget_bytes: Some(device_budget),
1504 host_budget_bytes: Some(host_budget),
1505 };
1506 let mut disk = resident.clone();
1507 disk.residency = ResidencyPlan::DenseDiskStream {
1508 device_budget_bytes: device_budget,
1509 host_budget_bytes: host_budget,
1510 host_lookahead: usize::from(host_budget > 0) * 2,
1511 background_queue: usize::from(host_budget > 0) * 2,
1512 };
1513 [resident, layerwise, disk]
1514}
1515
1516fn apply_bounded_probe<B: AutomaticPlanningBackend>(
1517 backend: &B,
1518 inspection: &B::Inspection,
1519 plan: &ExecutionPlan,
1520 budget: u64,
1521 admission: &mut CandidateAdmission,
1522 resources: &mut ModelResourceProfile,
1523 adjacent: bool,
1524) -> Result<(), AutomaticPlanningError> {
1525 if !admission.supported {
1526 return Ok(());
1527 }
1528 let requirement = backend.bounded_residency_requirement(inspection, plan)?;
1529 if requirement.required_bytes > budget {
1530 admission.supported = false;
1531 admission.rejection = Some(format!(
1532 "device budget {budget} bytes cannot contain {} pinned static bytes plus the depth-{} device window ({} bytes, {} total)",
1533 requirement.static_bytes,
1534 requirement.depth,
1535 requirement.window_bytes,
1536 requirement.required_bytes
1537 ));
1538 }
1539 resources.pinned_parameter_bytes =
1540 Observed::exact(requirement.static_bytes, "validated backend parameter plan");
1541 if adjacent {
1542 resources.largest_adjacent_execution_groups_bytes =
1543 Observed::exact(requirement.window_bytes, "validated backend parameter plan");
1544 } else {
1545 resources.largest_execution_group_bytes =
1546 Observed::exact(requirement.window_bytes, "validated backend parameter plan");
1547 }
1548 Ok(())
1549}
1550
1551fn rejection(admission: &CandidateAdmission) -> &str {
1552 admission.rejection.as_deref().unwrap_or("not admitted")
1553}
1554
1555fn with_expert_cache(mut plan: ExecutionPlan, policy: &AutomaticPlannerPolicy) -> ExecutionPlan {
1556 let split = |bytes: u64, percent: u8| bytes.saturating_mul(u64::from(percent)) / 100;
1557 let ordinary_share = 100 - policy.expert_cache_share_percent;
1558 let (device_budget, host_budget) = match &mut plan.residency {
1559 ResidencyPlan::FullyResident => (
1560 policy.device_memory_fallback_bytes,
1561 policy.host_memory_fallback_bytes,
1562 ),
1563 ResidencyPlan::LayerwiseHost {
1564 device_budget_bytes,
1565 host_budget_bytes,
1566 ..
1567 } => {
1568 let device = device_budget_bytes.unwrap_or(policy.device_memory_fallback_bytes);
1569 let host = host_budget_bytes.unwrap_or(policy.host_memory_fallback_bytes);
1570 *device_budget_bytes = Some(split(device, ordinary_share).max(1));
1571 *host_budget_bytes = Some(split(host, ordinary_share).max(1));
1572 (device, host)
1573 }
1574 ResidencyPlan::DenseDiskStream {
1575 device_budget_bytes,
1576 host_budget_bytes,
1577 ..
1578 } => {
1579 let (device, host) = (*device_budget_bytes, *host_budget_bytes);
1580 *device_budget_bytes = split(device, ordinary_share).max(1);
1581 *host_budget_bytes = split(host, ordinary_share).max(1);
1582 (device, host)
1583 }
1584 };
1585 let scratch = (1_u64 << 30).min(device_budget.max(1));
1586 plan.expert_cache = Some(ExpertCachePlan {
1587 device_budget_bytes: Some(split(device_budget, policy.expert_cache_share_percent).max(1)),
1588 host_budget_bytes: Some(split(host_budget, policy.expert_cache_share_percent).max(1)),
1589 scratch_bytes: scratch,
1590 prefill_bank_bytes: scratch,
1591 eviction_policy: crate::residency::CacheEvictionPolicy::LeastRecentlyUsed,
1592 });
1593 plan
1594}
1595
1596fn select_feedback_plan<B: AutomaticPlanningBackend>(
1597 backend: &B,
1598 inspection: &B::Inspection,
1599 request: &AutomaticPlanRequest,
1600 hardware: &HardwareProfile,
1601 resources: &ModelResourceProfile,
1602 policy: &AutomaticPlannerPolicy,
1603 embedded_layers: Option<usize>,
1604) -> Result<Option<(ExecutionPlan, usize, f64)>, AutomaticPlanningError> {
1605 let mut groups: Vec<(ExecutionPlan, Vec<f64>)> = Vec::new();
1606 for telemetry in &request.prior_telemetry {
1607 let (Some(plan), Some(prior_hardware), Some(prior_resources)) = (
1608 telemetry.plan.as_ref(),
1609 telemetry.hardware.as_ref(),
1610 telemetry.resources.as_ref(),
1611 ) else {
1612 continue;
1613 };
1614 if telemetry.schema_version != AUTOMATIC_SCHEMA_VERSION
1615 || telemetry.generated_tokens < policy.minimum_feedback_tokens
1616 || plan.device != request.device
1617 || prior_resources.path != resources.path
1618 || prior_resources.artifact_format != resources.artifact_format
1619 || prior_resources.model_family != resources.model_family
1620 || prior_hardware.operating_system != hardware.operating_system
1621 || prior_hardware.architecture != hardware.architecture
1622 || matches!(plan.drafting, DraftingPlan::External { .. })
1623 || (matches!(plan.drafting, DraftingPlan::Embedded { .. })
1624 && embedded_layers == Some(0))
1625 {
1626 continue;
1627 }
1628 let rate = telemetry
1629 .timing
1630 .decode_token_rate
1631 .filter(|value| value.is_finite() && *value > 0.0)
1632 .or_else(|| {
1633 (telemetry.timing.token_rate.is_finite() && telemetry.timing.token_rate > 0.0)
1634 .then_some(telemetry.timing.token_rate)
1635 });
1636 let Some(rate) = rate else { continue };
1637 if let Some((_, rates)) = groups.iter_mut().find(|(candidate, _)| candidate == plan) {
1638 rates.push(rate);
1639 } else {
1640 groups.push((plan.clone(), vec![rate]));
1641 }
1642 }
1643 let mut accepted = Vec::new();
1644 for (plan, mut rates) in groups {
1645 if !backend.admit_candidate(inspection, &plan)?.supported {
1646 continue;
1647 }
1648 rates.sort_by(f64::total_cmp);
1649 let middle = rates.len() / 2;
1650 let median = if rates.len() % 2 == 0 {
1651 (rates[middle - 1] + rates[middle]) / 2.0
1652 } else {
1653 rates[middle]
1654 };
1655 accepted.push((plan, rates.len(), median));
1656 }
1657 Ok(accepted
1658 .into_iter()
1659 .max_by(|left, right| left.2.total_cmp(&right.2)))
1660}
1661
1662#[cfg(test)]
1663mod tests {
1664 use super::*;
1665 use crate::execution::BackendId;
1666
1667 #[test]
1668 fn host_observation_preserves_injected_memory_and_backend_facts() {
1669 let physical = Observed::exact(4096, "foreign memory provider");
1670 let available = Observed::unavailable("not measured");
1671 let backend = HardwareBackendProfile {
1672 backend: BackendId::new("independent").unwrap(),
1673 available: false,
1674 detail: Some("no native context created".into()),
1675 devices: vec![],
1676 };
1677 let profile = HardwareProfile::observe_host(
1678 physical.clone(),
1679 available.clone(),
1680 HardwareMemorySemantics::SeparateTiers,
1681 vec![backend.clone()],
1682 );
1683 assert_eq!(profile.physical_memory_bytes, physical);
1684 assert_eq!(profile.available_memory_bytes, available);
1685 assert_eq!(profile.backends, vec![backend]);
1686 assert_eq!(
1687 profile.physical_memory_semantics,
1688 HardwareMemorySemantics::SeparateTiers
1689 );
1690 assert!(!profile.operating_system.is_empty());
1691 assert!(!profile.architecture.is_empty());
1692 }
1693
1694 #[test]
1695 fn physical_memory_semantics_preserve_unknown_and_separate_capacity() {
1696 use crate::capability::PhysicalMemorySemantics;
1697 for (physical, hardware) in [
1698 (
1699 PhysicalMemorySemantics::Unified,
1700 HardwareMemorySemantics::Unified,
1701 ),
1702 (
1703 PhysicalMemorySemantics::SeparateTiers,
1704 HardwareMemorySemantics::SeparateTiers,
1705 ),
1706 (
1707 PhysicalMemorySemantics::Unknown,
1708 HardwareMemorySemantics::Unknown,
1709 ),
1710 ] {
1711 assert_eq!(HardwareMemorySemantics::from(physical), hardware);
1712 }
1713 }
1714
1715 struct MockPlanningBackend {
1716 model_bytes: u64,
1717 embedded_layers: usize,
1718 }
1719
1720 impl Default for MockPlanningBackend {
1721 fn default() -> Self {
1722 Self {
1723 model_bytes: 2 << 30,
1724 embedded_layers: 0,
1725 }
1726 }
1727 }
1728
1729 impl AutomaticPlanningBackend for MockPlanningBackend {
1730 type Inspection = ();
1731
1732 fn backend_id(&self) -> BackendId {
1733 BackendId::new("mock").unwrap()
1734 }
1735
1736 fn discover_hardware(&self) -> Result<HardwareProfile, AutomaticPlanningError> {
1737 Ok(HardwareProfile {
1738 schema_version: AUTOMATIC_SCHEMA_VERSION,
1739 operating_system: "test".into(),
1740 architecture: "mock".into(),
1741 logical_cpu_count: Observed::exact(8, "fixture"),
1742 physical_memory_bytes: Observed::exact(32 << 30, "fixture"),
1743 available_memory_bytes: Observed::exact(24 << 30, "fixture"),
1744 physical_memory_semantics: HardwareMemorySemantics::SeparateTiers,
1745 backends: vec![HardwareBackendProfile {
1746 backend: BackendId::new("mock").unwrap(),
1747 available: true,
1748 detail: None,
1749 devices: vec![HardwareDeviceProfile {
1750 id: "gpu:0".into(),
1751 family: "gpu".into(),
1752 index: 0,
1753 total_memory_bytes: Observed::exact(16 << 30, "fixture"),
1754 available_memory_bytes: Observed::exact(12 << 30, "fixture"),
1755 }],
1756 }],
1757 })
1758 }
1759
1760 fn inspect_resources(
1761 &self,
1762 path: &std::path::Path,
1763 ) -> Result<(ModelResourceProfile, Self::Inspection), AutomaticPlanningError> {
1764 let mut profile =
1765 ModelResourceProfile::unmeasured(path.into(), ArtifactFormat::SafeTensors);
1766 profile.model_family = Some("llama".into());
1767 profile.embedded_draft_layers =
1768 Observed::exact(self.embedded_layers, "normalized architecture fixture");
1769 profile.stored_tensor_bytes = Observed::exact(self.model_bytes, "fixture");
1770 profile.materialized_parameter_bytes = Observed::exact(self.model_bytes, "fixture");
1771 Ok((profile, ()))
1772 }
1773
1774 fn admit_candidate(
1775 &self,
1776 _inspection: &Self::Inspection,
1777 _plan: &ExecutionPlan,
1778 ) -> Result<CandidateAdmission, AutomaticPlanningError> {
1779 Ok(CandidateAdmission {
1780 supported: true,
1781 rejection: None,
1782 })
1783 }
1784
1785 fn bounded_residency_requirement(
1786 &self,
1787 _inspection: &Self::Inspection,
1788 _plan: &ExecutionPlan,
1789 ) -> Result<BoundedResidencyRequirement, AutomaticPlanningError> {
1790 Ok(BoundedResidencyRequirement {
1791 static_bytes: 1 << 20,
1792 window_bytes: 2 << 20,
1793 required_bytes: 3 << 20,
1794 depth: 1,
1795 })
1796 }
1797 }
1798
1799 #[test]
1800 fn neutral_planner_selects_a_mock_backend_session_plan() {
1801 let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
1802 let report = AutomaticPlanner::default()
1803 .plan(&MockPlanningBackend::default(), &request)
1804 .unwrap();
1805 assert_eq!(report.plan.device.backend.as_str(), "mock");
1806 assert_eq!(report.plan.residency, ResidencyPlan::FullyResident);
1807 }
1808
1809 #[test]
1810 fn neutral_planner_selects_bounded_residency_and_embedded_drafting() {
1811 let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
1812 let report = AutomaticPlanner::default()
1813 .plan(
1814 &MockPlanningBackend {
1815 model_bytes: 10 << 30,
1816 embedded_layers: 2,
1817 },
1818 &request,
1819 )
1820 .unwrap();
1821 assert!(matches!(
1822 report.plan.residency,
1823 ResidencyPlan::LayerwiseHost { .. }
1824 ));
1825 assert!(matches!(
1826 report.plan.drafting,
1827 DraftingPlan::Embedded { .. }
1828 ));
1829 assert_eq!(
1830 observed_u64(&report.resources.pinned_parameter_bytes),
1831 Some(1 << 20)
1832 );
1833 }
1834
1835 #[test]
1836 fn selected_backend_identity_fails_closed() {
1837 let request =
1838 AutomaticPlanRequest::new("model", DevicePlan::new("other", "gpu:0").unwrap());
1839 assert!(matches!(
1840 AutomaticPlanner::default().plan(&MockPlanningBackend::default(), &request),
1841 Err(AutomaticPlanningError::Invalid(message))
1842 if message.contains("cannot plan device")
1843 ));
1844 }
1845
1846 #[test]
1847 fn documents_round_trip_without_an_accelerator_runtime() {
1848 let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
1849 let encoded = serde_json::to_vec(&request).unwrap();
1850 assert_eq!(
1851 serde_json::from_slice::<AutomaticPlanRequest>(&encoded).unwrap(),
1852 request
1853 );
1854 let unavailable = serde_json::to_value(Observed::<u64>::unavailable("unknown")).unwrap();
1855 assert!(unavailable.get("value").is_none());
1856 }
1857
1858 #[test]
1859 fn tokenizer_compatibility_requires_identical_vocabularies() {
1860 let fingerprint = [7; 32];
1861 let proof = TokenizerCompatibilityProof::prove(fingerprint, fingerprint).unwrap();
1862 assert_eq!(proof.fingerprint(), fingerprint);
1863 assert_eq!(proof.validate_target(fingerprint), Ok(()));
1864 assert_eq!(
1865 proof.validate_target([8; 32]),
1866 Err(TokenizerCompatibilityError)
1867 );
1868 assert_eq!(
1869 TokenizerCompatibilityProof::prove(fingerprint, [8; 32]),
1870 Err(TokenizerCompatibilityError)
1871 );
1872 }
1873
1874 struct RetainedPlanningBackend {
1875 inner: MockPlanningBackend,
1876 inspections: std::cell::Cell<usize>,
1877 admissions: std::cell::RefCell<Vec<ExecutionPlan>>,
1878 bounded_probes: std::cell::RefCell<Vec<ExecutionPlan>>,
1879 }
1880
1881 impl AutomaticPlanningBackend for RetainedPlanningBackend {
1882 type Inspection = usize;
1883
1884 fn backend_id(&self) -> BackendId {
1885 self.inner.backend_id()
1886 }
1887
1888 fn discover_hardware(&self) -> Result<HardwareProfile, AutomaticPlanningError> {
1889 self.inner.discover_hardware()
1890 }
1891
1892 fn inspect_resources(
1893 &self,
1894 path: &std::path::Path,
1895 ) -> Result<(ModelResourceProfile, Self::Inspection), AutomaticPlanningError> {
1896 self.inspections.set(self.inspections.get() + 1);
1897 self.inner
1898 .inspect_resources(path)
1899 .map(|(resources, ())| (resources, 7))
1900 }
1901
1902 fn admit_candidate(
1903 &self,
1904 inspection: &Self::Inspection,
1905 plan: &ExecutionPlan,
1906 ) -> Result<CandidateAdmission, AutomaticPlanningError> {
1907 assert_eq!(*inspection, 7, "every admission must reuse one inspection");
1908 self.admissions.borrow_mut().push(plan.clone());
1909 self.inner.admit_candidate(&(), plan)
1910 }
1911
1912 fn bounded_residency_requirement(
1913 &self,
1914 inspection: &Self::Inspection,
1915 plan: &ExecutionPlan,
1916 ) -> Result<BoundedResidencyRequirement, AutomaticPlanningError> {
1917 assert_eq!(
1918 *inspection, 7,
1919 "every bounded probe must reuse one inspection"
1920 );
1921 self.bounded_probes.borrow_mut().push(plan.clone());
1922 self.inner.bounded_residency_requirement(&(), plan)
1923 }
1924 }
1925
1926 #[test]
1927 fn automatic_planning_retains_one_inspection_and_exactly_reprobes_the_final_plan() {
1928 let backend = RetainedPlanningBackend {
1929 inner: MockPlanningBackend {
1930 model_bytes: 10 << 30,
1931 embedded_layers: 2,
1932 },
1933 inspections: std::cell::Cell::new(0),
1934 admissions: std::cell::RefCell::new(Vec::new()),
1935 bounded_probes: std::cell::RefCell::new(Vec::new()),
1936 };
1937 let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
1938 let retained = AutomaticPlanner::default()
1939 .plan_retained(&backend, &request)
1940 .unwrap();
1941 assert_eq!(backend.inspections.get(), 1);
1942 assert_eq!(
1943 backend.admissions.borrow().last(),
1944 Some(&retained.report().plan)
1945 );
1946 assert_eq!(
1947 backend.bounded_probes.borrow().last(),
1948 Some(&retained.report().plan),
1949 "drafting/expert/feedback mutations must be exact-probed, not only their base candidate"
1950 );
1951 assert!(matches!(
1952 retained.report().plan.drafting(),
1953 DraftingPlan::Embedded { .. }
1954 ));
1955 let (_, inspection) = retained.into_parts();
1956 assert_eq!(inspection, 7);
1957 }
1958
1959 #[test]
1960 fn automatic_feedback_cannot_select_an_uninspected_external_assistant() {
1961 let backend = MockPlanningBackend::default();
1962 let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
1963 let hardware = backend.discover_hardware().unwrap();
1964 let resources = backend.inspect_resources(&request.model_path).unwrap().0;
1965 let external = ExecutionPlan::fully_resident(request.device.clone()).with_drafting(
1966 DraftingPlan::External {
1967 model: "missing-assistant".into(),
1968 placement: crate::execution::DraftPlacementPlan::Target,
1969 max_draft_tokens: 2,
1970 lookahead: false,
1971 adaptive_lookahead: false,
1972 },
1973 );
1974 let telemetry = ExecutionTelemetry {
1975 schema_version: AUTOMATIC_SCHEMA_VERSION,
1976 effective_model_type: "fixture".into(),
1977 plan: Some(external),
1978 plan_explanation: None,
1979 hardware: Some(hardware),
1980 resources: Some(resources),
1981 prompt_tokens: 1,
1982 generated_tokens: 1,
1983 stop_reason: "length".into(),
1984 timing: TimingTelemetry::new(
1985 Duration::from_secs(1),
1986 Duration::from_secs(1),
1987 None,
1988 100,
1989 Duration::from_secs(2),
1990 ),
1991 allocator: None,
1992 residency: None,
1993 expert_cache: None,
1994 speculative: None,
1995 };
1996
1997 let report = AutomaticPlanner::default()
1998 .plan(&backend, &request.with_prior_telemetry([telemetry]))
1999 .unwrap();
2000
2001 assert!(matches!(report.plan.drafting(), DraftingPlan::Disabled));
2002 assert!(!report
2003 .explanation
2004 .entries
2005 .iter()
2006 .any(|entry| entry.code == "prior_telemetry_selected"));
2007 }
2008
2009 #[test]
2010 fn zero_duration_rates_are_finite() {
2011 let timing = TimingTelemetry::new(
2012 Duration::ZERO,
2013 Duration::ZERO,
2014 Some(Duration::ZERO),
2015 3,
2016 Duration::ZERO,
2017 );
2018 assert_eq!(timing.token_rate, 0.0);
2019 assert_eq!(timing.decode_token_rate, Some(0.0));
2020 }
2021}