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 pub external_names: Vec<String>,
29 pub dimensions: Vec<u64>,
30 pub encoding: WeightEncoding,
31 pub required: bool,
32}
33
34impl WeightComponentSpec {
35 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 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 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 pub tensor: ProgramTensorSpec,
1236 pub lifetime: StateLifetime,
1237 pub capacity_demand: StateCapacityDemand,
1238 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#[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#[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
1830pub 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 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
1863pub 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#[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#[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 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
2216pub 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 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}