ferrum_models/executor/vnext_executor/
composition.rs

1//! Metadata-only candidate compilation followed by initialization of the exact
2//! retained result. The borrow ties compilation to the actual registries used
3//! for binding; callers cannot substitute a reconstructed capability list.
4
5use super::*;
6
7pub struct VNextRuntimeComposition<R: DeviceRuntime> {
8    pub(super) runtime: Arc<R>,
9    pub(super) registry: OperationRuntimeRegistry<R>,
10    weight_materializers: WeightMaterializerRegistry,
11    pub(super) catalog: CapabilityCatalog,
12}
13
14/// No Deserialize or public field construction: only successful static
15/// compilation produces this initialization input. It borrows the exact
16/// prepared source and runtime composition, including on the Auto path.
17pub struct VNextCompiledModel<'a, R: DeviceRuntime> {
18    pub(super) composition: &'a VNextRuntimeComposition<R>,
19    pub(super) prepared: &'a PreparedProductionModel,
20    pub(super) info: ModelInfo,
21    pub(super) config: VNextExecutorConfig,
22    pub(super) compilation: ProgramPlanCompilation,
23    pub(super) language_io: VNextLanguageIoIds,
24    pub(super) checkpoint_selection: Option<VNextCheckpointSelection>,
25    pub(super) repetition_capacity: u64,
26    pub(super) executor_startup: StartupPhaseTimer,
27}
28
29impl<R: DeviceRuntime> VNextRuntimeComposition<R> {
30    pub fn new(
31        runtime: Arc<R>,
32        registry: OperationRuntimeRegistry<R>,
33        weight_materializers: WeightMaterializerRegistry,
34        catalog: CapabilityCatalog,
35    ) -> Self {
36        Self {
37            runtime,
38            registry,
39            weight_materializers,
40            catalog,
41        }
42    }
43
44    pub fn runtime(&self) -> &R {
45        self.runtime.as_ref()
46    }
47    pub fn catalog(&self) -> &CapabilityCatalog {
48        &self.catalog
49    }
50    pub fn weight_materializers(&self) -> &WeightMaterializerRegistry {
51        &self.weight_materializers
52    }
53
54    pub fn compile_model<'a>(
55        &'a self,
56        prepared: &'a PreparedProductionModel,
57        info: ModelInfo,
58        engine_config: &EngineConfig,
59        config: VNextExecutorConfig,
60        materializer_selection: WeightMaterializerSelection,
61    ) -> Result<VNextCompiledModel<'a, R>> {
62        if self.runtime.descriptor() != self.catalog.device() {
63            return Err(FerrumError::device(
64                "vNext static compilation catalog differs from the actual device runtime",
65            ));
66        }
67        let executor_startup = StartupPhaseTimer::start("executor_composition_total");
68        let checkpoint_selection = VNextCheckpointSelection::from_config(
69            engine_config.runtime.vnext_checkpoint_capture.as_ref(),
70        )?;
71        let family = prepared.family();
72        let language_io = VNextModelExecutor::<R>::resolve_language_io_ids(family.program())?;
73        let input_capacity = u64::try_from(config.maximum_model_tokens)
74            .map_err(|_| FerrumError::config("vNext model length exceeds u64"))?;
75        let vocabulary_size = u64::try_from(info.vocab_size)
76            .map_err(|_| FerrumError::config("vNext vocabulary exceeds u64"))?;
77        let repetition_capacity = input_capacity.min(vocabulary_size);
78        let tensor = |dimensions, element_type| ProgramTensorSpec {
79            dimensions,
80            element_type,
81            layout: ResolvedTensorLayout::Contiguous,
82        };
83        let mut options = ProgramPlanCompileOptions::new(BTreeMap::from([
84            (
85                language_io.token_input.clone(),
86                tensor(vec![input_capacity], ElementType::U32),
87            ),
88            (
89                language_io.token_mask_input.clone(),
90                tensor(vec![vocabulary_size], ElementType::U8),
91            ),
92            (
93                language_io.repetition_token_ids_input.clone(),
94                tensor(vec![repetition_capacity], ElementType::U32),
95            ),
96            (
97                language_io.repetition_offsets_input.clone(),
98                tensor(vec![2], ElementType::U32),
99            ),
100            (
101                language_io.repetition_penalty_input.clone(),
102                tensor(vec![1], ElementType::F32),
103            ),
104        ]))
105        .map_err(|error| FerrumError::model(format!("vNext compile input: {error}")))?;
106        config.plan_observation.apply(family, &mut options)?;
107        if let Some(selection) = &checkpoint_selection {
108            selection.retain_in(&mut options);
109        }
110        options.require_weight_materializer_selection(materializer_selection);
111        let compile_phase = StartupPhaseTimer::start("plan_compile");
112        let compilation = ProgramPlanCompiler::compile_with_weight_materializers(
113            family,
114            &self.catalog,
115            &config.runtime_policy,
116            &self.registry.planning(),
117            &self.weight_materializers,
118            &options,
119        )
120        .map_err(|error| FerrumError::model(format!("vNext plan compile: {error}")))?;
121        config
122            .plan_observation
123            .validate_compilation(family, &compilation)?;
124        compile_phase.finish();
125        Ok(VNextCompiledModel {
126            composition: self,
127            prepared,
128            info,
129            config,
130            compilation,
131            language_io,
132            checkpoint_selection,
133            repetition_capacity,
134            executor_startup,
135        })
136    }
137}
138
139impl<R: DeviceRuntime> VNextCompiledModel<'_, R> {
140    pub fn prepared(&self) -> &PreparedProductionModel {
141        self.prepared
142    }
143    pub fn compilation(&self) -> &ProgramPlanCompilation {
144        &self.compilation
145    }
146
147    /// Evaluate the retained compiled plan without allocating device memory.
148    pub fn startup_peak_bytes(&self, workload: ferrum_types::StartupWorkload) -> Result<u64> {
149        let (context, frontier, sequences, tokens) = match workload {
150            ferrum_types::StartupWorkload::Prefill {
151                context_tokens,
152                chunk_tokens,
153            } => (context_tokens, context_tokens, 1, chunk_tokens),
154            ferrum_types::StartupWorkload::Decode {
155                context_tokens,
156                active_sequences,
157            } => (context_tokens, 1, active_sequences, active_sequences),
158        };
159        let context = u64::try_from(context)
160            .map_err(|_| FerrumError::config("startup context exceeds u64"))?;
161        let frontier = u64::try_from(frontier)
162            .map_err(|_| FerrumError::config("startup sequence frontier exceeds u64"))?;
163        let sequences = u32::try_from(sequences)
164            .map_err(|_| FerrumError::config("startup sequence count exceeds u32"))?;
165        let tokens = u64::try_from(tokens)
166            .map_err(|_| FerrumError::config("startup token count exceeds u64"))?;
167        self.compilation
168            .executable()
169            .execution_plan()
170            .payload()
171            .memory()
172            .startup_workload_peak_bytes(context, frontier, sequences, tokens)
173            .map_err(|error| FerrumError::config(format!("compiled startup memory: {error}")))
174    }
175
176    /// Attach final evaluation evidence after confirming the report describes
177    /// this retained compilation's immutable policy.
178    pub fn set_startup_memory_plan(&mut self, plan: ferrum_types::StartupMemoryPlan) -> Result<()> {
179        let memory = self
180            .compilation
181            .executable()
182            .execution_plan()
183            .payload()
184            .memory();
185        if plan.selected.context_tokens != self.config.maximum_model_tokens
186            || plan.selected.max_sequences as u64 != u64::from(memory.maximum_active_sequences())
187            || plan.selected.max_batch_tokens as u64
188                != self.config.runtime_policy.maximum_scheduled_tokens()
189            || plan.request.usable_capacity_bytes != memory.usable_capacity_bytes()
190            || plan.context_peak_bytes > memory.usable_capacity_bytes()
191            || plan.decode_peak_bytes > memory.usable_capacity_bytes()
192        {
193            return Err(FerrumError::config(
194                "startup memory report differs from the retained compiled plan",
195            ));
196        }
197        self.config.startup_memory_plan = Some(plan);
198        Ok(())
199    }
200
201    pub fn initialize<F>(self, resolve_plan: F) -> Result<VNextModelExecutor<R>>
202    where
203        F: FnOnce(
204            &PreparedProductionModel,
205            &ResolvedRuntimePolicy,
206            &CapabilityCatalog,
207            &ProgramPlanCompilation,
208        ) -> Result<ResolvedModelPlan>,
209    {
210        VNextModelExecutor::from_compiled_model(self, resolve_plan)
211    }
212}
ferrum_models/executor/vnext_executor/composition.rs

ferrum_models/executor/vnext_executor/
composition.rs
