1use crate::{
4 backend::{DeviceCapabilities, SessionCapabilities},
5 residency::CacheEvictionPolicy,
6 topology::ParallelTopology,
7};
8use serde::{Deserialize, Serialize};
9
10pub const EXECUTION_PLAN_SCHEMA_VERSION: u32 = 4;
12
13pub const DEFAULT_MAX_CACHED_SHARDS: usize = 4;
15
16#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
21#[serde(transparent)]
22pub struct BackendId(String);
23
24impl BackendId {
25 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 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#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
57pub struct DevicePlan {
58 pub(crate) backend: BackendId,
60 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 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 pub const fn backend(&self) -> &BackendId {
98 &self.backend
99 }
100 pub fn device(&self) -> &str {
102 &self.device
103 }
104}
105
106#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
108#[serde(tag = "mode", rename_all = "snake_case")]
109#[non_exhaustive]
110pub enum ResidencyPlan {
111 FullyResident,
113 LayerwiseHost {
115 device_layer_window: usize,
117 #[serde(skip_serializing_if = "Option::is_none")]
119 device_budget_bytes: Option<u64>,
120 #[serde(skip_serializing_if = "Option::is_none")]
122 host_budget_bytes: Option<u64>,
123 },
124 DenseDiskStream {
126 device_budget_bytes: u64,
128 host_budget_bytes: u64,
130 host_lookahead: usize,
132 background_queue: usize,
134 },
135}
136
137#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
139pub struct ExpertCachePlan {
140 #[serde(skip_serializing_if = "Option::is_none")]
142 pub(crate) device_budget_bytes: Option<u64>,
143 #[serde(skip_serializing_if = "Option::is_none")]
145 pub(crate) host_budget_bytes: Option<u64>,
146 pub(crate) scratch_bytes: u64,
148 pub(crate) prefill_bank_bytes: u64,
150 pub(crate) eviction_policy: CacheEvictionPolicy,
152}
153
154impl ExpertCachePlan {
155 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 pub const fn device_budget_bytes(&self) -> Option<u64> {
173 self.device_budget_bytes
174 }
175 pub const fn host_budget_bytes(&self) -> Option<u64> {
177 self.host_budget_bytes
178 }
179 pub const fn scratch_bytes(&self) -> u64 {
181 self.scratch_bytes
182 }
183 pub const fn prefill_bank_bytes(&self) -> u64 {
185 self.prefill_bank_bytes
186 }
187 pub const fn eviction_policy(&self) -> CacheEvictionPolicy {
189 self.eviction_policy
190 }
191}
192
193#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
195#[serde(tag = "mode", rename_all = "snake_case")]
196#[non_exhaustive]
197pub enum DraftingPlan {
198 Disabled,
200 Embedded {
202 max_draft_tokens: usize,
204 lookahead: bool,
206 adaptive_lookahead: bool,
208 },
209 External {
211 model: String,
213 placement: DraftPlacementPlan,
215 max_draft_tokens: usize,
217 lookahead: bool,
219 adaptive_lookahead: bool,
221 },
222}
223
224impl DraftingPlan {
225 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 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 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#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
262#[serde(tag = "mode", rename_all = "snake_case")]
263#[non_exhaustive]
264pub enum DraftPlacementPlan {
265 Target,
267 Device {
269 device: DevicePlan,
271 },
272}
273
274impl DraftPlacementPlan {
275 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#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
289#[serde(tag = "mode", rename_all = "snake_case")]
290#[non_exhaustive]
291pub enum WeightTransformationPlan {
292 PreserveCheckpoint,
294 Affine {
296 bits: i32,
298 group_size: i32,
300 },
301 MxFp4,
303}
304
305#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
307pub struct ExecutionPlan {
308 pub(crate) schema_version: u32,
310 pub(crate) device: DevicePlan,
312 pub(crate) topology: ParallelTopology,
314 pub(crate) residency: ResidencyPlan,
316 pub(crate) weight_transformation: WeightTransformationPlan,
318 pub(crate) max_cached_shards: usize,
320 #[serde(skip_serializing_if = "Option::is_none")]
322 pub(crate) expert_cache: Option<ExpertCachePlan>,
323 pub(crate) drafting: DraftingPlan,
325 pub(crate) required_device_capabilities: DeviceCapabilities,
327 pub(crate) required_session_capabilities: SessionCapabilities,
329 #[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 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 pub const fn schema_version(&self) -> u32 {
358 self.schema_version
359 }
360 pub const fn device(&self) -> &DevicePlan {
362 &self.device
363 }
364 pub const fn topology(&self) -> &ParallelTopology {
366 &self.topology
367 }
368 pub const fn residency(&self) -> &ResidencyPlan {
370 &self.residency
371 }
372 pub const fn weight_transformation(&self) -> WeightTransformationPlan {
374 self.weight_transformation
375 }
376 pub const fn max_cached_shards(&self) -> usize {
378 self.max_cached_shards
379 }
380
381 pub const fn prompt_cache_persistence(&self) -> bool {
385 self.prompt_cache_persistence
386 }
387 pub const fn expert_cache(&self) -> Option<&ExpertCachePlan> {
389 self.expert_cache.as_ref()
390 }
391 pub const fn drafting(&self) -> &DraftingPlan {
393 &self.drafting
394 }
395 pub const fn required_device_capabilities(&self) -> &DeviceCapabilities {
397 &self.required_device_capabilities
398 }
399 pub const fn required_session_capabilities(&self) -> &SessionCapabilities {
401 &self.required_session_capabilities
402 }
403
404 pub fn with_topology(mut self, topology: ParallelTopology) -> Self {
406 self.topology = topology;
407 self
408 }
409 pub fn with_device(mut self, device: DevicePlan) -> Self {
411 self.device = device;
412 self
413 }
414 pub fn with_residency(mut self, residency: ResidencyPlan) -> Self {
416 self.residency = residency;
417 self
418 }
419 pub fn with_weight_transformation(mut self, transformation: WeightTransformationPlan) -> Self {
421 self.weight_transformation = transformation;
422 self
423 }
424 pub fn with_max_cached_shards(mut self, maximum: usize) -> Self {
426 self.max_cached_shards = maximum;
427 self
428 }
429
430 pub fn with_prompt_cache_persistence(mut self, required: bool) -> Self {
432 self.prompt_cache_persistence = required;
433 self
434 }
435 pub fn with_expert_cache(mut self, expert_cache: Option<ExpertCachePlan>) -> Self {
437 self.expert_cache = expert_cache;
438 self
439 }
440 pub fn with_drafting(mut self, drafting: DraftingPlan) -> Self {
442 self.drafting = drafting;
443 self
444 }
445 pub fn with_required_device_capabilities(mut self, capabilities: DeviceCapabilities) -> Self {
447 self.required_device_capabilities = capabilities;
448 self
449 }
450 pub fn with_required_session_capabilities(mut self, capabilities: SessionCapabilities) -> Self {
452 self.required_session_capabilities = capabilities;
453 self
454 }
455
456 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 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 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#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
540#[non_exhaustive]
541pub enum ExecutionPlanError {
542 #[error("execution-plan backend identifier must not be empty")]
544 EmptyBackendId,
545 #[error("execution-plan device identifier must not be empty")]
547 EmptyDeviceId,
548 #[error("unsupported execution-plan schema version {0}")]
550 Schema(u32),
551 #[error("execution plan requires unavailable capability {0}")]
553 Capability(&'static str),
554 #[error("execution-plan topology is invalid: {0}")]
556 Topology(String),
557 #[error("execution-plan max_cached_shards must be greater than zero")]
559 ZeroMappedShards,
560 #[error("execution-plan external draft model must not be empty")]
562 EmptyDraftModel,
563 #[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}