Skip to main content

ferrum_interfaces/vnext/
model.rs

1use serde::de::DeserializeOwned;
2use serde::{Deserialize, Deserializer, Serialize};
3use sha2::{Digest, Sha256};
4use std::collections::{BTreeMap, BTreeSet};
5
6use super::{
7    checked_elements, physical_component_ids, validate_physical_layout_budget, AttributeId,
8    AxisWeightComponent, BlockQuantizationSpec, CanonicalRational, CompositeWeightPart,
9    ContractVersion, ElementType, ExternalModelMetadataId, ModelFamilyId, NodeId, OperationId,
10    PhysicalStorageLayout, PhysicalWeightComponentBinding, PhysicalWeightLayout,
11    PhysicalWeightPadding, ProgramValueId, QuantizationGrouping, QuantizationPacking,
12    QuantizationSpec, ResolvedTensorLayout, ResolvedWeightBinding, ResolvedWeightComponentLayout,
13    ResolvedWeightLogicalValidation, SemanticValue, StateId, StateInitialization, TokenizerId,
14    VNextError, WeightComponentRole, WeightEncoding, WeightFormatId, WeightId, WeightLayoutId,
15    MAX_PHYSICAL_WEIGHT_LAYOUT_DEPTH, MAX_PHYSICAL_WEIGHT_LAYOUT_NODES,
16};
17
18#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
19pub struct WeightComponentSpec {
20    pub id: WeightId,
21    pub role: WeightComponentRole,
22    /// Ordered checkpoint tensors forming this physical component. Dense
23    /// multi-source components use `[source_count, source_shape...]` and stack
24    /// tensors in this exact order. Format adapters may instead define a typed
25    /// source recipe (for example values plus quantization sidecars), but must
26    /// validate every source identity, order, shape, and resulting byte count.
27    /// Aliases are resolved before the model family constructs this schema.
28    pub external_names: Vec<String>,
29    pub dimensions: Vec<u64>,
30    pub encoding: WeightEncoding,
31    pub required: bool,
32}
33
34impl WeightComponentSpec {
35    /// Exact bytes occupied by this physical component. Separate-component
36    /// quantized dimensions are byte dimensions. Block-quantized dimensions
37    /// are block-grid dimensions and are multiplied by the block ABI size.
38    /// Logical element counts live on `WeightTensorSpec` and are checked
39    /// against the corresponding physical layout contract.
40    pub fn physical_bytes(&self) -> Result<u64, VNextError> {
41        self.encoding.physical_bytes(&self.dimensions, &self.id)
42    }
43
44    pub fn dense_element_type(&self) -> Option<ElementType> {
45        self.encoding.dense_element_type()
46    }
47
48    pub fn physical_element_type(&self) -> ElementType {
49        self.dense_element_type().unwrap_or(ElementType::U8)
50    }
51}
52
53#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
54pub struct PhysicalWeightComponentRef {
55    pub component_id: WeightId,
56    pub physical_dimensions: Vec<u64>,
57    pub resource_bytes: u64,
58    pub element_type: ElementType,
59}
60
61#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
62pub struct WeightTensorSpec {
63    pub id: WeightId,
64    pub dimensions: Vec<u64>,
65    /// Dtype observed by the semantic program after decoding any physical
66    /// packing, indexing, or quantization components.
67    pub logical_element_type: ElementType,
68    pub physical_layout: PhysicalWeightLayout,
69    pub required: bool,
70}
71
72impl WeightTensorSpec {
73    pub fn logical_elements(&self) -> Result<u64, VNextError> {
74        checked_elements(&self.dimensions).ok_or_else(|| VNextError::InvalidExecutionPlan {
75            reason: format!("logical weight `{}` element count overflows u64", self.id),
76        })
77    }
78
79    pub fn logical_bytes(&self) -> Result<u64, VNextError> {
80        self.logical_elements()?
81            .checked_mul(self.logical_element_type.size_bytes())
82            .ok_or_else(|| VNextError::InvalidExecutionPlan {
83                reason: format!("logical weight `{}` byte size overflows u64", self.id),
84            })
85    }
86}
87
88#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
89pub struct WeightSchema {
90    pub format_id: WeightFormatId,
91    pub layout_id: WeightLayoutId,
92    pub version: ContractVersion,
93    pub components: Vec<WeightComponentSpec>,
94    pub tensors: Vec<WeightTensorSpec>,
95}
96
97impl WeightSchema {
98    pub(crate) fn normalize(&mut self) {
99        self.components
100            .sort_by(|left, right| left.id.cmp(&right.id));
101        for tensor in &mut self.tensors {
102            tensor.physical_layout.normalize();
103        }
104        self.tensors.sort_by(|left, right| left.id.cmp(&right.id));
105    }
106
107    pub fn validate(&self, family_id: &ModelFamilyId) -> Result<(), VNextError> {
108        if self.version.major == 0 || self.components.is_empty() || self.tensors.is_empty() {
109            return Err(VNextError::UnknownWeightLayout {
110                family_id: family_id.to_string(),
111                layout_id: self.layout_id.to_string(),
112            });
113        }
114        for tensor in &self.tensors {
115            validate_physical_layout_budget(&tensor.physical_layout).map_err(|reason| {
116                VNextError::InvalidModelConfig {
117                    family_id: family_id.to_string(),
118                    field: format!("weight_schema.tensors.{}.physical_layout", tensor.id),
119                    reason,
120                }
121            })?;
122        }
123        let mut component_ids = BTreeSet::new();
124        let mut names = BTreeSet::new();
125        let mut components = BTreeMap::new();
126        let mut quantization_abis = BTreeMap::new();
127        for component in &self.components {
128            if !component_ids.insert(component.id.clone())
129                || component.external_names.is_empty()
130                || component.dimensions.is_empty()
131                || component
132                    .external_names
133                    .iter()
134                    .any(|name| name.trim().is_empty() || !names.insert(name.clone()))
135                || component.dimensions.iter().any(|extent| *extent == 0)
136            {
137                return Err(VNextError::InvalidModelConfig {
138                    family_id: family_id.to_string(),
139                    field: "weight_schema.components".to_owned(),
140                    reason: "component identities, names, and dimensions must be valid and unique"
141                        .to_owned(),
142                });
143            }
144            if let WeightEncoding::Quantized(quantization) = &component.encoding {
145                quantization.validate()?;
146            }
147            if let WeightEncoding::BlockQuantized(quantization) = &component.encoding {
148                quantization.validate()?;
149            }
150            let quantization_format = match &component.encoding {
151                WeightEncoding::Quantized(spec) => Some(&spec.format_id),
152                WeightEncoding::BlockQuantized(spec) => Some(&spec.format_id),
153                WeightEncoding::Dense { .. } | WeightEncoding::DenseAffine { .. } => None,
154            };
155            if let Some(format_id) = quantization_format {
156                if let Some(existing) = quantization_abis.get(format_id) {
157                    if existing != &component.encoding {
158                        return Err(VNextError::InvalidModelConfig {
159                            family_id: family_id.to_string(),
160                            field: "weight_schema.components.encoding".to_owned(),
161                            reason: format!(
162                                "quantization format `{format_id}` maps to conflicting physical ABIs"
163                            ),
164                        });
165                    }
166                } else {
167                    quantization_abis.insert(format_id.clone(), component.encoding.clone());
168                }
169            }
170            if let WeightEncoding::DenseAffine { element_type, .. } = &component.encoding {
171                if component.role != WeightComponentRole::Values
172                    || !matches!(
173                        element_type,
174                        ElementType::F16 | ElementType::Bf16 | ElementType::F32
175                    )
176                {
177                    return Err(VNextError::InvalidModelConfig {
178                        family_id: family_id.to_string(),
179                        field: "weight_schema.components.encoding".to_owned(),
180                        reason: format!(
181                            "affine dense component `{}` must be a floating-point value component",
182                            component.id
183                        ),
184                    });
185                }
186            }
187            component
188                .physical_bytes()
189                .map_err(|error| VNextError::InvalidModelConfig {
190                    family_id: family_id.to_string(),
191                    field: "weight_schema.components.dimensions".to_owned(),
192                    reason: error.to_string(),
193                })?;
194            let role_encoding_valid = match component.role {
195                WeightComponentRole::Scales => matches!(
196                    component.encoding,
197                    WeightEncoding::Dense {
198                        element_type: ElementType::F16 | ElementType::Bf16 | ElementType::F32
199                    }
200                ),
201                WeightComponentRole::ZeroPoints
202                | WeightComponentRole::Indices
203                | WeightComponentRole::Permutation => matches!(
204                    component.encoding,
205                    WeightEncoding::Dense {
206                        element_type: ElementType::U8
207                            | ElementType::U32
208                            | ElementType::I8
209                            | ElementType::I32
210                    }
211                ),
212                WeightComponentRole::PackedValues => {
213                    matches!(
214                        component.encoding,
215                        WeightEncoding::Quantized(_) | WeightEncoding::BlockQuantized(_)
216                    )
217                }
218                _ => true,
219            };
220            if !role_encoding_valid {
221                return Err(VNextError::InvalidModelConfig {
222                    family_id: family_id.to_string(),
223                    field: "weight_schema.components.encoding".to_owned(),
224                    reason: format!(
225                        "component `{}` encoding is incompatible with its structural role",
226                        component.id
227                    ),
228                });
229            }
230            components.insert(component.id.clone(), component);
231        }
232
233        let mut tensor_ids = BTreeSet::new();
234        let mut referenced_components = BTreeSet::new();
235        for tensor in &self.tensors {
236            if !tensor_ids.insert(tensor.id.clone())
237                || tensor.dimensions.is_empty()
238                || tensor.dimensions.iter().any(|extent| *extent == 0)
239            {
240                return Err(VNextError::InvalidModelConfig {
241                    family_id: family_id.to_string(),
242                    field: "weight_schema.tensors".to_owned(),
243                    reason: "logical weight identities and dimensions must be valid and unique"
244                        .to_owned(),
245                });
246            }
247            self.validate_physical_layout(
248                family_id,
249                tensor,
250                &components,
251                &mut referenced_components,
252            )?;
253        }
254        if let Some(component) = self
255            .components
256            .iter()
257            .find(|component| component.required && !referenced_components.contains(&component.id))
258        {
259            return Err(VNextError::InvalidModelConfig {
260                family_id: family_id.to_string(),
261                field: "weight_schema.components".to_owned(),
262                reason: format!(
263                    "required component `{}` is not referenced by a logical weight",
264                    component.id
265                ),
266            });
267        }
268        Ok(())
269    }
270
271    pub fn fingerprint(&self) -> Result<String, VNextError> {
272        let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
273            context: "fingerprint weight schema",
274            message: error.to_string(),
275        })?;
276        Ok(format!("{:x}", Sha256::digest(bytes)))
277    }
278
279    pub fn quantization_formats(&self) -> BTreeSet<super::QuantizationFormatId> {
280        self.components
281            .iter()
282            .filter_map(|component| match &component.encoding {
283                WeightEncoding::Quantized(spec) => Some(spec.format_id.clone()),
284                WeightEncoding::BlockQuantized(spec) => Some(spec.format_id.clone()),
285                WeightEncoding::Dense { .. } | WeightEncoding::DenseAffine { .. } => None,
286            })
287            .collect()
288    }
289
290    fn validate_physical_layout(
291        &self,
292        family_id: &ModelFamilyId,
293        tensor: &WeightTensorSpec,
294        components: &BTreeMap<WeightId, &WeightComponentSpec>,
295        referenced: &mut BTreeSet<WeightId>,
296    ) -> Result<(), VNextError> {
297        let mut validator = PhysicalLayoutValidator {
298            family_id,
299            tensor_id: &tensor.id,
300            components,
301            referenced,
302            visited_nodes: 0,
303        };
304        validator.validate_layout(
305            &tensor.physical_layout,
306            &tensor.dimensions,
307            tensor.logical_element_type,
308            1,
309        )
310    }
311
312    pub fn tensor(&self, weight_id: &WeightId) -> Option<&WeightTensorSpec> {
313        self.tensors.iter().find(|tensor| &tensor.id == weight_id)
314    }
315
316    /// Ordered physical components needed to materialize one logical value.
317    /// The order is the schema's component order and is therefore stable for
318    /// fingerprints and resource-plan evidence.
319    pub fn physical_component_refs(
320        &self,
321        weight_id: &WeightId,
322    ) -> Result<Vec<&WeightComponentSpec>, VNextError> {
323        let tensor = self
324            .tensor(weight_id)
325            .ok_or_else(|| VNextError::InvalidExecutionPlan {
326                reason: format!("unknown logical weight `{weight_id}`"),
327            })?;
328        let required = physical_component_ids(&tensor.physical_layout).map_err(|reason| {
329            VNextError::InvalidExecutionPlan {
330                reason: format!(
331                    "logical weight `{weight_id}` has invalid physical layout: {reason}"
332                ),
333            }
334        })?;
335        let result = self
336            .components
337            .iter()
338            .filter(|component| required.contains(&component.id))
339            .collect::<Vec<_>>();
340        if result.len() != required.len() {
341            return Err(VNextError::InvalidExecutionPlan {
342                reason: format!("logical weight `{weight_id}` references an unknown component"),
343            });
344        }
345        Ok(result)
346    }
347
348    pub fn physical_bytes(&self, weight_id: &WeightId) -> Result<u64, VNextError> {
349        self.physical_component_refs(weight_id)?
350            .into_iter()
351            .try_fold(0_u64, |total, component| {
352                total
353                    .checked_add(component.physical_bytes()?)
354                    .ok_or_else(|| VNextError::InvalidExecutionPlan {
355                        reason: format!("logical weight `{weight_id}` physical bytes overflow u64"),
356                    })
357            })
358    }
359
360    pub fn physical_resource_requirements(
361        &self,
362        weight_id: &WeightId,
363    ) -> Result<Vec<PhysicalWeightComponentRef>, VNextError> {
364        self.physical_component_refs(weight_id)?
365            .into_iter()
366            .map(|component| {
367                Ok(PhysicalWeightComponentRef {
368                    component_id: component.id.clone(),
369                    physical_dimensions: component.dimensions.clone(),
370                    resource_bytes: component.physical_bytes()?,
371                    element_type: component.physical_element_type(),
372                })
373            })
374            .collect()
375    }
376}
377
378impl ResolvedWeightBinding {
379    pub fn from_schema(schema: &WeightSchema, weight_id: &WeightId) -> Result<Self, VNextError> {
380        let tensor = schema
381            .tensor(weight_id)
382            .ok_or_else(|| VNextError::InvalidExecutionPlan {
383                reason: format!("unknown logical weight `{weight_id}`"),
384            })?;
385        let mut components = schema
386            .physical_component_refs(weight_id)?
387            .into_iter()
388            .map(|component| {
389                ResolvedWeightComponentLayout::from_parts(
390                    component.id.clone(),
391                    component.role,
392                    component.dimensions.clone(),
393                    component.encoding.clone(),
394                )
395            })
396            .collect::<Vec<_>>();
397        components.sort_by(|left, right| left.component_id().cmp(right.component_id()));
398        let binding = Self::from_parts(
399            weight_id.clone(),
400            schema.format_id.clone(),
401            schema.layout_id.clone(),
402            schema.version,
403            tensor.physical_layout.clone(),
404            components,
405        )?;
406        binding.validate_logical(&tensor.dimensions, tensor.logical_element_type)?;
407        Ok(binding)
408    }
409
410    pub fn validate_logical(
411        &self,
412        logical_dimensions: &[u64],
413        logical_element_type: ElementType,
414    ) -> Result<(), VNextError> {
415        ResolvedWeightLogicalValidation::validate_logical_contract(
416            self,
417            logical_dimensions,
418            logical_element_type,
419        )
420    }
421}
422
423impl ResolvedWeightLogicalValidation for ResolvedWeightBinding {
424    fn validate_logical_contract(
425        &self,
426        logical_dimensions: &[u64],
427        logical_element_type: ElementType,
428    ) -> Result<(), VNextError> {
429        self.validate_structure()?;
430        let schema = WeightSchema {
431            format_id: self.schema_format_id().clone(),
432            layout_id: self.layout_id().clone(),
433            version: self.schema_version(),
434            components: self
435                .components()
436                .iter()
437                .map(|component| WeightComponentSpec {
438                    id: component.component_id().clone(),
439                    role: component.role(),
440                    external_names: vec![format!("resolved.{}", component.component_id())],
441                    dimensions: component.physical_dimensions().to_vec(),
442                    encoding: component.encoding().clone(),
443                    required: true,
444                })
445                .collect(),
446            tensors: vec![WeightTensorSpec {
447                id: self.weight_id().clone(),
448                dimensions: logical_dimensions.to_vec(),
449                logical_element_type,
450                physical_layout: self.physical_layout().clone(),
451                required: true,
452            }],
453        };
454        schema.validate(&ModelFamilyId::new("family.resolved-weight-binding")?)
455    }
456}
457
458struct PhysicalLayoutValidator<'schema, 'references> {
459    family_id: &'schema ModelFamilyId,
460    tensor_id: &'schema WeightId,
461    components: &'schema BTreeMap<WeightId, &'schema WeightComponentSpec>,
462    referenced: &'references mut BTreeSet<WeightId>,
463    visited_nodes: usize,
464}
465
466impl<'schema, 'references> PhysicalLayoutValidator<'schema, 'references> {
467    fn invalid(&self, reason: impl Into<String>) -> VNextError {
468        VNextError::InvalidModelConfig {
469            family_id: self.family_id.to_string(),
470            field: format!("weight_schema.tensors.{}.physical_layout", self.tensor_id),
471            reason: reason.into(),
472        }
473    }
474
475    fn visit_node(&mut self, depth: usize) -> Result<(), VNextError> {
476        if depth > MAX_PHYSICAL_WEIGHT_LAYOUT_DEPTH {
477            return Err(self.invalid(format!(
478                "physical layout depth exceeds {MAX_PHYSICAL_WEIGHT_LAYOUT_DEPTH}"
479            )));
480        }
481        self.visited_nodes = self
482            .visited_nodes
483            .checked_add(1)
484            .ok_or_else(|| self.invalid("physical layout node count overflows usize"))?;
485        if self.visited_nodes > MAX_PHYSICAL_WEIGHT_LAYOUT_NODES {
486            return Err(self.invalid(format!(
487                "physical layout node count exceeds {MAX_PHYSICAL_WEIGHT_LAYOUT_NODES}"
488            )));
489        }
490        Ok(())
491    }
492
493    fn component(
494        &self,
495        component_id: &WeightId,
496    ) -> Result<&'schema WeightComponentSpec, VNextError> {
497        self.components
498            .get(component_id)
499            .copied()
500            .ok_or_else(|| self.invalid(format!("unknown component `{component_id}`")))
501    }
502
503    fn bind_component(
504        &mut self,
505        binding: &PhysicalWeightComponentBinding,
506        semantic_dimensions: &[u64],
507        role: WeightComponentRole,
508        depth: usize,
509    ) -> Result<&'schema WeightComponentSpec, VNextError> {
510        self.visit_node(depth)?;
511        let component = self.component(&binding.component_id)?;
512        if component.role != role {
513            return Err(self.invalid(format!(
514                "component `{}` has role {:?}, expected {:?}",
515                component.id, component.role, role
516            )));
517        }
518        self.validate_storage(component, semantic_dimensions, &binding.storage)?;
519        if !self.referenced.insert(component.id.clone()) {
520            return Err(self.invalid(format!(
521                "component `{}` is referenced more than once in the physical layout tree",
522                component.id
523            )));
524        }
525        Ok(component)
526    }
527
528    fn validate_storage(
529        &self,
530        component: &WeightComponentSpec,
531        semantic_dimensions: &[u64],
532        storage: &PhysicalStorageLayout,
533    ) -> Result<(), VNextError> {
534        if semantic_dimensions.is_empty()
535            || semantic_dimensions.iter().any(|extent| *extent == 0)
536            || checked_elements(semantic_dimensions).is_none()
537        {
538            return Err(self.invalid(format!(
539                "component `{}` has an invalid or overflowing semantic shape",
540                component.id
541            )));
542        }
543        let raw_elements = checked_elements(&component.dimensions).ok_or_else(|| {
544            self.invalid(format!(
545                "component `{}` raw storage shape overflows u64",
546                component.id
547            ))
548        })?;
549        match storage {
550            PhysicalStorageLayout::Contiguous { padding } => {
551                let padded = self.resolve_padding(semantic_dimensions, padding)?;
552                if component.dimensions != padded {
553                    return Err(self.invalid(format!(
554                        "component `{}` contiguous shape {:?} differs from its explicit physical shape {:?}",
555                        component.id, padded, component.dimensions
556                    )));
557                }
558            }
559            PhysicalStorageLayout::Strided {
560                strides_in_elements,
561                padding,
562            } => {
563                let padded = self.resolve_padding(semantic_dimensions, padding)?;
564                let span = self.checked_strided_span(&padded, strides_in_elements, 1)?;
565                if span != raw_elements {
566                    return Err(self.invalid(format!(
567                        "component `{}` strided span {span} differs from its raw storage element count {raw_elements}",
568                        component.id
569                    )));
570                }
571            }
572            PhysicalStorageLayout::Tiled {
573                tile_shape,
574                axis_order,
575                tile_strides_in_elements,
576                padding,
577            } => {
578                let rank = semantic_dimensions.len();
579                if tile_shape.len() != rank
580                    || tile_shape.iter().any(|extent| *extent == 0)
581                    || !is_axis_permutation(axis_order, rank)
582                    || tile_strides_in_elements.len() != rank
583                {
584                    return Err(self.invalid(format!(
585                        "component `{}` tile shape, axis order, or strides do not match rank",
586                        component.id
587                    )));
588                }
589                let padded = self.resolve_padding(semantic_dimensions, padding)?;
590                let minimal_padded = semantic_dimensions
591                    .iter()
592                    .zip(tile_shape)
593                    .map(|(extent, tile)| checked_round_up(*extent, *tile))
594                    .collect::<Option<Vec<_>>>()
595                    .ok_or_else(|| {
596                        self.invalid(format!(
597                            "component `{}` tile padding overflows u64",
598                            component.id
599                        ))
600                    })?;
601                match padding {
602                    PhysicalWeightPadding::Exact if minimal_padded != semantic_dimensions => {
603                        return Err(self.invalid(format!(
604                            "component `{}` needs tile padding but declares exact storage",
605                            component.id
606                        )));
607                    }
608                    PhysicalWeightPadding::ZeroFill { .. } if padded != minimal_padded => {
609                        return Err(self.invalid(format!(
610                            "component `{}` tiled zero-fill shape is not the unique minimal padded shape",
611                            component.id
612                        )));
613                    }
614                    _ => {}
615                }
616                let semantic_grid = padded
617                    .iter()
618                    .zip(tile_shape)
619                    .map(|(extent, tile)| extent / tile)
620                    .collect::<Vec<_>>();
621                let physical_grid = axis_order
622                    .iter()
623                    .map(|axis| semantic_grid[*axis as usize])
624                    .collect::<Vec<_>>();
625                let tile_elements = checked_elements(tile_shape).ok_or_else(|| {
626                    self.invalid(format!(
627                        "component `{}` tile size overflows u64",
628                        component.id
629                    ))
630                })?;
631                let span = self.checked_strided_span(
632                    &physical_grid,
633                    tile_strides_in_elements,
634                    tile_elements,
635                )?;
636                if span != raw_elements {
637                    return Err(self.invalid(format!(
638                        "component `{}` tiled span {span} differs from its raw storage element count {raw_elements}",
639                        component.id
640                    )));
641                }
642            }
643        }
644        Ok(())
645    }
646
647    fn resolve_padding(
648        &self,
649        semantic_dimensions: &[u64],
650        padding: &PhysicalWeightPadding,
651    ) -> Result<Vec<u64>, VNextError> {
652        match padding {
653            PhysicalWeightPadding::Exact => Ok(semantic_dimensions.to_vec()),
654            PhysicalWeightPadding::ZeroFill { padded_dimensions } => {
655                if padded_dimensions.len() != semantic_dimensions.len()
656                    || padded_dimensions.iter().any(|extent| *extent == 0)
657                    || padded_dimensions
658                        .iter()
659                        .zip(semantic_dimensions)
660                        .any(|(padded, semantic)| padded < semantic)
661                    || padded_dimensions == semantic_dimensions
662                    || checked_elements(padded_dimensions).is_none()
663                {
664                    return Err(self.invalid(
665                        "zero-fill padding must explicitly enlarge a valid shape without shrinking any axis",
666                    ));
667                }
668                Ok(padded_dimensions.clone())
669            }
670        }
671    }
672
673    fn checked_strided_span(
674        &self,
675        dimensions: &[u64],
676        strides: &[u64],
677        base_span: u64,
678    ) -> Result<u64, VNextError> {
679        if dimensions.is_empty()
680            || dimensions.len() != strides.len()
681            || dimensions.iter().any(|extent| *extent == 0)
682            || strides.iter().any(|stride| *stride == 0)
683            || base_span == 0
684        {
685            return Err(self.invalid("strided storage dimensions and strides are invalid"));
686        }
687        let mut axes = dimensions
688            .iter()
689            .copied()
690            .zip(strides.iter().copied())
691            .filter(|(extent, _)| *extent > 1)
692            .collect::<Vec<_>>();
693        axes.sort_by_key(|(_, stride)| *stride);
694        let mut span = base_span;
695        for (extent, stride) in axes {
696            if stride < span {
697                return Err(
698                    self.invalid("strided storage aliases coordinates or overlaps physical tiles")
699                );
700            }
701            span = extent
702                .checked_sub(1)
703                .and_then(|count| count.checked_mul(stride))
704                .and_then(|addition| span.checked_add(addition))
705                .ok_or_else(|| self.invalid("strided storage span overflows u64"))?;
706        }
707        Ok(span)
708    }
709
710    fn grouped_dimensions(
711        &self,
712        semantic_dimensions: &[u64],
713        padding: &PhysicalWeightPadding,
714        group_axis: usize,
715        group_size: u64,
716    ) -> Result<Vec<u64>, VNextError> {
717        let axis_extent = semantic_dimensions[group_axis];
718        let minimal_axis = checked_round_up(axis_extent, group_size)
719            .ok_or_else(|| self.invalid("quantization group padding overflows u64"))?;
720        match padding {
721            PhysicalWeightPadding::Exact => {
722                if minimal_axis != axis_extent {
723                    return Err(self.invalid(
724                        "quantization groups require padding but exact storage was declared",
725                    ));
726                }
727                Ok(semantic_dimensions.to_vec())
728            }
729            PhysicalWeightPadding::ZeroFill { padded_dimensions } => {
730                if minimal_axis == axis_extent
731                    || padded_dimensions.len() != semantic_dimensions.len()
732                    || padded_dimensions.iter().enumerate().any(|(axis, extent)| {
733                        if axis == group_axis {
734                            *extent != minimal_axis
735                        } else {
736                            *extent != semantic_dimensions[axis]
737                        }
738                    })
739                {
740                    return Err(self.invalid(
741                        "quantization zero-fill must pad only the group axis to its unique minimal extent",
742                    ));
743                }
744                checked_elements(padded_dimensions)
745                    .is_some()
746                    .then(|| padded_dimensions.clone())
747                    .ok_or_else(|| self.invalid("quantization padded shape overflows u64"))
748            }
749        }
750    }
751
752    fn validate_dense_values(
753        &mut self,
754        binding: &PhysicalWeightComponentBinding,
755        semantic_dimensions: &[u64],
756        logical_element_type: ElementType,
757        depth: usize,
758    ) -> Result<(), VNextError> {
759        let component = self.bind_component(
760            binding,
761            semantic_dimensions,
762            WeightComponentRole::Values,
763            depth,
764        )?;
765        if component.dense_element_type() != Some(logical_element_type) {
766            return Err(self.invalid(format!(
767                "values component `{}` dtype differs from the logical tensor",
768                component.id
769            )));
770        }
771        Ok(())
772    }
773
774    fn validate_axis_component(
775        &mut self,
776        axis_component: &AxisWeightComponent,
777        semantic_dimensions: &[u64],
778        expected_axis: usize,
779        role: WeightComponentRole,
780        allow_narrow_integer: bool,
781        depth: usize,
782    ) -> Result<(), VNextError> {
783        if axis_component.axis as usize != expected_axis {
784            return Err(self.invalid(format!(
785                "axis component `{}` targets axis {}, expected {expected_axis}",
786                axis_component.component.component_id, axis_component.axis
787            )));
788        }
789        let axis_shape = [semantic_dimensions[expected_axis]];
790        let component = self.bind_component(&axis_component.component, &axis_shape, role, depth)?;
791        let integer_type_valid = component.dense_element_type().is_some_and(|element_type| {
792            if allow_narrow_integer {
793                matches!(
794                    element_type,
795                    ElementType::U8 | ElementType::U32 | ElementType::I8 | ElementType::I32
796                )
797            } else {
798                matches!(element_type, ElementType::U32 | ElementType::I32)
799            }
800        });
801        if !integer_type_valid {
802            return Err(self.invalid(format!(
803                "axis component `{}` must use an integer encoding valid for {:?}",
804                component.id, role
805            )));
806        }
807        Ok(())
808    }
809
810    fn validate_layout(
811        &mut self,
812        layout: &PhysicalWeightLayout,
813        semantic_dimensions: &[u64],
814        logical_element_type: ElementType,
815        depth: usize,
816    ) -> Result<(), VNextError> {
817        self.visit_node(depth)?;
818        if semantic_dimensions.is_empty()
819            || semantic_dimensions.iter().any(|extent| *extent == 0)
820            || checked_elements(semantic_dimensions).is_none()
821        {
822            return Err(self.invalid("logical layout shape is empty, zero, or overflowing"));
823        }
824        match layout {
825            PhysicalWeightLayout::Dense { component_id } => {
826                let binding =
827                    PhysicalWeightComponentBinding::exact_contiguous(component_id.clone());
828                self.validate_dense_values(
829                    &binding,
830                    semantic_dimensions,
831                    logical_element_type,
832                    depth,
833                )?;
834            }
835            PhysicalWeightLayout::Stored { component } => {
836                self.validate_dense_values(
837                    component,
838                    semantic_dimensions,
839                    logical_element_type,
840                    depth,
841                )?;
842            }
843            PhysicalWeightLayout::Composite { parts } => {
844                if parts.is_empty() {
845                    return Err(self.invalid("composite layout has no parts"));
846                }
847                let rank = semantic_dimensions.len();
848                let mut covered_elements = 0_u64;
849                for (index, part) in parts.iter().enumerate() {
850                    if part.logical_offsets.len() != rank
851                        || part.extents.len() != rank
852                        || part.extents.iter().any(|extent| *extent == 0)
853                        || part
854                            .logical_offsets
855                            .iter()
856                            .zip(&part.extents)
857                            .zip(semantic_dimensions)
858                            .any(|((offset, extent), logical)| {
859                                offset.checked_add(*extent).is_none_or(|end| end > *logical)
860                            })
861                    {
862                        return Err(self.invalid(format!(
863                            "composite part {index} has invalid semantic offsets or extents"
864                        )));
865                    }
866                    for previous in &parts[..index] {
867                        let overlaps = part
868                            .logical_offsets
869                            .iter()
870                            .zip(&part.extents)
871                            .zip(previous.logical_offsets.iter().zip(&previous.extents))
872                            .all(|((offset, extent), (other_offset, other_extent))| {
873                                offset.checked_add(*extent).is_some_and(|end| {
874                                    other_offset.checked_add(*other_extent).is_some_and(
875                                        |other_end| *offset < other_end && *other_offset < end,
876                                    )
877                                })
878                            });
879                        if overlaps {
880                            return Err(self.invalid("composite semantic placements overlap"));
881                        }
882                    }
883                    let part_elements = checked_elements(&part.extents)
884                        .ok_or_else(|| self.invalid("composite part size overflows u64"))?;
885                    covered_elements = covered_elements
886                        .checked_add(part_elements)
887                        .ok_or_else(|| self.invalid("composite coverage overflows u64"))?;
888                    self.validate_layout(
889                        &part.layout,
890                        &part.extents,
891                        logical_element_type,
892                        depth + 1,
893                    )?;
894                }
895                if covered_elements != checked_elements(semantic_dimensions).unwrap() {
896                    return Err(self.invalid(
897                        "composite semantic placements do not cover the logical tensor exactly",
898                    ));
899                }
900            }
901            PhysicalWeightLayout::Quantized {
902                packed_values,
903                packed_dimensions,
904                scales,
905                zero_points,
906                axis_indices,
907                permutation,
908                codebook,
909                group_axis,
910                group_padding,
911            } => {
912                if !matches!(
913                    logical_element_type,
914                    ElementType::F16 | ElementType::Bf16 | ElementType::F32
915                ) {
916                    return Err(
917                        self.invalid("quantized logical weight dtype must be floating point")
918                    );
919                }
920                let axis = *group_axis as usize;
921                if axis >= semantic_dimensions.len() {
922                    return Err(self.invalid("quantization group axis is out of range"));
923                }
924                let quantization = {
925                    let component = self.component(&packed_values.component_id)?;
926                    let WeightEncoding::Quantized(spec) = &component.encoding else {
927                        return Err(self.invalid(
928                            "packed-values component does not carry a quantization spec",
929                        ));
930                    };
931                    spec.clone()
932                };
933                let group_size = quantization
934                    .grouping
935                    .resolved_size(semantic_dimensions[axis]);
936                let grouped_dimensions =
937                    self.grouped_dimensions(semantic_dimensions, group_padding, axis, group_size)?;
938                let packed_bytes = checked_elements(&grouped_dimensions)
939                    .and_then(|elements| {
940                        elements.checked_mul(u64::from(quantization.bits_per_weight))
941                    })
942                    .and_then(|bits| bits.checked_add(7))
943                    .map(|bits| bits / 8)
944                    .ok_or_else(|| self.invalid("packed-values size overflows u64"))?;
945                if checked_elements(packed_dimensions) != Some(packed_bytes) {
946                    return Err(self.invalid(format!(
947                        "packed-values semantic shape contains {} storage bytes, expected {packed_bytes}",
948                        checked_elements(packed_dimensions)
949                            .map_or_else(|| "an overflowing number of".to_owned(), |value| value.to_string())
950                    )));
951                }
952                let packed = self.bind_component(
953                    packed_values,
954                    packed_dimensions,
955                    WeightComponentRole::PackedValues,
956                    depth,
957                )?;
958                if packed.encoding != WeightEncoding::Quantized(quantization.clone()) {
959                    return Err(self.invalid(
960                        "packed-values encoding changed while validating the quantized tree",
961                    ));
962                }
963
964                let mut group_shape = grouped_dimensions;
965                group_shape[axis] /= group_size;
966                let scales_component =
967                    self.bind_component(scales, &group_shape, WeightComponentRole::Scales, depth)?;
968                if scales_component.dense_element_type() != Some(quantization.scale_type) {
969                    return Err(
970                        self.invalid("scale component dtype differs from the quantization spec")
971                    );
972                }
973                match (quantization.zero_point_type, zero_points) {
974                    (Some(expected_type), Some(binding)) => {
975                        let component = self.bind_component(
976                            binding,
977                            &group_shape,
978                            WeightComponentRole::ZeroPoints,
979                            depth,
980                        )?;
981                        if component.dense_element_type() != Some(expected_type) {
982                            return Err(self.invalid(
983                                "zero-point component dtype differs from the quantization spec",
984                            ));
985                        }
986                    }
987                    (None, None) => {}
988                    _ => {
989                        return Err(self.invalid(
990                            "zero-point component presence differs from the quantization spec",
991                        ));
992                    }
993                }
994                if let Some(axis_indices) = axis_indices {
995                    self.validate_axis_component(
996                        axis_indices,
997                        semantic_dimensions,
998                        axis,
999                        WeightComponentRole::Indices,
1000                        true,
1001                        depth,
1002                    )?;
1003                }
1004                if let Some(permutation) = permutation {
1005                    self.validate_axis_component(
1006                        permutation,
1007                        semantic_dimensions,
1008                        axis,
1009                        WeightComponentRole::Permutation,
1010                        false,
1011                        depth,
1012                    )?;
1013                }
1014                if let Some(codebook) = codebook {
1015                    let entries = 1_u64
1016                        .checked_shl(u32::from(quantization.bits_per_weight))
1017                        .ok_or_else(|| self.invalid("codebook size overflows u64"))?;
1018                    let component = self.bind_component(
1019                        codebook,
1020                        &[entries],
1021                        WeightComponentRole::Codebook,
1022                        depth,
1023                    )?;
1024                    if component.dense_element_type() != Some(logical_element_type) {
1025                        return Err(
1026                            self.invalid("codebook dtype differs from the logical tensor dtype")
1027                        );
1028                    }
1029                }
1030            }
1031            PhysicalWeightLayout::BlockQuantized {
1032                blocks,
1033                block_axis,
1034                block_padding,
1035            } => {
1036                if !matches!(
1037                    logical_element_type,
1038                    ElementType::F16 | ElementType::Bf16 | ElementType::F32
1039                ) {
1040                    return Err(
1041                        self.invalid("block-quantized logical weight dtype must be floating point")
1042                    );
1043                }
1044                let axis = *block_axis as usize;
1045                if axis >= semantic_dimensions.len() {
1046                    return Err(self.invalid("block quantization axis is out of range"));
1047                }
1048                let quantization = {
1049                    let component = self.component(&blocks.component_id)?;
1050                    let WeightEncoding::BlockQuantized(spec) = &component.encoding else {
1051                        return Err(self
1052                            .invalid("block component does not carry a block quantization spec"));
1053                    };
1054                    spec.clone()
1055                };
1056                let mut block_dimensions = self.grouped_dimensions(
1057                    semantic_dimensions,
1058                    block_padding,
1059                    axis,
1060                    u64::from(quantization.logical_values_per_block),
1061                )?;
1062                block_dimensions[axis] /= u64::from(quantization.logical_values_per_block);
1063                let component = self.bind_component(
1064                    blocks,
1065                    &block_dimensions,
1066                    WeightComponentRole::PackedValues,
1067                    depth,
1068                )?;
1069                if component.encoding != WeightEncoding::BlockQuantized(quantization) {
1070                    return Err(
1071                        self.invalid("block encoding changed while validating the physical layout")
1072                    );
1073                }
1074            }
1075            PhysicalWeightLayout::AxisReshapePermutation {
1076                values,
1077                axis,
1078                logical_offset,
1079                extent,
1080                reshape,
1081                stored_axis_order,
1082            } => {
1083                let axis = *axis as usize;
1084                let end = logical_offset.checked_add(*extent);
1085                let reshape_rank = reshape.len();
1086                let order_is_permutation = stored_axis_order.len() == reshape_rank
1087                    && stored_axis_order
1088                        .iter()
1089                        .all(|axis| (*axis as usize) < reshape_rank)
1090                    && stored_axis_order
1091                        .iter()
1092                        .copied()
1093                        .collect::<BTreeSet<_>>()
1094                        .len()
1095                        == reshape_rank;
1096                let order_is_identity = stored_axis_order
1097                    .iter()
1098                    .filter(|stored| reshape[**stored as usize] > 1)
1099                    .copied()
1100                    .eq(reshape
1101                        .iter()
1102                        .enumerate()
1103                        .filter_map(|(axis, extent)| (*extent > 1).then_some(axis as u32)));
1104                if axis >= semantic_dimensions.len()
1105                    || *extent == 0
1106                    || end.is_none_or(|end| end > semantic_dimensions[axis])
1107                    || reshape_rank < 2
1108                    || reshape.iter().any(|dimension| *dimension == 0)
1109                    || checked_elements(reshape) != Some(*extent)
1110                    || !order_is_permutation
1111                    || order_is_identity
1112                {
1113                    return Err(self.invalid(
1114                        "axis reshape permutation has invalid range, shape, or stored axis order",
1115                    ));
1116                }
1117                self.validate_layout(values, semantic_dimensions, logical_element_type, depth + 1)?;
1118            }
1119            PhysicalWeightLayout::Indexed {
1120                indices,
1121                values,
1122                source_axis_extent,
1123            } => {
1124                let axis = indices.axis as usize;
1125                if axis >= semantic_dimensions.len() || *source_axis_extent == 0 {
1126                    return Err(self.invalid("indexed layout axis or source extent is invalid"));
1127                }
1128                self.validate_axis_component(
1129                    indices,
1130                    semantic_dimensions,
1131                    axis,
1132                    WeightComponentRole::Indices,
1133                    true,
1134                    depth,
1135                )?;
1136                let mut source_dimensions = semantic_dimensions.to_vec();
1137                source_dimensions[axis] = *source_axis_extent;
1138                checked_elements(&source_dimensions)
1139                    .ok_or_else(|| self.invalid("indexed source semantic shape overflows u64"))?;
1140                self.validate_layout(values, &source_dimensions, logical_element_type, depth + 1)?;
1141            }
1142            PhysicalWeightLayout::ExpertStack {
1143                experts,
1144                expert_axis,
1145            } => {
1146                let axis = *expert_axis as usize;
1147                if axis >= semantic_dimensions.len() {
1148                    return Err(self.invalid("expert stack axis is out of range"));
1149                }
1150                let expected_count = usize::try_from(semantic_dimensions[axis]).map_err(|_| {
1151                    self.invalid("expert stack count does not fit the platform usize")
1152                })?;
1153                if experts.is_empty() || experts.len() != expected_count {
1154                    return Err(self
1155                        .invalid("expert stack child count differs from its logical expert axis"));
1156                }
1157                let mut expert_dimensions = semantic_dimensions.to_vec();
1158                expert_dimensions.remove(axis);
1159                if expert_dimensions.is_empty() {
1160                    return Err(self.invalid(
1161                        "expert stack children must retain at least one tensor dimension",
1162                    ));
1163                }
1164                for expert in experts {
1165                    self.validate_layout(
1166                        expert,
1167                        &expert_dimensions,
1168                        logical_element_type,
1169                        depth + 1,
1170                    )?;
1171                }
1172            }
1173        }
1174        Ok(())
1175    }
1176}
1177
1178fn is_axis_permutation(axis_order: &[u32], rank: usize) -> bool {
1179    axis_order.len() == rank
1180        && axis_order.iter().all(|axis| (*axis as usize) < rank)
1181        && axis_order.iter().copied().collect::<BTreeSet<_>>().len() == rank
1182}
1183
1184fn checked_round_up(extent: u64, multiple: u64) -> Option<u64> {
1185    if extent == 0 || multiple == 0 {
1186        return None;
1187    }
1188    extent
1189        .checked_add(multiple.checked_sub(1)?)
1190        .map(|rounded| rounded / multiple * multiple)
1191}
1192
1193#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1194pub struct ProgramTensorSpec {
1195    pub dimensions: Vec<u64>,
1196    pub element_type: ElementType,
1197    pub layout: ResolvedTensorLayout,
1198}
1199
1200impl ProgramTensorSpec {
1201    pub fn validate(&self, field: &str) -> Result<(), VNextError> {
1202        super::ResolvedTensorSpec::new(
1203            self.dimensions.clone(),
1204            self.element_type,
1205            self.layout.clone(),
1206        )
1207        .map(|_| ())
1208        .map_err(|error| VNextError::InvalidExecutionPlan {
1209            reason: format!("{field} is invalid: {error}"),
1210        })
1211    }
1212
1213    pub fn byte_len(&self) -> Result<u64, VNextError> {
1214        checked_elements(&self.dimensions)
1215            .and_then(|elements| elements.checked_mul(self.element_type.size_bytes()))
1216            .ok_or_else(|| VNextError::InvalidExecutionPlan {
1217                reason: "program tensor byte size overflows u64".to_owned(),
1218            })
1219    }
1220}
1221
1222#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1223pub struct WeightReference {
1224    pub weight_id: WeightId,
1225    pub value_id: ProgramValueId,
1226    pub tensor: ProgramTensorSpec,
1227}
1228
1229#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1230pub struct StateSpec {
1231    pub id: StateId,
1232    pub value_id: ProgramValueId,
1233    /// Logical tensor view consumed by one operation. It does not select a
1234    /// physical allocator, addressability profile, or backing pool.
1235    pub tensor: ProgramTensorSpec,
1236    pub lifetime: StateLifetime,
1237    pub capacity_demand: StateCapacityDemand,
1238    /// Initial contents required when a new logical state scope acquires
1239    /// physical backing. This is semantic model state, not an allocator hint.
1240    pub initialization: StateInitialization,
1241}
1242
1243#[derive(Deserialize)]
1244#[serde(deny_unknown_fields)]
1245struct StateSpecWire {
1246    id: StateId,
1247    value_id: ProgramValueId,
1248    tensor: ProgramTensorSpec,
1249    lifetime: StateLifetime,
1250    capacity_demand: StateCapacityDemand,
1251    initialization: StateInitialization,
1252}
1253
1254impl<'de> Deserialize<'de> for StateSpec {
1255    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1256    where
1257        D: Deserializer<'de>,
1258    {
1259        let wire = StateSpecWire::deserialize(deserializer)?;
1260        wire.tensor
1261            .validate("state_spec.tensor")
1262            .and_then(|()| wire.capacity_demand.validate(wire.tensor.byte_len()?))
1263            .map_err(serde::de::Error::custom)?;
1264        Ok(Self {
1265            id: wire.id,
1266            value_id: wire.value_id,
1267            tensor: wire.tensor,
1268            lifetime: wire.lifetime,
1269            capacity_demand: wire.capacity_demand,
1270            initialization: wire.initialization,
1271        })
1272    }
1273}
1274
1275#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
1276#[serde(rename_all = "snake_case")]
1277pub enum StateLifetime {
1278    Request,
1279    Sequence,
1280    Step,
1281}
1282
1283/// Backend-neutral capacity formula for semantic state. This deliberately says
1284/// nothing about pages, blocks, allocator kind, or provider-visible regions.
1285/// Concrete physical storage is selected only while building an execution
1286/// plan from provider requirements, runtime offers, and typed policy.
1287#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1288#[serde(rename_all = "snake_case")]
1289pub enum StateCapacityDemand {
1290    FixedPerScope,
1291    TokenScaled {
1292        bytes_per_token: u64,
1293        maximum_tokens: u64,
1294    },
1295}
1296
1297#[derive(Deserialize)]
1298#[serde(rename_all = "snake_case", deny_unknown_fields)]
1299enum StateCapacityDemandWire {
1300    FixedPerScope,
1301    TokenScaled {
1302        bytes_per_token: u64,
1303        maximum_tokens: u64,
1304    },
1305}
1306
1307impl<'de> Deserialize<'de> for StateCapacityDemand {
1308    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1309    where
1310        D: Deserializer<'de>,
1311    {
1312        let demand = match StateCapacityDemandWire::deserialize(deserializer)? {
1313            StateCapacityDemandWire::FixedPerScope => Self::FixedPerScope,
1314            StateCapacityDemandWire::TokenScaled {
1315                bytes_per_token,
1316                maximum_tokens,
1317            } => Self::TokenScaled {
1318                bytes_per_token,
1319                maximum_tokens,
1320            },
1321        };
1322        demand.validate(1).map_err(serde::de::Error::custom)?;
1323        Ok(demand)
1324    }
1325}
1326
1327impl StateCapacityDemand {
1328    pub fn validate(self, tensor_minimum_bytes: u64) -> Result<(), VNextError> {
1329        let valid = match self {
1330            Self::FixedPerScope => tensor_minimum_bytes > 0,
1331            Self::TokenScaled {
1332                bytes_per_token,
1333                maximum_tokens,
1334            } => {
1335                bytes_per_token >= tensor_minimum_bytes
1336                    && maximum_tokens > 0
1337                    && bytes_per_token.checked_mul(maximum_tokens).is_some()
1338            }
1339        };
1340        if !valid {
1341            return Err(VNextError::InvalidExecutionPlan {
1342                reason: "state resource demand is zero, smaller than its tensor, or overflows u64"
1343                    .to_owned(),
1344            });
1345        }
1346        Ok(())
1347    }
1348
1349    pub fn minimum_bytes(self, tensor_minimum_bytes: u64) -> Result<u64, VNextError> {
1350        self.validate(tensor_minimum_bytes)?;
1351        Ok(match self {
1352            Self::FixedPerScope => tensor_minimum_bytes,
1353            Self::TokenScaled {
1354                bytes_per_token, ..
1355            } => bytes_per_token,
1356        })
1357    }
1358
1359    pub fn theoretical_bytes(self, tensor_minimum_bytes: u64) -> Result<u64, VNextError> {
1360        self.validate(tensor_minimum_bytes)?;
1361        match self {
1362            Self::FixedPerScope => Ok(tensor_minimum_bytes),
1363            Self::TokenScaled {
1364                bytes_per_token,
1365                maximum_tokens,
1366            } => bytes_per_token.checked_mul(maximum_tokens).ok_or_else(|| {
1367                VNextError::InvalidExecutionPlan {
1368                    reason: "token-scaled state demand overflows u64".to_owned(),
1369                }
1370            }),
1371        }
1372    }
1373}
1374
1375#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1376#[serde(rename_all = "snake_case")]
1377pub enum ProgramNodeWorkSpec {
1378    Fixed,
1379    Tokens { value_id: ProgramValueId, axis: u32 },
1380}
1381
1382impl ProgramNodeWorkSpec {
1383    pub fn tokens(value_id: ProgramValueId, axis: u32) -> Self {
1384        Self::Tokens { value_id, axis }
1385    }
1386}
1387
1388#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1389pub struct ProgramNode {
1390    pub id: NodeId,
1391    pub operation_id: OperationId,
1392    pub required_version: ContractVersion,
1393    pub work: ProgramNodeWorkSpec,
1394    pub inputs: Vec<ProgramValueId>,
1395    pub outputs: Vec<ProgramValueId>,
1396    pub attributes: BTreeMap<AttributeId, SemanticValue>,
1397}
1398
1399#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1400pub struct ProgramBlock {
1401    pub id: String,
1402    pub nodes: Vec<ProgramNode>,
1403}
1404
1405/// Backend-free semantic program for a model family.
1406#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1407pub struct ModelProgram {
1408    family_id: ModelFamilyId,
1409    inputs: Vec<ProgramValueId>,
1410    blocks: Vec<ProgramBlock>,
1411    states: Vec<StateSpec>,
1412    weights: Vec<WeightReference>,
1413    outputs: Vec<ProgramValueId>,
1414}
1415
1416#[derive(Deserialize)]
1417#[serde(deny_unknown_fields)]
1418struct ModelProgramWire {
1419    family_id: ModelFamilyId,
1420    inputs: Vec<ProgramValueId>,
1421    blocks: Vec<ProgramBlock>,
1422    states: Vec<StateSpec>,
1423    weights: Vec<WeightReference>,
1424    outputs: Vec<ProgramValueId>,
1425}
1426
1427impl<'de> Deserialize<'de> for ModelProgram {
1428    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1429    where
1430        D: Deserializer<'de>,
1431    {
1432        let wire = ModelProgramWire::deserialize(deserializer)?;
1433        Self::new(
1434            wire.family_id,
1435            wire.inputs,
1436            wire.blocks,
1437            wire.states,
1438            wire.weights,
1439            wire.outputs,
1440        )
1441        .map_err(serde::de::Error::custom)
1442    }
1443}
1444
1445impl ModelProgram {
1446    pub fn new(
1447        family_id: ModelFamilyId,
1448        inputs: Vec<ProgramValueId>,
1449        blocks: Vec<ProgramBlock>,
1450        mut states: Vec<StateSpec>,
1451        mut weights: Vec<WeightReference>,
1452        outputs: Vec<ProgramValueId>,
1453    ) -> Result<Self, VNextError> {
1454        if blocks.is_empty() {
1455            return Err(VNextError::InvalidModelConfig {
1456                family_id: family_id.to_string(),
1457                field: "program.blocks".to_owned(),
1458                reason: "at least one block is required".to_owned(),
1459            });
1460        }
1461        let mut known_values = BTreeSet::new();
1462        if inputs.is_empty()
1463            || inputs
1464                .iter()
1465                .any(|input| !known_values.insert(input.clone()))
1466        {
1467            return Err(VNextError::InvalidModelConfig {
1468                family_id: family_id.to_string(),
1469                field: "program.inputs".to_owned(),
1470                reason: "input identities must be non-empty and unique".to_owned(),
1471            });
1472        }
1473        let mut block_ids = BTreeSet::new();
1474        let mut node_ids = BTreeSet::new();
1475        for state in &states {
1476            let tensor_valid = state
1477                .tensor
1478                .validate(&format!("program.states.{}.tensor", state.id))
1479                .and_then(|()| state.capacity_demand.validate(state.tensor.byte_len()?));
1480            if tensor_valid.is_err() || !known_values.insert(state.value_id.clone()) {
1481                return Err(VNextError::InvalidModelConfig {
1482                    family_id: family_id.to_string(),
1483                    field: "program.states.value_id".to_owned(),
1484                    reason: format!("duplicate value `{}`", state.value_id),
1485                });
1486            }
1487        }
1488        let mut weight_ids = BTreeSet::new();
1489        for weight in &weights {
1490            if !weight_ids.insert(weight.weight_id.clone())
1491                || !known_values.insert(weight.value_id.clone())
1492            {
1493                return Err(VNextError::InvalidModelConfig {
1494                    family_id: family_id.to_string(),
1495                    field: "program.weights".to_owned(),
1496                    reason: format!(
1497                        "duplicate weight `{}` or value `{}`",
1498                        weight.weight_id, weight.value_id
1499                    ),
1500                });
1501            }
1502            weight
1503                .tensor
1504                .validate(&format!("program.weights.{}.tensor", weight.weight_id))?;
1505        }
1506        for block in &blocks {
1507            if block.id.is_empty() || block.nodes.is_empty() || !block_ids.insert(block.id.clone())
1508            {
1509                return Err(VNextError::InvalidModelConfig {
1510                    family_id: family_id.to_string(),
1511                    field: "program.blocks.id".to_owned(),
1512                    reason: "block identities must be non-empty and unique".to_owned(),
1513                });
1514            }
1515            for node in &block.nodes {
1516                if node.required_version.major == 0 || node.outputs.is_empty() {
1517                    return Err(VNextError::InvalidModelConfig {
1518                        family_id: family_id.to_string(),
1519                        field: "program.nodes.contract".to_owned(),
1520                        reason: format!("node `{}` has an invalid version or no outputs", node.id),
1521                    });
1522                }
1523                for value in node.attributes.values() {
1524                    value.validate(&format!("program node `{}` attributes", node.id))?;
1525                }
1526                if !node_ids.insert(node.id.clone()) {
1527                    return Err(VNextError::InvalidModelConfig {
1528                        family_id: family_id.to_string(),
1529                        field: "program.nodes.id".to_owned(),
1530                        reason: format!("duplicate node `{}`", node.id),
1531                    });
1532                }
1533                if let ProgramNodeWorkSpec::Tokens { value_id, .. } = &node.work {
1534                    let source_count = node
1535                        .inputs
1536                        .iter()
1537                        .chain(&node.outputs)
1538                        .filter(|candidate| *candidate == value_id)
1539                        .count();
1540                    let is_state_or_weight = states.iter().any(|state| state.value_id == *value_id)
1541                        || weights.iter().any(|weight| weight.value_id == *value_id);
1542                    if source_count != 1 || is_state_or_weight {
1543                        return Err(VNextError::InvalidModelConfig {
1544                            family_id: family_id.to_string(),
1545                            field: "program.nodes.work".to_owned(),
1546                            reason: format!(
1547                                "node `{}` token work source must identify one activation binding",
1548                                node.id
1549                            ),
1550                        });
1551                    }
1552                }
1553                if node
1554                    .inputs
1555                    .iter()
1556                    .any(|input| !known_values.contains(input))
1557                {
1558                    return Err(VNextError::InvalidModelConfig {
1559                        family_id: family_id.to_string(),
1560                        field: "program.nodes.inputs".to_owned(),
1561                        reason: format!("node `{}` references an unknown input", node.id),
1562                    });
1563                }
1564                for output in &node.outputs {
1565                    if !known_values.insert(output.clone()) {
1566                        return Err(VNextError::InvalidModelConfig {
1567                            family_id: family_id.to_string(),
1568                            field: "program.nodes.outputs".to_owned(),
1569                            reason: format!("value `{output}` has multiple producers"),
1570                        });
1571                    }
1572                }
1573            }
1574        }
1575        let mut state_ids = BTreeSet::new();
1576        if states
1577            .iter()
1578            .any(|state| !state_ids.insert(state.id.clone()))
1579        {
1580            return Err(VNextError::InvalidModelConfig {
1581                family_id: family_id.to_string(),
1582                field: "program.states.id".to_owned(),
1583                reason: "state identities must be unique".to_owned(),
1584            });
1585        }
1586        let mut output_ids = BTreeSet::new();
1587        if outputs.is_empty()
1588            || outputs
1589                .iter()
1590                .any(|output| !known_values.contains(output) || !output_ids.insert(output.clone()))
1591        {
1592            return Err(VNextError::InvalidModelConfig {
1593                family_id: family_id.to_string(),
1594                field: "program.outputs".to_owned(),
1595                reason: "program outputs must be non-empty, known, and unique".to_owned(),
1596            });
1597        }
1598        states.sort_by(|left, right| left.id.cmp(&right.id));
1599        weights.sort_by(|left, right| left.weight_id.cmp(&right.weight_id));
1600        Ok(Self {
1601            family_id,
1602            inputs,
1603            blocks,
1604            states,
1605            weights,
1606            outputs,
1607        })
1608    }
1609
1610    pub fn family_id(&self) -> &ModelFamilyId {
1611        &self.family_id
1612    }
1613
1614    pub fn inputs(&self) -> &[ProgramValueId] {
1615        &self.inputs
1616    }
1617
1618    pub fn blocks(&self) -> &[ProgramBlock] {
1619        &self.blocks
1620    }
1621
1622    pub fn states(&self) -> &[StateSpec] {
1623        &self.states
1624    }
1625
1626    pub fn weights(&self) -> &[WeightReference] {
1627        &self.weights
1628    }
1629
1630    pub fn outputs(&self) -> &[ProgramValueId] {
1631        &self.outputs
1632    }
1633
1634    pub fn fingerprint(&self) -> Result<String, VNextError> {
1635        let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
1636            context: "serialize model program",
1637            message: error.to_string(),
1638        })?;
1639        Ok(format!("{:x}", Sha256::digest(bytes)))
1640    }
1641}
1642
1643impl WeightSchema {
1644    pub fn validate_program_references(
1645        &self,
1646        family_id: &ModelFamilyId,
1647        program: &ModelProgram,
1648    ) -> Result<(), VNextError> {
1649        if program.family_id() != family_id {
1650            return Err(VNextError::InvalidModelConfig {
1651                family_id: family_id.to_string(),
1652                field: "program.family_id".to_owned(),
1653                reason: "program family does not match the weight schema owner".to_owned(),
1654            });
1655        }
1656        let schema_weights = self
1657            .tensors
1658            .iter()
1659            .map(|tensor| (&tensor.id, tensor.required))
1660            .collect::<BTreeMap<_, _>>();
1661        let referenced_weights = program
1662            .weights()
1663            .iter()
1664            .map(|reference| &reference.weight_id)
1665            .collect::<BTreeSet<_>>();
1666        if let Some(weight_id) = referenced_weights
1667            .iter()
1668            .find(|weight_id| !schema_weights.contains_key(**weight_id))
1669        {
1670            return Err(VNextError::InvalidModelConfig {
1671                family_id: family_id.to_string(),
1672                field: "program.weights".to_owned(),
1673                reason: format!("program references unknown weight `{weight_id}`"),
1674            });
1675        }
1676        if let Some(weight_id) = schema_weights.iter().find_map(|(weight_id, required)| {
1677            (*required && !referenced_weights.contains(weight_id)).then_some(*weight_id)
1678        }) {
1679            return Err(VNextError::InvalidModelConfig {
1680                family_id: family_id.to_string(),
1681                field: "program.weights".to_owned(),
1682                reason: format!("program does not reference required weight `{weight_id}`"),
1683            });
1684        }
1685        for reference in program.weights() {
1686            let tensor = self.tensor(&reference.weight_id).ok_or_else(|| {
1687                VNextError::InvalidModelConfig {
1688                    family_id: family_id.to_string(),
1689                    field: "program.weights".to_owned(),
1690                    reason: format!(
1691                        "program references unknown weight `{}`",
1692                        reference.weight_id
1693                    ),
1694                }
1695            })?;
1696            if reference.tensor.dimensions != tensor.dimensions
1697                || reference.tensor.element_type != tensor.logical_element_type
1698            {
1699                return Err(VNextError::InvalidModelConfig {
1700                    family_id: family_id.to_string(),
1701                    field: format!("program.weights.{}.tensor", reference.weight_id),
1702                    reason: "program value shape or dtype differs from the logical weight schema"
1703                        .to_owned(),
1704                });
1705            }
1706        }
1707        Ok(())
1708    }
1709}
1710
1711#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1712pub struct TemplateMetadata {
1713    pub template: String,
1714    pub source_file: String,
1715    pub sha256: String,
1716}
1717
1718#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
1719#[serde(rename_all = "snake_case")]
1720pub enum SpecialTokenRole {
1721    Bos,
1722    Eos,
1723    Pad,
1724    Stop,
1725}
1726
1727#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize)]
1728pub struct SpecialTokenCollision {
1729    first: SpecialTokenRole,
1730    second: SpecialTokenRole,
1731}
1732
1733#[derive(Deserialize)]
1734#[serde(deny_unknown_fields)]
1735struct SpecialTokenCollisionWire {
1736    first: SpecialTokenRole,
1737    second: SpecialTokenRole,
1738}
1739
1740impl<'de> Deserialize<'de> for SpecialTokenCollision {
1741    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1742    where
1743        D: Deserializer<'de>,
1744    {
1745        let wire = SpecialTokenCollisionWire::deserialize(deserializer)?;
1746        Self::new(wire.first, wire.second).map_err(serde::de::Error::custom)
1747    }
1748}
1749
1750impl SpecialTokenCollision {
1751    pub fn new(first: SpecialTokenRole, second: SpecialTokenRole) -> Result<Self, VNextError> {
1752        if first == second {
1753            return Err(VNextError::InvalidExecutionPlan {
1754                reason: "a special-token collision must name two different roles".to_owned(),
1755            });
1756        }
1757        let (first, second) = if first < second {
1758            (first, second)
1759        } else {
1760            (second, first)
1761        };
1762        Ok(Self { first, second })
1763    }
1764
1765    pub const fn first(&self) -> SpecialTokenRole {
1766        self.first
1767    }
1768
1769    pub const fn second(&self) -> SpecialTokenRole {
1770        self.second
1771    }
1772}
1773
1774#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1775pub struct SpecialTokenCollisionPolicy {
1776    allowed: BTreeSet<SpecialTokenCollision>,
1777}
1778
1779#[derive(Deserialize)]
1780#[serde(deny_unknown_fields)]
1781struct SpecialTokenCollisionPolicyWire {
1782    allowed: BTreeSet<SpecialTokenCollision>,
1783}
1784
1785impl<'de> Deserialize<'de> for SpecialTokenCollisionPolicy {
1786    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1787    where
1788        D: Deserializer<'de>,
1789    {
1790        let wire = SpecialTokenCollisionPolicyWire::deserialize(deserializer)?;
1791        Ok(Self::new(wire.allowed))
1792    }
1793}
1794
1795impl SpecialTokenCollisionPolicy {
1796    pub fn new(allowed: BTreeSet<SpecialTokenCollision>) -> Self {
1797        Self { allowed }
1798    }
1799
1800    pub fn require_distinct() -> Self {
1801        Self {
1802            allowed: BTreeSet::new(),
1803        }
1804    }
1805
1806    pub fn allows(&self, left: SpecialTokenRole, right: SpecialTokenRole) -> bool {
1807        SpecialTokenCollision::new(left, right)
1808            .is_ok_and(|collision| self.allowed.contains(&collision))
1809    }
1810
1811    pub fn allowed(&self) -> &BTreeSet<SpecialTokenCollision> {
1812        &self.allowed
1813    }
1814}
1815
1816#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1817pub struct SpecialTokenMetadata {
1818    pub bos_token_id: Option<u32>,
1819    pub eos_token_ids: BTreeSet<u32>,
1820    pub pad_token_id: Option<u32>,
1821    pub collision_policy: SpecialTokenCollisionPolicy,
1822}
1823
1824#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1825pub struct ModelSemanticMetadata {
1826    pub template: TemplateMetadata,
1827    pub special_tokens: SpecialTokenMetadata,
1828}
1829
1830/// Compile-time model family provider with a typed, validated configuration.
1831pub trait ModelFamilyProvider: Send + Sync {
1832    type Config: Clone + Send + Sync + Serialize + DeserializeOwned + 'static;
1833
1834    fn family_id(&self) -> &ModelFamilyId;
1835
1836    fn external_metadata_ids(&self) -> BTreeSet<ExternalModelMetadataId>;
1837
1838    fn validate_config_identity(
1839        &self,
1840        raw: &serde_json::Value,
1841        config: &Self::Config,
1842    ) -> Result<(), VNextError>;
1843
1844    /// Returns the exact external metadata identity represented by `config`.
1845    /// Every provider must make this selection explicit; core never assumes a
1846    /// singleton catalog row is the intended typed identity.
1847    fn validated_external_metadata_id(
1848        &self,
1849        raw: &serde_json::Value,
1850        config: &Self::Config,
1851    ) -> Result<ExternalModelMetadataId, VNextError>;
1852
1853    fn parse_config(&self, raw: &serde_json::Value) -> Result<Self::Config, VNextError>;
1854
1855    fn weight_schema(&self, config: &Self::Config) -> Result<WeightSchema, VNextError>;
1856
1857    fn semantic_program(&self, config: &Self::Config) -> Result<ModelProgram, VNextError>;
1858
1859    fn semantic_metadata(&self, config: &Self::Config)
1860        -> Result<ModelSemanticMetadata, VNextError>;
1861}
1862
1863/// Maximum raw JSON bytes accepted before decoding a prepared family package.
1864pub const MAX_PREPARED_MODEL_FAMILY_WIRE_BYTES: usize = 16 * 1024 * 1024;
1865
1866#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1867pub struct PreparedModelFamily {
1868    family_id: ModelFamilyId,
1869    external_metadata_id: ExternalModelMetadataId,
1870    canonical_config: serde_json::Value,
1871    config_fingerprint: String,
1872    weight_schema: WeightSchema,
1873    program: ModelProgram,
1874    metadata: ModelSemanticMetadata,
1875}
1876
1877/// Serialized prepared packages are evidence, not trusted runtime objects.
1878/// Rehydration must resolve the typed provider again and reproduce every field.
1879#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1880pub struct UnvalidatedPreparedModelFamily {
1881    family_id: ModelFamilyId,
1882    external_metadata_id: ExternalModelMetadataId,
1883    canonical_config: serde_json::Value,
1884    config_fingerprint: String,
1885    weight_schema: WeightSchema,
1886    program: ModelProgram,
1887    metadata: ModelSemanticMetadata,
1888}
1889
1890/// Crate-private serde shape used by both the top-level decoder and nested
1891/// resolved-plan wire. Public code can only obtain the explicit unvalidated
1892/// package through a byte-bounded decoder.
1893#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1894pub(crate) struct PreparedModelFamilyWire {
1895    family_id: ModelFamilyId,
1896    external_metadata_id: ExternalModelMetadataId,
1897    canonical_config: serde_json::Value,
1898    config_fingerprint: String,
1899    weight_schema: WeightSchema,
1900    program: ModelProgram,
1901    metadata: ModelSemanticMetadata,
1902}
1903
1904#[derive(Deserialize, Serialize)]
1905#[serde(deny_unknown_fields)]
1906struct PreparedModelFamilyWireFields {
1907    family_id: ModelFamilyId,
1908    external_metadata_id: ExternalModelMetadataId,
1909    canonical_config: serde_json::Value,
1910    config_fingerprint: String,
1911    weight_schema: WeightSchema,
1912    program: ModelProgram,
1913    metadata: ModelSemanticMetadata,
1914}
1915
1916impl<'de> Deserialize<'de> for PreparedModelFamilyWire {
1917    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1918    where
1919        D: Deserializer<'de>,
1920    {
1921        let raw = serde_json::Value::deserialize(deserializer)?;
1922        let fields =
1923            PreparedModelFamilyWireFields::deserialize(&raw).map_err(serde::de::Error::custom)?;
1924        let canonical = serde_json::to_value(&fields).map_err(serde::de::Error::custom)?;
1925        if canonical != raw {
1926            return Err(serde::de::Error::custom(
1927                "prepared model family wire contains unknown or non-canonical nested fields",
1928            ));
1929        }
1930        Ok(Self {
1931            family_id: fields.family_id,
1932            external_metadata_id: fields.external_metadata_id,
1933            canonical_config: fields.canonical_config,
1934            config_fingerprint: fields.config_fingerprint,
1935            weight_schema: fields.weight_schema,
1936            program: fields.program,
1937            metadata: fields.metadata,
1938        })
1939    }
1940}
1941
1942impl From<PreparedModelFamilyWire> for UnvalidatedPreparedModelFamily {
1943    fn from(wire: PreparedModelFamilyWire) -> Self {
1944        Self {
1945            family_id: wire.family_id,
1946            external_metadata_id: wire.external_metadata_id,
1947            canonical_config: wire.canonical_config,
1948            config_fingerprint: wire.config_fingerprint,
1949            weight_schema: wire.weight_schema,
1950            program: wire.program,
1951            metadata: wire.metadata,
1952        }
1953    }
1954}
1955
1956impl UnvalidatedPreparedModelFamily {
1957    pub fn revalidate(
1958        self,
1959        registry: &dyn ModelFamilyRegistry,
1960    ) -> Result<PreparedModelFamily, VNextError> {
1961        let registration = registry.resolve(&self.family_id)?;
1962        if registration.family_id() != &self.family_id {
1963            return Err(VNextError::InvalidModelConfig {
1964                family_id: self.family_id.to_string(),
1965                field: "registration.family_id".to_owned(),
1966                reason: "registry returned a registration for a different family".to_owned(),
1967            });
1968        }
1969        let metadata_registration = registry.resolve_external(&self.external_metadata_id)?;
1970        if !std::ptr::eq(registration, metadata_registration) {
1971            return Err(VNextError::InvalidModelConfig {
1972                family_id: self.family_id.to_string(),
1973                field: "external_metadata_id".to_owned(),
1974                reason: "external metadata identity resolves to a different family registration"
1975                    .to_owned(),
1976            });
1977        }
1978        let rebuilt = registration.prepare(&self.canonical_config)?;
1979        let exact_match = rebuilt.family_id == self.family_id
1980            && rebuilt.external_metadata_id == self.external_metadata_id
1981            && rebuilt.canonical_config == self.canonical_config
1982            && rebuilt.config_fingerprint == self.config_fingerprint
1983            && rebuilt.weight_schema == self.weight_schema
1984            && rebuilt.program == self.program
1985            && rebuilt.metadata == self.metadata;
1986        if !exact_match {
1987            return Err(VNextError::InvalidModelConfig {
1988                family_id: self.family_id.to_string(),
1989                field: "prepared_package".to_owned(),
1990                reason: "serialized package differs from the typed provider reconstruction"
1991                    .to_owned(),
1992            });
1993        }
1994        Ok(rebuilt)
1995    }
1996}
1997
1998impl PreparedModelFamily {
1999    fn from_canonical_config(
2000        family_id: ModelFamilyId,
2001        external_metadata_id: ExternalModelMetadataId,
2002        canonical_config: serde_json::Value,
2003        mut weight_schema: WeightSchema,
2004        program: ModelProgram,
2005        metadata: ModelSemanticMetadata,
2006    ) -> Result<Self, VNextError> {
2007        if !canonical_config.is_object()
2008            || canonicalize_json(canonical_config.clone()) != canonical_config
2009        {
2010            return Err(VNextError::InvalidModelConfig {
2011                family_id: family_id.to_string(),
2012                field: "config".to_owned(),
2013                reason: "prepared config must be a canonical JSON object".to_owned(),
2014            });
2015        }
2016        let config_bytes =
2017            serde_json::to_vec(&canonical_config).map_err(|error| VNextError::Serialization {
2018                context: "serialize canonical model family config",
2019                message: error.to_string(),
2020            })?;
2021        let config_fingerprint = format!("{:x}", Sha256::digest(config_bytes));
2022        weight_schema.validate(&family_id)?;
2023        weight_schema.normalize();
2024        weight_schema.validate(&family_id)?;
2025        if program.family_id() != &family_id {
2026            return Err(VNextError::InvalidModelConfig {
2027                family_id: family_id.to_string(),
2028                field: "program.family_id".to_owned(),
2029                reason: "program family does not match prepared family".to_owned(),
2030            });
2031        }
2032        weight_schema.validate_program_references(&family_id, &program)?;
2033        Self::validate_metadata(&family_id, &metadata)?;
2034        Ok(Self {
2035            family_id,
2036            external_metadata_id,
2037            canonical_config,
2038            config_fingerprint,
2039            weight_schema,
2040            program,
2041            metadata,
2042        })
2043    }
2044
2045    fn validate_metadata(
2046        family_id: &ModelFamilyId,
2047        metadata: &ModelSemanticMetadata,
2048    ) -> Result<(), VNextError> {
2049        let source = metadata.template.source_file.as_str();
2050        let valid_source = !source.is_empty()
2051            && !source.starts_with('/')
2052            && !source.contains('\\')
2053            && source
2054                .split('/')
2055                .all(|component| !matches!(component, "" | "." | ".."));
2056        if metadata.template.template.is_empty()
2057            || !valid_source
2058            || !is_canonical_sha256(&metadata.template.sha256)
2059            || metadata.special_tokens.eos_token_ids.is_empty()
2060        {
2061            return Err(VNextError::InvalidModelConfig {
2062                family_id: family_id.to_string(),
2063                field: "semantic_metadata".to_owned(),
2064                reason: "template, source, checksum, and end tokens must be explicit and valid"
2065                    .to_owned(),
2066            });
2067        }
2068        Ok(())
2069    }
2070
2071    pub fn family_id(&self) -> &ModelFamilyId {
2072        &self.family_id
2073    }
2074
2075    pub fn external_metadata_id(&self) -> &ExternalModelMetadataId {
2076        &self.external_metadata_id
2077    }
2078
2079    pub fn canonical_config(&self) -> &serde_json::Value {
2080        &self.canonical_config
2081    }
2082
2083    pub fn config_fingerprint(&self) -> &str {
2084        &self.config_fingerprint
2085    }
2086
2087    pub fn weight_schema(&self) -> &WeightSchema {
2088        &self.weight_schema
2089    }
2090
2091    pub fn program(&self) -> &ModelProgram {
2092        &self.program
2093    }
2094
2095    pub fn metadata(&self) -> &ModelSemanticMetadata {
2096        &self.metadata
2097    }
2098
2099    pub fn fingerprint(&self) -> Result<String, VNextError> {
2100        let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
2101            context: "serialize prepared model family",
2102            message: error.to_string(),
2103        })?;
2104        Ok(format!("{:x}", Sha256::digest(bytes)))
2105    }
2106
2107    pub fn decode_untrusted(bytes: &[u8]) -> Result<UnvalidatedPreparedModelFamily, VNextError> {
2108        if bytes.len() > MAX_PREPARED_MODEL_FAMILY_WIRE_BYTES {
2109            return Err(VNextError::Serialization {
2110                context: "decode untrusted prepared model family",
2111                message: format!(
2112                    "payload has {} bytes; maximum is {MAX_PREPARED_MODEL_FAMILY_WIRE_BYTES}",
2113                    bytes.len()
2114                ),
2115            });
2116        }
2117        serde_json::from_slice::<PreparedModelFamilyWire>(bytes)
2118            .map(Into::into)
2119            .map_err(|error| VNextError::Serialization {
2120                context: "decode untrusted prepared model family",
2121                message: error.to_string(),
2122            })
2123    }
2124
2125    pub fn from_json_validated(
2126        bytes: &[u8],
2127        registry: &dyn ModelFamilyRegistry,
2128    ) -> Result<Self, VNextError> {
2129        Self::decode_untrusted(bytes)?.revalidate(registry)
2130    }
2131}
2132
2133fn canonicalize_json(value: serde_json::Value) -> serde_json::Value {
2134    match value {
2135        serde_json::Value::Array(values) => {
2136            serde_json::Value::Array(values.into_iter().map(canonicalize_json).collect())
2137        }
2138        serde_json::Value::Object(values) => {
2139            let sorted = values
2140                .into_iter()
2141                .map(|(key, value)| (key, canonicalize_json(value)))
2142                .collect::<BTreeMap<_, _>>();
2143            serde_json::Value::Object(sorted.into_iter().collect())
2144        }
2145        other => other,
2146    }
2147}
2148
2149fn validate_raw_config_consumed(
2150    family_id: &ModelFamilyId,
2151    raw: &serde_json::Value,
2152    typed: &serde_json::Value,
2153) -> Result<(), VNextError> {
2154    // Typed serialization may add explicit provider defaults. Every caller-
2155    // supplied value must still survive at the same path with the same JSON
2156    // value; serde's default unknown-field behavior cannot erase input here.
2157    fn walk(raw: &serde_json::Value, typed: &serde_json::Value, path: &str) -> Option<String> {
2158        match (raw, typed) {
2159            (serde_json::Value::Object(raw), serde_json::Value::Object(typed)) => {
2160                for (key, raw_value) in raw {
2161                    let next = if path.is_empty() {
2162                        format!("/{key}")
2163                    } else {
2164                        format!("{path}/{key}")
2165                    };
2166                    let Some(typed_value) = typed.get(key) else {
2167                        return Some(next);
2168                    };
2169                    if let Some(rejected) = walk(raw_value, typed_value, &next) {
2170                        return Some(rejected);
2171                    }
2172                }
2173                None
2174            }
2175            (serde_json::Value::Array(raw), serde_json::Value::Array(typed))
2176                if raw.len() == typed.len() =>
2177            {
2178                raw.iter()
2179                    .zip(typed)
2180                    .enumerate()
2181                    .find_map(|(index, (raw, typed))| walk(raw, typed, &format!("{path}/{index}")))
2182            }
2183            _ if raw == typed => None,
2184            _ => Some(path.to_owned()),
2185        }
2186    }
2187
2188    if !raw.is_object() || !typed.is_object() {
2189        return Err(VNextError::InvalidModelConfig {
2190            family_id: family_id.to_string(),
2191            field: "config".to_owned(),
2192            reason: "raw and typed model configurations must be JSON objects".to_owned(),
2193        });
2194    }
2195    if let Some(path) = walk(raw, typed, "") {
2196        return Err(VNextError::InvalidModelConfig {
2197            family_id: family_id.to_string(),
2198            field: if path.is_empty() {
2199                "config".to_owned()
2200            } else {
2201                path
2202            },
2203            reason: "raw configuration field was ignored or changed by typed parsing".to_owned(),
2204        });
2205    }
2206    Ok(())
2207}
2208
2209fn is_canonical_sha256(value: &str) -> bool {
2210    value.len() == 64
2211        && value
2212            .bytes()
2213            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
2214}
2215
2216/// Object-safe loading-time trampoline for a heterogeneous family catalog.
2217/// The raw JSON exists only at the configuration boundary; a typed provider
2218/// validates it before producing the backend-free package.
2219pub trait ModelFamilyRegistration: Send + Sync {
2220    fn family_id(&self) -> &ModelFamilyId;
2221
2222    fn external_metadata_ids(&self) -> BTreeSet<ExternalModelMetadataId>;
2223
2224    fn prepare(&self, raw_config: &serde_json::Value) -> Result<PreparedModelFamily, VNextError>;
2225}
2226
2227pub struct TypedFamilyRegistration<P> {
2228    provider: P,
2229}
2230
2231impl<P> TypedFamilyRegistration<P> {
2232    pub fn new(provider: P) -> Self {
2233        Self { provider }
2234    }
2235}
2236
2237impl<P: ModelFamilyProvider> ModelFamilyRegistration for TypedFamilyRegistration<P> {
2238    fn family_id(&self) -> &ModelFamilyId {
2239        self.provider.family_id()
2240    }
2241
2242    fn external_metadata_ids(&self) -> BTreeSet<ExternalModelMetadataId> {
2243        self.provider.external_metadata_ids()
2244    }
2245
2246    fn prepare(&self, raw_config: &serde_json::Value) -> Result<PreparedModelFamily, VNextError> {
2247        let external_metadata_ids = self.provider.external_metadata_ids();
2248        if external_metadata_ids.is_empty() {
2249            return Err(VNextError::InvalidModelConfig {
2250                family_id: self.provider.family_id().to_string(),
2251                field: "external_metadata_ids".to_owned(),
2252                reason: "model family must declare at least one external metadata identity"
2253                    .to_owned(),
2254            });
2255        }
2256        let config = self.provider.parse_config(raw_config)?;
2257        let external_metadata_id = self
2258            .provider
2259            .validated_external_metadata_id(raw_config, &config)?;
2260        if !external_metadata_ids.contains(&external_metadata_id) {
2261            return Err(VNextError::InvalidModelConfig {
2262                family_id: self.provider.family_id().to_string(),
2263                field: "external_metadata_id".to_owned(),
2264                reason: format!(
2265                    "provider selected undeclared external metadata identity `{external_metadata_id}`"
2266                ),
2267            });
2268        }
2269        let typed_config = canonicalize_json(serde_json::to_value(&config).map_err(|error| {
2270            VNextError::Serialization {
2271                context: "serialize typed model configuration",
2272                message: error.to_string(),
2273            }
2274        })?);
2275        validate_raw_config_consumed(self.provider.family_id(), raw_config, &typed_config)?;
2276        let weight_schema = self.provider.weight_schema(&config)?;
2277        let program = self.provider.semantic_program(&config)?;
2278        let metadata = self.provider.semantic_metadata(&config)?;
2279        PreparedModelFamily::from_canonical_config(
2280            self.provider.family_id().clone(),
2281            external_metadata_id,
2282            typed_config,
2283            weight_schema,
2284            program,
2285            metadata,
2286        )
2287    }
2288}
2289
2290pub trait ModelFamilyRegistry: Send + Sync {
2291    /// Returns the complete trusted catalog. Core owns all identity lookup so a
2292    /// registry cannot silently fall back or hide ambiguous registrations.
2293    fn registrations(&self) -> Vec<&dyn ModelFamilyRegistration>;
2294}
2295
2296impl dyn ModelFamilyRegistry + '_ {
2297    pub fn resolve(
2298        &self,
2299        family_id: &ModelFamilyId,
2300    ) -> Result<&dyn ModelFamilyRegistration, VNextError> {
2301        let matches = self
2302            .registrations()
2303            .into_iter()
2304            .filter(|registration| registration.family_id() == family_id)
2305            .collect::<Vec<_>>();
2306        match matches.as_slice() {
2307            [] => Err(VNextError::UnknownModelFamily {
2308                family_id: family_id.to_string(),
2309            }),
2310            [registration] => Ok(*registration),
2311            _ => Err(VNextError::AmbiguousModelFamilyRegistration {
2312                identity_kind: "internal family",
2313                identity: family_id.to_string(),
2314                matches: matches.len(),
2315            }),
2316        }
2317    }
2318
2319    pub fn resolve_external(
2320        &self,
2321        metadata_id: &ExternalModelMetadataId,
2322    ) -> Result<&dyn ModelFamilyRegistration, VNextError> {
2323        let matches = self
2324            .registrations()
2325            .into_iter()
2326            .filter(|registration| registration.external_metadata_ids().contains(metadata_id))
2327            .collect::<Vec<_>>();
2328        match matches.as_slice() {
2329            [] => Err(VNextError::UnknownExternalModelMetadata {
2330                metadata_id: metadata_id.to_string(),
2331            }),
2332            [registration] => Ok(*registration),
2333            _ => Err(VNextError::AmbiguousModelFamilyRegistration {
2334                identity_kind: "external metadata",
2335                identity: metadata_id.to_string(),
2336                matches: matches.len(),
2337            }),
2338        }
2339    }
2340}
2341
2342#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
2343pub struct TokenizerDescriptor {
2344    pub tokenizer_id: TokenizerId,
2345    pub source_file: String,
2346    pub sha256: String,
2347    pub vocabulary_size: u64,
2348}