1use crate::{
8 artifact::ArtifactFormat,
9 backend::{BackendProvider, ModelLoadingBackend, ModelRuntime},
10 execution::{
11 DevicePlan, DraftingPlan, ExecutionPlan, ExpertCachePlan, ResidencyPlan,
12 DEFAULT_MAX_CACHED_SHARDS,
13 },
14 speculative::SpeculativeDraft,
15};
16use serde::{Deserialize, Serialize};
17use std::{path::PathBuf, time::Duration};
18
19pub const AUTOMATIC_SCHEMA_VERSION: u32 = 6;
21
22#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
24#[serde(rename_all = "snake_case")]
25pub enum ObservationKind {
26 Exact,
28 Conservative,
30 Observational,
32 Estimated,
34}
35
36#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
38#[serde(tag = "status", rename_all = "snake_case")]
39pub enum Observed<T> {
40 Available {
42 value: T,
44 kind: ObservationKind,
46 source: String,
48 },
49 Unsupported {
51 reason: String,
53 },
54 Unavailable {
56 reason: String,
58 },
59}
60
61impl<T> Observed<T> {
62 pub fn exact(value: T, source: impl Into<String>) -> Self {
64 Self::Available {
65 value,
66 kind: ObservationKind::Exact,
67 source: source.into(),
68 }
69 }
70
71 pub fn unavailable(reason: impl Into<String>) -> Self {
73 Self::Unavailable {
74 reason: reason.into(),
75 }
76 }
77
78 pub fn unsupported(reason: impl Into<String>) -> Self {
80 Self::Unsupported {
81 reason: reason.into(),
82 }
83 }
84
85 pub const fn value(&self) -> Option<&T> {
87 match self {
88 Self::Available { value, .. } => Some(value),
89 Self::Unsupported { .. } | Self::Unavailable { .. } => None,
90 }
91 }
92}
93
94fn unobserved_embedded_draft_layers() -> Observed<usize> {
95 Observed::unavailable("embedded drafting requires normalized architecture inspection")
96}
97
98#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
100pub struct ModelResourceProfile {
101 pub schema_version: u32,
103 pub path: PathBuf,
105 pub artifact_format: ArtifactFormat,
107 #[serde(skip_serializing_if = "Option::is_none")]
109 pub model_family: Option<String>,
110 #[serde(skip_serializing_if = "Option::is_none")]
112 pub architecture: Option<String>,
113 #[serde(skip_serializing_if = "Option::is_none")]
115 pub tensor_count: Option<usize>,
116 #[serde(skip_serializing_if = "Option::is_none")]
118 pub checkpoint_shards: Option<usize>,
119 #[serde(default = "unobserved_embedded_draft_layers")]
121 pub embedded_draft_layers: Observed<usize>,
122 pub stored_tensor_bytes: Observed<u64>,
124 pub largest_stored_tensor_bytes: Observed<u64>,
126 pub materialized_parameter_bytes: Observed<u64>,
128 pub pinned_parameter_bytes: Observed<u64>,
130 pub largest_execution_group_bytes: Observed<u64>,
132 pub largest_adjacent_execution_groups_bytes: Observed<u64>,
134 pub expert_parameter_bytes: Observed<u64>,
136}
137
138impl ModelResourceProfile {
139 pub fn unmeasured(path: PathBuf, artifact_format: ArtifactFormat) -> Self {
141 let unavailable = || {
142 Observed::unavailable("resource value requires a validated checkpoint parameter plan")
143 };
144 Self {
145 schema_version: AUTOMATIC_SCHEMA_VERSION,
146 path,
147 artifact_format,
148 model_family: None,
149 architecture: None,
150 tensor_count: None,
151 checkpoint_shards: None,
152 embedded_draft_layers: unobserved_embedded_draft_layers(),
153 stored_tensor_bytes: Observed::unavailable(
154 "checkpoint tensor catalog was not established",
155 ),
156 largest_stored_tensor_bytes: Observed::unavailable(
157 "checkpoint tensor catalog was not established",
158 ),
159 materialized_parameter_bytes: unavailable(),
160 pinned_parameter_bytes: unavailable(),
161 largest_execution_group_bytes: unavailable(),
162 largest_adjacent_execution_groups_bytes: unavailable(),
163 expert_parameter_bytes: unavailable(),
164 }
165 }
166}
167
168#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
170pub struct HardwareDeviceProfile {
171 pub id: String,
173 pub family: String,
175 pub index: usize,
177 pub total_memory_bytes: Observed<u64>,
179 pub available_memory_bytes: Observed<u64>,
181}
182
183#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
185pub struct HardwareBackendProfile {
186 pub backend: crate::execution::BackendId,
188 pub available: bool,
190 #[serde(skip_serializing_if = "Option::is_none")]
192 pub detail: Option<String>,
193 pub devices: Vec<HardwareDeviceProfile>,
195}
196
197#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
199pub struct HardwareProfile {
200 pub schema_version: u32,
202 pub operating_system: String,
204 pub architecture: String,
206 pub logical_cpu_count: Observed<u64>,
208 pub physical_memory_bytes: Observed<u64>,
210 pub available_memory_bytes: Observed<u64>,
212 pub physical_memory_semantics: HardwareMemorySemantics,
214 pub backends: Vec<HardwareBackendProfile>,
216}
217
218#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
220#[serde(rename_all = "snake_case")]
221pub enum HardwareMemorySemantics {
222 Unified,
224 SeparateTiers,
226 Unknown,
228}
229
230#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
232#[serde(rename_all = "snake_case")]
233pub enum PlanExplanationLevel {
234 Decision,
236 Warning,
238 Rejection,
240}
241
242#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
244pub struct PlanExplanationEntry {
245 pub level: PlanExplanationLevel,
247 pub code: String,
249 pub detail: String,
251}
252
253#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
255pub struct PlanExplanation {
256 pub summary: String,
258 pub entries: Vec<PlanExplanationEntry>,
260}
261
262#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
264pub struct ExecutionPlanReport {
265 pub schema_version: u32,
267 pub hardware: HardwareProfile,
269 pub resources: ModelResourceProfile,
271 pub plan: ExecutionPlan,
273 pub explanation: PlanExplanation,
275}
276
277#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
279#[serde(default)]
280#[non_exhaustive]
281pub struct AutomaticPlannerPolicy {
282 pub device_memory_fallback_bytes: u64,
284 pub host_memory_fallback_bytes: u64,
286 pub memory_headroom_percent: u8,
288 pub expert_cache_share_percent: u8,
290 pub device_layer_window: usize,
292 pub max_cached_shards: usize,
294 pub embedded_mtp_draft_tokens: usize,
296 pub minimum_feedback_tokens: usize,
298}
299
300impl Default for AutomaticPlannerPolicy {
301 fn default() -> Self {
302 Self {
303 device_memory_fallback_bytes: 4 << 30,
304 host_memory_fallback_bytes: 16 << 30,
305 memory_headroom_percent: 30,
306 expert_cache_share_percent: 40,
307 device_layer_window: 1,
308 max_cached_shards: DEFAULT_MAX_CACHED_SHARDS,
309 embedded_mtp_draft_tokens: 3,
310 minimum_feedback_tokens: 1,
311 }
312 }
313}
314
315#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
317pub struct TimingTelemetry {
318 pub load_seconds: f64,
320 pub generation_seconds: f64,
322 #[serde(skip_serializing_if = "Option::is_none")]
324 pub time_to_first_token_seconds: Option<f64>,
325 pub total_seconds: f64,
327 pub token_rate: f64,
329 #[serde(skip_serializing_if = "Option::is_none")]
331 pub decode_token_rate: Option<f64>,
332}
333
334impl TimingTelemetry {
335 pub fn new(
337 load: Duration,
338 generation: Duration,
339 time_to_first_token: Option<Duration>,
340 generated_tokens: usize,
341 total: Duration,
342 ) -> Self {
343 fn rate(tokens: usize, elapsed: Duration) -> f64 {
344 if elapsed.is_zero() {
345 0.0
346 } else {
347 tokens as f64 / elapsed.as_secs_f64()
348 }
349 }
350 Self {
351 load_seconds: load.as_secs_f64(),
352 generation_seconds: generation.as_secs_f64(),
353 time_to_first_token_seconds: time_to_first_token.map(|value| value.as_secs_f64()),
354 total_seconds: total.as_secs_f64(),
355 token_rate: rate(generated_tokens, generation),
356 decode_token_rate: time_to_first_token.map(|first| {
357 rate(
358 generated_tokens.saturating_sub(1),
359 generation.saturating_sub(first),
360 )
361 }),
362 }
363 }
364}
365
366#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
368pub struct AllocatorTelemetry {
369 pub peak_bytes: u64,
371 pub active_bytes: u64,
373 pub cache_bytes: u64,
375}
376
377#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
379pub struct ResidencyTelemetry {
380 pub planned_disk_bytes: u64,
382 pub planned_host_bytes: u64,
384 pub planned_device_bytes: u64,
386 pub current_host_bytes: u64,
388 pub current_device_bytes: u64,
390 pub peak_host_bytes: u64,
392 pub peak_device_bytes: u64,
394 pub transfers: Vec<TransferTelemetry>,
396}
397
398#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
400pub struct TransferTelemetry {
401 pub direction: String,
403 pub count: u64,
405 pub bytes: u64,
407 pub seconds: DurationSeconds,
409}
410
411#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
413#[serde(transparent)]
414pub struct DurationSeconds(pub f64);
415
416impl PartialEq for DurationSeconds {
417 fn eq(&self, other: &Self) -> bool {
418 self.0.to_bits() == other.0.to_bits()
419 }
420}
421impl Eq for DurationSeconds {}
422
423#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
425pub struct ExpertCacheTelemetry {
426 pub owned_experts: usize,
428 pub owned_bytes: u64,
430 pub host_resident_experts: usize,
432 pub device_resident_experts: usize,
434 pub host_resident_bytes: u64,
436 pub device_resident_bytes: u64,
438 pub peak_host_resident_bytes: u64,
440 pub peak_device_resident_bytes: u64,
442}
443
444#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
446pub struct SpeculativeDecodingTelemetry {
447 pub execution_topology: String,
449 pub target_tokens: usize,
451 pub draft_tokens: usize,
453 pub accepted_tokens: usize,
455 pub accept_rate: f64,
457 pub rounds: usize,
459 pub accept_lens: Vec<usize>,
461 pub emitted_tokens: usize,
463 pub optimistic_draft_tokens: usize,
465 pub reused_optimistic_tokens: usize,
467 pub discarded_optimistic_tokens: usize,
469 pub adaptive_lookahead_disabled: bool,
471 pub optimistic_draft_seconds: f64,
473 pub verification_in_flight_seconds: f64,
475}
476
477#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
479pub struct ExecutionTelemetry {
480 pub schema_version: u32,
482 pub effective_model_type: String,
484 #[serde(skip_serializing_if = "Option::is_none")]
486 pub plan: Option<ExecutionPlan>,
487 #[serde(skip_serializing_if = "Option::is_none")]
489 pub plan_explanation: Option<PlanExplanation>,
490 #[serde(skip_serializing_if = "Option::is_none")]
492 pub hardware: Option<HardwareProfile>,
493 #[serde(skip_serializing_if = "Option::is_none")]
495 pub resources: Option<ModelResourceProfile>,
496 pub prompt_tokens: usize,
498 pub generated_tokens: usize,
500 pub stop_reason: String,
502 pub timing: TimingTelemetry,
504 #[serde(skip_serializing_if = "Option::is_none")]
506 pub allocator: Option<AllocatorTelemetry>,
507 #[serde(skip_serializing_if = "Option::is_none")]
509 pub residency: Option<ResidencyTelemetry>,
510 #[serde(skip_serializing_if = "Option::is_none")]
512 pub expert_cache: Option<ExpertCacheTelemetry>,
513 #[serde(skip_serializing_if = "Option::is_none")]
515 pub speculative: Option<SpeculativeDecodingTelemetry>,
516}
517
518#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
520#[non_exhaustive]
521pub struct AutomaticPlanRequest {
522 pub schema_version: u32,
524 pub model_path: PathBuf,
526 pub device: DevicePlan,
528 #[serde(default, skip_serializing_if = "Vec::is_empty")]
530 pub prior_telemetry: Vec<ExecutionTelemetry>,
531}
532
533impl AutomaticPlanRequest {
534 pub fn new(model_path: impl Into<PathBuf>, device: DevicePlan) -> Self {
536 Self {
537 schema_version: AUTOMATIC_SCHEMA_VERSION,
538 model_path: model_path.into(),
539 device,
540 prior_telemetry: Vec::new(),
541 }
542 }
543
544 pub fn with_prior_telemetry(
546 mut self,
547 telemetry: impl IntoIterator<Item = ExecutionTelemetry>,
548 ) -> Self {
549 self.prior_telemetry.extend(telemetry);
550 self
551 }
552}
553
554#[derive(Debug, Clone, Eq, PartialEq)]
556pub struct CandidateAdmission {
557 pub supported: bool,
559 pub rejection: Option<String>,
561}
562
563#[derive(Debug, Clone, Copy, Eq, PartialEq)]
565pub struct BoundedResidencyRequirement {
566 pub static_bytes: u64,
568 pub window_bytes: u64,
570 pub required_bytes: u64,
572 pub depth: usize,
574}
575
576pub trait AutomaticPlanningBackend {
578 fn backend_id(&self) -> crate::execution::BackendId;
580 fn discover_hardware(&self) -> Result<HardwareProfile, AutomaticPlanningError>;
582 fn inspect_resources(
584 &self,
585 model_path: &std::path::Path,
586 ) -> Result<ModelResourceProfile, AutomaticPlanningError>;
587 fn admit_candidate(
589 &self,
590 model_path: &std::path::Path,
591 plan: &ExecutionPlan,
592 ) -> Result<CandidateAdmission, AutomaticPlanningError>;
593 fn bounded_residency_requirement(
595 &self,
596 model_path: &std::path::Path,
597 plan: &ExecutionPlan,
598 ) -> Result<BoundedResidencyRequirement, AutomaticPlanningError>;
599}
600
601pub struct ExecutionPlanTarget<B: ModelLoadingBackend> {
607 backend: B,
608 load_options: B::LoadOptions,
609}
610
611impl<B: ModelLoadingBackend> ExecutionPlanTarget<B> {
612 pub fn new(backend: B, load_options: B::LoadOptions) -> Self {
618 Self {
619 backend,
620 load_options,
621 }
622 }
623
624 pub const fn backend(&self) -> &B {
626 &self.backend
627 }
628
629 pub fn into_parts(self) -> (B, B::LoadOptions) {
631 (self.backend, self.load_options)
632 }
633}
634
635#[derive(Debug, Clone, Copy, Eq, PartialEq)]
641pub struct TokenizerCompatibilityProof {
642 fingerprint: [u8; 32],
643}
644
645impl TokenizerCompatibilityProof {
646 pub fn prove(
648 target_fingerprint: [u8; 32],
649 assistant_fingerprint: [u8; 32],
650 ) -> Result<Self, TokenizerCompatibilityError> {
651 if target_fingerprint != assistant_fingerprint {
652 return Err(TokenizerCompatibilityError);
653 }
654 Ok(Self {
655 fingerprint: target_fingerprint,
656 })
657 }
658
659 pub const fn fingerprint(self) -> [u8; 32] {
661 self.fingerprint
662 }
663
664 pub fn validate_target(
666 self,
667 target_fingerprint: [u8; 32],
668 ) -> Result<(), TokenizerCompatibilityError> {
669 if self.fingerprint != target_fingerprint {
670 return Err(TokenizerCompatibilityError);
671 }
672 Ok(())
673 }
674}
675
676#[derive(Debug, Clone, Copy, Eq, PartialEq, thiserror::Error)]
678#[error("assistant token-id vocabulary mapping does not match the target")]
679pub struct TokenizerCompatibilityError;
680
681#[derive(Debug, Clone, Eq, PartialEq)]
683pub struct ExternalDraftArtifact<P> {
684 pub preparation: P,
686 pub tokenizer_compatibility: TokenizerCompatibilityProof,
688}
689
690pub enum RealizedDrafting<D> {
692 Disabled,
694 Embedded,
696 External(D),
698}
699
700impl<D> RealizedDrafting<D> {
701 pub fn as_speculative_draft(&mut self) -> Option<SpeculativeDraft<'_, D>> {
703 match self {
704 Self::Disabled => None,
705 Self::Embedded => Some(SpeculativeDraft::Embedded),
706 Self::External(drafter) => Some(SpeculativeDraft::External(drafter)),
707 }
708 }
709
710 pub const fn is_external(&self) -> bool {
712 matches!(self, Self::External(_))
713 }
714}
715
716pub trait ExecutionPlanBackendFactory: AutomaticPlanningBackend {
723 type Backend: ModelLoadingBackend;
725 type DrafterPreparation;
727 type Drafter;
729
730 fn realize_target(
735 &self,
736 plan: &ExecutionPlan,
737 ) -> Result<ExecutionPlanTarget<Self::Backend>, AutomaticPlanningError>;
738
739 fn realize_drafting(
746 &self,
747 plan: &ExecutionPlan,
748 target: &ModelRuntime<Self::Backend>,
749 external_artifact: Option<ExternalDraftArtifact<Self::DrafterPreparation>>,
750 ) -> Result<RealizedDrafting<Self::Drafter>, AutomaticPlanningError>;
751}
752
753pub fn realize_execution_plan_target<F: ExecutionPlanBackendFactory>(
755 factory: &F,
756 plan: &ExecutionPlan,
757) -> Result<ExecutionPlanTarget<F::Backend>, AutomaticPlanningError> {
758 let expected_backend = factory.backend_id();
759 if plan.device.backend != expected_backend {
760 return Err(AutomaticPlanningError::Invalid(format!(
761 "execution plan selects backend {} but factory owns {}",
762 plan.device.backend, expected_backend
763 )));
764 }
765 plan.validate_structure()
766 .map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
767
768 let realization = factory.realize_target(plan)?;
769 let descriptor = realization.backend().descriptor();
770 if descriptor.name() != expected_backend.as_str() {
771 return Err(AutomaticPlanningError::Invalid(format!(
772 "factory identity {} does not match realized backend {}",
773 expected_backend,
774 descriptor.name()
775 )));
776 }
777 let devices =
778 realization
779 .backend()
780 .devices()
781 .map_err(|error| AutomaticPlanningError::Backend {
782 operation: "realize_execution_plan_devices",
783 message: error.to_string(),
784 })?;
785 let capabilities = devices
786 .iter()
787 .find_map(|(device, capabilities)| {
788 (device.id() == plan.device.device).then_some(capabilities)
789 })
790 .ok_or_else(|| {
791 AutomaticPlanningError::Invalid(format!(
792 "realized backend {} does not expose selected device {}",
793 expected_backend, plan.device.device
794 ))
795 })?;
796 plan.validate_device_capabilities(capabilities)
797 .map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
798 Ok(realization)
799}
800
801pub fn realize_execution_plan_drafting<F: ExecutionPlanBackendFactory>(
803 factory: &F,
804 plan: &ExecutionPlan,
805 target: &ModelRuntime<F::Backend>,
806 external_artifact: Option<ExternalDraftArtifact<F::DrafterPreparation>>,
807) -> Result<RealizedDrafting<F::Drafter>, AutomaticPlanningError> {
808 match (&plan.drafting, external_artifact.as_ref()) {
809 (DraftingPlan::External { .. }, None) => {
810 return Err(AutomaticPlanningError::Invalid(
811 "external drafting requires proven tokenizer compatibility".into(),
812 ));
813 }
814 (DraftingPlan::Disabled | DraftingPlan::Embedded { .. }, Some(_)) => {
815 return Err(AutomaticPlanningError::Invalid(
816 "tokenizer compatibility was supplied for a plan without an external assistant"
817 .into(),
818 ));
819 }
820 _ => {}
821 }
822 let drafting = factory.realize_drafting(plan, target, external_artifact)?;
823 let matches_plan = matches!(
824 (&plan.drafting, &drafting),
825 (DraftingPlan::Disabled, RealizedDrafting::Disabled)
826 | (DraftingPlan::Embedded { .. }, RealizedDrafting::Embedded)
827 | (DraftingPlan::External { .. }, RealizedDrafting::External(_))
828 );
829 if !matches_plan {
830 return Err(AutomaticPlanningError::Invalid(
831 "backend factory realized a drafting mode different from the execution plan".into(),
832 ));
833 }
834 Ok(drafting)
835}
836
837#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
839pub enum AutomaticPlanningError {
840 #[error("automatic planning error: {0}")]
842 Invalid(String),
843 #[error("automatic planning backend failed during {operation}: {message}")]
845 Backend {
846 operation: &'static str,
848 message: String,
850 },
851}
852
853#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Default)]
855pub struct AutomaticPlanner {
856 policy: AutomaticPlannerPolicy,
857}
858
859impl AutomaticPlanner {
860 pub fn new(policy: AutomaticPlannerPolicy) -> Self {
862 Self { policy }
863 }
864
865 pub fn policy(&self) -> &AutomaticPlannerPolicy {
867 &self.policy
868 }
869
870 pub fn plan<B: AutomaticPlanningBackend>(
872 &self,
873 backend: &B,
874 request: &AutomaticPlanRequest,
875 ) -> Result<ExecutionPlanReport, AutomaticPlanningError> {
876 validate_request(request, &self.policy)?;
877 let backend_id = backend.backend_id();
878 if request.device.backend != backend_id {
879 return Err(AutomaticPlanningError::Invalid(format!(
880 "selected planning backend {} cannot plan device owned by {}",
881 backend_id, request.device.backend
882 )));
883 }
884 let hardware = backend.discover_hardware()?;
885 validate_device(&hardware, &request.device)?;
886 let mut resources = backend.inspect_resources(&request.model_path)?;
887 let selected_device =
888 selected_device(&hardware, &request.device).expect("validated device is present");
889 let device_capacity = memory_basis(
890 observed_u64(&selected_device.available_memory_bytes),
891 observed_u64(&selected_device.total_memory_bytes)
892 .or_else(|| observed_u64(&hardware.physical_memory_bytes)),
893 hardware.physical_memory_semantics,
894 );
895 let host_capacity = memory_basis(
896 observed_u64(&hardware.available_memory_bytes),
897 observed_u64(&hardware.physical_memory_bytes),
898 hardware.physical_memory_semantics,
899 );
900 let device_budget = budget(
901 device_capacity,
902 self.policy.device_memory_fallback_bytes,
903 self.policy.memory_headroom_percent,
904 );
905 let host_budget = budget(
906 host_capacity,
907 self.policy.host_memory_fallback_bytes,
908 self.policy.memory_headroom_percent,
909 );
910 let model_bytes = observed_u64(&resources.materialized_parameter_bytes)
911 .or_else(|| observed_u64(&resources.stored_tensor_bytes));
912 let candidates = base_candidates(
913 request.device.clone(),
914 device_budget,
915 host_budget,
916 &self.policy,
917 );
918 let resident = backend.admit_candidate(&request.model_path, &candidates[0])?;
919 let mut layerwise = backend.admit_candidate(&request.model_path, &candidates[1])?;
920 let mut disk = backend.admit_candidate(&request.model_path, &candidates[2])?;
921 let resident_fits = model_bytes.is_some_and(|bytes| bytes <= device_budget);
922 let layerwise_host_fits = model_bytes.is_some_and(|bytes| {
923 if hardware.physical_memory_semantics == HardwareMemorySemantics::Unified {
924 bytes <= host_budget.saturating_mul(2)
925 } else {
926 bytes <= host_budget
927 }
928 });
929 if !resident_fits || !resident.supported {
930 apply_bounded_probe(
931 backend,
932 &request.model_path,
933 &candidates[1],
934 device_budget,
935 &mut layerwise,
936 &mut resources,
937 false,
938 )?;
939 apply_bounded_probe(
940 backend,
941 &request.model_path,
942 &candidates[2],
943 device_budget,
944 &mut disk,
945 &mut resources,
946 true,
947 )?;
948 }
949 let selected =
950 if resident_fits && resident.supported {
951 0
952 } else if layerwise_host_fits && layerwise.supported {
953 1
954 } else if disk.supported {
955 2
956 } else {
957 return Err(AutomaticPlanningError::Invalid(format!(
958 "no loadable single-device policy: resident: {}; layerwise: {}; disk-streamed: {}",
959 rejection(&resident), rejection(&layerwise), rejection(&disk)
960 )));
961 };
962 let mut plan = candidates[selected].clone();
963 let mut entries = vec![PlanExplanationEntry {
964 level: PlanExplanationLevel::Decision,
965 code: "single_device_scope".into(),
966 detail: format!(
967 "automatic planning is restricted to {}:{} with {}% memory headroom",
968 request.device.backend, request.device.device, self.policy.memory_headroom_percent
969 ),
970 }];
971 if selected > 0 {
972 entries.push(PlanExplanationEntry {
973 level: PlanExplanationLevel::Rejection,
974 code: "fully_resident_not_admitted".into(),
975 detail: resident
976 .rejection
977 .unwrap_or_else(|| "the model exceeds the device memory budget".into()),
978 });
979 }
980 if selected > 1 {
981 entries.push(PlanExplanationEntry {
982 level: PlanExplanationLevel::Rejection,
983 code: "layerwise_not_admitted".into(),
984 detail: layerwise
985 .rejection
986 .unwrap_or_else(|| "the model exceeds the host-backed admission budget".into()),
987 });
988 }
989 let mut summary = match selected {
990 0 => "selected fully resident execution for the lowest expected latency".to_string(),
991 1 => "selected host-backed layerwise execution with a validated bounded device window"
992 .to_string(),
993 _ => "selected bounded dense disk streaming because resident and layerwise admission failed"
994 .to_string(),
995 };
996
997 if selected > 0 {
998 let expert_plan = with_expert_cache(plan.clone(), &self.policy);
999 let expert = backend.admit_candidate(&request.model_path, &expert_plan)?;
1000 if expert.supported {
1001 plan = expert_plan;
1002 entries.push(PlanExplanationEntry {
1003 level: PlanExplanationLevel::Decision,
1004 code: "expert_cache_selected".into(),
1005 detail: "the backend admitted independent routed-expert caching".into(),
1006 });
1007 }
1008 }
1009
1010 let embedded_layers = resources.embedded_draft_layers.value().copied();
1011 if embedded_layers.is_some_and(|layers| layers > 0) {
1012 plan.drafting = DraftingPlan::Embedded {
1013 max_draft_tokens: self.policy.embedded_mtp_draft_tokens,
1014 lookahead: true,
1015 adaptive_lookahead: true,
1016 };
1017 entries.push(PlanExplanationEntry {
1018 level: PlanExplanationLevel::Decision,
1019 code: "embedded_mtp_selected".into(),
1020 detail: "checkpoint metadata advertises embedded prediction layers".into(),
1021 });
1022 }
1023
1024 if let Some((feedback, samples, median)) = select_feedback_plan(
1025 backend,
1026 request,
1027 &hardware,
1028 &resources,
1029 &self.policy,
1030 embedded_layers,
1031 )? {
1032 plan = feedback;
1033 summary = format!(
1034 "selected a previously observed plan at {median:.2} median decode tokens/s"
1035 );
1036 entries.push(PlanExplanationEntry {
1037 level: PlanExplanationLevel::Decision,
1038 code: "prior_telemetry_selected".into(),
1039 detail: format!("selected using {samples} matching runtime sample(s)"),
1040 });
1041 }
1042
1043 Ok(ExecutionPlanReport {
1044 schema_version: AUTOMATIC_SCHEMA_VERSION,
1045 hardware,
1046 resources,
1047 plan,
1048 explanation: PlanExplanation { summary, entries },
1049 })
1050 }
1051}
1052
1053fn observed_u64(value: &Observed<u64>) -> Option<u64> {
1054 value.value().copied()
1055}
1056
1057fn validate_request(
1058 request: &AutomaticPlanRequest,
1059 policy: &AutomaticPlannerPolicy,
1060) -> Result<(), AutomaticPlanningError> {
1061 if request.schema_version != AUTOMATIC_SCHEMA_VERSION {
1062 return Err(AutomaticPlanningError::Invalid(format!(
1063 "automatic request schema {} does not match supported schema {}",
1064 request.schema_version, AUTOMATIC_SCHEMA_VERSION
1065 )));
1066 }
1067 if policy.device_memory_fallback_bytes == 0 || policy.host_memory_fallback_bytes == 0 {
1068 return Err(AutomaticPlanningError::Invalid(
1069 "automatic fallback memory budgets must be greater than zero".into(),
1070 ));
1071 }
1072 if policy.memory_headroom_percent >= 100
1073 || policy.expert_cache_share_percent == 0
1074 || policy.expert_cache_share_percent >= 100
1075 || policy.device_layer_window == 0
1076 || policy.max_cached_shards == 0
1077 || policy.embedded_mtp_draft_tokens == 0
1078 || policy.minimum_feedback_tokens == 0
1079 {
1080 return Err(AutomaticPlanningError::Invalid(
1081 "automatic percentage and count policy values are outside their valid ranges".into(),
1082 ));
1083 }
1084 Ok(())
1085}
1086
1087fn selected_device<'a>(
1088 hardware: &'a HardwareProfile,
1089 device: &DevicePlan,
1090) -> Option<&'a HardwareDeviceProfile> {
1091 hardware
1092 .backends
1093 .iter()
1094 .find(|backend| backend.backend == device.backend && backend.available)
1095 .and_then(|backend| backend.devices.iter().find(|item| item.id == device.device))
1096}
1097
1098fn validate_device(
1099 hardware: &HardwareProfile,
1100 device: &DevicePlan,
1101) -> Result<(), AutomaticPlanningError> {
1102 selected_device(hardware, device)
1103 .map(|_| ())
1104 .ok_or_else(|| {
1105 AutomaticPlanningError::Invalid(format!(
1106 "hardware discovery did not report available {} device {}",
1107 device.backend, device.device
1108 ))
1109 })
1110}
1111
1112fn memory_basis(
1113 available: Option<u64>,
1114 physical: Option<u64>,
1115 semantics: HardwareMemorySemantics,
1116) -> Option<u64> {
1117 available.or_else(|| {
1118 (semantics == HardwareMemorySemantics::Unified)
1119 .then_some(physical)
1120 .flatten()
1121 })
1122}
1123
1124fn budget(available: Option<u64>, fallback: u64, headroom_percent: u8) -> u64 {
1125 available
1126 .map(|bytes| bytes.saturating_mul(u64::from(100 - headroom_percent)) / 100)
1127 .unwrap_or(fallback)
1128 .max(1)
1129}
1130
1131fn base_candidates(
1132 device: DevicePlan,
1133 device_budget: u64,
1134 host_budget: u64,
1135 policy: &AutomaticPlannerPolicy,
1136) -> [ExecutionPlan; 3] {
1137 let mut resident = ExecutionPlan::fully_resident(device);
1138 resident.max_cached_shards = policy.max_cached_shards;
1139 let mut layerwise = resident.clone();
1140 layerwise.residency = ResidencyPlan::LayerwiseHost {
1141 device_layer_window: policy.device_layer_window,
1142 device_budget_bytes: Some(device_budget),
1143 host_budget_bytes: Some(host_budget),
1144 };
1145 let mut disk = resident.clone();
1146 disk.residency = ResidencyPlan::DenseDiskStream {
1147 device_budget_bytes: device_budget,
1148 host_budget_bytes: host_budget,
1149 host_lookahead: usize::from(host_budget > 0) * 2,
1150 background_queue: usize::from(host_budget > 0) * 2,
1151 };
1152 [resident, layerwise, disk]
1153}
1154
1155fn apply_bounded_probe<B: AutomaticPlanningBackend>(
1156 backend: &B,
1157 path: &std::path::Path,
1158 plan: &ExecutionPlan,
1159 budget: u64,
1160 admission: &mut CandidateAdmission,
1161 resources: &mut ModelResourceProfile,
1162 adjacent: bool,
1163) -> Result<(), AutomaticPlanningError> {
1164 if !admission.supported {
1165 return Ok(());
1166 }
1167 let requirement = backend.bounded_residency_requirement(path, plan)?;
1168 if requirement.required_bytes > budget {
1169 admission.supported = false;
1170 admission.rejection = Some(format!(
1171 "device budget {budget} bytes cannot contain {} pinned static bytes plus the depth-{} device window ({} bytes, {} total)",
1172 requirement.static_bytes,
1173 requirement.depth,
1174 requirement.window_bytes,
1175 requirement.required_bytes
1176 ));
1177 }
1178 resources.pinned_parameter_bytes =
1179 Observed::exact(requirement.static_bytes, "validated backend parameter plan");
1180 if adjacent {
1181 resources.largest_adjacent_execution_groups_bytes =
1182 Observed::exact(requirement.window_bytes, "validated backend parameter plan");
1183 } else {
1184 resources.largest_execution_group_bytes =
1185 Observed::exact(requirement.window_bytes, "validated backend parameter plan");
1186 }
1187 Ok(())
1188}
1189
1190fn rejection(admission: &CandidateAdmission) -> &str {
1191 admission.rejection.as_deref().unwrap_or("not admitted")
1192}
1193
1194fn with_expert_cache(mut plan: ExecutionPlan, policy: &AutomaticPlannerPolicy) -> ExecutionPlan {
1195 let split = |bytes: u64, percent: u8| bytes.saturating_mul(u64::from(percent)) / 100;
1196 let ordinary_share = 100 - policy.expert_cache_share_percent;
1197 let (device_budget, host_budget) = match &mut plan.residency {
1198 ResidencyPlan::FullyResident => (
1199 policy.device_memory_fallback_bytes,
1200 policy.host_memory_fallback_bytes,
1201 ),
1202 ResidencyPlan::LayerwiseHost {
1203 device_budget_bytes,
1204 host_budget_bytes,
1205 ..
1206 } => {
1207 let device = device_budget_bytes.unwrap_or(policy.device_memory_fallback_bytes);
1208 let host = host_budget_bytes.unwrap_or(policy.host_memory_fallback_bytes);
1209 *device_budget_bytes = Some(split(device, ordinary_share).max(1));
1210 *host_budget_bytes = Some(split(host, ordinary_share).max(1));
1211 (device, host)
1212 }
1213 ResidencyPlan::DenseDiskStream {
1214 device_budget_bytes,
1215 host_budget_bytes,
1216 ..
1217 } => {
1218 let (device, host) = (*device_budget_bytes, *host_budget_bytes);
1219 *device_budget_bytes = split(device, ordinary_share).max(1);
1220 *host_budget_bytes = split(host, ordinary_share).max(1);
1221 (device, host)
1222 }
1223 };
1224 let scratch = (1_u64 << 30).min(device_budget.max(1));
1225 plan.expert_cache = Some(ExpertCachePlan {
1226 device_budget_bytes: Some(split(device_budget, policy.expert_cache_share_percent).max(1)),
1227 host_budget_bytes: Some(split(host_budget, policy.expert_cache_share_percent).max(1)),
1228 scratch_bytes: scratch,
1229 prefill_bank_bytes: scratch,
1230 eviction_policy: crate::residency::CacheEvictionPolicy::LeastRecentlyUsed,
1231 });
1232 plan
1233}
1234
1235fn select_feedback_plan<B: AutomaticPlanningBackend>(
1236 backend: &B,
1237 request: &AutomaticPlanRequest,
1238 hardware: &HardwareProfile,
1239 resources: &ModelResourceProfile,
1240 policy: &AutomaticPlannerPolicy,
1241 embedded_layers: Option<usize>,
1242) -> Result<Option<(ExecutionPlan, usize, f64)>, AutomaticPlanningError> {
1243 let mut groups: Vec<(ExecutionPlan, Vec<f64>)> = Vec::new();
1244 for telemetry in &request.prior_telemetry {
1245 let (Some(plan), Some(prior_hardware), Some(prior_resources)) = (
1246 telemetry.plan.as_ref(),
1247 telemetry.hardware.as_ref(),
1248 telemetry.resources.as_ref(),
1249 ) else {
1250 continue;
1251 };
1252 if telemetry.schema_version != AUTOMATIC_SCHEMA_VERSION
1253 || telemetry.generated_tokens < policy.minimum_feedback_tokens
1254 || plan.device != request.device
1255 || prior_resources.path != resources.path
1256 || prior_resources.artifact_format != resources.artifact_format
1257 || prior_resources.model_family != resources.model_family
1258 || prior_hardware.operating_system != hardware.operating_system
1259 || prior_hardware.architecture != hardware.architecture
1260 || (matches!(plan.drafting, DraftingPlan::Embedded { .. })
1261 && embedded_layers == Some(0))
1262 {
1263 continue;
1264 }
1265 let rate = telemetry
1266 .timing
1267 .decode_token_rate
1268 .filter(|value| value.is_finite() && *value > 0.0)
1269 .or_else(|| {
1270 (telemetry.timing.token_rate.is_finite() && telemetry.timing.token_rate > 0.0)
1271 .then_some(telemetry.timing.token_rate)
1272 });
1273 let Some(rate) = rate else { continue };
1274 if let Some((_, rates)) = groups.iter_mut().find(|(candidate, _)| candidate == plan) {
1275 rates.push(rate);
1276 } else {
1277 groups.push((plan.clone(), vec![rate]));
1278 }
1279 }
1280 let mut accepted = Vec::new();
1281 for (plan, mut rates) in groups {
1282 if !backend
1283 .admit_candidate(&request.model_path, &plan)?
1284 .supported
1285 {
1286 continue;
1287 }
1288 rates.sort_by(f64::total_cmp);
1289 let middle = rates.len() / 2;
1290 let median = if rates.len() % 2 == 0 {
1291 (rates[middle - 1] + rates[middle]) / 2.0
1292 } else {
1293 rates[middle]
1294 };
1295 accepted.push((plan, rates.len(), median));
1296 }
1297 Ok(accepted
1298 .into_iter()
1299 .max_by(|left, right| left.2.total_cmp(&right.2)))
1300}
1301
1302#[cfg(test)]
1303mod tests {
1304 use super::*;
1305 use crate::execution::BackendId;
1306
1307 struct MockPlanningBackend {
1308 model_bytes: u64,
1309 embedded_layers: usize,
1310 }
1311
1312 impl Default for MockPlanningBackend {
1313 fn default() -> Self {
1314 Self {
1315 model_bytes: 2 << 30,
1316 embedded_layers: 0,
1317 }
1318 }
1319 }
1320
1321 impl AutomaticPlanningBackend for MockPlanningBackend {
1322 fn backend_id(&self) -> BackendId {
1323 BackendId::new("mock").unwrap()
1324 }
1325
1326 fn discover_hardware(&self) -> Result<HardwareProfile, AutomaticPlanningError> {
1327 Ok(HardwareProfile {
1328 schema_version: AUTOMATIC_SCHEMA_VERSION,
1329 operating_system: "test".into(),
1330 architecture: "mock".into(),
1331 logical_cpu_count: Observed::exact(8, "fixture"),
1332 physical_memory_bytes: Observed::exact(32 << 30, "fixture"),
1333 available_memory_bytes: Observed::exact(24 << 30, "fixture"),
1334 physical_memory_semantics: HardwareMemorySemantics::SeparateTiers,
1335 backends: vec![HardwareBackendProfile {
1336 backend: BackendId::new("mock").unwrap(),
1337 available: true,
1338 detail: None,
1339 devices: vec![HardwareDeviceProfile {
1340 id: "gpu:0".into(),
1341 family: "gpu".into(),
1342 index: 0,
1343 total_memory_bytes: Observed::exact(16 << 30, "fixture"),
1344 available_memory_bytes: Observed::exact(12 << 30, "fixture"),
1345 }],
1346 }],
1347 })
1348 }
1349
1350 fn inspect_resources(
1351 &self,
1352 path: &std::path::Path,
1353 ) -> Result<ModelResourceProfile, AutomaticPlanningError> {
1354 let mut profile =
1355 ModelResourceProfile::unmeasured(path.into(), ArtifactFormat::SafeTensors);
1356 profile.model_family = Some("llama".into());
1357 profile.embedded_draft_layers =
1358 Observed::exact(self.embedded_layers, "normalized architecture fixture");
1359 profile.stored_tensor_bytes = Observed::exact(self.model_bytes, "fixture");
1360 profile.materialized_parameter_bytes = Observed::exact(self.model_bytes, "fixture");
1361 Ok(profile)
1362 }
1363
1364 fn admit_candidate(
1365 &self,
1366 _path: &std::path::Path,
1367 _plan: &ExecutionPlan,
1368 ) -> Result<CandidateAdmission, AutomaticPlanningError> {
1369 Ok(CandidateAdmission {
1370 supported: true,
1371 rejection: None,
1372 })
1373 }
1374
1375 fn bounded_residency_requirement(
1376 &self,
1377 _path: &std::path::Path,
1378 _plan: &ExecutionPlan,
1379 ) -> Result<BoundedResidencyRequirement, AutomaticPlanningError> {
1380 Ok(BoundedResidencyRequirement {
1381 static_bytes: 1 << 20,
1382 window_bytes: 2 << 20,
1383 required_bytes: 3 << 20,
1384 depth: 1,
1385 })
1386 }
1387 }
1388
1389 #[test]
1390 fn neutral_planner_selects_a_mock_backend_session_plan() {
1391 let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
1392 let report = AutomaticPlanner::default()
1393 .plan(&MockPlanningBackend::default(), &request)
1394 .unwrap();
1395 assert_eq!(report.plan.device.backend.as_str(), "mock");
1396 assert_eq!(report.plan.residency, ResidencyPlan::FullyResident);
1397 }
1398
1399 #[test]
1400 fn neutral_planner_selects_bounded_residency_and_embedded_drafting() {
1401 let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
1402 let report = AutomaticPlanner::default()
1403 .plan(
1404 &MockPlanningBackend {
1405 model_bytes: 10 << 30,
1406 embedded_layers: 2,
1407 },
1408 &request,
1409 )
1410 .unwrap();
1411 assert!(matches!(
1412 report.plan.residency,
1413 ResidencyPlan::LayerwiseHost { .. }
1414 ));
1415 assert!(matches!(
1416 report.plan.drafting,
1417 DraftingPlan::Embedded { .. }
1418 ));
1419 assert_eq!(
1420 observed_u64(&report.resources.pinned_parameter_bytes),
1421 Some(1 << 20)
1422 );
1423 }
1424
1425 #[test]
1426 fn selected_backend_identity_fails_closed() {
1427 let request =
1428 AutomaticPlanRequest::new("model", DevicePlan::new("other", "gpu:0").unwrap());
1429 assert!(matches!(
1430 AutomaticPlanner::default().plan(&MockPlanningBackend::default(), &request),
1431 Err(AutomaticPlanningError::Invalid(message))
1432 if message.contains("cannot plan device")
1433 ));
1434 }
1435
1436 #[test]
1437 fn documents_round_trip_without_an_accelerator_runtime() {
1438 let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
1439 let encoded = serde_json::to_vec(&request).unwrap();
1440 assert_eq!(
1441 serde_json::from_slice::<AutomaticPlanRequest>(&encoded).unwrap(),
1442 request
1443 );
1444 let unavailable = serde_json::to_value(Observed::<u64>::unavailable("unknown")).unwrap();
1445 assert!(unavailable.get("value").is_none());
1446 }
1447
1448 #[test]
1449 fn tokenizer_compatibility_requires_identical_vocabularies() {
1450 let fingerprint = [7; 32];
1451 let proof = TokenizerCompatibilityProof::prove(fingerprint, fingerprint).unwrap();
1452 assert_eq!(proof.fingerprint(), fingerprint);
1453 assert_eq!(proof.validate_target(fingerprint), Ok(()));
1454 assert_eq!(
1455 proof.validate_target([8; 32]),
1456 Err(TokenizerCompatibilityError)
1457 );
1458 assert_eq!(
1459 TokenizerCompatibilityProof::prove(fingerprint, [8; 32]),
1460 Err(TokenizerCompatibilityError)
1461 );
1462 }
1463
1464 #[test]
1465 fn zero_duration_rates_are_finite() {
1466 let timing = TimingTelemetry::new(
1467 Duration::ZERO,
1468 Duration::ZERO,
1469 Some(Duration::ZERO),
1470 3,
1471 Duration::ZERO,
1472 );
1473 assert_eq!(timing.token_rate, 0.0);
1474 assert_eq!(timing.decode_token_rate, Some(0.0));
1475 }
1476}