ferrum_interfaces/vnext/
numerical.rs

1//! Numerical ABI declarations are independent of source containers and devices.
2
3use std::collections::{BTreeMap, BTreeSet};
4
5pub use ferrum_types::{KvStorageFormat, NumericalExecutionPolicy, NumericalProfileId};
6use serde::{Deserialize, Serialize};
7use sha2::{Digest, Sha256};
8
9use super::{
10    ContractVersion, ElementType, ModelFamilyId, ModelProgram, OperationId, ProgramValueId,
11    ResolvedTensorSpec, StateSpec, VNextError,
12};
13
14mod kv_storage;
15pub use kv_storage::KvStateStorage;
16
17/// The referenced operation version defines its internal ordering and rounding
18/// points. Arithmetic types are explicit here; copying/indexing operations have
19/// no multiplication or accumulation. A fused operation's internal conversions
20/// remain part of its versioned operation contract.
21#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
22#[serde(deny_unknown_fields)]
23pub struct NumericalOperationContract {
24    pub operation_id: OperationId,
25    pub version: ContractVersion,
26    pub multiplication_type: Option<ElementType>,
27    pub accumulation_type: Option<ElementType>,
28}
29
30/// A complete numerical choice declared by a family before program generation.
31/// Public decoding alone does not confer trust: the family registration must
32/// reproduce this profile from its validated typed model definition.
33#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
34#[serde(deny_unknown_fields)]
35pub struct NumericalExecutionProfile {
36    pub id: NumericalProfileId,
37    pub version: ContractVersion,
38    pub family_id: ModelFamilyId,
39    /// The main activation is explicit so product metadata need not guess a
40    /// dtype from weights, a container, or a profile's human-readable name.
41    pub primary_activation: ProgramValueId,
42    /// Every operation output has a declared storage/rounding boundary. This
43    /// includes residuals, intermediate activations, logits and token outputs.
44    pub boundaries: BTreeMap<ProgramValueId, ElementType>,
45    /// State shapes, initialization and capacity are shared with the final
46    /// program, permitting startup sizing without constructing a dummy program.
47    pub states: Vec<StateSpec>,
48    /// Typed association of every causal KV payload and its optional scales.
49    /// Empty means the family declares no applicable KV state.
50    pub kv_storage: Vec<KvStateStorage>,
51    pub operations: Vec<NumericalOperationContract>,
52}
53
54fn invalid(family: &ModelFamilyId, reason: impl Into<String>) -> VNextError {
55    VNextError::InvalidModelConfig {
56        family_id: family.to_string(),
57        field: "numerical_profile".to_owned(),
58        reason: reason.into(),
59    }
60}
61
62fn floating(element_type: ElementType) -> bool {
63    matches!(
64        element_type,
65        ElementType::F16 | ElementType::Bf16 | ElementType::F32
66    )
67}
68
69impl NumericalExecutionProfile {
70    pub fn validate(&self) -> Result<(), VNextError> {
71        if self.version.major == 0 || self.operations.is_empty() {
72            return Err(invalid(
73                &self.family_id,
74                "profile version and operation contracts must be explicit",
75            ));
76        }
77        if !self
78            .boundaries
79            .get(&self.primary_activation)
80            .is_some_and(|dtype| floating(*dtype))
81        {
82            return Err(invalid(
83                &self.family_id,
84                "primary activation must name a floating-point boundary",
85            ));
86        }
87        let mut operation_ids = BTreeSet::new();
88        for operation in &self.operations {
89            if operation.version.major == 0
90                || !operation_ids.insert(&operation.operation_id)
91                || operation
92                    .multiplication_type
93                    .is_some_and(|dtype| !floating(dtype))
94                || operation
95                    .accumulation_type
96                    .is_some_and(|dtype| !floating(dtype))
97            {
98                return Err(invalid(&self.family_id, "operation contracts need unique identities, valid versions and floating-point arithmetic types"));
99            }
100        }
101        let mut state_ids = BTreeSet::new();
102        let mut state_values = BTreeSet::new();
103        for state in &self.states {
104            if !state_ids.insert(&state.id) || !state_values.insert(&state.value_id) {
105                return Err(invalid(
106                    &self.family_id,
107                    "state identities and values must be unique",
108                ));
109            }
110            state.tensor.validate("numerical_profile.state")?;
111            state.capacity_demand.validate(state.tensor.byte_len()?)?;
112        }
113        self.kv_storage_format()?;
114        Ok(())
115    }
116
117    pub(crate) fn normalize(&mut self) {
118        self.states.sort_by(|left, right| left.id.cmp(&right.id));
119        self.kv_storage
120            .sort_by(|left, right| left.payload_state().cmp(right.payload_state()));
121        self.operations
122            .sort_by(|left, right| left.operation_id.cmp(&right.operation_id));
123    }
124
125    pub fn activation_type(&self) -> Result<ElementType, VNextError> {
126        self.validate()?;
127        Ok(self.boundaries[&self.primary_activation])
128    }
129
130    pub fn kv_storage_format(&self) -> Result<Option<KvStorageFormat>, VNextError> {
131        kv_storage::validate_kv_storage(&self.family_id, &self.kv_storage, &self.states)
132    }
133
134    pub fn fingerprint(&self) -> Result<String, VNextError> {
135        self.validate()?;
136        let mut canonical = self.clone();
137        canonical.normalize();
138        let bytes = serde_json::to_vec(&canonical).map_err(|error| VNextError::Serialization {
139            context: "serialize numerical execution profile",
140            message: error.to_string(),
141        })?;
142        Ok(format!("{:x}", Sha256::digest(bytes)))
143    }
144
145    /// Preparation validates the declared graph contract; compiler inference
146    /// subsequently proves every boundary's actual dtype against real operations.
147    pub fn validate_program(&self, program: &ModelProgram) -> Result<(), VNextError> {
148        self.validate()?;
149        if program.family_id() != &self.family_id {
150            return Err(invalid(
151                &self.family_id,
152                "program belongs to another family",
153            ));
154        }
155        let contracts: BTreeMap<_, _> = self
156            .operations
157            .iter()
158            .map(|operation| (&operation.operation_id, operation.version))
159            .collect();
160        let mut used_operations = BTreeSet::new();
161        let mut outputs = BTreeSet::new();
162        for node in program.blocks().iter().flat_map(|block| &block.nodes) {
163            if contracts.get(&node.operation_id) != Some(&node.required_version) {
164                return Err(invalid(
165                    &self.family_id,
166                    format!(
167                        "node {} uses an undeclared numerical operation/version",
168                        node.id
169                    ),
170                ));
171            }
172            used_operations.insert(&node.operation_id);
173            outputs.extend(&node.outputs);
174        }
175        if used_operations.len() != contracts.len() || outputs != self.boundaries.keys().collect() {
176            return Err(invalid(
177                &self.family_id,
178                "profile must describe exactly the program's operations and output boundaries",
179            ));
180        }
181        let states: BTreeMap<_, _> = self.states.iter().map(|state| (&state.id, state)).collect();
182        let actual: BTreeMap<_, _> = program
183            .states()
184            .iter()
185            .map(|state| (&state.id, state))
186            .collect();
187        if states != actual {
188            return Err(invalid(
189                &self.family_id,
190                "program state ABI differs from the numerical profile",
191            ));
192        }
193        kv_storage::validate_program_kv_storage(self, program)?;
194        Ok(())
195    }
196
197    pub fn validate_inferred_boundaries(
198        &self,
199        values: &BTreeMap<ProgramValueId, ResolvedTensorSpec>,
200    ) -> Result<(), VNextError> {
201        for (value, dtype) in &self.boundaries {
202            if values.get(value).map(ResolvedTensorSpec::element_type) != Some(*dtype) {
203                return Err(invalid(
204                    &self.family_id,
205                    format!(
206                        "inferred dtype for {value} differs from the selected numerical profile"
207                    ),
208                ));
209            }
210        }
211        Ok(())
212    }
213}
214
215/// The preference order is a versioned family declaration, not registry order.
216/// Profiles absent from `auto_preference` remain available only to `Require`;
217/// this lets new numerical combinations be qualified before changing defaults.
218#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
219pub struct FamilyNumericalProfiles {
220    version: ContractVersion,
221    profiles: Vec<NumericalExecutionProfile>,
222    auto_preference: Vec<NumericalProfileId>,
223}
224
225impl FamilyNumericalProfiles {
226    pub fn new(
227        family_id: &ModelFamilyId,
228        version: ContractVersion,
229        mut profiles: Vec<NumericalExecutionProfile>,
230        auto_preference: Vec<NumericalProfileId>,
231    ) -> Result<Self, VNextError> {
232        if version.major == 0 || profiles.is_empty() {
233            return Err(invalid(
234                family_id,
235                "family must declare versioned numerical profiles",
236            ));
237        }
238        let mut ids = BTreeSet::new();
239        for profile in &mut profiles {
240            profile.validate()?;
241            profile.normalize();
242            if &profile.family_id != family_id || !ids.insert(profile.id.clone()) {
243                return Err(invalid(
244                    family_id,
245                    "numerical profiles must belong to this family and have unique identities",
246                ));
247            }
248        }
249        let mut automatic = BTreeSet::new();
250        if auto_preference
251            .iter()
252            .any(|id| !ids.contains(id) || !automatic.insert(id))
253        {
254            return Err(invalid(
255                family_id,
256                "Auto preferences must name unique declared profiles",
257            ));
258        }
259        profiles.sort_by(|left, right| left.id.cmp(&right.id));
260        Ok(Self {
261            version,
262            profiles,
263            auto_preference,
264        })
265    }
266
267    pub fn version(&self) -> ContractVersion {
268        self.version
269    }
270    pub fn profiles(&self) -> &[NumericalExecutionProfile] {
271        &self.profiles
272    }
273    pub fn auto_preference(&self) -> &[NumericalProfileId] {
274        &self.auto_preference
275    }
276
277    pub fn resolve(
278        &self,
279        id: &NumericalProfileId,
280    ) -> Result<&NumericalExecutionProfile, VNextError> {
281        self.profiles
282            .iter()
283            .find(|profile| &profile.id == id)
284            .ok_or_else(|| {
285                invalid(
286                    &self.profiles[0].family_id,
287                    format!("numerical profile {id} is not declared by this family"),
288                )
289            })
290    }
291
292    pub fn candidates(
293        &self,
294        policy: &NumericalExecutionPolicy,
295        kv_storage: KvStorageFormat,
296    ) -> Result<Vec<&NumericalExecutionProfile>, VNextError> {
297        let ids = match policy {
298            NumericalExecutionPolicy::Auto => self.auto_preference.iter().collect::<Vec<_>>(),
299            NumericalExecutionPolicy::Require(id) => vec![id],
300        };
301        if ids.is_empty() {
302            return Err(invalid(
303                &self.profiles[0].family_id,
304                "no numerical profile is qualified for Auto",
305            ));
306        }
307        let mut candidates = Vec::new();
308        for id in ids {
309            let profile = self.resolve(id)?;
310            // F16 is the existing default, including programs with no causal
311            // KV state. An explicit INT8 request must actually change KV storage.
312            let matches = profile
313                .kv_storage_format()?
314                .map_or(kv_storage == KvStorageFormat::F16, |actual| {
315                    actual == kv_storage
316                });
317            if matches {
318                candidates.push(profile);
319            } else if matches!(policy, NumericalExecutionPolicy::Require(_)) {
320                return Err(invalid(&profile.family_id, format!(
321                    "numerical profile {} conflicts with requested KV storage {kv_storage}; select a compatible profile and kv_dtype", profile.id
322                )));
323            }
324        }
325        if candidates.is_empty() {
326            return Err(invalid(&self.profiles[0].family_id, format!(
327                "no qualified numerical profile supports KV storage {kv_storage}; it may be unsupported or not applicable to this family; use fp16"
328            )));
329        }
330        Ok(candidates)
331    }
332}
ferrum_interfaces/vnext/numerical.rs

ferrum_interfaces/vnext/
numerical.rs
