eredu_core/
execution.rs

1//! Portable execution-plan and telemetry schemas.
2
3use crate::{
4    backend::{DeviceCapabilities, SessionCapabilities},
5    residency::CacheEvictionPolicy,
6    topology::ParallelTopology,
7};
8use serde::{Deserialize, Serialize};
9
10/// Schema version shared by execution-plan documents.
11pub const EXECUTION_PLAN_SCHEMA_VERSION: u32 = 4;
12
13/// Default bound for simultaneously open checkpoint payload sources.
14pub const DEFAULT_MAX_CACHED_SHARDS: usize = 4;
15
16/// Stable, extensible identity of an execution backend.
17///
18/// Core deliberately does not enumerate implementations. Values such as
19/// Concrete implementations are registered by their backend adapters.
20#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
21#[serde(transparent)]
22pub struct BackendId(String);
23
24impl BackendId {
25    /// Creates a non-empty backend identifier.
26    pub fn new(value: impl Into<String>) -> Result<Self, ExecutionPlanError> {
27        let value = value.into();
28        if value.trim().is_empty() {
29            return Err(ExecutionPlanError::EmptyBackendId);
30        }
31        Ok(Self(value))
32    }
33
34    /// Returns the stable adapter-defined identifier.
35    pub fn as_str(&self) -> &str {
36        &self.0
37    }
38}
39
40impl std::fmt::Display for BackendId {
41    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        formatter.write_str(&self.0)
43    }
44}
45
46impl<'de> Deserialize<'de> for BackendId {
47    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
48    where
49        D: serde::Deserializer<'de>,
50    {
51        Self::new(String::deserialize(deserializer)?).map_err(serde::de::Error::custom)
52    }
53}
54
55/// Backend and device selected for one complete model session.
56#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
57pub struct DevicePlan {
58    /// Backend adapter identity.
59    pub(crate) backend: BackendId,
60    /// Backend-stable device identifier.
61    pub(crate) device: String,
62}
63
64impl<'de> Deserialize<'de> for DevicePlan {
65    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
66    where
67        D: serde::Deserializer<'de>,
68    {
69        #[derive(Deserialize)]
70        struct RawDevicePlan {
71            backend: BackendId,
72            device: String,
73        }
74
75        let raw = RawDevicePlan::deserialize(deserializer)?;
76        Self::new(raw.backend.0, raw.device).map_err(serde::de::Error::custom)
77    }
78}
79
80impl DevicePlan {
81    /// Creates a validated backend/device selection.
82    pub fn new(
83        backend: impl Into<String>,
84        device: impl Into<String>,
85    ) -> Result<Self, ExecutionPlanError> {
86        let device = device.into();
87        if device.trim().is_empty() {
88            return Err(ExecutionPlanError::EmptyDeviceId);
89        }
90        Ok(Self {
91            backend: BackendId::new(backend)?,
92            device,
93        })
94    }
95
96    /// Selected backend identity.
97    pub const fn backend(&self) -> &BackendId {
98        &self.backend
99    }
100    /// Backend-stable device identifier.
101    pub fn device(&self) -> &str {
102        &self.device
103    }
104}
105
106/// Static weight placement selected by an execution plan.
107#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
108#[serde(tag = "mode", rename_all = "snake_case")]
109#[non_exhaustive]
110pub enum ResidencyPlan {
111    /// Retain all selected weights on the execution device.
112    FullyResident,
113    /// Retain repeated groups on the host and promote a bounded device window.
114    LayerwiseHost {
115        /// Maximum repeated groups resident on the device.
116        device_layer_window: usize,
117        /// Logical device parameter budget.
118        #[serde(skip_serializing_if = "Option::is_none")]
119        device_budget_bytes: Option<u64>,
120        /// Charged host-transfer budget.
121        #[serde(skip_serializing_if = "Option::is_none")]
122        host_budget_bytes: Option<u64>,
123    },
124    /// Stream repeated groups through disk, host, and device caches.
125    DenseDiskStream {
126        /// Finite logical device budget.
127        device_budget_bytes: u64,
128        /// Finite charged host budget.
129        host_budget_bytes: u64,
130        /// Protected host lookahead.
131        host_lookahead: usize,
132        /// Background materialization queue capacity.
133        background_queue: usize,
134    },
135}
136
137/// Optional independent routed-expert cache selected by a plan.
138#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
139pub struct ExpertCachePlan {
140    /// Logical device expert-cache budget.
141    #[serde(skip_serializing_if = "Option::is_none")]
142    pub(crate) device_budget_bytes: Option<u64>,
143    /// Charged host expert-cache budget.
144    #[serde(skip_serializing_if = "Option::is_none")]
145    pub(crate) host_budget_bytes: Option<u64>,
146    /// Hard compact-bank scratch bound.
147    pub(crate) scratch_bytes: u64,
148    /// Soft prefill compact-bank target.
149    pub(crate) prefill_bank_bytes: u64,
150    /// Deterministic eviction ordering for independently cached experts.
151    pub(crate) eviction_policy: CacheEvictionPolicy,
152}
153
154impl ExpertCachePlan {
155    /// Creates one independent routed-parameter cache plan.
156    pub const fn new(
157        device_budget_bytes: Option<u64>,
158        host_budget_bytes: Option<u64>,
159        scratch_bytes: u64,
160        prefill_bank_bytes: u64,
161        eviction_policy: CacheEvictionPolicy,
162    ) -> Self {
163        Self {
164            device_budget_bytes,
165            host_budget_bytes,
166            scratch_bytes,
167            prefill_bank_bytes,
168            eviction_policy,
169        }
170    }
171    /// Logical device cache budget.
172    pub const fn device_budget_bytes(&self) -> Option<u64> {
173        self.device_budget_bytes
174    }
175    /// Charged host cache budget.
176    pub const fn host_budget_bytes(&self) -> Option<u64> {
177        self.host_budget_bytes
178    }
179    /// Hard compact-bank scratch bound.
180    pub const fn scratch_bytes(&self) -> u64 {
181        self.scratch_bytes
182    }
183    /// Soft prefill compact-bank target.
184    pub const fn prefill_bank_bytes(&self) -> u64 {
185        self.prefill_bank_bytes
186    }
187    /// Deterministic eviction ordering.
188    pub const fn eviction_policy(&self) -> CacheEvictionPolicy {
189        self.eviction_policy
190    }
191}
192
193/// Speculative decoding selected by an execution plan.
194#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
195#[serde(tag = "mode", rename_all = "snake_case")]
196#[non_exhaustive]
197pub enum DraftingPlan {
198    /// Ordinary target-only decoding.
199    Disabled,
200    /// Use checkpoint-embedded prediction heads.
201    Embedded {
202        /// Maximum proposals per verification round.
203        max_draft_tokens: usize,
204        /// Whether same-request optimistic lookahead is enabled.
205        lookahead: bool,
206        /// Whether deterministic adaptive lookahead is enabled.
207        adaptive_lookahead: bool,
208    },
209    /// Use an explicitly supplied external assistant.
210    External {
211        /// Assistant artifact path or identifier.
212        model: String,
213        /// Backend/device placement used for assistant execution.
214        placement: DraftPlacementPlan,
215        /// Maximum proposals per verification round.
216        max_draft_tokens: usize,
217        /// Whether same-request optimistic lookahead is enabled.
218        lookahead: bool,
219        /// Whether deterministic adaptive lookahead is enabled.
220        adaptive_lookahead: bool,
221    },
222}
223
224/// External assistant placement selected by an execution plan.
225#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
226#[serde(tag = "mode", rename_all = "snake_case")]
227#[non_exhaustive]
228pub enum DraftPlacementPlan {
229    /// Reuse the target execution context.
230    Target,
231    /// Use an explicit backend/device selection.
232    Device {
233        /// Explicit process-local assistant device.
234        device: DevicePlan,
235    },
236}
237
238/// Optional load-time transformation applied to checkpoint weights.
239#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
240#[serde(tag = "mode", rename_all = "snake_case")]
241#[non_exhaustive]
242pub enum WeightTransformationPlan {
243    /// Preserve checkpoint-native weight encodings.
244    PreserveCheckpoint,
245    /// Convert eligible weights to grouped affine quantization while loading.
246    Affine {
247        /// Quantized bits per weight.
248        bits: i32,
249        /// Adjacent weights sharing quantization parameters.
250        group_size: i32,
251    },
252    /// Convert eligible weights to MXFP4 while loading.
253    MxFp4,
254}
255
256/// A concrete, backend-neutral set of model/session execution choices.
257#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
258pub struct ExecutionPlan {
259    /// Version of this serialized plan schema.
260    pub(crate) schema_version: u32,
261    /// Backend and process-local device selected for the whole session.
262    pub(crate) device: DevicePlan,
263    /// Distributed Cartesian topology.
264    pub(crate) topology: ParallelTopology,
265    /// Ordinary static-weight placement.
266    pub(crate) residency: ResidencyPlan,
267    /// Optional transformation applied while checkpoint weights are loaded.
268    pub(crate) weight_transformation: WeightTransformationPlan,
269    /// Maximum number of checkpoint shards or readers retained simultaneously.
270    pub(crate) max_cached_shards: usize,
271    /// Independent routed-expert cache, when enabled.
272    #[serde(skip_serializing_if = "Option::is_none")]
273    pub(crate) expert_cache: Option<ExpertCachePlan>,
274    /// Speculative decoding configuration.
275    pub(crate) drafting: DraftingPlan,
276    /// Capabilities which the selected device must provide.
277    pub(crate) required_device_capabilities: DeviceCapabilities,
278    /// Capabilities which the exact prepared session must provide.
279    pub(crate) required_session_capabilities: SessionCapabilities,
280}
281
282impl ExecutionPlan {
283    /// Creates the minimal fully-resident, target-only plan for one device.
284    pub fn fully_resident(device: DevicePlan) -> Self {
285        Self {
286            schema_version: EXECUTION_PLAN_SCHEMA_VERSION,
287            device,
288            topology: ParallelTopology::new(1, 1, 1, 1).expect("the singleton topology is valid"),
289            residency: ResidencyPlan::FullyResident,
290            weight_transformation: WeightTransformationPlan::PreserveCheckpoint,
291            max_cached_shards: DEFAULT_MAX_CACHED_SHARDS,
292            expert_cache: None,
293            drafting: DraftingPlan::Disabled,
294            required_device_capabilities: DeviceCapabilities::new(true, false, false),
295            required_session_capabilities: SessionCapabilities::default(),
296        }
297    }
298
299    /// Serialized schema version.
300    pub const fn schema_version(&self) -> u32 {
301        self.schema_version
302    }
303    /// Selected backend and device.
304    pub const fn device(&self) -> &DevicePlan {
305        &self.device
306    }
307    /// Distributed topology.
308    pub const fn topology(&self) -> &ParallelTopology {
309        &self.topology
310    }
311    /// Static weight placement.
312    pub const fn residency(&self) -> &ResidencyPlan {
313        &self.residency
314    }
315    /// Load-time weight transformation.
316    pub const fn weight_transformation(&self) -> WeightTransformationPlan {
317        self.weight_transformation
318    }
319    /// Maximum simultaneously retained checkpoint readers.
320    pub const fn max_cached_shards(&self) -> usize {
321        self.max_cached_shards
322    }
323    /// Optional independent routed-parameter cache.
324    pub const fn expert_cache(&self) -> Option<&ExpertCachePlan> {
325        self.expert_cache.as_ref()
326    }
327    /// Speculative drafting selection.
328    pub const fn drafting(&self) -> &DraftingPlan {
329        &self.drafting
330    }
331    /// Required device capabilities.
332    pub const fn required_device_capabilities(&self) -> &DeviceCapabilities {
333        &self.required_device_capabilities
334    }
335    /// Required prepared-session capabilities.
336    pub const fn required_session_capabilities(&self) -> &SessionCapabilities {
337        &self.required_session_capabilities
338    }
339
340    /// Replaces the distributed topology.
341    pub fn with_topology(mut self, topology: ParallelTopology) -> Self {
342        self.topology = topology;
343        self
344    }
345    /// Replaces the selected backend and device.
346    pub fn with_device(mut self, device: DevicePlan) -> Self {
347        self.device = device;
348        self
349    }
350    /// Replaces static weight placement.
351    pub fn with_residency(mut self, residency: ResidencyPlan) -> Self {
352        self.residency = residency;
353        self
354    }
355    /// Replaces load-time weight transformation.
356    pub fn with_weight_transformation(mut self, transformation: WeightTransformationPlan) -> Self {
357        self.weight_transformation = transformation;
358        self
359    }
360    /// Replaces the checkpoint reader bound.
361    pub fn with_max_cached_shards(mut self, maximum: usize) -> Self {
362        self.max_cached_shards = maximum;
363        self
364    }
365    /// Replaces the independent routed-parameter cache plan.
366    pub fn with_expert_cache(mut self, expert_cache: Option<ExpertCachePlan>) -> Self {
367        self.expert_cache = expert_cache;
368        self
369    }
370    /// Replaces speculative drafting selection.
371    pub fn with_drafting(mut self, drafting: DraftingPlan) -> Self {
372        self.drafting = drafting;
373        self
374    }
375    /// Replaces required device capabilities.
376    pub fn with_required_device_capabilities(mut self, capabilities: DeviceCapabilities) -> Self {
377        self.required_device_capabilities = capabilities;
378        self
379    }
380    /// Replaces required prepared-session capabilities.
381    pub fn with_required_session_capabilities(mut self, capabilities: SessionCapabilities) -> Self {
382        self.required_session_capabilities = capabilities;
383        self
384    }
385
386    /// Validates portable plan invariants and fail-closed capabilities.
387    pub fn validate_device_capabilities(
388        &self,
389        available: &DeviceCapabilities,
390    ) -> Result<(), ExecutionPlanError> {
391        self.validate_structure()?;
392        for (required, supported, name) in [
393            (
394                self.required_device_capabilities.exact_completion(),
395                available.exact_completion(),
396                "exact_completion",
397            ),
398            (
399                self.required_device_capabilities.transfers(),
400                available.transfers(),
401                "transfers",
402            ),
403            (
404                self.required_device_capabilities.collectives(),
405                available.collectives(),
406                "collectives",
407            ),
408        ] {
409            if required && !supported {
410                return Err(ExecutionPlanError::Capability(name));
411            }
412        }
413        Ok(())
414    }
415
416    /// Validates exact prepared-session requirements independently.
417    pub fn validate_session_capabilities(
418        &self,
419        available: &SessionCapabilities,
420    ) -> Result<(), ExecutionPlanError> {
421        self.validate_structure()?;
422        self.required_session_capabilities
423            .validate(available)
424            .map_err(|error| ExecutionPlanError::Capability(error.capability()))
425    }
426
427    /// Validates schema, topology, and portable resource invariants without a backend.
428    pub fn validate_structure(&self) -> Result<(), ExecutionPlanError> {
429        if self.schema_version != EXECUTION_PLAN_SCHEMA_VERSION {
430            return Err(ExecutionPlanError::Schema(self.schema_version));
431        }
432        if self.max_cached_shards == 0 {
433            return Err(ExecutionPlanError::ZeroMappedShards);
434        }
435        match &self.drafting {
436            DraftingPlan::Disabled => {}
437            DraftingPlan::Embedded {
438                max_draft_tokens, ..
439            } => {
440                if *max_draft_tokens == 0 {
441                    return Err(ExecutionPlanError::ZeroDraftTokens);
442                }
443            }
444            DraftingPlan::External {
445                model,
446                max_draft_tokens,
447                ..
448            } => {
449                if model.trim().is_empty() {
450                    return Err(ExecutionPlanError::EmptyDraftModel);
451                }
452                if *max_draft_tokens == 0 {
453                    return Err(ExecutionPlanError::ZeroDraftTokens);
454                }
455            }
456        }
457        ParallelTopology::new(
458            self.topology.tensor(),
459            self.topology.pipeline(),
460            self.topology.expert(),
461            self.topology.data(),
462        )
463        .map_err(|error| ExecutionPlanError::Topology(error.to_string()))?;
464        Ok(())
465    }
466}
467
468/// Execution plan validation error.
469#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
470#[non_exhaustive]
471pub enum ExecutionPlanError {
472    /// Backend identifier is empty.
473    #[error("execution-plan backend identifier must not be empty")]
474    EmptyBackendId,
475    /// Device identifier is empty.
476    #[error("execution-plan device identifier must not be empty")]
477    EmptyDeviceId,
478    /// Unsupported schema version.
479    #[error("unsupported execution-plan schema version {0}")]
480    Schema(u32),
481    /// Required capability is absent.
482    #[error("execution plan requires unavailable capability {0}")]
483    Capability(&'static str),
484    /// Topology is invalid.
485    #[error("execution-plan topology is invalid: {0}")]
486    Topology(String),
487    /// The checkpoint source bound is zero.
488    #[error("execution-plan max_cached_shards must be greater than zero")]
489    ZeroMappedShards,
490    /// An external assistant artifact path or identifier is empty.
491    #[error("execution-plan external draft model must not be empty")]
492    EmptyDraftModel,
493    /// Speculative execution has no proposal capacity.
494    #[error("execution-plan max_draft_tokens must be greater than zero")]
495    ZeroDraftTokens,
496}
497
498#[cfg(test)]
499mod tests {
500    use super::*;
501
502    #[test]
503    fn plan_round_trips_with_extensible_backend_identity() {
504        let plan = ExecutionPlan::fully_resident(DevicePlan::new("iree", "vulkan:2").unwrap());
505        let encoded = serde_json::to_vec(&plan).unwrap();
506        assert_eq!(
507            serde_json::from_slice::<serde_json::Value>(&encoded).unwrap()["schema_version"],
508            4
509        );
510        assert_eq!(
511            serde_json::from_slice::<ExecutionPlan>(&encoded).unwrap(),
512            plan
513        );
514    }
515
516    #[test]
517    fn plan_capabilities_fail_closed() {
518        let mut plan = ExecutionPlan::fully_resident(DevicePlan::new("mlx", "metal:0").unwrap());
519        assert_eq!(
520            plan.validate_device_capabilities(&DeviceCapabilities::default()),
521            Err(ExecutionPlanError::Capability("exact_completion"))
522        );
523
524        plan.required_session_capabilities = plan
525            .required_session_capabilities
526            .with_activation_inspection(true);
527        assert_eq!(
528            plan.validate_session_capabilities(&SessionCapabilities::default()),
529            Err(ExecutionPlanError::Capability("activation_inspection"))
530        );
531        assert!(plan
532            .validate_device_capabilities(&DeviceCapabilities::new(true, false, false))
533            .is_ok());
534        assert!(plan
535            .validate_session_capabilities(
536                &SessionCapabilities::default().with_activation_inspection(true),
537            )
538            .is_ok());
539    }
540
541    #[test]
542    fn backend_and_device_identifiers_fail_closed_during_deserialization() {
543        assert!(serde_json::from_str::<DevicePlan>(r#"{"backend":"","device":"cpu:0"}"#).is_err());
544        assert!(serde_json::from_str::<DevicePlan>(r#"{"backend":"mlx","device":""}"#).is_err());
545    }
546
547    #[test]
548    fn speculative_plan_structure_fails_closed() {
549        let mut plan = ExecutionPlan::fully_resident(DevicePlan::new("mock", "gpu:0").unwrap());
550        plan.drafting = DraftingPlan::Embedded {
551            max_draft_tokens: 0,
552            lookahead: false,
553            adaptive_lookahead: false,
554        };
555        assert_eq!(
556            plan.validate_structure(),
557            Err(ExecutionPlanError::ZeroDraftTokens)
558        );
559
560        plan.drafting = DraftingPlan::External {
561            model: "  ".into(),
562            placement: DraftPlacementPlan::Target,
563            max_draft_tokens: 1,
564            lookahead: false,
565            adaptive_lookahead: false,
566        };
567        assert_eq!(
568            plan.validate_structure(),
569            Err(ExecutionPlanError::EmptyDraftModel)
570        );
571    }
572}
eredu_core/execution.rs

eredu_core/
execution.rs
