use crate::{
artifact::ArtifactFormat,
backend::{BackendProvider, ModelLoadingBackend, ModelRuntime},
execution::{
DevicePlan, DraftingPlan, ExecutionPlan, ExpertCachePlan, ResidencyPlan,
DEFAULT_MAX_CACHED_SHARDS,
},
speculative::SpeculativeDraft,
};
use serde::{Deserialize, Serialize};
use std::{path::PathBuf, time::Duration};
pub const AUTOMATIC_SCHEMA_VERSION: u32 = 6;
#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ObservationKind {
Exact,
Conservative,
Observational,
Estimated,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
#[serde(tag = "status", rename_all = "snake_case")]
pub enum Observed<T> {
Available {
value: T,
kind: ObservationKind,
source: String,
},
Unsupported {
reason: String,
},
Unavailable {
reason: String,
},
}
impl<T> Observed<T> {
pub fn exact(value: T, source: impl Into<String>) -> Self {
Self::Available {
value,
kind: ObservationKind::Exact,
source: source.into(),
}
}
pub fn unavailable(reason: impl Into<String>) -> Self {
Self::Unavailable {
reason: reason.into(),
}
}
pub fn unsupported(reason: impl Into<String>) -> Self {
Self::Unsupported {
reason: reason.into(),
}
}
pub const fn value(&self) -> Option<&T> {
match self {
Self::Available { value, .. } => Some(value),
Self::Unsupported { .. } | Self::Unavailable { .. } => None,
}
}
}
fn unobserved_embedded_draft_layers() -> Observed<usize> {
Observed::unavailable("embedded drafting requires normalized architecture inspection")
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct ModelResourceProfile {
pub schema_version: u32,
pub path: PathBuf,
pub artifact_format: ArtifactFormat,
#[serde(skip_serializing_if = "Option::is_none")]
pub model_family: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub architecture: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tensor_count: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub checkpoint_shards: Option<usize>,
#[serde(default = "unobserved_embedded_draft_layers")]
pub embedded_draft_layers: Observed<usize>,
pub stored_tensor_bytes: Observed<u64>,
pub largest_stored_tensor_bytes: Observed<u64>,
pub materialized_parameter_bytes: Observed<u64>,
pub pinned_parameter_bytes: Observed<u64>,
pub largest_execution_group_bytes: Observed<u64>,
pub largest_adjacent_execution_groups_bytes: Observed<u64>,
pub expert_parameter_bytes: Observed<u64>,
}
impl ModelResourceProfile {
pub fn unmeasured(path: PathBuf, artifact_format: ArtifactFormat) -> Self {
let unavailable = || {
Observed::unavailable("resource value requires a validated checkpoint parameter plan")
};
Self {
schema_version: AUTOMATIC_SCHEMA_VERSION,
path,
artifact_format,
model_family: None,
architecture: None,
tensor_count: None,
checkpoint_shards: None,
embedded_draft_layers: unobserved_embedded_draft_layers(),
stored_tensor_bytes: Observed::unavailable(
"checkpoint tensor catalog was not established",
),
largest_stored_tensor_bytes: Observed::unavailable(
"checkpoint tensor catalog was not established",
),
materialized_parameter_bytes: unavailable(),
pinned_parameter_bytes: unavailable(),
largest_execution_group_bytes: unavailable(),
largest_adjacent_execution_groups_bytes: unavailable(),
expert_parameter_bytes: unavailable(),
}
}
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct HardwareDeviceProfile {
pub id: String,
pub family: String,
pub index: usize,
pub total_memory_bytes: Observed<u64>,
pub available_memory_bytes: Observed<u64>,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct HardwareBackendProfile {
pub backend: crate::execution::BackendId,
pub available: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub detail: Option<String>,
pub devices: Vec<HardwareDeviceProfile>,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct HardwareProfile {
pub schema_version: u32,
pub operating_system: String,
pub architecture: String,
pub logical_cpu_count: Observed<u64>,
pub physical_memory_bytes: Observed<u64>,
pub available_memory_bytes: Observed<u64>,
pub physical_memory_semantics: HardwareMemorySemantics,
pub backends: Vec<HardwareBackendProfile>,
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HardwareMemorySemantics {
Unified,
SeparateTiers,
Unknown,
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PlanExplanationLevel {
Decision,
Warning,
Rejection,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct PlanExplanationEntry {
pub level: PlanExplanationLevel,
pub code: String,
pub detail: String,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct PlanExplanation {
pub summary: String,
pub entries: Vec<PlanExplanationEntry>,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct ExecutionPlanReport {
pub schema_version: u32,
pub hardware: HardwareProfile,
pub resources: ModelResourceProfile,
pub plan: ExecutionPlan,
pub explanation: PlanExplanation,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
#[serde(default)]
#[non_exhaustive]
pub struct AutomaticPlannerPolicy {
pub device_memory_fallback_bytes: u64,
pub host_memory_fallback_bytes: u64,
pub memory_headroom_percent: u8,
pub expert_cache_share_percent: u8,
pub device_layer_window: usize,
pub max_cached_shards: usize,
pub embedded_mtp_draft_tokens: usize,
pub minimum_feedback_tokens: usize,
}
impl Default for AutomaticPlannerPolicy {
fn default() -> Self {
Self {
device_memory_fallback_bytes: 4 << 30,
host_memory_fallback_bytes: 16 << 30,
memory_headroom_percent: 30,
expert_cache_share_percent: 40,
device_layer_window: 1,
max_cached_shards: DEFAULT_MAX_CACHED_SHARDS,
embedded_mtp_draft_tokens: 3,
minimum_feedback_tokens: 1,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TimingTelemetry {
pub load_seconds: f64,
pub generation_seconds: f64,
#[serde(skip_serializing_if = "Option::is_none")]
pub time_to_first_token_seconds: Option<f64>,
pub total_seconds: f64,
pub token_rate: f64,
#[serde(skip_serializing_if = "Option::is_none")]
pub decode_token_rate: Option<f64>,
}
impl TimingTelemetry {
pub fn new(
load: Duration,
generation: Duration,
time_to_first_token: Option<Duration>,
generated_tokens: usize,
total: Duration,
) -> Self {
fn rate(tokens: usize, elapsed: Duration) -> f64 {
if elapsed.is_zero() {
0.0
} else {
tokens as f64 / elapsed.as_secs_f64()
}
}
Self {
load_seconds: load.as_secs_f64(),
generation_seconds: generation.as_secs_f64(),
time_to_first_token_seconds: time_to_first_token.map(|value| value.as_secs_f64()),
total_seconds: total.as_secs_f64(),
token_rate: rate(generated_tokens, generation),
decode_token_rate: time_to_first_token.map(|first| {
rate(
generated_tokens.saturating_sub(1),
generation.saturating_sub(first),
)
}),
}
}
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct AllocatorTelemetry {
pub peak_bytes: u64,
pub active_bytes: u64,
pub cache_bytes: u64,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct ResidencyTelemetry {
pub planned_disk_bytes: u64,
pub planned_host_bytes: u64,
pub planned_device_bytes: u64,
pub current_host_bytes: u64,
pub current_device_bytes: u64,
pub peak_host_bytes: u64,
pub peak_device_bytes: u64,
pub transfers: Vec<TransferTelemetry>,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct TransferTelemetry {
pub direction: String,
pub count: u64,
pub bytes: u64,
pub seconds: DurationSeconds,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(transparent)]
pub struct DurationSeconds(pub f64);
impl PartialEq for DurationSeconds {
fn eq(&self, other: &Self) -> bool {
self.0.to_bits() == other.0.to_bits()
}
}
impl Eq for DurationSeconds {}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct ExpertCacheTelemetry {
pub owned_experts: usize,
pub owned_bytes: u64,
pub host_resident_experts: usize,
pub device_resident_experts: usize,
pub host_resident_bytes: u64,
pub device_resident_bytes: u64,
pub peak_host_resident_bytes: u64,
pub peak_device_resident_bytes: u64,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SpeculativeDecodingTelemetry {
pub execution_topology: String,
pub target_tokens: usize,
pub draft_tokens: usize,
pub accepted_tokens: usize,
pub accept_rate: f64,
pub rounds: usize,
pub accept_lens: Vec<usize>,
pub emitted_tokens: usize,
pub optimistic_draft_tokens: usize,
pub reused_optimistic_tokens: usize,
pub discarded_optimistic_tokens: usize,
pub adaptive_lookahead_disabled: bool,
pub optimistic_draft_seconds: f64,
pub verification_in_flight_seconds: f64,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ExecutionTelemetry {
pub schema_version: u32,
pub effective_model_type: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub plan: Option<ExecutionPlan>,
#[serde(skip_serializing_if = "Option::is_none")]
pub plan_explanation: Option<PlanExplanation>,
#[serde(skip_serializing_if = "Option::is_none")]
pub hardware: Option<HardwareProfile>,
#[serde(skip_serializing_if = "Option::is_none")]
pub resources: Option<ModelResourceProfile>,
pub prompt_tokens: usize,
pub generated_tokens: usize,
pub stop_reason: String,
pub timing: TimingTelemetry,
#[serde(skip_serializing_if = "Option::is_none")]
pub allocator: Option<AllocatorTelemetry>,
#[serde(skip_serializing_if = "Option::is_none")]
pub residency: Option<ResidencyTelemetry>,
#[serde(skip_serializing_if = "Option::is_none")]
pub expert_cache: Option<ExpertCacheTelemetry>,
#[serde(skip_serializing_if = "Option::is_none")]
pub speculative: Option<SpeculativeDecodingTelemetry>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct AutomaticPlanRequest {
pub schema_version: u32,
pub model_path: PathBuf,
pub device: DevicePlan,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub prior_telemetry: Vec<ExecutionTelemetry>,
}
impl AutomaticPlanRequest {
pub fn new(model_path: impl Into<PathBuf>, device: DevicePlan) -> Self {
Self {
schema_version: AUTOMATIC_SCHEMA_VERSION,
model_path: model_path.into(),
device,
prior_telemetry: Vec::new(),
}
}
pub fn with_prior_telemetry(
mut self,
telemetry: impl IntoIterator<Item = ExecutionTelemetry>,
) -> Self {
self.prior_telemetry.extend(telemetry);
self
}
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct CandidateAdmission {
pub supported: bool,
pub rejection: Option<String>,
}
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct BoundedResidencyRequirement {
pub static_bytes: u64,
pub window_bytes: u64,
pub required_bytes: u64,
pub depth: usize,
}
pub trait AutomaticPlanningBackend {
fn backend_id(&self) -> crate::execution::BackendId;
fn discover_hardware(&self) -> Result<HardwareProfile, AutomaticPlanningError>;
fn inspect_resources(
&self,
model_path: &std::path::Path,
) -> Result<ModelResourceProfile, AutomaticPlanningError>;
fn admit_candidate(
&self,
model_path: &std::path::Path,
plan: &ExecutionPlan,
) -> Result<CandidateAdmission, AutomaticPlanningError>;
fn bounded_residency_requirement(
&self,
model_path: &std::path::Path,
plan: &ExecutionPlan,
) -> Result<BoundedResidencyRequirement, AutomaticPlanningError>;
}
pub struct ExecutionPlanTarget<B: ModelLoadingBackend> {
backend: B,
load_options: B::LoadOptions,
}
impl<B: ModelLoadingBackend> ExecutionPlanTarget<B> {
pub fn new(backend: B, load_options: B::LoadOptions) -> Self {
Self {
backend,
load_options,
}
}
pub const fn backend(&self) -> &B {
&self.backend
}
pub fn into_parts(self) -> (B, B::LoadOptions) {
(self.backend, self.load_options)
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct TokenizerCompatibilityProof {
fingerprint: [u8; 32],
}
impl TokenizerCompatibilityProof {
pub fn prove(
target_fingerprint: [u8; 32],
assistant_fingerprint: [u8; 32],
) -> Result<Self, TokenizerCompatibilityError> {
if target_fingerprint != assistant_fingerprint {
return Err(TokenizerCompatibilityError);
}
Ok(Self {
fingerprint: target_fingerprint,
})
}
pub const fn fingerprint(self) -> [u8; 32] {
self.fingerprint
}
pub fn validate_target(
self,
target_fingerprint: [u8; 32],
) -> Result<(), TokenizerCompatibilityError> {
if self.fingerprint != target_fingerprint {
return Err(TokenizerCompatibilityError);
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, thiserror::Error)]
#[error("assistant token-id vocabulary mapping does not match the target")]
pub struct TokenizerCompatibilityError;
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ExternalDraftArtifact<P> {
pub preparation: P,
pub tokenizer_compatibility: TokenizerCompatibilityProof,
}
pub enum RealizedDrafting<D> {
Disabled,
Embedded,
External(D),
}
impl<D> RealizedDrafting<D> {
pub fn as_speculative_draft(&mut self) -> Option<SpeculativeDraft<'_, D>> {
match self {
Self::Disabled => None,
Self::Embedded => Some(SpeculativeDraft::Embedded),
Self::External(drafter) => Some(SpeculativeDraft::External(drafter)),
}
}
pub const fn is_external(&self) -> bool {
matches!(self, Self::External(_))
}
}
pub trait ExecutionPlanBackendFactory: AutomaticPlanningBackend {
type Backend: ModelLoadingBackend;
type DrafterPreparation;
type Drafter;
fn realize_target(
&self,
plan: &ExecutionPlan,
) -> Result<ExecutionPlanTarget<Self::Backend>, AutomaticPlanningError>;
fn realize_drafting(
&self,
plan: &ExecutionPlan,
target: &ModelRuntime<Self::Backend>,
external_artifact: Option<ExternalDraftArtifact<Self::DrafterPreparation>>,
) -> Result<RealizedDrafting<Self::Drafter>, AutomaticPlanningError>;
}
pub fn realize_execution_plan_target<F: ExecutionPlanBackendFactory>(
factory: &F,
plan: &ExecutionPlan,
) -> Result<ExecutionPlanTarget<F::Backend>, AutomaticPlanningError> {
let expected_backend = factory.backend_id();
if plan.device.backend != expected_backend {
return Err(AutomaticPlanningError::Invalid(format!(
"execution plan selects backend {} but factory owns {}",
plan.device.backend, expected_backend
)));
}
plan.validate_structure()
.map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
let realization = factory.realize_target(plan)?;
let descriptor = realization.backend().descriptor();
if descriptor.name() != expected_backend.as_str() {
return Err(AutomaticPlanningError::Invalid(format!(
"factory identity {} does not match realized backend {}",
expected_backend,
descriptor.name()
)));
}
let devices =
realization
.backend()
.devices()
.map_err(|error| AutomaticPlanningError::Backend {
operation: "realize_execution_plan_devices",
message: error.to_string(),
})?;
let capabilities = devices
.iter()
.find_map(|(device, capabilities)| {
(device.id() == plan.device.device).then_some(capabilities)
})
.ok_or_else(|| {
AutomaticPlanningError::Invalid(format!(
"realized backend {} does not expose selected device {}",
expected_backend, plan.device.device
))
})?;
plan.validate_device_capabilities(capabilities)
.map_err(|error| AutomaticPlanningError::Invalid(error.to_string()))?;
Ok(realization)
}
pub fn realize_execution_plan_drafting<F: ExecutionPlanBackendFactory>(
factory: &F,
plan: &ExecutionPlan,
target: &ModelRuntime<F::Backend>,
external_artifact: Option<ExternalDraftArtifact<F::DrafterPreparation>>,
) -> Result<RealizedDrafting<F::Drafter>, AutomaticPlanningError> {
match (&plan.drafting, external_artifact.as_ref()) {
(DraftingPlan::External { .. }, None) => {
return Err(AutomaticPlanningError::Invalid(
"external drafting requires proven tokenizer compatibility".into(),
));
}
(DraftingPlan::Disabled | DraftingPlan::Embedded { .. }, Some(_)) => {
return Err(AutomaticPlanningError::Invalid(
"tokenizer compatibility was supplied for a plan without an external assistant"
.into(),
));
}
_ => {}
}
let drafting = factory.realize_drafting(plan, target, external_artifact)?;
let matches_plan = matches!(
(&plan.drafting, &drafting),
(DraftingPlan::Disabled, RealizedDrafting::Disabled)
| (DraftingPlan::Embedded { .. }, RealizedDrafting::Embedded)
| (DraftingPlan::External { .. }, RealizedDrafting::External(_))
);
if !matches_plan {
return Err(AutomaticPlanningError::Invalid(
"backend factory realized a drafting mode different from the execution plan".into(),
));
}
Ok(drafting)
}
#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
pub enum AutomaticPlanningError {
#[error("automatic planning error: {0}")]
Invalid(String),
#[error("automatic planning backend failed during {operation}: {message}")]
Backend {
operation: &'static str,
message: String,
},
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Default)]
pub struct AutomaticPlanner {
policy: AutomaticPlannerPolicy,
}
impl AutomaticPlanner {
pub fn new(policy: AutomaticPlannerPolicy) -> Self {
Self { policy }
}
pub fn policy(&self) -> &AutomaticPlannerPolicy {
&self.policy
}
pub fn plan<B: AutomaticPlanningBackend>(
&self,
backend: &B,
request: &AutomaticPlanRequest,
) -> Result<ExecutionPlanReport, AutomaticPlanningError> {
validate_request(request, &self.policy)?;
let backend_id = backend.backend_id();
if request.device.backend != backend_id {
return Err(AutomaticPlanningError::Invalid(format!(
"selected planning backend {} cannot plan device owned by {}",
backend_id, request.device.backend
)));
}
let hardware = backend.discover_hardware()?;
validate_device(&hardware, &request.device)?;
let mut resources = backend.inspect_resources(&request.model_path)?;
let selected_device =
selected_device(&hardware, &request.device).expect("validated device is present");
let device_capacity = memory_basis(
observed_u64(&selected_device.available_memory_bytes),
observed_u64(&selected_device.total_memory_bytes)
.or_else(|| observed_u64(&hardware.physical_memory_bytes)),
hardware.physical_memory_semantics,
);
let host_capacity = memory_basis(
observed_u64(&hardware.available_memory_bytes),
observed_u64(&hardware.physical_memory_bytes),
hardware.physical_memory_semantics,
);
let device_budget = budget(
device_capacity,
self.policy.device_memory_fallback_bytes,
self.policy.memory_headroom_percent,
);
let host_budget = budget(
host_capacity,
self.policy.host_memory_fallback_bytes,
self.policy.memory_headroom_percent,
);
let model_bytes = observed_u64(&resources.materialized_parameter_bytes)
.or_else(|| observed_u64(&resources.stored_tensor_bytes));
let candidates = base_candidates(
request.device.clone(),
device_budget,
host_budget,
&self.policy,
);
let resident = backend.admit_candidate(&request.model_path, &candidates[0])?;
let mut layerwise = backend.admit_candidate(&request.model_path, &candidates[1])?;
let mut disk = backend.admit_candidate(&request.model_path, &candidates[2])?;
let resident_fits = model_bytes.is_some_and(|bytes| bytes <= device_budget);
let layerwise_host_fits = model_bytes.is_some_and(|bytes| {
if hardware.physical_memory_semantics == HardwareMemorySemantics::Unified {
bytes <= host_budget.saturating_mul(2)
} else {
bytes <= host_budget
}
});
if !resident_fits || !resident.supported {
apply_bounded_probe(
backend,
&request.model_path,
&candidates[1],
device_budget,
&mut layerwise,
&mut resources,
false,
)?;
apply_bounded_probe(
backend,
&request.model_path,
&candidates[2],
device_budget,
&mut disk,
&mut resources,
true,
)?;
}
let selected =
if resident_fits && resident.supported {
0
} else if layerwise_host_fits && layerwise.supported {
1
} else if disk.supported {
2
} else {
return Err(AutomaticPlanningError::Invalid(format!(
"no loadable single-device policy: resident: {}; layerwise: {}; disk-streamed: {}",
rejection(&resident), rejection(&layerwise), rejection(&disk)
)));
};
let mut plan = candidates[selected].clone();
let mut entries = vec![PlanExplanationEntry {
level: PlanExplanationLevel::Decision,
code: "single_device_scope".into(),
detail: format!(
"automatic planning is restricted to {}:{} with {}% memory headroom",
request.device.backend, request.device.device, self.policy.memory_headroom_percent
),
}];
if selected > 0 {
entries.push(PlanExplanationEntry {
level: PlanExplanationLevel::Rejection,
code: "fully_resident_not_admitted".into(),
detail: resident
.rejection
.unwrap_or_else(|| "the model exceeds the device memory budget".into()),
});
}
if selected > 1 {
entries.push(PlanExplanationEntry {
level: PlanExplanationLevel::Rejection,
code: "layerwise_not_admitted".into(),
detail: layerwise
.rejection
.unwrap_or_else(|| "the model exceeds the host-backed admission budget".into()),
});
}
let mut summary = match selected {
0 => "selected fully resident execution for the lowest expected latency".to_string(),
1 => "selected host-backed layerwise execution with a validated bounded device window"
.to_string(),
_ => "selected bounded dense disk streaming because resident and layerwise admission failed"
.to_string(),
};
if selected > 0 {
let expert_plan = with_expert_cache(plan.clone(), &self.policy);
let expert = backend.admit_candidate(&request.model_path, &expert_plan)?;
if expert.supported {
plan = expert_plan;
entries.push(PlanExplanationEntry {
level: PlanExplanationLevel::Decision,
code: "expert_cache_selected".into(),
detail: "the backend admitted independent routed-expert caching".into(),
});
}
}
let embedded_layers = resources.embedded_draft_layers.value().copied();
if embedded_layers.is_some_and(|layers| layers > 0) {
plan.drafting = DraftingPlan::Embedded {
max_draft_tokens: self.policy.embedded_mtp_draft_tokens,
lookahead: true,
adaptive_lookahead: true,
};
entries.push(PlanExplanationEntry {
level: PlanExplanationLevel::Decision,
code: "embedded_mtp_selected".into(),
detail: "checkpoint metadata advertises embedded prediction layers".into(),
});
}
if let Some((feedback, samples, median)) = select_feedback_plan(
backend,
request,
&hardware,
&resources,
&self.policy,
embedded_layers,
)? {
plan = feedback;
summary = format!(
"selected a previously observed plan at {median:.2} median decode tokens/s"
);
entries.push(PlanExplanationEntry {
level: PlanExplanationLevel::Decision,
code: "prior_telemetry_selected".into(),
detail: format!("selected using {samples} matching runtime sample(s)"),
});
}
Ok(ExecutionPlanReport {
schema_version: AUTOMATIC_SCHEMA_VERSION,
hardware,
resources,
plan,
explanation: PlanExplanation { summary, entries },
})
}
}
fn observed_u64(value: &Observed<u64>) -> Option<u64> {
value.value().copied()
}
fn validate_request(
request: &AutomaticPlanRequest,
policy: &AutomaticPlannerPolicy,
) -> Result<(), AutomaticPlanningError> {
if request.schema_version != AUTOMATIC_SCHEMA_VERSION {
return Err(AutomaticPlanningError::Invalid(format!(
"automatic request schema {} does not match supported schema {}",
request.schema_version, AUTOMATIC_SCHEMA_VERSION
)));
}
if policy.device_memory_fallback_bytes == 0 || policy.host_memory_fallback_bytes == 0 {
return Err(AutomaticPlanningError::Invalid(
"automatic fallback memory budgets must be greater than zero".into(),
));
}
if policy.memory_headroom_percent >= 100
|| policy.expert_cache_share_percent == 0
|| policy.expert_cache_share_percent >= 100
|| policy.device_layer_window == 0
|| policy.max_cached_shards == 0
|| policy.embedded_mtp_draft_tokens == 0
|| policy.minimum_feedback_tokens == 0
{
return Err(AutomaticPlanningError::Invalid(
"automatic percentage and count policy values are outside their valid ranges".into(),
));
}
Ok(())
}
fn selected_device<'a>(
hardware: &'a HardwareProfile,
device: &DevicePlan,
) -> Option<&'a HardwareDeviceProfile> {
hardware
.backends
.iter()
.find(|backend| backend.backend == device.backend && backend.available)
.and_then(|backend| backend.devices.iter().find(|item| item.id == device.device))
}
fn validate_device(
hardware: &HardwareProfile,
device: &DevicePlan,
) -> Result<(), AutomaticPlanningError> {
selected_device(hardware, device)
.map(|_| ())
.ok_or_else(|| {
AutomaticPlanningError::Invalid(format!(
"hardware discovery did not report available {} device {}",
device.backend, device.device
))
})
}
fn memory_basis(
available: Option<u64>,
physical: Option<u64>,
semantics: HardwareMemorySemantics,
) -> Option<u64> {
available.or_else(|| {
(semantics == HardwareMemorySemantics::Unified)
.then_some(physical)
.flatten()
})
}
fn budget(available: Option<u64>, fallback: u64, headroom_percent: u8) -> u64 {
available
.map(|bytes| bytes.saturating_mul(u64::from(100 - headroom_percent)) / 100)
.unwrap_or(fallback)
.max(1)
}
fn base_candidates(
device: DevicePlan,
device_budget: u64,
host_budget: u64,
policy: &AutomaticPlannerPolicy,
) -> [ExecutionPlan; 3] {
let mut resident = ExecutionPlan::fully_resident(device);
resident.max_cached_shards = policy.max_cached_shards;
let mut layerwise = resident.clone();
layerwise.residency = ResidencyPlan::LayerwiseHost {
device_layer_window: policy.device_layer_window,
device_budget_bytes: Some(device_budget),
host_budget_bytes: Some(host_budget),
};
let mut disk = resident.clone();
disk.residency = ResidencyPlan::DenseDiskStream {
device_budget_bytes: device_budget,
host_budget_bytes: host_budget,
host_lookahead: usize::from(host_budget > 0) * 2,
background_queue: usize::from(host_budget > 0) * 2,
};
[resident, layerwise, disk]
}
fn apply_bounded_probe<B: AutomaticPlanningBackend>(
backend: &B,
path: &std::path::Path,
plan: &ExecutionPlan,
budget: u64,
admission: &mut CandidateAdmission,
resources: &mut ModelResourceProfile,
adjacent: bool,
) -> Result<(), AutomaticPlanningError> {
if !admission.supported {
return Ok(());
}
let requirement = backend.bounded_residency_requirement(path, plan)?;
if requirement.required_bytes > budget {
admission.supported = false;
admission.rejection = Some(format!(
"device budget {budget} bytes cannot contain {} pinned static bytes plus the depth-{} device window ({} bytes, {} total)",
requirement.static_bytes,
requirement.depth,
requirement.window_bytes,
requirement.required_bytes
));
}
resources.pinned_parameter_bytes =
Observed::exact(requirement.static_bytes, "validated backend parameter plan");
if adjacent {
resources.largest_adjacent_execution_groups_bytes =
Observed::exact(requirement.window_bytes, "validated backend parameter plan");
} else {
resources.largest_execution_group_bytes =
Observed::exact(requirement.window_bytes, "validated backend parameter plan");
}
Ok(())
}
fn rejection(admission: &CandidateAdmission) -> &str {
admission.rejection.as_deref().unwrap_or("not admitted")
}
fn with_expert_cache(mut plan: ExecutionPlan, policy: &AutomaticPlannerPolicy) -> ExecutionPlan {
let split = |bytes: u64, percent: u8| bytes.saturating_mul(u64::from(percent)) / 100;
let ordinary_share = 100 - policy.expert_cache_share_percent;
let (device_budget, host_budget) = match &mut plan.residency {
ResidencyPlan::FullyResident => (
policy.device_memory_fallback_bytes,
policy.host_memory_fallback_bytes,
),
ResidencyPlan::LayerwiseHost {
device_budget_bytes,
host_budget_bytes,
..
} => {
let device = device_budget_bytes.unwrap_or(policy.device_memory_fallback_bytes);
let host = host_budget_bytes.unwrap_or(policy.host_memory_fallback_bytes);
*device_budget_bytes = Some(split(device, ordinary_share).max(1));
*host_budget_bytes = Some(split(host, ordinary_share).max(1));
(device, host)
}
ResidencyPlan::DenseDiskStream {
device_budget_bytes,
host_budget_bytes,
..
} => {
let (device, host) = (*device_budget_bytes, *host_budget_bytes);
*device_budget_bytes = split(device, ordinary_share).max(1);
*host_budget_bytes = split(host, ordinary_share).max(1);
(device, host)
}
};
let scratch = (1_u64 << 30).min(device_budget.max(1));
plan.expert_cache = Some(ExpertCachePlan {
device_budget_bytes: Some(split(device_budget, policy.expert_cache_share_percent).max(1)),
host_budget_bytes: Some(split(host_budget, policy.expert_cache_share_percent).max(1)),
scratch_bytes: scratch,
prefill_bank_bytes: scratch,
eviction_policy: crate::residency::CacheEvictionPolicy::LeastRecentlyUsed,
});
plan
}
fn select_feedback_plan<B: AutomaticPlanningBackend>(
backend: &B,
request: &AutomaticPlanRequest,
hardware: &HardwareProfile,
resources: &ModelResourceProfile,
policy: &AutomaticPlannerPolicy,
embedded_layers: Option<usize>,
) -> Result<Option<(ExecutionPlan, usize, f64)>, AutomaticPlanningError> {
let mut groups: Vec<(ExecutionPlan, Vec<f64>)> = Vec::new();
for telemetry in &request.prior_telemetry {
let (Some(plan), Some(prior_hardware), Some(prior_resources)) = (
telemetry.plan.as_ref(),
telemetry.hardware.as_ref(),
telemetry.resources.as_ref(),
) else {
continue;
};
if telemetry.schema_version != AUTOMATIC_SCHEMA_VERSION
|| telemetry.generated_tokens < policy.minimum_feedback_tokens
|| plan.device != request.device
|| prior_resources.path != resources.path
|| prior_resources.artifact_format != resources.artifact_format
|| prior_resources.model_family != resources.model_family
|| prior_hardware.operating_system != hardware.operating_system
|| prior_hardware.architecture != hardware.architecture
|| (matches!(plan.drafting, DraftingPlan::Embedded { .. })
&& embedded_layers == Some(0))
{
continue;
}
let rate = telemetry
.timing
.decode_token_rate
.filter(|value| value.is_finite() && *value > 0.0)
.or_else(|| {
(telemetry.timing.token_rate.is_finite() && telemetry.timing.token_rate > 0.0)
.then_some(telemetry.timing.token_rate)
});
let Some(rate) = rate else { continue };
if let Some((_, rates)) = groups.iter_mut().find(|(candidate, _)| candidate == plan) {
rates.push(rate);
} else {
groups.push((plan.clone(), vec![rate]));
}
}
let mut accepted = Vec::new();
for (plan, mut rates) in groups {
if !backend
.admit_candidate(&request.model_path, &plan)?
.supported
{
continue;
}
rates.sort_by(f64::total_cmp);
let middle = rates.len() / 2;
let median = if rates.len() % 2 == 0 {
(rates[middle - 1] + rates[middle]) / 2.0
} else {
rates[middle]
};
accepted.push((plan, rates.len(), median));
}
Ok(accepted
.into_iter()
.max_by(|left, right| left.2.total_cmp(&right.2)))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::execution::BackendId;
struct MockPlanningBackend {
model_bytes: u64,
embedded_layers: usize,
}
impl Default for MockPlanningBackend {
fn default() -> Self {
Self {
model_bytes: 2 << 30,
embedded_layers: 0,
}
}
}
impl AutomaticPlanningBackend for MockPlanningBackend {
fn backend_id(&self) -> BackendId {
BackendId::new("mock").unwrap()
}
fn discover_hardware(&self) -> Result<HardwareProfile, AutomaticPlanningError> {
Ok(HardwareProfile {
schema_version: AUTOMATIC_SCHEMA_VERSION,
operating_system: "test".into(),
architecture: "mock".into(),
logical_cpu_count: Observed::exact(8, "fixture"),
physical_memory_bytes: Observed::exact(32 << 30, "fixture"),
available_memory_bytes: Observed::exact(24 << 30, "fixture"),
physical_memory_semantics: HardwareMemorySemantics::SeparateTiers,
backends: vec![HardwareBackendProfile {
backend: BackendId::new("mock").unwrap(),
available: true,
detail: None,
devices: vec![HardwareDeviceProfile {
id: "gpu:0".into(),
family: "gpu".into(),
index: 0,
total_memory_bytes: Observed::exact(16 << 30, "fixture"),
available_memory_bytes: Observed::exact(12 << 30, "fixture"),
}],
}],
})
}
fn inspect_resources(
&self,
path: &std::path::Path,
) -> Result<ModelResourceProfile, AutomaticPlanningError> {
let mut profile =
ModelResourceProfile::unmeasured(path.into(), ArtifactFormat::SafeTensors);
profile.model_family = Some("llama".into());
profile.embedded_draft_layers =
Observed::exact(self.embedded_layers, "normalized architecture fixture");
profile.stored_tensor_bytes = Observed::exact(self.model_bytes, "fixture");
profile.materialized_parameter_bytes = Observed::exact(self.model_bytes, "fixture");
Ok(profile)
}
fn admit_candidate(
&self,
_path: &std::path::Path,
_plan: &ExecutionPlan,
) -> Result<CandidateAdmission, AutomaticPlanningError> {
Ok(CandidateAdmission {
supported: true,
rejection: None,
})
}
fn bounded_residency_requirement(
&self,
_path: &std::path::Path,
_plan: &ExecutionPlan,
) -> Result<BoundedResidencyRequirement, AutomaticPlanningError> {
Ok(BoundedResidencyRequirement {
static_bytes: 1 << 20,
window_bytes: 2 << 20,
required_bytes: 3 << 20,
depth: 1,
})
}
}
#[test]
fn neutral_planner_selects_a_mock_backend_session_plan() {
let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
let report = AutomaticPlanner::default()
.plan(&MockPlanningBackend::default(), &request)
.unwrap();
assert_eq!(report.plan.device.backend.as_str(), "mock");
assert_eq!(report.plan.residency, ResidencyPlan::FullyResident);
}
#[test]
fn neutral_planner_selects_bounded_residency_and_embedded_drafting() {
let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
let report = AutomaticPlanner::default()
.plan(
&MockPlanningBackend {
model_bytes: 10 << 30,
embedded_layers: 2,
},
&request,
)
.unwrap();
assert!(matches!(
report.plan.residency,
ResidencyPlan::LayerwiseHost { .. }
));
assert!(matches!(
report.plan.drafting,
DraftingPlan::Embedded { .. }
));
assert_eq!(
observed_u64(&report.resources.pinned_parameter_bytes),
Some(1 << 20)
);
}
#[test]
fn selected_backend_identity_fails_closed() {
let request =
AutomaticPlanRequest::new("model", DevicePlan::new("other", "gpu:0").unwrap());
assert!(matches!(
AutomaticPlanner::default().plan(&MockPlanningBackend::default(), &request),
Err(AutomaticPlanningError::Invalid(message))
if message.contains("cannot plan device")
));
}
#[test]
fn documents_round_trip_without_an_accelerator_runtime() {
let request = AutomaticPlanRequest::new("model", DevicePlan::new("mock", "gpu:0").unwrap());
let encoded = serde_json::to_vec(&request).unwrap();
assert_eq!(
serde_json::from_slice::<AutomaticPlanRequest>(&encoded).unwrap(),
request
);
let unavailable = serde_json::to_value(Observed::<u64>::unavailable("unknown")).unwrap();
assert!(unavailable.get("value").is_none());
}
#[test]
fn tokenizer_compatibility_requires_identical_vocabularies() {
let fingerprint = [7; 32];
let proof = TokenizerCompatibilityProof::prove(fingerprint, fingerprint).unwrap();
assert_eq!(proof.fingerprint(), fingerprint);
assert_eq!(proof.validate_target(fingerprint), Ok(()));
assert_eq!(
proof.validate_target([8; 32]),
Err(TokenizerCompatibilityError)
);
assert_eq!(
TokenizerCompatibilityProof::prove(fingerprint, [8; 32]),
Err(TokenizerCompatibilityError)
);
}
#[test]
fn zero_duration_rates_are_finite() {
let timing = TimingTelemetry::new(
Duration::ZERO,
Duration::ZERO,
Some(Duration::ZERO),
3,
Duration::ZERO,
);
assert_eq!(timing.token_rate, 0.0);
assert_eq!(timing.decode_token_rate, Some(0.0));
}
}