eredu_runtime/
execution_plan.rs1use eredu_core::{
4 residency::OffloadConfig, ExecutionPlan, ResidencyPlan, WeightTransformationPlan,
5};
6
7use crate::{
8 DenseDiskStreamLoadOptions, LayerwiseLoadOptions, NormalizedLoadRequest,
9 NormalizedLoadRequestError, OrdinaryWeightResidency, ParallelLoadRequest,
10 ParameterBankLoadOptions, WeightResidency,
11};
12
13#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)]
17pub struct ResidencyDiagnostics {
18 backend_memory: bool,
19 process_memory: bool,
20}
21
22impl ResidencyDiagnostics {
23 pub const fn new(backend_memory: bool, process_memory: bool) -> Self {
25 Self {
26 backend_memory,
27 process_memory,
28 }
29 }
30}
31
32#[derive(Debug, thiserror::Error)]
34#[non_exhaustive]
35pub enum ExecutionPlanLoadError {
36 #[error(transparent)]
38 Plan(#[from] eredu_core::ExecutionPlanError),
39 #[error("single-device automatic plans require a 1x1x1 parallel topology; distributed plans require an explicit parallel load request")]
41 MissingParallelRequest,
42 #[error("parallel load request topology differs from the execution plan")]
44 ParallelTopologyMismatch,
45 #[error("unsupported execution-plan {0}")]
47 Unsupported(&'static str),
48 #[error(transparent)]
50 Offload(#[from] eredu_core::residency::OffloadError),
51 #[error(transparent)]
53 Residency(#[from] crate::WeightResidencyPolicyError),
54 #[error(transparent)]
56 Request(#[from] NormalizedLoadRequestError),
57}
58
59pub fn execution_plan_quantization(
61 transformation: WeightTransformationPlan,
62) -> Result<Option<eredu_core::QuantizationRequest>, ExecutionPlanLoadError> {
63 use eredu_core::QuantizationRequest;
64 let request = match transformation {
65 WeightTransformationPlan::PreserveCheckpoint => return Ok(None),
66 WeightTransformationPlan::Affine { bits, group_size } => QuantizationRequest::Affine {
67 group_size: u32::try_from(group_size).map_err(|_| {
68 NormalizedLoadRequestError::Quantization(format!(
69 "group_size must be non-negative, got {group_size}"
70 ))
71 })?,
72 bits: u8::try_from(bits).map_err(|_| {
73 NormalizedLoadRequestError::Quantization(format!("bits must fit in u8, got {bits}"))
74 })?,
75 },
76 WeightTransformationPlan::MxFp4 => QuantizationRequest::MxFp4,
77 _ => return Err(ExecutionPlanLoadError::Unsupported("weight transformation")),
78 };
79 NormalizedLoadRequest::with_quantization(request).weight_quantization()?;
80 Ok(Some(request))
81}
82
83impl NormalizedLoadRequest {
84 pub fn from_execution_plan(
90 plan: &ExecutionPlan,
91 diagnostics: ResidencyDiagnostics,
92 parallel: Option<ParallelLoadRequest>,
93 ) -> Result<Self, ExecutionPlanLoadError> {
94 plan.validate_structure()?;
95 match parallel {
96 Some(parallel) if parallel.rank().topology() != *plan.topology() => {
97 return Err(ExecutionPlanLoadError::ParallelTopologyMismatch);
98 }
99 None if !plan.topology().is_replicated() => {
100 return Err(ExecutionPlanLoadError::MissingParallelRequest);
101 }
102 _ => {}
103 }
104 let ordinary = match plan.residency() {
105 ResidencyPlan::FullyResident => OrdinaryWeightResidency::FullyResident,
106 ResidencyPlan::LayerwiseHost {
107 device_layer_window,
108 device_budget_bytes,
109 host_budget_bytes,
110 } => OrdinaryWeightResidency::LayerwiseHost(
111 LayerwiseLoadOptions::new(OffloadConfig::new(
112 *device_budget_bytes,
113 *host_budget_bytes,
114 *device_layer_window,
115 )?)
116 .with_max_cached_shards(plan.max_cached_shards())
117 .with_memory_sampling(diagnostics.backend_memory, diagnostics.process_memory),
118 ),
119 ResidencyPlan::DenseDiskStream {
120 device_budget_bytes,
121 host_budget_bytes,
122 host_lookahead,
123 background_queue,
124 } => OrdinaryWeightResidency::DenseDiskStream(
125 DenseDiskStreamLoadOptions::new(
126 *device_budget_bytes,
127 *host_budget_bytes,
128 *host_lookahead,
129 *background_queue,
130 )?
131 .with_max_cached_shards(plan.max_cached_shards())
132 .with_memory_sampling(diagnostics.backend_memory, diagnostics.process_memory),
133 ),
134 _ => return Err(ExecutionPlanLoadError::Unsupported("residency")),
135 };
136 let residency = match plan.expert_cache() {
137 Some(bank) => WeightResidency::with_independent_parameter_banks(
138 ordinary,
139 ParameterBankLoadOptions::new(
140 OffloadConfig::new(bank.device_budget_bytes(), bank.host_budget_bytes(), 1)?
141 .with_eviction_policy(bank.eviction_policy()),
142 bank.scratch_bytes(),
143 bank.prefill_bank_bytes(),
144 )?,
145 ),
146 None => WeightResidency::with_layers(ordinary.layers()),
147 };
148 let mut request = execution_plan_quantization(plan.weight_transformation())?
149 .map_or_else(Self::default, Self::with_quantization)
150 .with_weight_residency(residency)
151 .with_max_cached_shards(
152 std::num::NonZeroUsize::new(plan.max_cached_shards())
153 .expect("plan structure validates a positive reader limit"),
154 )
155 .with_required_session_capabilities(*plan.required_session_capabilities())
156 .with_prompt_cache_persistence(plan.prompt_cache_persistence())
157 .with_drafting_plan(plan.drafting())?;
158 if let Some(parallel) = parallel {
159 request = request.with_parallel_execution(parallel)?;
160 }
161 request.validate_model_preparation()?;
162 Ok(request)
163 }
164}
165
166#[cfg(test)]
167mod tests;