Skip to main content

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
224impl DraftingPlan {
225    /// Returns the maximum proposal count selected for each verification round.
226    pub const fn max_draft_tokens(&self) -> Option<usize> {
227        match self {
228            Self::Disabled => None,
229            Self::Embedded {
230                max_draft_tokens, ..
231            }
232            | Self::External {
233                max_draft_tokens, ..
234            } => Some(*max_draft_tokens),
235        }
236    }
237
238    /// Returns whether the selected plan enables same-request lookahead.
239    pub const fn lookahead(&self) -> bool {
240        match self {
241            Self::Disabled => false,
242            Self::Embedded { lookahead, .. } | Self::External { lookahead, .. } => *lookahead,
243        }
244    }
245
246    /// Returns whether deterministic adaptive lookahead is selected.
247    pub const fn adaptive_lookahead(&self) -> bool {
248        match self {
249            Self::Disabled => false,
250            Self::Embedded {
251                adaptive_lookahead, ..
252            }
253            | Self::External {
254                adaptive_lookahead, ..
255            } => *adaptive_lookahead,
256        }
257    }
258}
259
260/// External assistant placement selected by an execution plan.
261#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
262#[serde(tag = "mode", rename_all = "snake_case")]
263#[non_exhaustive]
264pub enum DraftPlacementPlan {
265    /// Reuse the target execution context.
266    Target,
267    /// Use an explicit backend/device selection.
268    Device {
269        /// Explicit process-local assistant device.
270        device: DevicePlan,
271    },
272}
273
274impl DraftPlacementPlan {
275    /// Resolves target/draft topology from the portable plan before queue construction.
276    pub fn execution_topology(&self, target: &DevicePlan) -> crate::SpeculativeExecutionTopology {
277        match self {
278            Self::Target => crate::SpeculativeExecutionTopology::Single,
279            Self::Device { device } if device == target => {
280                crate::SpeculativeExecutionTopology::SameDeviceSplit
281            }
282            Self::Device { .. } => crate::SpeculativeExecutionTopology::CrossDeviceSplit,
283        }
284    }
285}
286
287/// Optional load-time transformation applied to checkpoint weights.
288#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
289#[serde(tag = "mode", rename_all = "snake_case")]
290#[non_exhaustive]
291pub enum WeightTransformationPlan {
292    /// Preserve checkpoint-native weight encodings.
293    PreserveCheckpoint,
294    /// Convert eligible weights to grouped affine quantization while loading.
295    Affine {
296        /// Quantized bits per weight.
297        bits: i32,
298        /// Adjacent weights sharing quantization parameters.
299        group_size: i32,
300    },
301    /// Convert eligible weights to MXFP4 while loading.
302    MxFp4,
303}
304
305/// A concrete, backend-neutral set of model/session execution choices.
306#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
307pub struct ExecutionPlan {
308    /// Version of this serialized plan schema.
309    pub(crate) schema_version: u32,
310    /// Backend and process-local device selected for the whole session.
311    pub(crate) device: DevicePlan,
312    /// Distributed Cartesian topology.
313    pub(crate) topology: ParallelTopology,
314    /// Ordinary static-weight placement.
315    pub(crate) residency: ResidencyPlan,
316    /// Optional transformation applied while checkpoint weights are loaded.
317    pub(crate) weight_transformation: WeightTransformationPlan,
318    /// Maximum number of checkpoint shards or readers retained simultaneously.
319    pub(crate) max_cached_shards: usize,
320    /// Independent routed-expert cache, when enabled.
321    #[serde(skip_serializing_if = "Option::is_none")]
322    pub(crate) expert_cache: Option<ExpertCachePlan>,
323    /// Speculative decoding configuration.
324    pub(crate) drafting: DraftingPlan,
325    /// Capabilities which the selected device must provide.
326    pub(crate) required_device_capabilities: DeviceCapabilities,
327    /// Capabilities which the exact prepared session must provide.
328    pub(crate) required_session_capabilities: SessionCapabilities,
329    /// Explicitly requires persisted prompt-prefix import/export independently of cache residency.
330    #[serde(default, skip_serializing_if = "is_false")]
331    pub(crate) prompt_cache_persistence: bool,
332}
333
334fn is_false(value: &bool) -> bool {
335    !*value
336}
337
338impl ExecutionPlan {
339    /// Creates the minimal fully-resident, target-only plan for one device.
340    pub fn fully_resident(device: DevicePlan) -> Self {
341        Self {
342            schema_version: EXECUTION_PLAN_SCHEMA_VERSION,
343            device,
344            topology: ParallelTopology::new(1, 1, 1, 1).expect("the singleton topology is valid"),
345            residency: ResidencyPlan::FullyResident,
346            weight_transformation: WeightTransformationPlan::PreserveCheckpoint,
347            max_cached_shards: DEFAULT_MAX_CACHED_SHARDS,
348            expert_cache: None,
349            drafting: DraftingPlan::Disabled,
350            required_device_capabilities: DeviceCapabilities::new(true, false, false),
351            required_session_capabilities: SessionCapabilities::default(),
352            prompt_cache_persistence: false,
353        }
354    }
355
356    /// Serialized schema version.
357    pub const fn schema_version(&self) -> u32 {
358        self.schema_version
359    }
360    /// Selected backend and device.
361    pub const fn device(&self) -> &DevicePlan {
362        &self.device
363    }
364    /// Distributed topology.
365    pub const fn topology(&self) -> &ParallelTopology {
366        &self.topology
367    }
368    /// Static weight placement.
369    pub const fn residency(&self) -> &ResidencyPlan {
370        &self.residency
371    }
372    /// Load-time weight transformation.
373    pub const fn weight_transformation(&self) -> WeightTransformationPlan {
374        self.weight_transformation
375    }
376    /// Maximum simultaneously retained checkpoint readers.
377    pub const fn max_cached_shards(&self) -> usize {
378        self.max_cached_shards
379    }
380
381    /// Whether persisted prompt-prefix import/export is explicitly required.
382    ///
383    /// This is distinct from maintaining decode cache state between submissions.
384    pub const fn prompt_cache_persistence(&self) -> bool {
385        self.prompt_cache_persistence
386    }
387    /// Optional independent routed-parameter cache.
388    pub const fn expert_cache(&self) -> Option<&ExpertCachePlan> {
389        self.expert_cache.as_ref()
390    }
391    /// Speculative drafting selection.
392    pub const fn drafting(&self) -> &DraftingPlan {
393        &self.drafting
394    }
395    /// Required device capabilities.
396    pub const fn required_device_capabilities(&self) -> &DeviceCapabilities {
397        &self.required_device_capabilities
398    }
399    /// Required prepared-session capabilities.
400    pub const fn required_session_capabilities(&self) -> &SessionCapabilities {
401        &self.required_session_capabilities
402    }
403
404    /// Replaces the distributed topology.
405    pub fn with_topology(mut self, topology: ParallelTopology) -> Self {
406        self.topology = topology;
407        self
408    }
409    /// Replaces the selected backend and device.
410    pub fn with_device(mut self, device: DevicePlan) -> Self {
411        self.device = device;
412        self
413    }
414    /// Replaces static weight placement.
415    pub fn with_residency(mut self, residency: ResidencyPlan) -> Self {
416        self.residency = residency;
417        self
418    }
419    /// Replaces load-time weight transformation.
420    pub fn with_weight_transformation(mut self, transformation: WeightTransformationPlan) -> Self {
421        self.weight_transformation = transformation;
422        self
423    }
424    /// Replaces the checkpoint reader bound.
425    pub fn with_max_cached_shards(mut self, maximum: usize) -> Self {
426        self.max_cached_shards = maximum;
427        self
428    }
429
430    /// Requires persisted prompt-prefix import/export without changing cache residency.
431    pub fn with_prompt_cache_persistence(mut self, required: bool) -> Self {
432        self.prompt_cache_persistence = required;
433        self
434    }
435    /// Replaces the independent routed-parameter cache plan.
436    pub fn with_expert_cache(mut self, expert_cache: Option<ExpertCachePlan>) -> Self {
437        self.expert_cache = expert_cache;
438        self
439    }
440    /// Replaces speculative drafting selection.
441    pub fn with_drafting(mut self, drafting: DraftingPlan) -> Self {
442        self.drafting = drafting;
443        self
444    }
445    /// Replaces required device capabilities.
446    pub fn with_required_device_capabilities(mut self, capabilities: DeviceCapabilities) -> Self {
447        self.required_device_capabilities = capabilities;
448        self
449    }
450    /// Replaces required prepared-session capabilities.
451    pub fn with_required_session_capabilities(mut self, capabilities: SessionCapabilities) -> Self {
452        self.required_session_capabilities = capabilities;
453        self
454    }
455
456    /// Validates portable plan invariants and fail-closed capabilities.
457    pub fn validate_device_capabilities(
458        &self,
459        available: &DeviceCapabilities,
460    ) -> Result<(), ExecutionPlanError> {
461        self.validate_structure()?;
462        for (required, supported, name) in [
463            (
464                self.required_device_capabilities.exact_completion(),
465                available.exact_completion(),
466                "exact_completion",
467            ),
468            (
469                self.required_device_capabilities.transfers(),
470                available.transfers(),
471                "transfers",
472            ),
473            (
474                self.required_device_capabilities.collectives(),
475                available.collectives(),
476                "collectives",
477            ),
478        ] {
479            if required && !supported {
480                return Err(ExecutionPlanError::Capability(name));
481            }
482        }
483        Ok(())
484    }
485
486    /// Validates exact prepared-session requirements independently.
487    pub fn validate_session_capabilities(
488        &self,
489        available: &SessionCapabilities,
490    ) -> Result<(), ExecutionPlanError> {
491        self.validate_structure()?;
492        self.required_session_capabilities
493            .validate(available)
494            .map_err(|error| ExecutionPlanError::Capability(error.capability()))
495    }
496
497    /// Validates schema, topology, and portable resource invariants without a backend.
498    pub fn validate_structure(&self) -> Result<(), ExecutionPlanError> {
499        if self.schema_version != EXECUTION_PLAN_SCHEMA_VERSION {
500            return Err(ExecutionPlanError::Schema(self.schema_version));
501        }
502        if self.max_cached_shards == 0 {
503            return Err(ExecutionPlanError::ZeroMappedShards);
504        }
505        match &self.drafting {
506            DraftingPlan::Disabled => {}
507            DraftingPlan::Embedded {
508                max_draft_tokens, ..
509            } => {
510                if *max_draft_tokens == 0 {
511                    return Err(ExecutionPlanError::ZeroDraftTokens);
512                }
513            }
514            DraftingPlan::External {
515                model,
516                max_draft_tokens,
517                ..
518            } => {
519                if model.trim().is_empty() {
520                    return Err(ExecutionPlanError::EmptyDraftModel);
521                }
522                if *max_draft_tokens == 0 {
523                    return Err(ExecutionPlanError::ZeroDraftTokens);
524                }
525            }
526        }
527        ParallelTopology::new(
528            self.topology.tensor(),
529            self.topology.pipeline(),
530            self.topology.expert(),
531            self.topology.data(),
532        )
533        .map_err(|error| ExecutionPlanError::Topology(error.to_string()))?;
534        Ok(())
535    }
536}
537
538/// Execution plan validation error.
539#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
540#[non_exhaustive]
541pub enum ExecutionPlanError {
542    /// Backend identifier is empty.
543    #[error("execution-plan backend identifier must not be empty")]
544    EmptyBackendId,
545    /// Device identifier is empty.
546    #[error("execution-plan device identifier must not be empty")]
547    EmptyDeviceId,
548    /// Unsupported schema version.
549    #[error("unsupported execution-plan schema version {0}")]
550    Schema(u32),
551    /// Required capability is absent.
552    #[error("execution plan requires unavailable capability {0}")]
553    Capability(&'static str),
554    /// Topology is invalid.
555    #[error("execution-plan topology is invalid: {0}")]
556    Topology(String),
557    /// The checkpoint source bound is zero.
558    #[error("execution-plan max_cached_shards must be greater than zero")]
559    ZeroMappedShards,
560    /// An external assistant artifact path or identifier is empty.
561    #[error("execution-plan external draft model must not be empty")]
562    EmptyDraftModel,
563    /// Speculative execution has no proposal capacity.
564    #[error("execution-plan max_draft_tokens must be greater than zero")]
565    ZeroDraftTokens,
566}
567
568#[cfg(test)]
569mod tests {
570    use super::*;
571
572    #[test]
573    fn prompt_cache_persistence_serde_preserves_legacy_default_and_explicit_intent() {
574        let plan = ExecutionPlan::fully_resident(DevicePlan::new("foreign", "cpu:0").unwrap());
575        let legacy = serde_json::to_value(&plan).unwrap();
576        assert!(legacy.get("prompt_cache_persistence").is_none());
577        assert_eq!(
578            serde_json::from_value::<ExecutionPlan>(legacy).unwrap(),
579            plan
580        );
581        let explicit = plan.with_prompt_cache_persistence(true);
582        let encoded = serde_json::to_value(&explicit).unwrap();
583        assert_eq!(encoded["prompt_cache_persistence"], true);
584        assert_eq!(
585            serde_json::from_value::<ExecutionPlan>(encoded).unwrap(),
586            explicit
587        );
588    }
589
590    #[test]
591    fn plan_round_trips_with_extensible_backend_identity() {
592        let plan = ExecutionPlan::fully_resident(DevicePlan::new("iree", "vulkan:2").unwrap());
593        let encoded = serde_json::to_vec(&plan).unwrap();
594        assert_eq!(
595            serde_json::from_slice::<serde_json::Value>(&encoded).unwrap()["schema_version"],
596            4
597        );
598        assert_eq!(
599            serde_json::from_slice::<ExecutionPlan>(&encoded).unwrap(),
600            plan
601        );
602    }
603
604    #[test]
605    fn plan_capabilities_fail_closed() {
606        let mut plan = ExecutionPlan::fully_resident(DevicePlan::new("mlx", "metal:0").unwrap());
607        assert_eq!(
608            plan.validate_device_capabilities(&DeviceCapabilities::default()),
609            Err(ExecutionPlanError::Capability("exact_completion"))
610        );
611
612        plan.required_session_capabilities = plan
613            .required_session_capabilities
614            .with_activation_inspection(true);
615        assert_eq!(
616            plan.validate_session_capabilities(&SessionCapabilities::default()),
617            Err(ExecutionPlanError::Capability("activation_inspection"))
618        );
619        assert!(plan
620            .validate_device_capabilities(&DeviceCapabilities::new(true, false, false))
621            .is_ok());
622        assert!(plan
623            .validate_session_capabilities(
624                &SessionCapabilities::default().with_activation_inspection(true),
625            )
626            .is_ok());
627    }
628
629    #[test]
630    fn backend_and_device_identifiers_fail_closed_during_deserialization() {
631        assert!(serde_json::from_str::<DevicePlan>(r#"{"backend":"","device":"cpu:0"}"#).is_err());
632        assert!(serde_json::from_str::<DevicePlan>(r#"{"backend":"mlx","device":""}"#).is_err());
633    }
634
635    #[test]
636    fn speculative_plan_structure_fails_closed() {
637        let mut plan = ExecutionPlan::fully_resident(DevicePlan::new("mock", "gpu:0").unwrap());
638        plan.drafting = DraftingPlan::Embedded {
639            max_draft_tokens: 0,
640            lookahead: false,
641            adaptive_lookahead: false,
642        };
643        assert_eq!(
644            plan.validate_structure(),
645            Err(ExecutionPlanError::ZeroDraftTokens)
646        );
647
648        plan.drafting = DraftingPlan::External {
649            model: "  ".into(),
650            placement: DraftPlacementPlan::Target,
651            max_draft_tokens: 1,
652            lookahead: false,
653            adaptive_lookahead: false,
654        };
655        assert_eq!(
656            plan.validate_structure(),
657            Err(ExecutionPlanError::EmptyDraftModel)
658        );
659    }
660
661    #[test]
662    fn draft_topology_is_selected_from_the_portable_plan_before_queue_construction() {
663        let target = DevicePlan::new("mlx", "metal:0").unwrap();
664        assert_eq!(
665            DraftPlacementPlan::Target.execution_topology(&target),
666            crate::SpeculativeExecutionTopology::Single
667        );
668        assert_eq!(
669            DraftPlacementPlan::Device {
670                device: target.clone(),
671            }
672            .execution_topology(&target),
673            crate::SpeculativeExecutionTopology::SameDeviceSplit
674        );
675        assert_eq!(
676            DraftPlacementPlan::Device {
677                device: DevicePlan::new("mlx", "cpu:0").unwrap(),
678            }
679            .execution_topology(&target),
680            crate::SpeculativeExecutionTopology::CrossDeviceSplit
681        );
682    }
683}